Merge pull request #307 from boutproject/fix-gridfile-loading

Fix loading of grid files

Author: bendudson

examples/plot_grid.py

@@ -13,6 +13,6 @@
 grid["psi_poloidal"].bout.pcolormesh()
 grid["psi_poloidal"].bout.pcolormesh(shading="gouraud")
 
-grid["psi_poloidal"].bout.regions()
+grid["psi_poloidal"].bout.plot_regions()
 
 plt.show()

xbout/geometries.py

@@ -401,7 +401,7 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
     ds = _maybe_rename_dimension(ds, "y", coordinates["y"])
 
     # TODO automatically make this coordinate 1D in simplified cases?
-    ds = ds.rename(psixy=coordinates["x"])
+    ds[coordinates["x"]] = ds["psixy"]
     ds = ds.set_coords(coordinates["x"])
     ds[coordinates["x"]].attrs["units"] = "Wb"
 

xbout/load.py

@@ -64,6 +64,7 @@ def open_boutdataset(
     run_name=None,
     info=True,
     is_restart=None,
+    is_mms_dump=False,
     **kwargs,
 ):
     """Load a dataset from a set of BOUT output files, including the
@@ -179,12 +180,12 @@ def open_boutdataset(
         chunks = {}
 
     input_type = _check_dataset_type(datapath)
-
     if is_restart is None:
         is_restart = input_type == "restart"
     elif is_restart is True:
         input_type = "restart"
-
+    if is_mms_dump:
+        input_type = "dump"
     if "reload" in input_type:
         if input_type == "reload":
             if isinstance(datapath, Path):
@@ -358,7 +359,12 @@ def attrs_remove_section(obj, section):
                 ds.metadata[v] = grid[v].values
 
     # Update coordinates to match particular geometry of grid
-    ds = geometries.apply_geometry(ds, geometry, grid=grid)
+    if input_type == "grid":
+        # Specify coordinate names to avoid name clash
+        coordinates = {"x": "psi_poloidal", "y": "y", "z": "zeta"}
+    else:
+        coordinates = None
+    ds = geometries.apply_geometry(ds, geometry, grid=grid, coordinates=coordinates)
 
     if remove_yboundaries:
         ds = ds.bout.remove_yboundaries()
@@ -616,7 +622,6 @@ def _auto_open_mfboutdataset(
         nxpe, nype, mxg, myg, mxsub, mysub, is_squashed_doublenull = _read_splitting(
             filepaths[0], info, keep_yboundaries
         )
-
         if is_squashed_doublenull:
             # Need to remove y-boundaries after loading: (i) in case
             # we are loading a squashed data-set, in which case we

xbout/tests/test_boutdataset.py

@@ -1136,6 +1136,7 @@ def test_interpolate_parallel_all_variables_arg(self, bout_xyt_example_files):
                 "G3",
                 "J",
                 "Bxy",
+                "psixy",
             )
         )