feat(dfns): add model dependent variable info

autotest/dfns/test_migrate_v1_to_v2.py

@@ -1,4 +1,7 @@
+import pytest
+
 from modflow_devtools.dfn import schema as v1
+from modflow_devtools.dfns.migrate_v1_to_v2 import _DEPENDENT_VARS, _OC_RTYPE_VALID
 from modflow_devtools.dfns.migrate_v1_to_v2 import v1_to_v2 as v1_to_v2
 from modflow_devtools.dfns.schema import (
     Double,
@@ -10,6 +13,7 @@
     Record,
     Simulation,
     String,
+    Union,
 )
 
 
@@ -425,3 +429,71 @@ def test_components():
     result = v1_to_v2(dfn)
     assert isinstance(result, Package)
     assert result.subtype == "advanced"
+
+
+def _oc_dfn(prefix: str) -> v1.Dfn:
+    """Minimal OC-like v1 DFN for testing rtype.valid migration."""
+    return _v1_dfn(
+        name=f"{prefix}-oc",
+        parent=f"{prefix}-nam",
+        blocks={
+            "period": {
+                "saverecord": _v1_field(
+                    name="saverecord",
+                    type="record save rtype ocsetting",
+                    block="period",
+                    optional=True,
+                ),
+                "save": _v1_field(name="save", type="keyword", block="period", in_record=True),
+                "printrecord": _v1_field(
+                    name="printrecord",
+                    type="record print rtype ocsetting",
+                    block="period",
+                    optional=True,
+                ),
+                "print": _v1_field(name="print", type="keyword", block="period", in_record=True),
+                "rtype": _v1_field(
+                    name="rtype",
+                    type="string",
+                    block="period",
+                    in_record=True,
+                    tagged=False,
+                ),
+                "ocsetting": _v1_field(
+                    name="ocsetting",
+                    type="keystring all",
+                    block="period",
+                    in_record=True,
+                ),
+                "all": _v1_field(name="all", type="keyword", block="period", in_record=True),
+            }
+        },
+    )
+
+
+@pytest.mark.parametrize("prefix,expected", list(_OC_RTYPE_VALID.items()))
+def test_oc_rtype_valid(prefix, expected):
+    component = v1_to_v2(_oc_dfn(prefix))
+    assert isinstance(component, Package)
+    period = component.blocks["period"]
+    output = period.fields["output"]
+    assert isinstance(output, List)
+    assert isinstance(output.item, Union)
+    for arm in output.item.arms.values():
+        assert isinstance(arm, Record)
+        rtype = arm.fields["rtype"]
+        assert isinstance(rtype, String)
+        assert rtype.valid == expected
+
+
+@pytest.mark.parametrize("prefix,expected", list(_DEPENDENT_VARS.items()))
+def test_model_dependent_variable(prefix, expected):
+    result = v1_to_v2(_v1_dfn(name=f"{prefix}-nam"))
+    assert isinstance(result, Model)
+    assert result.dependent_variable == expected
+
+
+def test_model_no_dependent_variable():
+    result = v1_to_v2(_v1_dfn(name="prt-nam"))
+    assert isinstance(result, Model)
+    assert result.dependent_variable is None

docs/md/dfnspec.md

