feat(binaryfile): add write methods

autotest/benchmark_binaryfile_write.py

@@ -0,0 +1,151 @@
+"""Benchmark tests for binaryfile write methods."""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from flopy.mf6.utils import MfGrdFile
+from flopy.utils import CellBudgetFile, HeadFile
+
+
+@pytest.fixture
+def freyberg_hds_path(example_data_path):
+    return example_data_path / "freyberg_multilayer_transient" / "freyberg.hds"
+
+
+@pytest.fixture
+def freyberg_cbc_path(example_data_path):
+    return example_data_path / "freyberg_multilayer_transient" / "freyberg.cbc"
+
+
+@pytest.fixture
+def mfgrd_dis_path(example_data_path):
+    return example_data_path / "mf6-freyberg" / "freyberg.dis.grb"
+
+
+@pytest.fixture
+def mfgrd_disv_path(example_data_path):
+    return (
+        example_data_path
+        / "mf6"
+        / "test006_gwf3_disv"
+        / "expected_output"
+        / "flow.disv.grb"
+    )
+
+
+@pytest.mark.slow
+def test_headfile_write_benchmark(benchmark, freyberg_hds_path, tmp_path):
+    hds = HeadFile(freyberg_hds_path)
+    nsteps = min(100, len(hds.kstpkper))
+    kstpkper = hds.kstpkper[:nsteps]
+    output_file = tmp_path / "benchmark_output.hds"
+
+    def write_head():
+        hds.write(output_file, kstpkper=kstpkper)
+
+    benchmark(write_head)
+    assert output_file.exists()
+
+
+@pytest.mark.slow
+def test_cellbudgetfile_write_benchmark(benchmark, freyberg_cbc_path, tmp_path):
+    cbc = CellBudgetFile(freyberg_cbc_path)
+    nsteps = min(50, len(cbc.kstpkper))
+    kstpkper = cbc.kstpkper[:nsteps]
+    output_file = tmp_path / "benchmark_output.cbc"
+
+    def write_budget():
+        cbc.write(output_file, kstpkper=kstpkper)
+
+    benchmark(write_budget)
+    assert output_file.exists()
+
+
+@pytest.mark.slow
+def test_mfgrdfile_write_benchmark_dis(benchmark, tmp_path):
+    nlay, nrow, ncol = 10, 100, 100
+    nodes = nlay * nrow * ncol
+    nja = 7 * nodes - 2 * (nlay * nrow + nlay * ncol + nrow * ncol)
+    grb_data = {
+        "NCELLS": nodes,
+        "NLAY": nlay,
+        "NROW": nrow,
+        "NCOL": ncol,
+        "NJA": nja,
+        "XORIGIN": 0.0,
+        "YORIGIN": 0.0,
+        "ANGROT": 0.0,
+        "DELR": np.ones(ncol, dtype=np.float64) * 100.0,
+        "DELC": np.ones(nrow, dtype=np.float64) * 100.0,
+        "TOP": np.ones(nodes, dtype=np.float64) * 100.0,
+        "BOTM": np.arange(nodes, dtype=np.float64),
+        "IA": np.arange(nodes + 1, dtype=np.int32),
+        "JA": np.arange(nja, dtype=np.int32) % nodes,
+        "IDOMAIN": np.ones(nodes, dtype=np.int32),
+        "ICELLTYPE": np.ones(nodes, dtype=np.int32),
+    }
+
+    from flopy.utils.utils_def import FlopyBinaryData
+
+    temp_grb = tmp_path / "temp_input.grb"
+    writer = FlopyBinaryData()
+    writer.precision = "double"
+
+    with open(temp_grb, "wb") as f:
+        writer.file = f
+        writer.write_text("GRID DIS\n", 50)
+        writer.write_text("VERSION 1\n", 50)
+        writer.write_text("NTXT 16\n", 50)
+        writer.write_text("LENTXT 100\n", 50)
+        var_list = [
+            ("NCELLS", "INTEGER", 0, []),
+            ("NLAY", "INTEGER", 0, []),
+            ("NROW", "INTEGER", 0, []),
+            ("NCOL", "INTEGER", 0, []),
+            ("NJA", "INTEGER", 0, []),
+            ("XORIGIN", "DOUBLE", 0, []),
+            ("YORIGIN", "DOUBLE", 0, []),
+            ("ANGROT", "DOUBLE", 0, []),
+            ("DELR", "DOUBLE", 1, [ncol]),
+            ("DELC", "DOUBLE", 1, [nrow]),
+            ("TOP", "DOUBLE", 1, [nodes]),
+            ("BOTM", "DOUBLE", 1, [nodes]),
+            ("IA", "INTEGER", 1, [nodes + 1]),
+            ("JA", "INTEGER", 1, [nja]),
+            ("IDOMAIN", "INTEGER", 1, [nodes]),
+            ("ICELLTYPE", "INTEGER", 1, [nodes]),
+        ]
+
+        for name, dtype_str, ndim, dims in var_list:
+            if ndim == 0:
+                line = f"{name} {dtype_str} NDIM {ndim}\n"
+            else:
+                dims_str = " ".join(str(d) for d in dims[::-1])
+                line = f"{name} {dtype_str} NDIM {ndim} {dims_str}\n"
+            writer.write_text(line, 100)
+
+        for name, dtype_str, ndim, dims in var_list:
+            value = grb_data[name]
+            if ndim == 0:
+                if dtype_str == "INTEGER":
+                    writer.write_integer(int(value))
+                else:
+                    writer.write_real(float(value))
+            else:
+                arr = np.asarray(value)
+                if dtype_str == "INTEGER":
+                    arr = arr.astype(np.int32)
+                elif dtype_str == "DOUBLE":
+                    arr = arr.astype(np.float64)
+                writer.write_record(arr.flatten(order="F"), dtype=arr.dtype)
+
+    grb = MfGrdFile(str(temp_grb), verbose=False)
+    output_file = tmp_path / "benchmark_output.grb"
+
+    def write_grb():
+        grb.write(output_file, verbose=False)
+
+    benchmark(write_grb)
+    assert output_file.exists()

