diff --git a/xbout/calc/turbulence.py b/xbout/calc/turbulence.py index 926b0015..0ed0c334 100644 --- a/xbout/calc/turbulence.py +++ b/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? diff --git a/xbout/conftest.py b/xbout/conftest.py index d30f31ab..137479fd 100644 --- a/xbout/conftest.py +++ b/xbout/conftest.py @@ -1,6 +1,5 @@ import pytest - configure_is_imported = False diff --git a/xbout/geometries.py b/xbout/geometries.py index 24bbc937..35c541ea 100644 --- a/xbout/geometries.py +++ b/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 diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index e223983f..19582c79 100644 --- a/xbout/plotting/animate.py +++ b/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" diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 0f9dc51f..9c54d183 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -18,7 +18,6 @@ plot_targets, ) - if ( "pcolor.shading" in matplotlib.rcParams and matplotlib.rcParams["pcolor.shading"] == "flat" diff --git a/xbout/region.py b/xbout/region.py index 4db156a8..fcd63bce 100644 --- a/xbout/region.py +++ b/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 diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py index 4ef8c6fe..a03c4ed9 100644 --- a/xbout/tests/test_boutdataset.py +++ b/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" diff --git a/xbout/tests/test_region.py b/xbout/tests/test_region.py index 8ca6dc58..f3199d93 100644 --- a/xbout/tests/test_region.py +++ b/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), diff --git a/xbout/tests/utils_for_tests.py b/xbout/tests/utils_for_tests.py index 01f01159..6b2d1b51 100644 --- a/xbout/tests/utils_for_tests.py +++ b/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 = [