pystencils.plot2d refactored: more documentation & tests

gpucuda/cudajit.py

@@ -30,7 +30,7 @@ def make_python_function(kernel_function_node, argument_dict=None):
     code += "#define RESTRICT __restrict__\n\n"
     code += str(generate_c(kernel_function_node))
 
-    mod = SourceModule(code, options=["-w", "-std=c++11"])
+    mod = SourceModule(code, options=["-w", "-std=c++11", "-Wno-deprecated-gpu-targets"])
     func = mod.get_function(kernel_function_node.function_name)
 
     parameters = kernel_function_node.parameters

gpucuda/indexing.py

@@ -69,7 +69,7 @@ class BlockIndexing(AbstractIndexing):
     """
 
     def __init__(self, field, iteration_slice=None,
-                 block_size=(256, 8, 1), permute_block_size_dependent_on_layout=True):
+                 block_size=(16, 16, 1), permute_block_size_dependent_on_layout=True):
         if field.spatial_dimensions > 3:
             raise NotImplementedError("This indexing scheme supports at most 3 spatial dimensions")
 

plot2d.py

+"""
+This module extends the pyplot module with functions to show scalar and vector fields in the usual
+simulation coordinate system (y-axis goes up), instead of the "image coordinate system" (y axis goes down) that
+matplotlib normally uses.
+"""
 from matplotlib.pyplot import *
 
 
-def vector_field(field, step=2, **kwargs):
-    """
-    Plot given vector field as quiver (arrow) plot.
+def vector_field(array, step=2, **kwargs):
+    """Plots given vector field as quiver (arrow) plot.
+
+    Args:
+        array: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
+               coordinate should have shape 2 and stores the 2 velocity components
+        step: plots only every steps's cell, increase the step for high resolution arrays
+        kwargs: keyword arguments passed to :func:`matplotlib.pyplot.quiver`
 
-    :param field: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
-                  coordinate should have shape 2 and stores the 2 velocity components
-    :param step: plots only every steps's cell
-    :param kwargs: keyword arguments passed to :func:`matplotlib.pyplot.quiver`
+    Returns:
+        quiver plot object
     """
-    vel_n = field.swapaxes(0, 1)
+    assert len(array.shape) == 3, "Wrong shape of array - did you forget to slice your 3D domain first?"
+    assert array.shape[2] == 2, "Last array dimension is expected to store 2D vectors"
+    vel_n = array.swapaxes(0, 1)
     res = quiver(vel_n[::step, ::step, 0], vel_n[::step, ::step, 1], **kwargs)
     axis('equal')
     return res
 
 
-def vector_field_magnitude(field, **kwargs):
-    """
-    Plots the magnitude of a vector field as colormap
-    :param field: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
-                  coordinate should have shape 2 and stores the 2 velocity components
-    :param kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+def vector_field_magnitude(array, **kwargs):
+    """Plots the magnitude of a vector field as colormap.
+
+    Args:
+        array: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
+               coordinate should have shape 2 and stores the 2 velocity components
+        kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+
+    Returns:
+        imshow object
     """
+    assert len(array.shape) == 3, "Wrong shape of array - did you forget to slice your 3D domain first?"
+    assert array.shape[2] in (2, 3), "Wrong size of the last coordinate. Has to be a 2D or 3D vector field."
     from numpy.linalg import norm
-    norm = norm(field, axis=2, ord=2)
-    if hasattr(field, 'mask'):
-        norm = np.ma.masked_array(norm, mask=field.mask[:, :, 0])
+    norm = norm(array, axis=2, ord=2)
+    if hasattr(array, 'mask'):
+        norm = np.ma.masked_array(norm, mask=array.mask[:, :, 0])
     return scalar_field(norm, **kwargs)
 
 
-def scalar_field(field, **kwargs):
-    """
-    Plots field values as colormap
+def scalar_field(array, **kwargs):
+    """Plots field values as colormap.
+
+    Works just as imshow, but uses coordinate system where second coordinate (y) points upwards.
+
+    Args:
+        array: two dimensional numpy array
+        kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
 
-    :param field: two dimensional numpy array
-    :param kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+    Returns:
+        imshow object
     """
     import numpy
