20x faster 2D plots

xbout/plotting/plotfuncs.py

@@ -963,42 +963,28 @@ 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])
-    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:
@@ -1019,8 +1005,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)

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_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_region_aligned, da_inner_aligned = xr.align(da_region, da_inner)
+                x_sep = 0.5 * (
+                    da_region_aligned[x].isel(**{xcoord: 0}).values
+                    + da_inner_aligned[x].isel(**{xcoord: -1}).values
+                )
+                y_sep = 0.5 * (
+                    da_region_aligned[y].isel(**{xcoord: 0}).values
+                    + da_inner_aligned[y].isel(**{xcoord: -1}).values
+                )
+
             default_style = {"color": "black", "linestyle": "--"}
             if any(x for x in kwargs if x in ["c", "ls"]):
                 raise ValueError(

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