Dham/abstract reduced functional (firedrakeproject#3941)

Co-authored-by: Daiane Iglesia Dolci <63597005+Ig-dolci@users.noreply.github.com>
Co-authored-by: Josh Hope-Collins <joshua.hope-collins13@imperial.ac.uk>
Co-authored-by: Pablo Brubeck <brubeck@protonmail.com>

Author: 4 people

demos/full_waveform_inversion/full_waveform_inversion.py.rst

@@ -109,10 +109,12 @@ built over the ``my_ensemble.comm`` (spatial) communicator.
         dt = 0.03  # time step in seconds
         final_time = 0.6  # final time in seconds
         nx, ny = 15, 15
+        ftol = 0.9  # optimisation tolerance
     else:
         dt = 0.002  # time step in seconds
         final_time = 1.0  # final time in seconds
         nx, ny = 80, 80
+        ftol = 1e-2  # optimisation tolerance
 
     mesh = UnitSquareMesh(nx, ny, comm=my_ensemble.comm)
 
@@ -280,21 +282,28 @@ To have the step 4, we need first to tape the forward problem. That is done by c
 
 We now instantiate :class:`~.EnsembleReducedFunctional`::
 
-    J_hat = EnsembleReducedFunctional(J_val, Control(c_guess), my_ensemble)
+    J_hat = EnsembleReducedFunctional(J_val,
+                                      Control(c_guess, riesz_map="l2"),
+                                      my_ensemble)
 
 which enables us to recompute :math:`J` and its gradient :math:`\nabla_{\mathtt{c\_guess}} J`,
 where the :math:`J_s` and its gradients :math:`\nabla_{\mathtt{c\_guess}} J_s` are computed in parallel
 based on the ``my_ensemble`` configuration.
 
 
 **Steps 4-6**: The instance of the :class:`~.EnsembleReducedFunctional`, named ``J_hat``,
-is then passed as an argument to the ``minimize`` function::
+is then passed as an argument to the ``minimize`` function. The default ``minimize`` function
+uses ``scipy.minimize``, and wraps the ``ReducedFunctional`` in a ``ReducedFunctionalNumPy``
+that handles transferring data between Firedrake and numpy data structures. However, because
+we have a custom ``ReducedFunctional``, we need to do this ourselves::
 
-    c_optimised = minimize(J_hat, method="L-BFGS-B", options={"disp": True, "maxiter": 1},
-                            bounds=(1.5, 2.0), derivative_options={"riesz_representation": 'l2'}
-                            )
+    from pyadjoint.reduced_functional_numpy import ReducedFunctionalNumPy
+    Jnumpy = ReducedFunctionalNumPy(J_hat)
 
-The ``minimize`` function executes the optimisation algorithm until the stopping criterion (``maxiter``) is met.
+    c_optimised = minimize(Jnumpy, method="L-BFGS-B", options={"disp": True, "ftol": ftol},
+                           bounds=(1.5, 2.0))
+
+The ``minimize`` function executes the optimisation algorithm until the stopping criterion (``ftol``) is met.
 For 20 iterations, the predicted velocity model is shown in the following figure.
 
 .. image:: c_predicted.png
@@ -305,9 +314,7 @@ For 20 iterations, the predicted velocity model is shown in the following figure
 .. warning::
 
     The ``minimize`` function uses the SciPy library for optimisation. However, for scenarios that require higher
-    levels of spatial parallelism, you should assess whether SciPy is the most suitable option for your problem.
-    SciPy's optimisation algorithm is not inner-product-aware. Therefore, we configure the options with
-    ``derivative_options={"riesz_representation": 'l2'}`` to account for this requirement.
+    levels of spatial parallelism, you should assess whether SciPy is the most suitable option for your problem such as the pyadjoint's TAOSolver.
 
 .. note::
 

demos/shape_optimization/shape_optimization.py.rst

@@ -0,0 +1,109 @@
+Shape optimization
+==================
+
+Shape optimization is about modifying the shape of a domain :math:`\Omega` so
+that an objective function :math:`J(\Omega)` is minimized. In this demo, we
+consider an objective function constrained to a boundary value problem and
+implement a simple mesh-moving shape optimization strategy using Firedrake and
+pyadjoint.  This tutorial was contributed by `Alberto Paganini
+<mailto:apaganini@le.ac.uk>`__ with support from `Ado Farsi
+<mailto:ado.farsi@imperial.ac.uk>`__ and `Mirko Ciceri
+<mailto:mc5823@ic.ac.uk>`__, and was written during the `ccp-dcm hackaton
+<https://ccp-dcm.github.io/exeter_hackathon>`__ at Dartington Hall.
+
+Let
+
+.. math::
+
+   J(\Omega) = \int_\Omega \big(u(\mathbf{x}) - u_t(\mathbf{x})\big)^2 \,\mathrm{d}\mathbf{x}\,.
+
+measure the difference between a steady-state temperature profile
+:math:`u:\mathbb{R}^2\to\mathbb{R}` and a target steady-state temperature
+profile :math:`u_t:\mathbb{R}^2\to\mathbb{R}`. Specifically, the function
+:math:`u` is the solution to the steady-state heat equation
+
+.. math::
+
+    -\Delta u = 4 \quad \text{in }\Omega\,, \qquad u = 0 \quad \text{on } \partial\Omega
+
+
+and the target temperature profile :math:`u_t` is
+
+.. math::
+
+    u_t(x,y) = 1.21 - (x - 0.5)^2 - (y - 0.5)^2\,.
+
+Beside the empty set, the domain that minimizes :math:`J(\Omega)` is a disc of
+radius :math:`1.1` centered at :math:`(0.5,0.5)`.
+
+We can now proceed to set up the problem. We import firedrake and pyadjoint and
+choose an initial guess (in this case, a unit disc centred at the origin)::
+
+  from firedrake import *
+  from firedrake.adjoint import *
+  mesh = UnitDiskMesh(refinement_level=3)
+
+Then, we :ref:`start annotating <adjoint-taping>` and turn the mesh coordinates into a control variable::
+
+  continue_annotation()
+  Q = mesh.coordinates.function_space()
+  dT = Function(Q)
+  mesh.coordinates.assign(mesh.coordinates + dT)
+
+We can now implement the target function::
+
+  x, y = SpatialCoordinate(mesh)
+  u_t = Constant(1.21) - (x - Constant(0.5))**2 - (y - Constant(0.5))**2
+
+solve the weak form of the boundary value problem::
+
+  V = FunctionSpace(mesh, "CG", 1)
+  u = Function(V, name='state')
+  v = TestFunction(V)
+  F = (dot(grad(u), grad(v)) - 4 * v) * dx
+  bcs = DirichletBC(V, Constant(0.), "on_boundary")
+  solve(F == 0, u, bcs=bcs)
+
+and evaluate the objective function::
+
+  J = assemble((u - u_t)**2*dx)
+
+We now turn the objective function into a reduced function so that pyadjoint
+(and UFL shape differentiation capability) can automatically compute shape
+gradients, that is, directions of steepest ascent. We also set the relevant
+Riesz map for this problem::
+
+  Jred = ReducedFunctional(J, Control(dT, riesz_map="H1"))
+  stop_annotating()
+
+We now have all the ingredients to implement a basic steepest descent shape
+optimization algorithm with fixed step size.::
+
+  File = VTKFile("shape_iterates.pvd")
+  for ii in range(30):
+    print("J(ii =", ii, ") =", Jred(dT))
+    File.write(mesh.coordinates)
+
+    # compute the gradient (steepest ascent)
+    gradJ = Jred.derivative(apply_riesz=True)
+
+    # update domain
+    dT -= 0.2*gradJ
+
+  File.write(mesh.coordinates)
+  print("J(final) =", Jred(dT))
+
+.. only:: html
+
+  .. container:: youtube
+
+    .. vimeo:: 1083822714?loop=1
+       :width: 600px
+
+
+**Remark:** mesh-moving shape optimization can lead to mesh tangling, which
+invalidates finite element computations. For faster and more robust shape
+optimization, we recommend using Firedrake's shape optimization toolbox
+`Fireshape <https://github.com/fireshape/fireshape>`__.
+
+A python script version of this demo can be found :demo:`here <shape_optimization.py>`.

firedrake/adjoint/__init__.py

@@ -29,7 +29,7 @@
 from firedrake.adjoint_utils import get_solve_blocks  # noqa F401
 
 from pyadjoint.verification import taylor_test, taylor_to_dict  # noqa F401
-from pyadjoint.drivers import compute_gradient, compute_hessian  # noqa F401
+from pyadjoint.drivers import compute_gradient, compute_derivative, compute_hessian  # noqa F401
 from pyadjoint.adjfloat import AdjFloat  # noqa F401
 from pyadjoint.control import Control  # noqa F401
 from pyadjoint import IPOPTSolver, ROLSolver, MinimizationProblem, \

firedrake/adjoint/ensemble_reduced_functional.py

@@ -1,11 +1,12 @@
-from pyadjoint import ReducedFunctional
+from pyadjoint.reduced_functional import AbstractReducedFunctional, ReducedFunctional
 from pyadjoint.enlisting import Enlist
 from pyop2.mpi import MPI
 
-import firedrake
+from firedrake.function import Function
+from firedrake.cofunction import Cofunction
 
 
-class EnsembleReducedFunctional(ReducedFunctional):
+class EnsembleReducedFunctional(AbstractReducedFunctional):
     """Enable solving simultaneously reduced functionals in parallel.
 
     Consider a functional :math:`J` and its gradient :math:`\\dfrac{dJ}{dm}`,
@@ -34,7 +35,7 @@ class EnsembleReducedFunctional(ReducedFunctional):
 
     Parameters
     ----------
-    J : pyadjoint.OverloadedType
+    functional : pyadjoint.OverloadedType
         An instance of an OverloadedType, usually :class:`pyadjoint.AdjFloat`.
         This should be the functional that we want to reduce.
     control : pyadjoint.Control or list of pyadjoint.Control
@@ -86,28 +87,40 @@ class EnsembleReducedFunctional(ReducedFunctional):
     works, please refer to the `Firedrake manual
     <https://www.firedrakeproject.org/parallelism.html#ensemble-parallelism>`_.
     """
-    def __init__(self, J, control, ensemble, scatter_control=True,
-                 gather_functional=None, derivative_components=None,
-                 scale=1.0, tape=None, eval_cb_pre=lambda *args: None,
+    def __init__(self, functional, control, ensemble, scatter_control=True,
+                 gather_functional=None,
+                 derivative_components=None,
+                 scale=1.0, tape=None,
+                 eval_cb_pre=lambda *args: None,
                  eval_cb_post=lambda *args: None,
                  derivative_cb_pre=lambda controls: controls,
                  derivative_cb_post=lambda checkpoint, derivative_components, controls: derivative_components,
-                 hessian_cb_pre=lambda *args: None, hessian_cb_post=lambda *args: None):
-        super(EnsembleReducedFunctional, self).__init__(
-            J, control, derivative_components=derivative_components,
-            scale=scale, tape=tape, eval_cb_pre=eval_cb_pre,
-            eval_cb_post=eval_cb_post, derivative_cb_pre=derivative_cb_pre,
+                 hessian_cb_pre=lambda *args: None,
+                 hessian_cb_post=lambda *args: None):
+        self.local_reduced_functional = ReducedFunctional(
+            functional, control,
+            derivative_components=derivative_components,
+            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)
+            hessian_cb_pre=hessian_cb_pre,
+            hessian_cb_post=hessian_cb_post
+        )
 
         self.ensemble = ensemble
         self.scatter_control = scatter_control
         self.gather_functional = gather_functional
 
+    @property
+    def controls(self):
+        return self.local_reduced_functional.controls
+
     def _allgather_J(self, J):
         if isinstance(J, float):
             vals = self.ensemble.ensemble_comm.allgather(J)
-        elif isinstance(J, firedrake.Function):
+        elif isinstance(J, Function):
             #  allgather not implemented in ensemble.py
             vals = []
             for i in range(self.ensemble.ensemble_comm.size):
@@ -134,30 +147,31 @@ def __call__(self, values):
             The computed value. Typically of instance of :class:`pyadjoint.AdjFloat`.
 
         """
-        local_functional = super(EnsembleReducedFunctional, self).__call__(values)
+        local_functional = self.local_reduced_functional(values)
         ensemble_comm = self.ensemble.ensemble_comm
         if self.gather_functional:
             controls_g = self._allgather_J(local_functional)
             total_functional = self.gather_functional(controls_g)
         # if gather_functional is None then we do a sum
         elif isinstance(local_functional, float):
             total_functional = ensemble_comm.allreduce(sendobj=local_functional, op=MPI.SUM)
-        elif isinstance(local_functional, firedrake.Function):
+        elif isinstance(local_functional, Function):
             total_functional = type(local_functional)(local_functional.function_space())
             total_functional = self.ensemble.allreduce(local_functional, total_functional)
         else:
             raise NotImplementedError("This type of functional is not supported.")
         return total_functional
 
-    def derivative(self, adj_input=1.0, options=None):
+    def derivative(self, adj_input=1.0, apply_riesz=False):
         """Compute derivatives of a functional with respect to the control parameters.
 
         Parameters
         ----------
         adj_input : float
             The adjoint input.
-        options : dict
-            Additional options for the derivative computation.
+        apply_riesz: bool
+            If True, apply the Riesz map of each control in order to return
+            a primal gradient rather than a derivative in the dual space.
 
         Returns
         -------
@@ -171,29 +185,62 @@ def derivative(self, adj_input=1.0, options=None):
 
         if self.gather_functional:
             dJg_dmg = self.gather_functional.derivative(adj_input=adj_input,
-                                                        options=options)
+                                                        apply_riesz=False)
             i = self.ensemble.ensemble_comm.rank
             adj_input = dJg_dmg[i]
 
