|
22 | 22 | Util2d, |
23 | 23 | ) |
24 | 24 | from flopy.utils.binaryfile import get_headfile_precision |
| 25 | +from flopy.utils.gridutil import get_disu_kwargs, get_disv_kwargs |
25 | 26 |
|
26 | 27 |
|
27 | 28 | @pytest.fixture |
@@ -597,3 +598,180 @@ def test_read_mf2005_freyberg(example_data_path, function_tmpdir, compact): |
597 | 598 | assert len(cbb_data) == len(cbb_data_kstpkper) |
598 | 599 | for i in range(len(cbb_data)): |
599 | 600 | assert np.array_equal(cbb_data[i], cbb_data_kstpkper[i]) |
| 601 | + |
| 602 | + |
| 603 | +@pytest.fixture |
| 604 | +def dis_sim(function_tmpdir): |
| 605 | + from flopy.mf6 import ( |
| 606 | + MFSimulation, |
| 607 | + ModflowGwf, |
| 608 | + ModflowGwfchd, |
| 609 | + ModflowGwfdis, |
| 610 | + ModflowGwfic, |
| 611 | + ModflowGwfnpf, |
| 612 | + ModflowGwfoc, |
| 613 | + ModflowIms, |
| 614 | + ModflowTdis, |
| 615 | + ) |
| 616 | + |
| 617 | + sim_name = "test_ts_aux_vars" |
| 618 | + sim = MFSimulation(sim_name=sim_name, sim_ws=function_tmpdir, exe_name="mf6") |
| 619 | + tdis = ModflowTdis(sim, nper=2, perioddata=[(1.0, 1, 1.0), (1.0, 1, 1.0)]) |
| 620 | + ims = ModflowIms(sim) |
| 621 | + gwf = ModflowGwf(sim, modelname=sim_name, save_flows=True) |
| 622 | + nrow, ncol, nlay = 5, 5, 1 |
| 623 | + dis = ModflowGwfdis( |
| 624 | + gwf, |
| 625 | + nrow=nrow, |
| 626 | + ncol=ncol, |
| 627 | + nlay=nlay, |
| 628 | + delr=10.0, |
| 629 | + delc=10.0, |
| 630 | + top=10.0, |
| 631 | + botm=[0.0], |
| 632 | + ) |
| 633 | + ic = ModflowGwfic(gwf, strt=5.0) |
| 634 | + npf = ModflowGwfnpf(gwf, k=1.0, save_specific_discharge=True) |
| 635 | + chd_spd = [[(0, 0, 0), 10.0], [(0, 4, 4), 0.0]] |
| 636 | + chd = ModflowGwfchd(gwf, stress_period_data=chd_spd) |
| 637 | + budget_file = f"{sim_name}.cbc" |
| 638 | + head_file = f"{sim_name}.hds" |
| 639 | + oc = ModflowGwfoc( |
| 640 | + gwf, |
| 641 | + budget_filerecord=budget_file, |
| 642 | + head_filerecord=head_file, |
| 643 | + saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")], |
| 644 | + ) |
| 645 | + |
| 646 | + return sim |
| 647 | + |
| 648 | + |
| 649 | +@pytest.mark.requires_exe("mf6") |
| 650 | +def test_headfile_get_ts_disv_grid(dis_sim, function_tmpdir): |
| 651 | + """Test HeadFile.get_ts() with DISV grid using both new and old index formats.""" |
| 652 | + from flopy.mf6 import ModflowGwfchd, ModflowGwfdisv |
| 653 | + from flopy.utils import HeadFile |
| 654 | + |
| 655 | + sim = dis_sim |
| 656 | + gwf = sim.get_model() |
| 657 | + dis_grid = gwf.modelgrid |
| 658 | + |
| 659 | + # Create DISV model |
| 660 | + disv_kwargs = get_disv_kwargs( |
| 661 | + nlay=dis_grid.nlay, |
| 662 | + nrow=dis_grid.nrow, |
| 663 | + ncol=dis_grid.ncol, |
| 664 | + delr=dis_grid.delr, |
| 665 | + delc=dis_grid.delc, |
| 666 | + tp=dis_grid.top, |
| 667 | + botm=dis_grid.botm, |
| 668 | + ) |
| 669 | + gwf.remove_package("dis") |
| 670 | + gwf.remove_package("chd") |
| 671 | + disv = ModflowGwfdisv(gwf, **disv_kwargs) |
| 672 | + chd_spd = [[0, 0, 10.0], [0, 24, 0.0]] |
| 673 | + chd = ModflowGwfchd(gwf, stress_period_data=chd_spd) |
| 674 | + |
| 675 | + sim.set_sim_path(function_tmpdir / "disv_head") |
| 676 | + sim.write_simulation() |
| 677 | + success, _ = sim.run_simulation(silent=False) |
| 678 | + assert success |
| 679 | + |
| 680 | + # Open head file with modelgrid |
| 681 | + head_file = function_tmpdir / "disv_head" / f"{gwf.name}.hds" |
| 682 | + hds = HeadFile(head_file, modelgrid=gwf.modelgrid) |
| 683 | + |
| 684 | + # Test cell (layer=0, cellid=4) |
| 685 | + # NEW format: 2-tuple |
| 686 | + ts_new = hds.get_ts(idx=(0, 4)) |
| 687 | + |
| 688 | + # OLD format: 3-tuple with dummy middle value |
| 689 | + ts_old = hds.get_ts(idx=(0, 0, 4)) |
| 690 | + |
| 691 | + # Both formats should return identical results |
| 692 | + np.testing.assert_array_equal( |
| 693 | + ts_new, |
| 694 | + ts_old, |
| 695 | + err_msg="DISV HeadFile: old 3-tuple format should match new 2-tuple format", |
| 696 | + ) |
| 697 | + |
| 698 | + # Verify we got actual head values (not all zeros or NaN) |
| 699 | + assert ts_new.shape[0] > 0, "Should have at least one time step" |
| 700 | + assert ts_new.shape[1] == 2, "Should have time + 1 head column" |
| 701 | + assert not np.all(np.isnan(ts_new[:, 1])), "Head values should not be all NaN" |
| 702 | + |
| 703 | + # Test with list of cells |
| 704 | + ts_new_list = hds.get_ts(idx=[(0, 4), (0, 10)]) |
| 705 | + ts_old_list = hds.get_ts(idx=[(0, 0, 4), (0, 0, 10)]) |
| 706 | + |
| 707 | + np.testing.assert_array_equal( |
| 708 | + ts_new_list, |
| 709 | + ts_old_list, |
| 710 | + err_msg="DISV HeadFile: old list format should match new list format", |
| 711 | + ) |
| 712 | + |
| 713 | + |
| 714 | +@pytest.mark.requires_exe("mf6") |
| 715 | +def test_headfile_get_ts_disu_grid(dis_sim, function_tmpdir): |
| 716 | + """Test HeadFile.get_ts() with DISU grid using both new and old index formats.""" |
| 717 | + from flopy.mf6 import ModflowGwfchd, ModflowGwfdisu |
| 718 | + from flopy.utils import HeadFile |
| 719 | + |
| 720 | + sim = dis_sim |
| 721 | + gwf = sim.get_model() |
| 722 | + dis_grid = gwf.modelgrid |
| 723 | + |
| 724 | + # Create DISU model |
| 725 | + disu_kwargs = get_disu_kwargs( |
| 726 | + nlay=dis_grid.nlay, |
| 727 | + nrow=dis_grid.nrow, |
| 728 | + ncol=dis_grid.ncol, |
| 729 | + delr=dis_grid.delr, |
| 730 | + delc=dis_grid.delc, |
| 731 | + tp=dis_grid.top, |
| 732 | + botm=dis_grid.botm, |
| 733 | + return_vertices=True, |
| 734 | + ) |
| 735 | + gwf.remove_package("dis") |
| 736 | + gwf.remove_package("chd") |
| 737 | + disu = ModflowGwfdisu(gwf, **disu_kwargs) |
| 738 | + chd_spd = [[0, 10.0], [24, 0.0]] |
| 739 | + chd = ModflowGwfchd(gwf, stress_period_data=chd_spd) |
| 740 | + |
| 741 | + sim.set_sim_path(function_tmpdir / "disu_head") |
| 742 | + sim.write_simulation() |
| 743 | + success, _ = sim.run_simulation(silent=False) |
| 744 | + assert success |
| 745 | + |
| 746 | + # Open head file with modelgrid |
| 747 | + head_file = function_tmpdir / "disu_head" / f"{gwf.name}.hds" |
| 748 | + hds = HeadFile(head_file, modelgrid=gwf.modelgrid) |
| 749 | + |
| 750 | + # Test node 4 |
| 751 | + # NEW format: just the integer |
| 752 | + ts_new = hds.get_ts(idx=4) |
| 753 | + |
| 754 | + # OLD format: 3-tuple with dummy first two values |
| 755 | + ts_old = hds.get_ts(idx=(0, 0, 4)) |
| 756 | + |
| 757 | + # Both formats should return identical results |
| 758 | + np.testing.assert_array_equal( |
| 759 | + ts_new, |
| 760 | + ts_old, |
| 761 | + err_msg="DISU HeadFile: old 3-tuple format should match new integer format", |
| 762 | + ) |
| 763 | + |
| 764 | + # Verify we got actual head values (not all zeros or NaN) |
| 765 | + assert ts_new.shape[0] > 0, "Should have at least one time step" |
| 766 | + assert ts_new.shape[1] == 2, "Should have time + 1 head column" |
| 767 | + assert not np.all(np.isnan(ts_new[:, 1])), "Head values should not be all NaN" |
| 768 | + |
| 769 | + # Test with list of nodes |
| 770 | + ts_new_list = hds.get_ts(idx=[4, 10]) |
| 771 | + ts_old_list = hds.get_ts(idx=[(0, 0, 4), (0, 0, 10)]) |
| 772 | + |
| 773 | + np.testing.assert_array_equal( |
| 774 | + ts_new_list, |
| 775 | + ts_old_list, |
| 776 | + err_msg="DISU HeadFile: old list format should match new list format", |
| 777 | + ) |
0 commit comments