Repository: boutproject/xBOUT

Issue:
Plotting grids using `open_boutdataset` seems broken
There is a suggested example for plotting the fields of a `grid.nc` file https://github.com/boutproject/xBOUT/blob/master/examples/plot_grid.py. I have tried to use this code to plot the output from a simple Hyponotoad-generated grid. The result is that I received error messages which suggest that the advertised functionality is broken. 

To generate the grid I use (thanks to J. Omotani)
```
from hypnotoad.cases.circular import CircularEquilibrium
from hypnotoad.core.mesh import BoutMesh
import numpy as np

R0 = 2.3
B0 = 3.2
r_inner = 0.1
r_outer = 1.4
q = 4.1
settings = {
    "R0": R0,
    "B0": B0,
    "poloidal_spacing_method": "linear",
    "q_coefficients": [q],
    "r_inner": r_inner,
    "r_outer": r_outer,
    "refine_methods": "line",
    "refine_width": 1.0e-3,
}

equilib = CircularEquilibrium(settings, nonorthogonal_settings=settings)

mesh = BoutMesh(equilib, settings)
mesh.geometry()
mesh.writeGridfile("grid.nc")
```
I then run the following script to try to plot
```
from matplotlib import pyplot as plt
from xbout import open_boutdataset
#from xbout.load import _open_grid

gridfilepath = "grid.nc"
grid = open_boutdataset(gridfilepath, geometry="toroidal")
#grid = _open_grid(gridfilepath, chunks={},keep_xboundaries=True,keep_yboundaries=True)

grid["psi_poloidal"].bout.contourf()
grid["psi_poloidal"].bout.contour()
grid["psi_poloidal"].bout.pcolormesh()
grid["psi_poloidal"].bout.pcolormesh(shading="gouraud")
grid["psi_poloidal"].bout.regions()
plt.show()
```
With this script, I get the following error
```
$ python3 plot_test.py
NXPE not found, setting to 1
NYPE not found, setting to 1
MXG not found, setting to 2
MYG not found, setting to 0
MXSUB not found, setting to 2
MYSUB not found, setting to 8
Traceback (most recent call last):
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 1317, in _construct_dataarray
    variable = self._variables[name]
               ~~~~~~~~~~~~~~~^^^^^^
KeyError: 't_array'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "*/plot_test.py", line 11, in <module>
    grid = open_boutdataset(gridfilepath, geometry="toroidal")
  File "*/lib/python3.13/site-packages/xbout/load.py", line 279, in open_boutdataset
    ds, remove_yboundaries = _auto_open_mfboutdataset(
                             ~~~~~~~~~~~~~~~~~~~~~~~~^
        datapath=datapath,
        ^^^^^^^^^^^^^^^^^^
    ...<4 lines>...
        **kwargs,
        ^^^^^^^^^
    )
    ^
  File "*/lib/python3.13/site-packages/xbout/load.py", line 736, in _auto_open_mfboutdataset
    _, unique_indices = unique(ds["t_array"], return_index=True)
                               ~~^^^^^^^^^^^
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 1410, in __getitem__
    return self._construct_dataarray(key)
           ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 1319, in _construct_dataarray
    _, name, variable = _get_virtual_variable(self._variables, name, self.dims)
                        ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 175, in _get_virtual_variable
    raise KeyError(key)
KeyError: 't_array'
```
Trying to use `_open_grid` directly also failed. Fixing this would be very helpful for plotting Hypnotoad output.

Discussion:
@johnomotani: This problem was mostly caused by a change to hypnotoad which had resulted in a 't' dimension being created in grid files, which confused the function that detects whether the file being opened is 'dump', 'restart' or 'grid'. Also a name-conflict error.

Solution patch:
diff --git a/examples/plot_grid.py b/examples/plot_grid.py
index e8ea3b4..1b504d2 100644
--- a/examples/plot_grid.py
+++ b/examples/plot_grid.py
@@ -13,6 +13,6 @@ grid["psi_poloidal"].bout.contour()
 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 7223751..24bbc93 100644
--- a/xbout/geometries.py
+++ b/xbout/geometries.py
@@ -214,7 +214,7 @@ def apply_geometry(ds, geometry_name, *, coordinates=None, grid=None):
                 dz = updated_ds["dz"]
 
             z0 = 2 * np.pi * updated_ds.metadata["ZMIN"]
-            z1 = z0 + nz * dz.data[()]
+            z1 = z0 + nz * dz.data.flatten()[0]
             if not np.all(
                 np.isclose(
                     z1,
@@ -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 683c96d..1d9d4d9 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 @@ but we did load {grididfile}."""
                 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()
@@ -372,9 +378,9 @@ but we did load {grididfile}."""
     # by xBOUT, not by BOUT++
     ds.bout.fine_interpolation_factor = 8
 
-    if ("dump" in input_type or "restart" in input_type) and ds.metadata[
-        "BOUT_VERSION"
-    ] < 4.0:
+    if ("dump" in input_type or "restart" in input_type) and (
+        "BOUT_VERSION" in ds.metadata and ds.metadata["BOUT_VERSION"] < 4.0
+    ):
         # Add workarounds for missing information or different
         # conventions in data saved by BOUT++ v3.x.
         for v in ds:
@@ -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


Test patch:
diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py
index 6d20cd3..4ef8c6f 100644
--- a/xbout/tests/test_boutdataset.py
+++ b/xbout/tests/test_boutdataset.py
@@ -1136,6 +1136,7 @@ class TestBoutDatasetMethods:
                 "G3",
                 "J",
                 "Bxy",
+                "psixy",
             )
         )