diff --git a/plot2d.py b/plot2d.py index bb0b18feb69f0a166f67964946742c57e3ccf775..e0935c831186e378413613221f7dad25a227c9da 100644 --- a/plot2d.py +++ b/plot2d.py @@ -167,6 +167,13 @@ def sympy_function(expr, x_values=None, **kwargs): # ------------------------------------------- Animations --------------------------------------------------------------- +def __scale_array(arr): + from numpy.linalg import norm + norm_arr = norm(arr, axis=2, ord=2) + if isinstance(arr, np.ma.MaskedArray): + norm_arr = np.ma.masked_array(norm_arr, arr.mask[..., 0]) + return arr / norm_arr.max() + 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): @@ -188,15 +195,13 @@ def vector_field_animation(run_function, step=2, rescale=True, plot_setup_functi matplotlib animation object """ import matplotlib.animation as animation - from numpy.linalg import norm fig = gcf() im = None field = run_function() if rescale: - max_norm = np.max(norm(field, axis=2, ord=2)) - field = field / max_norm - kwargs.setdefault('scale', 1 / step) + field = __scale_array(field) + kwargs.setdefault('scale', 0.6) kwargs.setdefault('angles', 'xy') kwargs.setdefault('scale_units', 'xy') @@ -207,7 +212,7 @@ def vector_field_animation(run_function, step=2, rescale=True, plot_setup_functi f = run_function() f = np.swapaxes(f, 0, 1) if rescale: - f = f / np.max(norm(f, axis=2, ord=2)) + f = __scale_array(f) u, v = f[::step, ::step, 0], f[::step, ::step, 1] quiver_plot.set_UVC(u, v) plot_update_function(quiver_plot) @@ -216,7 +221,7 @@ def vector_field_animation(run_function, step=2, rescale=True, plot_setup_functi return animation.FuncAnimation(fig, update_figure, interval=interval, frames=frames) -def vector_field_magnitude_animation(run_function, plot_setup_function=lambda *_: None, +def vector_field_magnitude_animation(run_function, plot_setup_function=lambda *_: None, rescale=False, plot_update_function=lambda *_: None, interval=30, frames=180, **kwargs): """Animation of a vector field, showing the magnitude as colormap. @@ -228,11 +233,15 @@ def vector_field_magnitude_animation(run_function, plot_setup_function=lambda *_ fig = gcf() im = None field = run_function() + if rescale: + field = __scale_array(field) im = vector_field_magnitude(field, **kwargs) plot_setup_function(im) def update_figure(*_): f = run_function() + if rescale: + f = __scale_array(f) normed = norm(f, axis=2, ord=2) if hasattr(f, 'mask'): normed = np.ma.masked_array(normed, mask=f.mask[:, :, 0])