Fix concat_<side>_guards coords assignment

xbout/calc/turbulence.py

@@ -1,7 +1,6 @@
 import numpy as np
 import xarray as xr
 
-
 # TODO write a decorator to add functions as accessors to dataarrays?
 
 

xbout/conftest.py

@@ -1,6 +1,5 @@
 import pytest
 
-
 configure_is_imported = False
 
 

xbout/geometries.py

@@ -60,13 +60,9 @@ def apply_geometry(ds, geometry_name, *, coordinates=None, grid=None):
         try:
             add_geometry_coords = REGISTERED_GEOMETRIES[geometry_name]
         except KeyError:
-            message = dedent(
-                """{} is not a registered geometry. Inspect the global
+            message = dedent("""{} is not a registered geometry. Inspect the global
                              variable REGISTERED_GEOMETRIES to see which geometries
-                             have been registered.""".format(
-                    geometry_name
-                )
-            )
+                             have been registered.""".format(geometry_name))
             raise UnregisteredGeometryError(message)
 
         # User-registered functions may accept 'coordinates' and 'grid' arguments, but

xbout/plotting/animate.py

@@ -14,7 +14,6 @@
 )
 from matplotlib.animation import PillowWriter
 
-
 if (
     "pcolor.shading" in matplotlib.rcParams
     and matplotlib.rcParams["pcolor.shading"] == "flat"

xbout/plotting/plotfuncs.py

@@ -18,7 +18,6 @@
     plot_targets,
 )
 
-
 if (
     "pcolor.shading" in matplotlib.rcParams
     and matplotlib.rcParams["pcolor.shading"] == "flat"

xbout/region.py

@@ -1342,15 +1342,15 @@ def _concat_inner_guards(da, da_global, mxg):
         new_xcoord = da_global[xcoord].isel(**{xcoord: slices[xcoord]})
         new_ycoord = da_global[ycoord].isel(**{ycoord: slices[ycoord]})
 
-        # can't use commented out version, uncommented one works around xarray bug
-        # removing attrs
-        # https://github.com/pydata/xarray/issues/4415
-        # https://github.com/pydata/xarray/issues/4393
-        # da_inner = da_inner.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
-        da_inner[xcoord].data[...] = new_xcoord.data
-        da_inner = da_inner.reset_index(xcoord).set_xindex(xcoord)
-        da_inner[ycoord].data[...] = new_ycoord.data
-        da_inner = da_inner.reset_index(ycoord).set_xindex(ycoord)
+        # Preserve attributes when updating coordinates
+        # xarray coordinates are read-only by design, must use assign_coords()
+        xcoord_attrs = da_inner[xcoord].attrs.copy()
+        ycoord_attrs = da_inner[ycoord].attrs.copy()
+
+        # Assign new coordinate values
+        new_xcoord.attrs = xcoord_attrs
+        new_ycoord.attrs = ycoord_attrs
+        da_inner = da_inner.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
 
     save_regions = da.bout._regions
     da = xr.concat((da_inner, da), xcoord, join="exact")
@@ -1455,15 +1455,15 @@ def _concat_outer_guards(da, da_global, mxg):
         new_xcoord = da_global[xcoord].isel(**{xcoord: slices[xcoord]})
         new_ycoord = da_global[ycoord].isel(**{ycoord: slices[ycoord]})
 
-        # can't use commented out version, uncommented one works around xarray bug
-        # removing attrs
-        # https://github.com/pydata/xarray/issues/4415
-        # https://github.com/pydata/xarray/issues/4393
-        # da_outer = da_outer.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
-        da_outer[xcoord].data[...] = new_xcoord.data
-        da_outer = da_outer.reset_index(xcoord).set_xindex(xcoord)
-        da_outer[ycoord].data[...] = new_ycoord.data
-        da_outer = da_outer.reset_index(ycoord).set_xindex(ycoord)
+        # Preserve attributes when updating coordinates
+        # xarray coordinates are read-only by design, must use assign_coords()
+        xcoord_attrs = da_outer[xcoord].attrs.copy()
+        ycoord_attrs = da_outer[ycoord].attrs.copy()
+
+        # Assign new coordinate values
+        new_xcoord.attrs = xcoord_attrs
+        new_ycoord.attrs = ycoord_attrs
+        da_outer = da_outer.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
 
     save_regions = da.bout._regions
     da = xr.concat((da, da_outer), xcoord, join="exact")
@@ -1557,15 +1557,15 @@ def _concat_lower_guards(da, da_global, mxg, myg):
             )
             new_ycoord = new_ycoord + offset
 
