move PRMSChannel to new testing utils

autotest/test_prms_channel.py

@@ -3,10 +3,16 @@
 import pytest
 
 from pywatershed.base.control import Control
+from pywatershed.base.adapter import adapter_factory
 from pywatershed.base.parameters import Parameters
 from pywatershed.hydrology.prms_channel import PRMSChannel, has_prmschannel_f
 from pywatershed.parameters import PrmsParameters
-from pywatershed.utils.netcdf_utils import NetCdfCompare
+
+from utils_compare import compare_in_memory, compare_netcdfs
+
+# compare in memory (faster) or full output files?
+compare_output_files = False
+rtol = atol = 1.0e-8
 
 fail_fast = False
 
@@ -67,52 +73,39 @@ def test_compare_prms(
         budget_type="error",
         calc_method=calc_method,
     )
-    nc_parent = tmp_path / domain["domain_name"]
-    channel.initialize_netcdf(nc_parent)
-    # test that init netcdf twice raises a warning
-    with pytest.warns(UserWarning):
+
+    if compare_output_files:
+        nc_parent = tmp_path / domain["domain_name"]
         channel.initialize_netcdf(nc_parent)
+        # test that init netcdf twice raises a warning
+        with pytest.warns(UserWarning):
+            channel.initialize_netcdf(nc_parent)
+
+    else:
+        answers = {}
+        for var in PRMSChannel.get_variables():
+            var_pth = output_dir / f"{var}.nc"
+            answers[var] = adapter_factory(
+                var_pth, variable_name=var, control=control
+            )
 
     for istep in range(control.n_times):
         control.advance()
-
         channel.advance()
-
         channel.calculate(float(istep))
-
         channel.output()
+        if not compare_output_files:
+            compare_in_memory(channel, answers, atol=atol, rtol=rtol)
 
     channel.finalize()
 
-    output_compare = {}
-
-    for key in PRMSChannel.get_variables():
-        base_nc_path = output_dir / f"{key}.nc"
-        compare_nc_path = tmp_path / domain["domain_name"] / f"{key}.nc"
-        # PRMS does not output the storage change in the channel
-        if not base_nc_path.exists():
-            continue
-        output_compare[key] = (base_nc_path, compare_nc_path)
-
-    assert_error = False
-    for key, (base, compare) in output_compare.items():
-        print(f"\nbase_nc_path: {base}")
-        print(f"compare_nc_path: {compare}")
-        success, diff = NetCdfCompare(base, compare).compare()
-        if not success:
-            print(
-                f"comparison for {key} failed: "
-                + f"maximum error {diff[key][0]} "
-                + f"(maximum allowed error {diff[key][1]}) "
-                + f"in column {diff[key][2]}"
-            )
-            assert_error = True
-            if fail_fast:
-                assert False
-
-        else:
-            print(f"comparison for {key} passed")
-
-    assert not assert_error, "comparison failed"
+    if compare_output_files:
+        compare_netcdfs(
+            PRMSChannel.get_variables(),
+            tmp_path / domain["domain_name"],
+            output_dir,
+            atol=atol,
+            rtol=rtol,
+        )
 
     return

autotest/test_prms_groundwater.py

@@ -7,9 +7,9 @@
 from pywatershed.hydrology.prms_groundwater import has_prmsgroundwater_f
 from pywatershed.parameters import PrmsParameters
 
-from utils_compare import assert_allclose, compare_in_memory, compare_netcdfs
+from utils_compare import compare_in_memory, compare_netcdfs
 
-# compare in memory (faster) or full output files
+# compare in memory (faster) or full output files?
 compare_output_files = False
 rtol = atol = 1.0e-13
 

test_data/common/control.multi_hru

@@ -350,16 +350,17 @@ nsegmentOutON_OFF
 nsegmentOutVars
 1
 1
-9
+10
 ####
 nsegmentOut_freq
 1
 1
 1
 ####
 nsegmentOutVar_names
-9
+10
 4
+seg_inflow
 seg_outflow
 seg_gwflow
 seg_sroff

test_data/common/control.single_hru

@@ -348,16 +348,17 @@ nsegmentOutON_OFF
 nsegmentOutVars
 1
 1
