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