diff --git a/autotest/test_prms_channel.py b/autotest/test_prms_channel.py index 299db4a7..6aff98db 100644 --- a/autotest/test_prms_channel.py +++ b/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 diff --git a/autotest/test_prms_groundwater.py b/autotest/test_prms_groundwater.py index e3f7d880..c833334e 100644 --- a/autotest/test_prms_groundwater.py +++ b/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 diff --git a/test_data/common/control.multi_hru b/test_data/common/control.multi_hru index f29a2b2b..effd9667 100644 --- a/test_data/common/control.multi_hru +++ b/test_data/common/control.multi_hru @@ -350,7 +350,7 @@ nsegmentOutON_OFF nsegmentOutVars 1 1 -9 +10 #### nsegmentOut_freq 1 @@ -358,8 +358,9 @@ nsegmentOut_freq 1 #### nsegmentOutVar_names -9 +10 4 +seg_inflow seg_outflow seg_gwflow seg_sroff diff --git a/test_data/common/control.single_hru b/test_data/common/control.single_hru index 39252644..74ea87b3 100644 --- a/test_data/common/control.single_hru +++ b/test_data/common/control.single_hru @@ -348,7 +348,7 @@ nsegmentOutON_OFF nsegmentOutVars 1 1 -9 +10 #### nsegmentOut_freq 1 @@ -356,8 +356,9 @@ nsegmentOut_freq 1 #### nsegmentOutVar_names -9 +10 4 +seg_inflow seg_outflow seg_gwflow seg_sroff diff --git a/test_data/drb_2yr/control.test b/test_data/drb_2yr/control.test index f29a2b2b..effd9667 100644 --- a/test_data/drb_2yr/control.test +++ b/test_data/drb_2yr/control.test @@ -350,7 +350,7 @@ nsegmentOutON_OFF nsegmentOutVars 1 1 -9 +10 #### nsegmentOut_freq 1 @@ -358,8 +358,9 @@ nsegmentOut_freq 1 #### nsegmentOutVar_names -9 +10 4 +seg_inflow seg_outflow seg_gwflow seg_sroff diff --git a/test_data/generate/conftest.py b/test_data/generate/conftest.py index 83a4fcc3..31468dd7 100644 --- a/test_data/generate/conftest.py +++ b/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" + ) diff --git a/test_data/generate/convert_prms_output_to_nc.py b/test_data/generate/convert_prms_output_to_nc.py index eefebea3..af560a6d 100644 --- a/test_data/generate/convert_prms_output_to_nc.py +++ b/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): diff --git a/test_data/generate/prms_diagnostic_variables.py b/test_data/generate/prms_diagnostic_variables.py index 3591aa2b..12610eb3 100644 --- a/test_data/generate/prms_diagnostic_variables.py +++ b/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() diff --git a/test_data/hru_1/control.test b/test_data/hru_1/control.test index 39252644..74ea87b3 100644 --- a/test_data/hru_1/control.test +++ b/test_data/hru_1/control.test @@ -348,7 +348,7 @@ nsegmentOutON_OFF nsegmentOutVars 1 1 -9 +10 #### nsegmentOut_freq 1 @@ -356,8 +356,9 @@ nsegmentOut_freq 1 #### nsegmentOutVar_names -9 +10 4 +seg_inflow seg_outflow seg_gwflow seg_sroff diff --git a/test_data/ucb_2yr/control.test b/test_data/ucb_2yr/control.test index f29a2b2b..effd9667 100644 --- a/test_data/ucb_2yr/control.test +++ b/test_data/ucb_2yr/control.test @@ -350,7 +350,7 @@ nsegmentOutON_OFF nsegmentOutVars 1 1 -9 +10 #### nsegmentOut_freq 1 @@ -358,8 +358,9 @@ nsegmentOut_freq 1 #### nsegmentOutVar_names -9 +10 4 +seg_inflow seg_outflow seg_gwflow seg_sroff