cleaner coordinate aliasing (#116)

Take advantage of wpbonelli/xattree@54f07b0. This is a cleaner way to get lay/row/col/node indexing than a custom index, which is still pretty cutting-edge. We will probably still need a custom index to handle CRS etc but maybe best to wait for xarray to harden the api before developing one.

Author: wpbonelli

docs/examples/quickstart.py

@@ -29,6 +29,7 @@
 # check CHD
 assert chd.data["head"][0, 0].item() == 1.0
 assert chd.data.head.sel(per=0)[99].item() == 0.0
+assert chd.data.head.sel(per=0, node=99).item() == 0.0
 assert np.allclose(chd.data.head[:, 1:99], np.full(98, 1e30))
 
 # check DIS

flopy4/mf6/gwf/dis.py

@@ -4,11 +4,10 @@
 from xattree import array, dim, field, xattree
 
 from flopy4.mf6.converters import convert_array
-from flopy4.mf6.indexes import grid_index
 from flopy4.mf6.package import Package
 
 
-@xattree(index=grid_index, index_scope="gwf")
+@xattree
 class Dis(Package):
     length_units: str = field(
         default=None,
@@ -22,26 +21,23 @@ 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,
+        coord="lay",
         scope="gwf",
         default=1,
         metadata={
             "block": "dimensions",
         },
     )
     ncol: int = dim(
-        coord=False,
+        coord="col",
         scope="gwf",
         default=2,
         metadata={
             "block": "dimensions",
         },
     )
     nrow: int = dim(
-        coord=False,
+        coord="row",
         scope="gwf",
         default=2,
         metadata={
@@ -79,7 +75,7 @@ class Dis(Package):
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     nnodes: int = dim(
-        # coord=False,
+        coord="node",
         scope="gwf",
         init=False,
     )

flopy4/mf6/indexes.py

@@ -44,15 +44,6 @@ 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) -> MetaIndex:
     return MetaIndex(

flopy4/mf6/tdis.py

@@ -7,11 +7,10 @@
 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(index=time_index, index_scope=ROOT)
+@xattree
 class Tdis(Package):
     @define
     class PeriodData:
@@ -20,7 +19,7 @@ class PeriodData:
         tsmult: float
 
     nper: int = dim(
-        coord=False,
+        coord="per",
         default=1,
         scope=ROOT,
         metadata={"block": "dimensions"},

flopy4/utils.py

@@ -4,6 +4,6 @@
 
 def to_path(value: Any) -> Optional[Path]:
     """
-    Convert a value to a Path if it has a value, otherwise return None.
+    Try to convert a value to a Path if it's not None, otherwise return None.
     """
     return Path(value) if value else None