test(rcha): support writing layer array input package

flopy4/mf6/gwf/__init__.py

@@ -13,13 +13,15 @@
 from flopy4.mf6.gwf.ic import Ic
 from flopy4.mf6.gwf.npf import Npf
 from flopy4.mf6.gwf.oc import Oc
+from flopy4.mf6.gwf.rch import Rch
+from flopy4.mf6.gwf.rcha import Rcha
 from flopy4.mf6.gwf.wel import Wel
 from flopy4.mf6.model import Model
 from flopy4.mf6.spec import field, path
 from flopy4.mf6.utils import open_cbc, open_hds
 from flopy4.utils import to_path
 
-__all__ = ["Gwf", "Chd", "Dis", "Drn", "Ic", "Npf", "Oc", "Wel"]
+__all__ = ["Gwf", "Chd", "Dis", "Drn", "Ic", "Npf", "Oc", "Rch", "Rcha", "Wel"]
 
 
 def convert_grid(value):
@@ -76,6 +78,8 @@ def budget(self):
     oc: Oc = field(block="packages")
     npf: Npf = field(block="packages")
     chd: list[Chd] = field(block="packages")
+    rch: list[Rch] = field(block="packages")
+    rcha: list[Rcha] = field(block="packages")
     wel: list[Wel] = field(block="packages")
     drn: list[Drn] = field(block="packages")
     output: Output = attrs.field(

flopy4/mf6/gwf/rcha.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.utils import to_path
+
+
+@xattree
+class Rcha(Package):
+    multi_package: ClassVar[bool] = True
+    fixed_cell: bool = field(block="options", default=False)
+    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)
+    save_flows: bool = field(block="options", default=False)
+    ts_filerecord: Optional[Path] = path(
+        block="options", default=None, converter=to_path, inout="filein"
+    )
+    obs_filerecord: Optional[Path] = path(
+        block="options", default=None, converter=to_path, inout="fileout"
+    )
+    irch: Optional[NDArray[np.int64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            # "ncpl",
+            "nrow",
+            "ncol",
+        ),
+        default=1,
+        converter=Converter(structure_array, takes_self=True, takes_field=True),
+    )
+    recharge: Optional[NDArray[np.float64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            # "ncpl",
+            "nrow",
+            "ncol",
+        ),
+        default=None,
+        converter=Converter(structure_array, takes_self=True, takes_field=True),
+    )
+    aux: Optional[NDArray[np.float64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            # "ncpl",
+            "nrow",
+            "ncol",
+        ),
+        default=None,
+        converter=Converter(structure_array, takes_self=True, takes_field=True),
+    )

test/test_codec.py

@@ -1,10 +1,13 @@
 from pprint import pprint
 
+import numpy as np
 import pytest
 
 from flopy4.mf6.codec import dumps, loads
 from flopy4.mf6.converter import COMPONENT_CONVERTER
 
+DNODATA = 3.0e30
+
 
 def test_loads_dis_generic_simple():
     mf6_input = """
@@ -322,6 +325,80 @@ def test_dumps_chd_2():
     pprint(loaded)
 
 
+def test_dumps_rch():
+    from flopy4.mf6.gwf import Dis, Gwf, Rch
+
+    dis = Dis(nlay=1, nrow=20, ncol=30)
+    gwf = Gwf(dis=dis)
+
+    boundaries = {}
+    for row in range(5, 15):
+        boundaries[(0, row, 0)] = 100.0
+    for row in range(8, 12):
+        boundaries[(0, row, 29)] = 95.0
+    for col in range(10, 20):
+        boundaries[(0, 19, col)] = 98.0
+
+    rch = Rch(
+        parent=gwf, recharge={0: boundaries}, print_input=True, save_flows=True, dims={"nper": 1}
+    )
+
+    dumped = dumps(COMPONENT_CONVERTER.unstructure(rch))
+    print("RCH dump:")
+    print(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) == 24
+    assert "100.0" in dumped  # Left boundary
+    assert "95.0" in dumped  # Right boundary
+    assert "98.0" in dumped  # Bottom boundary
+    assert "1e+30" not in dumped
+    assert "1.0e+30" not in dumped
+
+    loaded = loads(dumped)
+    print("RCH load:")
+    pprint(loaded)
+
+
+def test_dumps_rcha():
+    from flopy4.mf6.gwf import Dis, Gwf, Rcha
+
+    dis = Dis(nlay=1, nrow=20, ncol=30)
+    gwf = Gwf(dis=dis)
+
+    boundaries = np.full((20, 30), DNODATA, dtype=float)
+    for row in range(5, 15):
+        boundaries[row, 0] = 100.0
+    for row in range(8, 12):
+        boundaries[row, 29] = 95.0
+    for col in range(10, 20):
+        boundaries[19, col] = 98.0
+
+    rch = Rcha(
+        parent=gwf, recharge={0: boundaries}, print_input=True, save_flows=True, dims={"nper": 1}
+    )
+
+    dumped = dumps(COMPONENT_CONVERTER.unstructure(rch))
+    print("RCH dump:")
+    print(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) == 24
+    assert "100.0" in dumped  # Left boundary
+    assert "95.0" in dumped  # Right boundary
+    assert "98.0" in dumped  # Bottom boundary
+    assert "1e+30" not in dumped
+    assert "1.0e+30" not in dumped
+
+    loaded = loads(dumped)
+    print("RCH load:")
+    pprint(loaded)
+
+
 def test_dumps_wel_with_aux():
     from flopy4.mf6.gwf import Dis, Gwf, Wel