-9
+10
 ####
 nsegmentOut_freq
 1
 1
 1
 ####
 nsegmentOutVar_names
-9
+10
 4
+seg_inflow
 seg_outflow
 seg_gwflow
 seg_sroff

test_data/drb_2yr/control.test

@@ -350,16 +350,17 @@ nsegmentOutON_OFF
 nsegmentOutVars
 1
 1
-9
+10
 ####
 nsegmentOut_freq
 1
 1
 1
 ####
 nsegmentOutVar_names
-9
+10
 4
+seg_inflow
 seg_outflow
 seg_gwflow
 seg_sroff

test_data/generate/conftest.py

@@ -160,3 +160,16 @@ def pytest_generate_tests(metafunc):
         ]
         ids = [pl.Path(kk).name for kk in simulations.keys()]
         metafunc.parametrize("simulation", sims, ids=ids, scope="session")
+
+    if "final_file" in metafunc.fixturenames:
+        final_var_names = ["through_rain", "seg_lateral_inflow"]
+        final_files = [
+            pl.Path(kk) / var
+            for kk in simulations.keys()
+            for var in final_var_names
+        ]
+        ids = [ff.parent.name + ":" + ff.name for ff in final_files]
+        # these are not really file names they are domain/key_var
+        metafunc.parametrize(
+            "final_file", final_files, ids=ids, scope="session"
+        )

test_data/generate/convert_prms_output_to_nc.py

@@ -11,7 +11,7 @@
 
 
 @pytest.fixture
-def netcdf_file(csv_file) -> Path:
+def netcdf_file(csv_file):
     """Convert CSV files from model output to NetCDF"""
 
     var_name = csv_file.stem
@@ -27,8 +27,8 @@ def make_netcdf_files(netcdf_file):
     print(f"Created NetCDF from CSV: {netcdf_file}")
 
 
-@pytest.fixture(scope="session")
-def soltab_netcdf_file(tmp_path_factory, soltab_file) -> Path:
+@pytest.fixture()
+def soltab_netcdf_file(tmp_path_factory, soltab_file):
     """Convert soltab files to NetCDF, one file for each variable"""
     domain_dir = soltab_file.parent
     output_dir = domain_dir / "output"
@@ -40,18 +40,18 @@ def make_soltab_netcdf_files(soltab_netcdf_file):
 
 
 @pytest.fixture(scope="session")
-def final_netcdf_file(tmp_path_factory, simulation) -> Path:
-    """Create the final NetCDF file (through_rain.nc) from other NetCDFs"""
+def final_netcdf_file(tmp_path_factory, final_file):
+    """Create NetCDF files that depend on multiple other NetCDFs"""
 
-    # Currently there is only a single final variable
-    var_name = "through_rain"
-    data_dir = Path(simulation["ws"]) / "output"
+    domain_dir = final_file.parent
+    var_name = final_file.name
+    data_dir = domain_dir / "output"
 
     root_tmpdir = tmp_path_factory.getbasetemp().parent
     with FileLock(root_tmpdir / "final_nc.lock"):
         yield  # do this in session cleanup
 
-        diagnose_final_vars_to_nc(var_name, data_dir)
+        diagnose_final_vars_to_nc(var_name, data_dir, domain_dir)
 
 
 def make_final_netcdf_files(final_netcdf_file):

test_data/generate/prms_diagnostic_variables.py

