From 756d73ebb35c60964421df6658a84c1eb570b73b Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 14 Jan 2021 20:11:37 +0000 Subject: [PATCH] Add default values for bout_tdim, etc., so that they always exist --- xbout/geometries.py | 16 +++++++++++----- xbout/utils.py | 7 +++++++ 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index 52745ba7..53a18ad6 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -209,12 +209,18 @@ def _set_default_toroidal_coordinates(coordinates, ds): coordinates = {} # Replace any values that have not been passed in with defaults - coordinates["t"] = coordinates.get("t", ds.metadata.get("bout_tdim", "t")) - coordinates["x"] = coordinates.get( - "x", ds.metadata.get("bout_xdim", "psi_poloidal") + coordinates["t"] = coordinates.get("t", ds.metadata["bout_tdim"]) + + default_x = ( + ds.metadata["bout_xdim"] if ds.metadata["bout_xdim"] != "x" else "psi_poloidal" ) - coordinates["y"] = coordinates.get("y", ds.metadata.get("bout_ydim", "theta")) - coordinates["z"] = coordinates.get("z", ds.metadata.get("bout_zdim", "zeta")) + coordinates["x"] = coordinates.get("x", default_x) + + default_y = ds.metadata["bout_ydim"] if ds.metadata["bout_ydim"] != "y" else "theta" + coordinates["y"] = coordinates.get("y", default_y) + + default_z = ds.metadata["bout_zdim"] if ds.metadata["bout_zdim"] != "z" else "zeta" + coordinates["z"] = coordinates.get("z", default_z) return coordinates diff --git a/xbout/utils.py b/xbout/utils.py index e2367ebf..cf0938e8 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -65,6 +65,13 @@ def _separate_metadata(ds): metadata_vals = [ds[var].values.item() for var in scalar_vars] metadata = dict(zip(scalar_vars, metadata_vals)) + # Add default values for dimensions to metadata. These may be modified later by + # apply_geometry() + metadata["bout_tdim"] = "t" + metadata["bout_xdim"] = "x" + metadata["bout_ydim"] = "y" + metadata["bout_zdim"] = "z" + return ds.drop_vars(scalar_vars), metadata