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_binaryfile.py b/autotest/test_binaryfile.py index 59558d98a..7001c1d10 100644 --- a/autotest/test_binaryfile.py +++ b/autotest/test_binaryfile.py @@ -775,3 +775,212 @@ 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_from_data_dict(function_tmpdir): + """Test HeadFile.from_data() with dict input format.""" + # Create test data for a 3-layer, 10-row, 20-col grid + nlay, nrow, ncol = 3, 10, 20 + head1 = np.random.rand(nlay, nrow, ncol).astype(np.float64) + head2 = np.random.rand(nlay, nrow, ncol).astype(np.float64) + + # Create using dict format + data = { + (1, 1): head1, + (1, 2): head2, + } + + # Test with auto-generated filename + hds = HeadFile.from_data(data, precision="double") + assert hds.nlay == nlay + assert hds.nrow == nrow + assert hds.ncol == ncol + assert len(hds.get_times()) == 2 + assert hds.get_times() == [1.0, 2.0] + assert hds.get_kstpkper() == [(0, 0), (0, 1)] + + # Verify data round-trips correctly + data_read1 = hds.get_data(kstpkper=(0, 0)) + data_read2 = hds.get_data(kstpkper=(0, 1)) + np.testing.assert_allclose(data_read1, head1, rtol=1e-10) + np.testing.assert_allclose(data_read2, head2, rtol=1e-10) + + # Test with explicit filename + outfile = function_tmpdir / "test.hds" + hds2 = HeadFile.from_data(data, filename=outfile, precision="double") + assert outfile.exists() + assert hds2.filename == outfile + + +def test_headfile_from_data_list(function_tmpdir): + """Test HeadFile.from_data() with list input format.""" + nlay, nrow, ncol = 2, 5, 10 + head1 = np.random.rand(nlay, nrow, ncol).astype(np.float32) + head2 = np.random.rand(nlay, nrow, ncol).astype(np.float32) + + # Create using list format with explicit metadata + data = [ + {"data": head1, "kstp": 1, "kper": 1, "totim": 10.0, "pertim": 10.0}, + {"data": head2, "kstp": 1, "kper": 2, "totim": 20.0, "pertim": 10.0}, + ] + + hds = HeadFile.from_data(data, precision="single") + assert hds.nlay == nlay + assert hds.nrow == nrow + assert hds.ncol == ncol + assert hds.get_times() == [10.0, 20.0] + + # Verify data + data_read1 = hds.get_data(totim=10.0) + data_read2 = hds.get_data(totim=20.0) + np.testing.assert_allclose(data_read1, head1, rtol=1e-6) + np.testing.assert_allclose(data_read2, head2, rtol=1e-6) + + +def test_headfile_from_data_2d(function_tmpdir): + """Test HeadFile.from_data() with 2D arrays (single layer).""" + nrow, ncol = 10, 20 + head1 = np.random.rand(nrow, ncol) + head2 = np.random.rand(nrow, ncol) + + data = { + (1, 1): head1, + (1, 2): head2, + } + + hds = HeadFile.from_data(data) + assert hds.nlay == 1 + assert hds.nrow == nrow + assert hds.ncol == ncol + + # Get data and check - should get 3D array back + data_read = hds.get_data(kstpkper=(0, 0)) + assert data_read.shape == (1, nrow, ncol) + np.testing.assert_allclose(data_read[0], head1, rtol=1e-10) + + +def test_headfile_from_data_custom_times(function_tmpdir): + """Test HeadFile.from_data() with custom time values.""" + nlay, nrow, ncol = 1, 5, 5 + head1 = np.ones((nlay, nrow, ncol)) + head2 = np.ones((nlay, nrow, ncol)) * 2 + + data = { + (1, 1): head1, + (2, 1): head2, + } + + # Custom totim and pertim + totim = {(1, 1): 5.5, (2, 1): 10.5} + pertim = {(1, 1): 5.5, (2, 1): 5.0} + + hds = HeadFile.from_data(data, totim=totim, pertim=pertim) + assert hds.get_times() == [5.5, 10.5] + + # Verify we can retrieve by totim + data_read = hds.get_data(totim=5.5) + np.testing.assert_allclose(data_read, head1) + + +def test_cellbudgetfile_from_data_dict(function_tmpdir): + """Test CellBudgetFile.from_data() with dict input format.""" + # Create test data for flow-ja-face (1D arrays) + nnodes = 600 # 3 * 10 * 20 + flow1 = np.random.rand(nnodes).astype(np.float64) + flow2 = np.random.rand(nnodes).astype(np.float64) + + data = { + (1, 1): flow1, + (1, 2): flow2, + } + + # Create with explicit dimensions + cbb = CellBudgetFile.from_data( + data, text="FLOW-JA-FACE", nlay=3, nrow=10, ncol=20, precision="double" + ) + assert cbb.nlay == 3 + assert cbb.nrow == 10 + assert cbb.ncol == 20 + assert len(cbb.get_times()) == 2 + assert cbb.get_times() == [1.0, 2.0] + + # Verify data round-trips + data_read1 = cbb.get_data(kstpkper=(0, 0), text="FLOW-JA-FACE")[0] + data_read2 = cbb.get_data(kstpkper=(0, 1), text="FLOW-JA-FACE")[0] + np.testing.assert_allclose(data_read1.flatten(), flow1, rtol=1e-10) + np.testing.assert_allclose(data_read2.flatten(), flow2, rtol=1e-10) + + +def test_cellbudgetfile_from_data_list(function_tmpdir): + """Test CellBudgetFile.from_data() with list input format.""" + nnodes = 100 + flow1 = np.random.rand(nnodes).astype(np.float32) + flow2 = np.random.rand(nnodes).astype(np.float32) + + data = [ + { + "data": flow1, + "kstp": 1, + "kper": 1, + "totim": 10.0, + "pertim": 10.0, + "text": "STORAGE", + }, + { + "data": flow2, + "kstp": 1, + "kper": 2, + "totim": 20.0, + "pertim": 10.0, + "text": "STORAGE", + }, + ] + + cbb = CellBudgetFile.from_data(data, nlay=1, nrow=1, ncol=100, precision="single") + assert cbb.get_times() == [10.0, 20.0] + + # Verify data + data_read1 = cbb.get_data(totim=10.0, text="STORAGE")[0] + data_read2 = cbb.get_data(totim=20.0, text="STORAGE")[0] + np.testing.assert_allclose(data_read1.flatten(), flow1, rtol=1e-6) + np.testing.assert_allclose(data_read2.flatten(), flow2, rtol=1e-6) + + +def test_headfile_from_data_errors(function_tmpdir): + """Test that HeadFile.from_data() raises appropriate errors.""" + # Empty data + with pytest.raises(ValueError, match="No data records"): + HeadFile.from_data({}) + + # 1D array (not allowed) + with pytest.raises(ValueError, match="at least 2D"): + HeadFile.from_data({(1, 1): np.array([1, 2, 3])}) + + # Inconsistent shapes + data = { + (1, 1): np.ones((10, 20)), + (1, 2): np.ones((10, 15)), # Different ncol + } + with pytest.raises(ValueError, match="Inconsistent array shapes"): + HeadFile.from_data(data) + + +def test_cellbudgetfile_from_data_errors(function_tmpdir): + """Test that CellBudgetFile.from_data() raises appropriate errors.""" + # Empty data + with pytest.raises(ValueError, match="No data records"): + CellBudgetFile.from_data({}) + + # imeth != 1 not supported yet + data = [{"data": np.ones(100), "kstp": 1, "kper": 1, "imeth": 6}] + with pytest.raises(NotImplementedError, match="Only imeth=1"): + CellBudgetFile.from_data(data) + + # Dimension mismatch + with pytest.raises(ValueError, match="Dimensions don't match"): + CellBudgetFile.from_data( + {(1, 1): np.ones(100)}, + nlay=2, + nrow=10, + ncol=10, # Should be 200 nodes + ) 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/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index 720388f25..9016d917a 100644 --- a/autotest/test_cellbudgetfile.py +++ b/autotest/test_cellbudgetfile.py @@ -1156,3 +1156,53 @@ def test_cellbudgetfile_get_ts_backwards_compatible_idx_format( ts_new_list, ts_old_list, ) + + +@pytest.mark.requires_exe("mf6") +def test_cellbudgetfile_write_preserves_aux_vars(dis_sim, function_tmpdir): + """Test that write() method preserves auxiliary variables in imeth=6 records.""" + from pathlib import Path + + import numpy as np + + sim = dis_sim + sim.write_simulation() + success, _ = sim.run_simulation(silent=True) + assert success + + gwf = sim.get_model() + cbc_orig = gwf.output.budget() + + # Get DATA-SPDIS which has aux vars (node, q, qx, qy, qz) + spdis_orig = cbc_orig.get_data(text="DATA-SPDIS") + assert len(spdis_orig) > 0 + + # Verify aux fields are present + for field in ["node", "q", "qx", "qy", "qz"]: + assert field in spdis_orig[0].dtype.names, ( + f"Field {field} not found in original data" + ) + + # Write to a new file + output_file = Path(function_tmpdir) / "test_aux_rewritten.cbc" + cbc_orig.write(output_file, kstpkper=cbc_orig.kstpkper[:2]) + + # Read back the written file + from flopy.utils import CellBudgetFile + + cbc_rewritten = CellBudgetFile(output_file) + spdis_rewritten = cbc_rewritten.get_data(text="DATA-SPDIS") + + # Verify aux vars were preserved + for field in ["node", "q", "qx", "qy", "qz"]: + assert field in spdis_rewritten[0].dtype.names, ( + f"Field {field} not found in rewritten data" + ) + + # Verify data matches + assert len(spdis_orig) >= len(spdis_rewritten) + for i in range(len(spdis_rewritten)): + for field in ["q", "qx", "qy", "qz"]: + assert np.allclose(spdis_orig[i][field], spdis_rewritten[i][field]), ( + f"Field {field} mismatch at timestep {i}" + ) 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..774bc9d49 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -550,6 +550,256 @@ def __init__( self.header_dtype = BinaryHeader.set_dtype(bintype="Head", precision=precision) super().__init__(filename, precision, verbose, **kwargs) + @classmethod + def from_data( + cls, + data, + nrow=None, + ncol=None, + nlay=None, + text="head", + precision="double", + totim=None, + pertim=None, + filename=None, + verbose=False, + ): + """ + Create a HeadFile from arrays without reading from an existing file. + + This factory method creates a binary head file from provided data arrays, + allowing programmatic creation of head files for testing or data generation. + + Parameters + ---------- + data : dict or list + Head data in one of two formats: + + 1. Dict mapping (kstp, kper) tuples to arrays: + {(kstp, kper): array, ...} + - Arrays should be 2D (nrow, ncol) or 3D (nlay, nrow, ncol) + - If totim/pertim not provided, totim defaults to kper, pertim to totim + + 2. List of dicts with full metadata: + [{'data': array, 'kstp': int, 'kper': int, + 'totim': float, 'pertim': float, 'ilay': int (optional)}, ...] + - Each dict represents one layer at one timestep + - ilay defaults to 1 if not provided + + nrow : int, optional + Number of rows. If None, inferred from data arrays. + ncol : int, optional + Number of columns. If None, inferred from data arrays. + nlay : int, optional + Number of layers. If None, inferred from data arrays. + text : str, default "head" + Text identifier for the head data (will be padded to 16 characters) + precision : str, default "double" + Precision of floating point data: 'single' or 'double' + totim : dict or list, optional + Total time values. If dict, maps (kstp, kper) to totim. + If list, should match order of data. If None, defaults to kper. + pertim : dict or list, optional + Period time values. If dict, maps (kstp, kper) to pertim. + If list, should match order of data. If None, defaults to totim. + filename : str or PathLike, optional + Path for the output file. If None, creates a temporary file. + verbose : bool, default False + Print progress messages + + Returns + ------- + HeadFile + Instance loaded from the created binary file + + Examples + -------- + >>> import numpy as np + >>> from flopy.utils import HeadFile + >>> + >>> # Create head data for two time steps + >>> head1 = np.random.rand(3, 10, 20) # 3 layers, 10 rows, 20 cols + >>> head2 = np.random.rand(3, 10, 20) + >>> data = { + ... (1, 1): head1, + ... (1, 2): head2, + ... } + >>> hds = HeadFile.from_data(data) + >>> hds.get_times() + [1.0, 2.0] + >>> + >>> # Or with explicit time values + >>> data_with_times = [ + ... {'data': head1, 'kstp': 1, 'kper': 1, 'totim': 10.0, 'pertim': 10.0}, + ... {'data': head2, 'kstp': 1, 'kper': 2, 'totim': 20.0, 'pertim': 10.0}, + ... ] + >>> hds = HeadFile.from_data(data_with_times) + """ + # Normalize data to list of record dicts + if isinstance(data, dict): + records = [] + for (kstp, kper), arr in sorted(data.items()): + # Ensure array is at least 2D + arr = np.asarray(arr) + if arr.ndim == 1: + raise ValueError( + "Data arrays must be at least 2D (nrow, ncol), got 1D array" + ) + + # Handle 2D vs 3D arrays + if arr.ndim == 2: + # Single layer + arr = arr.reshape(1, arr.shape[0], arr.shape[1]) + + nlayers, nrows, ncols = arr.shape + + # Get time values + if totim is None: + tot = float(kper) + elif isinstance(totim, dict): + tot = totim.get((kstp, kper), float(kper)) + else: + raise ValueError("totim must be None or dict when data is dict") + + if pertim is None: + per = tot + elif isinstance(pertim, dict): + per = pertim.get((kstp, kper), tot) + else: + raise ValueError("pertim must be None or dict when data is dict") + + # Create one record per layer + for ilay in range(nlayers): + records.append( + { + "data": arr[ilay], + "kstp": kstp, + "kper": kper, + "totim": tot, + "pertim": per, + "ilay": ilay + 1, + } + ) + elif isinstance(data, list): + records = [] + for rec in data: + arr = np.asarray(rec["data"]) + if arr.ndim == 1: + raise ValueError("Data arrays must be at least 2D") + + # Handle 2D vs 3D + if arr.ndim == 2: + # Single layer record + records.append( + { + "data": arr, + "kstp": rec["kstp"], + "kper": rec["kper"], + "totim": rec.get("totim", float(rec["kper"])), + "pertim": rec.get( + "pertim", rec.get("totim", float(rec["kper"])) + ), + "ilay": rec.get("ilay", 1), + } + ) + else: + # 3D array - create one record per layer + nlayers = arr.shape[0] + for ilay in range(nlayers): + records.append( + { + "data": arr[ilay], + "kstp": rec["kstp"], + "kper": rec["kper"], + "totim": rec.get("totim", float(rec["kper"])), + "pertim": rec.get( + "pertim", rec.get("totim", float(rec["kper"])) + ), + "ilay": ilay + 1, + } + ) + else: + raise ValueError("data must be dict or list") + + if len(records) == 0: + raise ValueError("No data records provided") + + # Infer dimensions from data if not provided + first_data = records[0]["data"] + if nrow is None: + nrow = first_data.shape[0] + if ncol is None: + ncol = first_data.shape[1] + if nlay is None: + nlay = max(rec["ilay"] for rec in records) + + # Validate dimensions + for rec in records: + if rec["data"].shape != (nrow, ncol): + raise ValueError( + f"Inconsistent array shapes: expected ({nrow}, {ncol}), " + f"got {rec['data'].shape}" + ) + + # Set precision dtype + realtype = np.float32 if precision == "single" else np.float64 + + # Pad text to 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)) + + # Create temporary file if no filename provided + if filename is None: + # Create a temp file that won't be auto-deleted + fd, filename = tempfile.mkstemp(suffix=".hds") + import os + + os.close(fd) # Close the file descriptor, we'll open it for writing + + # Write binary file + if verbose: + print(f"Writing binary head file: {filename}") + print(f" Text identifier: {text.decode().strip()}") + print(f" Precision: {precision}") + print(f" Dimensions: {nlay} layers, {nrow} rows, {ncol} columns") + print(f" Number of records: {len(records)}") + + # Use BinaryHeader.create() and write like Util2d.write_bin() does + with open(filename, "wb") as f: + for rec in records: + # Create header using BinaryHeader.create() + header = BinaryHeader.create( + bintype="Head", + precision=precision, + text=text.decode().strip(), + nrow=nrow, + ncol=ncol, + ilay=rec["ilay"], + pertim=rec["pertim"], + totim=rec["totim"], + kstp=rec["kstp"], + kper=rec["kper"], + ) + + # Write header and data + header.tofile(f) + rec["data"].astype(realtype).tofile(f) + + # Explicitly flush and sync to ensure data is written + f.flush() + import os + + os.fsync(f.fileno()) + + # Load and return HeadFile instance + return cls( + filename, text=text.decode().strip(), precision=precision, verbose=verbose + ) + def reverse(self, filename: Optional[PathLike] = None): """ Reverse the time order of the currently loaded binary head file. If a head @@ -639,6 +889,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): """ @@ -1021,6 +1413,303 @@ def __enter__(self): def __exit__(self, *exc): self.close() + @classmethod + def from_data( + cls, + data, + text="FLOW-JA-FACE", + imeth=1, + precision="double", + delt=1.0, + pertim=None, + totim=None, + nlay=None, + nrow=None, + ncol=None, + filename=None, + verbose=False, + ): + """ + Create a CellBudgetFile from arrays without reading from an existing file. + + This factory method creates a binary cell budget file from provided data arrays, + allowing programmatic creation of budget files for testing or data generation. + + Parameters + ---------- + data : dict or list + Budget data in one of two formats: + + 1. Dict mapping (kstp, kper) tuples to arrays: + {(kstp, kper): array, ...} + - For imeth=1: arrays should be 1D (flattened cell-by-cell data) + - If totim/pertim not provided, totim defaults to kper, pertim to totim + + 2. List of dicts with full metadata: + [{'data': array, 'kstp': int, 'kper': int, 'text': str (optional), + 'totim': float (optional), 'pertim': float (optional), + 'delt': float (optional), 'imeth': int (optional)}, ...] + - Allows per-record customization of all parameters + + text : str or list, default "FLOW-JA-FACE" + Budget text identifier (will be padded to 16 characters). + If list, must match length of data records. + imeth : int, default 1 + Method code: + - 1: Full 3D array (most common) + - 6: List-based budget (for MF6 advanced packages) + precision : str, default "double" + Precision of floating point data: 'single' or 'double' + delt : float or dict, default 1.0 + Time step length. If dict, maps (kstp, kper) to delt values. + pertim : float or dict, optional + Period time. If dict, maps (kstp, kper) to pertim. + If None, defaults to totim. + totim : float or dict, optional + Total simulation time. If dict, maps (kstp, kper) to totim. + If None, defaults to kper. + nlay : int, optional + Number of layers. Inferred if None (assumes nodes = nlay*nrow*ncol). + nrow : int, optional + Number of rows. Required for imeth=1 if nlay and ncol not provided. + ncol : int, optional + Number of columns. Required for imeth=1 if nlay and nrow not provided. + filename : str or PathLike, optional + Path for output file. If None, creates a temporary file. + verbose : bool, default False + Print progress messages + + Returns + ------- + CellBudgetFile + Instance loaded from the created binary file + + Examples + -------- + >>> import numpy as np + >>> from flopy.utils import CellBudgetFile + >>> + >>> # Create flow-ja-face data for two time steps + >>> # For a 3x10x20 grid, nja might be ~6000 connections + >>> flow1 = np.random.rand(6000) + >>> flow2 = np.random.rand(6000) + >>> data = { + ... (1, 1): flow1, + ... (1, 2): flow2, + ... } + >>> cbb = CellBudgetFile.from_data(data, text='FLOW-JA-FACE') + >>> + >>> # Or with explicit metadata + >>> data_with_meta = [ + ... {'data': flow1, 'kstp': 1, 'kper': 1, 'totim': 10.0, + ... 'text': 'FLOW-JA-FACE'}, + ... {'data': flow2, 'kstp': 1, 'kper': 2, 'totim': 20.0, + ... 'text': 'FLOW-JA-FACE'}, + ... ] + >>> cbb = CellBudgetFile.from_data(data_with_meta) + """ + # Normalize data to list of record dicts + if isinstance(data, dict): + records = [] + for (kstp, kper), arr in sorted(data.items()): + arr = np.asarray(arr).flatten() + + # Get time values + if totim is None: + tot = float(kper) + elif isinstance(totim, dict): + tot = totim.get((kstp, kper), float(kper)) + elif isinstance(totim, (int, float)): + tot = float(totim) + else: + raise ValueError("totim must be None, number, or dict") + + if pertim is None: + per = tot + elif isinstance(pertim, dict): + per = pertim.get((kstp, kper), tot) + elif isinstance(pertim, (int, float)): + per = float(pertim) + else: + raise ValueError("pertim must be None, number, or dict") + + if isinstance(delt, dict): + dt = delt.get((kstp, kper), 1.0) + elif isinstance(delt, (int, float)): + dt = float(delt) + else: + raise ValueError("delt must be number or dict") + + # Get text for this record + if isinstance(text, list): + rec_text = text[len(records)] + else: + rec_text = text + + records.append( + { + "data": arr, + "kstp": kstp, + "kper": kper, + "totim": tot, + "pertim": per, + "delt": dt, + "text": rec_text, + "imeth": imeth, + } + ) + elif isinstance(data, list): + records = [] + for rec in data: + arr = np.asarray(rec["data"]).flatten() + records.append( + { + "data": arr, + "kstp": rec["kstp"], + "kper": rec["kper"], + "totim": rec.get("totim", float(rec["kper"])), + "pertim": rec.get( + "pertim", rec.get("totim", float(rec["kper"])) + ), + "delt": rec.get("delt", 1.0), + "text": rec.get("text", text), + "imeth": rec.get("imeth", imeth), + } + ) + else: + raise ValueError("data must be dict or list") + + if len(records) == 0: + raise ValueError("No data records provided") + + # Only imeth=1 is supported for now + for rec in records: + if rec["imeth"] != 1: + raise NotImplementedError( + f"Only imeth=1 is currently supported, got imeth={rec['imeth']}" + ) + + # Infer dimensions from first record + first_data = records[0]["data"] + nnodes = len(first_data) + + # For imeth=1, we need nlay, nrow, ncol + # If all three provided, use them; otherwise try to infer + if nlay is not None and nrow is not None and ncol is not None: + if nlay * nrow * ncol != nnodes: + raise ValueError( + f"Dimensions don't match: nlay={nlay}, nrow={nrow}, ncol={ncol} " + f"gives {nlay * nrow * ncol} nodes but data has {nnodes}" + ) + else: + # Set defaults to make it work for common cases + if nlay is None: + nlay = 1 + if nrow is None: + nrow = 1 + if ncol is None: + ncol = nnodes + + # Set precision dtype + realtype = np.float32 if precision == "single" else np.float64 + + # Prepare dtypes for headers + 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), + ] + ) + + # Helper function to pad text + def pad_text(txt): + if isinstance(txt, str): + txt = txt.encode("ascii") + if len(txt) > 16: + return txt[:16] + elif len(txt) < 16: + return txt + b" " * (16 - len(txt)) + return txt + + # Create temporary file if no filename provided + if filename is None: + fd, filename = tempfile.mkstemp(suffix=".cbc") + import os + + os.close(fd) + + # Write binary file + if verbose: + print(f"Writing binary budget file: {filename}") + print(f" Precision: {precision}") + print(f" Dimensions: {nlay} layers, {nrow} rows, {ncol} columns") + print(f" Number of records: {len(records)}") + + with open(filename, "wb") as f: + for rec in records: + # Pad text + text_bytes = pad_text(rec["text"]) + + # Write header 1 + # Note: nlay must be negative to indicate a compact budget file + h1 = np.array( + (rec["kstp"], rec["kper"], text_bytes, ncol, nrow, -nlay), + dtype=h1dt, + ) + h1.tofile(f) + + # Write header 2 + h2 = np.array( + ( + rec["imeth"], + realtype(rec["delt"]), + realtype(rec["pertim"]), + realtype(rec["totim"]), + ), + dtype=h2dt, + ) + h2.tofile(f) + + # Write data + arr = rec["data"].astype(realtype) + if len(arr) != nnodes: + raise ValueError( + f"Inconsistent data sizes: expected {nnodes}, got {len(arr)}" + ) + arr.tofile(f) + + # Explicitly flush to ensure data is written + f.flush() + import os + + os.fsync(f.fileno()) + + # Load and return CellBudgetFile instance + obj = cls(filename, precision=precision, verbose=verbose) + + # If dimensions weren't set during _build_index() (e.g., all records were + # FLOW-JA-FACE which are skipped), set them explicitly from the provided values + if obj.nrow == 0: + obj.nrow = nrow + obj.ncol = ncol + obj.nlay = nlay + obj.shape = (nlay, nrow, ncol) + obj.nnodes = nlay * nrow * ncol + + return obj + def __len__(self): """ Return the number of records (headers) in the file. @@ -2278,6 +2967,251 @@ 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 fit_text(text): + """Pad or truncate text to 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 + + if kstpkper is None: + kstpkper = self.kstpkper + + 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] + + verbose = kwargs.get("verbose", False) + precision = kwargs.get("precision", self.precision) + realtype = np.float32 if precision == "single" else np.float64 + + # header dtypes + 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), + ] + ) + + sorted_kstpkper = sorted(kstpkper, key=lambda x: (int(x[0]), int(x[1]))) + + text_mapping = {} + for txt in textlist: + txt_str = txt.decode().strip() if isinstance(txt, bytes) else txt.strip() + txt_upper = txt_str.upper() + matching_records = [ + t for t in self.textlist if txt_upper in t.decode().strip().upper() + ] + text_mapping[txt] = matching_records + + 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 + + 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") + + with open(filename, "wb") as f: + for ksp in sorted_kstpkper: + kstp = int(ksp[0]) + kper = int(ksp[1]) + + # get_data() expects 0-based but kstpkper is 1-based + ksp_0based = (kstp - 1, kper - 1) + + if verbose: + print(f"\n Writing kstp={kstp}, kper={kper}") + + for txt in textlist: + for file_txt in text_mapping[txt]: + data_list = self.get_data(kstpkper=ksp_0based, text=file_txt) + if not data_list: + continue + data = data_list[0] + 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] + + if isinstance(data, np.recarray): + imeth = 6 # list + else: + imeth = 1 # array + + text_str = file_txt.decode().strip() + text_bytes = fit_text(text_str) + delt = float(record["delt"]) + pertim = float(record["pertim"]) + totim = float(record["totim"]) + + if verbose: + print(f" Writing {text_str}: imeth={imeth}") + + is_flowja = text_str.upper() == "FLOW-JA-FACE" + + # Determine dimensions based on data type + if is_flowja and imeth in [0, 1]: + # keep FLOW-JA-FACE flat/size NJA + nja = np.asarray(data).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}" + ) + # negative nlay -> 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) + + if imeth in [0, 1]: + arr = np.asarray(data, dtype=realtype) + # keep FLOW-JA-FACE flat/size NJA. + # reshape other variables to grid. + if is_flowja and arr.ndim != 1: + arr = arr.flatten() + elif arr.ndim == 1: + arr = arr.reshape(nlay, nrow, ncol) + + arr.tofile(f) + + elif imeth == 6: + # write model and package names + modelnam = record["modelnam"].decode().strip() + paknam = record["paknam"].decode().strip() + modelnam2 = record["modelnam2"].decode().strip() + paknam2 = record["paknam2"].decode().strip() + + for name in [modelnam, paknam, modelnam2, paknam2]: + name_bytes = fit_text(name) + f.write(name_bytes) + + # write naux and aux var names + standard_fields = {"node", "node2", "q"} + auxtxt = [ + name + for name in data.dtype.names + if name not in standard_fields + ] + naux = len(auxtxt) + np.array([naux + 1], dtype=np.int32).tofile(f) + for auxname in auxtxt: + f.write(fit_text(auxname)) + + if not (isinstance(data, np.ndarray) and data.dtype.names): + raise ValueError( + "For imeth=6, data must be a numpy recarray " + "with fields: node, node2, q, and optional " + "auxiliary fields" + ) + + # nrite nlist and list data + nlist = len(data) + np.array([nlist], dtype=np.int32).tofile(f) + 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) + 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 NotImplementedError( + "Expected imeth=1 (array) or imeth=6 (list)" + ) + + 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"): """