@@ -198,7 +198,7 @@ def diagnose_final_vars_to_nc(
         data = {}
         for vv in data_vars:
             data_file = data_dir / f"{vv}.nc"
-            data[vv] = xr.open_dataset(data_file)[vv]
+            data[vv] = xr.open_dataarray(data_file)
 
         nearzero = 1.0e-6
 
@@ -228,3 +228,105 @@ def diagnose_final_vars_to_nc(
             data[vv].close()
 
         assert out_file.exists()
+
+    if var_name in [
+        "channel_sroff_vol",
+        "channel_ssres_flow_vol",
+        "channel_gwres_flow_vol",
+        "seg_lateral_inflow",
+    ]:
+        if var_name != "seg_lateral_inflow":
+            return
+
+        data_vars = [
+            "sroff_vol",
+            "ssres_flow_vol",
+            "gwres_flow_vol",
+            "seg_outflow",
+            "seg_inflow",
+        ]
+        data = {}
+        for vv in data_vars:
+            data_file = data_dir / f"{vv}.nc"
+            data[vv] = xr.open_dataarray(data_file)
+
+        control = pws.Control.load(domain_dir / "control.test")
+        s_per_time = control.time_step_seconds
+        params = pws.parameters.PrmsParameters.load(
+            domain_dir / "myparam.param"
+        )
+
+        ntimes = control.n_times
+        nseg = params.dims["nsegment"]
+        nhru = params.dims["nhru"]
+        hru_segment = params.parameters["hru_segment"] - 1
+        to_segment = (params.parameters["tosegment"] - 1).astype("int64")
+
+        outflow_mask = np.full((len(to_segment)), False)
+        for iseg in range(nseg):
+            if to_segment[iseg] < 0:
+                outflow_mask[iseg] = True
+                continue
+
+        seg_stor_change = (
+            data["seg_inflow"] - data["seg_outflow"]
+        ) * s_per_time
+
+        channel_outflow_vol = (
+            np.where(outflow_mask, data["seg_outflow"], zero)
+        ) * s_per_time
+
+        seg_lateral_inflow = np.zeros((ntimes, nseg))
+        channel_sroff_vol = np.zeros((ntimes, nhru))
+        channel_ssres_flow_vol = np.zeros((ntimes, nhru))
+        channel_gwres_flow_vol = np.zeros((ntimes, nhru))
+
+        for ihru in range(nhru):
+            iseg = hru_segment[ihru]
+            if iseg < 0:
+                # This is bad, selective handling of fluxes is not cool,
+                # mass is being discarded in a way that has to be coordinated
+                # with other parts of the code.
+                # This code shuold be removed evenutally.
+                continue
+
+            else:
+                channel_sroff_vol[:, ihru] = data["sroff_vol"].values[:, ihru]
+                channel_ssres_flow_vol[:, ihru] = data[
+                    "ssres_flow_vol"
+                ].values[:, ihru]
+                channel_gwres_flow_vol[:, ihru] = data[
+                    "gwres_flow_vol"
+                ].values[:, ihru]
+
+            # cubicfeet to cfs
+            lateral_inflow = (
+                channel_sroff_vol[:, ihru]
+                + channel_ssres_flow_vol[:, ihru]
+                + channel_gwres_flow_vol[:, ihru]
+            ) / (s_per_time)
+
+            seg_lateral_inflow[:, iseg] += lateral_inflow
+
+        for hvar in [
+            "channel_sroff_vol",
+            "channel_ssres_flow_vol",
+            "channel_gwres_flow_vol",
+            "seg_lateral_inflow",
+            "seg_stor_change",
+            "channel_outflow_vol",
+        ]:
+            if hvar in [
+                "seg_lateral_inflow",
+                "seg_stor_change",
+                "channel_outflow_vol",
+            ]:
+                dum = data["seg_outflow"].rename(hvar) * np.nan
+            else:
+                dum = data["sroff_vol"].rename(hvar) * np.nan
+
+            dum[:, :] = locals()[hvar]
+            out_file = data_dir / f"{hvar}.nc"
+            dum.to_netcdf(out_file)
+            print(out_file)
+            assert out_file.exists()

test_data/hru_1/control.test

@@ -348,16 +348,17 @@ nsegmentOutON_OFF
 nsegmentOutVars
 1
 1
-9
+10
 ####
 nsegmentOut_freq
 1
 1
 1
 ####
 nsegmentOutVar_names
-9
+10
 4
+seg_inflow
 seg_outflow
 seg_gwflow
 seg_sroff

test_data/ucb_2yr/control.test

@@ -350,16 +350,17 @@ nsegmentOutON_OFF
 nsegmentOutVars
 1
 1
-9
+10
 ####
 nsegmentOut_freq
 1
 1
 1
 ####
 nsegmentOutVar_names
-9
+10
 4
+seg_inflow
 seg_outflow
 seg_gwflow
 seg_sroff