table = NDTable(Z, (x, y))
xi = yi = np.linspace(-6, 6, 200)
XI, YI = np.meshgrid(xi, yi, indexing='ij')
figure, axes = plt.subplots(ncols=2, nrows=2, sharex=True, sharey=True)
figure.canvas.set_window_title('Inter- & Extrapolation Methods')
figure.set_facecolor('white')
axes = axes.flatten()
ax = axes[0]
ax.set_title('original')
im = NonUniformImage(ax)
im.set_data(x, y, Z)
im.set_extent((-3, 3, -3, 3))
ax.images.append(im)
ax.set_xlim([-6, 6])
ax.set_ylim([-6, 6])
methods = [('nearest', 'hold'), ('linear', 'linear'), ('akima', 'linear')]
for ax, method in zip(axes[1:], methods):
ZI = table.evaluate((XI, YI), interp=method[0], extrap=method[1])
ax.set_title("interp='%s', extrap='%s'" % method)
im = NonUniformImage(ax)
im.set_data(xi, yi, ZI)
im.set_extent((-6, 6, -6, 6))
ax.images.append(im)
figure.tight_layout()
def animate_xz_density(trajectory,
xedges=None,
zedges=None,
figsize=(7, 7),
interval=1,
n_frames=10,
output_mode='animation',
axis_equal=True):
"""
Animates an density plot of a static simulation trajectory in a z-x projection. Still frames can also be rendered.
:param trajectory: Trajectory object with the particle trajectory data to be animated
:type trajectory: Trajectory
:param xedges: the edges of the bins of the density plot (2d histogram bins) in x direction,
if None the maximum extend is used with 50 bins, if a number n, the maximum extend is used with n bins
:type xedges: iterable or list / array or int
:param zedges: the edges of the bins of the density plot (2d histogram bins) in z direction,
if None the maximum extend is used with 50 bins, if a number n, the maximum extend is used with n bins
:type zedges: iterable or list / array or int
:param figsize: the figure size
:type figsize: tuple of two floats
:param interval: interval in terms of data frames in the input data between the animation frames
:type interval: int
:param n_frames: number of frames to render or the frame index to render if single frame mode
:type n_frames: int
:param output_mode: returns animation object when 'animation', 'singleFrame' returns a single frame figure
:type output_mode: str
:param axis_equal: if true, the axis are rendered with equal scaling
:type axis_equal: bool
:return: animation or figure
"""
if not trajectory.is_static_trajectory:
raise TypeError(
'XZ density animation is currently only implemented for static trajectories'
)
x_pos = trajectory.positions[:, 0, :]
z_pos = trajectory.positions[:, 2, :]
x_min = np.min(x_pos)
x_max = np.max(x_pos)
z_min = np.min(z_pos)
z_max = np.max(z_pos)
if xedges is None:
xedges = np.linspace(x_min, x_max, 50)
elif type(xedges) == int:
xedges = np.linspace(x_min, x_max, xedges)
if zedges is None:
zedges = np.linspace(z_min, z_max, 50)
elif type(zedges) == int:
zedges = np.linspace(z_min, z_max, zedges)
hist_vals, xed, zed = np.histogram2d(x_pos[:, 0],
z_pos[:, 0],
bins=(xedges, zedges))
hist_vals = hist_vals.T
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(111)
im = NonUniformImage(ax, interpolation='nearest')
xcenters = xed[:-1] + 0.5 * (xed[1:] - xed[:-1])
zcenters = zed[:-1] + 0.5 * (zed[1:] - zed[:-1])
im.set_data(xcenters, zcenters, hist_vals)
ax.images.append(im)
im.set_extent(im.get_extent(
)) # workaround for minor issue in matplotlib ocurring in jupyter lab
ax.set_xlim(xed[0], xed[-1])
ax.set_ylim(zed[0], zed[-1])
if axis_equal:
ax.set_aspect('equal')
def animate(i):
ts_number = i * interval
h_vals, _, _ = np.histogram2d(x_pos[:, ts_number],
z_pos[:, ts_number],
bins=(xedges, zedges))
h_vals = h_vals.T
im.set_data(xcenters, zcenters, h_vals)
if output_mode == 'animation':
anim = animation.FuncAnimation(fig,
animate,
frames=n_frames,
blit=False)
return anim
elif output_mode == 'singleFrame':
animate(n_frames)
return fig
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.image import NonUniformImage
from ndtable import NDTable
def peaks(x=np.linspace(-3, 3, 49), y=np.linspace(-3, 3, 49)):
X, Y = np.meshgrid(x, y)
Z = 3*(1-X)**2 * np.e**(-(X**2) - (Y+1)**2) - 10*(X/5 - X**3 - Y**5) * np.e**(-X**2-Y**2) - 1/3 * np.e**(-(X+1)**2 - Y**2)
return X, Y, Z
x = y = np.linspace(-3, 3, 20)
_, _, Z = peaks(x, y)
table = NDTable(Z, (x, y))
xi = yi = np.linspace(-10, 10, 100)
XI, YI = np.meshgrid(xi, yi)
figure = plt.figure(figsize=(10, 5))
figure.canvas.set_window_title('Extrapolation Methods')
for i, method in enumerate(['hold', 'linear']):
ZI = table.evaluate((XI, YI), interp='linear', extrap=method)
ax = figure.add_subplot(1,2,i+1)
ax.set_title(method)
im = NonUniformImage(ax)
im.set_data(xi, yi, ZI)
im.set_extent((-10, 10, -10, 10))
ax.images.append(im)
plt.show()