Commit 82bfa691 authored by Martin Bauer's avatar Martin Bauer
Browse files

Fixes in vector field animation rescaling

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