Commit 1fcc24b8 authored by Martin Bauer's avatar Martin Bauer
Browse files

pystencils.plot2d refactored: more documentation & tests

parent 703d4048
......@@ -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
......
......@@ -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")
......
"""
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)
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