diff --git a/xbout/geometries.py b/xbout/geometries.py index fa8a67b6..cf3a181c 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -10,6 +10,7 @@ _set_attrs_on_all_vars, _set_as_coord, _1d_coord_from_spacing, + _maybe_rename_dimension, ) REGISTERED_GEOMETRIES = {} @@ -386,12 +387,12 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): ], ) - if "t" in ds.dims: + if coordinates["t"] != "t": # Rename 't' if user requested it - ds = ds.rename(t=coordinates["t"]) + ds = _maybe_rename_dimension(ds, "t", coordinates["t"]) # Change names of dimensions to Orthogonal Toroidal ones - ds = ds.rename(y=coordinates["y"]) + ds = _maybe_rename_dimension(ds, "y", coordinates["y"]) # TODO automatically make this coordinate 1D in simplified cases? ds = ds.rename(psixy=coordinates["x"]) @@ -407,7 +408,7 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): # If full data (not just grid file) then toroidal dim will be present if "z" in ds.dims: - ds = ds.rename(z=coordinates["z"]) + ds = _maybe_rename_dimension(ds, "z", coordinates["z"]) # Record which dimension 'z' was renamed to. ds.metadata["bout_zdim"] = coordinates["z"] @@ -482,7 +483,7 @@ def add_s_alpha_geometry_coords(ds, *, coordinates=None, grid=None): ds["r"] = ds["hthe"].isel({ycoord: 0}).squeeze(drop=True) ds["r"].attrs["units"] = "m" ds = ds.set_coords("r") - ds = ds.rename(x="r") + ds = ds.swap_dims(x="r") ds.metadata["bout_xdim"] = "r" if hthe_from_grid: diff --git a/xbout/utils.py b/xbout/utils.py index 9587fcb4..d6f24537 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -880,3 +880,14 @@ def _set_as_coord(ds, name): except ValueError: pass return ds + + +def _maybe_rename_dimension(ds, old_name, new_name): + if old_name in ds.dims and new_name != old_name: + # Rename dimension + ds = ds.swap_dims({old_name: new_name}) + if old_name in ds: + # Rename coordinate if it exists + ds = ds.rename({old_name: new_name}) + + return ds