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/changelog/CHANGELOG.md b/changelog/CHANGELOG.md index 727e8f9..696802e 100644 --- a/changelog/CHANGELOG.md +++ b/changelog/CHANGELOG.md @@ -4,6 +4,19 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +### Breaking changes + +- `Mf6Adj.solve_adjoint()` and `PerfMeas.solve_adjoint()` no longer accept + `skip_solve`. The flag applied to every performance measure form, but a + transient `direct` or `residual` measure carries information backward from + one time step to the next, so skipping a time step returned incorrect + sensitivities with no indication that anything was wrong. Time steps with no + entries are now skipped automatically, and only for the `instantaneous` form, + where each time step is solved on its own and skipping is correct. + ## [1.1.0] - 2026-06-02 ### Changes diff --git a/mf6adj/adj.py b/mf6adj/adj.py index 3d31382..5b68aa3 100755 --- a/mf6adj/adj.py +++ b/mf6adj/adj.py @@ -116,6 +116,7 @@ def __init__( self._gwf = self._initialize_gwf(lib_name, self._flow_dir) self._gwf_version = self._get_gwf_version() self._hdf5_name = None + self._dt_dict = None self._structured_mg = None self.is_structured = is_structured @@ -171,27 +172,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)) != "direct": + 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_form' and is a " - + "residual 'pm_type', 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") @@ -562,9 +563,12 @@ def solve_forward_model( sp_package_data = None head_dict = None + dt_dict = None if pert_save: sp_package_data = {} head_dict = {} + # time step lengths, used to weight an instantaneous measure + dt_dict = {} while ctime < etime: sol_start = datetime.now() @@ -702,6 +706,7 @@ def solve_forward_model( data_dict["head"] = head if pert_save: head_dict[kperkstp] = head + dt_dict[kperkstp] = dt1 head_old = self._gwf.get_value( self._gwf.get_var_address("XOLD", self._gwf_name.upper()) @@ -896,12 +901,12 @@ def solve_forward_model( self._add_gwf_info_to_hdf(fhd) fhd.close() if pert_save: + self._dt_dict = dt_dict return head_dict, sp_package_data 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 +926,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 +984,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, @@ -1081,8 +1082,11 @@ def perturbation_method(self, pert_mult: float = 1.01) -> pd.DataFrame: # for d in org_sp_package_data["ghb6"][(0, 0)]: # # print(d) # tot += d["simval"] + dt_dict = self._dt_dict base_results = { - pm.name: pm._performance_measure_forward(org_head, org_sp_package_data) + pm.name: pm._performance_measure_forward( + org_head, org_sp_package_data, dt_dict + ) for pm in self._performance_measures } assert len(base_results) == len(self._performance_measures) @@ -1127,7 +1131,7 @@ def _compute_perturbation_results( """ return { pm.name: ( - pm._performance_measure_forward(pert_head, pert_sp_dict) + pm._performance_measure_forward(pert_head, pert_sp_dict, dt_dict) - base_results[pm.name] ) / epsilon diff --git a/mf6adj/pm.py b/mf6adj/pm.py index d659ad0..d962b62 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: @@ -204,7 +204,9 @@ def name(self) -> str: """ return str(self._name) - def _performance_measure_forward(self, head_dict, sp_package_dict) -> float: + def _performance_measure_forward( + self, head_dict, sp_package_dict, dt_dict=None + ) -> float: """Calculate the forward value of the performance measure. This helper is used during perturbation testing to evaluate the @@ -217,21 +219,37 @@ def _performance_measure_forward(self, head_dict, sp_package_dict) -> float: sp_package_dict : dict Nested mapping of stress-package results produced by :meth:`Mf6Adj.solve_forward_model` when ``pert_save=True``. + dt_dict : dict, optional + Mapping from ``(kper, kstp)`` tuples to time step lengths. Required + for an instantaneous measure. Returns ------- float Scalar forward value for the performance measure. Direct measures - are accumulated linearly and residual measures are accumulated as - weighted squared residuals. + are accumulated linearly, residual measures are accumulated as + weighted squared residuals, and instantaneous measures are the + time-weighted mean of the per-time-step values. """ - result = 0.0 + is_instantaneous = self._entries[0].pm_form == "instantaneous" + if is_instantaneous and dt_dict is None: + raise Exception( + "dt_dict is required to evaluate the instantaneous performance " + + f"measure {self.name}" + ) + + # accumulate by time step so an instantaneous measure can be averaged + # over time below. Every time step with an entry is represented, which + # matches the time steps the adjoint solution accumulates. + per_step = {pfr.kperkstp: 0.0 for pfr in self._entries} for pfr in self._entries: if pfr.pm_type == "head": - if pfr.pm_form == "direct": - result += pfr.weight * head_dict[pfr.kperkstp][pfr.inode] + if pfr.pm_form in ("direct", "instantaneous"): + per_step[pfr.kperkstp] += ( + pfr.weight * head_dict[pfr.kperkstp][pfr.inode] + ) elif pfr.pm_form == "residual": - result += ( + per_step[pfr.kperkstp] += ( pfr.weight * (head_dict[pfr.kperkstp][pfr.inode] - pfr.obsval) ) ** 2 else: @@ -245,13 +263,23 @@ 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": - result += pfr.weight * kk_d["simval"] + if pfr.pm_form in ("direct", "instantaneous"): + per_step[kk] += pfr.weight * kk_d["simval"] elif pfr.pm_form == "residual": - result += ( + per_step[kk] += ( pfr.weight * (kk_d["simval"] - pfr.obsval) ) ** 2 - return result + + if not is_instantaneous: + return float(sum(per_step.values())) + + # An instantaneous measure is the time-weighted mean of the per-step + # values, so the result does not depend on the time discretization. + # This matches the composite the adjoint solution reports. + wsum = sum(dt_dict[kk] for kk in per_step) + if wsum <= 0.0: + return 0.0 + return float(sum(dt_dict[kk] * val for kk, val in per_step.items()) / wsum) def _setup_block_jacobi_preconditioner(self, amat, block_size=5000): """Setup a block Jacobi preconditioner @@ -391,7 +419,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 +446,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 +621,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 +643,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 +681,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 +1004,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 +1023,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 +1032,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 +1054,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 +1098,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 +1614,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)