-        dJdm_local = super(EnsembleReducedFunctional, self).derivative(adj_input=adj_input, options=options)
+        dJdm_local = self.local_reduced_functional.derivative(adj_input=adj_input,
+                                                              apply_riesz=apply_riesz)
 
         if self.scatter_control:
             dJdm_local = Enlist(dJdm_local)
             dJdm_total = []
 
             for dJdm in dJdm_local:
-                if not isinstance(dJdm, (firedrake.Function, float)):
-                    raise NotImplementedError("This type of gradient is not supported.")
+                if not isinstance(dJdm, (Cofunction, Function, float)):
+                    raise NotImplementedError(
+                        f"Gradients of type {type(dJdm).__name__} are not supported.")
 
                 dJdm_total.append(
                     self.ensemble.allreduce(dJdm, type(dJdm)(dJdm.function_space()))
-                    if isinstance(dJdm, firedrake.Function)
+                    if isinstance(dJdm, (Cofunction, Function))
                     else self.ensemble.ensemble_comm.allreduce(sendobj=dJdm, op=MPI.SUM)
                 )
             return dJdm_local.delist(dJdm_total)
         return dJdm_local
 
-    def hessian(self, m_dot, options=None):
+    def tlm(self, m_dot):
+        """Return the action of the tangent linear model of the functional.
+
+        The tangent linear model is evaluated w.r.t. the control on a vector
+        m_dot, around the last supplied value of the control.
+
+        Parameters
+        ----------
+        m_dot : pyadjoint.OverloadedType
+            The direction in which to compute the action of the tangent linear model.
+
+        Returns
+        -------
+            pyadjoint.OverloadedType: The action of the tangent linear model in the
+            direction m_dot.  Should be an instance of the same type as the functional.
+        """
+        local_tlm = self.local_reduced_functional.tlm(m_dot)
+        ensemble_comm = self.ensemble.ensemble_comm
+        if self.gather_functional:
+            mdot_g = self._allgather_J(local_tlm)
+            total_tlm = self.gather_functional.tlm(mdot_g)
+        # if gather_functional is None then we do a sum
+        elif isinstance(local_tlm, float):
+            total_tlm = ensemble_comm.allreduce(sendobj=local_tlm, op=MPI.SUM)
+        elif isinstance(local_tlm, Function):
+            total_tlm = type(local_tlm)(local_tlm.function_space())
+            total_tlm = self.ensemble.allreduce(local_tlm, total_tlm)
+        else:
+            raise NotImplementedError("This type of functional is not supported.")
+        return total_tlm
+
+    def hessian(self, m_dot, hessian_input=None, evaluate_tlm=True, apply_riesz=False):
         """The Hessian is not yet implemented for ensemble reduced functional.
 
         Raises: