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..671db2e64 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -639,6 +639,160 @@ 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 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)]) + """ + + 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 + + # 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 + 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 + text_bytes = pad_text(text) + + 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(data_dict)}") + + # Write the file + with open(filename, "wb") as f: + write_records(f, data_dict, text_bytes, realtype, verbose) + + if verbose: + print(f"Successfully wrote {filename}") + class UcnFile(BinaryLayerFile): """ @@ -2278,6 +2432,332 @@ 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 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 + + 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)) + + # Sort by time step for consistent output + for kstpkper in sorted(time_steps.keys()): + kstp, kper = kstpkper + + if verbose: + print(f"\n Writing kstp={kstp}, kper={kper}") + + # Write each budget term for this time step + 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"] + + text_bytes = pad_text(text) + + 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 = 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 = term_data.get("auxtxt", []) + 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." + ) + + # 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}") + 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"): """