diff --git a/mfsetup/lakes.py b/mfsetup/lakes.py index dbb57c17..2d71cb35 100644 --- a/mfsetup/lakes.py +++ b/mfsetup/lakes.py @@ -134,22 +134,44 @@ def get_littoral_profundal_zones(lakzones): return zones -def get_flux_variable_from_config(variable, config, nper): +def get_flux_variable_from_config(variable, config, nlakes, nper): + # MODFLOW 6 models + if 'perioddata' in config: + config = config['perioddata'] data = config[variable] # copy to all stress periods # single scalar value if np.isscalar(data): - data = [data] * nper + values = [data] * nper * nlakes # flat list of global values by stress period elif isinstance(data, list) and np.isscalar(data[0]): + values = data.copy() if len(data) < nper: for i in range(nper - len(data)): - data.append(data[-1]) + values.append(data[-1]) + # repeat stress period data for all lakes + for lake_no in range(1, nlakes): + values += data + elif isinstance(data, dict): + values = [] + lake_numbers = list(sorted(data.keys())) + if not len(lake_numbers) == nlakes and all(np.diff(lake_numbers) == 1): + raise ValueError( + f"Lake Package {variable}: Dictionary-style input " + "requires consecutive lake numbers and an entry for each lake") + for lake_no in lake_numbers: + lake_values = data[lake_no] + if np.isscalar(lake_values): + lake_values = [lake_values] * nper + elif len(lake_values) < nper: + for i in range(nper - len(lake_values)): + lake_values.append(lake_values[-1]) + values += lake_values # dict of lists, or nested dict of values by lake, stress_period else: - raise NotImplementedError('Direct input of {} by lake.'.format(variable)) - assert isinstance(data, list) - return data + raise ValueError(f"Invalid Lake Package {variable}: input:\n{config}") + assert isinstance(values, list) + return values def setup_lake_info(model): @@ -239,23 +261,17 @@ def setup_lake_fluxes(model, block='lak'): df = pd.DataFrame(np.zeros((nlakes * model.nper, len(columns)), dtype=float), columns=columns) df['per'] = list(range(model.nper)) * nlakes + # lake dataframe sorted by lake id and then period df['lak_id'] = sorted(model.lake_info['lak_id'].tolist() * model.nper) # option 1; precip and evaporation specified directly # values are assumed to be in model units for variable in variables: - # MODFLOW 2005 models - if variable in model.cfg[block]: - config = model.cfg[block] - # MODFLOW 6 Lake Package - elif variable in model.cfg[block].get('perioddata', {}): - config = model.cfg[block]['perioddata'] - # no direct perioddata input - else: - continue - values = get_flux_variable_from_config(variable, config, model.nper) - # repeat values for each lake - df[variable] = values * len(model.lake_info['lak_id']) + # MODFLOW 2005 models or MODFLOW 6 Lake Package + if variable in model.cfg[block] or\ + variable in model.cfg[block].get('perioddata', {}): + values = get_flux_variable_from_config(variable, model.cfg[block], nlakes, model.nper) + df[variable] = values # option 2; precip and temp specified from PRISM output # compute evaporation from temp using Hamon method diff --git a/mfsetup/tests/test_lakes.py b/mfsetup/tests/test_lakes.py index c7269641..25dbd0af 100644 --- a/mfsetup/tests/test_lakes.py +++ b/mfsetup/tests/test_lakes.py @@ -194,15 +194,41 @@ def test_get_horizontal_connections(tmpdir, connection_info): assert ncon == np.sum(connections['k'] == k) -def test_get_flux_variable_from_config(get_pleasant_mf6_with_dis): - model = get_pleasant_mf6_with_dis - model.setup_lak() +@pytest.mark.parametrize('modflow_version', ('mf2005', 'mf6')) +@pytest.mark.parametrize('config, nlakes, nper, expected', ( + # single global value; gets repeated to all lakes and periods + ({'precipitation': 0.01}, 1, 1, [0.01]), + ({'precipitation': 0.01}, 2, 1, [0.01] * 2), + ({'precipitation': 0.01}, 2, 2, [0.01] * 4), + # steady value for each lake + ({'precipitation': {0: 0.01, 1: 0.02}}, 2, 2, [0.01]*2 + [0.02]*2), + ({'precipitation': [0.01, 0.02]}, 2, 2, [0.01, 0.02]*2), + #pytest.param({'precipitation': [0.01, 0.02]}, 2, 2, None, + # marks=pytest.mark.xfail(reason='multiple lakes require dictionary input')), + # time-varying values for one or more lakes + ({'precipitation': [0.01, 0.011]}, 1, 2, [0.01, 0.011]), + ({'precipitation': [0.01, 0.011]}, 1, 3, [0.01, 0.011, 0.011]), + pytest.param({'precipitation': [0.01, 0.011, 0.012]}, 1, 2, None, + marks=pytest.mark.xfail(reason='more values than periods')), + ({'precipitation': {0: [0.01, 0.011], + 1: [0.02, 0.021]}}, 2, 2, [0.01, 0.011, 0.02, 0.021]), + # last value gets repeated to remaining stress periods + ({'precipitation': {0: [0.01, 0.011], + 1: [0.02, 0.021]}}, 2, 3, + [0.01, 0.011, 0.011, 0.02, 0.021, 0.021]), + pytest.param({'precipitation': {0: [0.01, 0.011], + 1: [0.02, 0.021]}}, 3, 3, None, + marks=pytest.mark.xfail(reason='no values for lake 3')), + pytest.param({'precipitation': {0: [0.01, 0.011], + 2: [0.02, 0.021]}}, 2, 3, None, + marks=pytest.mark.xfail(reason='no values for lake 1 or nonconsecutive lake numbers')), + +)) +def test_get_flux_variable_from_config(modflow_version, config, nlakes, nper, expected + ): variable = 'precipitation' - config = { - 'perioddata': { - variable : [0.01] - } - } - results = get_flux_variable_from_config(variable, config) - # repeat values for each lake - df[variable] = results * len(model.lake_info['lak_id']) + if modflow_version == 'mf6': + config = {'perioddata': config} + results = get_flux_variable_from_config(variable, config, nlakes, nper) + # repeat values for each lake + assert results == expected