From 594ab43600c00559cbf277530ee038f9a754799d Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 6 Mar 2026 08:39:06 +0000 Subject: [PATCH 1/8] _from_region: avoid unnecessary deep copy This had a huge performance cost because of the recursive deep copies of the regions in .attrs.... which were unnecessarily copied once again manually a few lines later!! Instead, deep copy only the single required region, and copy the .attrs with the regions separately. --- xbout/region.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/xbout/region.py b/xbout/region.py index fcd63bce..17fd540c 100644 --- a/xbout/region.py +++ b/xbout/region.py @@ -1712,9 +1712,6 @@ def _concat_upper_guards(da, da_global, mxg, myg): def _from_region(ds_or_da, name, with_guards): - # ensure we do not modify the input - ds_or_da = ds_or_da.copy(deep=True) - region = ds_or_da.bout._regions[name] xcoord = ds_or_da.metadata["bout_xdim"] ycoord = ds_or_da.metadata["bout_ydim"] @@ -1730,10 +1727,10 @@ def _from_region(ds_or_da, name, with_guards): mxg = with_guards myg = with_guards - result = ds_or_da.isel(region.get_slices()).copy() + result = ds_or_da.isel(region.get_slices()).copy(deep=False) - # The returned result has only one region single_region = deepcopy(region) + result.attrs = dict(ds_or_da.attrs) result.attrs["regions"] = {name: single_region} # get inner x-guard cells for result from the global array From addd63e2320a74f0aa939b9e9d2dcff968de2104 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 6 Mar 2026 08:40:48 +0000 Subject: [PATCH 2/8] plot2d_polygon: vectorised polygon creation Significantly improves performance by reducing overhead --- xbout/plotting/plotfuncs.py | 33 +++++++++++---------------------- 1 file changed, 11 insertions(+), 22 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 9c54d183..024ecd9c 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -965,40 +965,29 @@ def plot2d_polygon( Nx = len(cell_r) Ny = len(cell_r[0]) - 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, - facecolor=None, - ) - patches.append(p) + + # Build polygon vertices vectorized instead of looping over each cell. + # idx selects corners in order: lower-left, lower-right, upper-right, upper-left. + # PolyCollection closes the polygon automatically. + idx = np.array([1, 2, 4, 3]) + # verts shape: (Nx, Ny, 4, 2) + verts = np.stack([cell_r[:, :, idx], cell_z[:, :, idx]], axis=-1).reshape(-1, 4, 2) 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=1, + colors = np.asarray(da.values).flatten() + polys = matplotlib.collections.PolyCollection( + verts, norm=norm, cmap=cmap, + alpha=1, antialiaseds=antialias, edgecolors=linecolor, linewidths=linewidth, joinstyle="bevel", ) - polys.set_array(colors) if add_colorbar: From 8596dd75f8e7c3a9a81e98b0808b62d2b9326665 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 6 Mar 2026 08:45:49 +0000 Subject: [PATCH 3/8] plot_separatrices: avoid xr.align Big performance boost: xr.align compares all coordinates for equality. Fall back on xr.align for edge cases like limiter configuration. --- xbout/plotting/utils.py | 50 ++++++++++++++++++++++++++--------------- 1 file changed, 32 insertions(+), 18 deletions(-) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 8a8d3c1a..3ad60002 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -98,26 +98,40 @@ def plot_separatrices(da, ax, *, x="R", y="Z", **kwargs): if inner in da_regions: da_inner = da_regions[inner] + # Extract boundary columns as numpy arrays to avoid expensive xarray + # coordinate-alignment (xr.align compares every coordinate, triggering + # dask computation of lazy arrays). + r_region = da_region[x].isel(**{xcoord: 0}) + r_inner = da_inner[x].isel(**{xcoord: -1}) + z_region = da_region[y].isel(**{xcoord: 0}) + z_inner = da_inner[y].isel(**{xcoord: -1}) + try: - da_region, da_inner = xr.align(da_region, da_inner) + x_sep = 0.5 * (r_region.values + r_inner.values) + y_sep = 0.5 * (z_region.values + z_inner.values) except ValueError: - # For geometries with a limiter, the closed field-line region may have - # guard cells while the open field line region does not. Also the - # closed-field line guard cells may (if the region is connected to - # itself) have duplicated coordinate values, which xr.align() cannot - # handle. Use np.unique() to remove the duplicated coordinate values - _, unique_yinds = np.unique(da_inner[ycoord], return_index=True) - da_inner = da_inner.isel(**{ycoord: unique_yinds}) - - # Put da_inner second as the unique_yinds selection may mess up the order of - # points. xarray will align the coordinates with the first argument (to the - # addition here). - x_sep = 0.5 * ( - da_region[x].isel(**{xcoord: 0}) + da_inner[x].isel(**{xcoord: -1}) - ) - y_sep = 0.5 * ( - da_region[y].isel(**{xcoord: 0}) + da_inner[y].isel(**{xcoord: -1}) - ) + # Arrays have different y-extents (e.g. limiter geometry where the + # closed field-line region has more guard cells than the open region, + # or duplicated coordinate values on a self-connected region). + # Fall back to xr.align to find the common coordinate intersection. + try: + da_r_al, da_i_al = xr.align(da_region, da_inner) + except ValueError: + # Duplicated coordinate values: deduplicate first, then align. + _, unique_yinds = np.unique( + da_inner[ycoord].values, return_index=True + ) + da_inner = da_inner.isel(**{ycoord: unique_yinds}) + da_r_al, da_i_al = xr.align(da_region, da_inner) + x_sep = 0.5 * ( + da_r_al[x].isel(**{xcoord: 0}).values + + da_i_al[x].isel(**{xcoord: -1}).values + ) + y_sep = 0.5 * ( + da_r_al[y].isel(**{xcoord: 0}).values + + da_i_al[y].isel(**{xcoord: -1}).values + ) + default_style = {"color": "black", "linestyle": "--"} if any(x for x in kwargs if x in ["c", "ls"]): raise ValueError( From 30617c0710b84af5bfc2bbc792447bd0eecdcebb Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 6 Mar 2026 08:46:57 +0000 Subject: [PATCH 4/8] plot2d_polygon: reduce coordinates before separatrix plot We should really take out corner coords from the coordinates. --- 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 024ecd9c..0a4959c0 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1008,8 +1008,16 @@ def plot2d_polygon( ax.set_xlim(cell_r.min(), cell_r.max()) ax.set_title(da.name) + if separatrix or targets: + # Drop the cell-corner coordinates (needed only for polygon construction) + # before decomposing regions to avoid deep-copying large arrays unnecessarily. + corner_coords = [ + c for c in da.coords if c.startswith("Rxy_") or c.startswith("Zxy_") + ] + da_minimal = da.drop_vars(corner_coords) if corner_coords else da + if separatrix: - plot_separatrices(da, ax, x="R", y="Z", **separatrix_kwargs) + plot_separatrices(da_minimal, ax, x="R", y="Z", **separatrix_kwargs) if targets: - plot_targets(da, ax, x="R", y="Z", hatching=add_limiter_hatching) + plot_targets(da_minimal, ax, x="R", y="Z", hatching=add_limiter_hatching) From a138f220065e00a87ed3f38974f1d6d1ae1ce07a Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 6 Mar 2026 09:01:32 +0000 Subject: [PATCH 5/8] plot_separatrices: clearer variable names --- xbout/plotting/utils.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 3ad60002..6a14f77d 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -115,21 +115,21 @@ def plot_separatrices(da, ax, *, x="R", y="Z", **kwargs): # or duplicated coordinate values on a self-connected region). # Fall back to xr.align to find the common coordinate intersection. try: - da_r_al, da_i_al = xr.align(da_region, da_inner) + da_region_aligned, da_inner_aligned = xr.align(da_region, da_inner) except ValueError: # Duplicated coordinate values: deduplicate first, then align. _, unique_yinds = np.unique( da_inner[ycoord].values, return_index=True ) da_inner = da_inner.isel(**{ycoord: unique_yinds}) - da_r_al, da_i_al = xr.align(da_region, da_inner) + da_region_aligned, da_inner_aligned = xr.align(da_region, da_inner) x_sep = 0.5 * ( - da_r_al[x].isel(**{xcoord: 0}).values - + da_i_al[x].isel(**{xcoord: -1}).values + da_region_aligned[x].isel(**{xcoord: 0}).values + + da_inner_aligned[x].isel(**{xcoord: -1}).values ) y_sep = 0.5 * ( - da_r_al[y].isel(**{xcoord: 0}).values - + da_i_al[y].isel(**{xcoord: -1}).values + da_region_aligned[y].isel(**{xcoord: 0}).values + + da_inner_aligned[y].isel(**{xcoord: -1}).values ) default_style = {"color": "black", "linestyle": "--"} From 2f389ef1cdd2bac7db0c8c4a39b6cfd6934e01eb Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Fri, 6 Mar 2026 09:07:29 +0000 Subject: [PATCH 6/8] Apply black formatting --- xbout/load.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index ab3e5a0d..2798c696 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -1146,9 +1146,11 @@ def _open_grid(datapath, chunks, keep_xboundaries, keep_yboundaries, mxg=2, **kw gridfilepath = Path(datapath) grid = xr.open_dataset( gridfilepath, - engine=file_engine - if file_engine is not None - else _check_filetype(gridfilepath), + engine=( + file_engine + if file_engine is not None + else _check_filetype(gridfilepath) + ), **kwargs, ) else: From 2c81cdc4ea0d19b561d6ed8eb3400a87ebc81f85 Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Fri, 6 Mar 2026 11:09:40 +0000 Subject: [PATCH 7/8] Remove unused variables --- xbout/plotting/plotfuncs.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 0a4959c0..d2169947 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -963,8 +963,6 @@ def plot2d_polygon( else: raise Exception("Cell corners not present in mesh, cannot do polygon plot") - Nx = len(cell_r) - Ny = len(cell_r[0]) # Build polygon vertices vectorized instead of looping over each cell. # idx selects corners in order: lower-left, lower-right, upper-right, upper-left. From 8853a47ade72067cfb845a9486833ba812e8b09d Mon Sep 17 00:00:00 2001 From: mikekryjak Date: Fri, 6 Mar 2026 11:09:20 +0000 Subject: [PATCH 8/8] Apply black formatting --- xbout/plotting/plotfuncs.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index d2169947..3796d543 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -963,7 +963,6 @@ def plot2d_polygon( else: raise Exception("Cell corners not present in mesh, cannot do polygon plot") - # Build polygon vertices vectorized instead of looping over each cell. # idx selects corners in order: lower-left, lower-right, upper-right, upper-left. # PolyCollection closes the polygon automatically.