#73: Automatically detect and store changing form coefficients

goalie/adjoint.py

@@ -381,6 +381,7 @@ def _solve_adjoint(
         adj_solver_kwargs=None,
         get_adj_values=False,
         test_checkpoint_qoi=False,
+        track_coefficients=False,
     ):
         """
         A generator for solving an adjoint problem on a sequence of subintervals.
@@ -402,6 +403,10 @@ def _solve_adjoint(
         :type get_adj_values: :class:`bool`
         :kwarg test_checkpoint_qoi: solve over the final subinterval when checkpointing
             so that the QoI value can be checked across runs
+        :kwarg: track_coefficients: if ``True``, the adjoint solver will track the
+            coefficients in the variational form to detect changes between export times.
+            Only relevant for goal-oriented error estimation on unsteady problems.
+        :type track_coefficients: :class:`bool`
         :yields: the solution data of the forward and adjoint solves
         :ytype: :class:`~.AdjointSolutionData`
         """
@@ -484,9 +489,20 @@ def wrapped_solver(subinterval, initial_condition_map, **kwargs):
             # Initialise the solver generator
             solver_gen = wrapped_solver(i, checkpoints[i], **solver_kwargs)
 
-            # Annotate tape on current subinterval
-            for _ in range(tp.num_timesteps_per_subinterval[i]):
-                next(solver_gen)
+            # Annotate tape on current subinterval.
+            # If we are using a goal-oriented approach on an unsteady problem, we need
+            # to keep track of the coefficients in the variational form to detect their
+            # changes between export times. In that case, we solve the forward problem
+            # sequentially between each export time and save changing coefficients.
+            # Otherwise, solve over the entire subinterval in one go.
+            if track_coefficients:
+                for _ in range(tp.num_exports_per_subinterval[i] - 1):
+                    for _ in range(tp.num_timesteps_per_export[i]):
+                        next(solver_gen)
+                    self._detect_changing_coefficients()
+            else:
+                for _ in range(tp.num_timesteps_per_subinterval[i]):
+                    next(solver_gen)
             pyadjoint.pause_annotation()
 
             # Final solution is used as the initial condition for the next subinterval

goalie/go_mesh_seq.py

@@ -3,6 +3,7 @@
 """
 
 from collections.abc import Iterable
+from copy import deepcopy
 
 import numpy as np
 import ufl
@@ -28,6 +29,9 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.estimator_values = []
         self._forms = None
+        self._init_form_coeffs = None
+        self._changed_form_coeffs = None
+        self._reset_changing_coefficients()
 
     def read_forms(self, forms_dictionary):
         """
@@ -65,6 +69,43 @@ def forms(self):
             )
         return self._forms
 
+    def _reset_changing_coefficients(self):
+        self._init_form_coeffs = None
+        self._changed_form_coeffs = {field: {} for field in self.fields}
+
+    @PETSc.Log.EventDecorator()
+    def _detect_changing_coefficients(self):
+        """
+        Detect whether coefficients other than the solution in the variational forms
+        change over time. If they do, store the changed coefficients so we can update
+        them in :meth:`~.GoalOrientedMeshSeq.indicate_errors`.
+        """
+        # Save a copy of the coefficients in the first export timestep
+        if self._init_form_coeffs is None:
+            self._init_form_coeffs = {
+                field: deepcopy(form.coefficients())
+                for field, form in self.forms.items()
+            }
+        # In latter export timesteps, detect and store coefficients that have changed
+        # since the first export timestep
+        else:
+            # coeffs = {field: form.coefficients() for field, form in self._forms.items()}
+            for field in self.fields:
+                coeffs = self.forms[field].coefficients()
+                for idx, (coeff, init_coeff) in enumerate(
+                    zip(coeffs, self._init_form_coeffs[field])
+                ):
+                    if coeff.name().split("_old")[0] in self.time_partition.field_names:
+                        continue
+                    if not np.array_equal(
+                        coeff.vector().array(), init_coeff.vector().array()
+                    ):
+                        if idx not in self._changed_form_coeffs[field]:
+                            self._changed_form_coeffs[field][idx] = [
+                                deepcopy(init_coeff)
+                            ]
+                        self._changed_form_coeffs[field][idx].append(deepcopy(coeff))
+
     @PETSc.Log.EventDecorator()
     def get_enriched_mesh_seq(self, enrichment_method="p", num_enrichments=1):
         """
@@ -192,15 +233,21 @@ def indicate_errors(
         self._create_indicators()
 
         # Initialise adjoint solver generators on the MeshSeq and its enriched version
-        adj_sol_gen = self._solve_adjoint(**solver_kwargs)
-        adj_sol_gen_enriched = enriched_mesh_seq._solve_adjoint(**solver_kwargs)
+        adj_sol_gen = self._solve_adjoint(track_coefficients=False, **solver_kwargs)
+        # We need to track changes in the form coefficients if the problem is unsteady
+        track_coefficients = not self.steady
+        adj_sol_gen_enriched = enriched_mesh_seq._solve_adjoint(
+            track_coefficients=track_coefficients, **solver_kwargs
+        )
 
         FWD, ADJ = "forward", "adjoint"
         FWD_OLD = "forward" if self.steady else "forward_old"
         ADJ_NEXT = "adjoint" if self.steady else "adjoint_next"
         P0_spaces = [FunctionSpace(mesh, "DG", 0) for mesh in self]
         # Loop over each subinterval in reverse
         for i in reversed(range(len(self))):
+            self._reset_changing_coefficients()
+
             # Solve the adjoint problem on the current subinterval
             next(adj_sol_gen)
             next(adj_sol_gen_enriched)
@@ -249,6 +296,11 @@ def indicate_errors(
                     )
                     u_star_e[f] -= u_star[f]
 
+                    # Update other time-dependent form coefficients
+                    if self._changed_form_coeffs[f]:
+                        for idx, coeffs in self._changed_form_coeffs[f].items():
+                            self.forms[f].coefficients()[idx].assign(coeffs[j])
+
                     # Evaluate error indicator
                     indi_e = indicator_fn(enriched_mesh_seq.forms[f], u_star_e[f])