def animatplot(canvas):
global anim
ax = canvas.figure.subplots()
x = np.linspace(0, 1, 50)
t = np.linspace(0, 1, 20)
X, T = np.meshgrid(x, t)
Y = np.sin(2 * np.pi * (X + T))
block = amp.blocks.Line(X, Y, ax=ax)
canvas.figure.subplots_adjust(
top=0.8) # squish the plot to make space for the controls
slider_ax = canvas.figure.add_axes([0.18, 0.89, 0.5,
0.03]) # the rect of the axis
button_ax = canvas.figure.add_axes([0.78, 0.87, 0.1,
0.07]) # x, y, width, height
anim = amp.Animation([block], fig=canvas.figure)
anim.toggle(ax=button_ax)
anim.timeline_slider(text="TIME",
ax=slider_ax,
color="red",
valfmt="%1.0f")
def to_gif(self, vmin=None, vmax=None, fps=10, dpi=100):
datlist = self.datlist
X, Y = self.X, self.Y
U3d, V3d = self.U3d, self.V3d
Ts = self.XYUVT['T']
if vmax is None:
vmax = np.stack(Ts).max()
if vmin is None:
vmin = np.stack(Ts).min()
self.vmin = vmin
self.vmax = vmax
fig, ax = plt.subplots(figsize=(8, 6), dpi=dpi)
arrows = amp.blocks.Quiver(X, Y, U3d, V3d, ax=ax, t_axis=2, color='k')
blocks = amp.blocks.Imshow(Ts,
ax=ax,
cmap='jet',
vmin=vmin,
vmax=vmax,
interpolation='none')
cbar = fig.colorbar(blocks.im, ax=ax)
timearray = np.array([predict_ICYCLE(s) for s in datlist], dtype='i')
timeline = amp.Timeline(timearray, fps=fps)
anim = amp.Animation([blocks, arrows], timeline)
ax.tick_params(labelbottom=False, bottom=False) # x軸の削除
ax.tick_params(labelleft=False, left=False) # y軸の削除
ax.set_xticklabels([])
fig.tight_layout()
anim.timeline_slider(text='ICYCLE', valfmt='%d')
return anim, fig, ax
def animate_scl(z, Re, T, title='', savename=None):
# get spatial and temporal number of steps
N, _, M = z.shape
# create grid
x = np.linspace(0, 1, N)
t = np.linspace(0, T, M)
Y, X, T = np.meshgrid(x, x, t)
# add title and labels
plt.gca().set_aspect('equal')
plt.gca().set_title(title)
plt.gca().set_ylabel(r'Coordinate $y$')
plt.gca().set_xlabel(r'Coordinate $x$')
# animate scalar field and add time axis
block = amp.blocks.Pcolormesh(X[:, :, 0],
Y[:, :, 0],
z,
t_axis=2,
cmap='RdBu')
plt.colorbar(block.quad)
anim = amp.Animation([block], amp.Timeline(t))
anim.controls()
# save and show
if savename is not None: anim.save('anims/' + savename, dpi=150)
plt.show()
def plot_animation(ref_df):
#X軸・Y軸のデータ取得
X_data = 0
Y_data = 0
#refの経路描画
time_data_np = np.array(ref_df["time"])
x_np = np.array(ref_df["x"])
y_np = np.array(ref_df["y"])
sensor_1_np = np.array(ref_df["sensor1"])
sensor_2_np = np.array(ref_df["sensor2"])
sensor_3_np = np.array(ref_df["sensor3"])
Xs_log = np.asarray([
x_np[t:t + 10] for t in range(len(time_data_np) - 10)
]) #X軸データ × 時間軸 分の配列
Ys_log = [y_np[t:t + 10]
for t in range(len(time_data_np) - 10)] #Y軸データ × 時間軸 分の配列
sensor_1_log = [
sensor_1_np[t:t + 10] for t in range(len(time_data_np) - 10)
]
sensor_2_log = [
sensor_2_np[t:t + 10] for t in range(len(time_data_np) - 10)
]
sensor_3_log = [
sensor_3_np[t:t + 10] for t in range(len(time_data_np) - 10)
]
Time_log = np.asarray(
[time_data_np[t:t + 10] for t in range(len(time_data_np) - 10)])
#subplotの描画 (X-Yの情報を3行分の画面で表示)
ax1 = plt.subplot2grid((3, 2), (0, 0), rowspan=3)
ax2 = plt.subplot2grid((3, 2), (0, 1))
ax3 = plt.subplot2grid((3, 2), (1, 1))
ax4 = plt.subplot2grid((3, 2), (2, 1))
ax1.set_xlim([x_np.min(), x_np.max()]) #描画範囲の設定
ax1.set_ylim([y_np.min(), y_np.max()]) #描画範囲の設定
block = amp.blocks.Scatter(Xs_log, Ys_log, label="X_Y", ax=ax1)
block2 = amp.blocks.Line(Time_log, sensor_1_log, label="sensor1", ax=ax2)
block3 = amp.blocks.Line(Time_log, sensor_2_log, label="sensor2", ax=ax3)
block4 = amp.blocks.Line(Time_log, sensor_3_log, label="sensor3", ax=ax4)
ax2.set_xlim([time_data_np.min(), time_data_np.max()]) #描画範囲の設定
ax2.set_ylim([sensor_1_np.min(), sensor_1_np.max()]) #描画範囲の設定
ax3.set_xlim([time_data_np.min(), time_data_np.max()]) #描画範囲の設定
ax3.set_ylim([sensor_1_np.min(), sensor_1_np.max()]) #描画範囲の設定
ax4.set_xlim([time_data_np.min(), time_data_np.max()]) #描画範囲の設定
ax4.set_ylim([sensor_1_np.min(), sensor_1_np.max()]) #描画範囲の設定
ax1.legend()
ax2.legend()
ax3.legend()
ax4.legend()
plt.subplots_adjust(wspace=0.4, hspace=0.6)
anim = amp.Animation([block, block2, block3, block4])
anim.controls()
anim.save_gif("result")
plt.show()
def test_Line(self):
x = np.linspace(0, 2 * np.pi, 20)
t = np.linspace(0, 2 * np.pi, 5)
X, T = np.meshgrid(x, t)
Y = np.sin(X + T)
block = amp.blocks.Line(X, Y)
return amp.Animation([block])
def test_Pcolormesh(self):
x = np.linspace(-2 * np.pi, 2 * np.pi, 100)
t = np.linspace(0, 2 * np.pi, 3)
X, Y, T = np.meshgrid(x, x, t)
Z = np.sin(X**2 + Y**2 - T)
block = amp.blocks.Pcolormesh(X[:, :, 0], Y[:, :, 0], Z, t_axis=2)
return amp.Animation([block])
def test_Nuke():
ax = plt.gca()
sizes = []
def animate(i):
sizes.append(i + 1)
ax.pie(sizes)
block = amp.blocks.Nuke(animate, length=3, ax=ax)
return amp.Animation([block])
def animation2(self):
x = np.linspace(0, 1, 50)
t = np.linspace(0, 1, 20)
X, T = np.meshgrid(x, t)
Y = np.tan(2 * np.pi * 2 * (X + T))
block = amp.blocks.Line(x, Y, marker=".", linestyle="-", color="r")
anim = amp.Animation([block])
plt.show()
def test_Imshow(self):
x = np.linspace(0, 1, 10)
X, Y = np.meshgrid(x, x)
U = []
for i in range(3):
U.append(X**2 + Y**2 + i)
block = amp.blocks.Imshow(U)
return amp.Animation([block])
def test_Quiver(self):
x = np.linspace(0, 1, 10)
X, Y = np.meshgrid(x, x)
U, V = [], []
for i in range(4):
U.append(X**2 + Y**2 + i)
V.append(X**2 + Y**2 + i)
block = amp.blocks.Quiver(X, Y, U, V)
return amp.Animation([block])
def test_controls():
x = np.linspace(0, 1, 5)
y = np.sin(np.pi * x)
t = np.linspace(0, 1, 5)
timeline = amp.Timeline(t, units='s', fps=5)
block = amp.blocks.ParametricLine(x, y)
block.ax.set_xlim([0, 1])
block.ax.set_ylim([0, 1])
anim = amp.Animation([block], timeline)
anim.controls()
return anim
def make_heatmap_gif(X, Y, f, T, colormap='PuBu'):
"""(X, Y)とz=f(x, y, t)となる、関数fとTを入力すれば、それの時間変化のヒートマップを描画する
:param X: xの一次元のベクトル
:param Y: yの一次元のベクトル
:param f: z=f(x, y, t)となる関数f
:param T: 時間tの一次元のベクトル(時間軸)
:return: なし
表記の約束事として、
x : 要素(値)
X : 一次元ベクトル
XX : Xの二次元メッシュ
XXX : Xの三次元メッシュ
"""
# TODO (nan) 19/2/21 create, not complete_method
XXX, YYY, TTT = np.meshgrid(X, Y, T)
# 高さ方向(ヒートの値)を関数に従って、時間軸で計算 ※meshはxとyが入れ替わるので注意
ZZZ = np.zeros((len(Y), len(X), len(T)))
total = len(Y) * len(X) * len(T)
count = 0
for t in T:
for y in Y:
for x in X:
ZZZ[y, x, t] = f(x, y, t)
count+= 1
print(count, ' / ', total)
# ZZZ = np.sin(XXX * XXX + YYY * YYY - TTT)
# 以下、gif画像描画用コード
print('描画中')
block = amp.blocks.Pcolormesh(XXX[:,:,0], YYY[:,:,0], ZZZ, t_axis=2, cmap=colormap)
plt.colorbar(block.quad)
# plt.gca().set_aspect('equal')
anim = amp.Animation([block], amp.Timeline(T))
anim.controls()
anim.save_gif('pcolormesh')
plt.show()
print('描画完了')
def test_save():
base = 'tests/output_images/'
if not os.path.exists(base):
os.mkdir(base)
x = np.linspace(0, 1, 5)
y = np.sin(np.pi * x)
block = amp.blocks.ParametricLine(x, y)
block.ax.set_xlim([0, 1])
block.ax.set_ylim([0, 1])
anim = amp.Animation([block])
anim.save_gif(base + 'save')
plt.close('all')
assert os.path.exists(base + 'save.gif')
def animate_lines(line_blocks, t, ymin, ymax, save_as=None, fps=10):
timeline = amp.Timeline(t, fps=fps)
animation = amp.Animation(line_blocks, timeline)
plt.xlabel("z")
plt.ylabel("f")
plt.ylim([ymin, ymax])
legend = plt.legend(loc="upper left")
legend.set_draggable(True)
animation.controls(timeline_slider_args={"text": "t"})
if save_as is not None:
from matplotlib.animation import PillowWriter
animation.save(save_as + ".gif", writer=PillowWriter(fps=fps))
return animation
def animate(self):
x = []
y = []
time = np.linspace(0, 1, 50)
for t in time:
self.iterate(15, t)
x.append(np.real(self.segment(self.S)))
y.append(np.imag(self.segment(self.S)))
self.reset()
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_title("Test")
ax.set_aspect("equal")
ax.set_ylim([-0.35, 1.75])
ax.set_xlim([-0.76, 2.5])
timeline = amp.Timeline(time, fps=10)
block = amp.blocks.Line(x, y)
anim = amp.Animation([block], timeline)
anim.controls()
anim.save_gif(f"{type(self).__name__}_15")
plt.show()
def animate(self, i: int):
x = [np.real(self.S0)]
y = [np.imag(self.S0)]
for _ in range(i):
self.iterate(1)
x, y = self.create_block(x, y)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_title("Test")
# ax.set_ylim(self.limits[2:])
# ax.set_xlim(self.limits[:2])
# ax.axis = self.limits
plt.axis(self.limits)
timeline = amp.Timeline(range(i + 1), units="iter", fps=2)
block = amp.blocks.Line(
np.array(x, dtype=object), np.array(y, dtype=object), ax=ax
)
anim = amp.Animation([block], timeline)
# anim.controls()
anim.timeline_slider()
# anim.save("levy2heighway", writer='ffmpeg', fps=timeline.fps)
plt.show()
def animate_f(t, z, f, save_as=None, fps=10):
fig, ax = plt.subplots()
line_block = amp.blocks.Line(z, f, ax=ax)
timeline = amp.Timeline(t, fps=fps)
animation = amp.Animation([line_block], timeline)
ax.set_xlabel("z")
ax.set_ylabel("f")
ax.set_ylim([f.min(), f.max()])
animation.controls(timeline_slider_args={"text": "t"})
if save_as is not None:
from matplotlib.animation import PillowWriter
animation.save(save_as + ".gif", writer=PillowWriter(fps=fps))
return animation
def animate_line(
data,
animate_over=None,
animate=True,
axis_coords=None,
vmin=None,
vmax=None,
fps=10,
save_as=None,
sep_pos=None,
ax=None,
aspect=None,
controls="both",
**kwargs,
):
"""
Plots a line plot which is animated with time.
Currently only supports 1D+1 data, which it plots with animatplot's Line animation.
Parameters
----------
data : xarray.DataArray
animate_over : str, optional
Dimension over which to animate, defaults to the time dimension
animate : bool, optional
If set to false, do not create the animation, just return the block
axis_coords : None, str, dict
Coordinates to use for axis labelling.
- None: Use the dimension coordinate for each axis, if it exists.
- "index": Use the integer index values.
- dict: keys are dimension names, values set axis_coords for each axis
separately. Values can be: None, "index", the name of a 1d variable or
coordinate (which must have the dimension given by 'key'), or a 1d
numpy array, dask array or DataArray whose length matches the length of
the dimension given by 'key'.
vmin : float, optional
Minimum value to use for colorbar. Default is to use minimum value of
data across whole timeseries.
vmax : float, optional
Maximum value to use for colorbar. Default is to use maximum value of
data across whole timeseries.
fps : int, optional
Frames per second of resulting gif
save_as : True or str, optional
If str is passed, save the animation as save_as+'.gif'.
If True is passed, save the animation with a default name,
'<variable name>_over_<animate_over>.gif'
sep_pos : int, optional
Radial position at which to plot the separatrix
ax : Axes, optional
A matplotlib axes instance to plot to. If None, create a new
figure and axes, and plot to that
aspect : str or None, optional
Argument to set_aspect(), defaults to "auto"
controls : string or None, default "both"
By default, add both the timeline and play/pause toggle to the animation. If
"timeline" is passed add only the timeline, if "toggle" is passed add only the
play/pause toggle. If None or an empty string is passed, add neither.
kwargs : dict, optional
Additional keyword arguments are passed on to the plotting function
animatplot.blocks.Line
Returns
-------
animation or block
If animate==True, returns an animatplot.Animation object, otherwise
returns an animatplot.blocks.Line instance.
"""
if animate_over is None:
animate_over = data.metadata.get("bout_tdim", "t")
if aspect is None:
aspect = "auto"
variable = data.name
# Check plot is the right orientation
t_read, x_read = data.dims
if t_read is animate_over:
x = x_read
else:
data = data.transpose(animate_over, t_read, transpose_coords=True)
x = t_read
# Load values eagerly otherwise for some reason the plotting takes
# 100's of times longer - for some reason animatplot does not deal
# well with dask arrays!
image_data = data.values
# If not specified, determine max and min values across entire data series
if vmax is None:
vmax = np.max(image_data)
if vmin is None:
vmin = np.min(image_data)
x_values, x_label = _parse_coord_option(x, axis_coords, data)
if not ax:
fig, ax = plt.subplots()
ax.set_aspect(aspect)
# set range of plot
ax.set_ylim([vmin, vmax])
line_block = amp.blocks.Line(x_values, image_data, ax=ax, **kwargs)
if animate:
t_values, t_label = _parse_coord_option(animate_over, axis_coords,
data)
t_values, t_suffix = _normalise_time_coord(t_values)
timeline = amp.Timeline(t_values, fps=fps, units=t_suffix)
anim = amp.Animation([line_block], timeline)
# Add title and axis labels
ax.set_title(variable)
ax.set_xlabel(x_label)
if "long_name" in data.attrs:
y_label = data.long_name
else:
y_label = variable
if "units" in data.attrs:
y_label = y_label + f" [{data.units}]"
ax.set_ylabel(y_label)
# Plot separatrix
if sep_pos:
ax.plot_vline(sep_pos, "--")
if animate:
_add_controls(anim, controls, t_label)
if save_as is not None:
if save_as is True:
save_as = "{}_over_{}".format(variable, animate_over)
anim.save(save_as + ".gif", writer=PillowWriter(fps=fps))
return anim
return line_block
def psi(t):
x = t
y = np.sin(t)
return x, y
t = np.linspace(0, 2 * np.pi, 25)
x, y = psi(t)
X, Y = amp.util.parametric_line(x, y)
timeline = amp.Timeline(t, "s", 24)
ax = plt.axes(xlim=[0, 2], ylim=[-1.1, 1.1])
block1 = amp.blocks.Line(X, Y, ax=ax)
# or equivalently
# block1 = amp.blocks.ParametricLine(x, y, ax=ax)
anim = amp.Animation([block1], timeline)
# Your standard matplotlib stuff
plt.title("Parametric Line")
plt.xlabel("x")
plt.ylabel(r"y")
# Create Interactive Elements
anim.toggle()
anim.timeline_slider()
anim.save("parametric.gif", writer=PillowWriter(fps=5))
plt.show()
import numpy as np
import matplotlib.pyplot as plt
import animatplot as amp
x = np.linspace(0, 1, 50)
t = np.linspace(0, 1, 20)
X, T = np.meshgrid(x, t)
Y = np.sin(2*np.pi*(X+T))
block = amp.blocks.Line(X, Y)
plt.subplots_adjust(top=0.8) # squish the plot to make space for the controls
slider_ax = plt.axes([.18, .89, .5, .03]) # the rect of the axis
button_ax = plt.axes([.78, .87, .1, .07]) # x, y, width, height
anim = amp.Animation([block])
anim.toggle(ax=button_ax)
anim.timeline_slider(text='TIME', ax=slider_ax, color='red', valfmt='%1.0f')
# equivalent to:
# anim.controls({'text':'TIME', 'ax':slider_ax, 'color':'red', 'valfmt':'%1.0f'},
# {'ax':button_axis})
# anim.save_gif('images/controls')
plt.show()
def animate_list(self,
variables,
animate_over='t',
save_as=None,
show=False,
fps=10,
nrows=None,
ncols=None,
poloidal_plot=False,
subplots_adjust=None,
vmin=None,
vmax=None,
logscale=None,
titles=None,
aspect='equal',
controls=True,
tight_layout=True,
**kwargs):
"""
Parameters
----------
variables : list of str or BoutDataArray
The variables to plot. For any string passed, the corresponding
variable in this DataSet is used - then the calling DataSet must
have only 3 dimensions. It is possible to pass BoutDataArrays to
allow more flexible plots, e.g. with different variables being
plotted against different axes.
animate_over : str, optional
Dimension over which to animate
save_as : str, optional
If passed, a gif is created with this filename
show : bool, optional
Call pyplot.show() to display the animation
fps : float, optional
Indicates the number of frames per second to play
nrows : int, optional
Specify the number of rows of plots
ncols : int, optional
Specify the number of columns of plots
poloidal_plot : bool or sequence of bool, optional
If set to True, make all 2D animations in the poloidal plane instead of using
grid coordinates, per variable if sequence is given
subplots_adjust : dict, optional
Arguments passed to fig.subplots_adjust()()
vmin : float or sequence of floats
Minimum value for color scale, per variable if a sequence is given
vmax : float or sequence of floats
Maximum value for color scale, per variable if a sequence is given
logscale : bool or float, sequence of bool or float, optional
If True, default to a logarithmic color scale instead of a linear one.
If a non-bool type is passed it is treated as a float used to set the linear
threshold of a symmetric logarithmic scale as
linthresh=min(abs(vmin),abs(vmax))*logscale, defaults to 1e-5 if True is
passed.
Per variable if sequence is given.
titles : sequence of str or None, optional
Custom titles for each plot. Pass None in the sequence to use the default for
a certain variable
aspect : str or None, or sequence of str or None, optional
Argument to set_aspect() for each plot
controls : bool, optional
If set to False, do not show the time-slider or pause button
tight_layout : bool or dict, optional
If set to False, don't call tight_layout() on the figure.
If a dict is passed, the dict entries are passed as arguments to
tight_layout()
**kwargs : dict, optional
Additional keyword arguments are passed on to each animation function
"""
nvars = len(variables)
if nrows is None and ncols is None:
ncols = int(np.ceil(np.sqrt(nvars)))
nrows = int(np.ceil(nvars / ncols))
elif nrows is None:
nrows = int(np.ceil(nvars / ncols))
elif ncols is None:
ncols = int(np.ceil(nvars / nrows))
else:
if nrows * ncols < nvars:
raise ValueError(
'Not enough rows*columns to fit all variables')
fig, axes = plt.subplots(nrows, ncols, squeeze=False)
axes = axes.flatten()
ncells = nrows * ncols
if nvars < ncells:
for index in range(ncells - nvars):
fig.delaxes(axes[ncells - index - 1])
if subplots_adjust is not None:
fig.subplots_adjust(**subplots_adjust)
def _expand_list_arg(arg, arg_name):
if isinstance(arg,
collections.Sequence) and not isinstance(arg, str):
if len(arg) != len(variables):
raise ValueError(
'if %s is a sequence, it must have the same '
'number of elements as "variables"' % arg_name)
else:
arg = [arg] * len(variables)
return arg
poloidal_plot = _expand_list_arg(poloidal_plot, 'poloidal_plot')
vmin = _expand_list_arg(vmin, 'vmin')
vmax = _expand_list_arg(vmax, 'vmax')
logscale = _expand_list_arg(logscale, 'logscale')
titles = _expand_list_arg(titles, 'titles')
aspect = _expand_list_arg(aspect, 'aspect')
blocks = []
for subplot_args in zip(variables, axes, poloidal_plot, vmin, vmax,
logscale, titles, aspect):
(v, ax, this_poloidal_plot, this_vmin, this_vmax, this_logscale,
this_title, this_aspect) = subplot_args
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.1)
ax.set_aspect(this_aspect)
if isinstance(v, str):
v = self.data[v]
data = v.bout.data
ndims = len(data.dims)
ax.set_title(data.name)
if ndims == 2:
blocks.append(
animate_line(data=data,
ax=ax,
animate_over=animate_over,
animate=False,
**kwargs))
elif ndims == 3:
if this_vmin is None:
this_vmin = data.min().values
if this_vmax is None:
this_vmax = data.max().values
norm = _create_norm(this_logscale, kwargs.get('norm', None),
this_vmin, this_vmax)
if this_poloidal_plot:
var_blocks = animate_poloidal(data,
ax=ax,
cax=cax,
animate_over=animate_over,
animate=False,
vmin=this_vmin,
vmax=this_vmax,
norm=norm,
aspect=this_aspect,
**kwargs)
for block in var_blocks:
blocks.append(block)
else:
blocks.append(
animate_pcolormesh(data=data,
ax=ax,
cax=cax,
animate_over=animate_over,
animate=False,
vmin=this_vmin,
vmax=this_vmax,
norm=norm,
**kwargs))
else:
raise ValueError("Unsupported number of dimensions " +
str(ndims) + ". Dims are " + str(v.dims))
if this_title is not None:
# Replace default title with user-specified one
ax.set_title(this_title)
timeline = amp.Timeline(np.arange(v.sizes[animate_over]), fps=fps)
anim = amp.Animation(blocks, timeline)
if tight_layout:
if subplots_adjust is not None:
warnings.warn(
'tight_layout argument to animate_list() is True, but '
'subplots_adjust argument is not None. subplots_adjust '
'is being ignored.')
if not isinstance(tight_layout, dict):
tight_layout = {}
fig.tight_layout(**tight_layout)
if controls:
anim.controls(timeline_slider_args={'text': animate_over})
if save_as is not None:
anim.save(save_as + '.gif', writer=PillowWriter(fps=fps))
if show:
plt.show()
return anim
def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True,
vmin=None, vmax=None, vsymmetric=False, fps=10, save_as=None,
ax=None, controls=True, **kwargs):
"""
Plots a color plot which is animated with time over the specified
coordinate.
Currently only supports 2D+1 data, which it plots with animatplotlib's
wrapping of matplotlib's pcolormesh.
Parameters
----------
data : xarray.DataArray
animate_over : str, optional
Dimension over which to animate
x : str, optional
Dimension to use on the x axis, default is None - then use the first spatial
dimension of the data
y : str, optional
Dimension to use on the y axis, default is None - then use the second spatial
dimension of the data
vmin : float, optional
Minimum value to use for colorbar. Default is to use minimum value of
data across whole timeseries.
vmax : float, optional
Maximum value to use for colorbar. Default is to use maximum value of
data across whole timeseries.
save_as: str, optional
Filename to give to the resulting gif
fps : int, optional
Frames per second of resulting gif
kwargs : dict, optional
Additional keyword arguments are passed on to the plotting function
(e.g. imshow for 2D plots).
"""
variable = data.name
# Check plot is the right orientation
spatial_dims = list(data.dims)
if len(data.dims) != 3:
raise ValueError('Data passed to animate_imshow must be 3-dimensional')
try:
spatial_dims.remove(animate_over)
except ValueError:
raise ValueError("Dimension animate_over={} is not present in the data"
.format(animate_over))
if x is None and y is None:
x, y = spatial_dims
elif x is None:
try:
spatial_dims.remove(y)
except ValueError:
raise ValueError("Dimension {} is not present in the data" .format(y))
x = spatial_dims[0]
elif y is None:
try:
spatial_dims.remove(x)
except ValueError:
raise ValueError("Dimension {} is not present in the data" .format(x))
y = spatial_dims[0]
data = data.transpose(animate_over, y, x)
# Load values eagerly otherwise for some reason the plotting takes
# 100's of times longer - for some reason animatplot does not deal
# well with dask arrays!
image_data = data.values
# If not specified, determine max and min values across entire data series
if vmax is None:
vmax = np.max(image_data)
if vmin is None:
vmin = np.min(image_data)
if vsymmetric:
vmax = max(np.abs(vmin), np.abs(vmax))
vmin = -vmax
if not ax:
fig, ax = plt.subplots()
# Note: animatplot's Pcolormesh gave strange outputs without passing
# explicitly x- and y-value arrays, although in principle these should not
# be necessary.
ny, nx = image_data.shape[1:]
pcolormesh_block = amp.blocks.Pcolormesh(np.arange(float(nx)), np.arange(float(ny)),
image_data, vmin=vmin, vmax=vmax, ax=ax,
**kwargs)
if animate:
timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps)
anim = amp.Animation([pcolormesh_block], timeline)
cbar = plt.colorbar(pcolormesh_block.quad, ax=ax)
cbar.ax.set_ylabel(variable)
# Add title and axis labels
ax.set_title(variable)
ax.set_xlabel(x)
ax.set_ylabel(y)
if animate:
if controls:
anim.controls(timeline_slider_args={'text': animate_over})
if not save_as:
save_as = "{}_over_{}".format(variable, animate_over)
anim.save(save_as + '.gif', writer=PillowWriter(fps=fps))
return pcolormesh_block
def animate_poloidal(
da,
*,
ax=None,
cax=None,
animate_over=None,
separatrix=True,
targets=True,
add_limiter_hatching=True,
cmap=None,
axis_coords=None,
vmin=None,
vmax=None,
logscale=False,
animate=True,
save_as=None,
fps=10,
controls="both",
aspect=None,
extend=None,
**kwargs,
):
"""
Make a 2D plot in R-Z coordinates using animatplotlib's Pcolormesh, taking into
account branch cuts (X-points).
Parameters
----------
da : xarray.DataArray
A 2D (x,y) DataArray of data to plot
ax : Axes, optional
A matplotlib axes instance to plot to. If None, create a new
figure and axes, and plot to that
cax : Axes, optional
Matplotlib axes instance where the colorbar will be plotted. If None, the default
position created by matplotlab.figure.Figure.colorbar() will be used.
animate_over : str, optional
Dimension over which to animate, defaults to the time dimension
separatrix : bool, optional
Add dashed lines showing separatrices
targets : bool, optional
Draw solid lines at the target surfaces
add_limiter_hatching : bool, optional
Draw hatched areas at the targets
cmap : matplotlib.colors.Colormap instance, optional
Colors to use for the plot
axis_coords : None, str, dict
Coordinates to use for axis labelling. Only affects time coordinate.
- None: Use the dimension coordinate for each axis, if it exists.
- "index": Use the integer index values.
- dict: keys are dimension names, values set axis_coords for each axis
separately. Values can be: None, "index", the name of a 1d variable or
coordinate (which must have the dimension given by 'key'), or a 1d
numpy array, dask array or DataArray whose length matches the length of
the dimension given by 'key'.
vmin : float, optional
Minimum value for the color scale
vmax : float, optional
Maximum value for the color scale
logscale : bool or float, optional
If True, default to a logarithmic color scale instead of a linear one.
If a non-bool type is passed it is treated as a float used to set the linear
threshold of a symmetric logarithmic scale as
linthresh=min(abs(vmin),abs(vmax))*logscale, defaults to 1e-5 if True is
passed.
animate : bool, optional
If set to false, do not create the animation, just return the blocks
save_as : True or str, optional
If str is passed, save the animation as save_as+'.gif'.
If True is passed, save the animation with a default name,
'<variable name>_over_<animate_over>.gif'
fps : float, optional
Frame rate for the animation
controls : string or None, default "both"
By default, add both the timeline and play/pause toggle to the animation. If
"timeline" is passed add only the timeline, if "toggle" is passed add only the
play/pause toggle. If None or an empty string is passed, add neither.
aspect : str or None, optional
Argument to set_aspect(), defaults to "equal"
extend : str or None, optional
Passed to fig.colorbar()
**kwargs : optional
Additional arguments are passed on to the animation method
animatplot.blocks.Pcolormesh
Returns
-------
animation or blocks
If animate==True, returns an animatplot.Animation object, otherwise
returns a list of animatplot.blocks.Pcolormesh instances.
"""
if animate_over is None:
animate_over = da.metadata.get("bout_tdim", "t")
if aspect is None:
aspect = "equal"
# TODO generalise this
x = kwargs.pop("x", "R")
y = kwargs.pop("y", "Z")
# Check plot is the right orientation
spatial_dims = list(da.dims)
try:
spatial_dims.remove(animate_over)
except ValueError:
raise ValueError(
"Dimension animate_over={} is not present in the data".format(
animate_over))
if len(da.dims) != 3:
raise ValueError("da must be 2+1D (t,x,y)")
if ax is None:
fig, ax = plt.subplots()
else:
fig = ax.get_figure()
if vmin is None:
vmin = da.min().values
if vmax is None:
vmax = da.max().values
if extend is None:
# Replicate default for older matplotlib that does not handle extend=None
# matplotlib-3.3 definitely does not need this. Not sure about 3.0, 3.1, 3.2.
extend = "neither"
# create colorbar
norm = _create_norm(logscale, kwargs.pop("norm", None), vmin, vmax)
sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
sm.set_array([])
cmap = sm.get_cmap()
cbar = fig.colorbar(sm, ax=ax, cax=cax, extend=extend)
if "long_name" in da.attrs:
cbar_label = da.long_name
else:
cbar_label = da.name
if "units" in da.attrs:
cbar_label += f" [{da.units}]"
cbar.ax.set_ylabel(cbar_label)
ax.set_aspect(aspect)
da_regions = _decompose_regions(da)
# Plot all regions on same axis
blocks = []
with warnings.catch_warnings():
# The coordinates we pass are a logically rectangular grid, so should be fine
# even if this warning is triggered by pcolor or pcolormesh
warnings.filterwarnings(
"ignore",
"The input coordinates to pcolormesh are interpreted as cell centers, but "
"are not monotonically increasing or decreasing. This may lead to "
"incorrectly calculated cell edges, in which case, please supply explicit "
"cell edges to pcolormesh.",
UserWarning,
)
for da_region in da_regions.values():
# Load values eagerly otherwise for some reason the plotting takes
# 100's of times longer - for some reason animatplot does not deal
# well with dask arrays!
blocks.append(
amp.blocks.Pcolormesh(
da_region.coords[x].values,
da_region.coords[y].values,
da_region.values,
ax=ax,
cmap=cmap,
norm=norm,
**kwargs,
))
ax.set_title(da.name)
ax.set_xlabel(x)
ax.set_ylabel(y)
if _is_core_only(da):
separatrix = False
targets = False
if separatrix:
plot_separatrices(da_regions, ax, x=x, y=y)
if targets:
plot_targets(da_regions, ax, x=x, y=y, hatching=add_limiter_hatching)
if animate:
t_values, t_label = _parse_coord_option(animate_over, axis_coords, da)
t_values, t_suffix = _normalise_time_coord(t_values)
timeline = amp.Timeline(t_values, fps=fps, units=t_suffix)
anim = amp.Animation(blocks, timeline)
_add_controls(anim, controls, t_label)
if save_as is not None:
if save_as is True:
save_as = "{}_over_{}".format(da.name, animate_over)
anim.save(save_as + ".gif", writer=PillowWriter(fps=fps))
return anim
return blocks
fig = figure()
ax = Axes3D(fig, azim=-40, elev=30)
sphere = Bloch(axes=ax)
def animate(i):
sphere.clear()
# sphere.add_states(i)
sphere.add_vectors([0, np.sin(i), np.sin(i) + np.cos(i)])
# sphere.add_points([sx[:i+1],sy[:i+1],sz[:i+1]])
sphere.make_sphere()
return ax
def init():
sphere.vector_color = ['r']
return ax
ani = animation.FuncAnimation(fig,
animate,
np.linspace(1, pi, 100),
init_func=init,
interval=1.0,
repeat=False)
anim = amp.Animation(ani)
anim.controls()
sphere.show()
def generate_video(self):
""" Generates a visualization of the simulated fields' dynamics.
| Output:
| - .avi format video.
"""
# Creating arrays of indices for x-y axis and time.
x_ind = range(0, len(self.x_axis))
y_ind = range(0, len(self.y_axis))
t_ind = range(0, len(self.data["time"]))
# Creating a grid with time and x-y axis data:
X, Y, _ = np.meshgrid(self.x_axis, self.y_axis, self.data["time"])
# Creating a figure for visualization
fig, (ax1, ax2, ax3) = plt.subplots(3, 1)
# Calculation of min and max of simulated fields.
dicmaxmins = {"Ex": ['valuese_x', [], []],\
"Ey": ['valuese_y', [], []],\
"Hz": ['values', [], []]}
t0 = 0
tf = 100
for i in dicmaxmins.values():
for time in self.data[i[0]][t0:tf]:
(i[1]).append(max([max(j) for j in time]))
(i[2]).append(min([min(j) for j in time]))
for i in dicmaxmins:
dicmaxmins[i][1] = max(dicmaxmins[i][1])
dicmaxmins[i][2] = min(dicmaxmins[i][2])
# Animation
fields = {
"Ex": ['valuese_x', ax1, [], []],
"Ey": ['valuese_y', ax2, [], []],
"Hz": ['values', ax3, [], []]
}
for i in fields.values():
i[1].set_ylabel('y')
fields["Hz"][1].set_xlabel('x')
fields["Ex"][1].set_title(r'$ {E_x} $')
fields["Ey"][1].set_title(r'$ {E_y} $')
fields["Hz"][1].set_title(r'$ H_z $')
for i in fields:
field = fields[i][0]
fields[i][2] = np.array([np.array([np.array([self.data[field][t][i][j]\
for t in t_ind]) for i in x_ind]) for j in y_ind])
fig.suptitle(r'${ {E_x} \ & \ {E_y} \ & \ H_z }$')
fig.subplots_adjust(left=None,
bottom=0.1,
right=None,
top=0.85,
wspace=None,
hspace=0.5)
# now we make our blocks
for i in fields:
fields[i][3] = amp.blocks.Pcolormesh(X[:, :, 0],
Y[:, :, 0],
fields[i][2],
ax=fields[i][1],
t_axis=2,
vmin=(dicmaxmins[i][2]) * 0.6,
vmax=(dicmaxmins[i][1]) * 0.6)
for i in fields:
plt.colorbar(fields[i][3].quad, ax=fields[i][1])
timeline = amp.Timeline([i * (10**9) for i in self.data["time"]],
fps=10,
units='ns')
# now to contruct the animation
anim = amp.Animation([
fields["Ex"][3],
fields["Ey"][3],
fields["Hz"][3],
], timeline)
anim.controls()
# Change if windows.
#anim.save_gif('videos/allfields')
anim.save('videos/allfields.avi')
plt.show()
# create the different plotting axes
fig, (ax1, ax2) = plt.subplots(1, 2)
for ax in [ax1, ax2]:
ax.set_aspect('equal')
ax.set_xlabel('x')
ax2.set_ylabel('y', labelpad=-5)
ax1.set_ylabel('z')
ax1.set_ylim([-1.1, 1.1])
fig.suptitle('Multiple blocks')
ax1.set_title('Cross Section: $y=0$')
ax2.set_title(r'$z=\sin(x^2+y^2-t)$')
# animatplot stuff
# now we make our blocks
line_block = amp.blocks.Line(X[0, :, :], line_data,
ax=ax1, t_axis=1)
pcolormesh_block = amp.blocks.Pcolormesh(X[:, :, 0], Y[:, :, 0], pcolormesh_data,
ax=ax2, t_axis=2, vmin=-1, vmax=1)
plt.colorbar(plt.cm.ScalarMappable())
timeline = amp.Timeline(t, fps=30)
# now to contruct the animation
anim = amp.Animation([pcolormesh_block, line_block], timeline)
anim.controls()
anim.save_gif('multiblock')
plt.show()
def animate_poloidal(da, *, ax=None, animate_over='t', separatrix=True, targets=True,
add_limiter_hatching=True, cmap=None, vmin=None, vmax=None,
animate=True, save_as=None, fps=10, controls=True, **kwargs):
"""
Make a 2D plot in R-Z coordinates using animatplotlib's Pcolormesh, taking into
account branch cuts (X-points).
Parameters
----------
da : xarray.DataArray
A 2D (x,y) DataArray of data to plot
ax : Axes, optional
A matplotlib axes instance to plot to. If None, create a new
figure and axes, and plot to that
separatrix : bool, optional
Add dashed lines showing separatrices
targets : bool, optional
Draw solid lines at the target surfaces
add_limiter_hatching : bool, optional
Draw hatched areas at the targets
**kwargs : optional
Additional arguments are passed on to method
###Returns
###-------
###artists
### List of the contourf instances
"""
# TODO generalise this
x = kwargs.pop('x', 'R')
y = kwargs.pop('y', 'Z')
# Check plot is the right orientation
spatial_dims = list(da.dims)
try:
spatial_dims.remove(animate_over)
except ValueError:
raise ValueError("Dimension animate_over={} is not present in the data"
.format(animate_over))
if len(da.dims) != 3:
raise ValueError("da must be 2+1D (t,x,y)")
if ax is None:
fig, ax = plt.subplots()
else:
fig = ax.get_figure()
if vmin is None:
vmin = da.min().values
if vmax is None:
vmax = da.max().values
# pass vmin and vmax through kwargs as they are not used for contour plots
kwargs['vmin'] = vmin
kwargs['vmax'] = vmax
# create colorbar
norm = (kwargs['norm'] if 'norm' in kwargs
else matplotlib.colors.Normalize(vmin=vmin, vmax=vmax))
sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
sm.set_array([])
cmap = sm.get_cmap()
fig.colorbar(sm, ax=ax)
ax.set_aspect('equal')
regions = _decompose_regions(da)
# Plot all regions on same axis
blocks = []
for region in regions:
# Load values eagerly otherwise for some reason the plotting takes
# 100's of times longer - for some reason animatplot does not deal
# well with dask arrays!
blocks.append(amp.blocks.Pcolormesh(region.coords[x].values,
region.coords[y].values, region.values, ax=ax, cmap=cmap,
**kwargs))
ax.set_title(da.name)
ax.set_xlabel(x)
ax.set_ylabel(y)
if _is_core_only(da):
separatrix = False
targets = False
if separatrix:
plot_separatrices(da, ax)
if targets:
plot_targets(da, ax, hatching=add_limiter_hatching)
if animate:
timeline = amp.Timeline(np.arange(da.sizes[animate_over]), fps=fps)
anim = amp.Animation(blocks, timeline)
if controls:
anim.controls(timeline_slider_args={'text': animate_over})
if not save_as:
save_as = "{}_over_{}".format(da.name, animate_over)
anim.save(save_as + '.gif', writer=PillowWriter(fps=fps))
return blocks
def animate_climate_fields(n_mapseries, lats, lons, suptitle='', fps=10,
save_path='', layout=None, shape=None, subtitles=[],
figsize=None, colorbar_clip_pct = 1,
share_colorbar=False):
"""Produces and shows a figure which is an inimation of n climate fields.
Saves figure out optionally.
Args:
n_mapseries (numpy array):
layout (numpy array): position of each index in grid, To leave
position blank set the layout value to nan for that position.
subtitles (list): subtitles for each respective element of n_mapseries"""
N = n_mapseries.shape[0]
# set plot layout
if layout is None:
if shape is None: shape = ((N+1)//2, 2)
if np.product(shape)<N:
raise ValueError('shape {} too small for {} modes'.format(layout, N))
layout = np.arange(np.product(shape), dtype=object)
layout[N:]=np.nan
else:
shape = layout.shape
fig = plt.figure(figsize = figsize)
ind_j, ind_i = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
# flatten and ignore some positions
layout=layout.flatten()
mask = [~np.isnan(x) for x in layout]
layout=layout[mask]
ind_j = ind_j.flatten()[mask]
ind_i = ind_i.flatten()[mask]
# default subtitles
if subtitles==[]: subtitles = ['mode {}'.format(i) for i in range(N)]
if share_colorbar:
vamp = np.abs(np.percentile(n_mapseries, [colorbar_clip_pct, 100-colorbar_clip_pct])).max()
# spatial and time coordinates
X,Y = np.meshgrid(lons,lats)
t = np.arange(n_mapseries.shape[1])
timeline = amp.Timeline(t, fps=fps)
#create axes and set blocks
blocks = []
for n in range(len(layout)):
m = int(layout[n])
ax = plt.subplot2grid(shape,(ind_i[n], ind_j[n]))
ax.set_aspect('equal')
ax.set_title(subtitles[m])
if not share_colorbar:
vamp = np.abs(np.percentile(n_mapseries[m], [colorbar_clip_pct, 100-colorbar_clip_pct])).max()
block = amp.blocks.Pcolormesh(X,Y, n_mapseries[m],
ax=ax, t_axis=0,
cmap=cmocean.cm.balance,
vmin=-vamp, vmax=vamp)
blocks.append(block)
#plt.colorbar(blocks[-1].quad)
# contruct the animation
anim = amp.Animation(blocks, timeline)
# other matplotlib modifications
fig.suptitle(suptitle)
plt.tight_layout()
# controls
anim.controls()
# save gif
if save_path!='':
anim.save_gif(save_path)
plt.show()
def animate_line(data, animate_over='t', animate=True,
vmin=None, vmax=None, fps=10, save_as=None, sep_pos=None, ax=None,
controls=True, **kwargs):
"""
Plots a line plot which is animated with time.
Currently only supports 1D+1 data, which it plots with animatplot's Line animation.
Parameters
----------
data : xarray.DataArray
animate_over : str, optional
Dimension over which to animate
vmin : float, optional
Minimum value to use for colorbar. Default is to use minimum value of
data across whole timeseries.
vmax : float, optional
Maximum value to use for colorbar. Default is to use maximum value of
data across whole timeseries.
sep_pos : int, optional
Radial position at which to plot the separatrix
save_as: str, optional
Filename to give to the resulting gif
fps : int, optional
Frames per second of resulting gif
kwargs : dict, optional
Additional keyword arguments are passed on to the plotting function
(e.g. imshow for 2D plots).
"""
variable = data.name
# Check plot is the right orientation
t_read, x_read = data.dims
if (t_read is animate_over):
pass
else:
data = data.transpose(animate_over, t_read)
# Load values eagerly otherwise for some reason the plotting takes
# 100's of times longer - for some reason animatplot does not deal
# well with dask arrays!
image_data = data.values
# If not specified, determine max and min values across entire data series
if vmax is None:
vmax = np.max(image_data)
if vmin is None:
vmin = np.min(image_data)
if not ax:
fig, ax = plt.subplots()
# set range of plot
ax.set_ylim([vmin, vmax])
line_block = amp.blocks.Line(image_data, ax=ax, **kwargs)
if animate:
timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps)
anim = amp.Animation([line_block], timeline)
# Add title and axis labels
ax.set_title(variable)
ax.set_xlabel(x_read)
ax.set_ylabel(variable)
# Plot separatrix
if sep_pos:
ax.plot_vline(sep_pos, '--')
if animate:
if controls:
anim.controls(timeline_slider_args={'text': animate_over})
if not save_as:
save_as = "{}_over_{}".format(variable, animate_over)
anim.save(save_as + '.gif', writer=PillowWriter(fps=fps))
return line_block
#minh = minh if minh>-np.inf else 0
#maxh = maxh if maxh<np.inf else 3
ax1.set_ylim(minu, maxu)
ax2.set_ylim(minh, maxh)
# time and space info for axis
t = np.arange(nframes)
x = np.linspace(0, diam, nx)
X = np.repeat(x[:, np.newaxis], nt // nsave, 1).T
# create variable blocks and timeline
u_block = amp.blocks.Line(X, u_all, ax=ax1)
h_block = amp.blocks.Line(X, h_all, ax=ax2)
timeline = amp.Timeline(t, units='s', fps=20)
# Make gif
anim = amp.Animation([u_block, h_block], timeline=timeline)
plt.tight_layout()
anim.controls()
anim.save_gif(gifname) # save animation for docs
plt.show()
print('DONE')
"""
funtion to test lengths of saved frames
def f(nt, nsave):
x0 = np.zeros((nt//nsave+1*((nt%nsave)!=0), nx)).shape[0]
x1 = [i//nsave for i in np.arange(nt) if i%nsave==0][-1]
print(x0,x1)
"""