basic index working

flopy4/mf6/gwf/dis.py

@@ -4,10 +4,11 @@
 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
+@xattree(index=grid_index, index_scope="gwf")
 class Dis(Package):
     length_units: str = field(
         default=None,
@@ -21,20 +22,23 @@ class Dis(Package):
         default=False, metadata={"block": "options"}
     )
     nlay: int = dim(
+        coord=False,
         scope="gwf",
         default=1,
         metadata={
             "block": "dimensions",
         },
     )
     ncol: int = dim(
+        coord=False,
         scope="gwf",
         default=2,
         metadata={
             "block": "dimensions",
         },
     )
     nrow: int = dim(
+        coord=False,
         scope="gwf",
         default=2,
         metadata={

flopy4/mf6/indexes.py

@@ -0,0 +1,51 @@
+import pandas as pd
+import xarray as xr
+from xarray.core.indexes import Index, PandasIndex
+from xarray.core.indexing import merge_sel_results
+
+
+def alias(dataset: xr.Dataset, old_name: str, new_name: str) -> PandasIndex:
+    """Alias a dimension coordinate to a coordinate with a different name."""
+    try:
+        size = dataset.sizes[old_name]
+    except KeyError:
+        try:
+            size = dataset.dims[old_name]
+        except KeyError:
+            size = dataset.attrs[old_name]
+    return PandasIndex(pd.RangeIndex(size, name=new_name), dim=old_name)
+
+
+class GridIndex(Index):
+    def __init__(self, indices):
+        self._indices = indices
+
+    @classmethod
+    def from_variables(cls, variables):
+        return {
+            k: PandasIndex.from_variables({k: v}) for k, v in variables.items()
+        }
+
+    def create_variables(self, variables=None):
+        idx_vars = {}
+        for index in self._indices.values():
+            idx_vars.update(index.create_variables(variables))
+        return idx_vars
+
+    def sel(self, labels):
+        results = []
+        for k, index in self._indices.items():
+            if k in labels:
+                results.append(index.sel({k: labels[k]}))
+        return merge_sel_results(results)
+
+
+def grid_index(dataset: xr.Dataset) -> GridIndex:
+    return GridIndex(
+        {
+            # k collides with npf.k so use "lay"
+            "lay": alias(dataset, "nlay", "lay"),
+            "col": alias(dataset, "ncol", "col"),
+            "row": alias(dataset, "nrow", "row"),
+        }
+    )