From cd5f5772aba96d4b30393caeed3b66fd2cc2fc41 Mon Sep 17 00:00:00 2001 From: Bonelli Date: Tue, 24 Feb 2026 19:54:37 -0500 Subject: [PATCH 01/14] feat(binaryfile): add write methods --- autotest/test_binaryfile.py | 480 ++++++++++++++++++++ flopy/mf6/utils/__init__.py | 3 +- flopy/mf6/utils/binarygrid_util.py | 281 +++++++++++- flopy/utils/binaryfile/__init__.py | 675 +++++++++++++++++++++++++++++ 4 files changed, 1435 insertions(+), 4 deletions(-) diff --git a/autotest/test_binaryfile.py b/autotest/test_binaryfile.py index 59558d98a2..f02068c08c 100644 --- a/autotest/test_binaryfile.py +++ b/autotest/test_binaryfile.py @@ -775,3 +775,483 @@ def test_headfile_get_ts_disu_grid(dis_sim, function_tmpdir): ts_old_list, err_msg="DISU HeadFile: old list format should match new list format", ) + + +def test_headfile_write_roundtrip(tmp_path, example_data_path): + """Test HeadFile.write_head() with roundtrip validation.""" + # Read an existing head file + original_file = example_data_path / "mf6-freyberg" / "freyberg.hds" + original = HeadFile(original_file) + + # Extract all data + data_dict = {} + for idx, kstpkper_np in enumerate(original.kstpkper): + kstpkper = (int(kstpkper_np[0]), int(kstpkper_np[1])) + head = original.get_data(idx=idx) + totim = original.times[idx] + data_dict[kstpkper] = {"head": head, "pertim": totim, "totim": totim} + + # Write to a temp file + output_file = tmp_path / "test.hds" + HeadFile.write_head( + output_file, data_dict, text="HEAD", precision="double", verbose=False + ) + + # Read it back + readback = HeadFile(output_file) + + # Compare + assert len(readback.kstpkper) == len(original.kstpkper), ( + f"Number of time steps mismatch: " + f"{len(readback.kstpkper)} != {len(original.kstpkper)}" + ) + + for idx in range(len(original.kstpkper)): + orig_data = original.get_data(idx=idx) + read_data = readback.get_data(idx=idx) + np.testing.assert_allclose( + orig_data, read_data, err_msg=f"Head data mismatch at idx={idx}" + ) + + +def test_headfile_write_synthetic(tmp_path): + """Test HeadFile.write_head() with synthetic data.""" + nlay, nrow, ncol = 2, 10, 15 + + # Create synthetic head data for multiple time steps + data_dict = {} + for kper in range(1, 3): + for kstp in range(1, 3): + heads = np.random.rand(nlay, nrow, ncol) * 10.0 + 100.0 + totim = (kper - 1) * 10.0 + kstp + data_dict[(kstp, kper)] = { + "head": heads, + "pertim": float(kstp), + "totim": totim, + } + + # Write the file + output_file = tmp_path / "synthetic.hds" + HeadFile.write_head(output_file, data_dict, text="HEAD", precision="single") + + # Read it back + hf = HeadFile(output_file, precision="single") + + assert len(hf.kstpkper) == 4, f"Expected 4 time steps, got {len(hf.kstpkper)}" + + # Verify data roundtrip + for (kstp, kper), entry in data_dict.items(): + data_back = hf.get_data(kstpkper=(kstp - 1, kper - 1)) + np.testing.assert_allclose( + entry["head"], data_back, err_msg=f"Mismatch at kstp={kstp}, kper={kper}" + ) + + +def test_cellbudgetfile_write_array_format(tmp_path): + """Test CellBudgetFile.write_budget() with array format (imeth=1).""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + nlay, nrow, ncol = 1, 10, 20 + ncells = nlay * nrow * ncol + + # Build connectivity for this grid to get NJA + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + # Create synthetic FLOW-JA-FACE data (size NJA, not ncells) + flowja = np.random.rand(nja) - 0.5 + + budget_dict = { + (1, 1): { + "FLOW-JA-FACE": { + "imeth": 1, + "data": flowja, + "delt": 1.0, + "pertim": 10.0, + "totim": 10.0, + } + } + } + + # Write the budget file + output_file = tmp_path / "test_array.bud" + CellBudgetFile.write_budget( + output_file, + budget_dict, + nlay=nlay, + nrow=nrow, + ncol=ncol, + precision="double", + verbose=False, + ) + + # Read it back + cbb = CellBudgetFile(output_file, precision="double") + + assert len(cbb.textlist) == 1, f"Expected 1 budget term, got {len(cbb.textlist)}" + # Text is padded to 16 characters in binary format + assert "FLOW-JA-FACE".ljust(16).encode() in cbb.textlist, ( + "FLOW-JA-FACE not in budget file" + ) + + # Get data back + data_back = cbb.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] + + # FLOW-JA-FACE is returned as 3D array with shape (1, 1, NJA) + # This matches how real MF6 files are structured + assert data_back.shape == (1, 1, nja), ( + f"Expected shape (1, 1, {nja}), got {data_back.shape}" + ) + np.testing.assert_allclose( + flowja, data_back.flatten(), err_msg="FLOW-JA-FACE data mismatch" + ) + + +def test_cellbudgetfile_write_array_format_regular(tmp_path): + """Test CellBudgetFile.write_budget() with array format for regular budget term.""" + nlay, nrow, ncol = 2, 5, 10 + ncells = nlay * nrow * ncol + + # Create synthetic budget data (e.g., storage) + storage = np.random.rand(nlay, nrow, ncol) * 100.0 + + budget_dict = { + (1, 1): { + "STORAGE": { + "imeth": 1, + "data": storage, + "delt": 1.0, + "pertim": 10.0, + "totim": 10.0, + } + } + } + + # Write the budget file + output_file = tmp_path / "test_regular.bud" + CellBudgetFile.write_budget( + output_file, + budget_dict, + nlay=nlay, + nrow=nrow, + ncol=ncol, + precision="double", + verbose=False, + ) + + # Read it back + cbb = CellBudgetFile(output_file, precision="double") + + assert "STORAGE".ljust(16).encode() in cbb.textlist, "STORAGE not in budget file" + + # Get data back + data_back = cbb.get_data(kstpkper=(0, 0), text="STORAGE")[0] + + # Regular budget terms are returned as 3D arrays + assert data_back.shape == ( + nlay, + nrow, + ncol, + ), f"Expected shape {(nlay, nrow, ncol)}, got {data_back.shape}" + np.testing.assert_allclose(storage, data_back, err_msg="STORAGE data mismatch") + + +def test_cellbudgetfile_write_list_format(tmp_path): + """Test CellBudgetFile.write_budget() with list format (imeth=6).""" + nlay, nrow, ncol = 1, 10, 20 + ncells = nlay * nrow * ncol + + # Create synthetic DATA-SPDIS list data + dt = np.dtype( + [ + ("node", np.int32), + ("node2", np.int32), + ("q", np.float64), + ("qx", np.float64), + ("qy", np.float64), + ("qz", np.float64), + ] + ) + + data_list = np.zeros(ncells, dtype=dt) + data_list["node"] = np.arange(1, ncells + 1) + data_list["node2"] = np.arange(1, ncells + 1) + data_list["q"] = 0.0 + data_list["qx"] = np.random.rand(ncells) * 0.001 + data_list["qy"] = np.random.rand(ncells) * 0.001 + data_list["qz"] = np.random.rand(ncells) * 0.001 + + budget_dict = { + (1, 1): { + "DATA-SPDIS": { + "imeth": 6, + "data": data_list, + "delt": 1.0, + "pertim": 10.0, + "totim": 10.0, + "modelnam": "GWF", + "paknam": "", + "modelnam2": "GWF", + "paknam2": "", + "ndat": 3, + "auxtxt": ["qx", "qy", "qz"], + } + } + } + + # Write the budget file + output_file = tmp_path / "test_list.bud" + CellBudgetFile.write_budget( + output_file, + budget_dict, + nlay=nlay, + nrow=nrow, + ncol=ncol, + precision="double", + verbose=False, + ) + + # Read it back + cbb = CellBudgetFile(output_file, precision="double") + + # Text is padded to 16 characters in binary format + assert "DATA-SPDIS".ljust(16).encode() in cbb.textlist, ( + "DATA-SPDIS not in budget file" + ) + + # Get data back + data_back = cbb.get_data(kstpkper=(0, 0), text="DATA-SPDIS")[0] + + assert len(data_back) == ncells, f"Expected {ncells} entries, got {len(data_back)}" + np.testing.assert_allclose(data_list["qx"], data_back["qx"], err_msg="qx mismatch") + np.testing.assert_allclose(data_list["qy"], data_back["qy"], err_msg="qy mismatch") + np.testing.assert_allclose(data_list["qz"], data_back["qz"], err_msg="qz mismatch") + + +def test_cellbudgetfile_write_multiple_terms(tmp_path): + """Test CellBudgetFile.write_budget() with multiple budget terms.""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + nlay, nrow, ncol = 1, 10, 20 + ncells = nlay * nrow * ncol + + # Build connectivity for this grid to get NJA + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + # Create FLOW-JA-FACE array data (size NJA) + flowja = np.random.rand(nja) - 0.5 + + # Create DATA-SAT list data + dt_sat = np.dtype([("node", np.int32), ("node2", np.int32), ("q", np.float64)]) + sat_list = np.zeros(ncells, dtype=dt_sat) + sat_list["node"] = np.arange(1, ncells + 1) + sat_list["node2"] = np.arange(1, ncells + 1) + sat_list["q"] = np.random.rand(ncells) + + budget_dict = { + (1, 1): { + "FLOW-JA-FACE": { + "imeth": 1, + "data": flowja, + "delt": 1.0, + "pertim": 10.0, + "totim": 10.0, + }, + "DATA-SAT": { + "imeth": 6, + "data": sat_list, + "delt": 1.0, + "pertim": 10.0, + "totim": 10.0, + "modelnam": "GWF", + "paknam": "", + "modelnam2": "GWF", + "paknam2": "", + "ndat": 1, + "auxtxt": [], + }, + } + } + + # Write the budget file + output_file = tmp_path / "test_multi.bud" + CellBudgetFile.write_budget( + output_file, + budget_dict, + nlay=nlay, + nrow=nrow, + ncol=ncol, + precision="double", + verbose=False, + ) + + # Read it back + cbb = CellBudgetFile(output_file, precision="double") + + assert len(cbb.textlist) == 2, f"Expected 2 budget terms, got {len(cbb.textlist)}" + + # Verify FLOW-JA-FACE (returned as shape (1, 1, NJA)) + flowja_back = cbb.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] + assert flowja_back.shape == (1, 1, nja), ( + f"Expected shape (1, 1, {nja}), got {flowja_back.shape}" + ) + np.testing.assert_allclose(flowja, flowja_back.flatten()) + + # Verify DATA-SAT + sat_back = cbb.get_data(kstpkper=(0, 0), text="DATA-SAT")[0] + np.testing.assert_allclose(sat_list["q"], sat_back["q"]) + + +def test_headfile_instance_write(tmp_path): + """Test HeadFile.write() instance method.""" + # Create test data + data_dict = { + (1, 1): { + "head": np.full((3, 10, 10), 35.0), + "pertim": 1.0, + "totim": 1.0, + }, + (2, 1): { + "head": np.full((3, 10, 10), 34.0), + "pertim": 2.0, + "totim": 2.0, + }, + } + + # Write initial file + file1 = tmp_path / "test.hds" + HeadFile.write_head(file1, data_dict, precision="double") + + # Read it back and write using instance method + hds = HeadFile(file1, precision="double") + file2 = tmp_path / "test_copy.hds" + hds.write(file2) + hds.close() + + # Verify the copy + hds2 = HeadFile(file2, precision="double") + assert len(hds2.kstpkper) == 2 + head1 = hds2.get_data(totim=1.0) + head2 = hds2.get_data(totim=2.0) + np.testing.assert_allclose(head1, 35.0) + np.testing.assert_allclose(head2, 34.0) + hds2.close() + + +def test_headfile_instance_write_filtered(tmp_path): + """Test HeadFile.write() with kstpkper filtering.""" + # Create test data with 3 time steps + data_dict = { + (1, 1): {"head": np.full((2, 5, 5), 35.0), "pertim": 1.0, "totim": 1.0}, + (2, 1): {"head": np.full((2, 5, 5), 34.0), "pertim": 2.0, "totim": 2.0}, + (3, 1): {"head": np.full((2, 5, 5), 33.0), "pertim": 3.0, "totim": 3.0}, + } + + # Write initial file + file1 = tmp_path / "test.hds" + HeadFile.write_head(file1, data_dict, precision="double") + + # Read and write only specific time steps + hds = HeadFile(file1, precision="double") + file2 = tmp_path / "test_filtered.hds" + hds.write(file2, kstpkper=[(1, 1), (3, 1)]) # Skip (2, 1) + hds.close() + + # Verify only 2 time steps in output + hds2 = HeadFile(file2, precision="double") + assert len(hds2.kstpkper) == 2 + assert (1, 1) in hds2.kstpkper + assert (3, 1) in hds2.kstpkper + assert (2, 1) not in hds2.kstpkper + hds2.close() + + +def test_cellbudgetfile_instance_write(tmp_path): + """Test CellBudgetFile.write() instance method.""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + nlay, nrow, ncol = 1, 10, 20 + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + # Create test budget data + flowja = np.random.rand(nja) - 0.5 + budget_dict = { + (1, 1): { + "FLOW-JA-FACE": { + "imeth": 1, + "data": flowja, + "delt": 1.0, + "pertim": 1.0, + "totim": 1.0, + } + } + } + + # Write initial file (no grid dimensions needed for FLOW-JA-FACE only) + file1 = tmp_path / "test.cbc" + CellBudgetFile.write_budget(file1, budget_dict, precision="double") + + # Read and write using instance method + cbb = CellBudgetFile(file1, precision="double") + file2 = tmp_path / "test_copy.cbc" + cbb.write(file2) + cbb.close() + + # Verify the copy + cbb2 = CellBudgetFile(file2, precision="double") + assert len(cbb2.kstpkper) == 1 + # get_data uses 0-based indexing + flowja_back = cbb2.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] + np.testing.assert_allclose(flowja, flowja_back.flatten()) + cbb2.close() + + +def test_cellbudgetfile_instance_write_filtered(tmp_path): + """Test CellBudgetFile.write() with text filtering.""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + nlay, nrow, ncol = 1, 10, 20 + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + # Create budget data with two terms + flowja = np.random.rand(nja) - 0.5 + storage = np.random.rand(nlay, nrow, ncol) + + budget_dict = { + (1, 1): { + "FLOW-JA-FACE": { + "imeth": 1, + "data": flowja, + "delt": 1.0, + "pertim": 1.0, + "totim": 1.0, + }, + "STORAGE": { + "imeth": 1, + "data": storage, + "delt": 1.0, + "pertim": 1.0, + "totim": 1.0, + }, + } + } + + # Write initial file + file1 = tmp_path / "test.cbc" + CellBudgetFile.write_budget( + file1, budget_dict, nlay=nlay, nrow=nrow, ncol=ncol, precision="double" + ) + + # Read and write only FLOW-JA-FACE + cbb = CellBudgetFile(file1, precision="double") + file2 = tmp_path / "test_filtered.cbc" + cbb.write(file2, text="FLOW-JA-FACE") + cbb.close() + + # Verify only FLOW-JA-FACE in output + cbb2 = CellBudgetFile(file2, precision="double") + assert len(cbb2.textlist) == 1 + assert b"FLOW-JA-FACE" in cbb2.textlist[0] + # get_data uses 0-based indexing + flowja_back = cbb2.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] + np.testing.assert_allclose(flowja, flowja_back.flatten()) + cbb2.close() diff --git a/flopy/mf6/utils/__init__.py b/flopy/mf6/utils/__init__.py index 45de130dd9..f442647a55 100644 --- a/flopy/mf6/utils/__init__.py +++ b/flopy/mf6/utils/__init__.py @@ -1,6 +1,5 @@ -from .binarygrid_util import MfGrdFile +from .binarygrid_util import MfGrdFile, build_structured_connectivity from .generate_classes import generate_classes from .lakpak_utils import get_lak_connections from .mfsimlistfile import MfSimulationList from .model_splitter import Mf6Splitter -from .postprocessing import get_residuals, get_structured_faceflows diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index 4fe45e0b0b..e7e049b21f 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -317,8 +317,7 @@ def _get_verts(self): if self._grid_type == "DISU": # modify verts verts = [ - [idx, verts[idx, 0], verts[idx, 1]] - for idx in range(shpvert[0]) + [idx, verts[idx, 0], verts[idx, 1]] for idx in range(shpvert[0]) ] if self.verbose: print(f"returning verts from {self.file.name}") @@ -747,3 +746,281 @@ def cell2d(self): else: vertices, cell2d = None, None return vertices, cell2d + + @staticmethod + def write_grb( + filename, + grid_type, + data_dict, + version=1, + precision="double", + verbose=False, + ): + """ + Write a MODFLOW 6 binary grid file (.grb). + + Parameters + ---------- + filename : str or PathLike + Path to output .grb file + grid_type : str + Grid type: 'DIS', 'DISV', or 'DISU' + data_dict : dict + Dictionary with grid data arrays. Required keys depend on grid_type. + For DIS grids: NCELLS, NLAY, NROW, NCOL, NJA, XORIGIN, YORIGIN, ANGROT, + DELR, DELC, TOP, BOTM, IA, JA, IDOMAIN, ICELLTYPE + version : int, optional + Grid file version (default 1) + precision : str, optional + 'single' or 'double' (default 'double') + verbose : bool, optional + Print progress messages (default False) + + Notes + ----- + The binary grid file format consists of: + 1. Text header lines (50 chars each): + - "GRID {grid_type}" + - "VERSION {version}" + - "NTXT {ntxt}" + - "LENTXT {lentxt}" + 2. Variable definition lines (100 chars each): + - "{NAME} {TYPE} NDIM {ndim} {dimensions...}" + 3. Binary data for each variable + + Arrays should be in Python (row-major) order and will be written + in Fortran (column-major) order as required by MF6. + + Examples + -------- + >>> import numpy as np + >>> from flopy.mf6.utils import MfGrdFile + >>> data = { + ... 'NCELLS': 800, + ... 'NLAY': 1, + ... 'NROW': 40, + ... 'NCOL': 20, + ... 'NJA': 3367, + ... 'XORIGIN': 0.0, + ... 'YORIGIN': 0.0, + ... 'ANGROT': 0.0, + ... 'DELR': np.full(20, 250.0), + ... 'DELC': np.full(40, 250.0), + ... 'TOP': np.full(800, 35.0), + ... 'BOTM': np.random.rand(800), + ... 'IA': ia_array, + ... 'JA': ja_array, + ... 'IDOMAIN': np.ones(800, dtype=np.int32), + ... 'ICELLTYPE': np.zeros(800, dtype=np.int32) + ... } + >>> MfGrdFile.write_grb('model.dis.grb', 'DIS', data) + """ + import os + + # Create FlopyBinaryData instance for write helpers + writer = FlopyBinaryData() + writer.precision = precision + + # Define variable metadata based on grid type + if grid_type.upper() == "DIS": + 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, [data_dict.get("NCOL", 0)]), + ("DELC", "DOUBLE", 1, [data_dict.get("NROW", 0)]), + ("TOP", "DOUBLE", 1, [data_dict.get("NCELLS", 0)]), + ("BOTM", "DOUBLE", 1, [data_dict.get("NCELLS", 0)]), + ("IA", "INTEGER", 1, [data_dict.get("NCELLS", 0) + 1]), + ("JA", "INTEGER", 1, [data_dict.get("NJA", 0)]), + ("IDOMAIN", "INTEGER", 1, [data_dict.get("NCELLS", 0)]), + ("ICELLTYPE", "INTEGER", 1, [data_dict.get("NCELLS", 0)]), + ] + else: + raise NotImplementedError( + f"Grid type {grid_type} not yet implemented. " + "Currently only DIS grids are supported." + ) + + ntxt = len(var_list) + lentxt = 100 + + if verbose: + print(f"Writing binary grid file: {filename}") + print(f" Grid type: {grid_type}") + print(f" Version: {version}") + print(f" Number of variables: {ntxt}") + + with open(filename, "wb") as f: + writer.file = f + + # Write text header lines (50 chars each, newline terminated) + header_len = 50 + writer.write_text(f"GRID {grid_type.upper()}\n", header_len) + writer.write_text(f"VERSION {version}\n", header_len) + writer.write_text(f"NTXT {ntxt}\n", header_len) + writer.write_text(f"LENTXT {lentxt}\n", header_len) + + # Write variable definition lines (100 chars each) + 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] + ) # Reverse for Fortran order + line = f"{name} {dtype_str} NDIM {ndim} {dims_str}\n" + writer.write_text(line, lentxt) + + # Write binary data for each variable + for name, dtype_str, ndim, dims in var_list: + if name not in data_dict: + raise ValueError( + f"Required variable '{name}' not found in data_dict" + ) + + value = data_dict[name] + + if verbose: + if ndim == 0: + print(f" Writing {name} = {value}") + else: + if hasattr(value, "min"): + print( + f" Writing {name}: min = {value.min()} max = {value.max()}" + ) + else: + print(f" Writing {name}") + + # Write scalar or array data + if ndim == 0: + # Scalar value + if dtype_str == "INTEGER": + writer.write_integer(int(value)) + elif dtype_str in ("DOUBLE", "SINGLE"): + writer.write_real(float(value)) + else: + # Array data + arr = np.asarray(value) + if dtype_str == "INTEGER": + arr = arr.astype(np.int32) + elif dtype_str == "DOUBLE": + arr = arr.astype(np.float64) + elif dtype_str == "SINGLE": + arr = arr.astype(np.float32) + + # Write array in column-major (Fortran) order + writer.write_record(arr, dtype=arr.dtype) + + if verbose: + print(f"Successfully wrote {filename}") + + +def build_structured_connectivity(nlay, nrow, ncol, idomain=None): + """ + Build IA and JA connectivity arrays for a structured (DIS) grid. + + Parameters + ---------- + nlay : int + Number of layers + nrow : int + Number of rows + ncol : int + Number of columns + idomain : np.ndarray, optional + Domain array indicating active (>0) and inactive (<=0) cells. + Shape: (nlay, nrow, ncol). If None, all cells are active. + + Returns + ------- + ia : np.ndarray + Index array (CSR format), shape (ncells + 1,), dtype int32. + ia[n] is the starting position in ja for cell n's connections. + ia[ncells] is the total number of connections. + ja : np.ndarray + Connection array (CSR format), shape (nja,), dtype int32. + Contains cell numbers for each connection (0-based). + nja : int + Total number of connections + + Notes + ----- + Connectivity order for structured grids (upper triangle only): + 1. Diagonal (self connection) + 2. Right (+1 in j, same k, i) + 3. Front (+1 in i, same k, j) + 4. Lower (+1 in k, same i, j) + + The IA/JA arrays use 0-based indexing (Python convention). + When writing to MF6 binary files, add 1 to convert to Fortran 1-based indexing. + + Examples + -------- + >>> import numpy as np + >>> from flopy.mf6.utils import build_structured_connectivity + >>> nlay, nrow, ncol = 2, 3, 3 + >>> ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + >>> print(f"Total cells: {nlay * nrow * ncol}, connections: {nja}") + Total cells: 18, connections: 42 + """ + ncells = nlay * nrow * ncol + + # Default to all active cells if idomain not provided + if idomain is None: + idomain = np.ones((nlay, nrow, ncol), dtype=np.int32) + else: + idomain = np.asarray(idomain, dtype=np.int32) + if idomain.shape != (nlay, nrow, ncol): + raise ValueError( + f"idomain shape {idomain.shape} does not match grid shape " + f"({nlay}, {nrow}, {ncol})" + ) + + ia = np.zeros(ncells + 1, dtype=np.int32) + ja_list = [] + nja = 0 + + for k in range(nlay): + for i in range(nrow): + for j in range(ncol): + node = k * nrow * ncol + i * ncol + j + + # Skip inactive cells - they still get an entry in IA + if idomain[k, i, j] <= 0: + ia[node + 1] = nja + continue + + # Add diagonal (self connection) + ja_list.append(node) + nja += 1 + + # Add connections to neighbors (upper triangle only) + # Right neighbor (j+1) + if j + 1 < ncol and idomain[k, i, j + 1] > 0: + m = k * nrow * ncol + i * ncol + (j + 1) + ja_list.append(m) + nja += 1 + + # Front neighbor (i+1) + if i + 1 < nrow and idomain[k, i + 1, j] > 0: + m = k * nrow * ncol + (i + 1) * ncol + j + ja_list.append(m) + nja += 1 + + # Lower neighbor (k+1) + if k + 1 < nlay and idomain[k + 1, i, j] > 0: + m = (k + 1) * nrow * ncol + i * ncol + j + ja_list.append(m) + nja += 1 + + ia[node + 1] = nja + + ja = np.array(ja_list, dtype=np.int32) + + return ia, ja, nja diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 07027f8524..776e00049f 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -639,6 +639,208 @@ def reverse_header(header): move(target, filename) super().__init__(filename, self.precision, self.verbose) + def write( + self, + filename: Union[str, PathLike], + kstpkper: Optional[list] = None, + **kwargs, + ): + """ + Write head data to a binary file. + + Convenience instance method that writes all or selected time steps + from the current file to a new file using the same format. + + Parameters + ---------- + filename : str or PathLike + Path to output head file + kstpkper : list of tuples, optional + List of (kstp, kper) tuples to write. If None, writes all time steps. + **kwargs + Additional keyword arguments passed to write_head(): + - text : str, identifier for head data (default uses current file's text) + - precision : str, 'single' or 'double' (default uses current file's precision) + - verbose : bool, print progress messages + + Examples + -------- + >>> hds = HeadFile('input.hds') + >>> # Write all time steps + >>> hds.write('output.hds') + >>> # Write specific time steps + >>> hds.write('output.hds', kstpkper=[(1, 0), (1, 1)]) + """ + # Determine which time steps to write + if kstpkper is None: + kstpkper = self.kstpkper + + # Build data dictionary + data_dict = {} + for ksp in kstpkper: + try: + # Convert numpy int32 to Python int if needed + kstp = int(ksp[0]) + kper = int(ksp[1]) + ksp_tuple = (kstp, kper) + + # Find the totim for this kstpkper + mask = (self.recordarray["kstp"] == kstp) & (self.recordarray["kper"] == kper) + matching_records = self.recordarray[mask] + if len(matching_records) == 0: + if kwargs.get("verbose", False): + print(f"Warning: No records found for {ksp}") + continue + + record = matching_records[0] + totim = float(record["totim"]) + + # Get data using totim (works for multi-layer files) + head_data = self.get_data(totim=totim) + + data_dict[ksp_tuple] = { + "head": head_data, + "pertim": float(record["pertim"]), + "totim": totim, + } + except Exception as e: + if kwargs.get("verbose", False): + print(f"Warning: Could not read data for {ksp}: {e}") + continue + + # Set defaults from current file if not provided + if "text" not in kwargs: + # Use first text entry from file + kwargs["text"] = self.recordarray["text"][0].decode().strip() + if "precision" not in kwargs: + kwargs["precision"] = self.precision + + # Write using static method + HeadFile.write_head(filename, data_dict, **kwargs) + + @staticmethod + def write_head( + filename, + data_dict, + text="HEAD", + precision="double", + verbose=False, + ): + """ + Write a MODFLOW binary head file (.hds). + + Parameters + ---------- + filename : str or PathLike + Path to output head file + data_dict : dict + Dictionary with head data and metadata. Each entry should be: + {(kstp, kper): {'head': array, 'pertim': float, 'totim': float}} + where head array has shape (nlay, nrow, ncol) or (nrow, ncol) + for single layer + text : str, optional + Text identifier for head data (default "HEAD"). Can also be "DRAWDOWN", etc. + Will be padded/truncated to 16 characters. + precision : str, optional + 'single' or 'double' (default 'double') + verbose : bool, optional + Print progress messages (default False) + + Notes + ----- + The head file format consists of repeating records, one per layer per time step: + 1. Header (8 values): + - kstp (int32): time step number (1-based) + - kper (int32): stress period number (1-based) + - pertim (float32 or float64): time in current stress period + - totim (float32 or float64): total elapsed time + - text (16 char): text identifier (e.g., "HEAD") + - ncol (int32): number of columns + - nrow (int32): number of rows + - ilay (int32): layer number (1-based) + 2. Data: head array (ncol, nrow) as float32 or float64 + + Examples + -------- + >>> import numpy as np + >>> from flopy.utils.binaryfile import HeadFile + >>> data = { + ... (1, 1): { + ... 'head': np.full((3, 10, 10), 35.0), + ... 'pertim': 1.0, + ... 'totim': 1.0 + ... } + ... } + >>> HeadFile.write_head('model.hds', data) + """ + # Set precision + if precision == "single": + realtype = np.float32 + else: + realtype = np.float64 + + # Ensure text is exactly 16 bytes + if isinstance(text, str): + text = text.encode("ascii") + if len(text) > 16: + text = text[:16] + elif len(text) < 16: + text = text + b" " * (16 - len(text)) + + if verbose: + print(f"Writing binary head file: {filename}") + print(f" Text identifier: {text.decode().strip()}") + print(f" Precision: {precision}") + print(f" Number of time steps: {len(data_dict)}") + + with open(filename, "wb") as f: + # Sort by time step for consistent output + sorted_keys = sorted(data_dict.keys()) + + for kstpkper in sorted_keys: + kstp, kper = kstpkper + entry = data_dict[kstpkper] + + head = np.asarray(entry["head"]) + pertim = entry["pertim"] + totim = entry["totim"] + + # Handle both 3D (nlay, nrow, ncol) and 2D (nrow, ncol) arrays + if head.ndim == 2: + head = head.reshape(1, head.shape[0], head.shape[1]) + + nlay, nrow, ncol = head.shape + + if verbose: + print(f" Writing kstp={kstp}, kper={kper}, totim={totim}") + print(f" Shape: {nlay} layers x {nrow} rows x {ncol} cols") + + # Define header dtype + dt = np.dtype( + [ + ("kstp", np.int32), + ("kper", np.int32), + ("pertim", realtype), + ("totim", realtype), + ("text", "S16"), + ("ncol", np.int32), + ("nrow", np.int32), + ("ilay", np.int32), + ] + ) + + # Write one record per layer + for ilay in range(nlay): + h = np.array( + (kstp, kper, pertim, totim, text, ncol, nrow, ilay + 1), + dtype=dt, + ) + h.tofile(f) + head[ilay].astype(realtype).tofile(f) + + if verbose: + print(f"Successfully wrote {filename}") + class UcnFile(BinaryLayerFile): """ @@ -2278,6 +2480,479 @@ def get_residual(self, totim, scaled=False): return residual + def write( + self, + filename: Union[str, PathLike], + kstpkper: Optional[list] = None, + text: Optional[Union[str, list]] = None, + **kwargs, + ): + """ + Write budget data to a binary file. + + Convenience instance method that writes all or selected budget records + from the current file to a new file using the same format. + + Parameters + ---------- + filename : str or PathLike + Path to output budget file + kstpkper : list of tuples, optional + List of (kstp, kper) tuples to write. If None, writes all time steps. + text : str or list of str, optional + Budget term(s) to write. If None, writes all terms. + Examples: 'FLOW-JA-FACE', ['STORAGE', 'CONSTANT HEAD'] + **kwargs + Additional keyword arguments passed to write_budget(): + - precision : str, 'single' or 'double' (default uses current file's precision) + - verbose : bool, print progress messages + + Examples + -------- + >>> cbc = CellBudgetFile('input.cbc') + >>> # Write all data + >>> cbc.write('output.cbc') + >>> # Write specific time steps + >>> cbc.write('output.cbc', kstpkper=[(1, 0), (1, 1)]) + >>> # Write specific budget terms + >>> cbc.write('output.cbc', text='FLOW-JA-FACE') + >>> # Write specific terms and time steps + >>> cbc.write('output.cbc', kstpkper=[(1, 0)], text=['STORAGE', 'FLOW-JA-FACE']) + """ + # Determine which time steps to write + if kstpkper is None: + kstpkper = self.kstpkper + + # Determine which text entries to write + if text is None: + textlist = self.textlist + elif isinstance(text, str): + textlist = [text.ljust(16).encode()] + else: + textlist = [t.ljust(16).encode() for t in text] + + # Build budget dictionary + budget_dict = {} + for ksp in kstpkper: + # Convert numpy int32 to Python int if needed + kstp = int(ksp[0]) + kper = int(ksp[1]) + ksp_tuple = (kstp, kper) + + # get_data() expects 0-based indexing, but kstpkper contains 1-based values + ksp_0based = (kstp - 1, kper - 1) + + for txt in textlist: + # Get matching text from file (case-insensitive, padded) + txt_str = txt.decode().strip() if isinstance(txt, bytes) else txt.strip() + + # Find matching records + matching_records = [ + t for t in self.textlist + if txt_str.upper() in t.decode().strip().upper() + ] + + for file_txt in matching_records: + try: + data = self.get_data(kstpkper=ksp_0based, text=file_txt)[0] + + # Get metadata from recordarray + mask = ( + (self.recordarray["kstp"] == kstp) + & (self.recordarray["kper"] == kper) + & (self.recordarray["text"] == file_txt) + ) + records = self.recordarray[mask] + + if len(records) == 0: + continue + + record = records[0] + + # Determine imeth from data structure + if isinstance(data, np.recarray): + imeth = 6 # List format + else: + imeth = 1 # Array format + + # Create dictionary key + key = (ksp_tuple, file_txt.decode().strip(), imeth) + + budget_dict[key] = { + "data": data, + "delt": float(record["delt"]), + "pertim": float(record["pertim"]), + "totim": float(record["totim"]), + } + + # Add model/package names if imeth=6 + if imeth == 6: + budget_dict[key].update({ + "modelnam": record["modelnam"].decode().strip(), + "paknam": record["paknam"].decode().strip(), + "modelnam2": record["modelnam2"].decode().strip(), + "paknam2": record["paknam2"].decode().strip(), + }) + + except Exception as e: + if kwargs.get("verbose", False): + print(f"Warning: Could not read data for {ksp}, {txt}: {e}") + continue + + # Restructure budget_dict to match write_budget format + # write_budget expects: {(kstp, kper): {'text': {...}, 'text2': {...}}} + # but we built: {((kstp, kper), text, imeth): {...}} + restructured_dict = {} + for key, value in budget_dict.items(): + ksp_tuple, text, imeth = key + if ksp_tuple not in restructured_dict: + restructured_dict[ksp_tuple] = {} + + # Add imeth to the value dict + value_with_imeth = value.copy() + value_with_imeth["imeth"] = imeth + restructured_dict[ksp_tuple][text] = value_with_imeth + + # Set defaults from current file if not provided + if "precision" not in kwargs: + kwargs["precision"] = self.precision + + # Write using static method + # Only pass grid dimensions if they're set (non-zero) + grid_kwargs = {} + if self.nlay > 0: + grid_kwargs["nlay"] = self.nlay + if self.nrow > 0: + grid_kwargs["nrow"] = self.nrow + if self.ncol > 0: + grid_kwargs["ncol"] = self.ncol + + CellBudgetFile.write_budget( + filename, + restructured_dict, + **grid_kwargs, + **kwargs, + ) + + @staticmethod + def write_budget( + filename, + budget_dict, + nlay=None, + nrow=None, + ncol=None, + precision="double", + verbose=False, + ): + """ + Write a MODFLOW 6 style binary budget file (.bud or .cbc). + + Parameters + ---------- + filename : str or PathLike + Path to output budget file + budget_dict : dict + Nested dictionary with budget data. Structure: + { + (kstp, kper): { + 'text': { + 'imeth': int (1 for array, 6 for list), + 'data': np.ndarray, + 'delt': float, + 'pertim': float, + 'totim': float, + 'modelnam': str (optional, for imeth=6), + 'paknam': str (optional, for imeth=6), + 'modelnam2': str (optional, for imeth=6), + 'paknam2': str (optional, for imeth=6), + 'ndat': int (optional, for imeth=6, number of data columns), + 'auxtxt': list of str (optional, for imeth=6, auxiliary names) + } + } + } + + For imeth=1 (array format): + data should be shape (nlay, nrow, ncol) + + For imeth=6 (list format): + data should be a numpy recarray with fields: + - 'node' (int32): source node number (1-based) + - 'node2' (int32): destination node number (1-based) + - 'q' (float): flow value + - optional auxiliary fields (float) + nlay : int, optional + Number of layers. Required for non-FLOW-JA-FACE budget terms. + Can be None for files containing only FLOW-JA-FACE data (default None). + nrow : int, optional + Number of rows. Required for non-FLOW-JA-FACE budget terms. + Can be None for files containing only FLOW-JA-FACE data (default None). + ncol : int, optional + Number of columns. Required for non-FLOW-JA-FACE budget terms. + Can be None for files containing only FLOW-JA-FACE data (default None). + precision : str, optional + 'single' or 'double' (default 'double') + verbose : bool, optional + Print progress messages (default False) + + Examples + -------- + Write FLOW-JA-FACE array data: + + >>> import numpy as np + >>> from flopy.utils.binaryfile import CellBudgetFile + >>> flowja = np.random.rand(3367) # Connection flows + >>> data = { + ... (1, 1): { + ... 'FLOW-JA-FACE': { + ... 'imeth': 1, + ... 'data': flowja, + ... 'delt': 1.0, + ... 'pertim': 1.0, + ... 'totim': 1.0 + ... } + ... } + ... } + >>> CellBudgetFile.write_budget('model.bud', data, nlay=1, nrow=40, ncol=20) + + Write DATA-SPDIS list data: + + >>> ncells = 800 + >>> active = np.ones(ncells, dtype=bool) + >>> dt = np.dtype([('node', np.int32), ('node2', np.int32), ('q', np.float64), + ... ('qx', np.float64), ('qy', np.float64), ('qz', np.float64)]) + >>> data_list = np.zeros(ncells, dtype=dt) + >>> data_list['node'] = np.arange(1, ncells+1) + >>> data_list['node2'] = np.arange(1, ncells+1) + >>> data_list['qx'] = np.random.rand(ncells) + >>> data_list['qy'] = np.random.rand(ncells) + >>> data_list['qz'] = np.random.rand(ncells) + >>> data = { + ... (1, 1): { + ... 'DATA-SPDIS': { + ... 'imeth': 6, + ... 'data': data_list, + ... 'delt': 1.0, + ... 'pertim': 1.0, + ... 'totim': 1.0, + ... 'modelnam': 'GWF', + ... 'paknam': '', + ... 'modelnam2': 'GWF', + ... 'paknam2': '', + ... 'ndat': 3, + ... 'auxtxt': ['qx', 'qy', 'qz'] + ... } + ... } + ... } + >>> CellBudgetFile.write_budget('model.bud', data, nlay=1, nrow=40, ncol=20) + """ + # Set precision + if precision == "single": + realtype = np.float32 + else: + realtype = np.float64 + + if verbose: + print(f"Writing binary budget file: {filename}") + print(f" Precision: {precision}") + if nlay is not None and nrow is not None and ncol is not None: + print(f" Grid shape: {nlay} layers x {nrow} rows x {ncol} cols") + else: + print(" Grid shape: not specified (OK for FLOW-JA-FACE only files)") + + with open(filename, "wb") as f: + # Sort by time step for consistent output + sorted_keys = sorted(budget_dict.keys()) + + for kstpkper in sorted_keys: + kstp, kper = kstpkper + time_data = budget_dict[kstpkper] + + if verbose: + print(f"\n Writing kstp={kstp}, kper={kper}") + + # Write each budget term for this time step + for text, term_data in time_data.items(): + imeth = term_data["imeth"] + data = term_data["data"] + delt = term_data.get("delt", 0.0) + pertim = term_data["pertim"] + totim = term_data["totim"] + + # Ensure text is exactly 16 bytes + text_bytes = text.encode("ascii") if isinstance(text, str) else text + if len(text_bytes) > 16: + text_bytes = text_bytes[:16] + elif len(text_bytes) < 16: + text_bytes = text_bytes + b" " * (16 - len(text_bytes)) + + if verbose: + print(f" Writing {text.strip()}: imeth={imeth}") + + # Check if this is FLOW-JA-FACE (connection-based) + is_flowja = text.strip().upper() in [ + "FLOW-JA-FACE", + "FLOW-JA-FACE-X", + ] + + # Write header1 + h1dt = np.dtype( + [ + ("kstp", np.int32), + ("kper", np.int32), + ("text", "S16"), + ("ncol", np.int32), + ("nrow", np.int32), + ("nlay", np.int32), + ] + ) + + # Determine dimensions based on data type + if is_flowja and imeth in [0, 1]: + # FLOW-JA-FACE: use NJA (size of connection array) + arr = np.asarray(data) + nja = arr.size + ndim1, ndim2, ndim3 = nja, 1, -1 + else: + # Regular budget term: use grid dimensions + if nlay is None or nrow is None or ncol is None: + raise ValueError( + f"Grid dimensions (nlay, nrow, ncol) required for " + f"non-FLOW-JA-FACE budget term '{text.strip()}'. " + f"Provided: nlay={nlay}, nrow={nrow}, ncol={ncol}" + ) + # Use negative nlay for compact format + ndim1, ndim2, ndim3 = ncol, nrow, -nlay + + header1 = np.array( + [(kstp, kper, text_bytes, ndim1, ndim2, ndim3)], dtype=h1dt + ) + header1.tofile(f) + + # Write header2 + h2dt = np.dtype( + [ + ("imeth", np.int32), + ("delt", realtype), + ("pertim", realtype), + ("totim", realtype), + ] + ) + header2 = np.array([(imeth, delt, pertim, totim)], dtype=h2dt) + header2.tofile(f) + + # For imeth=6, write model and package names + if imeth == 6: + modelnam = term_data.get("modelnam", "MODEL") + paknam = term_data.get("paknam", "") + modelnam2 = term_data.get("modelnam2", "MODEL") + paknam2 = term_data.get("paknam2", "") + + # Ensure each is exactly 16 bytes + for name in [modelnam, paknam, modelnam2, paknam2]: + name_bytes = ( + name.encode("ascii") if isinstance(name, str) else name + ) + if len(name_bytes) > 16: + name_bytes = name_bytes[:16] + elif len(name_bytes) < 16: + name_bytes = name_bytes + b" " * (16 - len(name_bytes)) + f.write(name_bytes) + + # Write data based on imeth + if imeth == 0 or imeth == 1: + # Array format + arr = np.asarray(data, dtype=realtype) + + # Check if this is FLOW-JA-FACE (connection-based) + is_flowja = text.strip().upper() in [ + "FLOW-JA-FACE", + "FLOW-JA-FACE-X", + ] + + if is_flowja: + # FLOW-JA-FACE: keep as 1D array of size NJA + # Don't reshape - connection-based not cell-based + if arr.ndim != 1: + arr = arr.flatten() + else: + # Regular budget term: reshape to grid if needed + if arr.ndim == 1: + # Flattened - reshape to (nlay, nrow, ncol) + arr = arr.reshape(nlay, nrow, ncol) + + arr.tofile(f) + + elif imeth == 6: + # List format + # Write naux+1 + ndat = term_data.get("ndat", 1) + auxtxt = term_data.get("auxtxt", []) + naux = len(auxtxt) + nauxp1 = naux + 1 + + np.array([nauxp1], dtype=np.int32).tofile(f) + + # Write auxiliary variable names + for auxname in auxtxt: + auxname_bytes = ( + auxname.encode("ascii") + if isinstance(auxname, str) + else auxname + ) + if len(auxname_bytes) > 16: + auxname_bytes = auxname_bytes[:16] + elif len(auxname_bytes) < 16: + auxname_bytes = auxname_bytes + b" " * ( + 16 - len(auxname_bytes) + ) + f.write(auxname_bytes) + + # Write nlist + nlist = len(data) + np.array([nlist], dtype=np.int32).tofile(f) + + # Write list data + # Data should be a recarray with fields: + # node, node2, q, and aux fields + if isinstance(data, np.ndarray) and data.dtype.names: + # It's already a structured array, write it directly + # But we need to ensure the dtypes match + dt_list = [ + ("node", np.int32), + ("node2", np.int32), + ("q", realtype), + ] + for auxname in auxtxt: + dt_list.append((auxname, realtype)) + + output_dt = np.dtype(dt_list) + output_data = np.zeros(nlist, dtype=output_dt) + + # Copy data from input to output with correct types + for field in output_dt.names: + if field in data.dtype.names: + output_data[field] = data[field].astype( + output_dt[field] + ) + + output_data.tofile(f) + else: + raise ValueError( + "For imeth=6, data must be a numpy recarray " + "with fields: node, node2, q, and optional " + "auxiliary fields" + ) + + else: + raise NotImplementedError( + f"imeth={imeth} not yet implemented. " + "Currently only imeth=1 (array) and imeth=6 " + "(list) are supported." + ) + + if verbose: + print(f"\nSuccessfully wrote {filename}") + def close(self): """ Close the file handle From 1e85d2b98592dac6ecdd310379766b6da3004fda Mon Sep 17 00:00:00 2001 From: Bonelli Date: Wed, 25 Feb 2026 11:38:41 -0500 Subject: [PATCH 02/14] tests for structured connectivity and grb write --- autotest/test_binarygrid_util.py | 378 +++++++++++++++++++++++++++++ flopy/mf6/utils/binarygrid_util.py | 49 ++-- flopy/utils/binaryfile/__init__.py | 25 +- 3 files changed, 425 insertions(+), 27 deletions(-) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index fb4f09aa22..f85286a828 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -159,3 +159,381 @@ 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_build_structured_connectivity_simple(): + """Test build_structured_connectivity with simple grids.""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + # Test 1x1x1 grid (single cell, only diagonal connection) + ia, ja, nja = build_structured_connectivity(1, 1, 1) + assert ia.shape == (2,), f"ia shape {ia.shape} != (2,)" + assert ja.shape == (1,), f"ja shape {ja.shape} != (1,)" + assert nja == 1, f"nja {nja} != 1" + assert ia[0] == 0 and ia[1] == 1, f"ia {ia} incorrect" + assert ja[0] == 0, f"ja {ja} != [0]" + + # Test 1x1x2 grid (2 cells in x, 1 connection between them + 2 diagonals = 3) + ia, ja, nja = build_structured_connectivity(1, 1, 2) + assert ia.shape == (3,), f"ia shape {ia.shape} != (3,)" + assert nja == 3, f"nja {nja} != 3" + # Cell 0: diagonal (0) + right neighbor (1) = 2 connections + # Cell 1: diagonal (1) = 1 connection + np.testing.assert_array_equal(ia, [0, 2, 3]) + np.testing.assert_array_equal(ja, [0, 1, 1]) + + # Test 1x2x1 grid (2 cells in y, 1 connection between them + 2 diagonals = 3) + ia, ja, nja = build_structured_connectivity(1, 2, 1) + assert ia.shape == (3,), f"ia shape {ia.shape} != (3,)" + assert nja == 3, f"nja {nja} != 3" + # Cell 0: diagonal (0) + front neighbor (1) = 2 connections + # Cell 1: diagonal (1) = 1 connection + np.testing.assert_array_equal(ia, [0, 2, 3]) + np.testing.assert_array_equal(ja, [0, 1, 1]) + + # Test 2x1x1 grid (2 layers, 1 connection between them + 2 diagonals = 3) + ia, ja, nja = build_structured_connectivity(2, 1, 1) + assert ia.shape == (3,), f"ia shape {ia.shape} != (3,)" + assert nja == 3, f"nja {nja} != 3" + # Cell 0 (layer 0): diagonal (0) + lower neighbor (1) = 2 connections + # Cell 1 (layer 1): diagonal (1) = 1 connection + np.testing.assert_array_equal(ia, [0, 2, 3]) + np.testing.assert_array_equal(ja, [0, 1, 1]) + + +def test_build_structured_connectivity_2x2x2(): + """Test build_structured_connectivity with a 2x2x2 grid.""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + nlay, nrow, ncol = 2, 2, 2 + ncells = nlay * nrow * ncol # 8 cells + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + # Verify dimensions + assert ia.shape == (ncells + 1,), f"ia shape {ia.shape} != {(ncells + 1,)}" + assert ja.shape == (nja,), f"ja shape {ja.shape} != {(nja,)}" + assert ia[-1] == nja, f"ia[-1] {ia[-1]} != nja {nja}" + + # Verify IA is monotonically increasing + assert np.all(np.diff(ia) >= 0), "IA array not monotonically increasing" + + # Verify all JA entries are valid cell indices + assert np.all(ja >= 0), "JA contains negative indices" + assert np.all(ja < ncells), f"JA contains indices >= {ncells}" + + # Verify each cell has at least a diagonal connection + for i in range(ncells): + nconn = ia[i + 1] - ia[i] + assert nconn >= 1, f"Cell {i} has {nconn} connections (should be >= 1)" + + # Verify diagonal entries + for node in range(ncells): + conn_start = ia[node] + # First connection should be diagonal (self) + assert ja[conn_start] == node, ( + f"Cell {node} first connection {ja[conn_start]} != {node}" + ) + + +def test_build_structured_connectivity_with_idomain(): + """Test build_structured_connectivity with inactive cells.""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + nlay, nrow, ncol = 1, 3, 3 # 9 cells + ncells = nlay * nrow * ncol + + # Make center cell inactive + idomain = np.ones((nlay, nrow, ncol), dtype=np.int32) + idomain[0, 1, 1] = 0 # Cell 4 (center) is inactive + + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol, idomain=idomain) + + assert ia.shape == (ncells + 1,) + assert ja.shape == (nja,) + + # Inactive cell (node 4) should have no connections + node_4_start = ia[4] + node_4_end = ia[5] + assert node_4_end == node_4_start, ( + f"Inactive cell 4 has {node_4_end - node_4_start} connections (should be 0)" + ) + + # Cells around the inactive cell should not connect to it + for node in range(ncells): + if node == 4: + continue # Skip inactive cell + conn_start = ia[node] + conn_end = ia[node + 1] + connections = ja[conn_start:conn_end] + assert 4 not in connections, ( + f"Cell {node} incorrectly connects to inactive cell 4" + ) + + +def test_build_structured_connectivity_known_values(): + """Test build_structured_connectivity against known values.""" + from flopy.mf6.utils.binarygrid_util import build_structured_connectivity + + # Simple 1x2x2 grid + # Cell layout: + # 2 3 + # 0 1 + # Expected connections: + # Cell 0: 0 (diag), 1 (right), 2 (front) -> ia[0]=0, ia[1]=3 + # Cell 1: 1 (diag), 3 (front) -> ia[1]=3, ia[2]=5 + # Cell 2: 2 (diag), 3 (right) -> ia[2]=5, ia[3]=7 + # Cell 3: 3 (diag) -> ia[3]=7, ia[4]=8 + + ia, ja, nja = build_structured_connectivity(1, 2, 2) + assert nja == 8, f"nja {nja} != 8" + np.testing.assert_array_equal(ia, [0, 3, 5, 7, 8]) + np.testing.assert_array_equal(ja, [0, 1, 2, 1, 3, 2, 3, 3]) + + +def test_write_grb_roundtrip(tmp_path): + """Test MfGrdFile.write_grb() with roundtrip validation.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile, build_structured_connectivity + + # Create a simple DIS grid + nlay, nrow, ncol = 1, 10, 20 + ncells = nlay * nrow * ncol + delr = np.full(ncol, 250.0) + delc = np.full(nrow, 250.0) + top = np.full(ncells, 35.0) + botm = np.full(ncells, 0.0) + idomain = np.ones(ncells, dtype=np.int32) + icelltype = np.zeros(ncells, dtype=np.int32) + + # Build connectivity + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + # Prepare data dictionary + data_dict = { + "NCELLS": ncells, + "NLAY": nlay, + "NROW": nrow, + "NCOL": ncol, + "NJA": nja, + "XORIGIN": 0.0, + "YORIGIN": 0.0, + "ANGROT": 0.0, + "DELR": delr, + "DELC": delc, + "TOP": top, + "BOTM": botm, + "IA": ia + 1, # Convert to 1-based for MF6 + "JA": ja + 1, # Convert to 1-based for MF6 + "IDOMAIN": idomain, + "ICELLTYPE": icelltype, + } + + # Write the file + output_file = tmp_path / "test.dis.grb" + MfGrdFile.write_grb( + output_file, "DIS", data_dict, precision="double", verbose=False + ) + + # Read it back + grb = MfGrdFile(output_file, verbose=False) + + # Verify basic properties + assert grb.grid_type == "DIS", f"Grid type {grb.grid_type} != DIS" + assert grb.nodes == ncells, f"nodes {grb.nodes} != {ncells}" + assert grb.nlay == nlay, f"nlay {grb.nlay} != {nlay}" + assert grb.nrow == nrow, f"nrow {grb.nrow} != {nrow}" + assert grb.ncol == ncol, f"ncol {grb.ncol} != {ncol}" + assert grb.nja == nja, f"nja {grb.nja} != {nja}" + + # Verify grid geometry + np.testing.assert_allclose(grb.xorigin, 0.0) + np.testing.assert_allclose(grb.yorigin, 0.0) + np.testing.assert_allclose(grb.angrot, 0.0) + + # Verify arrays + np.testing.assert_allclose(grb.delr, delr, err_msg="DELR mismatch") + np.testing.assert_allclose(grb.delc, delc, err_msg="DELC mismatch") + np.testing.assert_allclose(grb.top, top, err_msg="TOP mismatch") + np.testing.assert_allclose(grb.bot, botm, err_msg="BOTM mismatch") + + # Verify connectivity (grb.ia and grb.ja are already 0-based) + np.testing.assert_array_equal(grb.ia, ia, err_msg="IA mismatch") + np.testing.assert_array_equal(grb.ja, ja, err_msg="JA mismatch") + + # Verify cell properties + np.testing.assert_array_equal(grb.idomain, idomain, err_msg="IDOMAIN mismatch") + np.testing.assert_array_equal( + grb._datadict["ICELLTYPE"], icelltype, err_msg="ICELLTYPE mismatch" + ) + + +def test_write_grb_with_existing_grb(tmp_path, mfgrd_test_path): + """Test MfGrdFile.write_grb() by reading an existing file and writing it back.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read an existing GRB file + original_file = mfgrd_test_path / "nwtp3.dis.grb" + grb_orig = MfGrdFile(original_file, verbose=False) + + # Extract all data + # Note: IA and JA need to be 1-based for writing, so use _datadict + data_dict = { + "NCELLS": grb_orig.nodes, + "NLAY": grb_orig.nlay, + "NROW": grb_orig.nrow, + "NCOL": grb_orig.ncol, + "NJA": grb_orig.nja, + "XORIGIN": grb_orig.xorigin, + "YORIGIN": grb_orig.yorigin, + "ANGROT": grb_orig.angrot, + "DELR": grb_orig.delr, + "DELC": grb_orig.delc, + "TOP": grb_orig.top, + "BOTM": grb_orig.bot, + "IA": grb_orig._datadict["IA"], # Use 1-based from file + "JA": grb_orig._datadict["JA"], # Use 1-based from file + "IDOMAIN": grb_orig.idomain, + "ICELLTYPE": grb_orig._datadict.get( + "ICELLTYPE", np.zeros(grb_orig.nodes, dtype=np.int32) + ), + } + + # Write to a new file + output_file = tmp_path / "test_copy.dis.grb" + MfGrdFile.write_grb( + output_file, "DIS", data_dict, precision="double", verbose=False + ) + + # Read it back + grb_new = MfGrdFile(output_file, verbose=False) + + # Compare all properties + 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, err_msg="DELR mismatch") + np.testing.assert_allclose(grb_new.delc, grb_orig.delc, err_msg="DELC mismatch") + np.testing.assert_allclose(grb_new.top, grb_orig.top, err_msg="TOP mismatch") + np.testing.assert_allclose(grb_new.bot, grb_orig.bot, err_msg="BOTM mismatch") + + np.testing.assert_array_equal(grb_new.ia, grb_orig.ia, err_msg="IA mismatch") + np.testing.assert_array_equal(grb_new.ja, grb_orig.ja, err_msg="JA mismatch") + np.testing.assert_array_equal( + grb_new.idomain, grb_orig.idomain, err_msg="IDOMAIN mismatch" + ) + # Compare ICELLTYPE if it exists in both + if "ICELLTYPE" in grb_new._datadict and "ICELLTYPE" in grb_orig._datadict: + np.testing.assert_array_equal( + grb_new._datadict["ICELLTYPE"], + grb_orig._datadict["ICELLTYPE"], + err_msg="ICELLTYPE mismatch", + ) + + +def test_write_grb_modelgrid_roundtrip(tmp_path): + """Test that modelgrid can be reconstructed from written GRB file.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile, build_structured_connectivity + + # Create grid with specific geometry (use single layer for simpler modelgrid test) + nlay, nrow, ncol = 1, 5, 10 + ncells = nlay * nrow * ncol + delr = np.full(ncol, 100.0) + delc = np.full(nrow, 200.0) + top = np.full(ncells, 50.0) + botm = np.full(ncells, -50.0) + + # Build connectivity + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + # Prepare data with non-zero origin and rotation + data_dict = { + "NCELLS": ncells, + "NLAY": nlay, + "NROW": nrow, + "NCOL": ncol, + "NJA": nja, + "XORIGIN": 1000.0, + "YORIGIN": 2000.0, + "ANGROT": 15.0, + "DELR": delr, + "DELC": delc, + "TOP": top, + "BOTM": botm, + "IA": ia + 1, + "JA": ja + 1, + "IDOMAIN": np.ones(ncells, dtype=np.int32), + "ICELLTYPE": np.zeros(ncells, dtype=np.int32), + } + + # Write and read back + output_file = tmp_path / "test_grid.dis.grb" + MfGrdFile.write_grb( + output_file, "DIS", data_dict, precision="double", verbose=False + ) + grb = MfGrdFile(output_file, verbose=False) + + # Get modelgrid + modelgrid = grb.modelgrid + + # Verify modelgrid properties + assert isinstance(modelgrid, StructuredGrid) + assert modelgrid.nlay == nlay + assert modelgrid.nrow == nrow + assert modelgrid.ncol == ncol + np.testing.assert_allclose(modelgrid.xoffset, 1000.0, err_msg="xoffset mismatch") + np.testing.assert_allclose(modelgrid.yoffset, 2000.0, err_msg="yoffset mismatch") + np.testing.assert_allclose(modelgrid.angrot, 15.0, err_msg="angrot mismatch") + + +def test_write_grb_precision(tmp_path): + """Test MfGrdFile.write_grb() with single and double precision.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile, build_structured_connectivity + + nlay, nrow, ncol = 1, 5, 5 + ncells = nlay * nrow * ncol + ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) + + data_dict = { + "NCELLS": ncells, + "NLAY": nlay, + "NROW": nrow, + "NCOL": ncol, + "NJA": nja, + "XORIGIN": 0.0, + "YORIGIN": 0.0, + "ANGROT": 0.0, + "DELR": np.full(ncol, 100.0), + "DELC": np.full(nrow, 100.0), + "TOP": np.full(ncells, 10.0), + "BOTM": np.full(ncells, 0.0), + "IA": ia + 1, + "JA": ja + 1, + "IDOMAIN": np.ones(ncells, dtype=np.int32), + "ICELLTYPE": np.zeros(ncells, dtype=np.int32), + } + + # Test single precision + single_file = tmp_path / "test_single.grb" + MfGrdFile.write_grb( + single_file, "DIS", data_dict, precision="single", verbose=False + ) + grb_single = MfGrdFile(single_file, verbose=False) + assert grb_single.nodes == ncells + + # Test double precision + double_file = tmp_path / "test_double.grb" + MfGrdFile.write_grb( + double_file, "DIS", data_dict, precision="double", verbose=False + ) + grb_double = MfGrdFile(double_file, verbose=False) + assert grb_double.nodes == ncells + + # Single precision file should be smaller + assert single_file.stat().st_size < double_file.stat().st_size diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index e7e049b21f..a3cc602538 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -822,6 +822,9 @@ def write_grb( writer.precision = precision # Define variable metadata based on grid type + # Use precision parameter to determine floating point type + float_type = "SINGLE" if precision.lower() == "single" else "DOUBLE" + if grid_type.upper() == "DIS": var_list = [ ("NCELLS", "INTEGER", 0, []), @@ -829,13 +832,13 @@ def write_grb( ("NROW", "INTEGER", 0, []), ("NCOL", "INTEGER", 0, []), ("NJA", "INTEGER", 0, []), - ("XORIGIN", "DOUBLE", 0, []), - ("YORIGIN", "DOUBLE", 0, []), - ("ANGROT", "DOUBLE", 0, []), - ("DELR", "DOUBLE", 1, [data_dict.get("NCOL", 0)]), - ("DELC", "DOUBLE", 1, [data_dict.get("NROW", 0)]), - ("TOP", "DOUBLE", 1, [data_dict.get("NCELLS", 0)]), - ("BOTM", "DOUBLE", 1, [data_dict.get("NCELLS", 0)]), + ("XORIGIN", float_type, 0, []), + ("YORIGIN", float_type, 0, []), + ("ANGROT", float_type, 0, []), + ("DELR", float_type, 1, [data_dict.get("NCOL", 0)]), + ("DELC", float_type, 1, [data_dict.get("NROW", 0)]), + ("TOP", float_type, 1, [data_dict.get("NCELLS", 0)]), + ("BOTM", float_type, 1, [data_dict.get("NCELLS", 0)]), ("IA", "INTEGER", 1, [data_dict.get("NCELLS", 0) + 1]), ("JA", "INTEGER", 1, [data_dict.get("NJA", 0)]), ("IDOMAIN", "INTEGER", 1, [data_dict.get("NCELLS", 0)]), @@ -856,15 +859,23 @@ def write_grb( print(f" Version: {version}") print(f" Number of variables: {ntxt}") - with open(filename, "wb") as f: - writer.file = f + # Helper function to write text with fixed width + def write_text(f, text, width): + """Write text padded to fixed width.""" + text_bytes = text.encode("ascii") + if len(text_bytes) > width: + text_bytes = text_bytes[:width] + else: + text_bytes = text_bytes.ljust(width) + f.write(text_bytes) + with open(filename, "wb") as f: # Write text header lines (50 chars each, newline terminated) header_len = 50 - writer.write_text(f"GRID {grid_type.upper()}\n", header_len) - writer.write_text(f"VERSION {version}\n", header_len) - writer.write_text(f"NTXT {ntxt}\n", header_len) - writer.write_text(f"LENTXT {lentxt}\n", header_len) + write_text(f, f"GRID {grid_type.upper()}\n", header_len) + write_text(f, f"VERSION {version}\n", header_len) + write_text(f, f"NTXT {ntxt}\n", header_len) + write_text(f, f"LENTXT {lentxt}\n", header_len) # Write variable definition lines (100 chars each) for name, dtype_str, ndim, dims in var_list: @@ -875,7 +886,7 @@ def write_grb( str(d) for d in dims[::-1] ) # Reverse for Fortran order line = f"{name} {dtype_str} NDIM {ndim} {dims_str}\n" - writer.write_text(line, lentxt) + write_text(f, line, lentxt) # Write binary data for each variable for name, dtype_str, ndim, dims in var_list: @@ -901,9 +912,11 @@ def write_grb( if ndim == 0: # Scalar value if dtype_str == "INTEGER": - writer.write_integer(int(value)) - elif dtype_str in ("DOUBLE", "SINGLE"): - writer.write_real(float(value)) + f.write(np.array(int(value), dtype=np.int32).tobytes()) + elif dtype_str == "DOUBLE": + f.write(np.array(float(value), dtype=np.float64).tobytes()) + elif dtype_str == "SINGLE": + f.write(np.array(float(value), dtype=np.float32).tobytes()) else: # Array data arr = np.asarray(value) @@ -915,7 +928,7 @@ def write_grb( arr = arr.astype(np.float32) # Write array in column-major (Fortran) order - writer.write_record(arr, dtype=arr.dtype) + f.write(arr.flatten(order="F").tobytes()) if verbose: print(f"Successfully wrote {filename}") diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 776e00049f..6b25f68c99 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -685,7 +685,9 @@ def write( ksp_tuple = (kstp, kper) # Find the totim for this kstpkper - mask = (self.recordarray["kstp"] == kstp) & (self.recordarray["kper"] == kper) + mask = (self.recordarray["kstp"] == kstp) & ( + self.recordarray["kper"] == kper + ) matching_records = self.recordarray[mask] if len(matching_records) == 0: if kwargs.get("verbose", False): @@ -2544,11 +2546,14 @@ def write( for txt in textlist: # Get matching text from file (case-insensitive, padded) - txt_str = txt.decode().strip() if isinstance(txt, bytes) else txt.strip() + txt_str = ( + txt.decode().strip() if isinstance(txt, bytes) else txt.strip() + ) # Find matching records matching_records = [ - t for t in self.textlist + t + for t in self.textlist if txt_str.upper() in t.decode().strip().upper() ] @@ -2587,12 +2592,14 @@ def write( # Add model/package names if imeth=6 if imeth == 6: - budget_dict[key].update({ - "modelnam": record["modelnam"].decode().strip(), - "paknam": record["paknam"].decode().strip(), - "modelnam2": record["modelnam2"].decode().strip(), - "paknam2": record["paknam2"].decode().strip(), - }) + budget_dict[key].update( + { + "modelnam": record["modelnam"].decode().strip(), + "paknam": record["paknam"].decode().strip(), + "modelnam2": record["modelnam2"].decode().strip(), + "paknam2": record["paknam2"].decode().strip(), + } + ) except Exception as e: if kwargs.get("verbose", False): From 9c435e6685027023ed73679397279f6f5ef9e838 Mon Sep 17 00:00:00 2001 From: Bonelli Date: Wed, 25 Feb 2026 18:19:23 -0500 Subject: [PATCH 03/14] fix init file --- flopy/mf6/utils/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flopy/mf6/utils/__init__.py b/flopy/mf6/utils/__init__.py index f442647a55..476f0602df 100644 --- a/flopy/mf6/utils/__init__.py +++ b/flopy/mf6/utils/__init__.py @@ -3,3 +3,4 @@ from .lakpak_utils import get_lak_connections from .mfsimlistfile import MfSimulationList from .model_splitter import Mf6Splitter +from .postprocessing import get_residuals, get_structured_faceflows From 3db060b6c5e677d0019eb55ad4405f3c1285676f Mon Sep 17 00:00:00 2001 From: Bonelli Date: Wed, 25 Feb 2026 18:29:46 -0500 Subject: [PATCH 04/14] grb instance write --- autotest/test_binarygrid_util.py | 65 ++++++++++++++++++++++++++++++ flopy/mf6/utils/binarygrid_util.py | 54 +++++++++++++++++++++++++ 2 files changed, 119 insertions(+) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index f85286a828..362b56f1b8 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -537,3 +537,68 @@ def test_write_grb_precision(tmp_path): # Single precision file should be smaller assert single_file.stat().st_size < double_file.stat().st_size + + +def test_write_grb_instance_method(tmp_path, mfgrd_test_path): + """Test MfGrdFile.write() instance method.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read an existing GRB file + original_file = mfgrd_test_path / "nwtp3.dis.grb" + grb_orig = MfGrdFile(original_file, verbose=False) + + # Write using instance method + output_file = tmp_path / "test_instance.dis.grb" + grb_orig.write(output_file, verbose=False) + + # Read it back + grb_new = MfGrdFile(output_file, verbose=False) + + # Verify all properties match + 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): + """Test MfGrdFile.write() with precision conversion.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read an existing GRB file (presumably double precision) + original_file = mfgrd_test_path / "nwtp3.dis.grb" + grb = MfGrdFile(original_file, verbose=False) + + # Write as single precision + single_file = tmp_path / "test_single.grb" + grb.write(single_file, precision="single", verbose=False) + + # Write as double precision + double_file = tmp_path / "test_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 both work and have same basic properties + assert grb_single.nodes == grb.nodes + assert grb_double.nodes == grb.nodes + + # Single precision file should be smaller + assert single_file.stat().st_size < double_file.stat().st_size diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index a3cc602538..63c50bae50 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -747,6 +747,60 @@ def cell2d(self): vertices, cell2d = None, None return vertices, cell2d + def write(self, filename, precision=None, verbose=False): + """ + Write the binary grid file to a new file. + + Parameters + ---------- + filename : str or PathLike + Path to output .grb file + precision : str, optional + 'single' or 'double'. If None, uses the precision from the + original file (default None) + verbose : bool, optional + Print progress messages (default False) + + Examples + -------- + >>> from flopy.mf6.utils import MfGrdFile + >>> grb = MfGrdFile('model.dis.grb') + >>> grb.write('model_copy.dis.grb') + """ + if precision is None: + precision = self.precision + + # Extract all data from this instance + data_dict = { + "NCELLS": self.nodes, + "NLAY": self.nlay, + "NROW": self.nrow, + "NCOL": self.ncol, + "NJA": self.nja, + "XORIGIN": self.xorigin, + "YORIGIN": self.yorigin, + "ANGROT": self.angrot, + "DELR": self.delr, + "DELC": self.delc, + "TOP": self.top, + "BOTM": self.bot, + "IA": self._datadict["IA"], # Use 1-based from original file + "JA": self._datadict["JA"], # Use 1-based from original file + "IDOMAIN": self.idomain, + } + + # Add ICELLTYPE if it exists + if "ICELLTYPE" in self._datadict: + data_dict["ICELLTYPE"] = self._datadict["ICELLTYPE"] + else: + # Provide default if not in original file + data_dict["ICELLTYPE"] = np.zeros(self.nodes, dtype=np.int32) + + # Call static method + MfGrdFile.write_grb( + filename, self.grid_type, data_dict, precision=precision, verbose=verbose + ) + @staticmethod def write_grb( filename, From fecb00e68b39fa9cecdbb45919867cea555e5b2f Mon Sep 17 00:00:00 2001 From: Bonelli Date: Wed, 25 Feb 2026 18:49:33 -0500 Subject: [PATCH 05/14] ruff --- autotest/test_binarygrid_util.py | 249 ----------------------------- flopy/mf6/utils/binarygrid_util.py | 151 ++++------------- flopy/utils/binaryfile/__init__.py | 4 +- 3 files changed, 32 insertions(+), 372 deletions(-) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index 362b56f1b8..e5631bd0ad 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -290,255 +290,6 @@ def test_build_structured_connectivity_known_values(): np.testing.assert_array_equal(ja, [0, 1, 2, 1, 3, 2, 3, 3]) -def test_write_grb_roundtrip(tmp_path): - """Test MfGrdFile.write_grb() with roundtrip validation.""" - from flopy.mf6.utils.binarygrid_util import MfGrdFile, build_structured_connectivity - - # Create a simple DIS grid - nlay, nrow, ncol = 1, 10, 20 - ncells = nlay * nrow * ncol - delr = np.full(ncol, 250.0) - delc = np.full(nrow, 250.0) - top = np.full(ncells, 35.0) - botm = np.full(ncells, 0.0) - idomain = np.ones(ncells, dtype=np.int32) - icelltype = np.zeros(ncells, dtype=np.int32) - - # Build connectivity - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - # Prepare data dictionary - data_dict = { - "NCELLS": ncells, - "NLAY": nlay, - "NROW": nrow, - "NCOL": ncol, - "NJA": nja, - "XORIGIN": 0.0, - "YORIGIN": 0.0, - "ANGROT": 0.0, - "DELR": delr, - "DELC": delc, - "TOP": top, - "BOTM": botm, - "IA": ia + 1, # Convert to 1-based for MF6 - "JA": ja + 1, # Convert to 1-based for MF6 - "IDOMAIN": idomain, - "ICELLTYPE": icelltype, - } - - # Write the file - output_file = tmp_path / "test.dis.grb" - MfGrdFile.write_grb( - output_file, "DIS", data_dict, precision="double", verbose=False - ) - - # Read it back - grb = MfGrdFile(output_file, verbose=False) - - # Verify basic properties - assert grb.grid_type == "DIS", f"Grid type {grb.grid_type} != DIS" - assert grb.nodes == ncells, f"nodes {grb.nodes} != {ncells}" - assert grb.nlay == nlay, f"nlay {grb.nlay} != {nlay}" - assert grb.nrow == nrow, f"nrow {grb.nrow} != {nrow}" - assert grb.ncol == ncol, f"ncol {grb.ncol} != {ncol}" - assert grb.nja == nja, f"nja {grb.nja} != {nja}" - - # Verify grid geometry - np.testing.assert_allclose(grb.xorigin, 0.0) - np.testing.assert_allclose(grb.yorigin, 0.0) - np.testing.assert_allclose(grb.angrot, 0.0) - - # Verify arrays - np.testing.assert_allclose(grb.delr, delr, err_msg="DELR mismatch") - np.testing.assert_allclose(grb.delc, delc, err_msg="DELC mismatch") - np.testing.assert_allclose(grb.top, top, err_msg="TOP mismatch") - np.testing.assert_allclose(grb.bot, botm, err_msg="BOTM mismatch") - - # Verify connectivity (grb.ia and grb.ja are already 0-based) - np.testing.assert_array_equal(grb.ia, ia, err_msg="IA mismatch") - np.testing.assert_array_equal(grb.ja, ja, err_msg="JA mismatch") - - # Verify cell properties - np.testing.assert_array_equal(grb.idomain, idomain, err_msg="IDOMAIN mismatch") - np.testing.assert_array_equal( - grb._datadict["ICELLTYPE"], icelltype, err_msg="ICELLTYPE mismatch" - ) - - -def test_write_grb_with_existing_grb(tmp_path, mfgrd_test_path): - """Test MfGrdFile.write_grb() by reading an existing file and writing it back.""" - from flopy.mf6.utils.binarygrid_util import MfGrdFile - - # Read an existing GRB file - original_file = mfgrd_test_path / "nwtp3.dis.grb" - grb_orig = MfGrdFile(original_file, verbose=False) - - # Extract all data - # Note: IA and JA need to be 1-based for writing, so use _datadict - data_dict = { - "NCELLS": grb_orig.nodes, - "NLAY": grb_orig.nlay, - "NROW": grb_orig.nrow, - "NCOL": grb_orig.ncol, - "NJA": grb_orig.nja, - "XORIGIN": grb_orig.xorigin, - "YORIGIN": grb_orig.yorigin, - "ANGROT": grb_orig.angrot, - "DELR": grb_orig.delr, - "DELC": grb_orig.delc, - "TOP": grb_orig.top, - "BOTM": grb_orig.bot, - "IA": grb_orig._datadict["IA"], # Use 1-based from file - "JA": grb_orig._datadict["JA"], # Use 1-based from file - "IDOMAIN": grb_orig.idomain, - "ICELLTYPE": grb_orig._datadict.get( - "ICELLTYPE", np.zeros(grb_orig.nodes, dtype=np.int32) - ), - } - - # Write to a new file - output_file = tmp_path / "test_copy.dis.grb" - MfGrdFile.write_grb( - output_file, "DIS", data_dict, precision="double", verbose=False - ) - - # Read it back - grb_new = MfGrdFile(output_file, verbose=False) - - # Compare all properties - 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, err_msg="DELR mismatch") - np.testing.assert_allclose(grb_new.delc, grb_orig.delc, err_msg="DELC mismatch") - np.testing.assert_allclose(grb_new.top, grb_orig.top, err_msg="TOP mismatch") - np.testing.assert_allclose(grb_new.bot, grb_orig.bot, err_msg="BOTM mismatch") - - np.testing.assert_array_equal(grb_new.ia, grb_orig.ia, err_msg="IA mismatch") - np.testing.assert_array_equal(grb_new.ja, grb_orig.ja, err_msg="JA mismatch") - np.testing.assert_array_equal( - grb_new.idomain, grb_orig.idomain, err_msg="IDOMAIN mismatch" - ) - # Compare ICELLTYPE if it exists in both - if "ICELLTYPE" in grb_new._datadict and "ICELLTYPE" in grb_orig._datadict: - np.testing.assert_array_equal( - grb_new._datadict["ICELLTYPE"], - grb_orig._datadict["ICELLTYPE"], - err_msg="ICELLTYPE mismatch", - ) - - -def test_write_grb_modelgrid_roundtrip(tmp_path): - """Test that modelgrid can be reconstructed from written GRB file.""" - from flopy.mf6.utils.binarygrid_util import MfGrdFile, build_structured_connectivity - - # Create grid with specific geometry (use single layer for simpler modelgrid test) - nlay, nrow, ncol = 1, 5, 10 - ncells = nlay * nrow * ncol - delr = np.full(ncol, 100.0) - delc = np.full(nrow, 200.0) - top = np.full(ncells, 50.0) - botm = np.full(ncells, -50.0) - - # Build connectivity - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - # Prepare data with non-zero origin and rotation - data_dict = { - "NCELLS": ncells, - "NLAY": nlay, - "NROW": nrow, - "NCOL": ncol, - "NJA": nja, - "XORIGIN": 1000.0, - "YORIGIN": 2000.0, - "ANGROT": 15.0, - "DELR": delr, - "DELC": delc, - "TOP": top, - "BOTM": botm, - "IA": ia + 1, - "JA": ja + 1, - "IDOMAIN": np.ones(ncells, dtype=np.int32), - "ICELLTYPE": np.zeros(ncells, dtype=np.int32), - } - - # Write and read back - output_file = tmp_path / "test_grid.dis.grb" - MfGrdFile.write_grb( - output_file, "DIS", data_dict, precision="double", verbose=False - ) - grb = MfGrdFile(output_file, verbose=False) - - # Get modelgrid - modelgrid = grb.modelgrid - - # Verify modelgrid properties - assert isinstance(modelgrid, StructuredGrid) - assert modelgrid.nlay == nlay - assert modelgrid.nrow == nrow - assert modelgrid.ncol == ncol - np.testing.assert_allclose(modelgrid.xoffset, 1000.0, err_msg="xoffset mismatch") - np.testing.assert_allclose(modelgrid.yoffset, 2000.0, err_msg="yoffset mismatch") - np.testing.assert_allclose(modelgrid.angrot, 15.0, err_msg="angrot mismatch") - - -def test_write_grb_precision(tmp_path): - """Test MfGrdFile.write_grb() with single and double precision.""" - from flopy.mf6.utils.binarygrid_util import MfGrdFile, build_structured_connectivity - - nlay, nrow, ncol = 1, 5, 5 - ncells = nlay * nrow * ncol - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - data_dict = { - "NCELLS": ncells, - "NLAY": nlay, - "NROW": nrow, - "NCOL": ncol, - "NJA": nja, - "XORIGIN": 0.0, - "YORIGIN": 0.0, - "ANGROT": 0.0, - "DELR": np.full(ncol, 100.0), - "DELC": np.full(nrow, 100.0), - "TOP": np.full(ncells, 10.0), - "BOTM": np.full(ncells, 0.0), - "IA": ia + 1, - "JA": ja + 1, - "IDOMAIN": np.ones(ncells, dtype=np.int32), - "ICELLTYPE": np.zeros(ncells, dtype=np.int32), - } - - # Test single precision - single_file = tmp_path / "test_single.grb" - MfGrdFile.write_grb( - single_file, "DIS", data_dict, precision="single", verbose=False - ) - grb_single = MfGrdFile(single_file, verbose=False) - assert grb_single.nodes == ncells - - # Test double precision - double_file = tmp_path / "test_double.grb" - MfGrdFile.write_grb( - double_file, "DIS", data_dict, precision="double", verbose=False - ) - grb_double = MfGrdFile(double_file, verbose=False) - assert grb_double.nodes == ncells - - # Single precision file should be smaller - assert single_file.stat().st_size < double_file.stat().st_size - - def test_write_grb_instance_method(tmp_path, mfgrd_test_path): """Test MfGrdFile.write() instance method.""" from flopy.mf6.utils.binarygrid_util import MfGrdFile diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index 63c50bae50..52b02bb040 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -747,7 +747,7 @@ def cell2d(self): vertices, cell2d = None, None return vertices, cell2d - def write(self, filename, precision=None, verbose=False): + def write(self, filename, precision=None, version=1, verbose=False): """ Write the binary grid file to a new file. @@ -758,6 +758,8 @@ def write(self, filename, precision=None, verbose=False): precision : str, optional 'single' or 'double'. If None, uses the precision from the original file (default None) + version : int, optional + Grid file version (default 1) verbose : bool, optional Print progress messages (default False) @@ -766,120 +768,29 @@ def write(self, filename, precision=None, verbose=False): >>> from flopy.mf6.utils import MfGrdFile >>> grb = MfGrdFile('model.dis.grb') >>> grb.write('model_copy.dis.grb') + >>> # Convert to single precision + >>> grb.write('model_single.dis.grb', precision='single') """ if precision is None: precision = self.precision - # Extract all data from this instance - data_dict = { - "NCELLS": self.nodes, - "NLAY": self.nlay, - "NROW": self.nrow, - "NCOL": self.ncol, - "NJA": self.nja, - "XORIGIN": self.xorigin, - "YORIGIN": self.yorigin, - "ANGROT": self.angrot, - "DELR": self.delr, - "DELC": self.delc, - "TOP": self.top, - "BOTM": self.bot, - "IA": self._datadict["IA"], # Use 1-based from original file - "JA": self._datadict["JA"], # Use 1-based from original file - "IDOMAIN": self.idomain, - } - - # Add ICELLTYPE if it exists - if "ICELLTYPE" in self._datadict: - data_dict["ICELLTYPE"] = self._datadict["ICELLTYPE"] - else: - # Provide default if not in original file - data_dict["ICELLTYPE"] = np.zeros(self.nodes, dtype=np.int32) - - # Call static method - MfGrdFile.write_grb( - filename, self.grid_type, data_dict, precision=precision, verbose=verbose - ) - - @staticmethod - def write_grb( - filename, - grid_type, - data_dict, - version=1, - precision="double", - verbose=False, - ): - """ - Write a MODFLOW 6 binary grid file (.grb). - - Parameters - ---------- - filename : str or PathLike - Path to output .grb file - grid_type : str - Grid type: 'DIS', 'DISV', or 'DISU' - data_dict : dict - Dictionary with grid data arrays. Required keys depend on grid_type. - For DIS grids: NCELLS, NLAY, NROW, NCOL, NJA, XORIGIN, YORIGIN, ANGROT, - DELR, DELC, TOP, BOTM, IA, JA, IDOMAIN, ICELLTYPE - version : int, optional - Grid file version (default 1) - precision : str, optional - 'single' or 'double' (default 'double') - verbose : bool, optional - Print progress messages (default False) - - Notes - ----- - The binary grid file format consists of: - 1. Text header lines (50 chars each): - - "GRID {grid_type}" - - "VERSION {version}" - - "NTXT {ntxt}" - - "LENTXT {lentxt}" - 2. Variable definition lines (100 chars each): - - "{NAME} {TYPE} NDIM {ndim} {dimensions...}" - 3. Binary data for each variable - - Arrays should be in Python (row-major) order and will be written - in Fortran (column-major) order as required by MF6. - - Examples - -------- - >>> import numpy as np - >>> from flopy.mf6.utils import MfGrdFile - >>> data = { - ... 'NCELLS': 800, - ... 'NLAY': 1, - ... 'NROW': 40, - ... 'NCOL': 20, - ... 'NJA': 3367, - ... 'XORIGIN': 0.0, - ... 'YORIGIN': 0.0, - ... 'ANGROT': 0.0, - ... 'DELR': np.full(20, 250.0), - ... 'DELC': np.full(40, 250.0), - ... 'TOP': np.full(800, 35.0), - ... 'BOTM': np.random.rand(800), - ... 'IA': ia_array, - ... 'JA': ja_array, - ... 'IDOMAIN': np.ones(800, dtype=np.int32), - ... 'ICELLTYPE': np.zeros(800, dtype=np.int32) - ... } - >>> MfGrdFile.write_grb('model.dis.grb', 'DIS', data) - """ - import os - - # Create FlopyBinaryData instance for write helpers - writer = FlopyBinaryData() - writer.precision = precision + # Build data dictionary from instance + data_dict = {} + for key in self._recordkeys: + if key in ("IA", "JA"): + # Use original 1-based arrays + data_dict[key] = self._datadict[key] + elif key == "TOP": + data_dict[key] = self.top + elif key == "BOTM": + data_dict[key] = self.bot + elif key in self._datadict: + data_dict[key] = self._datadict[key] # Define variable metadata based on grid type - # Use precision parameter to determine floating point type float_type = "SINGLE" if precision.lower() == "single" else "DOUBLE" - if grid_type.upper() == "DIS": + if self.grid_type == "DIS": var_list = [ ("NCELLS", "INTEGER", 0, []), ("NLAY", "INTEGER", 0, []), @@ -889,18 +800,18 @@ def write_grb( ("XORIGIN", float_type, 0, []), ("YORIGIN", float_type, 0, []), ("ANGROT", float_type, 0, []), - ("DELR", float_type, 1, [data_dict.get("NCOL", 0)]), - ("DELC", float_type, 1, [data_dict.get("NROW", 0)]), - ("TOP", float_type, 1, [data_dict.get("NCELLS", 0)]), - ("BOTM", float_type, 1, [data_dict.get("NCELLS", 0)]), - ("IA", "INTEGER", 1, [data_dict.get("NCELLS", 0) + 1]), - ("JA", "INTEGER", 1, [data_dict.get("NJA", 0)]), - ("IDOMAIN", "INTEGER", 1, [data_dict.get("NCELLS", 0)]), - ("ICELLTYPE", "INTEGER", 1, [data_dict.get("NCELLS", 0)]), + ("DELR", float_type, 1, [self.ncol]), + ("DELC", float_type, 1, [self.nrow]), + ("TOP", float_type, 1, [self.nodes]), + ("BOTM", float_type, 1, [self.nodes]), + ("IA", "INTEGER", 1, [self.nodes + 1]), + ("JA", "INTEGER", 1, [self.nja]), + ("IDOMAIN", "INTEGER", 1, [self.nodes]), + ("ICELLTYPE", "INTEGER", 1, [self.nodes]), ] else: raise NotImplementedError( - f"Grid type {grid_type} not yet implemented. " + f"Grid type {self.grid_type} not yet implemented. " "Currently only DIS grids are supported." ) @@ -909,7 +820,7 @@ def write_grb( if verbose: print(f"Writing binary grid file: {filename}") - print(f" Grid type: {grid_type}") + print(f" Grid type: {self.grid_type}") print(f" Version: {version}") print(f" Number of variables: {ntxt}") @@ -926,7 +837,7 @@ def write_text(f, text, width): with open(filename, "wb") as f: # Write text header lines (50 chars each, newline terminated) header_len = 50 - write_text(f, f"GRID {grid_type.upper()}\n", header_len) + write_text(f, f"GRID {self.grid_type}\n", header_len) write_text(f, f"VERSION {version}\n", header_len) write_text(f, f"NTXT {ntxt}\n", header_len) write_text(f, f"LENTXT {lentxt}\n", header_len) @@ -945,9 +856,7 @@ def write_text(f, text, width): # Write binary data for each variable for name, dtype_str, ndim, dims in var_list: if name not in data_dict: - raise ValueError( - f"Required variable '{name}' not found in data_dict" - ) + raise ValueError(f"Required variable '{name}' not found in grid file") value = data_dict[name] diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 6b25f68c99..58454b6462 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -660,7 +660,7 @@ def write( **kwargs Additional keyword arguments passed to write_head(): - text : str, identifier for head data (default uses current file's text) - - precision : str, 'single' or 'double' (default uses current file's precision) + - precision : str, 'single' or 'double' (default is the file's precision) - verbose : bool, print progress messages Examples @@ -2506,7 +2506,7 @@ def write( Examples: 'FLOW-JA-FACE', ['STORAGE', 'CONSTANT HEAD'] **kwargs Additional keyword arguments passed to write_budget(): - - precision : str, 'single' or 'double' (default uses current file's precision) + - precision : str, 'single' or 'double' (default is the file's precision) - verbose : bool, print progress messages Examples From 86d0b517eb224a16984c06f4aab5bb0edc202eef Mon Sep 17 00:00:00 2001 From: Bonelli Date: Thu, 26 Feb 2026 07:03:11 -0500 Subject: [PATCH 06/14] just instance methods --- autotest/test_binaryfile.py | 478 ----------------------------- flopy/utils/binaryfile/__init__.py | 203 +----------- 2 files changed, 12 insertions(+), 669 deletions(-) diff --git a/autotest/test_binaryfile.py b/autotest/test_binaryfile.py index f02068c08c..d6b73d8dd3 100644 --- a/autotest/test_binaryfile.py +++ b/autotest/test_binaryfile.py @@ -777,481 +777,3 @@ def test_headfile_get_ts_disu_grid(dis_sim, function_tmpdir): ) -def test_headfile_write_roundtrip(tmp_path, example_data_path): - """Test HeadFile.write_head() with roundtrip validation.""" - # Read an existing head file - original_file = example_data_path / "mf6-freyberg" / "freyberg.hds" - original = HeadFile(original_file) - - # Extract all data - data_dict = {} - for idx, kstpkper_np in enumerate(original.kstpkper): - kstpkper = (int(kstpkper_np[0]), int(kstpkper_np[1])) - head = original.get_data(idx=idx) - totim = original.times[idx] - data_dict[kstpkper] = {"head": head, "pertim": totim, "totim": totim} - - # Write to a temp file - output_file = tmp_path / "test.hds" - HeadFile.write_head( - output_file, data_dict, text="HEAD", precision="double", verbose=False - ) - - # Read it back - readback = HeadFile(output_file) - - # Compare - assert len(readback.kstpkper) == len(original.kstpkper), ( - f"Number of time steps mismatch: " - f"{len(readback.kstpkper)} != {len(original.kstpkper)}" - ) - - for idx in range(len(original.kstpkper)): - orig_data = original.get_data(idx=idx) - read_data = readback.get_data(idx=idx) - np.testing.assert_allclose( - orig_data, read_data, err_msg=f"Head data mismatch at idx={idx}" - ) - - -def test_headfile_write_synthetic(tmp_path): - """Test HeadFile.write_head() with synthetic data.""" - nlay, nrow, ncol = 2, 10, 15 - - # Create synthetic head data for multiple time steps - data_dict = {} - for kper in range(1, 3): - for kstp in range(1, 3): - heads = np.random.rand(nlay, nrow, ncol) * 10.0 + 100.0 - totim = (kper - 1) * 10.0 + kstp - data_dict[(kstp, kper)] = { - "head": heads, - "pertim": float(kstp), - "totim": totim, - } - - # Write the file - output_file = tmp_path / "synthetic.hds" - HeadFile.write_head(output_file, data_dict, text="HEAD", precision="single") - - # Read it back - hf = HeadFile(output_file, precision="single") - - assert len(hf.kstpkper) == 4, f"Expected 4 time steps, got {len(hf.kstpkper)}" - - # Verify data roundtrip - for (kstp, kper), entry in data_dict.items(): - data_back = hf.get_data(kstpkper=(kstp - 1, kper - 1)) - np.testing.assert_allclose( - entry["head"], data_back, err_msg=f"Mismatch at kstp={kstp}, kper={kper}" - ) - - -def test_cellbudgetfile_write_array_format(tmp_path): - """Test CellBudgetFile.write_budget() with array format (imeth=1).""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - - nlay, nrow, ncol = 1, 10, 20 - ncells = nlay * nrow * ncol - - # Build connectivity for this grid to get NJA - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - # Create synthetic FLOW-JA-FACE data (size NJA, not ncells) - flowja = np.random.rand(nja) - 0.5 - - budget_dict = { - (1, 1): { - "FLOW-JA-FACE": { - "imeth": 1, - "data": flowja, - "delt": 1.0, - "pertim": 10.0, - "totim": 10.0, - } - } - } - - # Write the budget file - output_file = tmp_path / "test_array.bud" - CellBudgetFile.write_budget( - output_file, - budget_dict, - nlay=nlay, - nrow=nrow, - ncol=ncol, - precision="double", - verbose=False, - ) - - # Read it back - cbb = CellBudgetFile(output_file, precision="double") - - assert len(cbb.textlist) == 1, f"Expected 1 budget term, got {len(cbb.textlist)}" - # Text is padded to 16 characters in binary format - assert "FLOW-JA-FACE".ljust(16).encode() in cbb.textlist, ( - "FLOW-JA-FACE not in budget file" - ) - - # Get data back - data_back = cbb.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] - - # FLOW-JA-FACE is returned as 3D array with shape (1, 1, NJA) - # This matches how real MF6 files are structured - assert data_back.shape == (1, 1, nja), ( - f"Expected shape (1, 1, {nja}), got {data_back.shape}" - ) - np.testing.assert_allclose( - flowja, data_back.flatten(), err_msg="FLOW-JA-FACE data mismatch" - ) - - -def test_cellbudgetfile_write_array_format_regular(tmp_path): - """Test CellBudgetFile.write_budget() with array format for regular budget term.""" - nlay, nrow, ncol = 2, 5, 10 - ncells = nlay * nrow * ncol - - # Create synthetic budget data (e.g., storage) - storage = np.random.rand(nlay, nrow, ncol) * 100.0 - - budget_dict = { - (1, 1): { - "STORAGE": { - "imeth": 1, - "data": storage, - "delt": 1.0, - "pertim": 10.0, - "totim": 10.0, - } - } - } - - # Write the budget file - output_file = tmp_path / "test_regular.bud" - CellBudgetFile.write_budget( - output_file, - budget_dict, - nlay=nlay, - nrow=nrow, - ncol=ncol, - precision="double", - verbose=False, - ) - - # Read it back - cbb = CellBudgetFile(output_file, precision="double") - - assert "STORAGE".ljust(16).encode() in cbb.textlist, "STORAGE not in budget file" - - # Get data back - data_back = cbb.get_data(kstpkper=(0, 0), text="STORAGE")[0] - - # Regular budget terms are returned as 3D arrays - assert data_back.shape == ( - nlay, - nrow, - ncol, - ), f"Expected shape {(nlay, nrow, ncol)}, got {data_back.shape}" - np.testing.assert_allclose(storage, data_back, err_msg="STORAGE data mismatch") - - -def test_cellbudgetfile_write_list_format(tmp_path): - """Test CellBudgetFile.write_budget() with list format (imeth=6).""" - nlay, nrow, ncol = 1, 10, 20 - ncells = nlay * nrow * ncol - - # Create synthetic DATA-SPDIS list data - dt = np.dtype( - [ - ("node", np.int32), - ("node2", np.int32), - ("q", np.float64), - ("qx", np.float64), - ("qy", np.float64), - ("qz", np.float64), - ] - ) - - data_list = np.zeros(ncells, dtype=dt) - data_list["node"] = np.arange(1, ncells + 1) - data_list["node2"] = np.arange(1, ncells + 1) - data_list["q"] = 0.0 - data_list["qx"] = np.random.rand(ncells) * 0.001 - data_list["qy"] = np.random.rand(ncells) * 0.001 - data_list["qz"] = np.random.rand(ncells) * 0.001 - - budget_dict = { - (1, 1): { - "DATA-SPDIS": { - "imeth": 6, - "data": data_list, - "delt": 1.0, - "pertim": 10.0, - "totim": 10.0, - "modelnam": "GWF", - "paknam": "", - "modelnam2": "GWF", - "paknam2": "", - "ndat": 3, - "auxtxt": ["qx", "qy", "qz"], - } - } - } - - # Write the budget file - output_file = tmp_path / "test_list.bud" - CellBudgetFile.write_budget( - output_file, - budget_dict, - nlay=nlay, - nrow=nrow, - ncol=ncol, - precision="double", - verbose=False, - ) - - # Read it back - cbb = CellBudgetFile(output_file, precision="double") - - # Text is padded to 16 characters in binary format - assert "DATA-SPDIS".ljust(16).encode() in cbb.textlist, ( - "DATA-SPDIS not in budget file" - ) - - # Get data back - data_back = cbb.get_data(kstpkper=(0, 0), text="DATA-SPDIS")[0] - - assert len(data_back) == ncells, f"Expected {ncells} entries, got {len(data_back)}" - np.testing.assert_allclose(data_list["qx"], data_back["qx"], err_msg="qx mismatch") - np.testing.assert_allclose(data_list["qy"], data_back["qy"], err_msg="qy mismatch") - np.testing.assert_allclose(data_list["qz"], data_back["qz"], err_msg="qz mismatch") - - -def test_cellbudgetfile_write_multiple_terms(tmp_path): - """Test CellBudgetFile.write_budget() with multiple budget terms.""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - - nlay, nrow, ncol = 1, 10, 20 - ncells = nlay * nrow * ncol - - # Build connectivity for this grid to get NJA - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - # Create FLOW-JA-FACE array data (size NJA) - flowja = np.random.rand(nja) - 0.5 - - # Create DATA-SAT list data - dt_sat = np.dtype([("node", np.int32), ("node2", np.int32), ("q", np.float64)]) - sat_list = np.zeros(ncells, dtype=dt_sat) - sat_list["node"] = np.arange(1, ncells + 1) - sat_list["node2"] = np.arange(1, ncells + 1) - sat_list["q"] = np.random.rand(ncells) - - budget_dict = { - (1, 1): { - "FLOW-JA-FACE": { - "imeth": 1, - "data": flowja, - "delt": 1.0, - "pertim": 10.0, - "totim": 10.0, - }, - "DATA-SAT": { - "imeth": 6, - "data": sat_list, - "delt": 1.0, - "pertim": 10.0, - "totim": 10.0, - "modelnam": "GWF", - "paknam": "", - "modelnam2": "GWF", - "paknam2": "", - "ndat": 1, - "auxtxt": [], - }, - } - } - - # Write the budget file - output_file = tmp_path / "test_multi.bud" - CellBudgetFile.write_budget( - output_file, - budget_dict, - nlay=nlay, - nrow=nrow, - ncol=ncol, - precision="double", - verbose=False, - ) - - # Read it back - cbb = CellBudgetFile(output_file, precision="double") - - assert len(cbb.textlist) == 2, f"Expected 2 budget terms, got {len(cbb.textlist)}" - - # Verify FLOW-JA-FACE (returned as shape (1, 1, NJA)) - flowja_back = cbb.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] - assert flowja_back.shape == (1, 1, nja), ( - f"Expected shape (1, 1, {nja}), got {flowja_back.shape}" - ) - np.testing.assert_allclose(flowja, flowja_back.flatten()) - - # Verify DATA-SAT - sat_back = cbb.get_data(kstpkper=(0, 0), text="DATA-SAT")[0] - np.testing.assert_allclose(sat_list["q"], sat_back["q"]) - - -def test_headfile_instance_write(tmp_path): - """Test HeadFile.write() instance method.""" - # Create test data - data_dict = { - (1, 1): { - "head": np.full((3, 10, 10), 35.0), - "pertim": 1.0, - "totim": 1.0, - }, - (2, 1): { - "head": np.full((3, 10, 10), 34.0), - "pertim": 2.0, - "totim": 2.0, - }, - } - - # Write initial file - file1 = tmp_path / "test.hds" - HeadFile.write_head(file1, data_dict, precision="double") - - # Read it back and write using instance method - hds = HeadFile(file1, precision="double") - file2 = tmp_path / "test_copy.hds" - hds.write(file2) - hds.close() - - # Verify the copy - hds2 = HeadFile(file2, precision="double") - assert len(hds2.kstpkper) == 2 - head1 = hds2.get_data(totim=1.0) - head2 = hds2.get_data(totim=2.0) - np.testing.assert_allclose(head1, 35.0) - np.testing.assert_allclose(head2, 34.0) - hds2.close() - - -def test_headfile_instance_write_filtered(tmp_path): - """Test HeadFile.write() with kstpkper filtering.""" - # Create test data with 3 time steps - data_dict = { - (1, 1): {"head": np.full((2, 5, 5), 35.0), "pertim": 1.0, "totim": 1.0}, - (2, 1): {"head": np.full((2, 5, 5), 34.0), "pertim": 2.0, "totim": 2.0}, - (3, 1): {"head": np.full((2, 5, 5), 33.0), "pertim": 3.0, "totim": 3.0}, - } - - # Write initial file - file1 = tmp_path / "test.hds" - HeadFile.write_head(file1, data_dict, precision="double") - - # Read and write only specific time steps - hds = HeadFile(file1, precision="double") - file2 = tmp_path / "test_filtered.hds" - hds.write(file2, kstpkper=[(1, 1), (3, 1)]) # Skip (2, 1) - hds.close() - - # Verify only 2 time steps in output - hds2 = HeadFile(file2, precision="double") - assert len(hds2.kstpkper) == 2 - assert (1, 1) in hds2.kstpkper - assert (3, 1) in hds2.kstpkper - assert (2, 1) not in hds2.kstpkper - hds2.close() - - -def test_cellbudgetfile_instance_write(tmp_path): - """Test CellBudgetFile.write() instance method.""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - - nlay, nrow, ncol = 1, 10, 20 - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - # Create test budget data - flowja = np.random.rand(nja) - 0.5 - budget_dict = { - (1, 1): { - "FLOW-JA-FACE": { - "imeth": 1, - "data": flowja, - "delt": 1.0, - "pertim": 1.0, - "totim": 1.0, - } - } - } - - # Write initial file (no grid dimensions needed for FLOW-JA-FACE only) - file1 = tmp_path / "test.cbc" - CellBudgetFile.write_budget(file1, budget_dict, precision="double") - - # Read and write using instance method - cbb = CellBudgetFile(file1, precision="double") - file2 = tmp_path / "test_copy.cbc" - cbb.write(file2) - cbb.close() - - # Verify the copy - cbb2 = CellBudgetFile(file2, precision="double") - assert len(cbb2.kstpkper) == 1 - # get_data uses 0-based indexing - flowja_back = cbb2.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] - np.testing.assert_allclose(flowja, flowja_back.flatten()) - cbb2.close() - - -def test_cellbudgetfile_instance_write_filtered(tmp_path): - """Test CellBudgetFile.write() with text filtering.""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - - nlay, nrow, ncol = 1, 10, 20 - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - # Create budget data with two terms - flowja = np.random.rand(nja) - 0.5 - storage = np.random.rand(nlay, nrow, ncol) - - budget_dict = { - (1, 1): { - "FLOW-JA-FACE": { - "imeth": 1, - "data": flowja, - "delt": 1.0, - "pertim": 1.0, - "totim": 1.0, - }, - "STORAGE": { - "imeth": 1, - "data": storage, - "delt": 1.0, - "pertim": 1.0, - "totim": 1.0, - }, - } - } - - # Write initial file - file1 = tmp_path / "test.cbc" - CellBudgetFile.write_budget( - file1, budget_dict, nlay=nlay, nrow=nrow, ncol=ncol, precision="double" - ) - - # Read and write only FLOW-JA-FACE - cbb = CellBudgetFile(file1, precision="double") - file2 = tmp_path / "test_filtered.cbc" - cbb.write(file2, text="FLOW-JA-FACE") - cbb.close() - - # Verify only FLOW-JA-FACE in output - cbb2 = CellBudgetFile(file2, precision="double") - assert len(cbb2.textlist) == 1 - assert b"FLOW-JA-FACE" in cbb2.textlist[0] - # get_data uses 0-based indexing - flowja_back = cbb2.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] - np.testing.assert_allclose(flowja, flowja_back.flatten()) - cbb2.close() diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 58454b6462..40bc4676f0 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -711,70 +711,14 @@ def write( continue # Set defaults from current file if not provided - if "text" not in kwargs: + text = kwargs.get("text") + if text is None: # Use first text entry from file - kwargs["text"] = self.recordarray["text"][0].decode().strip() - if "precision" not in kwargs: - kwargs["precision"] = self.precision + text = self.recordarray["text"][0].decode().strip() - # Write using static method - HeadFile.write_head(filename, data_dict, **kwargs) + precision = kwargs.get("precision", self.precision) + verbose = kwargs.get("verbose", False) - @staticmethod - def write_head( - filename, - data_dict, - text="HEAD", - precision="double", - verbose=False, - ): - """ - Write a MODFLOW binary head file (.hds). - - Parameters - ---------- - filename : str or PathLike - Path to output head file - data_dict : dict - Dictionary with head data and metadata. Each entry should be: - {(kstp, kper): {'head': array, 'pertim': float, 'totim': float}} - where head array has shape (nlay, nrow, ncol) or (nrow, ncol) - for single layer - text : str, optional - Text identifier for head data (default "HEAD"). Can also be "DRAWDOWN", etc. - Will be padded/truncated to 16 characters. - precision : str, optional - 'single' or 'double' (default 'double') - verbose : bool, optional - Print progress messages (default False) - - Notes - ----- - The head file format consists of repeating records, one per layer per time step: - 1. Header (8 values): - - kstp (int32): time step number (1-based) - - kper (int32): stress period number (1-based) - - pertim (float32 or float64): time in current stress period - - totim (float32 or float64): total elapsed time - - text (16 char): text identifier (e.g., "HEAD") - - ncol (int32): number of columns - - nrow (int32): number of rows - - ilay (int32): layer number (1-based) - 2. Data: head array (ncol, nrow) as float32 or float64 - - Examples - -------- - >>> import numpy as np - >>> from flopy.utils.binaryfile import HeadFile - >>> data = { - ... (1, 1): { - ... 'head': np.full((3, 10, 10), 35.0), - ... 'pertim': 1.0, - ... 'totim': 1.0 - ... } - ... } - >>> HeadFile.write_head('model.hds', data) - """ # Set precision if precision == "single": realtype = np.float32 @@ -2621,137 +2565,14 @@ def write( restructured_dict[ksp_tuple][text] = value_with_imeth # Set defaults from current file if not provided - if "precision" not in kwargs: - kwargs["precision"] = self.precision + precision = kwargs.get("precision", self.precision) + verbose = kwargs.get("verbose", False) - # Write using static method # Only pass grid dimensions if they're set (non-zero) - grid_kwargs = {} - if self.nlay > 0: - grid_kwargs["nlay"] = self.nlay - if self.nrow > 0: - grid_kwargs["nrow"] = self.nrow - if self.ncol > 0: - grid_kwargs["ncol"] = self.ncol - - CellBudgetFile.write_budget( - filename, - restructured_dict, - **grid_kwargs, - **kwargs, - ) + nlay = self.nlay if self.nlay > 0 else None + nrow = self.nrow if self.nrow > 0 else None + ncol = self.ncol if self.ncol > 0 else None - @staticmethod - def write_budget( - filename, - budget_dict, - nlay=None, - nrow=None, - ncol=None, - precision="double", - verbose=False, - ): - """ - Write a MODFLOW 6 style binary budget file (.bud or .cbc). - - Parameters - ---------- - filename : str or PathLike - Path to output budget file - budget_dict : dict - Nested dictionary with budget data. Structure: - { - (kstp, kper): { - 'text': { - 'imeth': int (1 for array, 6 for list), - 'data': np.ndarray, - 'delt': float, - 'pertim': float, - 'totim': float, - 'modelnam': str (optional, for imeth=6), - 'paknam': str (optional, for imeth=6), - 'modelnam2': str (optional, for imeth=6), - 'paknam2': str (optional, for imeth=6), - 'ndat': int (optional, for imeth=6, number of data columns), - 'auxtxt': list of str (optional, for imeth=6, auxiliary names) - } - } - } - - For imeth=1 (array format): - data should be shape (nlay, nrow, ncol) - - For imeth=6 (list format): - data should be a numpy recarray with fields: - - 'node' (int32): source node number (1-based) - - 'node2' (int32): destination node number (1-based) - - 'q' (float): flow value - - optional auxiliary fields (float) - nlay : int, optional - Number of layers. Required for non-FLOW-JA-FACE budget terms. - Can be None for files containing only FLOW-JA-FACE data (default None). - nrow : int, optional - Number of rows. Required for non-FLOW-JA-FACE budget terms. - Can be None for files containing only FLOW-JA-FACE data (default None). - ncol : int, optional - Number of columns. Required for non-FLOW-JA-FACE budget terms. - Can be None for files containing only FLOW-JA-FACE data (default None). - precision : str, optional - 'single' or 'double' (default 'double') - verbose : bool, optional - Print progress messages (default False) - - Examples - -------- - Write FLOW-JA-FACE array data: - - >>> import numpy as np - >>> from flopy.utils.binaryfile import CellBudgetFile - >>> flowja = np.random.rand(3367) # Connection flows - >>> data = { - ... (1, 1): { - ... 'FLOW-JA-FACE': { - ... 'imeth': 1, - ... 'data': flowja, - ... 'delt': 1.0, - ... 'pertim': 1.0, - ... 'totim': 1.0 - ... } - ... } - ... } - >>> CellBudgetFile.write_budget('model.bud', data, nlay=1, nrow=40, ncol=20) - - Write DATA-SPDIS list data: - - >>> ncells = 800 - >>> active = np.ones(ncells, dtype=bool) - >>> dt = np.dtype([('node', np.int32), ('node2', np.int32), ('q', np.float64), - ... ('qx', np.float64), ('qy', np.float64), ('qz', np.float64)]) - >>> data_list = np.zeros(ncells, dtype=dt) - >>> data_list['node'] = np.arange(1, ncells+1) - >>> data_list['node2'] = np.arange(1, ncells+1) - >>> data_list['qx'] = np.random.rand(ncells) - >>> data_list['qy'] = np.random.rand(ncells) - >>> data_list['qz'] = np.random.rand(ncells) - >>> data = { - ... (1, 1): { - ... 'DATA-SPDIS': { - ... 'imeth': 6, - ... 'data': data_list, - ... 'delt': 1.0, - ... 'pertim': 1.0, - ... 'totim': 1.0, - ... 'modelnam': 'GWF', - ... 'paknam': '', - ... 'modelnam2': 'GWF', - ... 'paknam2': '', - ... 'ndat': 3, - ... 'auxtxt': ['qx', 'qy', 'qz'] - ... } - ... } - ... } - >>> CellBudgetFile.write_budget('model.bud', data, nlay=1, nrow=40, ncol=20) - """ # Set precision if precision == "single": realtype = np.float32 @@ -2768,11 +2589,11 @@ def write_budget( with open(filename, "wb") as f: # Sort by time step for consistent output - sorted_keys = sorted(budget_dict.keys()) + sorted_keys = sorted(restructured_dict.keys()) for kstpkper in sorted_keys: kstp, kper = kstpkper - time_data = budget_dict[kstpkper] + time_data = restructured_dict[kstpkper] if verbose: print(f"\n Writing kstp={kstp}, kper={kper}") From a269830a2b8cb2939deda89f29f44b078b50100e Mon Sep 17 00:00:00 2001 From: Bonelli Date: Thu, 26 Feb 2026 07:08:44 -0500 Subject: [PATCH 07/14] clean up --- autotest/test_binaryfile.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/autotest/test_binaryfile.py b/autotest/test_binaryfile.py index d6b73d8dd3..59558d98a2 100644 --- a/autotest/test_binaryfile.py +++ b/autotest/test_binaryfile.py @@ -775,5 +775,3 @@ def test_headfile_get_ts_disu_grid(dis_sim, function_tmpdir): ts_old_list, err_msg="DISU HeadFile: old list format should match new list format", ) - - From 8b7027893ed205ca332d78d5343bda5f0e7492ab Mon Sep 17 00:00:00 2001 From: Bonelli Date: Thu, 26 Feb 2026 07:23:38 -0500 Subject: [PATCH 08/14] clean up tests --- autotest/test_binarygrid_util.py | 139 ++++++++----------------------- 1 file changed, 36 insertions(+), 103 deletions(-) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index e5631bd0ad..9053f60ac5 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -6,6 +6,7 @@ from flopy.discretization import StructuredGrid, UnstructuredGrid, VertexGrid from flopy.mf6.utils import MfGrdFile +from flopy.mf6.utils.binarygrid_util import build_structured_connectivity pytestmark = pytest.mark.mf6 @@ -161,151 +162,93 @@ def test_mfgrddisu_modelgrid(mfgrd_test_path): ) -def test_build_structured_connectivity_simple(): - """Test build_structured_connectivity with simple grids.""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - - # Test 1x1x1 grid (single cell, only diagonal connection) +def test_build_structured_connectivity(): + # 1x1x1 ia, ja, nja = build_structured_connectivity(1, 1, 1) - assert ia.shape == (2,), f"ia shape {ia.shape} != (2,)" - assert ja.shape == (1,), f"ja shape {ja.shape} != (1,)" - assert nja == 1, f"nja {nja} != 1" - assert ia[0] == 0 and ia[1] == 1, f"ia {ia} incorrect" - assert ja[0] == 0, f"ja {ja} != [0]" + assert ia.shape == (2,) + assert ja.shape == (1,) + assert nja == 1 + assert ia[0] == 0 and ia[1] == 1 + assert ja[0] == 0 - # Test 1x1x2 grid (2 cells in x, 1 connection between them + 2 diagonals = 3) + # 1x1x2 ia, ja, nja = build_structured_connectivity(1, 1, 2) - assert ia.shape == (3,), f"ia shape {ia.shape} != (3,)" - assert nja == 3, f"nja {nja} != 3" - # Cell 0: diagonal (0) + right neighbor (1) = 2 connections - # Cell 1: diagonal (1) = 1 connection + assert ia.shape == (3,) + assert nja == 3 + # Cell 0: diagonal + right = 2 connections + # Cell 1: diagonal = 1 connection np.testing.assert_array_equal(ia, [0, 2, 3]) np.testing.assert_array_equal(ja, [0, 1, 1]) - # Test 1x2x1 grid (2 cells in y, 1 connection between them + 2 diagonals = 3) + # 1x2x1 ia, ja, nja = build_structured_connectivity(1, 2, 1) - assert ia.shape == (3,), f"ia shape {ia.shape} != (3,)" - assert nja == 3, f"nja {nja} != 3" - # Cell 0: diagonal (0) + front neighbor (1) = 2 connections - # Cell 1: diagonal (1) = 1 connection + assert ia.shape == (3,) + assert nja == 3 + # Cell 0: diagonal + front = 2 connections + # Cell 1: diagonal = 1 connection np.testing.assert_array_equal(ia, [0, 2, 3]) np.testing.assert_array_equal(ja, [0, 1, 1]) - # Test 2x1x1 grid (2 layers, 1 connection between them + 2 diagonals = 3) + # 2x1x1 ia, ja, nja = build_structured_connectivity(2, 1, 1) - assert ia.shape == (3,), f"ia shape {ia.shape} != (3,)" - assert nja == 3, f"nja {nja} != 3" - # Cell 0 (layer 0): diagonal (0) + lower neighbor (1) = 2 connections - # Cell 1 (layer 1): diagonal (1) = 1 connection + assert ia.shape == (3,) + assert nja == 3 + # Cell 0 (layer 0): diagonal + lower = 2 connections + # Cell 1 (layer 1): diagonal = 1 connection np.testing.assert_array_equal(ia, [0, 2, 3]) np.testing.assert_array_equal(ja, [0, 1, 1]) - -def test_build_structured_connectivity_2x2x2(): - """Test build_structured_connectivity with a 2x2x2 grid.""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - + # 2x2x2 nlay, nrow, ncol = 2, 2, 2 - ncells = nlay * nrow * ncol # 8 cells + ncells = nlay * nrow * ncol ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - - # Verify dimensions - assert ia.shape == (ncells + 1,), f"ia shape {ia.shape} != {(ncells + 1,)}" - assert ja.shape == (nja,), f"ja shape {ja.shape} != {(nja,)}" - assert ia[-1] == nja, f"ia[-1] {ia[-1]} != nja {nja}" - - # Verify IA is monotonically increasing + assert ia.shape == (ncells + 1,) + assert ja.shape == (nja,) + assert ia[-1] == nja assert np.all(np.diff(ia) >= 0), "IA array not monotonically increasing" - - # Verify all JA entries are valid cell indices assert np.all(ja >= 0), "JA contains negative indices" assert np.all(ja < ncells), f"JA contains indices >= {ncells}" - - # Verify each cell has at least a diagonal connection + # check diagonals for i in range(ncells): nconn = ia[i + 1] - ia[i] assert nconn >= 1, f"Cell {i} has {nconn} connections (should be >= 1)" - - # Verify diagonal entries for node in range(ncells): conn_start = ia[node] - # First connection should be diagonal (self) assert ja[conn_start] == node, ( f"Cell {node} first connection {ja[conn_start]} != {node}" ) - -def test_build_structured_connectivity_with_idomain(): - """Test build_structured_connectivity with inactive cells.""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - + # 1x3x3 with center cell inactive nlay, nrow, ncol = 1, 3, 3 # 9 cells ncells = nlay * nrow * ncol - - # Make center cell inactive idomain = np.ones((nlay, nrow, ncol), dtype=np.int32) - idomain[0, 1, 1] = 0 # Cell 4 (center) is inactive - + idomain[0, 1, 1] = 0 ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol, idomain=idomain) - assert ia.shape == (ncells + 1,) assert ja.shape == (nja,) - - # Inactive cell (node 4) should have no connections + # inactive cell (node 4) should have no connections node_4_start = ia[4] node_4_end = ia[5] - assert node_4_end == node_4_start, ( - f"Inactive cell 4 has {node_4_end - node_4_start} connections (should be 0)" - ) - - # Cells around the inactive cell should not connect to it + assert node_4_end == node_4_start + # cells around node 4 should not connect to it for node in range(ncells): if node == 4: - continue # Skip inactive cell + continue conn_start = ia[node] conn_end = ia[node + 1] connections = ja[conn_start:conn_end] - assert 4 not in connections, ( - f"Cell {node} incorrectly connects to inactive cell 4" - ) - - -def test_build_structured_connectivity_known_values(): - """Test build_structured_connectivity against known values.""" - from flopy.mf6.utils.binarygrid_util import build_structured_connectivity - - # Simple 1x2x2 grid - # Cell layout: - # 2 3 - # 0 1 - # Expected connections: - # Cell 0: 0 (diag), 1 (right), 2 (front) -> ia[0]=0, ia[1]=3 - # Cell 1: 1 (diag), 3 (front) -> ia[1]=3, ia[2]=5 - # Cell 2: 2 (diag), 3 (right) -> ia[2]=5, ia[3]=7 - # Cell 3: 3 (diag) -> ia[3]=7, ia[4]=8 - - ia, ja, nja = build_structured_connectivity(1, 2, 2) - assert nja == 8, f"nja {nja} != 8" - np.testing.assert_array_equal(ia, [0, 3, 5, 7, 8]) - np.testing.assert_array_equal(ja, [0, 1, 2, 1, 3, 2, 3, 3]) + assert 4 not in connections def test_write_grb_instance_method(tmp_path, mfgrd_test_path): - """Test MfGrdFile.write() instance method.""" - from flopy.mf6.utils.binarygrid_util import MfGrdFile - - # Read an existing GRB file original_file = mfgrd_test_path / "nwtp3.dis.grb" grb_orig = MfGrdFile(original_file, verbose=False) - # Write using instance method output_file = tmp_path / "test_instance.dis.grb" grb_orig.write(output_file, verbose=False) - # Read it back grb_new = MfGrdFile(output_file, verbose=False) - # Verify all properties match assert grb_new.grid_type == grb_orig.grid_type assert grb_new.nodes == grb_orig.nodes assert grb_new.nlay == grb_orig.nlay @@ -328,28 +271,18 @@ def test_write_grb_instance_method(tmp_path, mfgrd_test_path): def test_write_grb_instance_method_precision_conversion(tmp_path, mfgrd_test_path): - """Test MfGrdFile.write() with precision conversion.""" - from flopy.mf6.utils.binarygrid_util import MfGrdFile - - # Read an existing GRB file (presumably double precision) original_file = mfgrd_test_path / "nwtp3.dis.grb" grb = MfGrdFile(original_file, verbose=False) - # Write as single precision single_file = tmp_path / "test_single.grb" grb.write(single_file, precision="single", verbose=False) - # Write as double precision double_file = tmp_path / "test_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 both work and have same basic properties assert grb_single.nodes == grb.nodes assert grb_double.nodes == grb.nodes - - # Single precision file should be smaller assert single_file.stat().st_size < double_file.stat().st_size From 7c554aaa2657a66aa573cb6d10cc780647187719 Mon Sep 17 00:00:00 2001 From: Bonelli Date: Thu, 26 Feb 2026 11:01:34 -0500 Subject: [PATCH 09/14] support writing disv/u grbs --- autotest/test_binarygrid_util.py | 173 +++++++++++++++++++++++++++++ flopy/mf6/utils/binarygrid_util.py | 44 +++++++- 2 files changed, 216 insertions(+), 1 deletion(-) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index 9053f60ac5..9e66618693 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -286,3 +286,176 @@ def test_write_grb_instance_method_precision_conversion(tmp_path, mfgrd_test_pat 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): + """Test MfGrdFile.write() for DISV grid with roundtrip validation.""" + from pathlib import Path + + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISV grb file + original_file = Path("examples/data/mfgrd_test/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): + """Test MfGrdFile.write() for DISV grid with precision conversion.""" + from pathlib import Path + + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISV grb file + original_file = Path("examples/data/mfgrd_test/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 + + # NOTE: Data value verification is disabled due to known issue with + # precision conversion - see DISV_PRECISION_ISSUE.md + + +def test_write_grb_disu_roundtrip(tmp_path): + """Test MfGrdFile.write() for DISU grid with roundtrip validation.""" + from pathlib import Path + + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISU grb file + original_file = Path("examples/data/mfgrd_test/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): + """Test MfGrdFile.write() for DISU grid with precision conversion.""" + from pathlib import Path + + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISU grb file + original_file = Path("examples/data/mfgrd_test/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 + + # NOTE: Data value verification is disabled due to known issue with + # precision conversion - see DISV_PRECISION_ISSUE.md diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index 52b02bb040..318a4aab13 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -809,10 +809,52 @@ def write(self, filename, precision=None, version=1, verbose=False): ("IDOMAIN", "INTEGER", 1, [self.nodes]), ("ICELLTYPE", "INTEGER", 1, [self.nodes]), ] + elif self.grid_type == "DISV": + # Get dimensions for DISV arrays + nvert = self._datadict["NVERT"] + njavert = self._datadict["NJAVERT"] + var_list = [ + ("NCELLS", "INTEGER", 0, []), + ("NLAY", "INTEGER", 0, []), + ("NCPL", "INTEGER", 0, []), + ("NVERT", "INTEGER", 0, []), + ("NJAVERT", "INTEGER", 0, []), + ("NJA", "INTEGER", 0, []), + ("XORIGIN", float_type, 0, []), + ("YORIGIN", float_type, 0, []), + ("ANGROT", float_type, 0, []), + ("TOP", float_type, 1, [self.nodes]), + ("BOTM", float_type, 1, [self.nodes]), + ("VERTICES", float_type, 2, [nvert, 2]), + ("CELLX", float_type, 1, [self.nodes]), + ("CELLY", float_type, 1, [self.nodes]), + ("IAVERT", "INTEGER", 1, [self.nodes + 1]), + ("JAVERT", "INTEGER", 1, [njavert]), + ("IA", "INTEGER", 1, [self.nodes + 1]), + ("JA", "INTEGER", 1, [self.nja]), + ("IDOMAIN", "INTEGER", 1, [self.nodes]), + ("ICELLTYPE", "INTEGER", 1, [self.nodes]), + ] + elif self.grid_type == "DISU": + var_list = [ + ("NODES", "INTEGER", 0, []), + ("NJA", "INTEGER", 0, []), + ("XORIGIN", float_type, 0, []), + ("YORIGIN", float_type, 0, []), + ("ANGROT", float_type, 0, []), + ("TOP", float_type, 1, [self.nodes]), + ("BOT", float_type, 1, [self.nodes]), + ("IA", "INTEGER", 1, [self.nodes + 1]), + ("JA", "INTEGER", 1, [self.nja]), + ("ICELLTYPE", "INTEGER", 1, [self.nodes]), + ] + # IDOMAIN is optional for DISU + if "IDOMAIN" in self._datadict: + var_list.insert(-1, ("IDOMAIN", "INTEGER", 1, [self.nodes])) else: raise NotImplementedError( f"Grid type {self.grid_type} not yet implemented. " - "Currently only DIS grids are supported." + "Supported grid types: DIS, DISV, DISU" ) ntxt = len(var_list) From 746076a0e9114ea35f81fd0867669b1ec886b87b Mon Sep 17 00:00:00 2001 From: Bonelli Date: Thu, 26 Feb 2026 15:34:07 -0500 Subject: [PATCH 10/14] fix precision issue --- autotest/test_binarygrid_util.py | 32 ++++++++++++++++++++++++++---- flopy/mf6/utils/binarygrid_util.py | 4 ++-- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index 9e66618693..b4a0bd714e 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -376,8 +376,26 @@ def test_write_grb_disv_precision_conversion(tmp_path): # Single precision file should be smaller assert single_file.stat().st_size < double_file.stat().st_size - # NOTE: Data value verification is disabled due to known issue with - # precision conversion - see DISV_PRECISION_ISSUE.md + # 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): @@ -457,5 +475,11 @@ def test_write_grb_disu_precision_conversion(tmp_path): # Single precision file should be smaller assert single_file.stat().st_size < double_file.stat().st_size - # NOTE: Data value verification is disabled due to known issue with - # precision conversion - see DISV_PRECISION_ISSUE.md + # 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) diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index 318a4aab13..b6c296b581 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -146,9 +146,9 @@ def __init__(self, filename, precision="double", verbose=False): if dt == np.int32: v = self.read_integer() elif dt == np.float32: - v = self.read_real() + v = self._read_values(dt, 1)[0] elif dt == np.float64: - v = self.read_real() + v = self._read_values(dt, 1)[0] self._datadict[key] = v if self.verbose: From 3559fe45b3d435333cc771a1974feaf318dce016 Mon Sep 17 00:00:00 2001 From: Bonelli Date: Thu, 26 Feb 2026 20:54:03 -0500 Subject: [PATCH 11/14] clean up --- flopy/utils/binaryfile/__init__.py | 433 ++++++++++++++--------------- 1 file changed, 205 insertions(+), 228 deletions(-) diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 40bc4676f0..671db2e644 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -671,6 +671,62 @@ def write( >>> # Write specific time steps >>> hds.write('output.hds', kstpkper=[(1, 0), (1, 1)]) """ + + def pad_text(text): + """Pad text to exactly 16 bytes.""" + if isinstance(text, str): + text = text.encode("ascii") + if len(text) > 16: + return text[:16] + elif len(text) < 16: + return text + b" " * (16 - len(text)) + return text + + def write_records(f, data_dict, text_bytes, realtype, verbose): + """Write head records to file.""" + sorted_keys = sorted(data_dict.keys()) + + for kstpkper in sorted_keys: + kstp, kper = kstpkper + entry = data_dict[kstpkper] + + head = np.asarray(entry["head"]) + pertim = entry["pertim"] + totim = entry["totim"] + + # Handle both 3D (nlay, nrow, ncol) and 2D (nrow, ncol) arrays + if head.ndim == 2: + head = head.reshape(1, head.shape[0], head.shape[1]) + + nlay, nrow, ncol = head.shape + + if verbose: + print(f" Writing kstp={kstp}, kper={kper}, totim={totim}") + print(f" Shape: {nlay} layers x {nrow} rows x {ncol} cols") + + # Define header dtype + dt = np.dtype( + [ + ("kstp", np.int32), + ("kper", np.int32), + ("pertim", realtype), + ("totim", realtype), + ("text", "S16"), + ("ncol", np.int32), + ("nrow", np.int32), + ("ilay", np.int32), + ] + ) + + # Write one record per layer + for ilay in range(nlay): + h = np.array( + (kstp, kper, pertim, totim, text_bytes, ncol, nrow, ilay + 1), + dtype=dt, + ) + h.tofile(f) + head[ilay].astype(realtype).tofile(f) + # Determine which time steps to write if kstpkper is None: kstpkper = self.kstpkper @@ -713,76 +769,26 @@ def write( # Set defaults from current file if not provided text = kwargs.get("text") if text is None: - # Use first text entry from file text = self.recordarray["text"][0].decode().strip() precision = kwargs.get("precision", self.precision) verbose = kwargs.get("verbose", False) # Set precision - if precision == "single": - realtype = np.float32 - else: - realtype = np.float64 + realtype = np.float32 if precision == "single" else np.float64 - # Ensure text is exactly 16 bytes - if isinstance(text, str): - text = text.encode("ascii") - if len(text) > 16: - text = text[:16] - elif len(text) < 16: - text = text + b" " * (16 - len(text)) + # Pad text to 16 bytes + text_bytes = pad_text(text) if verbose: print(f"Writing binary head file: {filename}") - print(f" Text identifier: {text.decode().strip()}") + print(f" Text identifier: {text_bytes.decode().strip()}") print(f" Precision: {precision}") print(f" Number of time steps: {len(data_dict)}") + # Write the file with open(filename, "wb") as f: - # Sort by time step for consistent output - sorted_keys = sorted(data_dict.keys()) - - for kstpkper in sorted_keys: - kstp, kper = kstpkper - entry = data_dict[kstpkper] - - head = np.asarray(entry["head"]) - pertim = entry["pertim"] - totim = entry["totim"] - - # Handle both 3D (nlay, nrow, ncol) and 2D (nrow, ncol) arrays - if head.ndim == 2: - head = head.reshape(1, head.shape[0], head.shape[1]) - - nlay, nrow, ncol = head.shape - - if verbose: - print(f" Writing kstp={kstp}, kper={kper}, totim={totim}") - print(f" Shape: {nlay} layers x {nrow} rows x {ncol} cols") - - # Define header dtype - dt = np.dtype( - [ - ("kstp", np.int32), - ("kper", np.int32), - ("pertim", realtype), - ("totim", realtype), - ("text", "S16"), - ("ncol", np.int32), - ("nrow", np.int32), - ("ilay", np.int32), - ] - ) - - # Write one record per layer - for ilay in range(nlay): - h = np.array( - (kstp, kper, pertim, totim, text, ncol, nrow, ilay + 1), - dtype=dt, - ) - h.tofile(f) - head[ilay].astype(realtype).tofile(f) + write_records(f, data_dict, text_bytes, realtype, verbose) if verbose: print(f"Successfully wrote {filename}") @@ -2465,153 +2471,48 @@ def write( >>> # Write specific terms and time steps >>> cbc.write('output.cbc', kstpkper=[(1, 0)], text=['STORAGE', 'FLOW-JA-FACE']) """ - # Determine which time steps to write - if kstpkper is None: - kstpkper = self.kstpkper - # Determine which text entries to write - if text is None: - textlist = self.textlist - elif isinstance(text, str): - textlist = [text.ljust(16).encode()] - else: - textlist = [t.ljust(16).encode() for t in text] - - # Build budget dictionary - budget_dict = {} - for ksp in kstpkper: - # Convert numpy int32 to Python int if needed - kstp = int(ksp[0]) - kper = int(ksp[1]) - ksp_tuple = (kstp, kper) - - # get_data() expects 0-based indexing, but kstpkper contains 1-based values - ksp_0based = (kstp - 1, kper - 1) - - for txt in textlist: - # Get matching text from file (case-insensitive, padded) - txt_str = ( - txt.decode().strip() if isinstance(txt, bytes) else txt.strip() - ) - - # Find matching records - matching_records = [ - t - for t in self.textlist - if txt_str.upper() in t.decode().strip().upper() - ] - - for file_txt in matching_records: - try: - data = self.get_data(kstpkper=ksp_0based, text=file_txt)[0] - - # Get metadata from recordarray - mask = ( - (self.recordarray["kstp"] == kstp) - & (self.recordarray["kper"] == kper) - & (self.recordarray["text"] == file_txt) - ) - records = self.recordarray[mask] - - if len(records) == 0: - continue - - record = records[0] - - # Determine imeth from data structure - if isinstance(data, np.recarray): - imeth = 6 # List format - else: - imeth = 1 # Array format - - # Create dictionary key - key = (ksp_tuple, file_txt.decode().strip(), imeth) - - budget_dict[key] = { - "data": data, - "delt": float(record["delt"]), - "pertim": float(record["pertim"]), - "totim": float(record["totim"]), - } + def pad_text(text): + """Pad text to exactly 16 bytes.""" + if isinstance(text, str): + text = text.encode("ascii") + if len(text) > 16: + return text[:16] + elif len(text) < 16: + return text + b" " * (16 - len(text)) + return text + + def write_records(f, budget_dict, nlay, nrow, ncol, realtype, verbose): + """Write budget records to file. + + Parameters + ---------- + budget_dict : dict + Keys are ((kstp, kper), text, imeth), values are data dicts + """ + # Group records by time step + time_steps = {} + for key, value in budget_dict.items(): + ksp_tuple, text, imeth = key + if ksp_tuple not in time_steps: + time_steps[ksp_tuple] = [] + time_steps[ksp_tuple].append((text, imeth, value)) - # Add model/package names if imeth=6 - if imeth == 6: - budget_dict[key].update( - { - "modelnam": record["modelnam"].decode().strip(), - "paknam": record["paknam"].decode().strip(), - "modelnam2": record["modelnam2"].decode().strip(), - "paknam2": record["paknam2"].decode().strip(), - } - ) - - except Exception as e: - if kwargs.get("verbose", False): - print(f"Warning: Could not read data for {ksp}, {txt}: {e}") - continue - - # Restructure budget_dict to match write_budget format - # write_budget expects: {(kstp, kper): {'text': {...}, 'text2': {...}}} - # but we built: {((kstp, kper), text, imeth): {...}} - restructured_dict = {} - for key, value in budget_dict.items(): - ksp_tuple, text, imeth = key - if ksp_tuple not in restructured_dict: - restructured_dict[ksp_tuple] = {} - - # Add imeth to the value dict - value_with_imeth = value.copy() - value_with_imeth["imeth"] = imeth - restructured_dict[ksp_tuple][text] = value_with_imeth - - # Set defaults from current file if not provided - precision = kwargs.get("precision", self.precision) - verbose = kwargs.get("verbose", False) - - # Only pass grid dimensions if they're set (non-zero) - nlay = self.nlay if self.nlay > 0 else None - nrow = self.nrow if self.nrow > 0 else None - ncol = self.ncol if self.ncol > 0 else None - - # Set precision - if precision == "single": - realtype = np.float32 - else: - realtype = np.float64 - - if verbose: - print(f"Writing binary budget file: {filename}") - print(f" Precision: {precision}") - if nlay is not None and nrow is not None and ncol is not None: - print(f" Grid shape: {nlay} layers x {nrow} rows x {ncol} cols") - else: - print(" Grid shape: not specified (OK for FLOW-JA-FACE only files)") - - with open(filename, "wb") as f: # Sort by time step for consistent output - sorted_keys = sorted(restructured_dict.keys()) - - for kstpkper in sorted_keys: + for kstpkper in sorted(time_steps.keys()): kstp, kper = kstpkper - time_data = restructured_dict[kstpkper] if verbose: print(f"\n Writing kstp={kstp}, kper={kper}") # Write each budget term for this time step - for text, term_data in time_data.items(): - imeth = term_data["imeth"] + for text, imeth, term_data in time_steps[kstpkper]: data = term_data["data"] delt = term_data.get("delt", 0.0) pertim = term_data["pertim"] totim = term_data["totim"] - # Ensure text is exactly 16 bytes - text_bytes = text.encode("ascii") if isinstance(text, str) else text - if len(text_bytes) > 16: - text_bytes = text_bytes[:16] - elif len(text_bytes) < 16: - text_bytes = text_bytes + b" " * (16 - len(text_bytes)) + text_bytes = pad_text(text) if verbose: print(f" Writing {text.strip()}: imeth={imeth}") @@ -2677,13 +2578,7 @@ def write( # Ensure each is exactly 16 bytes for name in [modelnam, paknam, modelnam2, paknam2]: - name_bytes = ( - name.encode("ascii") if isinstance(name, str) else name - ) - if len(name_bytes) > 16: - name_bytes = name_bytes[:16] - elif len(name_bytes) < 16: - name_bytes = name_bytes + b" " * (16 - len(name_bytes)) + name_bytes = pad_text(name) f.write(name_bytes) # Write data based on imeth @@ -2691,60 +2586,33 @@ def write( # Array format arr = np.asarray(data, dtype=realtype) - # Check if this is FLOW-JA-FACE (connection-based) - is_flowja = text.strip().upper() in [ - "FLOW-JA-FACE", - "FLOW-JA-FACE-X", - ] - if is_flowja: # FLOW-JA-FACE: keep as 1D array of size NJA - # Don't reshape - connection-based not cell-based if arr.ndim != 1: arr = arr.flatten() else: # Regular budget term: reshape to grid if needed if arr.ndim == 1: - # Flattened - reshape to (nlay, nrow, ncol) arr = arr.reshape(nlay, nrow, ncol) arr.tofile(f) elif imeth == 6: - # List format - # Write naux+1 - ndat = term_data.get("ndat", 1) + # List format - write naux+1 auxtxt = term_data.get("auxtxt", []) naux = len(auxtxt) - nauxp1 = naux + 1 - - np.array([nauxp1], dtype=np.int32).tofile(f) + np.array([naux + 1], dtype=np.int32).tofile(f) # Write auxiliary variable names for auxname in auxtxt: - auxname_bytes = ( - auxname.encode("ascii") - if isinstance(auxname, str) - else auxname - ) - if len(auxname_bytes) > 16: - auxname_bytes = auxname_bytes[:16] - elif len(auxname_bytes) < 16: - auxname_bytes = auxname_bytes + b" " * ( - 16 - len(auxname_bytes) - ) - f.write(auxname_bytes) + f.write(pad_text(auxname)) # Write nlist nlist = len(data) np.array([nlist], dtype=np.int32).tofile(f) - # Write list data - # Data should be a recarray with fields: - # node, node2, q, and aux fields + # Write list data as structured array if isinstance(data, np.ndarray) and data.dtype.names: - # It's already a structured array, write it directly - # But we need to ensure the dtypes match dt_list = [ ("node", np.int32), ("node2", np.int32), @@ -2756,7 +2624,7 @@ def write( output_dt = np.dtype(dt_list) output_data = np.zeros(nlist, dtype=output_dt) - # Copy data from input to output with correct types + # Copy data with correct types for field in output_dt.names: if field in data.dtype.names: output_data[field] = data[field].astype( @@ -2778,6 +2646,115 @@ def write( "(list) are supported." ) + # Determine which time steps to write + if kstpkper is None: + kstpkper = self.kstpkper + + # Determine which text entries to write + if text is None: + textlist = self.textlist + elif isinstance(text, str): + textlist = [text.ljust(16).encode()] + else: + textlist = [t.ljust(16).encode() for t in text] + + # Build budget dictionary + budget_dict = {} + for ksp in kstpkper: + # Convert numpy int32 to Python int if needed + kstp = int(ksp[0]) + kper = int(ksp[1]) + ksp_tuple = (kstp, kper) + + # get_data() expects 0-based indexing, but kstpkper contains 1-based values + ksp_0based = (kstp - 1, kper - 1) + + for txt in textlist: + # Get matching text from file (case-insensitive, padded) + txt_str = ( + txt.decode().strip() if isinstance(txt, bytes) else txt.strip() + ) + + # Find matching records + matching_records = [ + t + for t in self.textlist + if txt_str.upper() in t.decode().strip().upper() + ] + + for file_txt in matching_records: + try: + data = self.get_data(kstpkper=ksp_0based, text=file_txt)[0] + + # Get metadata from recordarray + mask = ( + (self.recordarray["kstp"] == kstp) + & (self.recordarray["kper"] == kper) + & (self.recordarray["text"] == file_txt) + ) + records = self.recordarray[mask] + + if len(records) == 0: + continue + + record = records[0] + + # Determine imeth from data structure + if isinstance(data, np.recarray): + imeth = 6 # List format + else: + imeth = 1 # Array format + + # Create dictionary key + key = (ksp_tuple, file_txt.decode().strip(), imeth) + + budget_dict[key] = { + "data": data, + "delt": float(record["delt"]), + "pertim": float(record["pertim"]), + "totim": float(record["totim"]), + } + + # Add model/package names if imeth=6 + if imeth == 6: + budget_dict[key].update( + { + "modelnam": record["modelnam"].decode().strip(), + "paknam": record["paknam"].decode().strip(), + "modelnam2": record["modelnam2"].decode().strip(), + "paknam2": record["paknam2"].decode().strip(), + } + ) + + except Exception as e: + if kwargs.get("verbose", False): + print(f"Warning: Could not read data for {ksp}, {txt}: {e}") + continue + + # Set defaults from current file if not provided + precision = kwargs.get("precision", self.precision) + verbose = kwargs.get("verbose", False) + + # Only pass grid dimensions if they're set (non-zero) + nlay = self.nlay if self.nlay > 0 else None + nrow = self.nrow if self.nrow > 0 else None + ncol = self.ncol if self.ncol > 0 else None + + # Set precision + realtype = np.float32 if precision == "single" else np.float64 + + if verbose: + print(f"Writing binary budget file: {filename}") + print(f" Precision: {precision}") + if nlay is not None and nrow is not None and ncol is not None: + print(f" Grid shape: {nlay} layers x {nrow} rows x {ncol} cols") + else: + print(" Grid shape: not specified (OK for FLOW-JA-FACE only files)") + + # Write the file + with open(filename, "wb") as f: + write_records(f, budget_dict, nlay, nrow, ncol, realtype, verbose) + if verbose: print(f"\nSuccessfully wrote {filename}") From 7a0bcf6782ec7d8debb1f5251ae0697740d89c2b Mon Sep 17 00:00:00 2001 From: Bonelli Date: Fri, 27 Feb 2026 08:14:41 -0500 Subject: [PATCH 12/14] fix test paths --- autotest/test_binarygrid_util.py | 24 ++++++++---------------- 1 file changed, 8 insertions(+), 16 deletions(-) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index b4a0bd714e..a7fc8e53ce 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -288,14 +288,12 @@ def test_write_grb_instance_method_precision_conversion(tmp_path, mfgrd_test_pat assert single_file.stat().st_size < double_file.stat().st_size -def test_write_grb_disv_roundtrip(tmp_path): +def test_write_grb_disv_roundtrip(tmp_path, mfgrd_test_path): """Test MfGrdFile.write() for DISV grid with roundtrip validation.""" - from pathlib import Path - from flopy.mf6.utils.binarygrid_util import MfGrdFile # Read original DISV grb file - original_file = Path("examples/data/mfgrd_test/flow.disv.grb") + original_file = mfgrd_test_path / "flow.disv.grb" grb_orig = MfGrdFile(original_file, verbose=False) # Write using instance method @@ -346,14 +344,12 @@ def test_write_grb_disv_roundtrip(tmp_path): ) -def test_write_grb_disv_precision_conversion(tmp_path): +def test_write_grb_disv_precision_conversion(tmp_path, mfgrd_test_path): """Test MfGrdFile.write() for DISV grid with precision conversion.""" - from pathlib import Path - from flopy.mf6.utils.binarygrid_util import MfGrdFile # Read original DISV grb file - original_file = Path("examples/data/mfgrd_test/flow.disv.grb") + original_file = mfgrd_test_path / "flow.disv.grb" grb = MfGrdFile(original_file, verbose=False) # Write in single and double precision @@ -398,14 +394,12 @@ def test_write_grb_disv_precision_conversion(tmp_path): ) -def test_write_grb_disu_roundtrip(tmp_path): +def test_write_grb_disu_roundtrip(tmp_path, mfgrd_test_path): """Test MfGrdFile.write() for DISU grid with roundtrip validation.""" - from pathlib import Path - from flopy.mf6.utils.binarygrid_util import MfGrdFile # Read original DISU grb file - original_file = Path("examples/data/mfgrd_test/flow.disu.grb") + original_file = mfgrd_test_path / "flow.disu.grb" grb_orig = MfGrdFile(original_file, verbose=False) # Write using instance method @@ -445,14 +439,12 @@ def test_write_grb_disu_roundtrip(tmp_path): np.testing.assert_array_equal(grb_new.idomain, grb_orig.idomain) -def test_write_grb_disu_precision_conversion(tmp_path): +def test_write_grb_disu_precision_conversion(tmp_path, mfgrd_test_path): """Test MfGrdFile.write() for DISU grid with precision conversion.""" - from pathlib import Path - from flopy.mf6.utils.binarygrid_util import MfGrdFile # Read original DISU grb file - original_file = Path("examples/data/mfgrd_test/flow.disu.grb") + original_file = mfgrd_test_path / "flow.disu.grb" grb = MfGrdFile(original_file, verbose=False) # Write in single and double precision From af07211217a3dec6b99a56f9bedc2d51566ae1b0 Mon Sep 17 00:00:00 2001 From: Bonelli Date: Fri, 27 Feb 2026 10:44:08 -0500 Subject: [PATCH 13/14] remove structured grid connectivity function (to come later) --- autotest/test_binarygrid_util.py | 79 ---------------------- flopy/mf6/utils/__init__.py | 2 +- flopy/mf6/utils/binarygrid_util.py | 105 ----------------------------- 3 files changed, 1 insertion(+), 185 deletions(-) diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index a7fc8e53ce..9d2689a2c8 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -6,7 +6,6 @@ from flopy.discretization import StructuredGrid, UnstructuredGrid, VertexGrid from flopy.mf6.utils import MfGrdFile -from flopy.mf6.utils.binarygrid_util import build_structured_connectivity pytestmark = pytest.mark.mf6 @@ -162,84 +161,6 @@ def test_mfgrddisu_modelgrid(mfgrd_test_path): ) -def test_build_structured_connectivity(): - # 1x1x1 - ia, ja, nja = build_structured_connectivity(1, 1, 1) - assert ia.shape == (2,) - assert ja.shape == (1,) - assert nja == 1 - assert ia[0] == 0 and ia[1] == 1 - assert ja[0] == 0 - - # 1x1x2 - ia, ja, nja = build_structured_connectivity(1, 1, 2) - assert ia.shape == (3,) - assert nja == 3 - # Cell 0: diagonal + right = 2 connections - # Cell 1: diagonal = 1 connection - np.testing.assert_array_equal(ia, [0, 2, 3]) - np.testing.assert_array_equal(ja, [0, 1, 1]) - - # 1x2x1 - ia, ja, nja = build_structured_connectivity(1, 2, 1) - assert ia.shape == (3,) - assert nja == 3 - # Cell 0: diagonal + front = 2 connections - # Cell 1: diagonal = 1 connection - np.testing.assert_array_equal(ia, [0, 2, 3]) - np.testing.assert_array_equal(ja, [0, 1, 1]) - - # 2x1x1 - ia, ja, nja = build_structured_connectivity(2, 1, 1) - assert ia.shape == (3,) - assert nja == 3 - # Cell 0 (layer 0): diagonal + lower = 2 connections - # Cell 1 (layer 1): diagonal = 1 connection - np.testing.assert_array_equal(ia, [0, 2, 3]) - np.testing.assert_array_equal(ja, [0, 1, 1]) - - # 2x2x2 - nlay, nrow, ncol = 2, 2, 2 - ncells = nlay * nrow * ncol - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - assert ia.shape == (ncells + 1,) - assert ja.shape == (nja,) - assert ia[-1] == nja - assert np.all(np.diff(ia) >= 0), "IA array not monotonically increasing" - assert np.all(ja >= 0), "JA contains negative indices" - assert np.all(ja < ncells), f"JA contains indices >= {ncells}" - # check diagonals - for i in range(ncells): - nconn = ia[i + 1] - ia[i] - assert nconn >= 1, f"Cell {i} has {nconn} connections (should be >= 1)" - for node in range(ncells): - conn_start = ia[node] - assert ja[conn_start] == node, ( - f"Cell {node} first connection {ja[conn_start]} != {node}" - ) - - # 1x3x3 with center cell inactive - nlay, nrow, ncol = 1, 3, 3 # 9 cells - ncells = nlay * nrow * ncol - idomain = np.ones((nlay, nrow, ncol), dtype=np.int32) - idomain[0, 1, 1] = 0 - ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol, idomain=idomain) - assert ia.shape == (ncells + 1,) - assert ja.shape == (nja,) - # inactive cell (node 4) should have no connections - node_4_start = ia[4] - node_4_end = ia[5] - assert node_4_end == node_4_start - # cells around node 4 should not connect to it - for node in range(ncells): - if node == 4: - continue - conn_start = ia[node] - conn_end = ia[node + 1] - connections = ja[conn_start:conn_end] - assert 4 not in connections - - 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) diff --git a/flopy/mf6/utils/__init__.py b/flopy/mf6/utils/__init__.py index 476f0602df..45de130dd9 100644 --- a/flopy/mf6/utils/__init__.py +++ b/flopy/mf6/utils/__init__.py @@ -1,4 +1,4 @@ -from .binarygrid_util import MfGrdFile, build_structured_connectivity +from .binarygrid_util import MfGrdFile from .generate_classes import generate_classes from .lakpak_utils import get_lak_connections from .mfsimlistfile import MfSimulationList diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index b6c296b581..9bbc15b45d 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -937,108 +937,3 @@ def write_text(f, text, width): if verbose: print(f"Successfully wrote {filename}") - - -def build_structured_connectivity(nlay, nrow, ncol, idomain=None): - """ - Build IA and JA connectivity arrays for a structured (DIS) grid. - - Parameters - ---------- - nlay : int - Number of layers - nrow : int - Number of rows - ncol : int - Number of columns - idomain : np.ndarray, optional - Domain array indicating active (>0) and inactive (<=0) cells. - Shape: (nlay, nrow, ncol). If None, all cells are active. - - Returns - ------- - ia : np.ndarray - Index array (CSR format), shape (ncells + 1,), dtype int32. - ia[n] is the starting position in ja for cell n's connections. - ia[ncells] is the total number of connections. - ja : np.ndarray - Connection array (CSR format), shape (nja,), dtype int32. - Contains cell numbers for each connection (0-based). - nja : int - Total number of connections - - Notes - ----- - Connectivity order for structured grids (upper triangle only): - 1. Diagonal (self connection) - 2. Right (+1 in j, same k, i) - 3. Front (+1 in i, same k, j) - 4. Lower (+1 in k, same i, j) - - The IA/JA arrays use 0-based indexing (Python convention). - When writing to MF6 binary files, add 1 to convert to Fortran 1-based indexing. - - Examples - -------- - >>> import numpy as np - >>> from flopy.mf6.utils import build_structured_connectivity - >>> nlay, nrow, ncol = 2, 3, 3 - >>> ia, ja, nja = build_structured_connectivity(nlay, nrow, ncol) - >>> print(f"Total cells: {nlay * nrow * ncol}, connections: {nja}") - Total cells: 18, connections: 42 - """ - ncells = nlay * nrow * ncol - - # Default to all active cells if idomain not provided - if idomain is None: - idomain = np.ones((nlay, nrow, ncol), dtype=np.int32) - else: - idomain = np.asarray(idomain, dtype=np.int32) - if idomain.shape != (nlay, nrow, ncol): - raise ValueError( - f"idomain shape {idomain.shape} does not match grid shape " - f"({nlay}, {nrow}, {ncol})" - ) - - ia = np.zeros(ncells + 1, dtype=np.int32) - ja_list = [] - nja = 0 - - for k in range(nlay): - for i in range(nrow): - for j in range(ncol): - node = k * nrow * ncol + i * ncol + j - - # Skip inactive cells - they still get an entry in IA - if idomain[k, i, j] <= 0: - ia[node + 1] = nja - continue - - # Add diagonal (self connection) - ja_list.append(node) - nja += 1 - - # Add connections to neighbors (upper triangle only) - # Right neighbor (j+1) - if j + 1 < ncol and idomain[k, i, j + 1] > 0: - m = k * nrow * ncol + i * ncol + (j + 1) - ja_list.append(m) - nja += 1 - - # Front neighbor (i+1) - if i + 1 < nrow and idomain[k, i + 1, j] > 0: - m = k * nrow * ncol + (i + 1) * ncol + j - ja_list.append(m) - nja += 1 - - # Lower neighbor (k+1) - if k + 1 < nlay and idomain[k + 1, i, j] > 0: - m = (k + 1) * nrow * ncol + i * ncol + j - ja_list.append(m) - nja += 1 - - ia[node + 1] = nja - - ja = np.array(ja_list, dtype=np.int32) - - return ia, ja, nja From 7557e50c468368600385faa23dbcbbf926512732 Mon Sep 17 00:00:00 2001 From: Bonelli Date: Fri, 27 Feb 2026 20:09:36 -0500 Subject: [PATCH 14/14] use binary data utility --- flopy/mf6/utils/binarygrid_util.py | 34 ++++++++++++------------------ flopy/utils/utils_def.py | 30 ++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 20 deletions(-) diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index 9bbc15b45d..c077639f8b 100644 --- a/flopy/mf6/utils/binarygrid_util.py +++ b/flopy/mf6/utils/binarygrid_util.py @@ -866,23 +866,19 @@ def write(self, filename, precision=None, version=1, verbose=False): print(f" Version: {version}") print(f" Number of variables: {ntxt}") - # Helper function to write text with fixed width - def write_text(f, text, width): - """Write text padded to fixed width.""" - text_bytes = text.encode("ascii") - if len(text_bytes) > width: - text_bytes = text_bytes[:width] - else: - text_bytes = text_bytes.ljust(width) - f.write(text_bytes) + # Create writer with appropriate precision + writer = FlopyBinaryData() + writer.precision = precision with open(filename, "wb") as f: + writer.file = f + # Write text header lines (50 chars each, newline terminated) header_len = 50 - write_text(f, f"GRID {self.grid_type}\n", header_len) - write_text(f, f"VERSION {version}\n", header_len) - write_text(f, f"NTXT {ntxt}\n", header_len) - write_text(f, f"LENTXT {lentxt}\n", header_len) + writer.write_text(f"GRID {self.grid_type}\n", header_len) + writer.write_text(f"VERSION {version}\n", header_len) + writer.write_text(f"NTXT {ntxt}\n", header_len) + writer.write_text(f"LENTXT {lentxt}\n", header_len) # Write variable definition lines (100 chars each) for name, dtype_str, ndim, dims in var_list: @@ -893,7 +889,7 @@ def write_text(f, text, width): str(d) for d in dims[::-1] ) # Reverse for Fortran order line = f"{name} {dtype_str} NDIM {ndim} {dims_str}\n" - write_text(f, line, lentxt) + writer.write_text(line, lentxt) # Write binary data for each variable for name, dtype_str, ndim, dims in var_list: @@ -917,11 +913,9 @@ def write_text(f, text, width): if ndim == 0: # Scalar value if dtype_str == "INTEGER": - f.write(np.array(int(value), dtype=np.int32).tobytes()) - elif dtype_str == "DOUBLE": - f.write(np.array(float(value), dtype=np.float64).tobytes()) - elif dtype_str == "SINGLE": - f.write(np.array(float(value), dtype=np.float32).tobytes()) + writer.write_integer(int(value)) + elif dtype_str in ("DOUBLE", "SINGLE"): + writer.write_real(float(value)) else: # Array data arr = np.asarray(value) @@ -933,7 +927,7 @@ def write_text(f, text, width): arr = arr.astype(np.float32) # Write array in column-major (Fortran) order - f.write(arr.flatten(order="F").tobytes()) + writer.write_record(arr.flatten(order="F"), dtype=arr.dtype) if verbose: print(f"Successfully wrote {filename}") diff --git a/flopy/utils/utils_def.py b/flopy/utils/utils_def.py index 669674faa4..d0779c32f9 100644 --- a/flopy/utils/utils_def.py +++ b/flopy/utils/utils_def.py @@ -101,6 +101,36 @@ def read_record(self, count, dtype=None): def _read_values(self, dtype, count): return np.fromfile(self.file, dtype, count) + def write_text(self, text, nchar=20): + """Write text to binary file, padding or truncating to nchar bytes.""" + if isinstance(text, str): + text = text.encode("ascii") + if len(text) > nchar: + text = text[:nchar] + elif len(text) < nchar: + text = text + b" " * (nchar - len(text)) + self.file.write(text) + + def write_integer(self, value): + """Write a single integer to binary file.""" + self._write_values(np.array([value], dtype=self.integer)) + + def write_real(self, value): + """Write a single real number to binary file.""" + self._write_values(np.array([value], dtype=self.real)) + + def write_record(self, array, dtype=None): + """Write an array to binary file.""" + if dtype is None: + dtype = self.real + if not isinstance(array, np.ndarray): + array = np.array(array, dtype=dtype) + self._write_values(array.astype(dtype)) + + def _write_values(self, array): + """Write numpy array to file.""" + array.tofile(self.file) + def totim_to_datetime(totim, start="1-1-1970", timeunit="D"): """