feat(pm): Add instantaneous performance measure type

autotest/test_xd_box.py

@@ -5,6 +5,7 @@
 import sys
 
 import flopy
+import h5py
 import matplotlib
 
 matplotlib.use("Agg")
@@ -627,6 +628,156 @@ def test_xd_box():
     xd_box_compare(new_d, plot_compare)
     return
 
+
+def _parse_kperkstp(group_key):
+    # group keys look like 'solution_kper:00001_kstp:00000'
+    parts = group_key.split(":")
+    kper = int(parts[1].split("_")[0])
+    kstp = int(parts[2])
+    return (kper, kstp)
+
+
+def test_xd_box_instantaneous():
+    """Check the 'instantaneous' pm_form against the 'direct' pm_form.
+
+    An 'instantaneous' measure uses the same value as 'direct', but it looks at
+    each time step on its own and ignores how later time steps feed back into
+    earlier ones. The adjoint solve runs from the last time step to the first.
+    For a model with a steady-state period followed by transient periods, and
+    the same head observation, this means:
+
+      * steady-state step   -> same as direct (nothing feeds back anyway)
+      * last step solved    -> same as direct (nothing has fed back yet)
+      * transient, not last -> differs from direct (the feedback is dropped)
+
+    The overall (composite) result for an instantaneous measure is the average
+    of the per-step results (each weighted by its time-step length), not the
+    plain sum used for direct/residual. Time steps with no observation are
+    skipped, so a measure defined on only some time steps writes and averages
+    over just those steps.
+    """
+    new_d = pl.Path("xd_box_instantaneous_test")
+    nper = 3
+    sim = setup_xd_box_model(
+        new_d,
+        nper=nper,
+        include_sto=True,
+        include_id0=True,
+        nrow=5,
+        ncol=5,
+        nlay=3,
+        q=-3,
+        icelltype=1,
+        iconvert=1,
+        newton=True,
+        delr=10.0,
+        delc=10.0,
+        full_sat_bnd=False,
+        botm=[-10, -100, -1000],
+        alt_bnd="riv",
+        sp_len=10,
+    )
+
+    # Use one head observation. Write it as a direct measure and an
+    # instantaneous measure over every period, plus an instantaneous measure
+    # over only some periods to check that empty time steps are skipped.
+    k, i, j = 1, 2, 2
+    subset_kpers = [1, 2]  # zero-based periods 1 and 2 (skip steady period 0)
+    with open(new_d / "test.adj", "w") as f:
+        f.write("\nbegin options\nhdf5_name out.h5\nend options\n\n")
+        for form in ("direct", "instantaneous"):
+            f.write(f"begin performance_measure insttest_{form}\n")
+            for kper in range(nper):
+                f.write(
+                    f"{kper + 1} 1 {k + 1} {i + 1} {j + 1} "
+                    + f"head {form} 1.0 -1e+30\n"
+                )
+            f.write("end performance_measure\n\n")
+
+        f.write("begin performance_measure insttest_subset\n")
+        for kper in subset_kpers:
+            f.write(
+                f"{kper + 1} 1 {k + 1} {i + 1} {j + 1} head instantaneous 1.0 -1e+30\n"
+            )
+        f.write("end performance_measure\n\n")
+
+    adj = mf6adj.Mf6Adj(
+        "test.adj",
+        lib_name,
+        logging_level="WARNING",
+        working_directory=new_d,
+    )
+    adj.solve_forward_model()
+    adj.solve_adjoint(csv_summary=False)
+    adj.finalize()
+
+    # time-step length (dt) and the steady/transient flag for each time step
+    dts, iss = {}, {}
+    with h5py.File(new_d / "out.h5", "r") as hf:
+        for key in hf:
+            if key.startswith("solution_kper"):
+                kk = _parse_kperkstp(key)
+                dts[kk] = float(hf[key].attrs["dt"])
+                iss[kk] = int(hf[key]["iss"][0])
+
+    def load_adjoint(name):
+        lambdas = {}
+        with h5py.File(new_d / f"adjoint_solution_{name}.h5", "r") as hf:
+            for key in hf:
+                if key.startswith("solution_kper"):
+                    lambdas[_parse_kperkstp(key)] = hf[key]["lambda"][:]
+            composite = hf["composite"]["wel6_q"][:]
+        return lambdas, composite
+
+    direct_lam, _ = load_adjoint("insttest_direct")
+    inst_lam, inst_comp = load_adjoint("insttest_instantaneous")
+
+    steps = sorted(direct_lam.keys())
+    assert len(steps) == nper
+    last_solved = steps[-1]  # last in time == first solved going backward
+
+    saw_transient_diff = False
+    for kk in steps:
+        diff = np.max(np.abs(direct_lam[kk] - inst_lam[kk]))
+        should_match = iss[kk] == 1 or kk == last_solved
+        if should_match:
+            assert diff < 1e-12, (
+                f"step {kk} (iss={iss[kk]}) should match direct, max diff {diff}"
+            )
+        else:
+            assert diff >= 1e-12, (
+                f"transient step {kk} should differ from direct (coupling dropped)"
+            )
+            saw_transient_diff = True
+    assert saw_transient_diff, "no transient step exercised the time decoupling"
+
+    # the overall result is the per-step results averaged over time, each
+    # weighted by its time-step length
+    wsum = sum(dts[kk] for kk in steps)
+    expected_comp = sum(inst_lam[kk] * dts[kk] for kk in steps) / wsum
+    assert np.max(np.abs(inst_comp - expected_comp)) < 1e-10, (
+        "instantaneous composite is not the time-weighted mean of per-step states"
+    )
+
+    # the subset measure should only solve and write the periods it is defined
+    # on, and its overall result should be the average over just those steps
+    sub_lam, sub_comp = load_adjoint("insttest_subset")
+    sub_steps = sorted(sub_lam.keys())
+    assert sub_steps == [(kper, 0) for kper in subset_kpers], (
+        f"subset measure solved unexpected steps: {sub_steps}"
+    )
+    sub_wsum = sum(dts[kk] for kk in sub_steps)
+    sub_expected = sum(sub_lam[kk] * dts[kk] for kk in sub_steps) / sub_wsum
+    assert np.max(np.abs(sub_comp - sub_expected)) < 1e-10, (
+        "subset composite is not the time-weighted mean over its observation steps"
+    )
+    # the subset's per-step states match the full measure's at shared steps
+    for kk in sub_steps:
+        assert np.max(np.abs(sub_lam[kk] - inst_lam[kk])) < 1e-12, (
+            f"subset and full instantaneous states differ at shared step {kk}"
+        )
+
+
 def test_xd_box_unstruct():
     # workflow flags
     include_id0 = (

mf6adj/adj.py

@@ -186,12 +186,12 @@ def _add_performance_measure(
                 + f"{pm_name} has mixed 'pm_forms' ({pm_forms}), "
                 + "this is not supported"
             )
-        if next(iter(pm_types)) != "head" and next(iter(pm_forms)) != "direct":
+        if next(iter(pm_types)) != "head" and next(iter(pm_forms)) == "residual":
             raise Exception(
                 "performance measure"
                 + pm_name
-                + " has a flux 'pm_form' and is a "
-                + "residual 'pm_type', this is not supported"
+                + " has a flux 'pm_type' and is a "
+                + "residual 'pm_form', this is not supported"
             )
         if pm_name in [pm._name for pm in self._performance_measures]:
             raise Exception(f"PM {pm_name} multiply defined")
@@ -901,7 +901,6 @@ def solve_forward_model(
     def solve_adjoint(
         self,
         hdf5_adjoint_solution_fname: Optional[PathLike] = None,
-        skip_solve: bool = False,
         csv_summary: bool = False,
         linear_solver=None,
         linear_solver_kwargs: dict = {},
@@ -921,9 +920,6 @@ def solve_adjoint(
         hdf5_adjoint_solution_fname : PathLike, optional
             HDF5 file to write the adjoint solution. If omitted, a default
             name based on the performance-measure name is used.
-        skip_solve : bool, optional
-            Skip the adjoint solve for time steps with no performance-measure
-            entries.
         csv_summary : bool, optional
             Write a CSV summary of the sensitivity information.
         linear_solver : str or callable, optional
@@ -982,7 +978,6 @@ def solve_adjoint(
                 df = pm.solve_adjoint(
                     hdf5_forward_solution_fname=self._hdf5_name,
                     hdf5_adjoint_solution_fname=hdf5_adjoint_solution_fname,
-                    skip_solve=skip_solve,
                     csv_summary=csv_summary,
                     linear_solver=linear_solver,
                     linear_solver_kwargs=linear_solver_kwargs,

mf6adj/pm.py

@@ -118,10 +118,10 @@ def __init__(
         self.pm_type = pm_type.lower().strip()
 
         self.pm_form = pm_form.lower().strip()
-        if self.pm_form not in ["direct", "residual"]:
+        if self.pm_form not in ["direct", "residual", "instantaneous"]:
             raise Exception(
-                "PerfMeasRecord.pm_form must be 'direct' or 'residual', "
-                + f"not '{self.pm_form}'"
+                "PerfMeasRecord.pm_form must be 'direct', 'residual', or "
+                + f"'instantaneous', not '{self.pm_form}'"
             )
 
     def __repr__(self) -> str:
@@ -228,7 +228,7 @@ def _performance_measure_forward(self, head_dict, sp_package_dict) -> float:
         result = 0.0
         for pfr in self._entries:
             if pfr.pm_type == "head":
-                if pfr.pm_form == "direct":
+                if pfr.pm_form in ("direct", "instantaneous"):
                     result += pfr.weight * head_dict[pfr.kperkstp][pfr.inode]
                 elif pfr.pm_form == "residual":
                     result += (
@@ -245,7 +245,7 @@ def _performance_measure_forward(self, head_dict, sp_package_dict) -> float:
                                     kk_d["node"] == pfr.inode + 1
                                     and kk_d["packagename"] == pfr.pm_type
                                 ):
-                                    if pfr.pm_form == "direct":
+                                    if pfr.pm_form in ("direct", "instantaneous"):
                                         result += pfr.weight * kk_d["simval"]
                                     elif pfr.pm_form == "residual":
                                         result += (
@@ -391,7 +391,6 @@ def solve_adjoint(
         self,
         hdf5_forward_solution_fname: PathLike,
         hdf5_adjoint_solution_fname: Optional[PathLike] = None,
-        skip_solve: bool = False,
         csv_summary: bool = False,
         linear_solver=None,
         linear_solver_kwargs: Optional[dict] = None,
@@ -419,9 +418,6 @@ def solve_adjoint(
             HDF5 file to write the adjoint solution. If omitted, a default
             name based on the performance-measure name is used. If the target
             file already exists, it is removed before writing.
-        skip_solve : bool, optional
-            Skip the adjoint solve for time steps with no performance-measure
-            entries.
         csv_summary : bool, optional
             Write a CSV summary of the sensitivity information beside the
             adjoint HDF5 output file.
@@ -597,6 +593,16 @@ def solve_adjoint(
                 has_flux_pm = True
                 break
 
+        # An "instantaneous" measure looks at each time step on its own and
+        # ignores how later time steps would otherwise feed back into earlier
+        # ones. The result is the sensitivity at each time step by itself. All
+        # entries in a measure share the same form, so one flag covers the
+        # whole measure.
+        is_instantaneous = self._entries[0].pm_form == "instantaneous"
+        # running total of the time-step weights, used below to average the
+        # result of an instantaneous measure
+        wsum = 0.0
+
         comp_welq_sens = np.zeros(nnodes)
         comp_rch_sens = np.zeros((nnodes))
 
@@ -609,7 +615,12 @@ def solve_adjoint(
                         comp_bnd_results[pname + "_" + aname] = np.zeros(nnodes)
 
         for itime, kk in enumerate(kperkstp[::-1]):
-            if skip_solve and not self.__pm_available(kk):
+            # For an instantaneous measure, a time step with no observation
+            # adds nothing, so skip it. The average below then uses only the
+            # time steps that have an observation. Direct and residual measures
+            # pass information from one time step to the next, so every step
+            # must be solved and none are skipped.
+            if is_instantaneous and not self.__pm_available(kk):
                 continue
 
             data = {}
@@ -642,12 +653,26 @@ def solve_adjoint(
             data["dfdh"] = dfdh
             iss = hdf[sol_key]["iss"][0]
 
-            if iss == 0:  # transient
-                # get the derv of RHS WRT head
+            # Weight for this time step, used when combining the per-step
+            # results below. This does not change the old behavior: for direct
+            # and residual measures w is 1.0, so "value * w" is exactly the same
+            # as "value" and the divide by wsum at the end is skipped. Only an
+            # instantaneous measure uses a real weight (the length of the time
+            # step, dt) so its result can be averaged over time.
+            w = float(hdf[sol_key].attrs["dt"]) if is_instantaneous else 1.0
+            wsum += w
+
+            if iss == 0 and not is_instantaneous:  # transient
+                # drhsdh is the storage term that links this time step to the
+                # next. Multiplying it by the later step's result (lamb) carries
+                # that result backward in time into this step.
                 drhsdh = hdf[sol_key]["drhsdh"][:]
                 data["drhsdh"] = drhsdh
                 rhs = (drhsdh * lamb) - dfdh
             else:
+                # Either steady state or an instantaneous measure. In both
+                # cases there is no carryover from the later time step, so each
+                # time step is solved on its own.
                 rhs = -dfdh
 
             self.logger.logger.debug(
@@ -951,8 +976,8 @@ def solve_adjoint(
 
             data["k11"] = k_sens
             data["k33"] = k33_sens
-            comp_k_sens += k_sens
-            comp_k33_sens += k33_sens
+            comp_k_sens += k_sens * w
+            comp_k33_sens += k33_sens * w
             self.logger.logger.debug(
                 (
                     "Formulating lam_dresdk_h took: "
@@ -970,7 +995,7 @@ def solve_adjoint(
                 else:
                     ss_sens = np.zeros_like(lamb)
                 data["ss"] = ss_sens
-                comp_ss_sens += ss_sens
+                comp_ss_sens += ss_sens * w
                 self.logger.logger.debug(
                     (
                         "Formulating storage took: "
@@ -979,8 +1004,8 @@ def solve_adjoint(
                 )
 
             data["wel6_q"] = lamb
-            comp_welq_sens += lamb
-            comp_rch_sens += lamb
+            comp_welq_sens += lamb * w
+            comp_rch_sens += lamb * w
             data["rch6_recharge"] = lamb
 
             for ptype, pnames in gwf_package_dict.items():
@@ -1001,11 +1026,13 @@ def solve_adjoint(
                         sens_level, sens_cond = self.__lam_drhs_dbnd(
                             lamb, head, sp_bnd_dict, has_flux_pm
                         )
-                        comp_bnd_results[pname + "_" + bnd_dict[ptype][0]] += sens_level
+                        comp_bnd_results[pname + "_" + bnd_dict[ptype][0]] += (
+                            sens_level * w
+                        )
                         data[pname + "_" + bnd_dict[ptype][0]] = sens_level
                         if len(bnd_dict[ptype]) > 1:
                             comp_bnd_results[pname + "_" + bnd_dict[ptype][1]] += (
-                                sens_cond
+                                sens_cond * w
                             )
                             data[pname + "_" + bnd_dict[ptype][1]] = sens_cond
                         self.logger.logger.debug(
@@ -1043,6 +1070,19 @@ def solve_adjoint(
                 logger=self.logger.logger,
             )
         self.logger.logger.info("Formulate composite sensitivities")
+        # An instantaneous measure handles each time step on its own, so its
+        # overall result is the average of the per-step results (each weighted
+        # by its time-step length) instead of a plain sum. Direct and residual
+        # measures keep the plain sum.
+        if is_instantaneous and wsum > 0.0:
+            comp_k_sens /= wsum
+            comp_k33_sens /= wsum
+            comp_welq_sens /= wsum
+            comp_rch_sens /= wsum
+            if has_sto:
+                comp_ss_sens /= wsum
+            for name in comp_bnd_results:
+                comp_bnd_results[name] /= wsum
         data = {}
         data["k11"] = comp_k_sens
         data["k33"] = comp_k33_sens
@@ -1546,7 +1586,7 @@ def __dfdh(self, kk, sol_dataset) -> np.ndarray:
 
         for pfr in relevant_entries:
             if pfr.pm_type == "head":
-                if pfr.pm_form == "direct":
+                if pfr.pm_form in ("direct", "instantaneous"):
                     dfdh[pfr.inode] = pfr.weight
                 elif pfr.pm_form == "residual":
                     # Scalar math is fine here, but make sure head is a numpy array