From 2ad86d03d56ad73d216118178918db8c24104165 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 1 Mar 2023 11:08:04 +0000 Subject: [PATCH 01/25] Add cell corners to dataset coords --- xbout/geometries.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/xbout/geometries.py b/xbout/geometries.py index fa8a67b6..34c2f8b7 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -383,6 +383,14 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): "total_poloidal_distance", "zShift", "zShift_ylow", + "Rxy_corners", # Lower left corners + "Rxy_lower_right_corners", + "Rxy_upper_left_corners", + "Rxy_upper_right_corners", + "Zxy_corners", # Lower left corners + "Zxy_lower_right_corners", + "Zxy_upper_left_corners", + "Zxy_upper_right_corners" ], ) @@ -421,6 +429,23 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): ds = ds.set_coords(("R", "Z")) else: ds = ds.set_coords(("Rxy", "Zxy")) + + # Add cell corners as coordinates for polygon plotting + if "Rxy_lower_right_corners" in ds: + ds = ds.rename( + Rxy_corners = "Rxy_lower_left_corners", + Zxy_corners = "Zxy_lower_left_corners" + ) + ds = ds.set_coords(( + "Rxy_lower_left_corners", + "Rxy_lower_right_corners", + "Rxy_upper_left_corners", + "Rxy_upper_right_corners", + "Zxy_lower_left_corners", + "Zxy_lower_right_corners", + "Zxy_upper_left_corners", + "Zxy_upper_right_corners" + )) # Rename zShift_ylow if it was added from grid file, to be consistent with name if # it was added from dump file From 292044e73968a0870d17e1959f78115a957dfa35 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 1 Mar 2023 11:41:22 +0000 Subject: [PATCH 02/25] plotfuncs: add polygon_plot --- xbout/plotting/plotfuncs.py | 62 +++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 79b36207..2f3e62ef 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -843,3 +843,65 @@ def animation_func(): plt.show() else: raise ValueError(f"Unrecognised plot3d() 'engine' argument: {engine}") + +def plot2d_polygon( + da, + ax = None, + cax = None, + vmin = None, + vmax = None, + logscale = False, + cmap = "viridis", + +): + + if ax == None: + fig, ax = plt.subplots(figsize=(6, 10)) + + if "Rxy_lower_right_corners" in da.coords: + r_nodes = ["R", "Rxy_lower_left_corners", "Rxy_lower_right_corners", "Rxy_upper_left_corners", "Rxy_upper_right_corners"] + z_nodes = ["Z", "Zxy_lower_left_corners", "Zxy_lower_right_corners", "Zxy_upper_left_corners", "Zxy_upper_right_corners"] + cell_r = np.concatenate([np.expand_dims(da[x], axis = 2) for x in r_nodes], axis = 2) + cell_z = np.concatenate([np.expand_dims(da[x], axis = 2) for x in z_nodes], axis = 2) + else: + raise Exception("Cell corners not present in mesh, cannot do polygon plot") + + Nx = len(cell_r) + Ny = len(cell_r[0]) + edgecolor = "black" + patches = [] + + # https://matplotlib.org/2.0.2/examples/api/patch_collection.html + + idx = [np.array([1, 2, 4, 3, 1])] + patches = [] + for i in range(Nx): + for j in range(Ny): + p = matplotlib.patches.Polygon( + np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])).reshape(2, 5).T, + fill=False, + closed=True, + edgecolor=edgecolor, + ) + patches.append(p) + + + # create colorbar + norm = _create_norm(logscale, norm, vmin, vmax) + sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) + sm.set_array([]) + cmap = sm.get_cmap() + fig.colorbar(sm, ax=ax, extend=extend) + + colors = da.data.flatten() + polys = matplotlib.collections.PatchCollection(patches, alpha = 1, cmap = "Spectral_r", antialiaseds = True) + polys.set_array(colors) + + fig.colorbar(polys, ax = ax) + + ax.add_collection(polys) + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel("R [m]") + ax.set_ylabel("Z [m]") + ax.set_ylim(cell_z.min(), cell_z.max()) + ax.set_xlim(cell_r.min(), cell_r.max()) \ No newline at end of file From be110baf9375f757ea2936ea84a569180fcde861 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 1 Mar 2023 11:44:11 +0000 Subject: [PATCH 03/25] boutdataarray: add polygon wrapper --- xbout/boutdataarray.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index 7e57cf1e..b29ce8ec 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -1073,6 +1073,12 @@ def pcolormesh(self, ax=None, **kwargs): Colour-plot a radial-poloidal slice on the R-Z plane """ return plotfuncs.plot2d_wrapper(self.data, xr.plot.pcolormesh, ax=ax, **kwargs) + + def polygon(self, ax=None, **kwargs): + """ + Colour-plot of a radial-poloidal slice on the R-Z plane using polygons + """ + return plotfuncs.plot2d_polygon(self.data, ax=ax, **kwargs) def plot_regions(self, ax=None, **kwargs): """ From 9700b6caf53086e9c1fa2d3fc54947cdf9fae615 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 1 Mar 2023 14:21:47 +0000 Subject: [PATCH 04/25] plot2d_polygon: add customisation opts --- xbout/plotting/plotfuncs.py | 59 +++++++++++++++++++++++++++++-------- 1 file changed, 46 insertions(+), 13 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 2f3e62ef..e7bee37b 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -848,15 +848,32 @@ def plot2d_polygon( da, ax = None, cax = None, + cmap = "viridis", + norm = None, + logscale = False, vmin = None, vmax = None, - logscale = False, - cmap = "viridis", + extend = None, + separatrix = False, + targets = False, + add_limiter_hatching=True, + linewidth = 0, + linecolour = "black", ): if ax == None: - fig, ax = plt.subplots(figsize=(6, 10)) + fig, ax = plt.subplots(figsize=(3, 6), dpi = 120) + + if cax == None: + cax = ax + + if vmin is None: + vmin = da.min().values + + if vmax is None: + vmax = da.max().values + if "Rxy_lower_right_corners" in da.coords: r_nodes = ["R", "Rxy_lower_left_corners", "Rxy_lower_right_corners", "Rxy_upper_left_corners", "Rxy_upper_right_corners"] @@ -881,27 +898,43 @@ def plot2d_polygon( np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])).reshape(2, 5).T, fill=False, closed=True, - edgecolor=edgecolor, ) patches.append(p) # create colorbar norm = _create_norm(logscale, norm, vmin, vmax) - sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) - sm.set_array([]) - cmap = sm.get_cmap() - fig.colorbar(sm, ax=ax, extend=extend) colors = da.data.flatten() - polys = matplotlib.collections.PatchCollection(patches, alpha = 1, cmap = "Spectral_r", antialiaseds = True) - polys.set_array(colors) + polys = matplotlib.collections.PatchCollection( + patches, alpha = 1, norm = norm, cmap = cmap, + antialiaseds = True, + edgecolors = linecolour, + linewidths = linewidth, + joinstyle = "bevel") + + - fig.colorbar(polys, ax = ax) - ax.add_collection(polys) + polys.set_array(colors) + + fig.colorbar(polys, ax = cax) + ax.add_collection(polys) + ax.set_aspect("equal", adjustable="box") ax.set_xlabel("R [m]") ax.set_ylabel("Z [m]") ax.set_ylim(cell_z.min(), cell_z.max()) - ax.set_xlim(cell_r.min(), cell_r.max()) \ No newline at end of file + ax.set_xlim(cell_r.min(), cell_r.max()) + ax.set_title(da.name) + + if separatrix: + # plot_separatrices(da, ax, x = "R", y = "Z") + + m = da.metadata + sel = (m["ixseps1"]-m["MXG"], slice(0,m["ny"])) + + ax.plot(da["R"].data[sel], da["Z"].data[sel]) + + if targets: + plot_targets(da, ax, x = "R", y = "Z", hatching = add_limiter_hatching) \ No newline at end of file From a94ce4552909f552f73d8cfe5c0c62dfa600423e Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 1 Mar 2023 14:37:27 +0000 Subject: [PATCH 05/25] plot2d_polygon: fix if no ax given --- xbout/plotting/plotfuncs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index e7bee37b..481d35bd 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -864,6 +864,8 @@ def plot2d_polygon( if ax == None: fig, ax = plt.subplots(figsize=(3, 6), dpi = 120) + else: + fig = ax.get_figure() if cax == None: cax = ax From fa6105944dda484dc901672d20e631eab8a80f5e Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 15 Mar 2023 14:20:31 +0000 Subject: [PATCH 06/25] polygon: add new separatrix plotter and cbar --- xbout/plotting/plotfuncs.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 481d35bd..7226b556 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -851,9 +851,11 @@ def plot2d_polygon( cmap = "viridis", norm = None, logscale = False, + antialias = False, vmin = None, vmax = None, extend = None, + add_colorbar = True, separatrix = False, targets = False, add_limiter_hatching=True, @@ -910,7 +912,7 @@ def plot2d_polygon( colors = da.data.flatten() polys = matplotlib.collections.PatchCollection( patches, alpha = 1, norm = norm, cmap = cmap, - antialiaseds = True, + antialiaseds = antialias, edgecolors = linecolour, linewidths = linewidth, joinstyle = "bevel") @@ -920,7 +922,8 @@ def plot2d_polygon( polys.set_array(colors) - fig.colorbar(polys, ax = cax) + if add_colorbar: + fig.colorbar(polys, ax = cax) ax.add_collection(polys) ax.set_aspect("equal", adjustable="box") @@ -934,9 +937,11 @@ def plot2d_polygon( # plot_separatrices(da, ax, x = "R", y = "Z") m = da.metadata - sel = (m["ixseps1"]-m["MXG"], slice(0,m["ny"])) - - ax.plot(da["R"].data[sel], da["Z"].data[sel]) + lhs = (m["ixseps1"]-m["MXG"]-1, slice(0,m["ny_inner"])) + rhs = (m["ixseps1"]-m["MXG"]-1, slice(m["ny_inner"]+m["MYG"], None)) + + ax.plot(da["Rxy_lower_right_corners"].data[lhs], da["Zxy_lower_right_corners"].data[lhs], c = "white") + ax.plot(da["Rxy_lower_right_corners"].data[rhs], da["Zxy_lower_right_corners"].data[rhs], c = "white") if targets: plot_targets(da, ax, x = "R", y = "Z", hatching = add_limiter_hatching) \ No newline at end of file From 7c71c202b49a60f2ac2fe08fe36bbeaf9bd9146c Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Thu, 16 Mar 2023 14:43:32 +0000 Subject: [PATCH 07/25] Add topology dependent separatrix plot --- xbout/plotting/plotfuncs.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 2cb2b204..feb34ec8 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -863,11 +863,11 @@ def plot2d_polygon( vmax = None, extend = None, add_colorbar = True, - separatrix = False, + separatrix = True, targets = False, add_limiter_hatching=True, linewidth = 0, - linecolour = "black", + linecolor = "black", ): @@ -920,7 +920,7 @@ def plot2d_polygon( polys = matplotlib.collections.PatchCollection( patches, alpha = 1, norm = norm, cmap = cmap, antialiaseds = antialias, - edgecolors = linecolour, + edgecolors = linecolor, linewidths = linewidth, joinstyle = "bevel") @@ -943,12 +943,22 @@ def plot2d_polygon( if separatrix: # plot_separatrices(da, ax, x = "R", y = "Z") + color = "white" + ls = "-" + lw = 1 + m = da.metadata - lhs = (m["ixseps1"]-m["MXG"]-1, slice(0,m["ny_inner"])) - rhs = (m["ixseps1"]-m["MXG"]-1, slice(m["ny_inner"]+m["MYG"], None)) - ax.plot(da["Rxy_lower_right_corners"].data[lhs], da["Zxy_lower_right_corners"].data[lhs], c = "white") - ax.plot(da["Rxy_lower_right_corners"].data[rhs], da["Zxy_lower_right_corners"].data[rhs], c = "white") + if m["topology"] == "connected-double-null": + lhs = (m["ixseps1"]-m["MXG"]-1, slice(0,m["ny_inner"])) + rhs = (m["ixseps1"]-m["MXG"]-1, slice(m["ny_inner"], None)) + ax.plot(da["Rxy_lower_right_corners"].data[lhs], da["Zxy_lower_right_corners"].data[lhs], c = color, lw = lw, ls = ls) + ax.plot(da["Rxy_lower_right_corners"].data[rhs], da["Zxy_lower_right_corners"].data[rhs], c = color, lw = lw, ls = ls) + + if m["topology"] == "single-null": + points = (m["ixseps1"]-m["MXG"]-1, slice(0,None)) + ax.plot(da["Rxy_lower_right_corners"].data[points], da["Zxy_lower_right_corners"].data[points], c = color, lw = lw, ls = ls) + if targets: plot_targets(da, ax, x = "R", y = "Z", hatching = add_limiter_hatching) \ No newline at end of file From 00430c9cdc502580f0c20b673d2019df14944a3c Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 21 Jun 2023 17:00:43 +0100 Subject: [PATCH 08/25] Make separatrix plotter work with guards --- xbout/plotting/plotfuncs.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index feb34ec8..385c3be1 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -949,9 +949,17 @@ def plot2d_polygon( m = da.metadata + + if m["topology"] == "connected-double-null": - lhs = (m["ixseps1"]-m["MXG"]-1, slice(0,m["ny_inner"])) - rhs = (m["ixseps1"]-m["MXG"]-1, slice(m["ny_inner"], None)) + + # TODO: test this for single null + if m["MXG"] > 0 or m["MYG"] > 0: + lhs = (m["ixseps1"] + 1 - m["MXG"], slice(0,m["ny_inner"] + m["MYG"] * int(4/2))) + rhs = (m["ixseps1"] + 1 - m["MXG"], slice(m["ny_inner"]+m["MYG"]*(4 - 1), None)) + else: + lhs = (m["ixseps1"]-m["MXG"]-1, slice(0,m["ny_inner"])) + rhs = (m["ixseps1"]-m["MXG"]-1, slice(m["ny_inner"], None)) ax.plot(da["Rxy_lower_right_corners"].data[lhs], da["Zxy_lower_right_corners"].data[lhs], c = color, lw = lw, ls = ls) ax.plot(da["Rxy_lower_right_corners"].data[rhs], da["Zxy_lower_right_corners"].data[rhs], c = color, lw = lw, ls = ls) From 5391d64b65ffc476308c65535a017d6e32fc41e1 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Thu, 29 Jun 2023 15:40:54 +0100 Subject: [PATCH 09/25] polygon plots: expose alpha setting --- xbout/plotting/plotfuncs.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 385c3be1..5f504b51 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -868,6 +868,7 @@ def plot2d_polygon( add_limiter_hatching=True, linewidth = 0, linecolor = "black", + alpha = 1, ): @@ -918,7 +919,7 @@ def plot2d_polygon( colors = da.data.flatten() polys = matplotlib.collections.PatchCollection( - patches, alpha = 1, norm = norm, cmap = cmap, + patches, alpha = alpha, norm = norm, cmap = cmap, antialiaseds = antialias, edgecolors = linecolor, linewidths = linewidth, From 1a38a4952a10aa2aba24a0f933999a00a7e5c3be Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 22 Mar 2024 15:28:52 +0000 Subject: [PATCH 10/25] Polygon: ability to plot only grid and customise cbar label --- xbout/plotting/plotfuncs.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index d6049ab3..1e9a0dcd 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -862,12 +862,13 @@ def plot2d_polygon( vmax = None, extend = None, add_colorbar = True, + colorbar_label = None, separatrix = True, targets = False, add_limiter_hatching=True, + grid_only = False, linewidth = 0, linecolor = "black", - alpha = 1, ): @@ -880,10 +881,10 @@ def plot2d_polygon( cax = ax if vmin is None: - vmin = da.min().values + vmin = np.nanmin(da.values) if vmax is None: - vmax = da.max().values + vmax = np.nanmax(da.max().values) if "Rxy_lower_right_corners" in da.coords: @@ -909,6 +910,7 @@ def plot2d_polygon( np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])).reshape(2, 5).T, fill=False, closed=True, + facecolor = None, ) patches.append(p) @@ -916,9 +918,11 @@ def plot2d_polygon( # create colorbar norm = _create_norm(logscale, norm, vmin, vmax) + if grid_only is True: + cmap = matplotlib.colors.ListedColormap(["white"]) colors = da.data.flatten() polys = matplotlib.collections.PatchCollection( - patches, alpha = alpha, norm = norm, cmap = cmap, + patches, alpha = 1, norm = norm, cmap = cmap, antialiaseds = antialias, edgecolors = linecolor, linewidths = linewidth, @@ -930,7 +934,7 @@ def plot2d_polygon( polys.set_array(colors) if add_colorbar: - fig.colorbar(polys, ax = cax) + fig.colorbar(polys, ax = cax, label = colorbar_label) ax.add_collection(polys) ax.set_aspect("equal", adjustable="box") @@ -945,9 +949,9 @@ def plot2d_polygon( color = "white" ls = "-" - lw = 1 + lw = 2 - m = da.metadata + m = da.attrs["metadata"] From d76febe3dc1e00a61f84f026f62badfdba69d483 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 22 Mar 2024 15:43:35 +0000 Subject: [PATCH 11/25] Make polygon() use standard sepx plotter Also modified plotter to accept style kwargs --- xbout/plotting/plotfuncs.py | 28 ++-------------------------- xbout/plotting/utils.py | 12 +++++++++--- 2 files changed, 11 insertions(+), 29 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 1e9a0dcd..32007c3b 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -864,6 +864,7 @@ def plot2d_polygon( add_colorbar = True, colorbar_label = None, separatrix = True, + separatrix_kwargs = {"color":"white", "linestyle":"-", "linewidth":1}, targets = False, add_limiter_hatching=True, grid_only = False, @@ -945,32 +946,7 @@ def plot2d_polygon( ax.set_title(da.name) if separatrix: - # plot_separatrices(da, ax, x = "R", y = "Z") + plot_separatrices(da, ax, x = "R", y = "Z", **separatrix_kwargs) - color = "white" - ls = "-" - lw = 2 - - m = da.attrs["metadata"] - - - - if m["topology"] == "connected-double-null": - - # TODO: test this for single null - if m["MXG"] > 0 or m["MYG"] > 0: - lhs = (m["ixseps1"] + 1 - m["MXG"], slice(0,m["ny_inner"] + m["MYG"] * int(4/2))) - rhs = (m["ixseps1"] + 1 - m["MXG"], slice(m["ny_inner"]+m["MYG"]*(4 - 1), None)) - else: - lhs = (m["ixseps1"]-m["MXG"]-1, slice(0,m["ny_inner"])) - rhs = (m["ixseps1"]-m["MXG"]-1, slice(m["ny_inner"], None)) - ax.plot(da["Rxy_lower_right_corners"].data[lhs], da["Zxy_lower_right_corners"].data[lhs], c = color, lw = lw, ls = ls) - ax.plot(da["Rxy_lower_right_corners"].data[rhs], da["Zxy_lower_right_corners"].data[rhs], c = color, lw = lw, ls = ls) - - if m["topology"] == "single-null": - points = (m["ixseps1"]-m["MXG"]-1, slice(0,None)) - ax.plot(da["Rxy_lower_right_corners"].data[points], da["Zxy_lower_right_corners"].data[points], c = color, lw = lw, ls = ls) - - if targets: plot_targets(da, ax, x = "R", y = "Z", hatching = add_limiter_hatching) \ No newline at end of file diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index db70b74a..0e19331e 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -78,8 +78,10 @@ def _is_core_only(da): return ix1 >= nx and ix2 >= nx -def plot_separatrices(da, ax, *, x="R", y="Z"): - """Plot separatrices""" +def plot_separatrices(da, ax, *, x="R", y="Z", **kwargs): + """ + Plot separatrices. Kwargs are passed to ax.plot(). + """ if not isinstance(da, dict): da_regions = _decompose_regions(da) @@ -116,7 +118,11 @@ def plot_separatrices(da, ax, *, x="R", y="Z"): y_sep = 0.5 * ( da_region[y].isel(**{xcoord: 0}) + da_inner[y].isel(**{xcoord: -1}) ) - ax.plot(x_sep, y_sep, "k--") + default_style = {"color": "black", "linestyle": "--"} + if any(x for x in kwargs if x in ["c", "ls"]): + raise ValueError("When passing separatrix plot style kwargs, use 'color' and 'linestyle' instead lf 'c' and 'ls'") + style = {**default_style, **kwargs} + ax.plot(x_sep, y_sep, **style) def plot_targets(da, ax, *, x="R", y="Z", hatching=True): From d1fdc43aa17eb4e34f77f0509fe9cf28f334979d Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Fri, 22 Mar 2024 15:45:14 +0000 Subject: [PATCH 12/25] Apply black formatting --- xbout/boutdataarray.py | 2 +- xbout/geometries.py | 33 ++++---- xbout/plotting/plotfuncs.py | 117 ++++++++++++++++------------ xbout/plotting/utils.py | 4 +- xbout/tests/test_against_collect.py | 3 +- xbout/tests/test_load.py | 6 +- xbout/utils.py | 16 ++-- 7 files changed, 101 insertions(+), 80 deletions(-) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index b3a4edba..2a2c5ed6 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -1071,7 +1071,7 @@ def pcolormesh(self, ax=None, **kwargs): Colour-plot a radial-poloidal slice on the R-Z plane """ return plotfuncs.plot2d_wrapper(self.data, xr.plot.pcolormesh, ax=ax, **kwargs) - + def polygon(self, ax=None, **kwargs): """ Colour-plot of a radial-poloidal slice on the R-Z plane using polygons diff --git a/xbout/geometries.py b/xbout/geometries.py index 985422ee..2299f54b 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -381,14 +381,14 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): "total_poloidal_distance", "zShift", "zShift_ylow", - "Rxy_corners", # Lower left corners + "Rxy_corners", # Lower left corners "Rxy_lower_right_corners", "Rxy_upper_left_corners", "Rxy_upper_right_corners", - "Zxy_corners", # Lower left corners + "Zxy_corners", # Lower left corners "Zxy_lower_right_corners", "Zxy_upper_left_corners", - "Zxy_upper_right_corners" + "Zxy_upper_right_corners", ], ) @@ -427,23 +427,24 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): ds = ds.set_coords(("R", "Z")) else: ds = ds.set_coords(("Rxy", "Zxy")) - + # Add cell corners as coordinates for polygon plotting if "Rxy_lower_right_corners" in ds: ds = ds.rename( - Rxy_corners = "Rxy_lower_left_corners", - Zxy_corners = "Zxy_lower_left_corners" + Rxy_corners="Rxy_lower_left_corners", Zxy_corners="Zxy_lower_left_corners" + ) + ds = ds.set_coords( + ( + "Rxy_lower_left_corners", + "Rxy_lower_right_corners", + "Rxy_upper_left_corners", + "Rxy_upper_right_corners", + "Zxy_lower_left_corners", + "Zxy_lower_right_corners", + "Zxy_upper_left_corners", + "Zxy_upper_right_corners", ) - ds = ds.set_coords(( - "Rxy_lower_left_corners", - "Rxy_lower_right_corners", - "Rxy_upper_left_corners", - "Rxy_upper_right_corners", - "Zxy_lower_left_corners", - "Zxy_lower_right_corners", - "Zxy_upper_left_corners", - "Zxy_upper_right_corners" - )) + ) # Rename zShift_ylow if it was added from grid file, to be consistent with name if # it was added from dump file diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 32007c3b..1ea10bc1 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -752,9 +752,9 @@ def create_or_update_plot(plot_objects=None, tind=None, this_save_as=None): X, Y, Z, scalars=data, vmin=vmin, vmax=vmax, **kwargs ) else: - plot_objects[ - region_name + str(i) - ].mlab_source.scalars = data + plot_objects[region_name + str(i)].mlab_source.scalars = ( + data + ) if mayavi_view is not None: mlab.view(*mayavi_view) @@ -850,49 +850,64 @@ def animation_func(): else: raise ValueError(f"Unrecognised plot3d() 'engine' argument: {engine}") + def plot2d_polygon( da, - ax = None, - cax = None, - cmap = "viridis", - norm = None, - logscale = False, - antialias = False, - vmin = None, - vmax = None, - extend = None, - add_colorbar = True, - colorbar_label = None, - separatrix = True, - separatrix_kwargs = {"color":"white", "linestyle":"-", "linewidth":1}, - targets = False, + ax=None, + cax=None, + cmap="viridis", + norm=None, + logscale=False, + antialias=False, + vmin=None, + vmax=None, + extend=None, + add_colorbar=True, + colorbar_label=None, + separatrix=True, + separatrix_kwargs={"color": "white", "linestyle": "-", "linewidth": 1}, + targets=False, add_limiter_hatching=True, - grid_only = False, - linewidth = 0, - linecolor = "black", - + grid_only=False, + linewidth=0, + linecolor="black", ): - + if ax == None: - fig, ax = plt.subplots(figsize=(3, 6), dpi = 120) + fig, ax = plt.subplots(figsize=(3, 6), dpi=120) else: fig = ax.get_figure() - + if cax == None: - cax = ax - + cax = ax + if vmin is None: vmin = np.nanmin(da.values) - + if vmax is None: vmax = np.nanmax(da.max().values) - if "Rxy_lower_right_corners" in da.coords: - r_nodes = ["R", "Rxy_lower_left_corners", "Rxy_lower_right_corners", "Rxy_upper_left_corners", "Rxy_upper_right_corners"] - z_nodes = ["Z", "Zxy_lower_left_corners", "Zxy_lower_right_corners", "Zxy_upper_left_corners", "Zxy_upper_right_corners"] - cell_r = np.concatenate([np.expand_dims(da[x], axis = 2) for x in r_nodes], axis = 2) - cell_z = np.concatenate([np.expand_dims(da[x], axis = 2) for x in z_nodes], axis = 2) + r_nodes = [ + "R", + "Rxy_lower_left_corners", + "Rxy_lower_right_corners", + "Rxy_upper_left_corners", + "Rxy_upper_right_corners", + ] + z_nodes = [ + "Z", + "Zxy_lower_left_corners", + "Zxy_lower_right_corners", + "Zxy_upper_left_corners", + "Zxy_upper_right_corners", + ] + cell_r = np.concatenate( + [np.expand_dims(da[x], axis=2) for x in r_nodes], axis=2 + ) + cell_z = np.concatenate( + [np.expand_dims(da[x], axis=2) for x in z_nodes], axis=2 + ) else: raise Exception("Cell corners not present in mesh, cannot do polygon plot") @@ -908,14 +923,15 @@ def plot2d_polygon( for i in range(Nx): for j in range(Ny): p = matplotlib.patches.Polygon( - np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])).reshape(2, 5).T, + np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])) + .reshape(2, 5) + .T, fill=False, closed=True, - facecolor = None, + facecolor=None, ) patches.append(p) - # create colorbar norm = _create_norm(logscale, norm, vmin, vmax) @@ -923,30 +939,31 @@ def plot2d_polygon( cmap = matplotlib.colors.ListedColormap(["white"]) colors = da.data.flatten() polys = matplotlib.collections.PatchCollection( - patches, alpha = 1, norm = norm, cmap = cmap, - antialiaseds = antialias, - edgecolors = linecolor, - linewidths = linewidth, - joinstyle = "bevel") - - - + patches, + alpha=1, + norm=norm, + cmap=cmap, + antialiaseds=antialias, + edgecolors=linecolor, + linewidths=linewidth, + joinstyle="bevel", + ) polys.set_array(colors) if add_colorbar: - fig.colorbar(polys, ax = cax, label = colorbar_label) - ax.add_collection(polys) - + fig.colorbar(polys, ax=cax, label=colorbar_label) + ax.add_collection(polys) + ax.set_aspect("equal", adjustable="box") ax.set_xlabel("R [m]") ax.set_ylabel("Z [m]") ax.set_ylim(cell_z.min(), cell_z.max()) ax.set_xlim(cell_r.min(), cell_r.max()) ax.set_title(da.name) - + if separatrix: - plot_separatrices(da, ax, x = "R", y = "Z", **separatrix_kwargs) - + plot_separatrices(da, ax, x="R", y="Z", **separatrix_kwargs) + if targets: - plot_targets(da, ax, x = "R", y = "Z", hatching = add_limiter_hatching) \ No newline at end of file + plot_targets(da, ax, x="R", y="Z", hatching=add_limiter_hatching) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 0e19331e..8a8d3c1a 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -120,7 +120,9 @@ def plot_separatrices(da, ax, *, x="R", y="Z", **kwargs): ) default_style = {"color": "black", "linestyle": "--"} if any(x for x in kwargs if x in ["c", "ls"]): - raise ValueError("When passing separatrix plot style kwargs, use 'color' and 'linestyle' instead lf 'c' and 'ls'") + raise ValueError( + "When passing separatrix plot style kwargs, use 'color' and 'linestyle' instead lf 'c' and 'ls'" + ) style = {**default_style, **kwargs} ax.plot(x_sep, y_sep, **style) diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py index 5f22cf97..a1bf3da3 100644 --- a/xbout/tests/test_against_collect.py +++ b/xbout/tests/test_against_collect.py @@ -220,5 +220,4 @@ def test_new_collect_indexing_slice(self, tmp_path_factory): @pytest.mark.skip -class test_speed_against_old_collect: - ... +class test_speed_against_old_collect: ... diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index d8766236..bb4c917e 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -472,8 +472,7 @@ def test_combine_along_y(self, tmp_path_factory, bout_xyt_example_files): xrt.assert_identical(actual, fake) @pytest.mark.skip - def test_combine_along_t(self): - ... + def test_combine_along_t(self): ... @pytest.mark.parametrize( "bout_v5,metric_3D", [(False, False), (True, False), (True, True)] @@ -623,8 +622,7 @@ def test_drop_vars(self, tmp_path_factory, bout_xyt_example_files): assert "n" in ds.keys() @pytest.mark.skip - def test_combine_along_tx(self): - ... + def test_combine_along_tx(self): ... def test_restarts(self): datapath = Path(__file__).parent.joinpath( diff --git a/xbout/utils.py b/xbout/utils.py index 32be7edc..66dd9593 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -167,12 +167,16 @@ def _1d_coord_from_spacing(spacing, dim, ds=None, *, origin_at=None): ) point_to_use = { - spacing.metadata["bout_xdim"]: spacing.metadata.get("MXG", 0) - if spacing.metadata["keep_xboundaries"] - else 0, - spacing.metadata["bout_ydim"]: spacing.metadata.get("MYG", 0) - if spacing.metadata["keep_yboundaries"] - else 0, + spacing.metadata["bout_xdim"]: ( + spacing.metadata.get("MXG", 0) + if spacing.metadata["keep_xboundaries"] + else 0 + ), + spacing.metadata["bout_ydim"]: ( + spacing.metadata.get("MYG", 0) + if spacing.metadata["keep_yboundaries"] + else 0 + ), spacing.metadata["bout_zdim"]: spacing.metadata.get("MZG", 0), } From a74a197c267981555793761bf9d59ddf2bac5c00 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Mon, 25 Mar 2024 10:42:47 +0000 Subject: [PATCH 13/25] Polygon: add docstring, tidy up settings --- xbout/plotting/plotfuncs.py | 54 +++++++++++++++++++++++++++++++++---- 1 file changed, 49 insertions(+), 5 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 32007c3b..d34fb0e7 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -860,18 +860,63 @@ def plot2d_polygon( antialias = False, vmin = None, vmax = None, - extend = None, + extend = "neither", add_colorbar = True, colorbar_label = None, separatrix = True, - separatrix_kwargs = {"color":"white", "linestyle":"-", "linewidth":1}, - targets = False, + separatrix_kwargs = {"color":"black", "linestyle":"--", "linewidth":2}, + targets = True, add_limiter_hatching=True, grid_only = False, linewidth = 0, linecolor = "black", ): + """ + Nice looking 2D plots which have no visual artifacts around the X-point. + + Parameters + ---------- + da : xarray.DataArray + A 2D (x,y) DataArray of data to plot + ax : Axes, optional + Axes to plot on. If not provided, will make its own. + cax : Axes, optional + Axes to plot colorbar on. If not provided, will plot on the same axes as the plot. + cmap : str or matplotlib.colors.Colormap, default "viridis" + Colormap to use for the plot + norm : matplotlib.colors.Normalize, optional + Normalization to use for the color scale + logscale : bool, default False + If True, use a symlog color scale + antialias : bool, default False + Enables antialiasing. Note: this also shows mesh cell edges - it's unclear how to disable this. + vmin : float, optional + Minimum value for the color scale + vmax : float, optional + Maximum value for the color scale + extend : str, optional, default "neither" + Extend the colorbar. Options are "neither", "both", "min", "max" + add_colorbar : bool, default True + Enable colorbar in figure? + colorbar_label : str, optional + Label for the colorbar + separatrix : bool, default True + Add lines showing separatrices + separatrix_kwargs : dict + Keyword arguments to pass custom style to the separatrices plot + targets : bool, default True + Draw solid lines at the target surfaces + add_limiter_hatching : bool, default True + Draw hatched areas at the targets + grid_only : bool, default False + Only plot the grid, not the data. This sets all the polygons to have a white face. + linewidth : float, default 0 + Width of the gridlines on cell edges + linecolor : str, default "black" + Color of the gridlines on cell edges + """ + if ax == None: fig, ax = plt.subplots(figsize=(3, 6), dpi = 120) @@ -898,7 +943,6 @@ def plot2d_polygon( Nx = len(cell_r) Ny = len(cell_r[0]) - edgecolor = "black" patches = [] # https://matplotlib.org/2.0.2/examples/api/patch_collection.html @@ -935,7 +979,7 @@ def plot2d_polygon( polys.set_array(colors) if add_colorbar: - fig.colorbar(polys, ax = cax, label = colorbar_label) + fig.colorbar(polys, ax = cax, label = colorbar_label, extend = extend) ax.add_collection(polys) ax.set_aspect("equal", adjustable="box") From 59f82e45bf1cf88e3e176f4a2c7952b5b08a3741 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Mon, 25 Mar 2024 10:43:03 +0000 Subject: [PATCH 14/25] Black formatting --- xbout/plotting/plotfuncs.py | 120 ++++++++++++++++++++---------------- 1 file changed, 68 insertions(+), 52 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index d34fb0e7..894c556a 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -752,9 +752,9 @@ def create_or_update_plot(plot_objects=None, tind=None, this_save_as=None): X, Y, Z, scalars=data, vmin=vmin, vmax=vmax, **kwargs ) else: - plot_objects[ - region_name + str(i) - ].mlab_source.scalars = data + plot_objects[region_name + str(i)].mlab_source.scalars = ( + data + ) if mayavi_view is not None: mlab.view(*mayavi_view) @@ -850,31 +850,31 @@ def animation_func(): else: raise ValueError(f"Unrecognised plot3d() 'engine' argument: {engine}") + def plot2d_polygon( da, - ax = None, - cax = None, - cmap = "viridis", - norm = None, - logscale = False, - antialias = False, - vmin = None, - vmax = None, - extend = "neither", - add_colorbar = True, - colorbar_label = None, - separatrix = True, - separatrix_kwargs = {"color":"black", "linestyle":"--", "linewidth":2}, - targets = True, + ax=None, + cax=None, + cmap="viridis", + norm=None, + logscale=False, + antialias=False, + vmin=None, + vmax=None, + extend="neither", + add_colorbar=True, + colorbar_label=None, + separatrix=True, + separatrix_kwargs={"color": "black", "linestyle": "--", "linewidth": 2}, + targets=True, add_limiter_hatching=True, - grid_only = False, - linewidth = 0, - linecolor = "black", - + grid_only=False, + linewidth=0, + linecolor="black", ): """ Nice looking 2D plots which have no visual artifacts around the X-point. - + Parameters ---------- da : xarray.DataArray @@ -916,28 +916,42 @@ def plot2d_polygon( linecolor : str, default "black" Color of the gridlines on cell edges """ - - + if ax == None: - fig, ax = plt.subplots(figsize=(3, 6), dpi = 120) + fig, ax = plt.subplots(figsize=(3, 6), dpi=120) else: fig = ax.get_figure() - + if cax == None: - cax = ax - + cax = ax + if vmin is None: vmin = np.nanmin(da.values) - + if vmax is None: vmax = np.nanmax(da.max().values) - if "Rxy_lower_right_corners" in da.coords: - r_nodes = ["R", "Rxy_lower_left_corners", "Rxy_lower_right_corners", "Rxy_upper_left_corners", "Rxy_upper_right_corners"] - z_nodes = ["Z", "Zxy_lower_left_corners", "Zxy_lower_right_corners", "Zxy_upper_left_corners", "Zxy_upper_right_corners"] - cell_r = np.concatenate([np.expand_dims(da[x], axis = 2) for x in r_nodes], axis = 2) - cell_z = np.concatenate([np.expand_dims(da[x], axis = 2) for x in z_nodes], axis = 2) + r_nodes = [ + "R", + "Rxy_lower_left_corners", + "Rxy_lower_right_corners", + "Rxy_upper_left_corners", + "Rxy_upper_right_corners", + ] + z_nodes = [ + "Z", + "Zxy_lower_left_corners", + "Zxy_lower_right_corners", + "Zxy_upper_left_corners", + "Zxy_upper_right_corners", + ] + cell_r = np.concatenate( + [np.expand_dims(da[x], axis=2) for x in r_nodes], axis=2 + ) + cell_z = np.concatenate( + [np.expand_dims(da[x], axis=2) for x in z_nodes], axis=2 + ) else: raise Exception("Cell corners not present in mesh, cannot do polygon plot") @@ -952,14 +966,15 @@ def plot2d_polygon( for i in range(Nx): for j in range(Ny): p = matplotlib.patches.Polygon( - np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])).reshape(2, 5).T, + np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])) + .reshape(2, 5) + .T, fill=False, closed=True, - facecolor = None, + facecolor=None, ) patches.append(p) - # create colorbar norm = _create_norm(logscale, norm, vmin, vmax) @@ -967,30 +982,31 @@ def plot2d_polygon( cmap = matplotlib.colors.ListedColormap(["white"]) colors = da.data.flatten() polys = matplotlib.collections.PatchCollection( - patches, alpha = 1, norm = norm, cmap = cmap, - antialiaseds = antialias, - edgecolors = linecolor, - linewidths = linewidth, - joinstyle = "bevel") - - - + patches, + alpha=1, + norm=norm, + cmap=cmap, + antialiaseds=antialias, + edgecolors=linecolor, + linewidths=linewidth, + joinstyle="bevel", + ) polys.set_array(colors) if add_colorbar: - fig.colorbar(polys, ax = cax, label = colorbar_label, extend = extend) - ax.add_collection(polys) - + fig.colorbar(polys, ax=cax, label=colorbar_label, extend=extend) + ax.add_collection(polys) + ax.set_aspect("equal", adjustable="box") ax.set_xlabel("R [m]") ax.set_ylabel("Z [m]") ax.set_ylim(cell_z.min(), cell_z.max()) ax.set_xlim(cell_r.min(), cell_r.max()) ax.set_title(da.name) - + if separatrix: - plot_separatrices(da, ax, x = "R", y = "Z", **separatrix_kwargs) - + plot_separatrices(da, ax, x="R", y="Z", **separatrix_kwargs) + if targets: - plot_targets(da, ax, x = "R", y = "Z", hatching = add_limiter_hatching) \ No newline at end of file + plot_targets(da, ax, x="R", y="Z", hatching=add_limiter_hatching) From 49790ce4338765752999461fbf59055c231acca6 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Mon, 25 Mar 2024 10:51:37 +0000 Subject: [PATCH 15/25] Replace == None with is None --- xbout/plotting/plotfuncs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 894c556a..37cc27ac 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -917,12 +917,12 @@ def plot2d_polygon( Color of the gridlines on cell edges """ - if ax == None: + if ax is None: fig, ax = plt.subplots(figsize=(3, 6), dpi=120) else: fig = ax.get_figure() - if cax == None: + if cax is None: cax = ax if vmin is None: From f6aa96f816a24d1c6079a91f364e99bee60ffa4b Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Wed, 22 May 2024 10:44:37 +0100 Subject: [PATCH 16/25] Add colorbar label --- xbout/plotting/plotfuncs.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 37cc27ac..04d2c242 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1,6 +1,7 @@ from collections.abc import Sequence import matplotlib import matplotlib.pyplot as plt +from mpl_toolkits.axes_grid1 import make_axes_locatable import numpy as np from pathlib import Path from tempfile import TemporaryDirectory @@ -930,6 +931,16 @@ def plot2d_polygon( if vmax is None: vmax = np.nanmax(da.max().values) + + if colorbar_label is not None: + cbar_label = colorbar_label + else: + if "long_name" in da.attrs: + cbar_label = da.long_name + else: + cbar_label = da.name + if "units" in da.attrs: + cbar_label += f" [{da.units}]" if "Rxy_lower_right_corners" in da.coords: r_nodes = [ @@ -975,7 +986,6 @@ def plot2d_polygon( ) patches.append(p) - # create colorbar norm = _create_norm(logscale, norm, vmin, vmax) if grid_only is True: From c09fc4517ad8364135a09a3652e1fee66e001ed5 Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Thu, 23 May 2024 08:49:30 +0000 Subject: [PATCH 17/25] Apply black formatting --- xbout/plotting/plotfuncs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 04d2c242..7d00fbfb 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -931,7 +931,7 @@ def plot2d_polygon( if vmax is None: vmax = np.nanmax(da.max().values) - + if colorbar_label is not None: cbar_label = colorbar_label else: From 529c4e356bc0758c95d16e7af73ae12d9ca00de1 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Sat, 25 May 2024 10:08:44 +0100 Subject: [PATCH 18/25] Make polygon plot colorbar always the height of plot - This becomes a "foolproof" colorbar that won't mess up the aspect ratio of the figure. --- xbout/plotting/plotfuncs.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 7d00fbfb..14972e95 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1005,7 +1005,14 @@ def plot2d_polygon( polys.set_array(colors) if add_colorbar: - fig.colorbar(polys, ax=cax, label=colorbar_label, extend=extend) + # This produces a "foolproof" colorbar which + # is always the height of the plot + # From https://joseph-long.com/writing/colorbars/ + divider = make_axes_locatable(ax) + cax = divider.append_axes("right", size="5%", pad=0.05) + fig.colorbar(polys, cax = cax, label=colorbar_label, extend=extend) + cax.grid(visible = False) + ax.add_collection(polys) ax.set_aspect("equal", adjustable="box") From 491be954dda7eb63deb38b748c62b44f10e1fd6a Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Sat, 25 May 2024 10:09:38 +0100 Subject: [PATCH 19/25] Change polygon plot defaults for a cleaner look --- xbout/plotting/plotfuncs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 14972e95..cf2ed472 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -866,9 +866,9 @@ def plot2d_polygon( add_colorbar=True, colorbar_label=None, separatrix=True, - separatrix_kwargs={"color": "black", "linestyle": "--", "linewidth": 2}, - targets=True, - add_limiter_hatching=True, + separatrix_kwargs={"color": "white", "linestyle": "-", "linewidth": 2}, + targets=False, + add_limiter_hatching=False, grid_only=False, linewidth=0, linecolor="black", From edfbfb83cef413893048902cd62d2b451b413ec7 Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Sat, 25 May 2024 09:11:00 +0000 Subject: [PATCH 20/25] Apply black formatting --- xbout/plotting/plotfuncs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index cf2ed472..9105070b 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1005,13 +1005,13 @@ def plot2d_polygon( polys.set_array(colors) if add_colorbar: - # This produces a "foolproof" colorbar which + # This produces a "foolproof" colorbar which # is always the height of the plot # From https://joseph-long.com/writing/colorbars/ divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) - fig.colorbar(polys, cax = cax, label=colorbar_label, extend=extend) - cax.grid(visible = False) + fig.colorbar(polys, cax=cax, label=colorbar_label, extend=extend) + cax.grid(visible=False) ax.add_collection(polys) From 068b3dcdf958f4d4af8fd1093fce37b9fb48776e Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Sat, 25 May 2024 10:18:55 +0100 Subject: [PATCH 21/25] polygon: fix grid showing in colorbar --- xbout/plotting/plotfuncs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 9105070b..e1ac3007 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1011,7 +1011,7 @@ def plot2d_polygon( divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) fig.colorbar(polys, cax=cax, label=colorbar_label, extend=extend) - cax.grid(visible=False) + cax.grid(which = "both", visible=False) ax.add_collection(polys) From d1c5db6cd696cd11cfb2754d4c17e2551dc2d4fc Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Sat, 25 May 2024 09:30:23 +0000 Subject: [PATCH 22/25] Apply black formatting --- xbout/plotting/plotfuncs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index e1ac3007..16f84499 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1011,7 +1011,7 @@ def plot2d_polygon( divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) fig.colorbar(polys, cax=cax, label=colorbar_label, extend=extend) - cax.grid(which = "both", visible=False) + cax.grid(which="both", visible=False) ax.add_collection(polys) From 40781f1ac0349ec0a8a0fef59b7c68bdad2f387d Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Sat, 25 May 2024 13:02:38 +0100 Subject: [PATCH 23/25] Polygon plots: fix cbar label not appearing, change to short name --- xbout/plotting/plotfuncs.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index e1ac3007..fa9fbd19 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -932,15 +932,13 @@ def plot2d_polygon( if vmax is None: vmax = np.nanmax(da.max().values) - if colorbar_label is not None: - cbar_label = colorbar_label - else: - if "long_name" in da.attrs: - cbar_label = da.long_name + if colorbar_label is None: + if "short_name" in da.attrs: + colorbar_label = da.short_name else: - cbar_label = da.name + colorbar_label = da.name if "units" in da.attrs: - cbar_label += f" [{da.units}]" + colorbar_label += f" [{da.units}]" if "Rxy_lower_right_corners" in da.coords: r_nodes = [ @@ -1011,7 +1009,7 @@ def plot2d_polygon( divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) fig.colorbar(polys, cax=cax, label=colorbar_label, extend=extend) - cax.grid(which = "both", visible=False) + cax.grid(which="both", visible=False) ax.add_collection(polys) From afbeaf013de5189457ff2e41b70f669ffb215f41 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Sat, 25 May 2024 15:45:46 +0100 Subject: [PATCH 24/25] Polygon plots: improve handling of colorbar label - Now handles cases when there is no variable name, e.g. when you plot a derived variable like ds["Ne"] * ds["Te"] --- xbout/plotting/plotfuncs.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index fa9fbd19..db0a1e67 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -932,13 +932,16 @@ def plot2d_polygon( if vmax is None: vmax = np.nanmax(da.max().values) - if colorbar_label is None: + if colorbar_label == None: if "short_name" in da.attrs: - colorbar_label = da.short_name - else: + colorbar_label = da.attrs["short_name"] + elif da.name != None: colorbar_label = da.name - if "units" in da.attrs: - colorbar_label += f" [{da.units}]" + else: + colorbar_label = "" + + if "units" in da.attrs: + colorbar_label += f" [{da.attrs['units']}]" if "Rxy_lower_right_corners" in da.coords: r_nodes = [ From ba35f4673791af90fd277b35e4b4eae6b9d8a913 Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Tue, 28 May 2024 16:22:43 +0100 Subject: [PATCH 25/25] Expose separatrix kwargs setting in animate_poloidal --- xbout/plotting/animate.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 68536a97..e223983f 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -97,6 +97,7 @@ def animate_poloidal( cax=None, animate_over=None, separatrix=True, + separatrix_kwargs=dict(), targets=True, add_limiter_hatching=True, cmap=None, @@ -130,6 +131,8 @@ def animate_poloidal( Dimension over which to animate, defaults to the time dimension separatrix : bool, optional Add dashed lines showing separatrices + separatrix_kwargs : dict, optional + Options to pass to the separatrix plotter (e.g. line color) targets : bool, optional Draw solid lines at the target surfaces add_limiter_hatching : bool, optional @@ -277,7 +280,7 @@ def animate_poloidal( targets = False if separatrix: - plot_separatrices(da_regions, ax, x=x, y=y) + plot_separatrices(da_regions, ax, x=x, y=y, **separatrix_kwargs) if targets: plot_targets(da_regions, ax, x=x, y=y, hatching=add_limiter_hatching)