Commit 99a5792e authored by Cynthia Crowley's avatar Cynthia Crowley

Support basis arg for forcing workspace paths

parent b20b5f31
......@@ -35,17 +35,29 @@ class TestWorkspacePaths(unittest.TestCase):
)
def test_forcing(self):
# Observed data
# Observed data:
self.assertEqual(
join(self.root, 'forcing', 'forcing_1mo_201612.nc'),
self.ws.forcing(yearmon='201612', window=1)
)
self.assertEqual(
join(self.root, 'forcing_integrated', 'forcing_24mo_201612.nc'),
self.ws.forcing(yearmon='201612', window=24)
)
# Observed basins:
self.assertEqual(
join(self.root, 'basin_forcing', 'basin_forcing_1mo_201612.nc'),
self.ws.forcing(yearmon='201612', window=1, basis=Basis.BASIN)
)
# Forecast data
self.assertEqual(
join(self.root, 'forcing', 'forcing_1mo_201612_trgt201703_fcstCFS13.nc'),
self.ws.forcing(yearmon='201612', target='201703', member='CFS13', window=1)
)
def test_results(self):
......
......@@ -155,8 +155,8 @@ def time_integrate(workspace: DefaultWorkspace,
window_forecast = []
if forcing:
prev = [workspace.forcing(yearmon=x, window=1) for x in window_observed] + \
[workspace.forcing(yearmon=yearmon, member=member, target=x, window=1) for x in window_forecast]
prev = [workspace.forcing(yearmon=x, window=1, basis=basis) for x in window_observed] + \
[workspace.forcing(yearmon=yearmon, member=member, target=x, window=1, basis=basis) for x in window_forecast]
else:
prev = [workspace.results(yearmon=x, window=1, basis=basis) for x in window_observed] + \
[workspace.results(yearmon=yearmon, member=member, target=x, window=1, basis=basis) for x in window_forecast]
......@@ -175,7 +175,8 @@ def time_integrate(workspace: DefaultWorkspace,
basis=basis) if not forcing else workspace.forcing(yearmon=yearmon,
window=window,
member=member,
target=target)
target=target,
basis=basis)
)
]
......@@ -381,9 +382,7 @@ def fit_var(config: ConfigBase,
param_to_read = param
if param in config.forcing_rp_vars():
assert basis is None
infile = config.workspace().forcing(yearmon=input_range, window=window)
infile = config.workspace().forcing(yearmon=input_range, window=window, basis=basis)
elif param in config.state_rp_vars():
assert window == 1
......
......@@ -138,7 +138,10 @@ class ConfigBase(metaclass=abc.ABCMeta):
]
if basis == Basis.BASIN:
return []
return [
'T',
'Pr'
]
assert False
......@@ -180,6 +183,20 @@ class ConfigBase(metaclass=abc.ABCMeta):
if basis == Basis.BASIN:
return []
@classmethod
def forcing_integrated_vars(cls, basis: Optional[Basis]=None) -> Dict[str, List[str]]:
"""
Provides a dictionary whose keys are forcing variables to be time-integrated, and whose
values are lists of stats to apply to each of those variables (min, max, ave, etc.)
"""
if not basis or basis == Basis.BASIN:
return {
'Pr' : ['sum'],
'T' : ['ave']
}
assert False
@classmethod
def lsm_integrated_vars(cls, basis: Optional[Basis]=None) -> Dict[str, List[str]]:
"""
......@@ -228,20 +245,6 @@ class ConfigBase(metaclass=abc.ABCMeta):
"""
return [var + '_' + stat for var, stats in cls.lsm_integrated_vars(basis=basis).items() for stat in stats]
@classmethod
def forcing_integrated_vars(cls, basis: Optional[Basis]=None) -> Dict[str, List[str]]:
"""
Provides a dictionary whose keys are forcing variables to be time-integrated, and whose
values are lists of stats to apply to each of those variables (min, max, ave, etc.)
"""
if not basis:
return {
'Pr' : ['sum'],
'T' : ['ave'],
}
assert False
@classmethod
def forcing_integrated_stats(cls, basis: Optional[Basis]=None) -> Dict[str, List[str]]:
......
......@@ -303,14 +303,14 @@ def time_integrate_forcing(config:Config, window: int, *, basis: Optional[Basis]
yearmons_out = yearmons_in[window-1:]
integrate = wsim_integrate(
inputs= [read_vars(config.workspace().forcing(window=1, yearmon=date_range(yearmons_in)),
inputs= [read_vars(config.workspace().forcing(window=1, yearmon=date_range(yearmons_in), basis=basis),
*config.forcing_integrated_vars(basis=basis).keys())
],
window=window,
stats=[stat + '::' + ','.join(varname) for stat, varname in config.forcing_integrated_stats(basis=basis).items()],
attrs=[attrs.integration_window(var='*', months=window)],
output=config.workspace().forcing(yearmon=date_range(yearmons_out),
window=window)
window=window, basis=basis)
)
tag_name = config.workspace().tag('{}spinup_{}mo_forcing'.format((basis.value + '_' if basis else ''), window))
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment