From 82d7d5610d056ad770fb902b7eec7106e98d506b Mon Sep 17 00:00:00 2001 From: James McCreight Date: Thu, 14 Nov 2024 19:52:32 -0700 Subject: [PATCH 1/4] asv tweaks --- asv_benchmarks/asv.conf.json | 3 ++- asv_benchmarks/benchmarks/prms.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/asv_benchmarks/asv.conf.json b/asv_benchmarks/asv.conf.json index 23162b6b..ada26271 100644 --- a/asv_benchmarks/asv.conf.json +++ b/asv_benchmarks/asv.conf.json @@ -30,7 +30,8 @@ // List of branches to benchmark. If not provided, defaults to "master" // (for git) or "default" (for mercurial). - "branches": ["main"], // for git + // "branches": ["main"], // for git + "branches": ["develop"], // for git // "branches": ["default"], // for mercurial // The DVCS being used. If not set, it will be automatically diff --git a/asv_benchmarks/benchmarks/prms.py b/asv_benchmarks/benchmarks/prms.py index 13484c4d..c5b416f3 100644 --- a/asv_benchmarks/benchmarks/prms.py +++ b/asv_benchmarks/benchmarks/prms.py @@ -61,7 +61,7 @@ def time_prms_parameter_read(self, domain): (domains), ) def time_prms_control_read(self, domain): - control_file = test_data_dir / f"{domain}/control.test" + control_file = test_data_dir / f"{domain}/nhm.control" if pws.__version__ == "0.2.0": _ = pws.Control.load(control_file) else: @@ -78,7 +78,7 @@ def setup(self, *args): self.tag = args[1] self.processes = model_tests[self.tag] - self.control_file = test_data_dir / f"{self.domain}/control.test" + self.control_file = test_data_dir / f"{self.domain}/nhm.control" self.parameter_file = test_data_dir / f"{self.domain}/myparam.param" # backwards compatability pre pywatershed From 51fbbccabe201eb187d5f72b241cf16ba07b53be Mon Sep 17 00:00:00 2001 From: James McCreight Date: Mon, 25 Nov 2024 15:23:44 -0700 Subject: [PATCH 2/4] rename HruSegmentFlowExchange to HruNodeFlowExchange --- autotest/test_prms_channel_flow_graph.py | 9 +++++---- pywatershed/__init__.py | 4 ++-- pywatershed/hydrology/__init__.py | 4 ++-- .../hydrology/prms_channel_flow_graph.py | 18 ++++++++---------- 4 files changed, 17 insertions(+), 18 deletions(-) diff --git a/autotest/test_prms_channel_flow_graph.py b/autotest/test_prms_channel_flow_graph.py index d1e5c294..7141d825 100644 --- a/autotest/test_prms_channel_flow_graph.py +++ b/autotest/test_prms_channel_flow_graph.py @@ -19,7 +19,7 @@ from pywatershed.constants import nan, zero from pywatershed.hydrology.prms_channel_flow_graph import ( HruSegmentFlowAdapter, - HruSegmentFlowExchange, + HruNodeFlowExchange, PRMSChannelFlowNodeMaker, ) from pywatershed.parameters import PrmsParameters @@ -203,7 +203,7 @@ def test_prms_channel_flow_graph_compare_prms( flow_graph.finalize() -exchange_types = ("hrusegmentflowexchange", "inflowexchangefactory") +exchange_types = ("hrunodeflowexchange", "inflowexchangefactory") @pytest.mark.parametrize("exchange_type", exchange_types) @@ -215,8 +215,8 @@ def test_hru_segment_flow_exchange( tmp_path, exchange_type, ): - if exchange_type == "hrusegmentflowexchange": - Exchange = HruSegmentFlowExchange + if exchange_type == "hrunodeflowexchange": + Exchange = HruNodeFlowExchange else: # else we implement the exchange by-hand here. # this is also implemented in channel_flow_graph_to_model_dict @@ -354,6 +354,7 @@ def calculation(self) -> None: ) model = Model(model_dict) + for istep in range(control.n_times): model.advance() model.calculate() diff --git a/pywatershed/__init__.py b/pywatershed/__init__.py index fd2b8c32..60fb3ab3 100644 --- a/pywatershed/__init__.py +++ b/pywatershed/__init__.py @@ -20,7 +20,7 @@ from .hydrology.prms_channel import PRMSChannel from .hydrology.prms_channel_flow_graph import ( HruSegmentFlowAdapter, - HruSegmentFlowExchange, + HruNodeFlowExchange, PRMSChannelFlowNode, PRMSChannelFlowNodeMaker, prms_channel_flow_graph_postprocess, @@ -54,7 +54,7 @@ "PRMSChannelFlowNode", "PRMSChannelFlowNodeMaker", "HruSegmentFlowAdapter", - "HruSegmentFlowExchange", + "HruNodeFlowExchange", "ModelGraph", "ColorBrewer", "PRMSAtmosphere", diff --git a/pywatershed/hydrology/__init__.py b/pywatershed/hydrology/__init__.py index 5edfe2b7..32b03293 100644 --- a/pywatershed/hydrology/__init__.py +++ b/pywatershed/hydrology/__init__.py @@ -2,7 +2,7 @@ from .prms_channel import PRMSChannel from .prms_channel_flow_graph import ( HruSegmentFlowAdapter, - HruSegmentFlowExchange, + HruNodeFlowExchange, PRMSChannelFlowNode, PRMSChannelFlowNodeMaker, prms_channel_flow_graph_postprocess, @@ -22,7 +22,7 @@ "PRMSChannelFlowNode", "PRMSChannelFlowNodeMaker", "HruSegmentFlowAdapter", - "HruSegmentFlowExchange", + "HruNodeFlowExchange", "PRMSCanopy", "PRMSChannel", "PRMSGroundwater", diff --git a/pywatershed/hydrology/prms_channel_flow_graph.py b/pywatershed/hydrology/prms_channel_flow_graph.py index 017e157c..a7a99e33 100644 --- a/pywatershed/hydrology/prms_channel_flow_graph.py +++ b/pywatershed/hydrology/prms_channel_flow_graph.py @@ -405,7 +405,7 @@ class HruSegmentFlowAdapter(Adapter): """Adapt volumetric flows from HRUs to lateral inflows on PRMS segments/nodes. This class specifically maps from PRMS HRU outflows to PRMS segment inflows - using the parameters known to `PRMSChannel`. . + using the parameters known to `PRMSChannel`. This class is a subclass of :class:`Adapter` which means that it makes existing or known flows available over time (but dosent calculate a @@ -490,14 +490,12 @@ def _calculate_segment_lateral_inflows(self): return -class HruSegmentFlowExchange(ConservativeProcess): - """Process to map PRMS HRU outflows to lateral inflows on segments/nodes. - - This class specifically maps from PRMS HRU outflows to PRMS segment inflows - using the parameters known to `PRMSChannel`. +class HruNodeFlowExchange(ConservativeProcess): + """Process to map PRMS HRU outflows to lateral inflows on nodes. - This class is meant to take flows from "upstream" :class:`Process`\ es and - provide flows to a :class:`FlowGraph` in the context of a :class:`Model`. + This class maps PRMS HRU outflows to :class:`FlowGraph` node inflows in + the context of a :class:`Model`. To map HRU outflows to PRMS segments, see + :class:`HruSegmentFlowExchange`. """ def __init__( @@ -511,7 +509,7 @@ def __init__( budget_type: Literal["defer", None, "warn", "error"] = "defer", verbose: bool = None, ) -> None: - """Instantiate a HruSegmentFlowExchange. + """Instantiate a HruNodeFlowExchange. Args: control: A :class:`Control` object. @@ -534,7 +532,7 @@ def __init__( discretization=discretization, parameters=parameters, ) - self.name = "HruSegmentFlowExchange" + self.name = "HruNodeFlowExchange" self._set_inputs(locals()) self._set_options(locals()) From 9490e00a16d64d3762b399e203ca357244cd2a74 Mon Sep 17 00:00:00 2001 From: James McCreight Date: Mon, 25 Nov 2024 20:20:10 -0700 Subject: [PATCH 3/4] prms_segment_lateral_inflow_components_to_netcdf --- autotest/test_prms_channel_flow_graph.py | 24 +++- pywatershed/__init__.py | 4 +- pywatershed/base/adapter.py | 4 +- pywatershed/hydrology/__init__.py | 4 +- .../hydrology/prms_channel_flow_graph.py | 135 +++++++++++++++++- 5 files changed, 160 insertions(+), 11 deletions(-) diff --git a/autotest/test_prms_channel_flow_graph.py b/autotest/test_prms_channel_flow_graph.py index 7141d825..1f621127 100644 --- a/autotest/test_prms_channel_flow_graph.py +++ b/autotest/test_prms_channel_flow_graph.py @@ -10,6 +10,7 @@ PRMSRunoff, PRMSSoilzone, prms_channel_flow_graph_to_model_dict, + prms_segment_lateral_inflow_components_to_netcdf, ) from pywatershed.base.adapter import AdapterNetcdf, adapter_factory from pywatershed.base.control import Control @@ -18,8 +19,8 @@ from pywatershed.base.parameters import Parameters from pywatershed.constants import nan, zero from pywatershed.hydrology.prms_channel_flow_graph import ( - HruSegmentFlowAdapter, HruNodeFlowExchange, + HruSegmentFlowAdapter, PRMSChannelFlowNodeMaker, ) from pywatershed.parameters import PrmsParameters @@ -167,6 +168,7 @@ def test_prms_channel_flow_graph_compare_prms( # check exchange lateral_inflow_answers.advance() + np.testing.assert_allclose( inflow_prms.current, lateral_inflow_answers.current, @@ -508,3 +510,23 @@ def test_prms_channel_flow_graph_to_model_dict( assert da.node_maker_index[-3:].values.tolist() == check_indices assert da.node_maker_name[-3:].values.tolist() == check_names assert da.node_maker_id[-3:].values.tolist() == check_ids + + +def test_prms_segment_lateral_inflow_components_to_netcdf( + simulation, control, parameters, tmp_path +): + nc_out_file_path = tmp_path / "segment_lateral_inflows.nc" + prms_segment_lateral_inflow_components_to_netcdf( + control, + parameters, + input_dir=simulation["output_dir"], + nc_out_file_path=nc_out_file_path, + output_sum=True, + ) + + results = xr.load_dataset(nc_out_file_path).lateral_inflow_vol + answers = xr.load_dataarray( + simulation["output_dir"] / "seg_lateral_inflow.nc" + ) + + xr.testing.assert_allclose(answers, results) diff --git a/pywatershed/__init__.py b/pywatershed/__init__.py index 60fb3ab3..c8699615 100644 --- a/pywatershed/__init__.py +++ b/pywatershed/__init__.py @@ -19,12 +19,13 @@ from .hydrology.prms_canopy import PRMSCanopy from .hydrology.prms_channel import PRMSChannel from .hydrology.prms_channel_flow_graph import ( - HruSegmentFlowAdapter, HruNodeFlowExchange, + HruSegmentFlowAdapter, PRMSChannelFlowNode, PRMSChannelFlowNodeMaker, prms_channel_flow_graph_postprocess, prms_channel_flow_graph_to_model_dict, + prms_segment_lateral_inflow_components_to_netcdf, ) from .hydrology.prms_et import PRMSEt from .hydrology.prms_groundwater import PRMSGroundwater @@ -51,6 +52,7 @@ __all__ = ( "prms_channel_flow_graph_postprocess", "prms_channel_flow_graph_to_model_dict", + "prms_segment_lateral_inflow_components_to_netcdf", "PRMSChannelFlowNode", "PRMSChannelFlowNodeMaker", "HruSegmentFlowAdapter", diff --git a/pywatershed/base/adapter.py b/pywatershed/base/adapter.py index d7890b55..18d3e289 100644 --- a/pywatershed/base/adapter.py +++ b/pywatershed/base/adapter.py @@ -127,9 +127,11 @@ def __init__( self, data: np.ndarray, variable: str, + control: Control = None, ) -> None: super().__init__(variable) self.name = "AdapterOnedarray" + self.control = control self._current_value = data return @@ -174,7 +176,7 @@ def adapter_factory( elif isinstance(var, np.ndarray) and len(var.shape) == 1: # Adapt 1-D np.ndarrays - return AdapterOnedarray(var, variable=variable_name) + return AdapterOnedarray(var, variable=variable_name, control=control) elif isinstance(var, TimeseriesArray): # Adapt TimeseriesArrays as is. diff --git a/pywatershed/hydrology/__init__.py b/pywatershed/hydrology/__init__.py index 32b03293..1326b140 100644 --- a/pywatershed/hydrology/__init__.py +++ b/pywatershed/hydrology/__init__.py @@ -1,12 +1,13 @@ from .prms_canopy import PRMSCanopy from .prms_channel import PRMSChannel from .prms_channel_flow_graph import ( - HruSegmentFlowAdapter, HruNodeFlowExchange, + HruSegmentFlowAdapter, PRMSChannelFlowNode, PRMSChannelFlowNodeMaker, prms_channel_flow_graph_postprocess, prms_channel_flow_graph_to_model_dict, + prms_segment_lateral_inflow_components_to_netcdf, ) from .prms_groundwater import PRMSGroundwater from .prms_groundwater_no_dprst import PRMSGroundwaterNoDprst @@ -19,6 +20,7 @@ __all__ = ( "prms_channel_flow_graph_postprocess", "prms_channel_flow_graph_to_model_dict", + "prms_segment_lateral_inflow_components_to_netcdf", "PRMSChannelFlowNode", "PRMSChannelFlowNodeMaker", "HruSegmentFlowAdapter", diff --git a/pywatershed/hydrology/prms_channel_flow_graph.py b/pywatershed/hydrology/prms_channel_flow_graph.py index a7a99e33..d125296e 100644 --- a/pywatershed/hydrology/prms_channel_flow_graph.py +++ b/pywatershed/hydrology/prms_channel_flow_graph.py @@ -4,8 +4,14 @@ import numba as nb import numpy as np +import xarray as xr -from pywatershed.base.adapter import Adapter, AdapterNetcdf, adaptable +from pywatershed.base.adapter import ( + Adapter, + AdapterNetcdf, + adaptable, + adapter_factory, +) from pywatershed.base.conservative_process import ConservativeProcess from pywatershed.base.control import Control from pywatershed.base.flow_graph import ( @@ -402,19 +408,21 @@ def _calculate_subtimestep_numpy( class HruSegmentFlowAdapter(Adapter): - """Adapt volumetric flows from HRUs to lateral inflows on PRMS segments/nodes. + """Adapt volumetric flows from HRUs to lateral inflows on PRMS segments. This class specifically maps from PRMS HRU outflows to PRMS segment inflows - using the parameters known to `PRMSChannel`. + using the parameters known to `PRMSChannel`. This reproduces the PRMS + variable hru_streamflow_out. This class is a subclass of :class:`Adapter` which means that it makes existing or known flows available over time (but dosent calculate a - process). This class is meant to force a stand-alone :class:`FlowGraph` runoff - outside the context of a :class:`Model`. The calculated lateral flows - (in cubic feet per second) are availble from the current_value property. + process). This class is meant to force a stand-alone :class:`FlowGraph` + runoff outside the context of a :class:`Model`. The calculated lateral + flows (in cubic feet per second) are availble from the current_value + property. See :class:`FlowGraph` for discussion and an example. - """ # noqa: E501 + """ def __init__( self, @@ -1015,3 +1023,116 @@ def _build_flow_graph_inputs( } | new_nodes_maker_dict return (params_flow_graph, node_maker_dict) + + +def prms_segment_lateral_inflow_components_to_netcdf( + control: Control, + parameters: Parameters, + input_dir: pl.Path, + nc_out_file_path: str, + output_sum: bool = False, +) -> None: + """Write to NetCDf the components of lateral flow on PRMS segments. + + This helper function takes the PRMS lateral flows from HRUs (sroff_vol, + ssres_flow_vol, and gwres_flow_vol) and maps them individually to the + PRMS segments. + + Args: + control: A Control object for the input files selected. + parameters: A Parameters object with both nhru and nsegment dimensions. + input_dir: The directory to look for the inputs files: sroff_vol.nc, + ssres_flow_vol.nc, and gwres_flow_vol.nc. + nc_out_file_path: The path of the output netcdf file. + output_sum: Also include the sum of the lateral flow components in the + output file (named 'lateral_inflow_vol'). + + """ + time = np.arange( + control.start_time, + control.end_time + control.time_step, # per arange construction + control.time_step, + ).astype("datetime64[ns]") + nhm_seg = parameters.parameters["nhm_seg"] + nhm_id = parameters.parameters["nhm_id"] + + ntime = len(time) + nsegment = len(nhm_seg) + nhru = len(nhm_id) + + zero_adapter = adapter_factory(np.zeros(nhru), "zero", control) + + input_adapters = {} + for vv in ["sroff_vol", "ssres_flow_vol", "gwres_flow_vol"]: + nc_path = input_dir / f"{vv}.nc" + input_adapters[vv] = AdapterNetcdf(nc_path, vv, control) + + exch_sroff = HruSegmentFlowAdapter( + parameters=parameters, + sroff_vol=input_adapters["sroff_vol"], + ssres_flow_vol=zero_adapter, + gwres_flow_vol=zero_adapter, + ) + + exch_ssres = HruSegmentFlowAdapter( + parameters=parameters, + sroff_vol=zero_adapter, + ssres_flow_vol=input_adapters["ssres_flow_vol"], + gwres_flow_vol=zero_adapter, + ) + + exch_gwres = HruSegmentFlowAdapter( + parameters=parameters, + sroff_vol=zero_adapter, + ssres_flow_vol=zero_adapter, + gwres_flow_vol=input_adapters["gwres_flow_vol"], + ) + + data_vars = dict( + sroff_vol=( + ["time", "nhm_seg"], + np.zeros((ntime, nsegment)) * np.nan, + ), + ssres_flow_vol=( + ["time", "nhm_seg"], + np.zeros((ntime, nsegment)) * np.nan, + ), + gwres_flow_vol=( + ["time", "nhm_seg"], + np.zeros((ntime, nsegment)) * np.nan, + ), + ) + if output_sum: + data_vars["lateral_inflow_vol"] = ( + ["time", "nhm_seg"], + np.zeros((ntime, nsegment)) * np.nan, + ) + + seg_lateral_inflow = xr.Dataset( + data_vars=data_vars, + coords=dict( + time=(["time"], time), + nhm_seg=(["nhm_seg"], nhm_seg), + ), + attrs=dict( + description="Daily lateral inflow volumes to PRMS Channel", + units="cubic feet", + ), + ) + for tt in range(control.n_times): + control.advance() + exch_sroff.advance() + exch_ssres.advance() + exch_gwres.advance() + seg_lateral_inflow["sroff_vol"][tt, :] = exch_sroff.current + seg_lateral_inflow["ssres_flow_vol"][tt, :] = exch_ssres.current + seg_lateral_inflow["gwres_flow_vol"][tt, :] = exch_gwres.current + # sum + if output_sum: + seg_lateral_inflow["lateral_inflow_vol"][tt, :] = ( + exch_sroff.current + exch_ssres.current + exch_gwres.current + ) + + seg_lateral_inflow.to_netcdf(nc_out_file_path) + + return None From a0c2b231fb3dfd1fe3b0f4ef0b420d49348039a4 Mon Sep 17 00:00:00 2001 From: James McCreight Date: Tue, 26 Nov 2024 13:53:42 -0700 Subject: [PATCH 4/4] docs for prms_segment_lateral_inflow_components_to_netcdf --- doc/api/flow_graph.rst | 1 + .../hydrology/prms_channel_flow_graph.py | 51 ++++++++++++++++++- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/doc/api/flow_graph.rst b/doc/api/flow_graph.rst index 248bd76a..f86c4b89 100644 --- a/doc/api/flow_graph.rst +++ b/doc/api/flow_graph.rst @@ -23,6 +23,7 @@ FlowGraph base classes and subclasses. See :class:`FlowGraph` for an overview of PRMSChannelFlowNodeMaker prms_channel_flow_graph_to_model_dict prms_channel_flow_graph_postprocess + prms_segment_lateral_inflow_components_to_netcdf HruSegmentFlowAdapter HruSegmentFlowExchange diff --git a/pywatershed/hydrology/prms_channel_flow_graph.py b/pywatershed/hydrology/prms_channel_flow_graph.py index d125296e..80fd33c4 100644 --- a/pywatershed/hydrology/prms_channel_flow_graph.py +++ b/pywatershed/hydrology/prms_channel_flow_graph.py @@ -1047,7 +1047,56 @@ def prms_segment_lateral_inflow_components_to_netcdf( output_sum: Also include the sum of the lateral flow components in the output file (named 'lateral_inflow_vol'). - """ + Examples + --------- + + This example will work if you have an editable install of the repository + (not installed from pypi) and you have generated the test data for the + drb_2yr domain. + + >>> import pathlib as pl + >>> + >>> import pywatershed as pws + >>> import xarray as xr + >>> + >>> domain_dir = ( + ... pws.constants.__pywatershed_root__ / "../test_data/drb_2yr" + ... ) + >>> + >>> control = pws.Control.load_prms(domain_dir / "nhm.control") + >>> params = pws.parameters.PrmsParameters.load( + ... domain_dir / control.options["parameter_file"] + ... ) + >>> + >>> nc_out_file_path = pl.Path(".") / "segment_lateral_inflows.nc" + >>> pws.prms_segment_lateral_inflow_components_to_netcdf( + ... control, + ... params, + ... input_dir=domain_dir + ... / "output", # PRMS/pywatershed outputs are here + ... nc_out_file_path=nc_out_file_path, + ... output_sum=True, + ... ) + >>> + >>> lat_inflow_vols = xr.load_dataset(nc_out_file_path) + >>> display(lat_inflow_vols) + Size: 11MB + Dimensions: (time: 731, nhm_seg: 456) + Coordinates: + * time (time) datetime64[ns] 6kB 1979-01-01 ... 1980-12-31 + * nhm_seg (nhm_seg) int64 4kB 2048 4182 4183 ... 2045 2046 2047 + Data variables: + sroff_vol (time, nhm_seg) float64 3MB 10.35 10.96 ... 0.0 0.0 + ssres_flow_vol (time, nhm_seg) float64 3MB 0.0 0.0 ... 0.01131 0.01401 + gwres_flow_vol (time, nhm_seg) float64 3MB 4.181 1.993 ... 1.762 2.33 + lateral_inflow_vol (time, nhm_seg) float64 3MB 14.53 12.95 ... 1.774 2.344 + Attributes : + description: Daily lateral inflow volumes to PRMS Channel + units: cubic feet + + + """ # noqa: E501 + time = np.arange( control.start_time, control.end_time + control.time_step, # per arange construction