@@ -16,6 +16,7 @@ This document describes the MODFLOW 6 component definition (DFN) system. This sy
     - [Model](#model)
       - [Type-specific attributes](#type-specific-attributes)
         - [`solution`](#solution)
+        - [`dependent_variable`](#dependent_variable)
     - [Package](#package)
       - [Type-specific attributes](#type-specific-attributes-1)
         - [`multi`](#multi)
@@ -160,6 +161,10 @@ A model represents a hydrologic process. Models are managed and solved by the si
 
 `"ims" | "ems" | "sln-ims" | "sln-ems" | null (default: null)`. MF6 supports different solution schemes: implicit solutions (solve systems of coupled equations iteratively) and explicit solutions (used when closed-form solutions are available). A model declares which solution type it requires with the optional `solution` attribute. Solution packages do not redundantly declare which model types they support; compatibility is determined entirely from the model side.
 
+###### `dependent_variable`
+
+`string | null (default: null)`. The dependent variable this model type computes, e.g. `"head"` for GWF. The model's OC package `rtype` string field should also specify this variable's name as a `valid` value.
+
 #### Package
 
 A package is any component that is not a simulation or a model.

modflow_devtools/dfns/__init__.py

@@ -21,6 +21,7 @@
     Model,
     Package,
     Record,
+    Scalar,
     Simulation,
     String,
     Union,
@@ -57,6 +58,7 @@
     "Package",
     "Record",
     "RemoteDfnRegistry",
+    "Scalar",
     "Simulation",
     "String",
     "Union",

modflow_devtools/dfns/migrate_v1_to_v2.py

@@ -10,6 +10,26 @@
 _IDENT_RE = re.compile(r"^[A-Za-z_]\w*$")
 _LOOKUP_RE = re.compile(r"^(\w+)\.(\w+)\((\w+)\)$")
 
+_DEPENDENT_VARS: dict[str, str] = {
+    "gwf": "head",
+    "gwt": "concentration",
+    "gwe": "temperature",
+    "chf": "stage",
+    "olf": "stage",
+    "swf": "stage",
+    # prt: particle tracking; no scalar dependent variable
+}
+
+_OC_RTYPE_VALID: dict[str, list[str]] = {
+    "gwf": ["HEAD", "BUDGET"],
+    "gwt": ["CONCENTRATION", "BUDGET"],
+    "gwe": ["TEMPERATURE", "BUDGET"],
+    "chf": ["STAGE", "BUDGET"],
+    "olf": ["STAGE", "BUDGET"],
+    "swf": ["STAGE", "BUDGET"],
+    "prt": ["BUDGET"],
+}
+
 
 def _scope_for(
     parent: "str | list[str] | None",
@@ -348,6 +368,61 @@ def _collapse_sto_keywords(
     return result
 
 
+def _patch_oc_rtype(
+    name: str,
+    blocks: dict[str, v2.Block],
+) -> dict[str, v2.Block]:
+    """Set valid values on rtype string fields in OC packages."""
+    if not name.endswith("-oc"):
+        return blocks
+    prefix = name.split("-")[0]
+    valid = _OC_RTYPE_VALID.get(prefix)
+    if not valid:
+        return blocks
+
+    def _patch(fields: dict) -> tuple[dict, bool]:
+        new_fields = {}
+        changed = False
+        for fname, field in fields.items():
+            if isinstance(field, v2.String) and fname == "rtype":
+                field = field.model_copy(update={"valid": valid})
+                changed = True
+            elif isinstance(field, v2.Record):
+                patched, c = _patch(field.fields)
+                if c:
+                    field = field.model_copy(update={"fields": patched})
+                    changed = True
+            elif isinstance(field, v2.Union):
+                patched, c = _patch(field.arms)
+                if c:
+                    field = field.model_copy(update={"arms": patched})
+                    changed = True
+            elif isinstance(field, v2.List):
+                item = field.item
+                if isinstance(item, v2.Record):
+                    patched, c = _patch(item.fields)
+                    if c:
+                        field = field.model_copy(
+                            update={"item": item.model_copy(update={"fields": patched})}
+                        )
+                        changed = True
+                elif isinstance(item, v2.Union):
+                    patched, c = _patch(item.arms)
+                    if c:
+                        field = field.model_copy(
+                            update={"item": item.model_copy(update={"arms": patched})}
+                        )
+                        changed = True
+            new_fields[fname] = field
+        return new_fields, changed
+
+    result = {}
+    for bname, block in blocks.items():
+        new_fields, changed = _patch(block.fields)
+        result[bname] = block.model_copy(update={"fields": new_fields}) if changed else block
+    return result
+
+
 def v1_to_v2(dfn: v1.Dfn) -> v2.Component:
     """Map a component definition from the v1 schema to v2."""
 
@@ -754,6 +829,7 @@ def _record_fields() -> dict:
     blocks = _fill_period_list_shapes(blocks, explicit_dims)
     blocks = _wrap_oc_period_records(blocks)
     blocks = _collapse_sto_keywords(blocks)
+    blocks = _patch_oc_rtype(name, blocks)
     dims = {**explicit_dims, **array_dims} or None
 
     d: dict[str, Any] = {
@@ -766,7 +842,8 @@ def _record_fields() -> dict:
     if name == "sim-nam":
         return v2.Simulation(**d)
     if name.endswith("-nam"):
-        return v2.Model(**d)
+        prefix = name.split("-")[0]
+        return v2.Model(**d, dependent_variable=_DEPENDENT_VARS.get(prefix))
 
     subtype: Literal["solution", "exchange", "stress", "advanced", "utility"] | None = None
     if name.startswith("sln-"):

modflow_devtools/dfns/schema.py

@@ -413,6 +413,7 @@ class Simulation(ComponentBase):
 class Model(ComponentBase):
     type: Literal["model"] = "model"
     solution: Literal["ims", "ems", "sln-ims", "sln-ems"] | None = None
+    dependent_variable: str | None = None
 
 
 class Package(ComponentBase):