From e43fc05e0ae2471e438ebe7111cf1a6d07f5d5cf Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 17 Jul 2019 23:27:40 +0100 Subject: [PATCH 01/42] Method to make multiple animations in one figure --- xbout/boutdataset.py | 49 ++++++++++++++++++++++++++++++++++++++- xbout/plotting/animate.py | 6 ++--- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 6726f9e5..e4ec1517 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -2,12 +2,15 @@ from pprint import pformat import configparser +import animatplot as amp +from matplotlib import pyplot as plt +import numpy as np from xarray import register_dataset_accessor, \ save_mfdataset, set_options, merge from dask.diagnostics import ProgressBar - from .load import _auto_open_mfboutdataset +from .plotting.animate import animate_imshow, animate_line # This code should run whenever any function from this module is imported @@ -239,6 +242,50 @@ def to_restart(self, savepath='.', nxpe=None, nype=None, save_mfdataset(restart_datasets, paths, compute=True) return + def animate_list(self, variables, animate_over='t', save_as=None, show=False, fps=10, + nrows=None, ncols=None): + 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: + assert nrows*ncols >= nvars, 'nrows*ncols not as big as nvar' + + fig, axes = plt.subplots(nrows, ncols, squeeze=False) + + blocks = [] + for v, ax in zip(variables, axes.flatten()): + 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)) + elif ndims == 3: + blocks.append(animate_imshow(data=data, ax=ax, animate_over=animate_over, + animate=False)) + else: + raise ValueError("Unsupported number of dimensions " + + str(ndims) + ". Dims are " + str(v.dims)) + + timeline = amp.Timeline(np.arange(variables[0].sizes[animate_over]), fps=fps) + anim = amp.Animation(blocks, timeline) + anim.controls(timeline_slider_args={'text': animate_over}) + + if save_as is not None: + anim.save(save_as + '.gif', writer='imagemagick') + + if show: + plt.show() + + return anim + def _find_time_dependent_vars(data): evolving_vars = set(var for var in data.data_vars if 't' in data[var].dims) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 64731215..bb2d9758 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -71,9 +71,8 @@ def animate_imshow(data, animate_over='t', x='x', y='y', animate=True, 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: + timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps) anim = amp.Animation([imshow_block], timeline) cbar = plt.colorbar(imshow_block.im, ax=ax) @@ -159,9 +158,8 @@ def animate_line(data, animate_over='t', animate=True, line_block = amp.blocks.Line(image_data, ax=ax, **kwargs) - timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps) - if animate: + timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps) anim = amp.Animation([line_block], timeline) # Add title and axis labels From 39e8fb9a3d6b9dee3f3bb86c05f12dc805a2795c Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 18 Jul 2019 16:57:16 +0100 Subject: [PATCH 02/42] Shorten titles in animate_imshow and animate_line Previously contained 'animate over <>', but the variable being animated over is shown on the time-slider anyway. Shortening the titles makes things clearer when multiple plots are animated at once. --- xbout/plotting/animate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index bb2d9758..a28ffd6a 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -79,7 +79,7 @@ def animate_imshow(data, animate_over='t', x='x', y='y', animate=True, cbar.ax.set_ylabel(variable) # Add title and axis labels - ax.set_title("{} variation over {}".format(variable, animate_over)) + ax.set_title(variable) ax.set_xlabel(x) ax.set_ylabel(y) @@ -163,7 +163,7 @@ def animate_line(data, animate_over='t', animate=True, anim = amp.Animation([line_block], timeline) # Add title and axis labels - ax.set_title("{} variation over {}".format(variable, animate_over)) + ax.set_title(variable) ax.set_xlabel(x_read) ax.set_ylabel(variable) From 5b27452d6af50dc905c0dff4019970bdca63f111 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 18 Jul 2019 17:51:51 +0100 Subject: [PATCH 03/42] Option to use a color-scale symmetric around 0 --- xbout/boutdataset.py | 6 +++--- xbout/plotting/animate.py | 5 ++++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index e4ec1517..17735b75 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -243,7 +243,7 @@ def to_restart(self, savepath='.', nxpe=None, nype=None, return def animate_list(self, variables, animate_over='t', save_as=None, show=False, fps=10, - nrows=None, ncols=None): + nrows=None, ncols=None, **kwargs): nvars = len(variables) if nrows is None and ncols is None: @@ -266,10 +266,10 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp if ndims == 2: blocks.append(animate_line(data=data, ax=ax, animate_over=animate_over, - animate=False)) + animate=False, **kwargs)) elif ndims == 3: blocks.append(animate_imshow(data=data, ax=ax, animate_over=animate_over, - animate=False)) + animate=False, **kwargs)) else: raise ValueError("Unsupported number of dimensions " + str(ndims) + ". Dims are " + str(v.dims)) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index a28ffd6a..b1874ed4 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -7,7 +7,7 @@ def animate_imshow(data, animate_over='t', x='x', y='y', animate=True, - vmin=None, vmax=None, fps=10, save_as=None, + vmin=None, vmax=None, vsymmetric=False, fps=10, save_as=None, sep_pos=None, ax=None, **kwargs): """ Plots a color plot which is animated with time over the specified @@ -64,6 +64,9 @@ def animate_imshow(data, animate_over='t', x='x', y='y', animate=True, 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() From c96a2da0f154cd939f16d4d02b7b4d7d058db002 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 19 Jul 2019 16:18:43 +0100 Subject: [PATCH 04/42] Raise exception instead of using 'assert' --- xbout/boutdataset.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 17735b75..4ee444dc 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -254,7 +254,8 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp elif ncols is None: ncols = int(np.ceil(nvars/nrows)) else: - assert nrows*ncols >= nvars, 'nrows*ncols not as big as nvar' + if nrows*ncols < nvars: + raise ValueError('Not enough rows*columns to fit all variables') fig, axes = plt.subplots(nrows, ncols, squeeze=False) From 13b4c9f2c73185b66f344df93d2caa1c9b122396 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 27 Jul 2019 19:42:19 +0200 Subject: [PATCH 05/42] Don't pass 'sep_pos' to animate_imshow --- xbout/boutdataarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index 49f066bc..8caf7d76 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -79,7 +79,7 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, print("{} data passed has {} dimensions - will use " "animatplot.blocks.Imshow()".format(variable, str(n_dims))) imshow_block = animate_imshow(data=data, animate_over=animate_over, - x=x, y=y, sep_pos=sep_pos, + x=x, y=y, animate=animate, fps=fps, save_as=save_as, ax=ax, **kwargs) return imshow_block From 8ddf637be5ae3a8813a08b5a5c26a375776eb7b3 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 13:49:38 +0100 Subject: [PATCH 06/42] Use pcolormesh instead of imshow for 2d animations --- xbout/__init__.py | 2 +- xbout/boutdataarray.py | 18 +++++++++--------- xbout/plotting/animate.py | 19 +++++++++---------- 3 files changed, 19 insertions(+), 20 deletions(-) diff --git a/xbout/__init__.py b/xbout/__init__.py index dbdd4ebf..0871017b 100644 --- a/xbout/__init__.py +++ b/xbout/__init__.py @@ -6,7 +6,7 @@ from .boutdataset import BoutDatasetAccessor from .boutdataarray import BoutDataArrayAccessor -from .plotting.animate import animate_imshow +from .plotting.animate import animate_pcolormesh from .plotting.utils import plot_separatrix from ._version import __version__ diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index 8caf7d76..fd00fefc 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -4,7 +4,7 @@ import xarray as xr from xarray import register_dataarray_accessor -from .plotting.animate import animate_imshow, animate_line +from .plotting.animate import animate_pcolormesh, animate_line from .plotting import plotfuncs @@ -44,8 +44,8 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, 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. + Currently only supports 2D+1 data, which it plots with animatplot's + wrapping of matplotlib's pcolormesh. Parameters ---------- @@ -77,12 +77,12 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, n_dims = len(data.dims) if n_dims == 3: print("{} data passed has {} dimensions - will use " - "animatplot.blocks.Imshow()".format(variable, str(n_dims))) - imshow_block = animate_imshow(data=data, animate_over=animate_over, - x=x, y=y, - animate=animate, fps=fps, - save_as=save_as, ax=ax, **kwargs) - return imshow_block + "animatplot.blocks.Pcolormesh()".format(variable, str(n_dims))) + pcolormesh_block = animate_pcolormesh(data=data, animate_over=animate_over, + x=x, y=y, + animate=animate, fps=fps, + save_as=save_as, ax=ax, **kwargs) + return pcolormesh_block else: raise ValueError( "Data passed has an unsupported number of dimensions " diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 94e71d53..ce8c6674 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -7,15 +7,15 @@ from matplotlib.animation import PillowWriter -def animate_imshow(data, animate_over='t', x=None, y=None, animate=True, +def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, vmin=None, vmax=None, vsymmetric=False, 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. + Currently only supports 2D+1 data, which it plots with animatplotlib's + wrapping of matplotlib's pcolormesh. Parameters ---------- @@ -74,8 +74,7 @@ def animate_imshow(data, animate_over='t', x=None, y=None, animate=True, raise ValueError("Dimension {} is not present in the data" .format(x)) y = spatial_dims[0] - # Use (y, x) here so we transpose by default for imshow - data = data.transpose(animate_over, y, x) + data = data.transpose(animate_over, x, y) # Load values eagerly otherwise for some reason the plotting takes # 100's of times longer - for some reason animatplot does not deal @@ -94,14 +93,14 @@ def animate_imshow(data, animate_over='t', x=None, y=None, animate=True, if not ax: fig, ax = plt.subplots() - imshow_block = amp.blocks.Imshow(image_data, vmin=vmin, vmax=vmax, - ax=ax, origin='lower', **kwargs) + pcolormesh_block = amp.blocks.Pcolormesh(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([imshow_block], timeline) + anim = amp.Animation([pcolormesh_block], timeline) - cbar = plt.colorbar(imshow_block.im, ax=ax) + cbar = plt.colorbar(pcolormesh_block.im, ax=ax) cbar.ax.set_ylabel(variable) # Add title and axis labels @@ -120,7 +119,7 @@ def animate_imshow(data, animate_over='t', x=None, y=None, animate=True, save_as = "{}_over_{}".format(variable, animate_over) anim.save(save_as + '.gif', writer=PillowWriter(fps=fps)) - return imshow_block + return pcolormesh_block def animate_line(data, animate_over='t', animate=True, From 34af96357b8370e49b23f9fc2bc96c717a474e72 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 14:00:32 +0100 Subject: [PATCH 07/42] Use animate_pcolormesh in animate_list --- xbout/boutdataset.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 7e262a22..b656ddfe 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -8,7 +8,7 @@ import numpy as np from dask.diagnostics import ProgressBar -from .plotting.animate import animate_imshow, animate_line +from .plotting.animate import animate_pcolormesh, animate_line @register_dataset_accessor('bout') @@ -189,10 +189,11 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp if ndims == 2: blocks.append(animate_line(data=data, ax=ax, animate_over=animate_over, - animate=False, **kwargs)) + animate=False, **kwargs)) elif ndims == 3: - blocks.append(animate_imshow(data=data, ax=ax, animate_over=animate_over, - animate=False, **kwargs)) + blocks.append(animate_pcolormesh(data=data, ax=ax, + animate_over=animate_over, + animate=False, **kwargs)) else: raise ValueError("Unsupported number of dimensions " + str(ndims) + ". Dims are " + str(v.dims)) From 408e5711741c6df1077039acf17fbc307743c338 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 19:32:28 +0200 Subject: [PATCH 08/42] Add support for passing time-varying data to _decompose_regions --- xbout/plotting/utils.py | 71 ++++++++++++++++++++++++----------------- 1 file changed, 41 insertions(+), 30 deletions(-) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 1064ce25..8cd3d992 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -41,18 +41,21 @@ def _decompose_regions(da): j11, j12, j21, j22, ix1, ix2, nin, _, ny, y_boundary_guards = _get_seps(da) regions = [] + x, y = da.dims[-2:] + other_dims = da.dims[:-2] + ystart = 0 # Y index to start the next section if j11 >= 0: # plot lower inner leg - region1 = da[:, ystart:(j11 + 1)] + region1 = da.isel(**{y:slice(ystart,(j11 + 1))}) yind = [j11, j22 + 1] - region2 = da[:ix1, yind] + region2 = da.isel(**{x:slice(0,ix1), y:yind}) - region3 = da[ix1:, j11: (j11 + 2)] + region3 = da.isel(**{x:slice(ix1, None), y:slice(j11, (j11 + 2))}) yind = [j22, j11 + 1] - region4 = da[:ix1, yind] + region4 = da.isel(**{x:slice(0,ix1), y:yind}) regions.extend([region1, region2, region3, region4]) @@ -60,7 +63,7 @@ def _decompose_regions(da): if j21 + 1 > ystart: # Inner SOL - region5 = da[:, ystart:(j21 + 1)] + region5 = da.isel(**{y:slice(ystart, (j21 + 1))}) regions.append(region5) ystart = j21 + 1 @@ -69,19 +72,19 @@ def _decompose_regions(da): # Contains upper PF region # Inner leg - region6 = da[ix1:, j21:(j21 + 2)] - region7 = da[:, ystart:nin] + region6 = da.isel(**{x:slice(ix1, None), y:slice(j21, (j21 + 2))}) + region7 = da.isel(**{y:slice(ystart, nin)}) # Outer leg - region8 = da[:, nin:(j12 + 1)] - region9 = da[ix1:, j12:(j12 + 2)] + region8 = da.isel(**{y:slice(nin, (j12 + 1))}) + region9 = da.isel(**{x:slice(ix1, None), y:slice(j12, (j12 + 2))}) yind = [j21, j12 + 1] - region10 = da[:ix1, yind] + region10 = da.isel(**{x:slice(0, ix1), y:yind}) yind = [j21 + 1, j12] - region11 = da[:ix1, yind] + region11 = da.isel(**{x:slice(0, ix1), y:yind}) regions.extend([region6, region7, region8, region9, region10, region11]) @@ -92,57 +95,65 @@ def _decompose_regions(da): if j22 + 1 > ystart: # Outer SOL - region12 = da[:, ystart:(j22 + 1)] + region12 = da.isel(**{y:slice(ystart, (j22 + 1))}) regions.append(region12) ystart = j22 + 1 if j22 + 1 < ny: # Outer leg - region13 = da[ix1:, j22:(j22 + 2)] - region14 = da[:, ystart:ny] + region13 = da.isel(**{x:slice(ix1, None), y:slice(j22, (j22 + 2))}) + region14 = da.isel(**{y:slice(ystart, ny)}) # X-point regions - corner1 = da[ix1 - 1, j11] - corner2 = da[ix1, j11] - corner3 = da[ix1, j11 + 1] - corner4 = da[ix1 - 1, j11 + 1] + corner1 = da.isel(**{x:ix1-1, y:j11}) + corner2 = da.isel(**{x:ix1, y:j11}) + corner3 = da.isel(**{x:ix1, y:j11+1}) + corner4 = da.isel(**{x:ix1-1, y:j11+1}) xregion_lower = xr.concat([corner1, corner2, corner3, corner4], dim='dim1') - corner5 = da[ix1 - 1, j22 + 1] - corner6 = da[ix1, j22 + 1] - corner7 = da[ix1, j22] - corner8 = da[ix1 - 1, j22] + corner5 = da.isel(**{x:ix1-1, y:j22+1}) + corner6 = da.isel(**{x:ix1, y:j22+1}) + corner7 = da.isel(**{x:ix1, y:j22}) + corner8 = da.isel(**{x:ix1-1, y:j22}) xregion_upper = xr.concat([corner5, corner6, corner7, corner8], dim='dim1') region15 = xr.concat([xregion_lower, xregion_upper], dim='dim2') + # re-arrange dimensions so that the new 'dim1' and 'dim2' are at the + # end - ensures that a time dimension stays at the beginning + region15 = region15.transpose(*other_dims, 'dim2', 'dim1') + regions.extend([region13, region14, region15]) if j21 > j11 and j12 > j21 and j22 > j12: # X-point regions - corner1 = da[ix1-1, j12] - corner2 = da[ix1, j12] - corner3 = da[ix1, j12+1] - corner4 = da[ix1-1, j12+1] + corner1 = da.isel(**{x:ix1-1, y:j12}) + corner2 = da.isel(**{x:ix1, y:j12}) + corner3 = da.isel(**{x:ix1, y:j12+1}) + corner4 = da.isel(**{x:ix1-1, y:j12+1}) xregion_lower = xr.concat([corner1, corner2, corner3, corner4], dim='dim1') - corner5 = da[ix1 - 1, j21 + 1] - corner6 = da[ix1, j21+1] - corner7 = da[ix1, j21] - corner8 = da[ix1 - 1, j21] + corner5 = da.isel(**{x:ix1-1, y:j21+1}) + corner6 = da.isel(**{x:ix1, y:j21+1}) + corner7 = da.isel(**{x:ix1, y:j21}) + corner8 = da.isel(**{x:ix1-1, y:j21}) xregion_upper = xr.concat([corner5, corner6, corner7, corner8], dim='dim1') region16 = xr.concat([xregion_lower, xregion_upper], dim='dim2') + # re-arrange dimensions so that the new 'dim1' and 'dim2' are at the + # end - ensures that a time dimension stays at the beginning + region16 = region16.transpose(*other_dims, 'dim2', 'dim1') + regions.append(region16) return regions From 7e3619d67adeb656d4bac9198e281fae7d452a83 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 13:54:25 +0100 Subject: [PATCH 09/42] Add option/function to make animated poloidal plots (in R,Z coordinates) --- xbout/__init__.py | 2 +- xbout/boutdataarray.py | 37 ++++++++----- xbout/plotting/animate.py | 107 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 13 deletions(-) diff --git a/xbout/__init__.py b/xbout/__init__.py index 0871017b..02ae11a6 100644 --- a/xbout/__init__.py +++ b/xbout/__init__.py @@ -6,7 +6,7 @@ from .boutdataset import BoutDatasetAccessor from .boutdataarray import BoutDataArrayAccessor -from .plotting.animate import animate_pcolormesh +from .plotting.animate import animate_pcolormesh, animate_poloidal from .plotting.utils import plot_separatrix from ._version import __version__ diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index fd00fefc..b664f4e0 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -4,7 +4,7 @@ import xarray as xr from xarray import register_dataarray_accessor -from .plotting.animate import animate_pcolormesh, animate_line +from .plotting.animate import animate_poloidal, animate_pcolormesh, animate_line from .plotting import plotfuncs @@ -38,8 +38,8 @@ def __str__(self): text += "Options:\n{}".format(styled(self.options)) return text - def animate2D(self, animate_over='t', x=None, y=None, animate=True, - fps=10, save_as=None, sep_pos=None, ax=None, **kwargs): + def animate2D(self, animate_over='t', x=None, y=None, animate=True, fps=10, + save_as=None, ax=None, poloidal_plot=False, **kwargs): """ Plots a color plot which is animated with time over the specified coordinate. @@ -63,10 +63,11 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, Frames per second of resulting gif save_as : str, optional Filename to give to the resulting gif - sep_pos : int, optional - Position along the 'x' dimension to plot the separatrix ax : matplotlib.pyplot.axes object, optional Axis on which to plot the gif + poloidal_plot : bool, optional + Use animate_poloidal to make a plot in R-Z coordinates (input field must be + (t,x,y)) kwargs : dict, optional Additional keyword arguments are passed on to the plotting function (e.g. imshow for 2D plots). @@ -76,13 +77,25 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, variable = data.name n_dims = len(data.dims) if n_dims == 3: - print("{} data passed has {} dimensions - will use " - "animatplot.blocks.Pcolormesh()".format(variable, str(n_dims))) - pcolormesh_block = animate_pcolormesh(data=data, animate_over=animate_over, - x=x, y=y, - animate=animate, fps=fps, - save_as=save_as, ax=ax, **kwargs) - return pcolormesh_block + if poloidal_plot: + print("{} data passed has {} dimensions - making poloidal plot with " + "animate_poloidal()".format(variable, str(n_dims))) + if x is not None: + kwargs['x'] = x + if y is not None: + kwargs['y'] = y + poloidal_blocks = animate_poloidal(data, animate_over=animate_over, + animate=animate, fps=fps, + save_as=save_as, ax=ax, **kwargs) + return poloidal_blocks + else: + print("{} data passed has {} dimensions - will use " + "animatplot.blocks.Pcolormesh()".format(variable, str(n_dims))) + pcolormesh_block = animate_pcolormesh(data=data, animate_over=animate_over, + x=x, y=y, + animate=animate, fps=fps, + save_as=save_as, ax=ax, **kwargs) + return pcolormesh_block else: raise ValueError( "Data passed has an unsupported number of dimensions " diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index ce8c6674..fe131ab7 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -1,8 +1,115 @@ import numpy as np +import matplotlib import matplotlib.pyplot as plt import animatplot as amp +from .utils import _decompose_regions, plot_separatrices, plot_targets + + +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, **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)") + + # TODO work out how to auto-set the aspect ratio of the plot correctly + height = da.coords[y].max() - da.coords[y].min() + width = da.coords[x].max() - da.coords[x].min() + aspect = height / width + + if ax is None: + fig, ax = plt.subplots() + + 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 = 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) + + 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 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) + + 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='imagemagick') + + return blocks + from .utils import plot_separatrix from matplotlib.animation import PillowWriter From 2ab9fea7e716d4db3aab322b9a1ac11fc515f26e Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 18:39:26 +0100 Subject: [PATCH 10/42] Support for poloidal plots in animate_list --- xbout/boutdataset.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index b656ddfe..74e124b4 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -165,7 +165,7 @@ def to_restart(self, savepath='.', nxpe=None, nype=None, return def animate_list(self, variables, animate_over='t', save_as=None, show=False, fps=10, - nrows=None, ncols=None, **kwargs): + nrows=None, ncols=None, poloidal_plot=False, **kwargs): nvars = len(variables) if nrows is None and ncols is None: @@ -191,9 +191,14 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp blocks.append(animate_line(data=data, ax=ax, animate_over=animate_over, animate=False, **kwargs)) elif ndims == 3: - blocks.append(animate_pcolormesh(data=data, ax=ax, - animate_over=animate_over, - animate=False, **kwargs)) + if poloidal_plot: + blocks.append(animate_poloidal(data, ax=ax, + animate_over=animate_over, + animate=False, **kwargs)) + else: + blocks.append(animate_pcolormesh(data=data, ax=ax, + animate_over=animate_over, + animate=False, **kwargs)) else: raise ValueError("Unsupported number of dimensions " + str(ndims) + ". Dims are " + str(v.dims)) From 2647995b011ee290feb726e145c606b8bd6b0525 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 22:04:58 +0200 Subject: [PATCH 11/42] Fix animate_pcolormesh animatplot's Pcolormesh seems to need x- and y-value arrays to be passed explicitly, otherwise the animations looked strange. --- xbout/plotting/animate.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index fe131ab7..1201a4ce 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -181,7 +181,7 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, raise ValueError("Dimension {} is not present in the data" .format(x)) y = spatial_dims[0] - data = data.transpose(animate_over, x, y) + 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 @@ -200,7 +200,12 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, if not ax: fig, ax = plt.subplots() - pcolormesh_block = amp.blocks.Pcolormesh(image_data, vmin=vmin, vmax=vmax, ax=ax, + # 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: From dabf6a69bd6adaed3b391b4c91379673bb4bded8 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 22:28:15 +0100 Subject: [PATCH 12/42] Import animate_poloidal to fix poloidal_plots option of animate_list --- xbout/boutdataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 74e124b4..72ea1cd3 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -8,7 +8,7 @@ import numpy as np from dask.diagnostics import ProgressBar -from .plotting.animate import animate_pcolormesh, animate_line +from .plotting.animate import animate_poloidal, animate_pcolormesh, animate_line @register_dataset_accessor('bout') From 9eebf99a1f52812532433729a8d046a1fbcd9664 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 22:29:15 +0100 Subject: [PATCH 13/42] Fix colorbar of animate_poloidal when passing in axes object as ax --- xbout/plotting/animate.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 1201a4ce..fe35537a 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -59,6 +59,8 @@ def animate_poloidal(da, *, ax=None, animate_over = 't', separatrix=True, target if ax is None: fig, ax = plt.subplots() + else: + fig = ax.get_figure() if vmin is None: vmin = da.min().values @@ -74,7 +76,7 @@ def animate_poloidal(da, *, ax=None, animate_over = 't', separatrix=True, target sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) sm.set_array([]) cmap = sm.get_cmap() - fig.colorbar(sm) + fig.colorbar(sm, ax=ax) regions = _decompose_regions(da) From 25d4836b9d8773ea5fdc8657a643995d94c5a026 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 22:31:55 +0100 Subject: [PATCH 14/42] Allow variable names as well as BoutDataArrays in animate_list variables --- xbout/boutdataset.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 72ea1cd3..f1518433 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -166,6 +166,17 @@ def to_restart(self, savepath='.', nxpe=None, nype=None, def animate_list(self, variables, animate_over='t', save_as=None, show=False, fps=10, nrows=None, ncols=None, poloidal_plot=False, **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. + """ + nvars = len(variables) if nrows is None and ncols is None: @@ -183,6 +194,10 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp blocks = [] for v, ax in zip(variables, axes.flatten()): + + if isinstance(v, str): + v = self.data[v] + data = v.bout.data ndims = len(data.dims) ax.set_title(data.name) @@ -203,7 +218,7 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp raise ValueError("Unsupported number of dimensions " + str(ndims) + ". Dims are " + str(v.dims)) - timeline = amp.Timeline(np.arange(variables[0].sizes[animate_over]), fps=fps) + timeline = amp.Timeline(np.arange(v.sizes[animate_over]), fps=fps) anim = amp.Animation(blocks, timeline) anim.controls(timeline_slider_args={'text': animate_over}) From aff3e8a388db2f089bb81a3a1a2c7dc61d81fe7f Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Aug 2019 23:13:11 +0100 Subject: [PATCH 15/42] Fix appending of blocks when using animate_poloidal in animate_list --- xbout/boutdataset.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index f1518433..52c4827d 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -207,9 +207,11 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp animate=False, **kwargs)) elif ndims == 3: if poloidal_plot: - blocks.append(animate_poloidal(data, ax=ax, - animate_over=animate_over, - animate=False, **kwargs)) + var_blocks = animate_poloidal(data, ax=ax, + animate_over=animate_over, + animate=False, **kwargs) + for block in var_blocks: + blocks.append(block) else: blocks.append(animate_pcolormesh(data=data, ax=ax, animate_over=animate_over, From a9294453f8df1619204527e0047066d3fc5deab8 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 6 Aug 2019 00:04:09 +0100 Subject: [PATCH 16/42] Set equal aspect ratio in animate_poloidal --- xbout/plotting/animate.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index fe35537a..a7cf8967 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -78,6 +78,8 @@ def animate_poloidal(da, *, ax=None, animate_over = 't', separatrix=True, target cmap = sm.get_cmap() fig.colorbar(sm, ax=ax) + ax.set_aspect('equal') + regions = _decompose_regions(da) # Plot all regions on same axis From 0c4df3e76c4783866c44a63d9e57c35ff8c52489 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 6 Aug 2019 00:04:30 +0100 Subject: [PATCH 17/42] subplots_adjust option for animate_list --- xbout/boutdataset.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 52c4827d..8812224f 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -165,7 +165,8 @@ def to_restart(self, savepath='.', nxpe=None, nype=None, return def animate_list(self, variables, animate_over='t', save_as=None, show=False, fps=10, - nrows=None, ncols=None, poloidal_plot=False, **kwargs): + nrows=None, ncols=None, poloidal_plot=False, + subplots_adjust=None, **kwargs): """ Parameters ---------- @@ -192,6 +193,9 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp fig, axes = plt.subplots(nrows, ncols, squeeze=False) + if subplots_adjust is not None: + fig.subplots_adjust(**subplots_adjust) + blocks = [] for v, ax in zip(variables, axes.flatten()): From 9eca9a2ce12a6f56ad722d9653c84cc68732ba3a Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 6 Aug 2019 00:44:20 +0100 Subject: [PATCH 18/42] Arguments to make controls optional in animation methods --- xbout/plotting/animate.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index a7cf8967..c385c532 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -9,7 +9,7 @@ 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, **kwargs): + 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). @@ -106,7 +106,8 @@ def animate_poloidal(da, *, ax=None, animate_over = 't', separatrix=True, target timeline = amp.Timeline(np.arange(da.sizes[animate_over]), fps=fps) anim = amp.Animation(blocks, timeline) - anim.controls(timeline_slider_args={'text': animate_over}) + if controls: + anim.controls(timeline_slider_args={'text': animate_over}) if not save_as: save_as = "{}_over_{}".format(da.name, animate_over) @@ -120,7 +121,7 @@ def animate_poloidal(da, *, ax=None, animate_over = 't', separatrix=True, target def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, vmin=None, vmax=None, vsymmetric=False, fps=10, save_as=None, - sep_pos=None, ax=None, **kwargs): + ax=None, controls=True, **kwargs): """ Plots a color plot which is animated with time over the specified coordinate. @@ -229,7 +230,8 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, ax = plot_separatrix(data, sep_pos, ax) if animate: - anim.controls(timeline_slider_args={'text': animate_over}) + if controls: + anim.controls(timeline_slider_args={'text': animate_over}) if not save_as: save_as = "{}_over_{}".format(variable, animate_over) @@ -240,7 +242,7 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, def animate_line(data, animate_over='t', animate=True, vmin=None, vmax=None, fps=10, save_as=None, sep_pos=None, ax=None, - **kwargs): + controls=True, **kwargs): """ Plots a line plot which is animated with time. @@ -311,7 +313,8 @@ def animate_line(data, animate_over='t', animate=True, ax.plot_vline(sep_pos, '--') if animate: - anim.controls(timeline_slider_args={'text': animate_over}) + if controls: + anim.controls(timeline_slider_args={'text': animate_over}) if not save_as: save_as = "{}_over_{}".format(variable, animate_over) From 40801968c165eb279a3107372c6c52c98443c81c Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 6 Aug 2019 11:14:59 +0100 Subject: [PATCH 19/42] Fix colorbar creation in animate_pcolormesh The block returned from animatplot's Pcolormesh has a 'quad' attribute instead of 'im' ('im' is used by Imshow). --- xbout/plotting/animate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index c385c532..945fc244 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -217,7 +217,7 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps) anim = amp.Animation([pcolormesh_block], timeline) - cbar = plt.colorbar(pcolormesh_block.im, ax=ax) + cbar = plt.colorbar(pcolormesh_block.quad, ax=ax) cbar.ax.set_ylabel(variable) # Add title and axis labels From f63c0acaddacc2d94ccc7fc8ce51926ac580fbee Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 20 Aug 2019 12:59:39 +0100 Subject: [PATCH 20/42] Remove sep_pos from animate_pcolormesh --- xbout/plotting/animate.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 945fc244..598c1115 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -146,8 +146,6 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, 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 @@ -225,10 +223,6 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, ax.set_xlabel(x) ax.set_ylabel(y) - # Plot separatrix - if sep_pos: - ax = plot_separatrix(data, sep_pos, ax) - if animate: if controls: anim.controls(timeline_slider_args={'text': animate_over}) From 5286636420f0d23faae32f102257d9123a248036 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 23 Aug 2019 13:37:19 +0100 Subject: [PATCH 21/42] Option for plotting gridlines in plot2d_wrapper --- xbout/plotting/plotfuncs.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 430e939d..4d3a476c 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1,3 +1,5 @@ +import collections + import matplotlib import matplotlib.pyplot as plt import numpy as np @@ -35,8 +37,8 @@ def regions(da, ax=None, **kwargs): def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, - add_limiter_hatching=True, cmap=None, vmin=None, vmax=None, - aspect=None, **kwargs): + add_limiter_hatching=True, gridlines=None, cmap=None, vmin=None, + vmax=None, aspect=None, **kwargs): """ Make a 2D plot using an xarray method, taking into account branch cuts (X-points). @@ -57,6 +59,9 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, Draw solid lines at the target surfaces add_limiter_hatching : bool, optional Draw hatched areas at the targets + gridlines : bool or int, optional + If True, draw grid lines on the plot. If an int is passed, it is used as the + stride when plotting grid lines (to reduce the number on the plot) cmap : Matplotlib colormap, optional Color map to use for the plot vmin : float, optional @@ -169,6 +174,24 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, extend = kwargs.get('extend', 'neither') fig.colorbar(artists[0], ax=ax, extend=extend) + if gridlines is not None: + if gridlines is True: + gridlines = (1, 1) + if not isinstance(gridlines, collections.Sequence): + gridlines = (gridlines, gridlines) + R_global = da['R'] + R_global.attrs['metadata'] = da.metadata + + Z_global = da['Z'] + Z_global.attrs['metadata'] = da.metadata + + R_regions = _decompose_regions(da['R']) + Z_regions = _decompose_regions(da['Z']) + + for R, Z in zip(R_regions, Z_regions): + plt.plot(R[::gridlines[0], :].T, Z[::gridlines[0], :].T, color='k', lw=0.1) + plt.plot(R[:, ::gridlines[1]], Z[:, ::gridlines[1]], color='k', lw=0.1) + ax.set_title(da.name) if separatrix: From 35deefa63773b28fffd96d1eb964b80b278903b2 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 28 Nov 2019 10:56:11 +0000 Subject: [PATCH 22/42] Allow per-variable vmin and vmax arguments to animate_list --- xbout/boutdataset.py | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 8812224f..8057a01b 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -166,7 +166,7 @@ def to_restart(self, savepath='.', nxpe=None, nype=None, 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, **kwargs): + subplots_adjust=None, vmin=None, vmax=None, **kwargs): """ Parameters ---------- @@ -176,6 +176,10 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp 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. + 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 """ nvars = len(variables) @@ -196,8 +200,22 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp if subplots_adjust is not None: fig.subplots_adjust(**subplots_adjust) + try: + if len(vmin) != len(variables): + raise ValueError('if vmin is a sequence, it must have the same number ' + 'of elements as "variables"') + except TypeError: + vmin = [vmin] * len(variables) + + try: + if len(vmax) != len(variables): + raise ValueError('if vmin is a sequence, it must have the same number ' + 'of elements as "variables"') + except TypeError: + vmax = [vmax] * len(variables) + blocks = [] - for v, ax in zip(variables, axes.flatten()): + for v, ax, this_vmin, this_vmax in zip(variables, axes.flatten(), vmin, vmax): if isinstance(v, str): v = self.data[v] @@ -213,13 +231,15 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp if poloidal_plot: var_blocks = animate_poloidal(data, ax=ax, animate_over=animate_over, - animate=False, **kwargs) + animate=False, vmin=this_vmin, + vmax=this_vmax, **kwargs) for block in var_blocks: blocks.append(block) else: blocks.append(animate_pcolormesh(data=data, ax=ax, animate_over=animate_over, - animate=False, **kwargs)) + animate=False, vmin=this_vmin, + vmax=this_vmax, **kwargs)) else: raise ValueError("Unsupported number of dimensions " + str(ndims) + ". Dims are " + str(v.dims)) From 43c2990ba5e602b3da54a254d8aa732520d9fac3 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 4 Dec 2019 13:08:38 +0000 Subject: [PATCH 23/42] Fix test_animate2D Needs to check for Pcolormesh instead of Imshow. --- xbout/tests/test_animate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xbout/tests/test_animate.py b/xbout/tests/test_animate.py index 6a6a6f2e..4a222f07 100644 --- a/xbout/tests/test_animate.py +++ b/xbout/tests/test_animate.py @@ -4,7 +4,7 @@ from xbout.boutdataarray import BoutDataArrayAccessor from .test_load import create_bout_ds_list -from animatplot.blocks import Imshow, Line +from animatplot.blocks import Pcolormesh, Line @pytest.fixture @@ -34,7 +34,7 @@ def test_animate2D(self, create_test_file): save_dir, ds = create_test_file animation = ds['n'].isel(y=1).bout.animate2D(y='z', save_as="%s/test" % save_dir) - assert isinstance(animation, Imshow) + assert isinstance(animation, Pcolormesh) def test_animate1D(self, create_test_file): From 37c895811ad13c719f10106545d21251a0818e9b Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 4 Dec 2019 13:14:41 +0000 Subject: [PATCH 24/42] PEP8 fixes --- xbout/boutdataarray.py | 4 +-- xbout/plotting/animate.py | 6 ++-- xbout/plotting/utils.py | 60 +++++++++++++++++++-------------------- 3 files changed, 35 insertions(+), 35 deletions(-) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index b664f4e0..ef8b1137 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -91,8 +91,8 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, fps=10, else: print("{} data passed has {} dimensions - will use " "animatplot.blocks.Pcolormesh()".format(variable, str(n_dims))) - pcolormesh_block = animate_pcolormesh(data=data, animate_over=animate_over, - x=x, y=y, + pcolormesh_block = animate_pcolormesh(data=data, + animate_over=animate_over, x=x, y=y, animate=animate, fps=fps, save_as=save_as, ax=ax, **kwargs) return pcolormesh_block diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 598c1115..e5f170da 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -7,7 +7,7 @@ from .utils import _decompose_regions, plot_separatrices, plot_targets -def animate_poloidal(da, *, ax=None, animate_over = 't', separatrix=True, targets=True, +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): """ @@ -120,8 +120,8 @@ def animate_poloidal(da, *, ax=None, animate_over = 't', separatrix=True, target 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): + 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. diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 8cd3d992..eeaecb6b 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -47,15 +47,15 @@ def _decompose_regions(da): ystart = 0 # Y index to start the next section if j11 >= 0: # plot lower inner leg - region1 = da.isel(**{y:slice(ystart,(j11 + 1))}) + region1 = da.isel(**{y: slice(ystart, (j11 + 1))}) yind = [j11, j22 + 1] - region2 = da.isel(**{x:slice(0,ix1), y:yind}) + region2 = da.isel(**{x: slice(0, ix1), y:yind}) - region3 = da.isel(**{x:slice(ix1, None), y:slice(j11, (j11 + 2))}) + region3 = da.isel(**{x: slice(ix1, None), y: slice(j11, (j11 + 2))}) yind = [j22, j11 + 1] - region4 = da.isel(**{x:slice(0,ix1), y:yind}) + region4 = da.isel(**{x: slice(0, ix1), y: yind}) regions.extend([region1, region2, region3, region4]) @@ -63,7 +63,7 @@ def _decompose_regions(da): if j21 + 1 > ystart: # Inner SOL - region5 = da.isel(**{y:slice(ystart, (j21 + 1))}) + region5 = da.isel(**{y: slice(ystart, (j21 + 1))}) regions.append(region5) ystart = j21 + 1 @@ -72,19 +72,19 @@ def _decompose_regions(da): # Contains upper PF region # Inner leg - region6 = da.isel(**{x:slice(ix1, None), y:slice(j21, (j21 + 2))}) - region7 = da.isel(**{y:slice(ystart, nin)}) + region6 = da.isel(**{x: slice(ix1, None), y: slice(j21, (j21 + 2))}) + region7 = da.isel(**{y: slice(ystart, nin)}) # Outer leg - region8 = da.isel(**{y:slice(nin, (j12 + 1))}) - region9 = da.isel(**{x:slice(ix1, None), y:slice(j12, (j12 + 2))}) + region8 = da.isel(**{y: slice(nin, (j12 + 1))}) + region9 = da.isel(**{x: slice(ix1, None), y: slice(j12, (j12 + 2))}) yind = [j21, j12 + 1] - region10 = da.isel(**{x:slice(0, ix1), y:yind}) + region10 = da.isel(**{x: slice(0, ix1), y: yind}) yind = [j21 + 1, j12] - region11 = da.isel(**{x:slice(0, ix1), y:yind}) + region11 = da.isel(**{x: slice(0, ix1), y: yind}) regions.extend([region6, region7, region8, region9, region10, region11]) @@ -95,29 +95,29 @@ def _decompose_regions(da): if j22 + 1 > ystart: # Outer SOL - region12 = da.isel(**{y:slice(ystart, (j22 + 1))}) + region12 = da.isel(**{y: slice(ystart, (j22 + 1))}) regions.append(region12) ystart = j22 + 1 if j22 + 1 < ny: # Outer leg - region13 = da.isel(**{x:slice(ix1, None), y:slice(j22, (j22 + 2))}) - region14 = da.isel(**{y:slice(ystart, ny)}) + region13 = da.isel(**{x: slice(ix1, None), y: slice(j22, (j22 + 2))}) + region14 = da.isel(**{y: slice(ystart, ny)}) # X-point regions - corner1 = da.isel(**{x:ix1-1, y:j11}) - corner2 = da.isel(**{x:ix1, y:j11}) - corner3 = da.isel(**{x:ix1, y:j11+1}) - corner4 = da.isel(**{x:ix1-1, y:j11+1}) + corner1 = da.isel(**{x: ix1-1, y: j11}) + corner2 = da.isel(**{x: ix1, y: j11}) + corner3 = da.isel(**{x: ix1, y: j11+1}) + corner4 = da.isel(**{x: ix1-1, y: j11+1}) xregion_lower = xr.concat([corner1, corner2, corner3, corner4], dim='dim1') - corner5 = da.isel(**{x:ix1-1, y:j22+1}) - corner6 = da.isel(**{x:ix1, y:j22+1}) - corner7 = da.isel(**{x:ix1, y:j22}) - corner8 = da.isel(**{x:ix1-1, y:j22}) + corner5 = da.isel(**{x: ix1-1, y: j22+1}) + corner6 = da.isel(**{x: ix1, y: j22+1}) + corner7 = da.isel(**{x: ix1, y: j22}) + corner8 = da.isel(**{x: ix1-1, y: j22}) xregion_upper = xr.concat([corner5, corner6, corner7, corner8], dim='dim1') @@ -132,18 +132,18 @@ def _decompose_regions(da): if j21 > j11 and j12 > j21 and j22 > j12: # X-point regions - corner1 = da.isel(**{x:ix1-1, y:j12}) - corner2 = da.isel(**{x:ix1, y:j12}) - corner3 = da.isel(**{x:ix1, y:j12+1}) - corner4 = da.isel(**{x:ix1-1, y:j12+1}) + corner1 = da.isel(**{x: ix1-1, y: j12}) + corner2 = da.isel(**{x: ix1, y: j12}) + corner3 = da.isel(**{x: ix1, y: j12+1}) + corner4 = da.isel(**{x: ix1-1, y: j12+1}) xregion_lower = xr.concat([corner1, corner2, corner3, corner4], dim='dim1') - corner5 = da.isel(**{x:ix1-1, y:j21+1}) - corner6 = da.isel(**{x:ix1, y:j21+1}) - corner7 = da.isel(**{x:ix1, y:j21}) - corner8 = da.isel(**{x:ix1-1, y:j21}) + corner5 = da.isel(**{x: ix1-1, y: j21+1}) + corner6 = da.isel(**{x: ix1, y: j21+1}) + corner7 = da.isel(**{x: ix1, y: j21}) + corner8 = da.isel(**{x: ix1-1, y: j21}) xregion_upper = xr.concat([corner5, corner6, corner7, corner8], dim='dim1') From 8ed2268ad4a02e826259063c62c9d94beb40099f Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 4 Dec 2019 22:02:29 +0000 Subject: [PATCH 25/42] Remove lines calculating unused 'aspect' --- xbout/plotting/animate.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index e5f170da..de67fef6 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -52,11 +52,6 @@ def animate_poloidal(da, *, ax=None, animate_over='t', separatrix=True, targets= if len(da.dims) != 3: raise ValueError("da must be 2+1D (t,x,y)") - # TODO work out how to auto-set the aspect ratio of the plot correctly - height = da.coords[y].max() - da.coords[y].min() - width = da.coords[x].max() - da.coords[x].min() - aspect = height / width - if ax is None: fig, ax = plt.subplots() else: From bfba4379db1ce63b3e5e47aac7ea6013e545394c Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 4 Dec 2019 22:15:27 +0000 Subject: [PATCH 26/42] Fix merge of test_animate2D --- xbout/tests/test_animate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xbout/tests/test_animate.py b/xbout/tests/test_animate.py index 76243af8..7b3fef16 100644 --- a/xbout/tests/test_animate.py +++ b/xbout/tests/test_animate.py @@ -35,7 +35,7 @@ def test_animate2D(self, create_test_file): animation = ds['n'].isel(x=1).bout.animate2D(save_as="%s/testyz" % save_dir) - assert isinstance(animation, Imshow) + assert isinstance(animation, Pcolormesh) assert animation.ax.get_xlabel() == 'y' assert animation.ax.get_ylabel() == 'z' @@ -48,7 +48,7 @@ def test_animate2D(self, create_test_file): animation = ds['n'].isel(z=3).bout.animate2D(save_as="%s/testxy" % save_dir) - assert isinstance(animation, Imshow) + assert isinstance(animation, Pcolormesh) assert animation.ax.get_xlabel() == 'x' assert animation.ax.get_ylabel() == 'y' From 54e533dc18f19c8657ace79db33e3f43a4998f38 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 5 Dec 2019 21:15:15 +0000 Subject: [PATCH 27/42] Arguments to make logarithmic color-scale in 2d plots 'norm' option allows any matplotlib.colors.Normalize instance to be passed, for general control of the color-scale. For convenience, also adds a 'logscale' option which can be set to True to create a standard log-scale if vmin and vmax are both positive or both negative, and a symmetric log-scale with a linear threshold of min(abs(vmin),abs(vmax))*1.e-5 if vmin and vmax have opposite signs. If a float is passed for 'logscale' and vmin and vmax have opposite signs then the value of 'logscale' is used instead of the default 1.e-5 to set the linear threshold. --- xbout/plotting/plotfuncs.py | 51 ++++++++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 08f5649b..c29db9ea 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -37,8 +37,9 @@ def regions(da, ax=None, **kwargs): def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, - add_limiter_hatching=True, gridlines=None, cmap=None, vmin=None, - vmax=None, aspect=None, **kwargs): + add_limiter_hatching=True, gridlines=None, cmap=None, + norm = None, logscale=None, vmin=None, vmax=None, + aspect=None, **kwargs): """ Make a 2D plot using an xarray method, taking into account branch cuts (X-points). @@ -64,6 +65,15 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, stride when plotting grid lines (to reduce the number on the plot) cmap : Matplotlib colormap, optional Color map to use for the plot + norm : matplotlib.colors.Normalize instance, optional + Normalization to use for the color scale. + Cannot be set at the same time as 'logscale' + 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. + Cannot be set at the same time as 'norm' vmin : float, optional Minimum value for the color scale vmax : float, optional @@ -103,6 +113,28 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, if vmax is None: vmax = da.max().values + def create_norm(norm): + if logscale: + if norm is not None: + raise ValueError("norm and logscale cannot both be passed at the same " + "time.") + if vmin*vmax > 0: + # vmin and vmax have the same sign, so can use standard log-scale + norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax) + else: + # vmin and vmax have opposite signs, so use symmetrical logarithmic scale + if not isinstance(logscale, bool): + linear_scale = logscale + else: + linear_scale = 1.e-5 + linear_threshold = min(abs(vmin), abs(vmax)) * linear_scale + norm = matplotlib.colors.SymLogNorm(linear_threshold, vmin=vmin, + vmax=vmax) + elif norm is None: + norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) + + return norm + # set up 'levels' if needed if method is xr.plot.contourf or method is xr.plot.contour: levels = kwargs.get('levels', 7) @@ -116,12 +148,23 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, vmin = np.min(levels) vmax = np.max(levels) + levels = kwargs.get('levels', 7) + if isinstance(levels, np.int): + levels = np.linspace(vmin, vmax, levels, endpoint=True) + # put levels back into kwargs + kwargs['levels'] = levels + else: + levels = np.array(list(levels)) + kwargs['levels'] = levels + vmin = np.min(levels) + vmax = np.max(levels) + # Need to create a colorscale that covers the range of values in the whole array. # Using the add_colorbar argument would create a separate color scale for each # separate region, which would not make sense. if method is xr.plot.contourf: # create colorbar - norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) + norm = create_norm(norm) sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) # make colorbar have only discrete levels # average the levels so that colors in the colorbar represent the intervals @@ -146,7 +189,7 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, kwargs['vmax'] = vmax # create colorbar - norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) + norm = create_norm(norm) sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) sm.set_array([]) cmap = sm.get_cmap() From 122704bba46de2e851d29a478ca75c69714e0834 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 5 Dec 2019 22:02:56 +0000 Subject: [PATCH 28/42] Implement logscale option for 2D animation methods --- xbout/boutdataarray.py | 14 +++++++++++++- xbout/boutdataset.py | 24 ++++++++++++++++++++---- xbout/plotting/animate.py | 3 ++- xbout/plotting/plotfuncs.py | 29 ++++------------------------- xbout/plotting/utils.py | 24 ++++++++++++++++++++++++ 5 files changed, 63 insertions(+), 31 deletions(-) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index ef8b1137..3711bae1 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -6,6 +6,7 @@ from .plotting.animate import animate_poloidal, animate_pcolormesh, animate_line from .plotting import plotfuncs +from .plotting.utils import _create_norm @register_dataarray_accessor('bout') @@ -39,7 +40,7 @@ def __str__(self): return text def animate2D(self, animate_over='t', x=None, y=None, animate=True, fps=10, - save_as=None, ax=None, poloidal_plot=False, **kwargs): + save_as=None, ax=None, poloidal_plot=False, logscale=None, **kwargs): """ Plots a color plot which is animated with time over the specified coordinate. @@ -68,6 +69,12 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, fps=10, poloidal_plot : bool, optional Use animate_poloidal to make a plot in R-Z coordinates (input field must be (t,x,y)) + 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. kwargs : dict, optional Additional keyword arguments are passed on to the plotting function (e.g. imshow for 2D plots). @@ -76,7 +83,12 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True, fps=10, data = self.data variable = data.name n_dims = len(data.dims) + if n_dims == 3: + vmin = kwargs['vmin'] if 'vmin' in kwargs else data.min().values + vmax = kwargs['vmax'] if 'vmax' in kwargs else data.max().values + kwargs['norm'] = _create_norm(logscale, kwargs.get('norm', None), vmin, vmax) + if poloidal_plot: print("{} data passed has {} dimensions - making poloidal plot with " "animate_poloidal()".format(variable, str(n_dims))) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 8057a01b..6f860052 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -9,6 +9,7 @@ from dask.diagnostics import ProgressBar from .plotting.animate import animate_poloidal, animate_pcolormesh, animate_line +from .plotting.utils import _create_norm @register_dataset_accessor('bout') @@ -165,8 +166,8 @@ def to_restart(self, savepath='.', nxpe=None, nype=None, return 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, **kwargs): + nrows=None, ncols=None, poloidal_plot=False, subplots_adjust=None, + vmin=None, vmax=None, logscale=None, **kwargs): """ Parameters ---------- @@ -180,6 +181,12 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp 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, 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. """ nvars = len(variables) @@ -228,18 +235,27 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp 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(logscale, kwargs.get('norm', None), this_vmin, + this_vmax) + if poloidal_plot: var_blocks = animate_poloidal(data, ax=ax, animate_over=animate_over, animate=False, vmin=this_vmin, - vmax=this_vmax, **kwargs) + vmax=this_vmax, norm=norm, **kwargs) for block in var_blocks: blocks.append(block) else: blocks.append(animate_pcolormesh(data=data, ax=ax, animate_over=animate_over, animate=False, vmin=this_vmin, - vmax=this_vmax, **kwargs)) + vmax=this_vmax, norm=norm, + **kwargs)) else: raise ValueError("Unsupported number of dimensions " + str(ndims) + ". Dims are " + str(v.dims)) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index de67fef6..be68c47a 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -67,7 +67,8 @@ def animate_poloidal(da, *, ax=None, animate_over='t', separatrix=True, targets= kwargs['vmax'] = vmax # create colorbar - norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) + 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() diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index c29db9ea..2413477a 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -6,7 +6,8 @@ import xarray as xr -from .utils import _decompose_regions, _is_core_only, plot_separatrices, plot_targets +from .utils import (_create_norm, _decompose_regions, _is_core_only, plot_separatrices, + plot_targets) def regions(da, ax=None, **kwargs): @@ -113,28 +114,6 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, if vmax is None: vmax = da.max().values - def create_norm(norm): - if logscale: - if norm is not None: - raise ValueError("norm and logscale cannot both be passed at the same " - "time.") - if vmin*vmax > 0: - # vmin and vmax have the same sign, so can use standard log-scale - norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax) - else: - # vmin and vmax have opposite signs, so use symmetrical logarithmic scale - if not isinstance(logscale, bool): - linear_scale = logscale - else: - linear_scale = 1.e-5 - linear_threshold = min(abs(vmin), abs(vmax)) * linear_scale - norm = matplotlib.colors.SymLogNorm(linear_threshold, vmin=vmin, - vmax=vmax) - elif norm is None: - norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) - - return norm - # set up 'levels' if needed if method is xr.plot.contourf or method is xr.plot.contour: levels = kwargs.get('levels', 7) @@ -164,7 +143,7 @@ def create_norm(norm): # separate region, which would not make sense. if method is xr.plot.contourf: # create colorbar - norm = create_norm(norm) + norm = _create_norm(logscale, norm, vmin, vmax) sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) # make colorbar have only discrete levels # average the levels so that colors in the colorbar represent the intervals @@ -189,7 +168,7 @@ def create_norm(norm): kwargs['vmax'] = vmax # create colorbar - norm = create_norm(norm) + norm = _create_norm(logscale, norm, vmin, vmax) sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) sm.set_array([]) cmap = sm.get_cmap() diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index d3de9b8a..4168ec66 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -1,9 +1,33 @@ import warnings +import matplotlib as mpl import numpy as np import xarray as xr +def _create_norm(logscale, norm, vmin, vmax): + if logscale: + if norm is not None: + raise ValueError("norm and logscale cannot both be passed at the same " + "time.") + if vmin*vmax > 0: + # vmin and vmax have the same sign, so can use standard log-scale + norm = mpl.colors.LogNorm(vmin=vmin, vmax=vmax) + else: + # vmin and vmax have opposite signs, so use symmetrical logarithmic scale + if not isinstance(logscale, bool): + linear_scale = logscale + else: + linear_scale = 1.e-5 + linear_threshold = min(abs(vmin), abs(vmax)) * linear_scale + norm = mpl.colors.SymLogNorm(linear_threshold, vmin=vmin, + vmax=vmax) + elif norm is None: + norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) + + return norm + + def plot_separatrix(da, sep_pos, ax, radial_coord='x'): """ Plots the separatrix as a black dotted line. From 625a8e2a659978941965223fd5b8b0268e99ff7c Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 5 Dec 2019 22:10:55 +0000 Subject: [PATCH 29/42] Document all arguments of animate_list() --- xbout/boutdataset.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 6f860052..f75a6266 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -177,6 +177,23 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp 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, optional + If set to True, make all 2D animations in the poloidal plane instead of using + grid coordinates + 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 @@ -187,6 +204,8 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp threshold of a symmetric logarithmic scale as linthresh=min(abs(vmin),abs(vmax))*logscale, defaults to 1e-5 if True is passed. + **kwargs : dict, optional + Additional keyword arguments are passed on to each animation function """ nvars = len(variables) From 91efd51c72a5ace53878c405ba45725b027fc6db Mon Sep 17 00:00:00 2001 From: johnomotani Date: Sun, 8 Dec 2019 20:15:29 +0100 Subject: [PATCH 30/42] Use PillowWriter to save .gif in animate_list --- xbout/plotting/animate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index be68c47a..e4092f39 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -107,7 +107,7 @@ def animate_poloidal(da, *, ax=None, animate_over='t', separatrix=True, targets= if not save_as: save_as = "{}_over_{}".format(da.name, animate_over) - anim.save(save_as + '.gif', writer='imagemagick') + anim.save(save_as + '.gif', writer=PillowWriter(fps=fps)) return blocks From 85b0cc6ebe816604b5994662e99c5c33ac35d6f0 Mon Sep 17 00:00:00 2001 From: johnomotani Date: Sun, 8 Dec 2019 20:24:21 +0100 Subject: [PATCH 31/42] Tidy up imports in animate.py Had been mixed up in an automatic merge. --- xbout/plotting/animate.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index e4092f39..03d5c581 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -5,6 +5,7 @@ import animatplot as amp from .utils import _decompose_regions, plot_separatrices, plot_targets +from matplotlib.animation import PillowWriter def animate_poloidal(da, *, ax=None, animate_over='t', separatrix=True, targets=True, @@ -111,9 +112,6 @@ def animate_poloidal(da, *, ax=None, animate_over='t', separatrix=True, targets= return blocks -from .utils import plot_separatrix -from matplotlib.animation import PillowWriter - def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True, vmin=None, vmax=None, vsymmetric=False, fps=10, save_as=None, From 9d18e9d2d02fdb135c17352f39c1ced144a37015 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Dec 2019 21:45:03 +0000 Subject: [PATCH 32/42] Allow logscale to be passed as a list to animate_list Allows logscale to be passed per-variable. --- xbout/boutdataset.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index f75a6266..fc042234 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -198,12 +198,13 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp 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, optional + 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. **kwargs : dict, optional Additional keyword arguments are passed on to each animation function """ @@ -240,8 +241,16 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp except TypeError: vmax = [vmax] * len(variables) + try: + if len(logscale) != len(variables): + raise ValueError('if logscale is a sequence, it must have the ' + 'same number of elements as "variables"') + except TypeError: + logscale = [logscale] * len(variables) + blocks = [] - for v, ax, this_vmin, this_vmax in zip(variables, axes.flatten(), vmin, vmax): + for v, ax, this_vmin, this_vmax, this_logscale in zip(variables, axes.flatten(), + vmin, vmax, logscale): if isinstance(v, str): v = self.data[v] @@ -259,7 +268,7 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp if this_vmax is None: this_vmax = data.max().values - norm = _create_norm(logscale, kwargs.get('norm', None), this_vmin, + norm = _create_norm(this_logscale, kwargs.get('norm', None), this_vmin, this_vmax) if poloidal_plot: From 26f1e1957cae17a005e2111953b46722b8d7ec7b Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Dec 2019 21:58:09 +0000 Subject: [PATCH 33/42] Fix use of animate_over in animate_line() --- xbout/plotting/animate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 03d5c581..409e7154 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -266,7 +266,7 @@ def animate_line(data, animate_over='t', animate=True, if (t_read is animate_over): pass else: - data = data.transpose(x_read, animate_over) + 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 From ee98a709c6afd50546c9d91176ce88f71cbfd9ff Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Dec 2019 22:10:26 +0000 Subject: [PATCH 34/42] Allow poloidal_plot to be passed per variable to animate_list --- xbout/boutdataset.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index fc042234..3cac3f3f 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -189,9 +189,9 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp Specify the number of rows of plots ncols : int, optional Specify the number of columns of plots - poloidal_plot : bool, optional + 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 + grid coordinates, per variable if sequence is given subplots_adjust : dict, optional Arguments passed to fig.subplots_adjust()() vmin : float or sequence of floats @@ -227,6 +227,13 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp if subplots_adjust is not None: fig.subplots_adjust(**subplots_adjust) + try: + if len(poloidal_plot) != len(variables): + raise ValueError('if poloidal_plot is a sequence, it must have the same ' + 'number of elements as "variables"') + except TypeError: + poloidal_plot = [poloidal_plot] * len(variables) + try: if len(vmin) != len(variables): raise ValueError('if vmin is a sequence, it must have the same number ' @@ -249,8 +256,9 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp logscale = [logscale] * len(variables) blocks = [] - for v, ax, this_vmin, this_vmax, this_logscale in zip(variables, axes.flatten(), - vmin, vmax, logscale): + for this in zip(variables, axes.flatten(), poloidal_plot, vmin, vmax, logscale): + + v, ax, this_poloidal_plot, this_vmin, this_vmax, this_logscale = this if isinstance(v, str): v = self.data[v] @@ -271,7 +279,7 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp norm = _create_norm(this_logscale, kwargs.get('norm', None), this_vmin, this_vmax) - if poloidal_plot: + if this_poloidal_plot: var_blocks = animate_poloidal(data, ax=ax, animate_over=animate_over, animate=False, vmin=this_vmin, From 1bf86d96a1be3c414c8cae391f226c4a2c9aa29e Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Dec 2019 23:05:41 +0000 Subject: [PATCH 35/42] Tests for animate_list --- xbout/tests/test_animate.py | 227 ++++++++++++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) diff --git a/xbout/tests/test_animate.py b/xbout/tests/test_animate.py index 7b3fef16..b64e6a9e 100644 --- a/xbout/tests/test_animate.py +++ b/xbout/tests/test_animate.py @@ -1,4 +1,6 @@ import pytest +import numpy as np +import xarray as xr from xbout import open_boutdataset from xbout.boutdataarray import BoutDataArrayAccessor @@ -58,3 +60,228 @@ def test_animate1D(self, create_test_file): animation = ds['n'].isel(y=2, z=0).bout.animate1D(save_as="%s/test" % save_dir) assert isinstance(animation, Line) + + def test_animate_list(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)]) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_1d_default(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(y=2, z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)]) + + assert isinstance(animations.blocks[0], Line) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_animate_over(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(t=2), + ds['n'].isel(y=1, z=2)], + animate_over='x') + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_save_as(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + save_as="%s/test" % save_dir) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_fps(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + fps=42) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_nrows(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + nrows=2) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_ncols(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + ncols=3) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_not_enough_nrowsncols(self, create_test_file): + + save_dir, ds = create_test_file + + with pytest.raises(ValueError): + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + nrows=2, ncols=1) + + @pytest.mark.skip(reason='test data for plot_poloidal needs more work') + def test_animate_list_poloidal_plot(self, create_test_file): + + save_dir, ds = create_test_file + + metadata = ds.metadata + metadata['ixseps1'] = 2 + metadata['ixseps2'] = 4 + metadata['jyseps1_1'] = 1 + metadata['jyseps2_1'] = 2 + metadata['ny_inner'] = 3 + metadata['jyseps1_2'] = 4 + metadata['jyseps2_2'] = 5 + from ..geometries import apply_geometry + from ..utils import _set_attrs_on_all_vars + ds = _set_attrs_on_all_vars(ds, 'metadata', metadata) + + nx = ds.metadata['nx'] + ny = ds.metadata['ny'] + R = xr.DataArray(np.ones([nx, ny])*np.linspace(0, 1, nx)[:, np.newaxis], + dims=['x', 'y']) + Z = xr.DataArray(np.ones([nx, ny])*np.linspace(0, 1, ny)[np.newaxis, :], + dims=['x', 'y']) + ds['psixy'] = R + ds['Rxy'] = R + ds['Zxy'] = Z + + ds = apply_geometry(ds, 'toroidal') + + animations = ds.isel(zeta=3).bout.animate_list(['n', ds['T'].isel(zeta=3), + ds['n'].isel(theta=1, zeta=2)], + poloidal_plot=True) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_subplots_adjust(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + subplots_adjust={'hspace': 4, + 'wspace': 5}) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_vmin(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + vmin=-0.1) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_vmin_list(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + vmin=[0., 0.1, 0.2]) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_vmax(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + vmax=1.1) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_vmax_list(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + vmax=[1., 1.1, 1.2]) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_logscale(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + logscale=True) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_logscale_float(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + logscale=1.e-2) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) + + def test_animate_list_logscale_list(self, create_test_file): + + save_dir, ds = create_test_file + + animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + ds['n'].isel(y=1, z=2)], + logscale=[True, 1.e-2, False]) + + assert isinstance(animations.blocks[0], Pcolormesh) + assert isinstance(animations.blocks[1], Pcolormesh) + assert isinstance(animations.blocks[2], Line) From ef9947de9beb753a1a39d0eea1d2537e4f126c0c Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 10 Dec 2019 11:49:43 +0000 Subject: [PATCH 36/42] Don't rely on order of dims in _decompose_regions() Can now use 'bout_xdim' and 'bout_ydim' metadata to find the renamed x- and y-dimensions. --- xbout/plotting/utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 4168ec66..2a5dd87d 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -65,8 +65,11 @@ def _decompose_regions(da): j11, j12, j21, j22, ix1, ix2, nin, _, ny, y_boundary_guards = _get_seps(da) regions = [] - x, y = da.dims[-2:] - other_dims = da.dims[:-2] + x = da.dims[da.metadata['bout_xdim']] + y = da.dims[da.metadata['bout_ydim']] + other_dims = list(da.dims) + other_dims.remove(x) + other_dims.remove(y) ystart = 0 # Y index to start the next section if j11 >= 0: From 9f8470fde153681ad326589cac87390abb836985 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 10 Dec 2019 11:54:55 +0000 Subject: [PATCH 37/42] Replace redundant argument unpacking in plotting.utils --- xbout/plotting/utils.py | 60 ++++++++++++++++++++--------------------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 2a5dd87d..a2c3c797 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -74,15 +74,15 @@ def _decompose_regions(da): ystart = 0 # Y index to start the next section if j11 >= 0: # plot lower inner leg - region1 = da.isel(**{y: slice(ystart, (j11 + 1))}) + region1 = da.isel(y=slice(ystart, (j11 + 1))) yind = [j11, j22 + 1] - region2 = da.isel(**{x: slice(0, ix1), y:yind}) + region2 = da.isel(x=slice(0, ix1), y=yind) - region3 = da.isel(**{x: slice(ix1, None), y: slice(j11, (j11 + 2))}) + region3 = da.isel(x=slice(ix1, None), y=slice(j11, (j11 + 2))) yind = [j22, j11 + 1] - region4 = da.isel(**{x: slice(0, ix1), y: yind}) + region4 = da.isel(x=slice(0, ix1), y=yind) regions.extend([region1, region2, region3, region4]) @@ -90,7 +90,7 @@ def _decompose_regions(da): if j21 + 1 > ystart: # Inner SOL - region5 = da.isel(**{y: slice(ystart, (j21 + 1))}) + region5 = da.isel(y=slice(ystart, (j21 + 1))) regions.append(region5) ystart = j21 + 1 @@ -99,19 +99,19 @@ def _decompose_regions(da): # Contains upper PF region # Inner leg - region6 = da.isel(**{x: slice(ix1, None), y: slice(j21, (j21 + 2))}) - region7 = da.isel(**{y: slice(ystart, nin)}) + region6 = da.isel(x=slice(ix1, None), y=slice(j21, (j21 + 2))) + region7 = da.isel(y=slice(ystart, nin)) # Outer leg - region8 = da.isel(**{y: slice(nin, (j12 + 1))}) - region9 = da.isel(**{x: slice(ix1, None), y: slice(j12, (j12 + 2))}) + region8 = da.isel(y=slice(nin, (j12 + 1))) + region9 = da.isel(x=slice(ix1, None), y=slice(j12, (j12 + 2))) yind = [j21, j12 + 1] - region10 = da.isel(**{x: slice(0, ix1), y: yind}) + region10 = da.isel(x=slice(0, ix1), y=yind) yind = [j21 + 1, j12] - region11 = da.isel(**{x: slice(0, ix1), y: yind}) + region11 = da.isel(x=slice(0, ix1), y=yind) regions.extend([region6, region7, region8, region9, region10, region11]) @@ -122,29 +122,29 @@ def _decompose_regions(da): if j22 + 1 > ystart: # Outer SOL - region12 = da.isel(**{y: slice(ystart, (j22 + 1))}) + region12 = da.isel(y=slice(ystart, (j22 + 1))) regions.append(region12) ystart = j22 + 1 if j22 + 1 < ny: # Outer leg - region13 = da.isel(**{x: slice(ix1, None), y: slice(j22, (j22 + 2))}) - region14 = da.isel(**{y: slice(ystart, ny)}) + region13 = da.isel(x=slice(ix1, None), y=slice(j22, (j22 + 2))) + region14 = da.isel(y=slice(ystart, ny)) # X-point regions - corner1 = da.isel(**{x: ix1-1, y: j11}) - corner2 = da.isel(**{x: ix1, y: j11}) - corner3 = da.isel(**{x: ix1, y: j11+1}) - corner4 = da.isel(**{x: ix1-1, y: j11+1}) + corner1 = da.isel(x=ix1-1, y=j11) + corner2 = da.isel(x=ix1, y=j11) + corner3 = da.isel(x=ix1, y=j11+1) + corner4 = da.isel(x=ix1-1, y=j11+1) xregion_lower = xr.concat([corner1, corner2, corner3, corner4], dim='dim1') - corner5 = da.isel(**{x: ix1-1, y: j22+1}) - corner6 = da.isel(**{x: ix1, y: j22+1}) - corner7 = da.isel(**{x: ix1, y: j22}) - corner8 = da.isel(**{x: ix1-1, y: j22}) + corner5 = da.isel(x=ix1-1, y=j22+1) + corner6 = da.isel(x=ix1, y=j22+1) + corner7 = da.isel(x=ix1, y=j22) + corner8 = da.isel(x=ix1-1, y=j22) xregion_upper = xr.concat([corner5, corner6, corner7, corner8], dim='dim1') @@ -159,18 +159,18 @@ def _decompose_regions(da): if j21 > j11 and j12 > j21 and j22 > j12: # X-point regions - corner1 = da.isel(**{x: ix1-1, y: j12}) - corner2 = da.isel(**{x: ix1, y: j12}) - corner3 = da.isel(**{x: ix1, y: j12+1}) - corner4 = da.isel(**{x: ix1-1, y: j12+1}) + corner1 = da.isel(x=ix1-1, y=j12) + corner2 = da.isel(x=ix1, y=j12) + corner3 = da.isel(x=ix1, y=j12+1) + corner4 = da.isel(x=ix1-1, y=j12+1) xregion_lower = xr.concat([corner1, corner2, corner3, corner4], dim='dim1') - corner5 = da.isel(**{x: ix1-1, y: j21+1}) - corner6 = da.isel(**{x: ix1, y: j21+1}) - corner7 = da.isel(**{x: ix1, y: j21}) - corner8 = da.isel(**{x: ix1-1, y: j21}) + corner5 = da.isel(x=ix1-1, y=j21+1) + corner6 = da.isel(x=ix1, y=j21+1) + corner7 = da.isel(x=ix1, y=j21) + corner8 = da.isel(x=ix1-1, y=j21) xregion_upper = xr.concat([corner5, corner6, corner7, corner8], dim='dim1') From 3176809377b3a6e1e39a427307990527a36ddf52 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 10 Dec 2019 12:06:26 +0000 Subject: [PATCH 38/42] Simplify checking for sequences in arguments of animate_list --- xbout/boutdataset.py | 40 +++++++++++++--------------------------- 1 file changed, 13 insertions(+), 27 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 3cac3f3f..6dfc93fd 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -1,3 +1,4 @@ +import collections from pprint import pformat as prettyformat from functools import partial @@ -227,33 +228,18 @@ def animate_list(self, variables, animate_over='t', save_as=None, show=False, fp if subplots_adjust is not None: fig.subplots_adjust(**subplots_adjust) - try: - if len(poloidal_plot) != len(variables): - raise ValueError('if poloidal_plot is a sequence, it must have the same ' - 'number of elements as "variables"') - except TypeError: - poloidal_plot = [poloidal_plot] * len(variables) - - try: - if len(vmin) != len(variables): - raise ValueError('if vmin is a sequence, it must have the same number ' - 'of elements as "variables"') - except TypeError: - vmin = [vmin] * len(variables) - - try: - if len(vmax) != len(variables): - raise ValueError('if vmin is a sequence, it must have the same number ' - 'of elements as "variables"') - except TypeError: - vmax = [vmax] * len(variables) - - try: - if len(logscale) != len(variables): - raise ValueError('if logscale is a sequence, it must have the ' - 'same number of elements as "variables"') - except TypeError: - logscale = [logscale] * len(variables) + def _expand_list_arg(arg, arg_name): + if isinstance(arg, collections.Sequence): + 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') blocks = [] for this in zip(variables, axes.flatten(), poloidal_plot, vmin, vmax, logscale): From 230d27ceee3ec4e92223ce94b3417e47e47f3f4e Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 10 Dec 2019 12:11:11 +0000 Subject: [PATCH 39/42] Rename 'this' -> 'subplot_args' --- xbout/boutdataset.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 6dfc93fd..f706a432 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -242,9 +242,10 @@ def _expand_list_arg(arg, arg_name): logscale = _expand_list_arg(logscale, 'logscale') blocks = [] - for this in zip(variables, axes.flatten(), poloidal_plot, vmin, vmax, logscale): + for subplot_args in zip(variables, axes.flatten(), poloidal_plot, vmin, vmax, + logscale): - v, ax, this_poloidal_plot, this_vmin, this_vmax, this_logscale = this + v, ax, this_poloidal_plot, this_vmin, this_vmax, this_logscale = subplot_args if isinstance(v, str): v = self.data[v] From 7ea1c7019eb2689a0dcd430cb22764fe2d82601e Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 10 Dec 2019 12:25:54 +0000 Subject: [PATCH 40/42] Make specification of 'gridlines' argument to plot2d_wrapper nicer Also specify in docstring what happens if a dict is passed. --- xbout/plotting/plotfuncs.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 2413477a..1d567688 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -61,9 +61,12 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, Draw solid lines at the target surfaces add_limiter_hatching : bool, optional Draw hatched areas at the targets - gridlines : bool or int, optional + gridlines : bool, int or slice or dict of bool, int or slice, optional If True, draw grid lines on the plot. If an int is passed, it is used as the - stride when plotting grid lines (to reduce the number on the plot) + stride when plotting grid lines (to reduce the number on the plot). If a slice is + passed it is used to select the grid lines to plot. + If a dict is passed, the 'x' entry (bool, int or slice) is used for the radial + grid-lines and the 'y' entry for the poloidal grid lines. cmap : Matplotlib colormap, optional Color map to use for the plot norm : matplotlib.colors.Normalize instance, optional @@ -197,10 +200,20 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, fig.colorbar(artists[0], ax=ax, extend=extend) if gridlines is not None: - if gridlines is True: - gridlines = (1, 1) - if not isinstance(gridlines, collections.Sequence): - gridlines = (gridlines, gridlines) + # convert gridlines to dict + if not isinstance(gridlines, dict): + gridlines = {'x': gridlines, 'y': gridlines} + + for key in gridlines: + value = gridlines[key] + if value is True: + gridlines[key] = slice(None) + elif isinstance(value, int): + gridlines[key] = slice(0, None, value) + elif not value is None: + if not isinstance(value, slice): + raise ValueError('Argument passed to gridlines must be bool, int or ' + 'slice. Got ' + str(value)) R_global = da['R'] R_global.attrs['metadata'] = da.metadata @@ -211,8 +224,11 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True, Z_regions = _decompose_regions(da['Z']) for R, Z in zip(R_regions, Z_regions): - plt.plot(R[::gridlines[0], :].T, Z[::gridlines[0], :].T, color='k', lw=0.1) - plt.plot(R[:, ::gridlines[1]], Z[:, ::gridlines[1]], color='k', lw=0.1) + if 'x' in gridlines and gridlines['x'] is not None: + plt.plot(R[:, gridlines['y']], Z[:, gridlines['y']], color='k', lw=0.1) + if 'y' in gridlines and gridlines['y'] is not None: + plt.plot(R[gridlines['y'], :].T, Z[gridlines['y'], :].T, color='k', + lw=0.1) ax.set_title(da.name) From 3bad703d896a1a0adcb06de961edd4522c89407d Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 10 Dec 2019 12:28:42 +0000 Subject: [PATCH 41/42] Rename animations->animation in animate_list tests Returned object is a single animatplot 'animation' with multiple blocks, not a list of things. --- xbout/tests/test_animate.py | 130 ++++++++++++++++++------------------ 1 file changed, 65 insertions(+), 65 deletions(-) diff --git a/xbout/tests/test_animate.py b/xbout/tests/test_animate.py index b64e6a9e..ab98f09f 100644 --- a/xbout/tests/test_animate.py +++ b/xbout/tests/test_animate.py @@ -65,90 +65,90 @@ def test_animate_list(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)]) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_1d_default(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(y=2, z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(y=2, z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)]) - assert isinstance(animations.blocks[0], Line) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Line) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_animate_over(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(t=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(t=2), ds['n'].isel(y=1, z=2)], animate_over='x') - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_save_as(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], save_as="%s/test" % save_dir) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_fps(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], fps=42) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_nrows(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], nrows=2) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_ncols(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], ncols=3) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_not_enough_nrowsncols(self, create_test_file): save_dir, ds = create_test_file with pytest.raises(ValueError): - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], nrows=2, ncols=1) @@ -181,107 +181,107 @@ def test_animate_list_poloidal_plot(self, create_test_file): ds = apply_geometry(ds, 'toroidal') - animations = ds.isel(zeta=3).bout.animate_list(['n', ds['T'].isel(zeta=3), + animation = ds.isel(zeta=3).bout.animate_list(['n', ds['T'].isel(zeta=3), ds['n'].isel(theta=1, zeta=2)], poloidal_plot=True) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_subplots_adjust(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], subplots_adjust={'hspace': 4, 'wspace': 5}) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_vmin(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], vmin=-0.1) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_vmin_list(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], vmin=[0., 0.1, 0.2]) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_vmax(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], vmax=1.1) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_vmax_list(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], vmax=[1., 1.1, 1.2]) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_logscale(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], logscale=True) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_logscale_float(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], logscale=1.e-2) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) def test_animate_list_logscale_list(self, create_test_file): save_dir, ds = create_test_file - animations = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), + animation = ds.isel(z=3).bout.animate_list(['n', ds['T'].isel(x=2), ds['n'].isel(y=1, z=2)], logscale=[True, 1.e-2, False]) - assert isinstance(animations.blocks[0], Pcolormesh) - assert isinstance(animations.blocks[1], Pcolormesh) - assert isinstance(animations.blocks[2], Line) + assert isinstance(animation.blocks[0], Pcolormesh) + assert isinstance(animation.blocks[1], Pcolormesh) + assert isinstance(animation.blocks[2], Line) From dba0831c992bc07f692c2d57016a32d171d00f0f Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 10 Dec 2019 12:29:48 +0000 Subject: [PATCH 42/42] Check fps was actually set it test_animate_list_fps --- xbout/tests/test_animate.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xbout/tests/test_animate.py b/xbout/tests/test_animate.py index ab98f09f..8991dd0c 100644 --- a/xbout/tests/test_animate.py +++ b/xbout/tests/test_animate.py @@ -118,6 +118,7 @@ def test_animate_list_fps(self, create_test_file): assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) + assert animation.timeline.fps == 42 def test_animate_list_nrows(self, create_test_file):