From 0b154e54cb824db4ebd7379bdaac00bf0ce02bc1 Mon Sep 17 00:00:00 2001 From: Joseph Hughes Date: Mon, 15 Jun 2026 09:07:36 -0500 Subject: [PATCH 1/2] feat(pm): Add instantaneous performance measure type --- autotest/test_xd_box.py | 151 ++++++++++++++++++++++++++++++++++ mf6adj/adj.py | 11 +-- mf6adj/pm.py | 80 +++++++++++++----- mf6adj/utils/utils_pm_read.py | 24 ++++-- 4 files changed, 232 insertions(+), 34 deletions(-) diff --git a/autotest/test_xd_box.py b/autotest/test_xd_box.py index 25a0d30..cb45632 100644 --- a/autotest/test_xd_box.py +++ b/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 = ( diff --git a/mf6adj/adj.py b/mf6adj/adj.py index 3d31382..f77887e 100755 --- a/mf6adj/adj.py +++ b/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, diff --git a/mf6adj/pm.py b/mf6adj/pm.py index d659ad0..b66ccc7 100755 --- a/mf6adj/pm.py +++ b/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 diff --git a/mf6adj/utils/utils_pm_read.py b/mf6adj/utils/utils_pm_read.py index e39c4a6..eebe614 100644 --- a/mf6adj/utils/utils_pm_read.py +++ b/mf6adj/utils/utils_pm_read.py @@ -111,8 +111,8 @@ def read_adj_file( The input file structure closely follows MODFLOW 6 input files. Each performance measure is defined in a ``performance_measure`` block. Entries provide time location, spatial location, output type (``head`` or a flux - package name), form (``direct`` or ``residual``), weight, and optionally - an observed value. + package name), form (``direct``, ``residual``, or ``instantaneous``), + weight, and optionally an observed value. Example entries for a structured model: @@ -120,10 +120,22 @@ def read_adj_file( ``25 3 3 10 34 head direct 1.0 -999`` - Residual head measure: ``25 3 3 10 34 head residual 1.0 123.45`` - - At present, forms (``direct`` vs ``residual``) cannot be mixed within a - single performance measure, and types (``head`` vs flux package) also - cannot be mixed within a single performance measure. + - Instantaneous head measure: + ``25 3 3 10 34 head instantaneous 1.0 -999`` + + An ``instantaneous`` measure uses the same value as a ``direct`` + measure, but it looks at each time step on its own instead of letting + later time steps feed back into earlier ones. The result is the + sensitivity at each time step by itself, rather than the total + sensitivity over the whole run that ``direct`` gives. Time steps with no + entry are skipped, and the overall (composite) result is the average + over the observation times (each weighted by its time-step length) + instead of the plain sum used for ``direct`` and ``residual`` measures. + An ``instantaneous`` form works for both ``head`` and flux measures. + + At present, forms (``direct``, ``residual``, ``instantaneous``) cannot be + mixed within a single performance measure, and types (``head`` vs flux + package) also cannot be mixed within a single performance measure. """ hdf5_name = None if current_hdf5_name is None else pl.Path(current_hdf5_name) From 30d312d825eed68db762b1ac37dbf3188e227095 Mon Sep 17 00:00:00 2001 From: Joseph Hughes Date: Fri, 31 Jul 2026 10:49:27 -0500 Subject: [PATCH 2/2] fix(adj): make the performance measure validation deterministic The flux/residual check read an arbitrary element of the pm_type set, so a measure combining head and flux entries could pass or fail depending on set iteration order. Test for any flux pm_type instead. Correct the docstring, which claimed head and flux entries cannot be combined, and add the missing spaces in the error messages. --- mf6adj/adj.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/mf6adj/adj.py b/mf6adj/adj.py index f77887e..c604da1 100755 --- a/mf6adj/adj.py +++ b/mf6adj/adj.py @@ -171,27 +171,27 @@ def _add_performance_measure( Notes ----- - A performance measure must contain at least one entry, cannot mix - ``direct`` and ``residual`` forms, and cannot mix ``head`` entries with - flux-package entries. + A performance measure must contain at least one entry and cannot mix + ``pm_form`` values. ``head`` and flux-package entries may be combined, + but a flux entry cannot be part of a ``residual`` measure. """ if len(pm_entries) == 0: raise Exception(f"no entries found for PM {pm_name}") - pm_types = {entry.pm_type for entry in pm_entries} pm_forms = {entry.pm_form for entry in pm_entries} if len(pm_forms) > 1: raise Exception( - "performance measure" - + f"{pm_name} has mixed 'pm_forms' ({pm_forms}), " - + "this is not supported" + f"performance measure {pm_name} has mixed 'pm_forms' " + + f"({pm_forms}), this is not supported" ) - if next(iter(pm_types)) != "head" and next(iter(pm_forms)) == "residual": + pm_form = next(iter(pm_forms)) + + flux_types = {entry.pm_type for entry in pm_entries} - {"head"} + if flux_types and pm_form == "residual": raise Exception( - "performance measure" - + pm_name - + " has a flux 'pm_type' and is a " - + "residual 'pm_form', this is not supported" + f"performance measure {pm_name} has flux 'pm_types' " + + f"({flux_types}) 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")