add drng

Author: mjreno

flopy4/mf6/constants.py

@@ -2,5 +2,5 @@
 
 MF6 = "mf6"
 FILL_DEFAULT = np.nan
-FILL_DNODATA = 3.0e30
+FILL_DNODATA = 3e30
 LENBOUNDNAME = 40

flopy4/mf6/converter/unstructure.py

@@ -259,7 +259,7 @@ def _unstructure_grid_component(value: Component) -> dict[str, Any]:
                 if not np.all(val == FILL_DNODATA):
                     if key not in blocks:
                         blocks[key] = {}
-                    blocks[f"period {kper + 1}"][arr_name] = val
+                    blocks[key][arr_name] = val
 
     return {name: block for name, block in blocks.items() if name != "period"}
 

flopy4/mf6/gwf/__init__.py

@@ -11,6 +11,7 @@
 from flopy4.mf6.gwf.chdg import Chdg
 from flopy4.mf6.gwf.dis import Dis
 from flopy4.mf6.gwf.drn import Drn
+from flopy4.mf6.gwf.drng import Drng
 from flopy4.mf6.gwf.ic import Ic
 from flopy4.mf6.gwf.npf import Npf
 from flopy4.mf6.gwf.oc import Oc
@@ -21,7 +22,7 @@
 from flopy4.mf6.utils import open_cbc, open_hds
 from flopy4.utils import to_path
 
-__all__ = ["Gwf", "Chd", "Chdg", "Dis", "Drn", "Ic", "Npf", "Oc", "Sto", "Wel"]
+__all__ = ["Gwf", "Chd", "Chdg", "Dis", "Drn", "Drng", "Ic", "Npf", "Oc", "Sto", "Wel"]
 
 
 def convert_grid(value):
@@ -82,6 +83,7 @@ def budget(self):
     chdg: list[Chdg] = field(block="packages")
     wel: list[Wel] = field(block="packages")
     drn: list[Drn] = field(block="packages")
+    drng: list[Drng] = field(block="packages")
     output: Output = attrs.field(
         default=attrs.Factory(lambda self: Gwf.Output(self), takes_self=True)
     )

flopy4/mf6/gwf/drng.py

@@ -0,0 +1,62 @@
+from pathlib import Path
+from typing import ClassVar, Optional
+
+import numpy as np
+
+# from attrs import Converter
+from numpy.typing import NDArray
+from xattree import xattree
+
+# from flopy4.mf6.converter import structure_array
+from flopy4.mf6.package import Package
+from flopy4.mf6.spec import array, field, path
+from flopy4.mf6.utils.grid import update_maxbound
+from flopy4.utils import to_path
+
+
+@xattree
+class Drng(Package):
+    multi_package: ClassVar[bool] = True
+    auxiliary: Optional[list[str]] = array(block="options", default=None)
+    auxmultname: Optional[str] = field(block="options", default=None)
+    print_input: bool = field(block="options", default=False)
+    print_flows: bool = field(block="options", default=False)
+    readarraygrid: bool = field(block="options", default=True)
+    save_flows: bool = field(block="options", default=False)
+    obs_filerecord: Optional[Path] = path(
+        block="options", default=None, converter=to_path, inout="fileout"
+    )
+    export_array_netcdf: bool = field(block="options", default=False)
+    mover: bool = field(block="options", default=False)
+    dev_cubic_scaling: bool = field(default=False, block="options")
+    maxbound: Optional[int] = field(block="dimensions", default=None, init=False)
+    elev: Optional[NDArray[np.float64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            "nodes",
+        ),
+        default=None,
+        # converter=Converter(structure_array, takes_self=True, takes_field=True),
+        on_setattr=update_maxbound,
+    )
+    cond: Optional[NDArray[np.float64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            "nodes",
+        ),
+        default=None,
+        # converter=Converter(structure_array, takes_self=True, takes_field=True),
+        on_setattr=update_maxbound,
+    )
+    aux: Optional[NDArray[np.float64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            "nodes",
+        ),
+        default=None,
+        # converter=Converter(structure_array, takes_self=True, takes_field=True),
+        on_setattr=update_maxbound,
+    )

test/test_mf6_codec.py

@@ -2,7 +2,10 @@
 
 from pprint import pprint
 
+import numpy as np
+
 from flopy4.mf6.codec import dumps, loads
+from flopy4.mf6.constants import FILL_DNODATA
 from flopy4.mf6.converter import COMPONENT_CONVERTER
 
 
@@ -236,6 +239,49 @@ def test_dumps_chd():
     pprint(loaded)
 
 
+def test_dumps_chdg():
+    from flopy4.mf6.gwf import Chdg, Dis, Gwf
+
+    nlay = 1
+    nrow = 10
+    ncol = 10
+
+    dis = Dis(nlay=nlay, nrow=nrow, ncol=ncol)
+    gwf = Gwf(dis=dis)
+
+    head = np.full((nlay, nrow, ncol), FILL_DNODATA, dtype=float)
+    head[0, 0, 0] = 1.0
+    head[0, 9, 9] = 0.0
+    chd = Chdg(
+        parent=gwf,
+        head=np.expand_dims(head.ravel(), axis=0),
+        save_flows=True,
+        print_input=True,
+        dims={"nper": 1},
+    )
+
+    dumped = dumps(COMPONENT_CONVERTER.unstructure(chd))
+    print("CHD dump:")
+    print(dumped)
+
+    assert "BEGIN PERIOD 1" in dumped
+    assert "END PERIOD 1" in dumped
+
+    period_section = dumped.split("BEGIN PERIOD 1")[1].split("END PERIOD 1")[0].strip()
+    lines = [line.strip() for line in period_section.split("\n") if line.strip()]
+
+    assert len(lines) == 12
+    assert "READARRAYGRID" in dumped
+    assert "MAXBOUND 2" in dumped
+    dump_data = [[float(x) for x in line.split()] for line in lines[2:12]]
+    dump_head = np.array(dump_data)
+    assert np.allclose(head, dump_head)
+
+    loaded = loads(dumped)
+    print("CHDG load:")
+    pprint(loaded)
+
+
 def test_dumps_wel():
     from flopy4.mf6.gwf import Dis, Gwf, Wel