diff --git a/autotest/benchmark_binaryfile_write.py b/autotest/benchmark_binaryfile_write.py new file mode 100644 index 000000000..4c3760f79 --- /dev/null +++ b/autotest/benchmark_binaryfile_write.py @@ -0,0 +1,151 @@ +"""Benchmark tests for binaryfile write methods.""" + +from pathlib import Path + +import numpy as np +import pytest + +from flopy.mf6.utils import MfGrdFile +from flopy.utils import CellBudgetFile, HeadFile + + +@pytest.fixture +def freyberg_hds_path(example_data_path): + return example_data_path / "freyberg_multilayer_transient" / "freyberg.hds" + + +@pytest.fixture +def freyberg_cbc_path(example_data_path): + return example_data_path / "freyberg_multilayer_transient" / "freyberg.cbc" + + +@pytest.fixture +def mfgrd_dis_path(example_data_path): + return example_data_path / "mf6-freyberg" / "freyberg.dis.grb" + + +@pytest.fixture +def mfgrd_disv_path(example_data_path): + return ( + example_data_path + / "mf6" + / "test006_gwf3_disv" + / "expected_output" + / "flow.disv.grb" + ) + + +@pytest.mark.slow +def test_headfile_write_benchmark(benchmark, freyberg_hds_path, tmp_path): + hds = HeadFile(freyberg_hds_path) + nsteps = min(100, len(hds.kstpkper)) + kstpkper = hds.kstpkper[:nsteps] + output_file = tmp_path / "benchmark_output.hds" + + def write_head(): + hds.write(output_file, kstpkper=kstpkper) + + benchmark(write_head) + assert output_file.exists() + + +@pytest.mark.slow +def test_cellbudgetfile_write_benchmark(benchmark, freyberg_cbc_path, tmp_path): + cbc = CellBudgetFile(freyberg_cbc_path) + nsteps = min(50, len(cbc.kstpkper)) + kstpkper = cbc.kstpkper[:nsteps] + output_file = tmp_path / "benchmark_output.cbc" + + def write_budget(): + cbc.write(output_file, kstpkper=kstpkper) + + benchmark(write_budget) + assert output_file.exists() + + +@pytest.mark.slow +def test_mfgrdfile_write_benchmark_dis(benchmark, tmp_path): + nlay, nrow, ncol = 10, 100, 100 + nodes = nlay * nrow * ncol + nja = 7 * nodes - 2 * (nlay * nrow + nlay * ncol + nrow * ncol) + grb_data = { + "NCELLS": nodes, + "NLAY": nlay, + "NROW": nrow, + "NCOL": ncol, + "NJA": nja, + "XORIGIN": 0.0, + "YORIGIN": 0.0, + "ANGROT": 0.0, + "DELR": np.ones(ncol, dtype=np.float64) * 100.0, + "DELC": np.ones(nrow, dtype=np.float64) * 100.0, + "TOP": np.ones(nodes, dtype=np.float64) * 100.0, + "BOTM": np.arange(nodes, dtype=np.float64), + "IA": np.arange(nodes + 1, dtype=np.int32), + "JA": np.arange(nja, dtype=np.int32) % nodes, + "IDOMAIN": np.ones(nodes, dtype=np.int32), + "ICELLTYPE": np.ones(nodes, dtype=np.int32), + } + + from flopy.utils.utils_def import FlopyBinaryData + + temp_grb = tmp_path / "temp_input.grb" + writer = FlopyBinaryData() + writer.precision = "double" + + with open(temp_grb, "wb") as f: + writer.file = f + writer.write_text("GRID DIS\n", 50) + writer.write_text("VERSION 1\n", 50) + writer.write_text("NTXT 16\n", 50) + writer.write_text("LENTXT 100\n", 50) + var_list = [ + ("NCELLS", "INTEGER", 0, []), + ("NLAY", "INTEGER", 0, []), + ("NROW", "INTEGER", 0, []), + ("NCOL", "INTEGER", 0, []), + ("NJA", "INTEGER", 0, []), + ("XORIGIN", "DOUBLE", 0, []), + ("YORIGIN", "DOUBLE", 0, []), + ("ANGROT", "DOUBLE", 0, []), + ("DELR", "DOUBLE", 1, [ncol]), + ("DELC", "DOUBLE", 1, [nrow]), + ("TOP", "DOUBLE", 1, [nodes]), + ("BOTM", "DOUBLE", 1, [nodes]), + ("IA", "INTEGER", 1, [nodes + 1]), + ("JA", "INTEGER", 1, [nja]), + ("IDOMAIN", "INTEGER", 1, [nodes]), + ("ICELLTYPE", "INTEGER", 1, [nodes]), + ] + + for name, dtype_str, ndim, dims in var_list: + if ndim == 0: + line = f"{name} {dtype_str} NDIM {ndim}\n" + else: + dims_str = " ".join(str(d) for d in dims[::-1]) + line = f"{name} {dtype_str} NDIM {ndim} {dims_str}\n" + writer.write_text(line, 100) + + for name, dtype_str, ndim, dims in var_list: + value = grb_data[name] + if ndim == 0: + if dtype_str == "INTEGER": + writer.write_integer(int(value)) + else: + writer.write_real(float(value)) + else: + arr = np.asarray(value) + if dtype_str == "INTEGER": + arr = arr.astype(np.int32) + elif dtype_str == "DOUBLE": + arr = arr.astype(np.float64) + writer.write_record(arr.flatten(order="F"), dtype=arr.dtype) + + grb = MfGrdFile(str(temp_grb), verbose=False) + output_file = tmp_path / "benchmark_output.grb" + + def write_grb(): + grb.write(output_file, verbose=False) + + benchmark(write_grb) + assert output_file.exists() diff --git a/autotest/test_binarygrid_util.py b/autotest/test_binarygrid_util.py index fb4f09aa2..9d2689a2c 100644 --- a/autotest/test_binarygrid_util.py +++ b/autotest/test_binarygrid_util.py @@ -159,3 +159,240 @@ def test_mfgrddisu_modelgrid(mfgrd_test_path): assert nvert == verts.shape[0], ( f"number of vertex (x, y) pairs ({verts.shape[0]}) does not equal {nvert}" ) + + +def test_write_grb_instance_method(tmp_path, mfgrd_test_path): + original_file = mfgrd_test_path / "nwtp3.dis.grb" + grb_orig = MfGrdFile(original_file, verbose=False) + + output_file = tmp_path / "test_instance.dis.grb" + grb_orig.write(output_file, verbose=False) + + grb_new = MfGrdFile(output_file, verbose=False) + + assert grb_new.grid_type == grb_orig.grid_type + assert grb_new.nodes == grb_orig.nodes + assert grb_new.nlay == grb_orig.nlay + assert grb_new.nrow == grb_orig.nrow + assert grb_new.ncol == grb_orig.ncol + assert grb_new.nja == grb_orig.nja + + np.testing.assert_allclose(grb_new.xorigin, grb_orig.xorigin) + np.testing.assert_allclose(grb_new.yorigin, grb_orig.yorigin) + np.testing.assert_allclose(grb_new.angrot, grb_orig.angrot) + + np.testing.assert_allclose(grb_new.delr, grb_orig.delr) + np.testing.assert_allclose(grb_new.delc, grb_orig.delc) + np.testing.assert_allclose(grb_new.top, grb_orig.top) + np.testing.assert_allclose(grb_new.bot, grb_orig.bot) + + np.testing.assert_array_equal(grb_new.ia, grb_orig.ia) + np.testing.assert_array_equal(grb_new.ja, grb_orig.ja) + np.testing.assert_array_equal(grb_new.idomain, grb_orig.idomain) + + +def test_write_grb_instance_method_precision_conversion(tmp_path, mfgrd_test_path): + original_file = mfgrd_test_path / "nwtp3.dis.grb" + grb = MfGrdFile(original_file, verbose=False) + + single_file = tmp_path / "test_single.grb" + grb.write(single_file, precision="single", verbose=False) + + double_file = tmp_path / "test_double.grb" + grb.write(double_file, precision="double", verbose=False) + + grb_single = MfGrdFile(single_file, verbose=False) + grb_double = MfGrdFile(double_file, verbose=False) + + assert grb_single.nodes == grb.nodes + assert grb_double.nodes == grb.nodes + assert single_file.stat().st_size < double_file.stat().st_size + + +def test_write_grb_disv_roundtrip(tmp_path, mfgrd_test_path): + """Test MfGrdFile.write() for DISV grid with roundtrip validation.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISV grb file + original_file = mfgrd_test_path / "flow.disv.grb" + grb_orig = MfGrdFile(original_file, verbose=False) + + # Write using instance method + output_file = tmp_path / "test_disv.grb" + grb_orig.write(output_file, verbose=False) + + # Read it back + grb_new = MfGrdFile(output_file, verbose=False) + + # Verify grid type and dimensions + assert grb_new.grid_type == "DISV" + assert grb_new.grid_type == grb_orig.grid_type + assert grb_new.nodes == grb_orig.nodes + assert grb_new.nlay == grb_orig.nlay + assert grb_new.ncpl == grb_orig.ncpl + assert grb_new.nja == grb_orig.nja + + # Verify coordinates + np.testing.assert_allclose(grb_new.xorigin, grb_orig.xorigin) + np.testing.assert_allclose(grb_new.yorigin, grb_orig.yorigin) + np.testing.assert_allclose(grb_new.angrot, grb_orig.angrot) + + # Verify elevation arrays + np.testing.assert_allclose(grb_new.top, grb_orig.top) + np.testing.assert_allclose(grb_new.bot, grb_orig.bot) + + # Verify cell connectivity + np.testing.assert_array_equal(grb_new.ia, grb_orig.ia) + np.testing.assert_array_equal(grb_new.ja, grb_orig.ja) + + # Verify DISV-specific data + assert grb_new._datadict["NVERT"] == grb_orig._datadict["NVERT"] + assert grb_new._datadict["NJAVERT"] == grb_orig._datadict["NJAVERT"] + np.testing.assert_allclose( + grb_new._datadict["VERTICES"], grb_orig._datadict["VERTICES"] + ) + np.testing.assert_allclose(grb_new._datadict["CELLX"], grb_orig._datadict["CELLX"]) + np.testing.assert_allclose(grb_new._datadict["CELLY"], grb_orig._datadict["CELLY"]) + np.testing.assert_array_equal( + grb_new._datadict["IAVERT"], grb_orig._datadict["IAVERT"] + ) + np.testing.assert_array_equal( + grb_new._datadict["JAVERT"], grb_orig._datadict["JAVERT"] + ) + np.testing.assert_array_equal(grb_new.idomain, grb_orig.idomain) + np.testing.assert_array_equal( + grb_new._datadict["ICELLTYPE"], grb_orig._datadict["ICELLTYPE"] + ) + + +def test_write_grb_disv_precision_conversion(tmp_path, mfgrd_test_path): + """Test MfGrdFile.write() for DISV grid with precision conversion.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISV grb file + original_file = mfgrd_test_path / "flow.disv.grb" + grb = MfGrdFile(original_file, verbose=False) + + # Write in single and double precision + single_file = tmp_path / "test_disv_single.grb" + grb.write(single_file, precision="single", verbose=False) + + double_file = tmp_path / "test_disv_double.grb" + grb.write(double_file, precision="double", verbose=False) + + # Read them back + grb_single = MfGrdFile(single_file, verbose=False) + grb_double = MfGrdFile(double_file, verbose=False) + + # Verify dimensions are preserved + assert grb_single.nodes == grb.nodes + assert grb_double.nodes == grb.nodes + assert grb_single.grid_type == "DISV" + assert grb_double.grid_type == "DISV" + + # Single precision file should be smaller + assert single_file.stat().st_size < double_file.stat().st_size + + # Verify data values are preserved (with appropriate tolerances) + # Single precision has ~7 decimal digits of precision + np.testing.assert_allclose(grb_single.top, grb.top, rtol=1e-6) + np.testing.assert_allclose(grb_single.bot, grb.bot, rtol=1e-6) + np.testing.assert_allclose( + grb_single._datadict["VERTICES"], grb._datadict["VERTICES"], rtol=1e-6 + ) + np.testing.assert_allclose( + grb_single._datadict["CELLX"], grb._datadict["CELLX"], rtol=1e-6 + ) + np.testing.assert_allclose( + grb_single._datadict["CELLY"], grb._datadict["CELLY"], rtol=1e-6 + ) + + # Double precision should match exactly (same precision as original) + np.testing.assert_allclose(grb_double.top, grb.top, rtol=1e-12) + np.testing.assert_allclose(grb_double.bot, grb.bot, rtol=1e-12) + np.testing.assert_allclose( + grb_double._datadict["VERTICES"], grb._datadict["VERTICES"], rtol=1e-12 + ) + + +def test_write_grb_disu_roundtrip(tmp_path, mfgrd_test_path): + """Test MfGrdFile.write() for DISU grid with roundtrip validation.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISU grb file + original_file = mfgrd_test_path / "flow.disu.grb" + grb_orig = MfGrdFile(original_file, verbose=False) + + # Write using instance method + output_file = tmp_path / "test_disu.grb" + grb_orig.write(output_file, verbose=False) + + # Read it back + grb_new = MfGrdFile(output_file, verbose=False) + + # Verify grid type and dimensions + assert grb_new.grid_type == "DISU" + assert grb_new.grid_type == grb_orig.grid_type + assert grb_new.nodes == grb_orig.nodes + assert grb_new.nja == grb_orig.nja + + # Verify coordinates + np.testing.assert_allclose(grb_new.xorigin, grb_orig.xorigin) + np.testing.assert_allclose(grb_new.yorigin, grb_orig.yorigin) + np.testing.assert_allclose(grb_new.angrot, grb_orig.angrot) + + # Verify elevation arrays (note: DISU uses TOP/BOT not TOP/BOTM) + np.testing.assert_allclose(grb_new.top, grb_orig.top) + np.testing.assert_allclose(grb_new._datadict["BOT"], grb_orig._datadict["BOT"]) + + # Verify cell connectivity + np.testing.assert_array_equal(grb_new.ia, grb_orig.ia) + np.testing.assert_array_equal(grb_new.ja, grb_orig.ja) + + # Verify DISU-specific data + np.testing.assert_array_equal( + grb_new._datadict["ICELLTYPE"], grb_orig._datadict["ICELLTYPE"] + ) + + # IDOMAIN is optional in DISU - check if present + if "IDOMAIN" in grb_orig._datadict: + assert "IDOMAIN" in grb_new._datadict + np.testing.assert_array_equal(grb_new.idomain, grb_orig.idomain) + + +def test_write_grb_disu_precision_conversion(tmp_path, mfgrd_test_path): + """Test MfGrdFile.write() for DISU grid with precision conversion.""" + from flopy.mf6.utils.binarygrid_util import MfGrdFile + + # Read original DISU grb file + original_file = mfgrd_test_path / "flow.disu.grb" + grb = MfGrdFile(original_file, verbose=False) + + # Write in single and double precision + single_file = tmp_path / "test_disu_single.grb" + grb.write(single_file, precision="single", verbose=False) + + double_file = tmp_path / "test_disu_double.grb" + grb.write(double_file, precision="double", verbose=False) + + # Read them back + grb_single = MfGrdFile(single_file, verbose=False) + grb_double = MfGrdFile(double_file, verbose=False) + + # Verify dimensions are preserved + assert grb_single.nodes == grb.nodes + assert grb_double.nodes == grb.nodes + assert grb_single.grid_type == "DISU" + assert grb_double.grid_type == "DISU" + + # Single precision file should be smaller + assert single_file.stat().st_size < double_file.stat().st_size + + # Verify data values are preserved (with appropriate tolerances) + # Single precision has ~7 decimal digits of precision + np.testing.assert_allclose(grb_single.top, grb.top, rtol=1e-6) + np.testing.assert_allclose(grb_single.bot, grb.bot, rtol=1e-6) + + # Double precision should match exactly (same precision as original) + np.testing.assert_allclose(grb_double.top, grb.top, rtol=1e-12) + np.testing.assert_allclose(grb_double.bot, grb.bot, rtol=1e-12) diff --git a/flopy/mf6/utils/binarygrid_util.py b/flopy/mf6/utils/binarygrid_util.py index 4fe45e0b0..c077639f8 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: @@ -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,188 @@ def cell2d(self): else: vertices, cell2d = None, None return vertices, cell2d + + def write(self, filename, precision=None, version=1, 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) + version : int, optional + Grid file version (default 1) + 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') + >>> # Convert to single precision + >>> grb.write('model_single.dis.grb', precision='single') + """ + if precision is None: + precision = self.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 + float_type = "SINGLE" if precision.lower() == "single" else "DOUBLE" + + if self.grid_type == "DIS": + var_list = [ + ("NCELLS", "INTEGER", 0, []), + ("NLAY", "INTEGER", 0, []), + ("NROW", "INTEGER", 0, []), + ("NCOL", "INTEGER", 0, []), + ("NJA", "INTEGER", 0, []), + ("XORIGIN", float_type, 0, []), + ("YORIGIN", float_type, 0, []), + ("ANGROT", float_type, 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]), + ] + 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. " + "Supported grid types: DIS, DISV, DISU" + ) + + ntxt = len(var_list) + lentxt = 100 + + if verbose: + print(f"Writing binary grid file: {filename}") + print(f" Grid type: {self.grid_type}") + print(f" Version: {version}") + print(f" Number of variables: {ntxt}") + + # 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 + 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: + 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 grid file") + + 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.flatten(order="F"), dtype=arr.dtype) + + if verbose: + print(f"Successfully wrote {filename}") diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 07027f852..934433707 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -639,6 +639,148 @@ 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. + + Parameters + ---------- + filename : str or PathLike + Path to output head file + kstpkper : list of tuples, optional + Subset 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 is the 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 + + # Set defaults from current file if not provided + text = kwargs.get("text") + if text is None: + text = self.recordarray["text"][0].decode().strip() + + precision = kwargs.get("precision", self.precision) + verbose = kwargs.get("verbose", False) + + # Set precision + realtype = np.float32 if precision == "single" else np.float64 + + # Pad text to 16 bytes + 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 + + text_bytes = pad_text(text) + + # Pre-allocate header dtype outside loop for better performance + dt = np.dtype( + [ + ("kstp", np.int32), + ("kper", np.int32), + ("pertim", realtype), + ("totim", realtype), + ("text", "S16"), + ("ncol", np.int32), + ("nrow", np.int32), + ("ilay", np.int32), + ] + ) + + # Sort kstpkper upfront for correct output order + sorted_kstpkper = sorted(kstpkper, key=lambda x: (int(x[0]), int(x[1]))) + + if verbose: + print(f"Writing binary head file: {filename}") + print(f" Text identifier: {text_bytes.decode().strip()}") + print(f" Precision: {precision}") + print(f" Number of time steps: {len(sorted_kstpkper)}") + + # Write the file + with open(filename, "wb") as f: + for ksp in sorted_kstpkper: + try: + # Convert numpy int32 to Python int if needed + kstp = int(ksp[0]) + kper = int(ksp[1]) + + # 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 verbose: + print(f"Warning: No records found for {ksp}") + continue + + record = matching_records[0] + totim = float(record["totim"]) + pertim = float(record["pertim"]) + + # Get data using totim (works for multi-layer files) + head = np.asarray(self.get_data(totim=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") + + # 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) + + except Exception as e: + if verbose: + print(f"Warning: Could not read data for {ksp}: {e}") + continue + + if verbose: + print(f"Successfully wrote {filename}") + class UcnFile(BinaryLayerFile): """ @@ -2278,6 +2420,296 @@ 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. + + Parameters + ---------- + filename : str or PathLike + Path to output budget file + kstpkper : list of tuples, optional + Subset 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 is the 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']) + """ + + 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 + + # 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] + + # 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 + + # Pre-allocate header dtypes outside loops for better performance + h1dt = np.dtype( + [ + ("kstp", np.int32), + ("kper", np.int32), + ("text", "S16"), + ("ncol", np.int32), + ("nrow", np.int32), + ("nlay", np.int32), + ] + ) + h2dt = np.dtype( + [ + ("imeth", np.int32), + ("delt", realtype), + ("pertim", realtype), + ("totim", realtype), + ] + ) + + # Sort kstpkper upfront for correct output order + sorted_kstpkper = sorted(kstpkper, key=lambda x: (int(x[0]), int(x[1]))) + + # Pre-compute text matching to avoid redundant string operations + text_mapping = {} + for txt in textlist: + txt_str = txt.decode().strip() if isinstance(txt, bytes) else txt.strip() + txt_upper = txt_str.upper() + + # Find matching records from file + matching_records = [ + t for t in self.textlist if txt_upper in t.decode().strip().upper() + ] + text_mapping[txt] = matching_records + + 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: + for ksp in sorted_kstpkper: + # Convert numpy int32 to Python int if needed + kstp = int(ksp[0]) + kper = int(ksp[1]) + + if verbose: + print(f"\n Writing kstp={kstp}, kper={kper}") + + # get_data() expects 0-based indexing, but kstpkper + # contains 1-based values + ksp_0based = (kstp - 1, kper - 1) + + for txt in textlist: + # Use pre-computed text matching + for file_txt in text_mapping[txt]: + 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 + + # Decode text once and reuse + text_str = file_txt.decode().strip() + text_bytes = pad_text(text_str) + delt = float(record["delt"]) + pertim = float(record["pertim"]) + totim = float(record["totim"]) + + if verbose: + print(f" Writing {text_str}: imeth={imeth}") + + # Check if this is FLOW-JA-FACE (connection-based) + is_flowja = text_str.upper() in [ + "FLOW-JA-FACE", + "FLOW-JA-FACE-X", + ] + + # 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) " + f"required for non-FLOW-JA-FACE " + f"budget term '{text_str}'. " + f"Provided: nlay={nlay}, nrow={nrow}, " + f"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) + + header2 = np.array( + [(imeth, delt, pertim, totim)], dtype=h2dt + ) + header2.tofile(f) + + # For imeth=6, write model and package names + if imeth == 6: + modelnam = record["modelnam"].decode().strip() + paknam = record["paknam"].decode().strip() + modelnam2 = record["modelnam2"].decode().strip() + paknam2 = record["paknam2"].decode().strip() + + # Ensure each is exactly 16 bytes + for name in [modelnam, paknam, modelnam2, paknam2]: + name_bytes = pad_text(name) + f.write(name_bytes) + + # Write data based on imeth + if imeth == 0 or imeth == 1: + # Array format + arr = np.asarray(data, dtype=realtype) + + if is_flowja: + # FLOW-JA-FACE: keep as 1D array of size NJA + if arr.ndim != 1: + arr = arr.flatten() + else: + # Regular budget term: reshape to grid if needed + if arr.ndim == 1: + arr = arr.reshape(nlay, nrow, ncol) + + arr.tofile(f) + + elif imeth == 6: + # List format - write naux+1 + auxtxt = [] + if "auxtxt" in record.dtype.names: + # Extract auxiliary text names if available + pass + naux = len(auxtxt) + np.array([naux + 1], dtype=np.int32).tofile(f) + + # Write auxiliary variable names + for auxname in auxtxt: + f.write(pad_text(auxname)) + + # Write nlist + nlist = len(data) + np.array([nlist], dtype=np.int32).tofile(f) + + # Write list data as structured array + if isinstance(data, np.ndarray) and data.dtype.names: + 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 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." + ) + + except Exception as e: + if verbose: + print( + f"Warning: Could not read data for " + f"{ksp}, {txt}: {e}" + ) + continue + + if verbose: + print(f"\nSuccessfully wrote {filename}") + def close(self): """ Close the file handle diff --git a/flopy/utils/utils_def.py b/flopy/utils/utils_def.py index 669674faa..d0779c32f 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"): """