Skip to content

Commit

Permalink
clean up model docs and test its examples
Browse files Browse the repository at this point in the history
  • Loading branch information
@jmccreight
jmccreight committed Dec 15, 2023
1 parent 9765eaf commit 01fe647
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 216 deletions.
20 changes: 4 additions & 16 deletions pywatershed/atmosphere/prms_atmosphere.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,12 +82,6 @@ class PRMSAtmosphere(Process):
radiation on a horizontal plane
verbose: Print extra information or not?
netcdf_output_dir: A directory to write netcdf outpuf files
netcdf_output_vars: A list of variables to output via netcdf.
netcdf_separate_files: Separate or a single netcdf output file
load_n_time_batches: How often to load from disk (not-implemented?)
n_time_chunk: the inverse of load_n_time_batches, the number of
times in a chunk/batch (implemented?)
"""

Expand All @@ -102,20 +96,14 @@ def __init__(
soltab_potsw: adaptable,
soltab_horad_potsw: adaptable,
verbose: bool = False,
# netcdf_output_dir: [str, pl.Path] = None,
# netcdf_output_vars: list = None,
# netcdf_separate_files: bool = None,
# from_file_dir: [str, pl.Path] = None,
n_time_chunk: int = -1,
load_n_time_batches: int = 1,
):
# Defering handling batch handling of time chunks but self.n_time_chunk
# is a dimension used in the metadata/variables dimensions.
# TODO: make time chunking options work (esp with output)
if n_time_chunk <= 0:
self.n_time_chunk = control.n_times
else:
self.n_time_chunk = n_time_chunk
# if n_time_chunk <= 0:
# self.n_time_chunk = control.n_times
# else:
# self.n_time_chunk = n_time_chunk

# Initialize full time with nans
self._time = np.full(control.n_times, nan, dtype="datetime64[s]")
Expand Down
2 changes: 0 additions & 2 deletions pywatershed/atmosphere/prms_solar_geometry.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,6 @@ class PRMSSolarGeometry(Process):
verbose: Print extra information or not?
from_prms_file: Load from a PRMS output file?
from_nc_files_dir: [str, pl.Path] = None,
load_n_time_batches: How often to load from disk (not-implemented?)
"""

Expand All @@ -66,7 +65,6 @@ def __init__(
verbose: bool = False,
from_prms_file: [str, pl.Path] = None,
from_nc_files_dir: [str, pl.Path] = None,
load_n_time_batches: int = 1,
):
# self._time is needed by Process for timeseries arrays
# TODO: this is redundant because the parameter doy is set
Expand Down
209 changes: 11 additions & 198 deletions pywatershed/base/model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
from copy import deepcopy
from datetime import datetime
from typing import Union
from warnings import warn

from tqdm.auto import tqdm

