Enable adjoint

gusto/core/configuration.py

@@ -1,7 +1,7 @@
 """Some simple tools for configuring the model."""
-from abc import ABCMeta, abstractproperty
+from abc import ABCMeta, abstractmethod
 from enum import Enum
-from firedrake import sqrt, Constant
+from firedrake import sqrt, Constant, Function, FunctionSpace
 
 
 __all__ = [
@@ -12,7 +12,7 @@
     "EmbeddedDGOptions", "ConservativeEmbeddedDGOptions", "RecoveryOptions",
     "ConservativeRecoveryOptions", "SUPGOptions", "MixedFSOptions",
     "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters",
-    "SubcyclingOptions"
+    "SubcyclingOptions", "convert_parameters_to_real_space"
 ]
 
 
@@ -167,7 +167,7 @@ class ShallowWaterParameters(Configuration):
 class WrapperOptions(Configuration, metaclass=ABCMeta):
     """Base class for specifying options for a transport scheme."""
 
-    @abstractproperty
+    @abstractmethod
     def name(self):
         pass
 
@@ -308,3 +308,17 @@ def check_options(self):
             raise ValueError(
                 "Cannot provide both fixed_subcycles and subcycle_by_courant"
                 + "parameters.")
+
+
+def convert_parameters_to_real_space(parameters, mesh):
+    """Convert parameters to functions in real space.
+
+    Args:
+        parameters (:class:`Configuration`): the configuration object
+        containing the parameters to convert
+        mesh (:class:`firedrake.Mesh`): the mesh object to use for the real space.
+    """
+    R = FunctionSpace(mesh, 'R', 0)
+    for name, value in vars(parameters).items():
+        if isinstance(value, (float, Constant)):
+            setattr(parameters, name, Function(R, val=float(value)))

gusto/core/domain.py

@@ -64,15 +64,16 @@ def __init__(self, mesh, dt, family, degree=None,
         # -------------------------------------------------------------------- #
 
         # Store central dt for use in the rest of the model
+        R = FunctionSpace(mesh, "R", 0)
         if type(dt) is Constant:
-            self.dt = dt
+            self.dt = Function(R, val=float(dt))
         elif type(dt) in (float, int):
-            self.dt = Constant(dt)
+            self.dt = Function(R, val=dt)
         else:
             raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}')
 
         # Make a placeholder for the time
-        self.t = Constant(0.0)
+        self.t = Function(R, val=0.0)
 
         # -------------------------------------------------------------------- #
         # Build compatible function spaces

gusto/equations/boussinesq_equations.py

@@ -6,6 +6,7 @@
     time_derivative, transport, prognostic, linearisation,
     pressure_gradient, coriolis, divergence, gravity, incompressible
 )
+from gusto.core.configuration import convert_parameters_to_real_space
 from gusto.equations.common_forms import (
     advection_form, vector_invariant_form,
     kinetic_energy_form, advection_equation_circulation_form,
@@ -105,6 +106,11 @@ def __init__(self, domain, parameters,
                          active_tracers=active_tracers)
 
         self.parameters = parameters
+        # Convert the attributes of type ``float`` or ``firedrake.Constant``
+        # in the parameters to a function in real space. This conversion is a
+        # preventive to avoid issues with adjoint computations, particularly
+        # when the parameters  are used as controls in sensitivity analyses.
+        convert_parameters_to_real_space(parameters, domain.mesh)
         self.compressible = compressible
 
         w, phi, gamma = self.tests[0:3]
@@ -168,10 +174,12 @@ def __init__(self, domain, parameters,
         # -------------------------------------------------------------------- #
         if compressible:
             cs = parameters.cs
+            # On assuming ``cs`` as a constant, it is right keep it out of the
+            # integration.
             linear_div_form = divergence(subject(
-                prognostic(cs**2 * phi * div(u_trial) * dx, 'p'), self.X))
+                prognostic(cs**2 * (phi * div(u_trial) * dx), 'p'), self.X))
             divergence_form = divergence(linearisation(
-                subject(prognostic(cs**2 * phi * div(u) * dx, 'p'), self.X),
+                subject(prognostic(cs**2 * (phi * div(u) * dx), 'p'), self.X),
                 linear_div_form))
         else:
             # This enforces that div(u) = 0

gusto/equations/compressible_euler_equations.py

@@ -17,7 +17,7 @@
 )
 from gusto.equations.active_tracers import Phases, TracerVariableType
 from gusto.equations.prognostic_equations import PrognosticEquationSet
-
+from gusto.core.configuration import convert_parameters_to_real_space
 __all__ = ["CompressibleEulerEquations", "HydrostaticCompressibleEulerEquations"]
 
 
