Commit 60d1b071 authored by Martin Bauer's avatar Martin Bauer
Browse files

More flexible plotting of field access (AA, EsoTwist..)

parent 81c6c262
......@@ -202,10 +202,10 @@ class EsoTwistEvenTimeStepAccessor(PdfFieldAccessor):
# -------------------------------------------- Visualization -----------------------------------------------------------
def visualize_field_mapping(axes, stencil, field_mapping, color='b'):
def visualize_field_mapping(axes, stencil, field_mapping, inverted=False, color='b'):
from lbmpy.plot import LbGrid
grid = LbGrid(3, 3)
grid.fill_with_default_arrows()
grid.fill_with_default_arrows(inverted=inverted)
for field_access, direction in zip(field_mapping, stencil):
field_position = stencil[field_access.index[0]]
neighbor = field_access.offsets
......@@ -214,7 +214,8 @@ def visualize_field_mapping(axes, stencil, field_mapping, color='b'):
grid.draw(axes)
def visualize_pdf_field_accessor(pdf_field_accessor, figure=None):
def visualize_pdf_field_accessor(pdf_field_accessor, title=True, read_plot_params={}, write_plot_params={},
figure=None):
from lbmpy.stencils import get_stencil
if figure is None:
......@@ -232,8 +233,14 @@ def visualize_pdf_field_accessor(pdf_field_accessor, figure=None):
ax_left = figure.add_subplot(1, 2, 1)
ax_right = figure.add_subplot(1, 2, 2)
visualize_field_mapping(ax_left, stencil, pre_collision_accesses, color='k')
visualize_field_mapping(ax_right, stencil, post_collision_accesses, color='r')
if 'color' not in read_plot_params:
read_plot_params['color'] = 'k'
if 'color' not in write_plot_params:
write_plot_params['color'] = 'r'
ax_left.set_title("Read")
ax_right.set_title("Write")
visualize_field_mapping(ax_left, stencil, pre_collision_accesses, **read_plot_params)
visualize_field_mapping(ax_right, stencil, post_collision_accesses, **write_plot_params)
if title:
ax_left.set_title("Read")
ax_right.set_title("Write")
......@@ -139,14 +139,17 @@ class LbGrid:
self.annotations[(cell, arrow_position)] = Text(arrow_end[0], arrow_end[1], annotation,
**annotation_kwargs)
def fill_with_default_arrows(self, **kwargs):
def fill_with_default_arrows(self, inverted=False, **kwargs):
"""Fills the complete grid with the default pdf arrows"""
for x in range(self._xCells):
for y in range(self._yCells):
for dx in [-1, 0, 1]:
for dy in [-1, 0, 1]:
kwargs.setdefault('color', '#bbbbbb')
self.add_arrow((x, y), (dx, dy), (dx, dy), **kwargs)
if not inverted:
self.add_arrow((x, y), (dx, dy), (dx, dy), **kwargs)
else:
self.add_arrow((x, y), (dx, dy), (-dx, -dy), **kwargs)
def draw(self, ax):
"""Draw the grid into a given matplotlib axes object"""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment