add from_tokens interface for record types

Author: mjreno

flopy4/mf6/codec/writer/filters.py

@@ -243,12 +243,16 @@ def dataset2list(value: xr.Dataset):
                 yield (*name.split("_"), val)  # type: ignore
 
         else:
-            vals = []
-            for name in value.data_vars.keys():
-                val = value[name]
-                val = val.item() if val.shape == () else val
-                vals.append(val)
-            yield tuple(vals)
+            row: list[Any] = []
+            for name, da in value.data_vars.items():
+                val = da.item() if da.shape == () else da
+                if kw := da.attrs.get("row_keyword", False):
+                    if val:
+                        row.append(kw if isinstance(kw, str) else str(name).upper())
+                else:
+                    row.extend(da.attrs.get("prefix", ()))
+                    row.append(val)
+            yield tuple(row)
         return
 
     combined_mask: Any = None
@@ -266,15 +270,20 @@ def dataset2list(value: xr.Dataset):
             for name in value.data_vars.keys():
                 val = value[name][tuple(idx[i] for idx in indices)]
                 val = val.item() if val.shape == () else val
-                yield (*name.split("_"), val)  # type: ignore
+                yield (*str(name).split("_"), val)  # type: ignore
         else:
-            vals = []
-            for name in value.data_vars.keys():
-                val = value[name][tuple(idx[i] for idx in indices)]
+            row2: list[Any] = []
+            for name, da in value.data_vars.items():
+                val = da[tuple(idx[i] for idx in indices)]
                 val = val.item() if val.shape == () else val
-                vals.append(val)
+                if kw := da.attrs.get("row_keyword", False):
+                    if val:
+                        row2.append(kw if isinstance(kw, str) else str(name).upper())
+                else:
+                    row2.extend(da.attrs.get("prefix", ()))
+                    row2.append(val)
             if has_spatial_dims:
                 cellid = tuple(idx[i] + 1 for idx in indices)
-                yield cellid + tuple(vals)
+                yield tuple(cellid) + tuple(row2)
             else:
-                yield tuple(vals)
+                yield tuple(row2)

flopy4/mf6/converter/egress/unstructure.py