-        # can't use commented out version, uncommented one works around xarray bug
-        # removing attrs
-        # https://github.com/pydata/xarray/issues/4415
-        # https://github.com/pydata/xarray/issues/4393
-        # da_lower = da_lower.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
-        da_lower[xcoord].data[...] = new_xcoord.data
-        da_lower = da_lower.reset_index(xcoord).set_xindex(xcoord)
-        da_lower[ycoord].data[...] = new_ycoord.data
-        da_lower = da_lower.reset_index(ycoord).set_xindex(ycoord)
+        # Preserve attributes when updating coordinates
+        # xarray coordinates are read-only by design, must use assign_coords()
+        xcoord_attrs = da_lower[xcoord].attrs.copy()
+        ycoord_attrs = da_lower[ycoord].attrs.copy()
+
+        # Assign new coordinate values
+        new_xcoord.attrs = xcoord_attrs
+        new_ycoord.attrs = ycoord_attrs
+        da_lower = da_lower.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
 
     if "poloidal_distance" in da.coords and myg > 0:
         # Special handling for core regions to deal with branch cut
@@ -1675,15 +1675,15 @@ def _concat_upper_guards(da, da_global, mxg, myg):
             )
             new_ycoord = new_ycoord + offset
 
-        # can't use commented out version, uncommented one works around xarray bug
-        # removing attrs
-        # https://github.com/pydata/xarray/issues/4415
-        # https://github.com/pydata/xarray/issues/4393
-        # da_upper = da_upper.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
-        da_upper[xcoord].data[...] = new_xcoord.data
-        da_upper = da_upper.reset_index(xcoord).set_xindex(xcoord)
-        da_upper[ycoord].data[...] = new_ycoord.data
-        da_upper = da_upper.reset_index(ycoord).set_xindex(ycoord)
+        # Preserve attributes when updating coordinates
+        # xarray coordinates are read-only by design, must use assign_coords()
+        xcoord_attrs = da_upper[xcoord].attrs.copy()
+        ycoord_attrs = da_upper[ycoord].attrs.copy()
+
+        # Assign new coordinate values
+        new_xcoord.attrs = xcoord_attrs
+        new_ycoord.attrs = ycoord_attrs
+        da_upper = da_upper.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord})
 
     if "poloidal_distance" in da.coords and myg > 0:
         # Special handling for core regions to deal with branch cut

xbout/tests/test_boutdataset.py

@@ -20,7 +20,6 @@
 from xbout.utils import _set_attrs_on_all_vars, _1d_coord_from_spacing
 from xbout.tests.utils_for_tests import set_geometry_from_input_file
 
-
 EXAMPLE_OPTIONS_FILE_PATH = "./xbout/tests/data/options/BOUT.inp"
 
 

xbout/tests/test_region.py

@@ -6,7 +6,6 @@
 
 from xbout import open_boutdataset
 
-
 params_guards = "guards"
 params_guards_values = [
     pytest.param({"x": 0, "y": 0}, marks=pytest.mark.long),

xbout/tests/utils_for_tests.py

@@ -6,7 +6,6 @@
 from pathlib import Path
 from xarray import DataArray, Dataset
 
-
 # Note, MYPE, PE_XIND and PE_YIND not included, since they are different for each
 # processor and so are dropped when loading datasets.
 METADATA_VARS = [