-    field = numpy.swapaxes(field, 0, 1)
-    res = imshow(field, origin='lower', **kwargs)
+    array = numpy.swapaxes(array, 0, 1)
+    res = imshow(array, origin='lower', **kwargs)
     axis('equal')
     return res
 
 
-def scalar_field_alpha_value(field, color, clip=False, **kwargs):
+def scalar_field_alpha_value(array, color, clip=False, **kwargs):
+    """Plots an image with same color everywhere, using the array values as transparency.
+
+    Array is supposed to have values between 0 and 1 (if this is not the case it is normalized).
+    An image is plotted that has the same color everywhere, the passed array determines the transparency.
+    Regions where the array is 1  are fully opaque, areas with 0 are fully transparent.
+
+    Args:
+        array: 2D array with alpha values
+        color: fill color
+        clip: if True, all values in the array larger than 1 are set to 1, all values smaller than 0 are set to zero
+              if False, the array is linearly scaled to the [0, 1] interval
+        **kwargs: arguments passed to imshow
+
+    Returns:
+        imshow object
+    """
     import numpy
     import matplotlib
-    field = numpy.swapaxes(field, 0, 1)
-    color = matplotlib.colors.to_rgba(color)
-
-    field_to_plot = numpy.empty(field.shape + (4,))
-    for i in range(3):
-        field_to_plot[:, :, i] = color[i]
+    assert len(array.shape) == 2, "Wrong shape of array - did you forget to slice your 3D domain first?"
+    array = numpy.swapaxes(array, 0, 1)
 
     if clip:
-        normalized_field = field.copy()
+        normalized_field = array.copy()
         normalized_field[normalized_field < 0] = 0
         normalized_field[normalized_field > 1] = 1
     else:
-        minimum, maximum = numpy.min(field), numpy.max(field)
-        normalized_field = (field - minimum) / (maximum - minimum)
+        minimum, maximum = numpy.min(array), numpy.max(array)
+        normalized_field = (array - minimum) / (maximum - minimum)
+
+    color = matplotlib.colors.to_rgba(color)
+    field_to_plot = numpy.empty(array.shape + (4,))
+    # set the complete array to the color
+    for i in range(3):
+        field_to_plot[:, :, i] = color[i]
+    # only the alpha channel varies using the array values
     field_to_plot[:, :, 3] = normalized_field
 
     res = imshow(field_to_plot, origin='lower', **kwargs)
@@ -68,33 +108,85 @@ def scalar_field_alpha_value(field, color, clip=False, **kwargs):
     return res
 
 
-def scalar_field_contour(field, **kwargs):
-    field = np.swapaxes(field, 0, 1)
-    res = contour(field, **kwargs)
+def scalar_field_contour(array, **kwargs):
+    """Small wrapper around contour to transform the coordinate system.
+
+    For details see  :func:`matplotlib.pyplot.imshow`
+    """
+    array = np.swapaxes(array, 0, 1)
+    res = contour(array, **kwargs)
     axis('equal')
     return res
 
 
-def multiple_scalar_fields(field, **_):
-    sub_plots = field.shape[-1]
+def multiple_scalar_fields(array, **kwargs):
+    """Plots a 3D array by slicing the last dimension and creates on plot for each entry of the last dimension.
+
+    Args:
+        array: 3D array to plot.
+        **kwargs: passed along to imshow
+    """
+    assert len(array.shape) == 3
+    sub_plots = array.shape[-1]
     for i in range(sub_plots):
         subplot(1, sub_plots, i + 1)
         title(str(i))
-        scalar_field(field[..., i])
+        scalar_field(array[..., i], **kwargs)
         colorbar()
 
 
-def sympy_function(f, var, bounds, **kwargs):
+def sympy_function(expr, x_values=None, **kwargs):
+    """Plots the graph of a sympy term that depends on one symbol only.
+
+    Args:
+        expr: sympy term that depends on one symbol only, which is plotted on the x axis
+        x_values: describes sampling of x axis. Possible values are:
+                     * tuple of (start, stop) or (start, stop, nr_of_steps)
+                     * None, then start=0, stop=1, nr_of_steps=100
+                     * 1D numpy array with x values
+        **kwargs: passed on to :func:`matplotlib.pyplot.plot`
+
+    Returns:
+        plot object
+    """
     import sympy as sp
