Fix lak direct flux input

mfsetup/lakes.py

@@ -134,22 +134,43 @@ def get_littoral_profundal_zones(lakzones):
     return zones
 
 
-def get_flux_variable_from_config(variable, model):
-    data = model.cfg['lak'][variable]
+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] * model.nper
+        values = [data] * nper * nlakes
     # flat list of global values by stress period
     elif isinstance(data, list) and np.isscalar(data[0]):
-        if len(data) < model.nper:
-            for i in range(model.nper - len(data)):
-                data.append(data[-1])
+        values = data.copy()
+        if len(data) < nper:
+            for i in range(nper - len(data)):
+                values.append(data[-1])
+        # repeat stress period data for all lakes
+        values = values * nlakes
+    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,15 +260,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:
-        if variable in model.cfg[block]:
-            values = get_flux_variable_from_config(variable, model)
-            # 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

mfsetup/tests/test_lakes.py

@@ -10,6 +10,7 @@
 from mfsetup.fileio import load_array
 from mfsetup.lakes import (
     PrismSourceData,
+    get_flux_variable_from_config,
     get_horizontal_connections,
     setup_lake_connectiondata,
     setup_lake_info,
@@ -191,3 +192,43 @@ def test_get_horizontal_connections(tmpdir, connection_info):
 
         ncon = np.sum(np.abs(sobel_x) > 1) + np.sum(np.abs(sobel_y) > 1)
         assert ncon == np.sum(connections['k'] == k)
+
+
+@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'
+    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