update indexing in quickstart (#115)

followup on #114, fix some things and update the quickstart to demo the new indexing syntax this allows

Author: wpbonelli

docs/examples/quickstart.py

@@ -28,13 +28,12 @@
 
 # check CHD
 assert chd.data["head"][0, 0].item() == 1.0
-assert chd.data["head"][0, 99].item() == 0.0
-assert np.allclose(chd.data["head"][:, 1:99].data.todense(), np.full(98, 1e30))
+assert chd.data.head.sel(per=0)[99].item() == 0.0
+assert np.allclose(chd.data.head[:, 1:99], np.full(98, 1e30))
 
-# TODO: xarray index aliasing nlay/ncol/nrow to k/i/j?
-# assert chd.data["head"].loc(dict(k=0, i=0, j=0)) == 1.
-# assert chd.data["head"].loc(dict(k=0, i=9, j=9)) == 0.
+# check DIS
+assert dis.data.botm.sel(lay=0, col=0, row=0) == 0.0
 
 # check OC
-assert oc.data["save_head"][0].item() == "all"
-assert oc.data["save_budget"][0].item() == "all"
+assert oc.data["save_head"][0] == "all"
+assert oc.data.save_head.sel(per=0) == "all"

flopy4/mf6/gwf/dis.py

@@ -22,6 +22,9 @@ class Dis(Package):
         default=False, metadata={"block": "options"}
     )
     nlay: int = dim(
+        # disable the otherwise automatic coordinate variable
+        # because we're going to create another one for this
+        # dimension with a different name via a custom index
         coord=False,
         scope="gwf",
         default=1,
@@ -46,36 +49,40 @@ class Dis(Package):
         },
     )
     delr: NDArray[np.floating] = array(
-        dims=("ncol",),
         default=1.0,
+        dims=("ncol",),
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     delc: NDArray[np.floating] = array(
-        dims=("nrow",),
         default=1.0,
+        dims=("nrow",),
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     top: NDArray[np.floating] = array(
-        dims=("ncol", "nrow"),
         default=1.0,
+        dims=("ncol", "nrow"),
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     botm: NDArray[np.floating] = array(
-        dims=("ncol", "nrow", "nlay"),
         default=0.0,
+        dims=("ncol", "nrow", "nlay"),
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     idomain: NDArray[np.integer] = array(
-        dims=("ncol", "nrow", "nlay"),
         default=1,
+        dims=("ncol", "nrow", "nlay"),
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
-    nnodes: int = dim(scope="gwf", init=False)
+    nnodes: int = dim(
+        # coord=False,
+        scope="gwf",
+        init=False,
+    )
 
     def __attrs_post_init__(self):
         self.nnodes = self.ncol * self.nrow * self.nlay

flopy4/mf6/gwf/oc.py

@@ -53,27 +53,31 @@ class Period:
     format: Optional[Format] = field(
         default=None, init=False, metadata={"block": "options"}
     )
-    save_head: Optional[NDArray[np.str_ | np.integer]] = array(
-        dims=("nper",),
+    save_head: Optional[NDArray[np.object_]] = array(
+        Steps,
         default="all",
+        dims=("nper",),
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
-    save_budget: Optional[NDArray[np.str_ | np.integer]] = array(
-        dims=("nper",),
+    save_budget: Optional[NDArray[np.object_]] = array(
+        Steps,
         default="all",
+        dims=("nper",),
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
-    print_head: Optional[NDArray[np.str_ | np.integer]] = array(
-        dims=("nper",),
+    print_head: Optional[NDArray[np.object_]] = array(
+        Steps,
         default="all",
+        dims=("nper",),
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
-    print_budget: Optional[NDArray[np.str_ | np.integer]] = array(
-        dims=("nper",),
+    print_budget: Optional[NDArray[np.object_]] = array(
+        Steps,
         default="all",
+        dims=("nper",),
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )

flopy4/mf6/indexes.py

@@ -5,18 +5,23 @@
 
 
 def alias(dataset: xr.Dataset, old_name: str, new_name: str) -> PandasIndex:
-    """Alias a dimension coordinate to a coordinate with a different name."""
+    """
+    Alias a dimension coordinate to a coordinate with a different name.
+    Suggested by Benoit Bovy https://github.com/pydata/xarray/pull/10076#issuecomment-2809041994.
+    """
     try:
         size = dataset.sizes[old_name]
     except KeyError:
-        try:
-            size = dataset.dims[old_name]
-        except KeyError:
-            size = dataset.attrs[old_name]
+        size = dataset.attrs[old_name]
     return PandasIndex(pd.RangeIndex(size, name=new_name), dim=old_name)
 
 
-class GridIndex(Index):
+class MetaIndex(Index):
+    """
+    Combine multiple indexes into a single index.
+    Adapted from https://docs.xarray.dev/en/stable/internals/how-to-create-custom-index.html#meta-indexes.
+    """
+
     def __init__(self, indices):
         self._indices = indices
 
@@ -39,13 +44,29 @@ def sel(self, labels):
                 results.append(index.sel({k: labels[k]}))
         return merge_sel_results(results)
 
+    def to_pandas_index(self) -> pd.Index:
+        # from https://github.com/corteva/rioxarray/pull/846/files#diff-917105823f61e63ef4afde8bed408a6c249e375690e56bc800406676f02551d8R418
+        if len(self._indices) == 1:
+            index = next(iter(self._indices.values()))
+            if isinstance(index, PandasIndex):
+                return index.to_pandas_index()
+
+        raise ValueError("Cannot convert MetaIndex to pandas.Index")
+
 
-def grid_index(dataset: xr.Dataset) -> GridIndex:
-    return GridIndex(
+def grid_index(dataset: xr.Dataset) -> MetaIndex:
+    return MetaIndex(
         {
-            # k collides with npf.k so use "lay"
+            # TODO add 'per' (stress period)
             "lay": alias(dataset, "nlay", "lay"),
             "col": alias(dataset, "ncol", "col"),
             "row": alias(dataset, "nrow", "row"),
+            # "node": alias(dataset, "nnodes", "node"),
+            # TODO: adding node breaks the other three.
+            # and just having node by itself works. why?
         }
     )
+
+
+def time_index(dataset: xr.Dataset) -> PandasIndex:
+    return alias(dataset, "nper", "per")

flopy4/mf6/tdis.py

@@ -7,10 +7,11 @@
 from xattree import ROOT, array, dim, field, xattree
 
 from flopy4.mf6.converters import convert_array
+from flopy4.mf6.indexes import time_index
 from flopy4.mf6.package import Package
 
 
-@xattree
+@xattree(index=time_index, index_scope=ROOT)
 class Tdis(Package):
     @define
     class PeriodData:
@@ -19,6 +20,7 @@ class PeriodData:
         tsmult: float
 
     nper: int = dim(
+        coord=False,
         default=1,
         scope=ROOT,
         metadata={"block": "dimensions"},