-    x_arr = np.linspace(bounds[0], bounds[1], 101)
-    y_arr = sp.lambdify(var, f)(x_arr)
-    plot(x_arr, y_arr, **kwargs)
+    if x_values is None:
+        x_arr = np.linspace(0, 1, 100)
+    elif type(x_values) is tuple:
+        x_arr = np.linspace(*x_values)
+    elif isinstance(x_values, np.ndarray):
+        assert len(x_values.shape) == 1
+        x_arr = x_values
+    else:
+        raise ValueError("Invalid value for parameter x_values")
+    symbols = expr.atoms(sp.Symbol)
+    assert len(symbols) == 1, "Sympy expression may only depend on one variable only. Depends on " + str(symbols)
+    y_arr = sp.lambdify(symbols.pop(), expr)(x_arr)
+    return plot(x_arr, y_arr, **kwargs)
+
 
 # ------------------------------------------- Animations ---------------------------------------------------------------
 
 
-def vector_field_animation(run_function, step=2, rescale=True, plot_setup_function=lambda: None,
-                           plot_update_function=lambda: None, interval=30, frames=180, **kwargs):
+def vector_field_animation(run_function, step=2, rescale=True, plot_setup_function=lambda *_: None,
+                           plot_update_function=lambda *_: None, interval=200, frames=180, **kwargs):
+    """Creates a matplotlib animation of a vector field using a quiver plot.
+
+    Args:
+        run_function: callable without arguments, returning a 2D vector field i.e. numpy array with len(shape)==3
+        step: see documentation of vector_field function
+        rescale: if True, the length of the arrows is rescaled in every time step
+        plot_setup_function: optional callable with the quiver object as argument,
+                             that can be used to set up the plot (title, legend,..)
+        plot_update_function: optional callable with the quiver object as argument
+                              that is called of the quiver object was updated
+        interval: delay between frames in milliseconds (see matplotlib.FuncAnimation)
+        frames: how many frames should be generated, see matplotlib.FuncAnimation
+        **kwargs: passed to quiver plot
+
+    Returns:
+        matplotlib animation object
+    """
     import matplotlib.animation as animation
     from numpy.linalg import norm
 
@@ -108,7 +200,7 @@ def vector_field_animation(run_function, step=2, rescale=True, plot_setup_functi
             kwargs['scale'] = 1.0
 
     quiver_plot = vector_field(field, step=step, **kwargs)
-    plot_setup_function()
+    plot_setup_function(quiver_plot)
 
     def update_figure(*_):
         f = run_function()
@@ -117,14 +209,18 @@ def vector_field_animation(run_function, step=2, rescale=True, plot_setup_functi
             f = f / np.max(norm(f, axis=2, ord=2))
         u, v = f[::step, ::step, 0], f[::step, ::step, 1]
         quiver_plot.set_UVC(u, v)
-        plot_update_function()
+        plot_update_function(quiver_plot)
         return im,
 
     return animation.FuncAnimation(fig, update_figure, interval=interval, frames=frames)
 
 
-def vector_field_magnitude_animation(run_function, plot_setup_function=lambda: None,
-                                     plot_update_function=lambda: None, interval=30, frames=180, **kwargs):
+def vector_field_magnitude_animation(run_function, plot_setup_function=lambda *_: None,
+                                     plot_update_function=lambda *_: None, interval=30, frames=180, **kwargs):
+    """Animation of a vector field, showing the magnitude as colormap.
+
+    For arguments, see vector_field_animation
+    """
     import matplotlib.animation as animation
     from numpy.linalg import norm
 
@@ -132,7 +228,7 @@ def vector_field_magnitude_animation(run_function, plot_setup_function=lambda: N
     im = None
     field = run_function()
     im = vector_field_magnitude(field, **kwargs)
-    plot_setup_function()
+    plot_setup_function(im)
 
     def update_figure(*_):
         f = run_function()
@@ -141,7 +237,7 @@ def vector_field_magnitude_animation(run_function, plot_setup_function=lambda: N
             normed = np.ma.masked_array(normed, mask=f.mask[:, :, 0])
         normed = np.swapaxes(normed, 0, 1)
         im.set_array(normed)
-        plot_update_function()
+        plot_update_function(im)
         return im,
 
     return animation.FuncAnimation(fig, update_figure, interval=interval, frames=frames)