@@ -45,7 +45,7 @@ def __init__(self, domain, parameters, sponge_options=None,
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
                 mesh and the compatible function spaces.
-            parameters (:class:`Configuration`, optional): an object containing
+            x (:class:`Configuration`, optional): an object containing
                 the model's physical parameters.
             sponge_options (:class:`SpongeLayerParameters`, optional): any
                 parameters for applying a sponge layer to the upper boundary.
@@ -101,6 +101,11 @@ def __init__(self, domain, parameters, sponge_options=None,
                          active_tracers=active_tracers)
 
         self.parameters = parameters
+        # Convert the attributes of type ``float`` or ``firedrake.Constant``
+        # in the parameters to a function in real space. This conversion is a
+        # preventive to avoid issues with adjoint computations, particularly
+        # when the parameters  are used as controls in sensitivity analyses.
+        convert_parameters_to_real_space(parameters, domain.mesh)
         g = parameters.g
         cp = parameters.cp
 
@@ -109,6 +114,7 @@ def __init__(self, domain, parameters, sponge_options=None,
         u_trial = split(self.trials)[0]
         _, rho_bar, theta_bar = split(self.X_ref)[0:3]
         zero_expr = Constant(0.0)*theta
+        # Check this for adjoints
         exner = exner_pressure(parameters, rho, theta)
         n = FacetNormal(domain.mesh)
 

gusto/equations/shallow_water_equations.py

@@ -12,6 +12,7 @@
     linear_continuity_form, linear_advection_form
 )
 from gusto.equations.prognostic_equations import PrognosticEquationSet
+from gusto.core.configuration import convert_parameters_to_real_space
 
 
 __all__ = ["ShallowWaterEquations", "LinearShallowWaterEquations",
@@ -87,6 +88,11 @@ def __init__(self, domain, parameters, fexpr=None, topog_expr=None,
 
         self.parameters = parameters
         self.domain = domain
+        # Convert the attributes of type ``float`` or ``firedrake.Constant``
+        # in the parameters to a function in real space. This conversion is a
+        # preventive to avoid issues with adjoint computations, particularly
+        # when the parameters  are used as controls in sensitivity analyses.
+        convert_parameters_to_real_space(parameters, self.domain.mesh)
         self.active_tracers = active_tracers
 
         self._setup_residual(fexpr, topog_expr, u_transport_option)
@@ -163,8 +169,9 @@ def _setup_residual(self, fexpr, topog_expr, u_transport_option):
         # -------------------------------------------------------------------- #
         # Pressure Gradient Term
         # -------------------------------------------------------------------- #
+        # On assuming ``g``, it is right to keep it out of the integral.
         pressure_gradient_form = pressure_gradient(
-            subject(prognostic(-g*div(w)*D*dx, 'u'), self.X))
+            subject(prognostic(-g*(div(w)*D*dx), 'u'), self.X))
 
         residual = (mass_form + adv_form + pressure_gradient_form)
 
@@ -418,6 +425,7 @@ def _setup_residual(self, fexpr, topog_expr, u_transport_option):
         # provide linearisation
         if self.equivalent_buoyancy:
             beta2 = self.parameters.beta2
+
             qsat_expr = self.compute_saturation(self.X)
             qv = conditional(qt < qsat_expr, qt, qsat_expr)
             qvbar = conditional(qtbar < qsat_expr, qtbar, qsat_expr)
@@ -647,6 +655,11 @@ def __init__(self, domain, parameters,
                          active_tracers=active_tracers)
 
         self.parameters = parameters
+        # Convert the attributes of type ``float`` or ``firedrake.Constant``
+        # in the parameters to a function in real space. This conversion is a
+        # preventive to avoid issues with adjoint computations, particularly
+        # when the parameters  are used as controls in sensitivity analyses.
+        convert_parameters_to_real_space(parameters, domain.mesh)
         g = parameters.g
         H = parameters.H
 

gusto/physics/shallow_water_microphysics.py

@@ -3,7 +3,7 @@
 """
 
 from firedrake import (
-    conditional, Function, dx, min_value, max_value, Constant, assemble
+    conditional, Function, dx, min_value, max_value, FunctionSpace, assemble
 )
 from firedrake.__future__ import interpolate
 from firedrake.fml import subject
@@ -94,13 +94,14 @@ def __init__(self, equation, saturation_curve,
         self.water_v = Function(Vv)
         self.source = Function(Vv)
 
+        R = FunctionSpace(equation.domain.mesh, "R", 0)
         # tau is the timescale for conversion (may or may not be the timestep)
         if tau is not None:
             self.set_tau_to_dt = False
-            self.tau = tau
+            self.tau = Function(R).assign(tau)
         else:
             self.set_tau_to_dt = True
-            self.tau = Constant(0)
+            self.tau = Function(R)
             logger.info("Timescale for rain conversion has been set to dt. If this is not the intention then provide a tau parameter as an argument to InstantRain.")
 
         if self.time_varying_saturation:
@@ -270,12 +271,13 @@ def __init__(self, equation, saturation_curve,
             V_idxs.append(self.Vb_idx)
 
         # tau is the timescale for condensation/evaporation (may or may not be the timestep)
+        R = FunctionSpace(equation.domain.mesh, "R", 0)
         if tau is not None:
             self.set_tau_to_dt = False
-            self.tau = tau
+            self.tau = Function(R).assign(tau)
         else:
             self.set_tau_to_dt = True
-            self.tau = Constant(0)
+            self.tau = Function(R)
             logger.info("Timescale for moisture conversion between vapour and cloud has been set to dt. If this is not the intention then provide a tau parameter as an argument to SWSaturationAdjustment.")
 
         if self.time_varying_saturation:

gusto/spatial_methods/diffusion_methods.py

@@ -160,4 +160,5 @@ def __init__(self, equation, variable, diffusion_parameters):
             kappa = diffusion_parameters.kappa
             self.form = diffusion(kappa * self.test.dx(0) * self.field.dx(0) * dx)
         else:
-            raise NotImplementedError("CG diffusion only implemented in 1D")
+            kappa = diffusion_parameters.kappa
+            self.form = diffusion(kappa * inner(grad(self.test), grad(self.field)) * dx)

gusto/time_discretisation/time_discretisation.py

@@ -9,8 +9,8 @@
 import math
 
 from firedrake import (Function, TestFunction, TestFunctions, DirichletBC,
-                       Constant, NonlinearVariationalProblem,
-                       NonlinearVariationalSolver)
+                       NonlinearVariationalProblem, NonlinearVariationalSolver,
+                       FunctionSpace)
 from firedrake.fml import (replace_subject, replace_test_function, Term,
                            all_terms, drop)
 from firedrake.formmanipulation import split_form
@@ -88,10 +88,10 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
         self.domain = domain
         self.field_name = field_name
         self.equation = None
-
-        self.dt = Constant(0.0)
+        R = FunctionSpace(domain.mesh, "R", 0)
+        self.dt = Function(R, val=0.0)
         self.dt.assign(domain.dt)
-        self.original_dt = Constant(0.0)
+        self.original_dt = Function(R, val=0.0)
         self.original_dt.assign(self.dt)
         self.options = options
         self.limiter = limiter

gusto/timestepping/timestepper.py

@@ -215,7 +215,7 @@ def run(self, t, tmax, pick_up=False):
 
             self.timestep()
 
-            self.t.assign(self.t + self.dt)
+            self.t.assign(float(self.t) + float(self.dt))
             self.step += 1
 
             with timed_stage("Dump output"):

integration-tests/adjoints/test_diffusion_sensitivity.py

@@ -0,0 +1,78 @@
+import pytest
+import numpy as np
+
+from firedrake import *
+from firedrake.adjoint import *
+from pyadjoint import get_working_tape
+from gusto import *
+
+
+@pytest.fixture(autouse=True)
+def handle_taping():
+    yield
+    tape = get_working_tape()
+    tape.clear_tape()
+
+
+@pytest.fixture(autouse=True, scope="module")
+def handle_annotation():
+    from firedrake.adjoint import annotate_tape, continue_annotation
+    if not annotate_tape():
+        continue_annotation()
+    yield
+    # Ensure annotation is paused when we finish.
+    annotate = annotate_tape()
+    if annotate:
+        pause_annotation()
+
+
+@pytest.mark.parametrize("nu_is_control", [True, False])
+def test_diffusion_sensitivity(nu_is_control, tmpdir):
+    assert get_working_tape()._blocks == []
+    n = 30
+    mesh = PeriodicUnitSquareMesh(n, n)
+    output = OutputParameters(dirname=str(tmpdir))
+    dt = 0.01
+    domain = Domain(mesh, 10*dt, family="BDM", degree=1)
+    io = IO(domain, output)
+
+    V = VectorFunctionSpace(mesh, "CG", 2)
+    domain.spaces.add_space("vecCG", V)
+
+    R = FunctionSpace(mesh, "R", 0)
+    # We need to define nu as a function in order to have a control variable.
+    nu = Function(R, val=0.0001)
+    diffusion_params = DiffusionParameters(kappa=nu)
+    eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params)
+
+    diffusion_scheme = BackwardEuler(domain)
+    diffusion_methods = [CGDiffusion(eqn, "f", diffusion_params)]
+    timestepper = Timestepper(eqn, diffusion_scheme, io, spatial_methods=diffusion_methods)
+
+    x = SpatialCoordinate(mesh)
+    fexpr = as_vector((sin(2*pi*x[0]), cos(2*pi*x[1])))
+    timestepper.fields("f").interpolate(fexpr)
+
+    end = 0.1
+    timestepper.run(0., end)
+
+    u = timestepper.fields("f")
+    J = assemble(inner(u, u)*dx)
+
+    if nu_is_control:
+        control = Control(nu)
+        h = Function(R, val=0.0001)  # the direction of the perturbation
+    else:
+        control = Control(u)
+        # the direction of the perturbation
+        h = Function(V).interpolate(fexpr * np.random.rand())
+
+    # the functional as a pure function of nu
+    Jhat = ReducedFunctional(J, control)
+
+    if nu_is_control:
+        assert np.allclose(J, Jhat(nu))
+        assert taylor_test(Jhat, nu, h) > 1.95
+    else:
+        assert np.allclose(J, Jhat(u))
+        assert taylor_test(Jhat, u, h) > 1.95