Expand Down Expand Up @@ -134,33 +133,16 @@ class Model:
Construct a PRMS-legacy based model:
..
import pywatershed as pws
test_data_dir = pws.constants.__pywatershed_root__ / "../test_data"
domain_dir = test_data_dir / "drb_2yr"
# A PRMS-native control file
control_file = domain_dir / "control.test"
# PRMS-native parameter file
parameter_file = domain_dir / "myparam.param"
control = pws.Control.load(control_file)
control.options['input_dir'] = domain_dir / "output"
params = pws.parameters.PrmsParameters.load(parameter_file)
model_procs = [pws.PRMSGroundwater, pws.PRMSChannel,]
model = pws.Model(
model_procs,
control=control,
parameters=params,
)
model.run()
>>> import pywatershed as pws
>>> test_data_dir = pws.constants.__pywatershed_root__ / "../test_data"
>>> domain_dir = test_data_dir / "drb_2yr"
>>> # A PRMS-native control file
>>> control_file = domain_dir / "control.test"
>>> # PRMS-native parameter file
>>> parameter_file = domain_dir / "myparam.param"
>>> control = pws.Control.load(control_file)
>>> control = pws.Control.load_prms(
... control_file, warn_unused_options=False
... )
>>> control.options["input_dir"] = domain_dir / "output"
>>> params = pws.parameters.PrmsParameters.load(parameter_file)
>>> model_procs = [
Expand All @@ -181,61 +163,16 @@ class Model:
Construct a model the pywatershed-centric way, in memory:
..
import pywatershed as pws
test_data_dir = pws.constants.__pywatershed_root__ / "../test_data"
domain_dir = test_data_dir / "drb_2yr"
dis_hru = pws.Parameters.from_netcdf(
domain_dir / "parameters_dis_hru.nc", encoding=False
)
control_file = domain_dir / "control.test"
control = pws.Control.load(control_file)
control.options['input_dir'] = domain_dir
params = {}
for proc in ["SolarGeometry", "Atmosphere", "Canopy", "Snow"]:
param_file = domain_dir / f"parameters_PRMS{proc}.nc"
params[proc.lower()] = pws.Parameters.from_netcdf(param_file)
model_dict = {
'control': control,
'dis_hru': dis_hru,
'model_order': [
'prmssolargeometry',
'prmsatmosphere',
'prmscanopy',
'prmssnow',
],
'prmssolargeometry': {
'class': pws.PRMSSolarGeometry,
'parameters': params['solargeometry'],
'dis': 'dis_hru',
},
'prmsatmosphere': {
'class': pws.PRMSAtmosphere,
'parameters': params['atmosphere'],
'dis': 'dis_hru',
},
'prmscanopy': {
'class': pws.PRMSCanopy,
'parameters': params['canopy'],
'dis': 'dis_hru',
},
'prmssnow': {
'class': pws.PRMSSnow,
'parameters': params['snow'],
'dis': 'dis_hru',
}
}
model = pws.Model(model_dict)
model.run()
>>> import pywatershed as pws
>>> test_data_dir = pws.constants.__pywatershed_root__ / "../test_data"
>>> domain_dir = test_data_dir / "drb_2yr"
>>> dis_hru = pws.Parameters.from_netcdf(
... domain_dir / "parameters_dis_hru.nc", encoding=False
... )
>>> control_file = domain_dir / "control.test"
>>> control = pws.Control.load(control_file)
>>> control = pws.Control.load_prms(
... control_file, warn_unused_options=False
... )
>>> control.options["input_dir"] = domain_dir
>>> params = {}
>>> for proc in ["SolarGeometry", "Atmosphere", "Canopy", "Snow"]:
Expand Down Expand Up @@ -276,68 +213,11 @@ class Model:
PRMSCanopy jit compiling with numba
PRMSSnow jit compiling with numba
>>> model.run()
0%| | 0/731 [00:00<?, ?it/s]
pywatershed/pywatershed/hydrology/prms_snow.py:1086: NumbaExperimentalFeatureWarning: Use of isinstance() detected. This is an experimental feature.
through_rain[:] = np.where(wh_through, net_rain, zero)
100%|███████████████████████████████████████████████████████████| 731/731 [00:07<00:00, 96.76it/s]
100%|███████████████████████████████████████████████████████████| 731/731 [00:07<00:00, 60.76it/s]
model.run(): finalizing
Construct a model the pywatershed-centric way, from a yaml file definition:
..
import yaml
import pywatershed as pws
test_data_dir = pws.constants.__pywatershed_root__ / "../test_data"
domain_dir = test_data_dir / "drb_2yr"
control = {
"start_time": "1979-01-01T00:00:00",
"end_time": "1980-12-31T00:00:00",
"time_step": 24,
"time_step_units": "h",
"verbosity": 0,
"budget_type": "warn",
"init_vars_from_file": 0,
"input_dir": str(domain_dir),
}
control_file = domain_dir / "example_control.yml"
model_dict = {
"control": str(control_file),
"dis_hru": "parameters_dis_hru.nc",
"dis_both": "parameters_dis_both.nc",
"solargeometry": {
"class": "PRMSSolarGeometry",
"parameters": "parameters_PRMSSolarGeometry.nc",
"dis": "dis_hru",
},
"atmosphere": {
"class": "PRMSAtmosphere",
"parameters": "parameters_PRMSAtmosphere.nc",
"dis": "dis_hru",
},
"canopy": {
"class": "PRMSCanopy",
"parameters": "parameters_PRMSCanopy.nc",
"dis": "dis_hru",
},
"snow": {
"class": "PRMSSnow",
"parameters": "parameters_PRMSSnow.nc",
"dis": "dis_hru",
},
"model_order": ["solargeometry", "atmosphere", "canopy", "snow"],
}
model_dict_file = domain_dir / "example_model_dict.yml"
dump_dict = {control_file: control, model_dict_file: model_dict}
for key, val in dump_dict.items():
with open(key, "w") as file:
documents = yaml.dump(val, file)
model = pws.Model.from_yaml(model_dict_file)
model.run()
control_file.unlink()
model_dict_file.unlink()
>>> import yaml
>>> import pywatershed as pws
>>> test_data_dir = pws.constants.__pywatershed_root__ / "../test_data"
Expand All @@ -349,10 +229,9 @@ class Model:
... "time_step_units": "h",
... "verbosity": 0,
... "budget_type": "warn",
... "init_vars_from_file": 0,
... "input_dir": str(domain_dir),
... }
>>> control_file = domain_dir / "example_control.yml"
>>> control_file = domain_dir / "example_control.yaml"
>>> model_dict = {
... "control": str(control_file),
... "dis_hru": "parameters_dis_hru.nc",
Expand All @@ -379,23 +258,17 @@ class Model:
... },
... "model_order": ["solargeometry", "atmosphere", "canopy", "snow"],
... }
>>> model_dict_file = domain_dir / "example_model_dict.yml"
>>> model_dict_file = domain_dir / "example_model_dict.yaml"
>>> dump_dict = {control_file: control, model_dict_file: model_dict}
>>> for key, val in dump_dict.items():
... with open(key, "w") as file:
... documents = yaml.dump(val, file)
...
>>> model = pws.Model.from_yml(model_dict_file)
>>> model = pws.Model.from_yaml(model_dict_file)
PRMSCanopy jit compiling with numba
PRMSSnow jit compiling with numba
>>> model.run()
0%| | 0/731 [00:00<?, ?it/s]
pywatershed/pywatershed/hydrology/prms_snow.py:1086: NumbaExperimentalFeatureWarning: Use of isinstance() detected. This is an experimental feature.
through_rain[:] = np.where(wh_through, net_rain, zero)
0%| | 1/731 [00:05<1:08:27, 5.63s/it]
/pywatershed/pywatershed/base/budget.py:317: UserWarning: The flux unit balance not equal to the change in unit storage: PRMSSnow
warn(msg, UserWarning)
100%|██████████████████████████████████████████████████████████| 731/731 [00:06<00:00, 119.95it/s]
100%|████████████████████████████████████████████████████████████| 731/731 [00:07<00:00, 92.30it/s]
model.run(): finalizing
>>> control_file.unlink()
>>> model_dict_file.unlink()
Expand Down Expand Up @@ -784,16 +657,6 @@ def initialize_netcdf(
if not self._found_input_files:
self._find_input_files()

# (
# output_dir,
# output_vars,
# separate_files,
# ) = self._reconcile_nc_args_w_control_opts(
# output_dir, output_vars, separate_files
# )

# self._netcdf_dir = pl.Path(output_dir)

for cls in self.process_order:
self.processes[cls].initialize_netcdf(
output_dir=output_dir,
Expand Down Expand Up @@ -829,9 +692,6 @@ def run(
n_time_steps: the number of timesteps to run
output_vars: the vars to output to the netcdf_dir
"""
# Can supply options ton initialize netcdf on .run but not with
# .advance. However, the first advance takes care of finding
# the input files.
if netcdf_dir or (
not self._netcdf_initialized
and self._default_nc_out_dir is not None
Expand Down Expand Up @@ -885,50 +745,3 @@ def finalize(self):
for cls in self.process_order:
self.processes[cls].finalize()
return

# def _reconcile_nc_args_w_control_opts(
# self, output_dir, output_vars, separate_files
# ):
# # can treat the other args but they are not yet in the available opts
# arg_opt_name_map = {
# "output_dir": "netcdf_output_dir",
# "output_vars": "netcdf_output_var_names",
# "separate_files": "netcdf_output_separate_files",
# }

# args = {
# "output_dir": output_dir,
# "output_vars": output_vars,
# "separate_files": separate_files,
# }

# for vv in args.keys():
# arg_val = args[vv]
# opt_name = arg_opt_name_map[vv]
# opts = self.control.options
# if opt_name in opts.keys():
# opt_val = opts[opt_name]
# else:
# opt_val = None

# # set the arg vals to return

# if opt_val is None and arg_val is None:
# pass

# elif opt_val is None:
# pass

# elif arg_val is None:
# args[vv] = opt_val

# else:
# msg = (
# f"control.option '{opt_name}' being superceeded by "
# f"model.initialize_netcdf argument {vv}"
# )
# # TODO: should this edit control? and then model writes control
# # at the end of run to the output dir?
# warn(msg)

# return args["output_dir"], args["output_vars"], args["separate_files"]

0 comments on commit 01fe647

Please sign in to comment.