diff --git a/examples/plot_grid.py b/examples/plot_grid.py index e8ea3b4a..1b504d2e 100644 --- a/examples/plot_grid.py +++ b/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() diff --git a/xbout/geometries.py b/xbout/geometries.py index 7223751b..dd1b6692 100644 --- a/xbout/geometries.py +++ b/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" diff --git a/xbout/load.py b/xbout/load.py index 683c96d8..5c3d8a61 100644 --- a/xbout/load.py +++ b/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 diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py index 6d20cd30..4ef8c6fe 100644 --- a/xbout/tests/test_boutdataset.py +++ b/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", ) )