refactor: harmonize OC setup between MF6 and MFNWT; add support for e…

…mpty periods to turn off output writing

docs/source/concepts/oc.rst

@@ -0,0 +1,6 @@
+===========================================================
+MODFLOW Output Control
+===========================================================
+
+Features and Limitations
+-------------------------

mfsetup/mf6_defaults.yml

@@ -183,7 +183,7 @@ chd:
 oc:
   head_fileout_fmt: '{}.hds'
   budget_fileout_fmt: '{}.cbc'
-  # example of using MODFLOW 6 text input
+  # example of using MODFLOW 6-style text input
   period_options: {0: ['save head last',
                        'save budget last']
   }

mfsetup/mf6model.py

@@ -40,6 +40,7 @@
 from mfsetup.mfmodel import MFsetupMixin
 from mfsetup.mover import get_mover_sfr_package_input
 from mfsetup.obs import setup_head_observations
+from mfsetup.oc import parse_oc_period_input
 from mfsetup.tdis import add_date_comments_to_tdis
 from mfsetup.tmr import TmrNew
 from mfsetup.units import convert_time_units
@@ -911,25 +912,9 @@ def setup_oc(self):
         kwargs = self.cfg[package]
         kwargs['budget_filerecord'] = self.cfg[package]['budget_fileout_fmt'].format(self.name)
         kwargs['head_filerecord'] = self.cfg[package]['head_fileout_fmt'].format(self.name)
-        # parse both flopy and mf6-style input into flopy input
-        for rec in ['printrecord', 'saverecord']:
-            if rec in kwargs:
-                data = kwargs[rec]
-                mf6_input = {}
-                for kper, words in data.items():
-                    mf6_input[kper] = []
-                    for var, instruction in words.items():
-                        mf6_input[kper].append((var, instruction))
-                kwargs[rec] = mf6_input
-            elif 'period_options' in kwargs:
-                mf6_input = defaultdict(list)
-                for kper, options in kwargs['period_options'].items():
-                    for words in options:
-                        type, var, instruction = words.split()
-                        if type == rec.replace('record', ''):
-                            mf6_input[kper].append((var, instruction))
-                if len(mf6_input) > 0:
-                    kwargs[rec] = mf6_input
+
+        period_input = parse_oc_period_input(kwargs)
+        kwargs.update(period_input)
 
         kwargs = get_input_arguments(kwargs, mf6.ModflowGwfoc)
         oc = mf6.ModflowGwfoc(self, **kwargs)

mfsetup/mfnwt_defaults.yml

@@ -123,8 +123,8 @@ mnw:
 oc:
   head_fileout_fmt: '{}.hds'
   budget_fileout_fmt: '{}.cbc'
-  period_options: {0: ['save head',
-                       'save budget']
+  period_options: {0: ['save head last',
+                       'save budget last']
   }
 
 hyd:

mfsetup/mfnwtmodel.py

@@ -38,6 +38,7 @@
 )
 from mfsetup.mfmodel import MFsetupMixin
 from mfsetup.obs import read_observation_data, setup_head_observations
+from mfsetup.oc import parse_oc_period_input
 from mfsetup.tdis import get_parent_stress_periods, setup_perioddata_group
 from mfsetup.tmr import TmrNew
 from mfsetup.units import convert_length_units, itmuni_text, lenuni_text
@@ -324,12 +325,18 @@ def setup_oc(self):
         package = 'oc'
         print('\nSetting up {} package...'.format(package.upper()))
         t0 = time.time()
-        stress_period_data = {}
-        for i, r in self.perioddata.iterrows():
-            stress_period_data[(r.per, r.nstp -1)] = r.oc
-
+        #stress_period_data = {}
+        #for i, r in self.perioddata.iterrows():
+        #    stress_period_data[(r.per, r.nstp -1)] = r.oc
+
+        # use stress_period_data if supplied
+        # (instead of period_input defaults)
+        if 'stress_period_data' in self.cfg['oc']:
+            del self.cfg['oc']['period_options']
         kwargs = self.cfg['oc']
-        kwargs['stress_period_data'] = stress_period_data
+        period_input = parse_oc_period_input(kwargs, nstp=self.perioddata.nstp,
+                                             output_fmt='mfnwt')
+        kwargs.update(period_input)
         kwargs = get_input_arguments(kwargs, fm.ModflowOc)
         oc = fm.ModflowOc(model=self, **kwargs)
         print("finished in {:.2f}s\n".format(time.time() - t0))

mfsetup/tdis.py

