Adjoint

gusto/core/configuration.py

@@ -1,7 +1,7 @@
 """Some simple tools for configuring the model."""
 from abc import ABCMeta, abstractproperty
 from enum import Enum
-from firedrake import sqrt, Constant
+from firedrake import sqrt
 
 
 __all__ = [
@@ -80,7 +80,11 @@ def __setattr__(self, name, value):
         # Almost all parameters should be Constants -- but there are some
         # specific exceptions which should be kept as integers
         if type(value) in [float, int] and name not in ['dumpfreq', 'pddumpfreq', 'chkptfreq']:
-            object.__setattr__(self, name, Constant(value))
+            object.__setattr__(self, name, value)
+            # DO NOT MERGE
+            # Adjoint is a bit twitchy with Constants so let's not make them
+            # for now
+            # object.__setattr__(self, name, Constant(value))
         else:
             object.__setattr__(self, name, value)
 

gusto/core/domain.py

@@ -63,15 +63,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/physics/shallow_water_microphysics.py

@@ -3,7 +3,7 @@
 """
 
 from firedrake import (
-    Interpolator, conditional, Function, dx, min_value, max_value, Constant
+    Interpolator, conditional, Function, dx, min_value, max_value, FunctionSpace
 )
 from firedrake.fml import subject
 from gusto.core.logging import logger
@@ -93,13 +93,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:
@@ -269,12 +270,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,7 +9,7 @@
 import math
 
 from firedrake import (Function, TestFunction, TestFunctions, DirichletBC,
-                       Constant, NonlinearVariationalProblem,
+                       NonlinearVariationalProblem,
                        NonlinearVariationalSolver)
 from firedrake.fml import (replace_subject, replace_test_function, Term,
                            all_terms, drop)
@@ -71,10 +71,8 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
         self.field_name = field_name
         self.equation = None
 
-        self.dt = Constant(0.0)
-        self.dt.assign(domain.dt)
-        self.original_dt = Constant(0.0)
-        self.original_dt.assign(self.dt)
+        self.dt = Function(domain.dt.function_space(), val=domain.dt)
+        self.original_dt = Function(domain.dt.function_space(), val=self.dt)
         self.options = options
         self.limiter = limiter
         self.courant_max = None

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