feat(binaryfile): add write methods (#2722)

Support writing binary head, budget and grid files, both brand-new via classmethods, and copying existing files to new paths via instance methods.

Add write() classmethods to HeadFile and CellBudgetFile. These write a new file with the given data and return an instance with it open.

There are a few syntax variants for head files.

# dict keyed by (kstp, kper)
hds = HeadFile.write(
    'output.hds',
    # totim/pertim inferred (dt = 1.0 / time step)
    data={
        (1, 1): h_t1,
        (1, 2): h_t2,
    }
)

# list
hds = HeadFile.write('output.hds', data=[
    {'data': h_t2, 'kstp': 1, 'kper': 1, 'totim': 10.0, 'pertim': 10.0},
    {'data': h_t2, 'kstp': 1, 'kper': 2, 'totim': 20.0, 'pertim': 10.0},
])

# array/list with time as first dimension
# defaults to sequential stress periods: (1,1), (1,2), (1,3), ...
heads = [h_t1, h_t2, h_t3]  # or np.array([h_t1, h_t2, h_t3])
hds = HeadFile.write('output.hds', heads)

# array/list with custom tdis
hds = HeadFile.write('output.hds', heads, kstpkper=[(1, 1), (2, 1), (3, 1)])

Again there are variants for budget files. The typical case is to use data with a list, where each entry has "text", "kper", "kstp", and "data" entries, and "data" is an array, typically grid-shaped. But it can be convenient to use text and pass a time-indexed dictionary to data to create a file with a single variable.

# just face flows, dict keyed by (kstp, kper)
cbc = CellBudgetFile.write(
    'output.cbc',
    text='FLOW-JA-FACE',
    nlay=3,
    nrow=10,
    ncol=20,
    data={
        (1, 1): q_t1,
        (1, 2): q_t2,
    }
)

# multiple variables, list
cbc = CellBudgetFile.write(
    'output.cbc',
    data=[
        {'data': q_t1, 'kstp': 1, 'kper': 1, 'totim': 10.0,
         'text': 'FLOW-JA-FACE'},
        {'data': ...},
        ...
    ]
)

# array/list with time dimension (grid-shaped data like storage)
# defaults to sequential stress periods: (1,1), (1,2), (1,3), ...
# grid dimensions inferred from array shape
storage = [s_t1, s_t2, s_t3]  # nlay x nrow x ncol arrays
cbc = CellBudgetFile.write('output.cbc', storage, text='STORAGE')

# for face flows, grid dimensions are required since data is 1D
flows = [q_t1, q_t2]  # 1D arrays
cbc = CellBudgetFile.write(
    'output.cbc',
    flows,
    text='FLOW-JA-FACE',
    nlay=3, nrow=10, ncol=20
)

If only face flows are provided, the grid shape must be specified with nlay/nrow/ncol, nlay/ncpl, or nnodes. If grid-shaped variables are provided, the grid's shape will be inferred.
Instance methods

Add instance export() methods to MfGrdFile, HeadFile, and CellBudgetFile. These copy the contents of an open file to another, optionally filtering by variable and/or time step, or changing the precision. There is no write() method for MfGrdFile, as its signature would have been long and complicated to accommodate all grid types.

from flopy.mf6.utils.binarygrid_util import MfGrdFile

grb = MfGrdFile("model.grb")
grb.export("copy.grb") # copy to another file
grb.export("diff_prec.grb", precision="single") # different precision

from flopy.utils.binaryfile import HeadFile, CellBudgetFile

hds = HeadFile("model.hds")
hds.export("copy.hds")  # copy to another file
hds.export("filtered.hds", kstpkper=[(1, 0), (1, 1)])  # filter time steps
hds.export("diff_prec.hds", precision="single")  # different precision

cbc = CellBudgetFile("model.cbc")
cbc.export("copy.cbc")
cbc.export("flowja.cbc", text="FLOW-JA-FACE")
cbc.export("bndpkgs.cbc", text=["STORAGE", "CONSTANT HEAD"])
cbc.export("filtered.cbc", kstpkper=[(1, 0)], text="FLOW-JA-FACE")

Author: wpbonelli

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.export(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.export(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.export(output_file, verbose=False)
+
+    benchmark(write_grb)
+    assert output_file.exists()