autotest/test_binarygrid_util.py

@@ -159,3 +159,240 @@ def test_mfgrddisu_modelgrid(mfgrd_test_path):
     assert nvert == verts.shape[0], (
         f"number of vertex (x, y) pairs ({verts.shape[0]}) does not equal {nvert}"
     )
+
+
+def test_write_grb_instance_method(tmp_path, mfgrd_test_path):
+    original_file = mfgrd_test_path / "nwtp3.dis.grb"
+    grb_orig = MfGrdFile(original_file, verbose=False)
+
+    output_file = tmp_path / "test_instance.dis.grb"
+    grb_orig.write(output_file, verbose=False)
+
+    grb_new = MfGrdFile(output_file, verbose=False)
+
+    assert grb_new.grid_type == grb_orig.grid_type
+    assert grb_new.nodes == grb_orig.nodes
+    assert grb_new.nlay == grb_orig.nlay
+    assert grb_new.nrow == grb_orig.nrow
+    assert grb_new.ncol == grb_orig.ncol
+    assert grb_new.nja == grb_orig.nja
+
+    np.testing.assert_allclose(grb_new.xorigin, grb_orig.xorigin)
+    np.testing.assert_allclose(grb_new.yorigin, grb_orig.yorigin)
+    np.testing.assert_allclose(grb_new.angrot, grb_orig.angrot)
+
+    np.testing.assert_allclose(grb_new.delr, grb_orig.delr)
+    np.testing.assert_allclose(grb_new.delc, grb_orig.delc)
+    np.testing.assert_allclose(grb_new.top, grb_orig.top)
+    np.testing.assert_allclose(grb_new.bot, grb_orig.bot)
+
+    np.testing.assert_array_equal(grb_new.ia, grb_orig.ia)
+    np.testing.assert_array_equal(grb_new.ja, grb_orig.ja)
+    np.testing.assert_array_equal(grb_new.idomain, grb_orig.idomain)
+
+
+def test_write_grb_instance_method_precision_conversion(tmp_path, mfgrd_test_path):
+    original_file = mfgrd_test_path / "nwtp3.dis.grb"
+    grb = MfGrdFile(original_file, verbose=False)
+
+    single_file = tmp_path / "test_single.grb"
+    grb.write(single_file, precision="single", verbose=False)
+
+    double_file = tmp_path / "test_double.grb"
+    grb.write(double_file, precision="double", verbose=False)
+
+    grb_single = MfGrdFile(single_file, verbose=False)
+    grb_double = MfGrdFile(double_file, verbose=False)
+
+    assert grb_single.nodes == grb.nodes
+    assert grb_double.nodes == grb.nodes
+    assert single_file.stat().st_size < double_file.stat().st_size
+
+
+def test_write_grb_disv_roundtrip(tmp_path, mfgrd_test_path):
+    """Test MfGrdFile.write() for DISV grid with roundtrip validation."""
+    from flopy.mf6.utils.binarygrid_util import MfGrdFile
+
+    # Read original DISV grb file
+    original_file = mfgrd_test_path / "flow.disv.grb"
+    grb_orig = MfGrdFile(original_file, verbose=False)
+
+    # Write using instance method
+    output_file = tmp_path / "test_disv.grb"
+    grb_orig.write(output_file, verbose=False)
+
+    # Read it back
+    grb_new = MfGrdFile(output_file, verbose=False)
+
+    # Verify grid type and dimensions
+    assert grb_new.grid_type == "DISV"
+    assert grb_new.grid_type == grb_orig.grid_type
+    assert grb_new.nodes == grb_orig.nodes
+    assert grb_new.nlay == grb_orig.nlay
+    assert grb_new.ncpl == grb_orig.ncpl
+    assert grb_new.nja == grb_orig.nja
+
+    # Verify coordinates
+    np.testing.assert_allclose(grb_new.xorigin, grb_orig.xorigin)
+    np.testing.assert_allclose(grb_new.yorigin, grb_orig.yorigin)
+    np.testing.assert_allclose(grb_new.angrot, grb_orig.angrot)
+
+    # Verify elevation arrays
+    np.testing.assert_allclose(grb_new.top, grb_orig.top)
+    np.testing.assert_allclose(grb_new.bot, grb_orig.bot)
+
+    # Verify cell connectivity
+    np.testing.assert_array_equal(grb_new.ia, grb_orig.ia)
+    np.testing.assert_array_equal(grb_new.ja, grb_orig.ja)
+
+    # Verify DISV-specific data
+    assert grb_new._datadict["NVERT"] == grb_orig._datadict["NVERT"]
+    assert grb_new._datadict["NJAVERT"] == grb_orig._datadict["NJAVERT"]
+    np.testing.assert_allclose(
+        grb_new._datadict["VERTICES"], grb_orig._datadict["VERTICES"]
+    )
+    np.testing.assert_allclose(grb_new._datadict["CELLX"], grb_orig._datadict["CELLX"])
+    np.testing.assert_allclose(grb_new._datadict["CELLY"], grb_orig._datadict["CELLY"])
+    np.testing.assert_array_equal(
+        grb_new._datadict["IAVERT"], grb_orig._datadict["IAVERT"]
+    )
+    np.testing.assert_array_equal(
+        grb_new._datadict["JAVERT"], grb_orig._datadict["JAVERT"]
+    )
+    np.testing.assert_array_equal(grb_new.idomain, grb_orig.idomain)
+    np.testing.assert_array_equal(
+        grb_new._datadict["ICELLTYPE"], grb_orig._datadict["ICELLTYPE"]
+    )
+
+
+def test_write_grb_disv_precision_conversion(tmp_path, mfgrd_test_path):
+    """Test MfGrdFile.write() for DISV grid with precision conversion."""
+    from flopy.mf6.utils.binarygrid_util import MfGrdFile
+
+    # Read original DISV grb file
+    original_file = mfgrd_test_path / "flow.disv.grb"
+    grb = MfGrdFile(original_file, verbose=False)
+
+    # Write in single and double precision
+    single_file = tmp_path / "test_disv_single.grb"
+    grb.write(single_file, precision="single", verbose=False)
+
+    double_file = tmp_path / "test_disv_double.grb"
+    grb.write(double_file, precision="double", verbose=False)
+
+    # Read them back
+    grb_single = MfGrdFile(single_file, verbose=False)
+    grb_double = MfGrdFile(double_file, verbose=False)
+
+    # Verify dimensions are preserved
+    assert grb_single.nodes == grb.nodes
+    assert grb_double.nodes == grb.nodes
+    assert grb_single.grid_type == "DISV"
+    assert grb_double.grid_type == "DISV"
+
+    # Single precision file should be smaller
+    assert single_file.stat().st_size < double_file.stat().st_size
+
+    # Verify data values are preserved (with appropriate tolerances)
+    # Single precision has ~7 decimal digits of precision
+    np.testing.assert_allclose(grb_single.top, grb.top, rtol=1e-6)
+    np.testing.assert_allclose(grb_single.bot, grb.bot, rtol=1e-6)
+    np.testing.assert_allclose(
+        grb_single._datadict["VERTICES"], grb._datadict["VERTICES"], rtol=1e-6
+    )
+    np.testing.assert_allclose(
+        grb_single._datadict["CELLX"], grb._datadict["CELLX"], rtol=1e-6
+    )
+    np.testing.assert_allclose(
+        grb_single._datadict["CELLY"], grb._datadict["CELLY"], rtol=1e-6
+    )
+
+    # Double precision should match exactly (same precision as original)
+    np.testing.assert_allclose(grb_double.top, grb.top, rtol=1e-12)
+    np.testing.assert_allclose(grb_double.bot, grb.bot, rtol=1e-12)
+    np.testing.assert_allclose(
+        grb_double._datadict["VERTICES"], grb._datadict["VERTICES"], rtol=1e-12
+    )
+
+
+def test_write_grb_disu_roundtrip(tmp_path, mfgrd_test_path):
+    """Test MfGrdFile.write() for DISU grid with roundtrip validation."""
+    from flopy.mf6.utils.binarygrid_util import MfGrdFile
+
+    # Read original DISU grb file
+    original_file = mfgrd_test_path / "flow.disu.grb"
+    grb_orig = MfGrdFile(original_file, verbose=False)
+
+    # Write using instance method
+    output_file = tmp_path / "test_disu.grb"
+    grb_orig.write(output_file, verbose=False)
+
+    # Read it back
+    grb_new = MfGrdFile(output_file, verbose=False)
+
+    # Verify grid type and dimensions
+    assert grb_new.grid_type == "DISU"
+    assert grb_new.grid_type == grb_orig.grid_type
+    assert grb_new.nodes == grb_orig.nodes
+    assert grb_new.nja == grb_orig.nja
+
+    # Verify coordinates
+    np.testing.assert_allclose(grb_new.xorigin, grb_orig.xorigin)
+    np.testing.assert_allclose(grb_new.yorigin, grb_orig.yorigin)
+    np.testing.assert_allclose(grb_new.angrot, grb_orig.angrot)
+
+    # Verify elevation arrays (note: DISU uses TOP/BOT not TOP/BOTM)
+    np.testing.assert_allclose(grb_new.top, grb_orig.top)
+    np.testing.assert_allclose(grb_new._datadict["BOT"], grb_orig._datadict["BOT"])
+
+    # Verify cell connectivity
+    np.testing.assert_array_equal(grb_new.ia, grb_orig.ia)
+    np.testing.assert_array_equal(grb_new.ja, grb_orig.ja)
+
+    # Verify DISU-specific data
+    np.testing.assert_array_equal(
+        grb_new._datadict["ICELLTYPE"], grb_orig._datadict["ICELLTYPE"]
+    )
+
+    # IDOMAIN is optional in DISU - check if present
+    if "IDOMAIN" in grb_orig._datadict:
+        assert "IDOMAIN" in grb_new._datadict
+        np.testing.assert_array_equal(grb_new.idomain, grb_orig.idomain)
+
+
+def test_write_grb_disu_precision_conversion(tmp_path, mfgrd_test_path):
+    """Test MfGrdFile.write() for DISU grid with precision conversion."""
+    from flopy.mf6.utils.binarygrid_util import MfGrdFile
+
+    # Read original DISU grb file
+    original_file = mfgrd_test_path / "flow.disu.grb"
+    grb = MfGrdFile(original_file, verbose=False)
+
+    # Write in single and double precision
+    single_file = tmp_path / "test_disu_single.grb"
+    grb.write(single_file, precision="single", verbose=False)
+
+    double_file = tmp_path / "test_disu_double.grb"
+    grb.write(double_file, precision="double", verbose=False)
+
+    # Read them back
+    grb_single = MfGrdFile(single_file, verbose=False)
+    grb_double = MfGrdFile(double_file, verbose=False)
+
+    # Verify dimensions are preserved
+    assert grb_single.nodes == grb.nodes
+    assert grb_double.nodes == grb.nodes
+    assert grb_single.grid_type == "DISU"
+    assert grb_double.grid_type == "DISU"
+
+    # Single precision file should be smaller
+    assert single_file.stat().st_size < double_file.stat().st_size
+
+    # Verify data values are preserved (with appropriate tolerances)
+    # Single precision has ~7 decimal digits of precision
+    np.testing.assert_allclose(grb_single.top, grb.top, rtol=1e-6)
+    np.testing.assert_allclose(grb_single.bot, grb.bot, rtol=1e-6)
+
+    # Double precision should match exactly (same precision as original)
+    np.testing.assert_allclose(grb_double.top, grb.top, rtol=1e-12)
+    np.testing.assert_allclose(grb_double.bot, grb.bot, rtol=1e-12)