@@ -248,6 +248,14 @@ def _unstructure_block_param(
                     for n, v in dat.items():
                         period_data[n] = v
             else:
+                arr_spec = xatspec.arrays.get(field_name)
+                field_meta = (arr_spec.metadata or {}) if arr_spec is not None else {}
+                if "prefix" in field_meta or "row_keyword" in field_meta:
+                    field_value = field_value.copy()
+                    if "prefix" in field_meta:
+                        field_value.attrs["prefix"] = field_meta["prefix"]
+                    if "row_keyword" in field_meta:
+                        field_value.attrs["row_keyword"] = field_meta["row_keyword"]
                 blocks[block_name][field_name] = field_value
         case _:
             blocks[block_name][field_name] = field_value
@@ -306,15 +314,17 @@ def _unstructure_array_component(value: Component) -> dict[str, Any]:
 
 # Block names that MF6 rejects if present but empty.
 # These blocks should only be written when they contain data.
-_SKIP_IF_EMPTY = frozenset({"dimensions", "tracktimes"})
+_SKIP_IF_EMPTY = frozenset({"dimensions", "fileinput", "tracktimes"})
 
 # Block names whose fields are list columns (one array per column, same dim)
 # rather than independent grid arrays.  Only these blocks are auto-combined
 # into an xr.Dataset for row-per-record output.  griddata-style blocks must
 # NOT be in this set — their fields are written individually with
 # INTERNAL/CONSTANT/NETCDF format.
 # "sources" is the SSM sources block (pname/srctype/auxname per-row tabular input).
-_LIST_BLOCK_NAMES = frozenset({"packagedata", "packages", "perioddata", "sources", "table"})
+_LIST_BLOCK_NAMES = frozenset(
+    {"packagedata", "packages", "perioddata", "sources", "fileinput", "table"}
+)
 
 
 def _unstructure_component(value: Component) -> dict[str, Any]:

flopy4/mf6/gwe/oc.py

@@ -11,14 +11,15 @@
 
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
+from flopy4.mf6.record import Record
 from flopy4.mf6.spec import array, field, path
 from flopy4.utils import to_path
 
 
 @xattree(kw_only=True)
 class Oc(Package):
     @attrs.define
-    class Temperatureprint:
+    class Temperatureprint(Record):
         _keyword: ClassVar[str] = "temperature"
         _extra_tokens: ClassVar[tuple[str, ...]] = ("PRINT_FORMAT",)
 

flopy4/mf6/gwe/ssm.py

@@ -20,7 +20,6 @@ class Ssm(Package):
     save_flows: bool = field(
         block="options", default=False, longname="save calculated flows to budget file"
     )
-    # TODO: fileinput — could not determine list column dimension
     nsources: Optional[int] = dim(block="__dim__", coord=False, default=None)
     pname: Optional[NDArray[np.object_]] = array(
         block="sources",
@@ -43,3 +42,26 @@ class Ssm(Package):
         converter=Converter(structure_array, takes_self=True, takes_field=True),
         longname="auxiliary variable name",
     )
+    nfileinput: Optional[int] = dim(block="__dim__", coord=False, default=None)
+    fi_pname: Optional[NDArray[np.object_]] = array(
+        block="fileinput",
+        dims=("nfileinput",),
+        default=None,
+        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        longname="package name",
+    )
+    fi_spc6_filename: Optional[NDArray[np.object_]] = array(
+        block="fileinput",
+        dims=("nfileinput",),
+        default=None,
+        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        longname="spc6 file name",
+        prefix=("SPC6", "FILEIN"),
+    )
+    fi_mixed: Optional[NDArray[np.bool_]] = array(
+        block="fileinput",
+        dims=("nfileinput",),
+        default=None,
+        longname="mixed keyword",
+        row_keyword="MIXED",
+    )

flopy4/mf6/gwf/npf.py

@@ -11,26 +11,27 @@
 
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
+from flopy4.mf6.record import Record
 from flopy4.mf6.spec import array, field, path
 from flopy4.utils import to_path
 
 
 @xattree(kw_only=True)
 class Npf(Package):
     @attrs.define
-    class Cvoptions:
+    class Cvoptions(Record):
         _keyword: ClassVar[str] = "variablecv"
         dewatered: Optional[bool] = attrs.field(default=None)
 
     @attrs.define
-    class Rewet:
+    class Rewet(Record):
         _keyword: ClassVar[str] = "rewet"
         wetfct: float = attrs.field(metadata={"tagged": True})
         iwetit: int = attrs.field(metadata={"tagged": True})
         ihdwet: int = attrs.field(metadata={"tagged": True})
 
     @attrs.define
-    class Xt3doptions:
+    class Xt3doptions(Record):
         _keyword: ClassVar[str] = "xt3d"
         rhs: Optional[bool] = attrs.field(default=None)
 

flopy4/mf6/gwf/oc.py

@@ -11,14 +11,15 @@
 
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
+from flopy4.mf6.record import Record
 from flopy4.mf6.spec import array, field, path
 from flopy4.utils import to_path
 
 
 @xattree(kw_only=True)
 class Oc(Package):
     @attrs.define
-    class Headprint:
+    class Headprint(Record):
         _keyword: ClassVar[str] = "head"
         _extra_tokens: ClassVar[tuple[str, ...]] = ("PRINT_FORMAT",)
         format_: str = attrs.field()

flopy4/mf6/gwt/ist.py

@@ -11,6 +11,7 @@
 
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
+from flopy4.mf6.record import Record
 from flopy4.mf6.spec import array, field, path
 from flopy4.utils import to_path
 
@@ -20,7 +21,7 @@ class Ist(Package):
     multi_package: ClassVar[bool] = True
 
     @attrs.define
-    class Cimprint:
+    class Cimprint(Record):
         _keyword: ClassVar[str] = "cim"
         _extra_tokens: ClassVar[tuple[str, ...]] = ("PRINT_FORMAT",)
 

flopy4/mf6/gwt/oc.py

@@ -11,14 +11,15 @@
 
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
+from flopy4.mf6.record import Record
 from flopy4.mf6.spec import array, field, path
 from flopy4.utils import to_path
 
 
 @xattree(kw_only=True)
 class Oc(Package):
     @attrs.define
-    class Concentrationprint:
+    class Concentrationprint(Record):
         _keyword: ClassVar[str] = "concentration"
         _extra_tokens: ClassVar[tuple[str, ...]] = ("PRINT_FORMAT",)
 

flopy4/mf6/gwt/ssm.py

@@ -20,7 +20,6 @@ class Ssm(Package):
     save_flows: bool = field(
         block="options", default=False, longname="save calculated flows to budget file"
     )
-    # TODO: fileinput — could not determine list column dimension
     nsources: Optional[int] = dim(block="__dim__", coord=False, default=None)
     pname: Optional[NDArray[np.object_]] = array(
         block="sources",
@@ -43,3 +42,26 @@ class Ssm(Package):
         converter=Converter(structure_array, takes_self=True, takes_field=True),
         longname="auxiliary variable name",
     )
+    nfileinput: Optional[int] = dim(block="__dim__", coord=False, default=None)
+    fi_pname: Optional[NDArray[np.object_]] = array(
+        block="fileinput",
+        dims=("nfileinput",),
+        default=None,
+        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        longname="package name",
+    )
+    fi_spc6_filename: Optional[NDArray[np.object_]] = array(
+        block="fileinput",
+        dims=("nfileinput",),
+        default=None,
+        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        longname="spc6 file name",
+        prefix=("SPC6", "FILEIN"),
+    )
+    fi_mixed: Optional[NDArray[np.bool_]] = array(
+        block="fileinput",
+        dims=("nfileinput",),
+        default=None,
+        longname="mixed keyword",
+        row_keyword="MIXED",
+    )

flopy4/mf6/ims.py

@@ -6,6 +6,7 @@
 import attrs
 from xattree import xattree
 
+from flopy4.mf6.record import Record
 from flopy4.mf6.solution import Solution
 from flopy4.mf6.spec import field, path
 from flopy4.utils import to_path
@@ -16,12 +17,12 @@ class Ims(Solution):
     slntype: ClassVar[str] = "ims"
 
     @attrs.define
-    class NoPtc:
+    class NoPtc(Record):
         _keyword: ClassVar[str] = "no_ptc"
         no_ptc_option: Optional[str] = attrs.field(default=None)
 
     @attrs.define
-    class Rclose:
+    class Rclose(Record):
         _keyword: ClassVar[str] = ""
         inner_rclose: float = attrs.field(metadata={"tagged": True})
         rclose_option: Optional[str] = attrs.field(default=None)