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 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 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_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(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 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()
import matplotlib.pyplot as plt
from matplotlib.animation import PillowWriter
import animatplot as amp
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()
triangle_points = [[1 / np.sqrt(3), 0], [0, 1], [-1 / np.sqrt(3), 0]]
x_start = 0
y_start = 0
start_point = [x_start, y_start]
n = 10000
point_list = np.zeros([n, 2])
point_list = calc_points(n, start_point, triangle_points, point_list)
t = np.linspace(1, n, n)
x = point_list[:, 0]
y = point_list[:, 1]
X, Y = aplt.util.parametric_line(x, y)
# Animatplot animation setup
timeline = aplt.Timeline(t, r'$\ pts$', fps=60)
ax = plt.axes()
for j in range(0, len(triangle_points)):
plt.plot(triangle_points[j][0], triangle_points[j][1], 'ro')
plt.plot(x_start, y_start, 'bo')
plt.axis('equal')
plt.grid()
block1 = aplt.blocks.Line(X,
Y,
ax,
marker='o',
color='k',
linestyle='',
markersize='0.5')
# 式(3.35)
Fr[i][j] = (R * EXP(1j * (w * t -
(krz * z + krx * x)))).real
if x < 0:
# 式(3.37)
Ft[i][j] = (T * EXP(1j * (w * t -
(ktz * z + ktx * x)))).real
F = Fi + Fr + Ft
images = [Fi, F, Fr + Ft]
titles = [
'incident field', 'total field',
'reflected and \ntransmitted fields'
]
for i in range(0, len(images)):
ax[i].imshow(images[i], vmin=-2, vmax=2, cmap='jet')
ax[i].invert_yaxis()
ax[i].set_title(titles[i])
ax[i].set_xlabel('z')
ax[i].set_xticks([])
ax[i].set_ylabel('x')
ax[i].set_yticks([])
# アニメーション表示
timeline = amp.Timeline(range(0, step_per_period))
block = amp.blocks.Nuke(update, length=len(timeline))
anim = amp.Animation([block], timeline)
anim.save_gif(os.path.basename(__file__).split('.')[0])
plt.show()
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()
# standard matplotlib stuff
# create the different plotting axes
plt.figure()
plt.xlabel('x')
plt.ylabel('y')
plt.title(r'$z=Ae^{b(\vec r-\vec v t)^2}$')
x_ind = range(0,Nx+1)
y_ind = range(0,Ny+1)
t_ind = range(0,Nt)
X_ind, Y_ind, T_ind = np.meshgrid(x_ind, y_ind, t_ind)
pcolormesh_data2 = np.array([np.array([np.array([Z[t][i][j] for t in t_ind]) for i in x_ind]) for j in y_ind])
# animatplot stuff
# now we make our blocks
pcolormesh_block = amp.blocks.Pcolormesh(X[:,:,0], Y[:,:,0], pcolormesh_data2,
t_axis=2, vmin=-1, vmax=1)
plt.colorbar(pcolormesh_block.quad)
timeline = amp.Timeline(t, fps=10)
# now to contruct the animation
anim = amp.Animation([pcolormesh_block], timeline)
anim.controls()
anim.save_gif('Gaussian')
plt.show()
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_imshow(data,
animate_over='t',
x='x',
y='y',
animate=True,
vmin='min',
vmax='max',
fps=10,
save_as=None,
sep_pos=None,
ax=None,
**kwargs):
"""
Plots a color plot which is animated with time over the specified
coordinate.
Currently only supports 2D+1 data, which it plots with xarray's
wrapping of matplotlib's imshow.
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 'x'
y : str, optional
Dimension to use on the y axis, default is 'y'
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, y_read, x_read = data.dims
if (x_read is x) & (y_read is y):
pass
elif (x_read is y) & (y_read is x):
data = data.transpose(animate_over, y, x)
else:
raise ValueError(
"Dimensions {} or {} are not present in the data".format(x, y))
# 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 'max':
vmax = np.max(image_data)
if vmin is 'min':
vmin = np.min(image_data)
if not ax:
fig, ax = plt.subplots()
imshow_block = amp.blocks.Imshow(image_data,
vmin=vmin,
vmax=vmax,
ax=ax,
origin='lower',
**kwargs)
timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps)
if animate:
anim = amp.Animation([imshow_block], timeline)
cbar = plt.colorbar(imshow_block.im, ax=ax)
cbar.ax.set_ylabel(variable)
# Add title and axis labels
ax.set_title("{} variation over {}".format(variable, animate_over))
ax.set_xlabel(x)
ax.set_ylabel(y)
# Plot separatrix
if sep_pos:
ax = plot_separatrix(data, sep_pos, ax)
if animate:
anim.controls(timeline_slider_args={'text': animate_over})
if not save_as:
save_as = "{}_over_{}".format(variable, animate_over)
# TODO save using PillowWriter instead once matplotlib 3.1 comes out
# see https://github.com/t-makaro/animatplot/issues/24
anim.save(save_as + '.gif', writer='imagemagick')
return imshow_block
def make_average_temp(self,
output_filename='average_temp',
tmin=None,
tmax=None,
dpi=100,
**kwargs):
datlist = self.datlist
X, Y = self.X, self.Y
U3d, V3d = self.U3d, self.V3d
Ts = self.XYUVT['T']
prev_rect = set()
list_xy = []
count = []
for_average_plot = {}
def onclick(event):
count.append(1)
x, y = int(round(event.xdata)), int(round(event.ydata))
print(event.button, event.x, event.y, x, y)
if event.button == 3:
while prev_rect:
prev_rect.pop().remove()
while list_xy:
list_xy.pop()
while count:
count.pop()
return
if len(count) == 1:
r = patches.Rectangle((x - 0.5, y - 0.5),
width=1,
height=1,
ec='gray',
fill=False,
lw=3)
ax.add_patch(r)
fig.canvas.draw()
prev_rect.add(r)
list_xy.append((x, y))
elif len(count) == 2:
prev_x, prev_y = list_xy[0]
if prev_x < x:
from_x = prev_x - 0.5
to_x = x + 0.5
else:
from_x = prev_x + 0.5
to_x = x - 0.5
if prev_y < y:
from_y = prev_y - 0.5
to_y = y + 0.5
else:
from_y = prev_y + 0.5
to_y = y - 0.5
r = patches.Rectangle((from_x, from_y),
width=to_x - from_x,
height=to_y - from_y,
ec='k',
fill=False,
lw=3)
prev_rect.pop().remove()
ax.add_patch(r)
fig.canvas.draw()
prev_rect.add(r)
list_xy.append((x, y))
elif len(count) == 3:
xs = sorted([l[0] for l in list_xy])
ys = sorted([l[1] for l in list_xy])
average_temp = [
np.average(T[ys[0]:ys[1] + 1, xs[0]:xs[1] + 1]) for T in Ts
]
for_average_plot['temps'] = average_temp
r = prev_rect.pop()
x, y, w, h = r.get_x(), r.get_y(), r.get_width(), r.get_height(
)
r.remove()
r = patches.Rectangle((x, y),
width=w,
height=h,
ec='red',
fill=False,
lw=3)
ax.add_patch(r)
fig.canvas.draw()
prev_rect.add(r)
for_average_plot['rect'] = patches.Rectangle((x, y),
width=w,
height=h,
ec='black',
fill=False,
lw=3)
fig.canvas.mpl_disconnect(cid)
plt.close()
anim, fig, ax = self.to_gif(dpi=70, **kwargs)
vmin, vmax = self.vmin, self.vmax
fps = anim.timeline.fps
cid = fig.canvas.mpl_connect('button_press_event', onclick)
plt.show()
try:
average_temp = for_average_plot['temps']
r = for_average_plot['rect']
except KeyError:
print('Canceled')
return
# plot Another Figure
fig2, (ax2, ax3) = plt.subplots(1, 2, figsize=(10, 5), dpi=dpi)
ax2.add_patch(r)
arrows = amp.blocks.Quiver(X, Y, U3d, V3d, ax=ax2, t_axis=2, color='k')
blocks = amp.blocks.Imshow(Ts,
ax=ax2,
cmap='jet',
vmin=vmin,
vmax=vmax,
interpolation='none')
divider = mpl_toolkits.axes_grid1.make_axes_locatable(ax2)
cax = divider.append_axes('right', '5%', pad='3%')
cbar = fig2.colorbar(blocks.im, ax=ax2, cax=cax)
timeline2 = amp.Timeline(anim.timeline.t, fps=fps)
timearray2 = timeline2.t
ax2.tick_params(labelbottom=False, bottom=False) # x軸の削除
ax2.tick_params(labelleft=False, left=False) # y軸の削除
ax2.set_xticklabels([])
fig2.tight_layout()
times, temps = amp.util.parametric_line(timearray2, average_temp)
average_blocks = amp.blocks.Line(times, temps, ax=ax3, color='k', lw=3)
ax3.set_xlim(timearray2[0], timearray2[-1])
if tmin is None:
tmin = vmin
if tmax is None:
tmax = vmax
ax3.set_ylim(tmin, tmax)
ax3.set_xlabel('ICYCLE')
ax3.set_ylabel('T average')
ax3.grid(True)
plt.subplots_adjust(wspace=0.4)
anim3 = amp.Animation([arrows, blocks, average_blocks], timeline2)
anim3.timeline_slider(text='ICYCLE', valfmt='%d')
anim3.save_gif(output_filename)
def AbrirGraficaLateralIntensidad (tamañoLadoX, x0, númeroTérminosN, tamañoLadoY, y0, númeroTérminosM, cómoCalor):
"""
Se genera una nueva ventana en la que se despliegan los resultados gráficos, tanto una vista lateral, como una gráfica de
intensidad del cálculo de difusión establecido.
Parámetros de la función:
------------------------
tamañoLadoX y tamañoLadoY: Dimensiones de la placa.
x0 y y0: Puntos donde se aplica el calor a la placa.
númeroTérminosN y númeroTérminosM: Cantidad de términos con los que se realizará el cálculo por series de Fourier.
cómoCalor: Variable que indica la forma en que se está aplicando el calor a la placa.
Salida de la función:
---------------------
Grafica con vista lateral y gráfica de intensidad de los cálculos realizados.
"""
# Se definen los puntos en los que se va a evaluar la ecuación.
puntosMalla = 30
# Se crea un arreglo con los valores para cada variable.
x = np.linspace (0, tamañoLadoX, puntosMalla)
y = np.linspace (0, tamañoLadoY, puntosMalla)
t = np.linspace (0,5,50)
# Se unifican las tres variables para la graficación de intensidad.
# Luego se unifican las variables posicionales para la graficación lateral.
X, Y, T= np.meshgrid (x, y, t)
X1, Y1 = np.meshgrid (x, y)
# Se define una variable pcolormesh_data la cual contiene el resultado de la difusión.
pcolormesh_data = CalculoDifusión (T, X, x0, tamañoLadoX, Y, y0, tamañoLadoY, númeroTérminosN, númeroTérminosM, cómoCalor)
fig, (ax1, ax2) = plt.subplots(1, 2)
# Graficación de intensidad.
for ax in [ax1, ax2]:
ax.set_aspect('equal')
ax.set_xlabel('x (m)')
ax2.set_ylabel('y (m)', labelpad=-5)
fig.suptitle(r'Evolución de la difusión '
"\n"
r'a través del tiempo')
ax2.set_title('Gráfica de Intensidad')
pcolormesh_block = amp.blocks.Pcolormesh(X[:,:,0], Y[:,:,0], pcolormesh_data,
ax=ax2, t_axis=2, vmin=0, vmax=5)
plt.colorbar(pcolormesh_block.quad)
timeline = amp.Timeline(t, fps=10)
anim = amp.Animation([pcolormesh_block], timeline)
# Graficación de vista lateral.
def actualizar (i):
ax1.clear ()
difusion = CalculoDifusión(t[i], X1, x0, tamañoLadoX, Y1, y0, tamañoLadoY, númeroTérminosN, númeroTérminosM, cómoCalor)
ax1.plot (x, difusion)
ax1.set_title('Vista lateral')
ax1.set_xlabel('y (m)')
ax1.set_ylabel('u (x, y, t)')
ax1.set_ylim([0,5])
ani = animation.FuncAnimation (fig,actualizar,range (len(t)), interval = 1)
plt.get_current_fig_manager().window.showMaximized ()
plt.show()
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
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 animate_pcolormesh(
data,
animate_over=None,
x=None,
y=None,
animate=True,
axis_coords=None,
vmin=None,
vmax=None,
vsymmetric=False,
logscale=False,
fps=10,
save_as=None,
ax=None,
cax=None,
aspect=None,
extend=None,
controls="both",
**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, defaults to the time dimension
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
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.
vsymmetric : bool, optional
If set to true, make the color-scale symmetric
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.
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'
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.
aspect : str or None, optional
Argument to set_aspect(), defaults to "auto"
extend : str or None, optional
Passed to fig.colorbar()
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 animation function
animatplot.blocks.Pcolormesh
Returns
-------
animation or block
If animate==True, returns an animatplot.Animation object, otherwise
returns an animatplot.blocks.Pcolormesh 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
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]
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"
x_values, x_label = _parse_coord_option(x, axis_coords, data)
y_values, y_label = _parse_coord_option(y, axis_coords, data)
data = data.transpose(animate_over, y, x, transpose_coords=True)
# 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
kwargs["norm"] = _create_norm(logscale, kwargs.get("norm", None), vmin,
vmax)
if not ax:
fig, ax = plt.subplots()
ax.set_aspect(aspect)
# 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:]
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,
)
pcolormesh_block = amp.blocks.Pcolormesh(x_values,
y_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([pcolormesh_block], timeline)
cbar = plt.colorbar(pcolormesh_block.quad, ax=ax, cax=cax, extend=extend)
if "long_name" in data.attrs:
cbar_label = data.long_name
else:
cbar_label = variable
if "units" in data.attrs:
cbar_label += f" [{data.units}]"
cbar.ax.set_ylabel(cbar_label)
# Add title and axis labels
ax.set_title(variable)
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
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 pcolormesh_block
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
#minu = minu if minu>-np.inf else -1
#maxu = maxu if maxu<np.inf else 1
#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)
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
# 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_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
def animate_list(
self,
variables,
animate_over=None,
save_as=None,
show=False,
fps=10,
nrows=None,
ncols=None,
poloidal_plot=False,
axis_coords=None,
subplots_adjust=None,
vmin=None,
vmax=None,
logscale=None,
titles=None,
aspect=None,
extend=None,
controls="both",
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, defaults to the time dimension
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
axis_coords : None, str, dict or list of None, str or 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'.
Only affects time coordinate for plots with poloidal_plot=True.
If a list is passed, it must have the same length as 'variables' and gives
the axis_coords setting for each plot individually.
The setting to use for the 'animate_over' coordinate can be passed in one or
more dict values, but must be the same in all dicts if given more than once.
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. Defaults to "equal" for poloidal
plots and "auto" for others.
extend : str or None, optional
Passed to fig.colorbar()
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.
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, per
variable if a sequence is given.
Returns
-------
animation
An animatplot.Animation object.
"""
if animate_over is None:
animate_over = self.metadata.get("bout_tdim", "t")
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.abc.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")
extend = _expand_list_arg(extend, "extend")
axis_coords = _expand_list_arg(axis_coords, "axis_coords")
for k in kwargs:
kwargs[k] = _expand_list_arg(kwargs[k], k)
blocks = []
def is_list(variable):
return (
isinstance(variable, list)
or isinstance(variable, tuple)
or isinstance(variable, set)
)
for i, subplot_args in enumerate(
zip(
variables,
axes,
poloidal_plot,
axis_coords,
vmin,
vmax,
logscale,
titles,
aspect,
extend,
)
):
(
v,
ax,
this_poloidal_plot,
this_axis_coords,
this_vmin,
this_vmax,
this_logscale,
this_title,
this_aspect,
this_extend,
) = subplot_args
this_kwargs = {k: v[i] for k, v in kwargs.items()}
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.1)
if is_list(v):
for i in range(len(v)):
if isinstance(v[i], str):
v[i] = self.data[v[i]]
# list of variables for one subplot only supported for line plots with 1
# dimension plus time
ndims = 2
dims = v[0].dims
if len(dims) != 2:
raise ValueError(
"Variables in sublist must be 2d - can only overlay line plots"
)
for w in v:
if not w.dims == dims:
raise ValueError(
f"All variables in sub-list must have same dimensions."
f"{v[0].name} had {v[0].dims} but {w.name} had {w.dims}."
)
else:
if isinstance(v, str):
v = self.data[v]
data = v.bout.data
ndims = len(data.dims)
ax.set_title(data.name)
if ndims == 2:
if not is_list(v):
blocks.append(
animate_line(
data=data,
ax=ax,
animate_over=animate_over,
animate=False,
axis_coords=this_axis_coords,
aspect=this_aspect,
**this_kwargs,
)
)
else:
for w in v:
blocks.append(
animate_line(
data=w,
ax=ax,
animate_over=animate_over,
animate=False,
axis_coords=this_axis_coords,
aspect=this_aspect,
label=w.name,
**this_kwargs,
)
)
legend = ax.legend()
legend.set_draggable(True)
# set 'v' to use for the timeline below
v = v[0]
elif ndims == 3:
if this_poloidal_plot:
var_blocks = animate_poloidal(
data,
ax=ax,
cax=cax,
animate_over=animate_over,
animate=False,
axis_coords=this_axis_coords,
vmin=this_vmin,
vmax=this_vmax,
logscale=this_logscale,
aspect=this_aspect,
extend=this_extend,
**this_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,
axis_coords=this_axis_coords,
vmin=this_vmin,
vmax=this_vmax,
logscale=this_logscale,
aspect=this_aspect,
extend=this_extend,
**this_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)
if np.all([a == "index" for a in axis_coords]):
time_opt = "index"
elif np.any([isinstance(a, dict) and animate_over in a for a in axis_coords]):
given_values = [
a[animate_over]
for a in axis_coords
if isinstance(a, dict) and animate_over in a
]
time_opt = given_values[0]
if len(given_values) > 1 and not np.all(
[v == time_opt for v in given_values[1:]]
):
raise ValueError(
f"Inconsistent axis_coords values given for animate_over "
f"coordinate ({animate_over}). Got {given_values}."
)
else:
time_opt = None
time_values, time_label = _parse_coord_option(animate_over, time_opt, self.data)
time_values, time_suffix = _normalise_time_coord(time_values)
timeline = amp.Timeline(time_values, fps=fps, units=time_suffix)
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)
_add_controls(anim, controls, time_label)
if save_as is not None:
anim.save(save_as + ".gif", writer=PillowWriter(fps=fps))
if show:
plt.show()
return anim
