Implementation of a Parametrised Reduced Functional

pyadjoint/optimization/optimization.py

@@ -229,7 +229,7 @@ def minimize(rf, method='L-BFGS-B', scale=1.0, **kwargs):
     rf.scale = scale
     if isinstance(rf, ReducedFunctionalNumPy):
         rf_np = rf
-    elif isinstance(rf, ReducedFunctional):
+    elif isinstance(rf, (ReducedFunctional)):
         rf_np = ReducedFunctionalNumPy(rf)
     else:
         # Assume the user knows what he is doing - he might for example written

pyadjoint/reduced_functional.py

@@ -168,6 +168,16 @@ class ReducedFunctional(AbstractReducedFunctional):
     A reduced functional maps a control value to the provided functional.
     It may also be used to compute the derivative of the functional with
     respect to the control.
+    When the `parameters` argument is provided, it represents an object which encompasses computations of the form::
+
+        Jhat(m; p) = J(u(m; p), m; p)
+
+    Where `u` is the system state and `m` is a `pyadjoint.Control` or list of
+    `pyadjoint.Control`, `p` is a list of parameters, `J` is an overloaded type providing
+    the functional value, and `Jhat` is a reduced functional where the explicit dependence
+    on `u` has been eliminated. The parameters `p` can be updated between evaluations, after which evaluation
+    of the reduced functional is performed with the updated parameters, but they are not
+    included in the derivative calculations.
 
     Args:
         functional (:obj:`OverloadedType`): An instance of an OverloadedType,
@@ -179,7 +189,10 @@ class ReducedFunctional(AbstractReducedFunctional):
         derivative_components (tuple of int): The indices of the controls with
             respect to which to take the derivative. By default, the derivative
             is taken with respect to all controls. If present, it overwrites
-            derivative_cb_pre and derivative_cb_post.
+            derivative_cb_pre and derivative_cb_post. Will be deprecated in future in favour
+            of the parameters argument.
+        parameters (list): A list of parameters, which are updated, but not included in the derivative.
+        Cannot be specified at the same time as derivative_components.
         scale (float): A scaling factor applied to the functional and its
             gradient with respect to the control.
         tape (Tape): A tape object that the reduced functional will use to