@@ -169,7 +169,7 @@ def setup_perioddata_group(start_date_time, end_date_time=None,
     oc_saverecord : dict
         Dictionary with zero-based stress periods as keys and output control options as values.
         Similar to MODFLOW-6 input, the information specified for a period will
-        continue to apply until information for another perior is specified.
+        continue to apply until information for another period is specified.
 
     Returns
     -------

mfsetup/tests/data/shellmound.yml

@@ -318,7 +318,7 @@ oc:
   # MODFLOW 6-style text input can also be used
   # e.g.
   # period_options: {0: ['save head last',
-  #                     'save budget last' ]
+  #              'save budget last' ]
 
 obs:
   source_data:

mfsetup/tests/test_mf6_shellmound.py

@@ -487,15 +487,19 @@ def test_obs_setup(shellmound_model_with_dis, config):
                 break
 
 
-@pytest.mark.parametrize('options', [{'saverecord': {0: {'head': 'last',
-                                                         'budget': 'last'}}},
-                                     {'period_options': {0: ['save head last',
-                                                             'save budget last']}}
+@pytest.mark.parametrize('input', [
+    # flopy-style input
+    {'saverecord': {0: {'head': 'last', 'budget': 'last'}}},
+    # MODFLOW 6-style input
+    {'period_options': {0: ['save head last', 'save budget last']}},
+    # blank period to skip subsequent periods
+    {'period_options': {0: ['save head last', 'save budget last'],
+                        1: []}}
                                         ])
-def test_oc_setup(shellmound_model_with_dis, options):
+def test_oc_setup(shellmound_model_with_dis, input):
     cfg = {'head_fileout_fmt': '{}.hds',
            'budget_fileout_fmt': '{}.cbc'}
-    cfg.update(options)
+    cfg.update(input)
     m = shellmound_model_with_dis  # deepcopy(model)
     m.cfg['oc'] = cfg
     oc = m.setup_oc()
@@ -506,10 +510,15 @@ def test_oc_setup(shellmound_model_with_dis, options):
     options = read_mf6_block(ocfile, 'options')
     options = {k: ' '.join(v).lower() for k, v in options.items()}
     perioddata = read_mf6_block(ocfile, 'period')
+    # convert back to zero-based
+    perioddata = {k-1:v for k, v in perioddata.items()}
     assert 'fileout' in options['budget'] and '.cbc' in options['budget']
     assert 'fileout' in options['head'] and '.hds' in options['head']
-    assert 'save head last' in perioddata[1]
-    assert 'save budget last' in perioddata[1]
+    if 'saverecord' in input:
+        assert 'save head last' in perioddata[0]
+        assert 'save budget last' in perioddata[0]
+    else:
+        assert perioddata == input['period_options']
 
 
 def test_rch_setup(shellmound_model_with_dis):
@@ -728,7 +737,7 @@ def test_sfr_inflows_from_csv(model_with_sfr):
     pd.testing.assert_series_equal(left, right, check_names=False, check_freq=False)
 
 
-#@pytest.mark.xfail(reason='flopy remove_package() issue')
+@pytest.mark.xfail(reason='flopy remove_package() issue')
 def test_idomain_above_sfr(model_with_sfr):
     m = model_with_sfr
     sfr = m.sfr

mfsetup/tests/test_mf6_shellmound_rot_grid.py

@@ -117,7 +117,7 @@ def test_rotated_grid(shellmound_cfg, shellmound_simulation, mf6_exe):
     cfg['setup_grid']['yoff'] = yoff
     cfg['setup_grid']['rotation'] = rotation
     cfg['dis']['dimensions']['nrow'] = nrow
-    cfg['dis']['dimensions']['ncol'] = 25
+    cfg['dis']['dimensions']['ncol'] = ncol
 
     cfg = MF6model._parse_model_kwargs(cfg)
     kwargs = get_input_arguments(cfg['model'], mf6.ModflowGwf,

mfsetup/tests/test_oc.py

@@ -0,0 +1,64 @@
+"""Tests for the oc.py module
+"""
+import pytest
+
+from mfsetup.oc import fill_oc_stress_period_data, parse_oc_period_input
+
+
+@pytest.mark.parametrize('input,expected,output_fmt', [
+    # dictionary-based flopy-like input to (mf6) flopy-style input
+    ({'saverecord': {0: {'head': 'last', 'budget': 'last'}}},
+     {'saverecord': {0: [('head', 'last'), ('budget', 'last')]}},
+     'mf6'),
+    # mf6-style input to (mf6) flopy-style input
+    ({'period_options': {0: ['save head last', 'save budget last']}},
+     {'saverecord': {0: [('head', 'last'), ('budget', 'last')]}},
+     'mf6'),
+    # mf6-style input to flopy-style input
+    ({'period_options': {0: ['save head first', 'save budget first']}},
+     {'stress_period_data': {(0, 0): ['save head', 'save budget']}},
+     'mfnwt'),
+    ({'period_options': {0: ['save head last', 'save budget last']}},
+     {'stress_period_data': {(0, 9): ['save head', 'save budget']}},
+     'mfnwt'),
+    ({'period_options': {0: ['save head frequency 5', 'save budget frequency 5']}},
+     {'stress_period_data': {(0, 0): ['save head', 'save budget'],
+                             (0, 5): ['save head', 'save budget']}},
+     'mfnwt'),
+    ({'period_options': {0: ['save head steps 2 3', 'save budget steps 2 3']}},
+     {'stress_period_data': {(0, 2): ['save head', 'save budget'],
+                             (0, 3): ['save head', 'save budget']}},
+     'mfnwt'),
+    ({'period_options': {0: ['save head all', 'save budget all']}},
+     None, 'mfnwt'
+    ),
+    # input already in flopy format
+    ({'stress_period_data': {(0, 2): ['save head', 'save budget'],
+                             (0, 3): ['save head', 'save budget']}},
+     {},
+    'mfnwt')
+                                        ],
+                         )
+def test_parse_oc_period_input(input, expected, output_fmt):
+    results = parse_oc_period_input(input, nstp=[10], output_fmt=output_fmt)
+    # kludge for testing 'all'
+    if expected is None:
+        expected = {'stress_period_data': {(0, i): ['save head', 'save budget']
+                                           for i in range(10)}}
+    assert results == expected
+
+
+@pytest.mark.parametrize('stress_period_data,nper',
+                         [({(0, 0): ['save head', 'save budget'],
+                            (3, 0): ['save head']
+                            },
+                           5)]
+                         )
+def test_fill_oc_stress_period_data(stress_period_data, nper):
+    results = fill_oc_stress_period_data(stress_period_data, nper)
+    expected = {(0, 0): ['save head', 'save budget'],
+                (1, 0): ['save head', 'save budget'],
+                (2, 0): ['save head', 'save budget'],
+                (3, 0): ['save head'],
+                (4, 0): ['save head']}
+    assert results == expected

mfsetup/tests/test_pfl_mfnwt_inset.py

@@ -515,7 +515,6 @@ def test_lak_setup(pfl_nwt_with_dis):
 
 
 def test_nwt_setup(pfl_nwt, project_root_path):
-
     m = pfl_nwt  #deepcopy(pfl_nwt)
     m.cfg['nwt']['use_existing_file'] = project_root_path + '/mfsetup/tests/data/RGN_rjh_3_23_18.NWT'
     nwt = m.setup_nwt()
@@ -525,12 +524,23 @@ def test_nwt_setup(pfl_nwt, project_root_path):
     nwt.write_file()
 
 
-def test_oc_setup(pfl_nwt):
+@pytest.mark.parametrize('input,expected', [
+    # MODFLOW 6-style input
+    ({'period_options': {0: ['save head last', 'save budget last'],
+                         1: []}},
+     {'stress_period_data': {(0, 0): ['save head', 'save budget'],
+                            (0, 1): []}}),
+    # MODFLOW 2005-style input
+    ({'stress_period_data': {(0, 0): ['save head', 'save budget'],
+                            (0, 1): []}},
+     {'stress_period_data': {(0, 0): ['save head', 'save budget'],
+                            (0, 1): []}})
+])
+def test_oc_setup(pfl_nwt, input, expected):
     m = pfl_nwt
+    m.cfg['oc'].update(input)
     oc = m.setup_oc()
-    for (kper, kstp), words in oc.stress_period_data.items():
-        assert kstp == m.perioddata.loc[kper, 'nstp'] - 1
-        assert words == m.perioddata.loc[kper, 'oc']
+    assert oc.stress_period_data == expected['stress_period_data']
 
     # TODO: add datetime comments to OC file
 

mfsetup/tests/test_pleasant_mfnwt_inset.py

@@ -103,6 +103,12 @@ def test_wel_setup(get_pleasant_nwt_with_dis_bas6):
         assert len(spd) >= nwells0 + n_added_wels
 
 
+def test_oc_setup(get_pleasant_nwt_with_dis_bas6):
+    m = get_pleasant_nwt_with_dis_bas6  # deepcopy(model)
+    oc = m.setup_oc()
+    # oc stress period data should be filled
+    assert len(oc.stress_period_data) == m.nper
+
 def test_model_setup(full_pleasant_nwt):
     m = full_pleasant_nwt
     assert isinstance(m, MFnwtModel)