Merge pull request #73 from hydrocomputing/stress_period_end

Fix number of stress period end states for multi-model simulations

Author: jlarsen-usgs

autotest/test_mf6_examples.py

@@ -1,4 +1,5 @@
 import modflow_devtools.models as models
+import pandas as pd
 import pytest
 
 from modflowapi import run_simulation
@@ -12,6 +13,61 @@
     "ex-prt-mp7-p04",
 ]
 
+temporary_skip_state_count = set(
+    [
+        # Skip state count tests for these model for now.
+        # There seems to be a different problem.
+        "ex-gwf-fhb",
+        "ex-gwf-nwt-p02a",
+        "ex-gwf-nwt-p02b",
+        "ex-gwe-geotherm",
+        "ex-gwe-radial-slow",
+    ]
+)
+
+
+class StateCollector:
+    """
+    Callback that collects the model state sequence.
+
+    THe attribute `states` is a list of tuples `(model_name, state_name)`,
+    where `state_name` is `callback_step.name`.
+    """
+
+    def __init__(self):
+        self.states = []
+
+    def __call__(self, sim, callback_step):
+        self.states.append((sim.model_names[0], callback_step.name))
+
+
+def count_states(states):
+    """
+    Count states per model.
+
+    The numbers of `start_state` and `end_state` should match.
+    """
+    states = pd.DataFrame(states, columns=["model", "state"])
+    counts = states.groupby(["model", "state"])[["state"]].agg("count")
+    counts.columns = ["counts"]
+    return counts
+
+
+def compare_counts(counts):
+    matches = [
+        ("iteration_start", "iteration_end"),
+        ("stress_period_start", "stress_period_end"),
+        ("timestep_start", "timestep_end"),
+    ]
+    for model_name in counts.index.levels[0]:
+        model = counts.loc[model_name]
+        for state1, state2 in matches:
+            count1, count2 = model.loc[state1].iloc[0], model.loc[state2].iloc[0]
+            assert count1 == count2, (
+                f"{state1}: {int(count1)}",
+                f"{state2}: {int(count2)}",
+            )
+
 
 @pytest.mark.parametrize("example_name", examples.keys())
 def test_example(function_tmpdir, example_name):
@@ -21,4 +77,8 @@ def test_example(function_tmpdir, example_name):
         model_relpath = model_name.rpartition(example_name + "/")[-1]
         model_workspace = function_tmpdir / model_relpath
         models.copy_to(model_workspace, model_name, verbose=True)
-        run_simulation(dll, model_workspace, lambda sim, step: None, verbose=True)
+        callback = StateCollector()
+        run_simulation(dll, model_workspace, callback=callback, verbose=True)
+        if model_name in temporary_skip_state_count:
+            counts = count_states(callback.states)
+            compare_counts(counts)

manualtest/collect_states.py

@@ -0,0 +1,70 @@
+"""Test helper for sequence of model states."""
+
+from collections import defaultdict
+
+import pandas as pd
+
+import modflowapi
+
+
+class StateCollector:
+    """
+    Callback that collects the model state sequence.
+
+    THe attribute `states` is a list of tuples `(model_name, state_name)`,
+    where `state_name` is `callback_step.name`.
+    """
+
+    def __init__(self):
+        self.states = []
+
+    def __call__(self, sim, callback_step):
+        self.states.append((sim.model_names[0], callback_step.name))
+
+
+def run(dll, sim_path):
+    """Run a model and collect its sat sequence."""
+    callback = StateCollector()
+    modflowapi.run_simulation(dll=dll, sim_path=sim_path, callback=callback)
+    return callback.states
+
+
+def visualize_states(states, limit=None):
+    """
+    Visualize the state sequence by model.
+
+    This prints one tree-like sequence of state names per model.
+    Different state groups have different indentations to visualize the
+    flow.
+    """
+    indents = {
+        "initialize": 0,
+        "finalize": 0,
+        "stress_period_start": 4,
+        "stress_period_end": 4,
+        "timestep_start": 8,
+        "timestep_end": 8,
+        "iteration_start": 12,
+        "iteration_end": 12,
+    }
+    models = defaultdict(list)
+    for model_name, state in states:
+        models[model_name].append(state)
+    for model_name, model_states in models.items():
+        print(model_name)
+        for count, state in enumerate(model_states):
+            if limit and limit == count:
+                break
+            print(" " * indents[state], state)
+
+
+def count_states(states):
+    """
+    Count states per model.
+
+    The numbers of `start_state` and `end_state` should match.
+    """
+    states = pd.DataFrame(states, columns=["model", "state"])
+    counts = states.groupby(["model", "state"])[["state"]].agg("count")
+    counts.columns = ["counts"]
+    return counts

modflowapi/extensions/runner.py

@@ -108,16 +108,15 @@ def run_simulation(dll, sim_path, callback, verbose=False, _develop=False):
 
             callback(sim_grp, Callbacks.timestep_end)
             mf6.finalize_solve(sol_id)
+            if sim_grp.nstp == sim_grp.kstp + 1:
+                callback(sim_grp, Callbacks.stress_period_end)
 
         mf6.finalize_time_step()
         current_time = mf6.get_current_time()
 
         if not has_converged:
             print(f"Simulation group: {sim_grp} DID NOT CONVERGE")
 
-        if sim_grp.nstp == sim_grp.kstp + 1:
-            callback(sim_grp, Callbacks.stress_period_end)
-
     try:
         callback(sim, Callbacks.finalize)
         mf6.finalize()