@@ -209,6 +222,7 @@ class ReducedFunctional(AbstractReducedFunctional):
 
     def __init__(self, functional, controls,
                  derivative_components=None,
+                 parameters=None,
                  scale=1.0, tape=None,
                  eval_cb_pre=lambda *args: None,
                  eval_cb_post=lambda *args: None,
@@ -219,12 +233,18 @@ def __init__(self, functional, controls,
                  hessian_cb_post=lambda *args: None,
                  tlm_cb_pre=lambda *args: None,
                  tlm_cb_post=lambda *args: None):
+        if derivative_components is not None and parameters is not None:
+            raise ValueError(
+                "Cannot specify both derivative_components and parameters. "
+                "Please specify only one of these, or neither.")
+        elif parameters is not None and len(Enlist(parameters)) == 0:
+            raise ValueError("Parameters list cannot be empty")
         if not isinstance(functional, OverloadedType):
             raise TypeError("Functional must be an OverloadedType.")
+
         self.functional = functional
         self.tape = get_working_tape() if tape is None else tape
         self._controls = Enlist(controls)
-        self.derivative_components = derivative_components
         self.scale = scale
         self.eval_cb_pre = eval_cb_pre
         self.eval_cb_post = eval_cb_post
@@ -235,7 +255,26 @@ def __init__(self, functional, controls,
         self.tlm_cb_pre = tlm_cb_pre
         self.tlm_cb_post = tlm_cb_post
 
-        if self.derivative_components:
+        if parameters is not None:
+            self._parameters = Enlist(parameters)
+            self.n_opt = len(self._controls)
+            self._all_controls = self._controls + Enlist([Control(p) for p in self._parameters])
+            self._reduced_functional = ReducedFunctional(
+                functional=functional,
+                controls=self._all_controls,
+                scale=scale,
+                tape=tape,
+                eval_cb_pre=eval_cb_pre,
+                eval_cb_post=eval_cb_post,
+                derivative_cb_pre=derivative_cb_pre,
+                derivative_cb_post=derivative_cb_post,
+                hessian_cb_pre=hessian_cb_pre,
+                hessian_cb_post=hessian_cb_post,
+                tlm_cb_pre=tlm_cb_pre,
+                tlm_cb_post=tlm_cb_post,
+            )
+        elif derivative_components is not None:
+            self.derivative_components = derivative_components
             # pre callback
             self.derivative_cb_pre = _get_extract_derivative_components(
                 derivative_components)
@@ -247,6 +286,24 @@ def __init__(self, functional, controls,
     def controls(self) -> list[Control]:
         return self._controls
 
+    @property
+    def parameters(self) -> list[OverloadedType]:
+        if hasattr(self, "_parameters"):
+            return self._parameters
+        else:
+            raise AttributeError("This ReducedFunctional does not have parameters.")
+
+    @no_annotations
+    def update_parameters(self, new_parameters):
+        if not hasattr(self, "_parameters"):
+            raise AttributeError("This ReducedFunctional does not have parameters.")
+        elif hasattr(self, "_parameters") and len(Enlist(new_parameters)) != len(self._parameters):
+            raise ValueError(
+                """new_parameters should be a list of the same
+                length as parameters."""
+            )
+        self._parameters = Enlist(new_parameters)
+
     @no_annotations
     def derivative(self, adj_input=1.0, apply_riesz=False):
         values = [c.tape_value() for c in self.controls]
@@ -278,7 +335,11 @@ def derivative(self, adj_input=1.0, apply_riesz=False):
             for derivative_cb_post has changed. It should now return a
             list of derivatives, usually the same list as input.""")
 
-        return self.controls.delist(derivatives)
+        if not hasattr(self, "_parameters"):
+            return self.controls.delist(derivatives)
+        else:
+            derivatives_all = self._reduced_functional.derivative(adj_input=adj_input, apply_riesz=apply_riesz)
+            return self.controls.delist(Enlist(derivatives_all)[:self.n_opt])
 
     @no_annotations
     def hessian(self, m_dot, hessian_input=None, evaluate_tlm=True, apply_riesz=False):
@@ -295,7 +356,21 @@ def hessian(self, m_dot, hessian_input=None, evaluate_tlm=True, apply_riesz=Fals
                              self.controls.delist(r),
                              self.controls.delist(values))
 
-        return self.controls.delist(r)
+        if not hasattr(self, "_parameters"):
+            return self.controls.delist(r)
+        else:
+            # self._reduced_functional.hessian will expect len(m_dot)
+            # = len(self._all_controls), so we pad it with zeros.
+            m_dot_all = Enlist(m_dot) + [p._ad_init_zero() for p in self._parameters]
+            hessian_all = self._reduced_functional.hessian(
+                m_dot_all,
+                hessian_input=hessian_input,
+                evaluate_tlm=evaluate_tlm,
+                apply_riesz=apply_riesz,
+            )
+
+            # Return only the hessian components corresponding to optimization controls.
+            return self.controls.delist(Enlist(hessian_all)[:self.n_opt])
 
     @no_annotations
     def tlm(self, m_dot):
@@ -308,13 +383,22 @@ def tlm(self, m_dot):
         # Call callback
         self.tlm_cb_post(self.functional.block_variable.checkpoint,
                          tlm, self.controls.delist(values))
-
-        return tlm
+        if not hasattr(self, "_parameters"):
+            return tlm
+        else:
+            # self._reduced_functional.tlm will expect len(m_dot) = len(self._all_controls), so we pad it with zeros.
+            m_dot_all = Enlist(m_dot) + [p._ad_init_zero() for p in self._parameters]
+            tlm_all = self._reduced_functional.tlm(m_dot_all)
+            return tlm_all
 
     @no_annotations
     def __call__(self, values):
         values = Enlist(values)
-        if len(values) != len(self.controls):
+        if hasattr(self, "_parameters") and len(values) != self.n_opt:
+            raise ValueError(
+                f"values should be a list of same length as optimization controls, which is {self.n_opt}."
+            )
+        elif len(values) != len(self.controls):
             raise ValueError(
                 "values should be a list of same length as controls."
             )
@@ -361,7 +445,11 @@ def __call__(self, values):
         # Call callback
         self.eval_cb_post(func_value, self.controls.delist(values))
 
-        return func_value
+        if not hasattr(self, "_parameters"):
+            return func_value
+        else:
+            full_values = values + self._parameters
+            return self._reduced_functional(full_values)
 
     def optimize_tape(self):
         self.tape.optimize(