diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index fe0e49a28..110405818 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -11,13 +11,6 @@ on:
     # Scheduled build at 0330 UTC on Monday mornings to detect bitrot.
     - cron:  '30 3 * * 1'
 
-concurrency:
-  # Cancels jobs running if new commits are pushed
-  group: >
-    ${{ github.workflow }}-
-    ${{ github.event.pull_request.number || github.ref }}
-  cancel-in-progress: true
-
 jobs:
   build:
     name: "Build Gusto"
diff --git a/docs/notebook/shallow_water_adjoint.ipynb b/docs/notebook/shallow_water_adjoint.ipynb
new file mode 100644
index 000000000..d0a8c0062
--- /dev/null
+++ b/docs/notebook/shallow_water_adjoint.ipynb
@@ -0,0 +1,768 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Sensitivities computation via automated adjoint for a 2D shallow water model\n",
+    "\n",
+    "This notebook demonstrates how to compute the sensitivities (gradient) of a cost function with respect to control parameters using automated adjoints, a tool provided by the Firedrake library. The model used in this notebook is a shallow water model on a sphere. \n",
+    "\n",
+    "## Cost function and model\n",
+    "The cost function measures the summation of the kinetic and the potential energy of the fluid, which is modelled by the following expression:\n",
+    "\\begin{equation}\n",
+    "J = \\int_{\\Omega} \\left( \\frac{1}{2} \\textbf{u}(\\mathbf{x}, t_f) \\cdot \\textbf{u}(\\mathbf{x}, t_f) + \\frac{1}{2} g D^2(\\mathbf{x}, t_f) \\right) \\, dx,\n",
+    "\\end{equation}\n",
+    "where $\\textbf{u} (\\mathbf{x}, t_f)$ is the velocity of the fluid at the final step $t_f$, $D(\\mathbf{x}, t_f)$ is the fluid depth at the final step, and $g$ is the gravitational acceleration. The model is governed by the following set of equations:\n",
+    "\n",
+    "\\begin{align}\n",
+    "    \\textbf{u}_t + (\\textbf{u} \\cdot \\nabla) \\textbf{u} + f \\textbf{u}^{\\perp} + g \\nabla (D + b)  &= 0, \\tag{2}\\\\\n",
+    "    D_t + \\nabla \\cdot (D \\textbf{u}) &= 0, \\tag{3}\n",
+    "\\end{align}\n",
+    "where $f$ is the Coriolis parameter, the fluid depth is given by $D = H+h-b$, $H$ is the mean fluid depth, $h$ is the free surface height and $b$ is the topography. No boundary conditions are required as we are solving on a spherical domain, $\\Omega$."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Solve the model with Gusto\n",
+    "We first import the libraries required for the computations: ``gusto`` and ``firedrake``."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 80,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from firedrake import *\n",
+    "from gusto import *"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We then define the shallow water parameter $H$ (The fluid depth)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 81,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "H = 5960.0\n",
+    "parameters = ShallowWaterParameters(H=H)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We will use one of the spherical meshes provided by Firedrake: the ``IcosahedralSphereMesh``. As the spherical domain we are solving over is the Earth we specify the radius as 6371220m. The refinement level, ``ref_level``, specifies the number of times the base icosahedron is refined. The argument ``degree`` specifies the polynomial degree of the function space used to represent the coordinates."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 82,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "R = 6371220. # Earth radius in meters\n",
+    "mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "A domain is created with the ``Domain`` Gusto object, which holds the mesh and the function spaces defined on it. It also holds the model’s time interval."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 83,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dt = 900.\n",
+    "domain = Domain(mesh, dt, 'BDM', 1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can now set up the finite element form of the shallow water equations by passing ``ShallowWaterEquations`` Gusto class the following arguments: ``domain``, ``parameters``,  the expressions ``fexpr`` for the Coriolis parameter and ``bexpr`` for the bottom surface of the fluid."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 84,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = SpatialCoordinate(mesh)\n",
+    "Omega = parameters.Omega  # Spatial domain\n",
+    "fexpr = 2*Omega*x[2]/R  # Expression for the Coriolis parameter\n",
+    "# ``lonlatr_from_xyz`` is returning the required spherical longitude ``lamda``, latitude ``theta``.\n",
+    "lamda, theta, _ = lonlatr_from_xyz(x[0], x[1], x[2])\n",
+    "R0 = pi/9. # radius of mountain (rad)\n",
+    "R0sq = R0**2  \n",
+    "lamda_c = -pi/2. # longitudinal centre of mountain (rad)\n",
+    "lsq = (lamda - lamda_c)**2 \n",
+    "theta_c = pi/6.  # latitudinal centre of mountain (rad)\n",
+    "thsq = (theta - theta_c)**2\n",
+    "rsq = min_value(R0sq, lsq+thsq)\n",
+    "r = sqrt(rsq)\n",
+    "bexpr = 2000 * (1 - r/R0)  # An expression for the bottom surface of the fluid.\n",
+    "eqn = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=bexpr)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We specify instructions regarding the output using the ``OutputParameters`` class. The directory name, ``dirname``, must be specified. To prevent losing hard-earned simulation data, Gusto will not allow existing files to be overwritten. Hence, if one wishes to re-run a simulation with the same output filename, the existing results file needs to be moved or deleted first. This is also the place to specify the output frequency (in timesteps) of VTU files. The default is ``dumpfreq=1``, which outputs VTU files at every timestep (very useful when first setting up a problem!). Below we set ``dumpfreq=5``.\n",
+    "\n",
+    "We can specify which diagnostics to record over a simulation. The list of available diagnostics in the gusto source code: [diagnostics](https://github.com/firedrakeproject/gusto/blob/main/gusto/diagnostics/diagnostics.py). Since this flow should be in a steady state, it is also instructive to output the steady state error fields for both $\\textbf{u}$ and $D$ as this will tell us how close the simulation is to be correct. The errors should not grow in time. They should reduce as the mesh and timestep are refined. We pass these diagnostics into the ``IO`` class, which controls the input and output and stores the fields which will be updated at each timestep."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 85,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output = OutputParameters(dirname=\"adjoint_sw\", log_courant=False)\n",
+    "io = IO(domain, output)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We now start to build a setup for time discretisation. We use the ``SemiImplicitQuasiNewton`` approach, which splits the equation into `transport` terms and 'forcing' terms (i.e. everything that does not transport) and solves each separately. That allows for different time-steppers used to transport the velocity and depth fields. We are employing an Implicit Midpoint method for the velocity and an explicit strong stability preserving RK3 (SSPRK3) method for the depth. Since the Courant number for a stable SSPRK3 scheme is lower than that for the Implicit Midpoint method, we do two subcycles of the SSPRK3 scheme per timestep, allowing us to use a longer timestep overall. The full list of time-stepping methods is available at the [time discretisation](https://github.com/firedrakeproject/gusto/blob/main/gusto/time_discretisation.py) python file."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 86,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO     Physical parameters that take non-default values:\n",
+      "INFO:gusto:Physical parameters that take non-default values:\n",
+      "INFO     H: 5960.0\n",
+      "INFO:gusto:H: 5960.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Transport schemes\n",
+    "transported_fields = [TrapeziumRule(domain, \"u\"), SSPRK3(domain, \"D\")]\n",
+    "transport_methods = [DGUpwind(eqn, \"u\"), DGUpwind(eqn, \"D\")]\n",
+    "\n",
+    "# Time stepper\n",
+    "stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, transport_methods)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We are now ready to specify the initial conditions:\n",
+    "\\begin{align}\n",
+    "    \\textbf{u}_0 &= \\frac{u_{max}}{R} [-y,x,0], \\tag{4}\\\\\n",
+    "    D_0 &= H - \\frac{\\Omega u_{max} z^2}{g R} \\tag{5}\n",
+    "\\end{align}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 87,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Coefficient(WithGeometry(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x1682cee90>, FiniteElement('Discontinuous Lagrange', triangle, 1), name='L2', index=1, component=None), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 2), dim=3), 175601)), 248029)"
+      ]
+     },
+     "execution_count": 87,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "u0 = stepper.fields('u')\n",
+    "D0 = stepper.fields('D')\n",
+    "u_max = 20.   # Maximum amplitude of the zonal wind (m/s)\n",
+    "uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])\n",
+    "g = parameters.g  # acceleration due to gravity (m/s^2)\n",
+    "Rsq = R**2\n",
+    "Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*x[2]**2/Rsq)/g - bexpr\n",
+    "\n",
+    "u0.project(uexpr)\n",
+    "D0.interpolate(Dexpr)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "When using the `SemiImplicitQuasiNewton` time ``stepper``, we also have to set up any non-zero reference profiles."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 88,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Dbar = Function(D0.function_space()).assign(H)\n",
+    "stepper.set_reference_profiles([('D', Dbar)])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We are almost ready to run the model. Before we do so, we have to tape the model to compute the sensitivities via automated adjoint. Hence, at this point, we import ``firedrake.adjoint`` and enable the tape of the forward solver (Shallow Water model) with ``continue_annotation()``."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 89,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 89,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from firedrake.adjoint import *\n",
+    "continue_annotation()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We are finally ready to run the model!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 90,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "DEBUG:gusto:Dumping output to disk\n",
+      "INFO     \n",
+      "INFO:gusto:\n",
+      "INFO     ========================================\n",
+      "INFO:gusto:========================================\n",
+      "INFO     at start of timestep 1, t=0.0, dt=900.0\n",
+      "INFO:gusto:at start of timestep 1, t=0.0, dt=900.0\n",
+      "INFO     Semi-implicit Quasi Newton: Explicit forcing\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Explicit forcing\n",
+      "DEBUG:gusto:    Residual norms for ExplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 3.717272816438e+05 true resid norm 3.717272816438e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.620107299083e-10 true resid norm 1.620107299083e-10 ||r(i)||/||b|| 4.358322294556e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.585397212708e+05 true resid norm 6.127594999762e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.064757082310e+04 true resid norm 6.297866206201e+03 ||r(i)||/||b|| 1.027787607772e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 1.984167553564e+02 true resid norm 2.154172949430e+02 ||r(i)||/||b|| 3.515527624645e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 4.896138420464e+00 true resid norm 5.372370990865e+00 ||r(i)||/||b|| 8.767503386031e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 4.896138420921e+00 true resid norm 5.372370992372e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.213180858366e-01 true resid norm 1.505771328497e-01 ||r(i)||/||b|| 2.802805931748e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.386110502201e-03 true resid norm 7.233829275614e-03 ||r(i)||/||b|| 1.346487293205e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.888081698314e-04 true resid norm 2.155980867138e-04 ||r(i)||/||b|| 4.013090068052e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.508958609457e-06 true resid norm 8.412262624625e-06 ||r(i)||/||b|| 1.565837995285e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 5.901594446440e-07 true resid norm 5.901594446440e-07 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 3.718314622419e-22 true resid norm 3.718314622419e-22 ||r(i)||/||b|| 6.300525487076e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 3.478820869067e+05 true resid norm 3.478820869067e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.994472980877e-10 true resid norm 1.994472980877e-10 ||r(i)||/||b|| 5.733186777771e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.587569440754e+05 true resid norm 6.129105611037e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.065740127148e+04 true resid norm 6.297890541694e+03 ||r(i)||/||b|| 1.027538264368e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 1.985356265347e+02 true resid norm 2.156115172251e+02 ||r(i)||/||b|| 3.517830021347e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 4.881602606427e+00 true resid norm 5.364127741803e+00 ||r(i)||/||b|| 8.751893150843e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 4.881602606361e+00 true resid norm 5.364127745118e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.189409731723e-01 true resid norm 1.489319454693e-01 ||r(i)||/||b|| 2.776442928767e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.287887155869e-03 true resid norm 7.132861396961e-03 ||r(i)||/||b|| 1.329733693135e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.876689010863e-04 true resid norm 2.147829601784e-04 ||r(i)||/||b|| 4.004061245072e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.484304727953e-06 true resid norm 8.385391455195e-06 ||r(i)||/||b|| 1.563234854507e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.375704065564e+05 true resid norm 3.375200844565e+05 ||r(i)||/||b|| 2.453435247486e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 5.515968938792e-11 true resid norm 5.515968938792e-11 ||r(i)||/||b|| 4.009560687406e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.066019019371e+05 true resid norm 2.423054418141e+05 ||r(i)||/||b|| 2.272993609037e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 3.961061610235e-11 true resid norm 3.961061610235e-11 ||r(i)||/||b|| 3.715751349890e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "DEBUG:gusto:Leaving Semi-implicit Quasi-Newton timestep method\n",
+      "INFO     \n",
+      "INFO:gusto:\n",
+      "INFO     ========================================\n",
+      "INFO:gusto:========================================\n",
+      "INFO     at start of timestep 2, t=900.0, dt=900.0\n",
+      "INFO:gusto:at start of timestep 2, t=900.0, dt=900.0\n",
+      "INFO     Semi-implicit Quasi Newton: Explicit forcing\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Explicit forcing\n",
+      "DEBUG:gusto:    Residual norms for ExplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 5.362120374912e+04 true resid norm 3.859973858259e+05 ||r(i)||/||b|| 7.198596055991e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.893877795969e-11 true resid norm 1.893877795969e-11 ||r(i)||/||b|| 3.531956881890e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.604766785161e+05 true resid norm 6.152090559168e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.068399186592e+04 true resid norm 6.329441734408e+03 ||r(i)||/||b|| 1.028827790088e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.000708822109e+02 true resid norm 2.162420601441e+02 ||r(i)||/||b|| 3.514936232886e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 4.931342828052e+00 true resid norm 5.396506875620e+00 ||r(i)||/||b|| 8.771826135717e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 4.931342827824e+00 true resid norm 5.396506873935e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.227791166505e-01 true resid norm 1.521523930563e-01 ||r(i)||/||b|| 2.819460747677e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.471748032847e-03 true resid norm 7.301465459939e-03 ||r(i)||/||b|| 1.352998454464e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.898529502076e-04 true resid norm 2.167792649846e-04 ||r(i)||/||b|| 4.017029349701e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.555933266040e-06 true resid norm 8.456006918931e-06 ||r(i)||/||b|| 1.566940822363e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 5.922290117006e-07 true resid norm 1.066019019371e+05 ||r(i)||/||b|| 1.800011479190e+11\n",
+      "DEBUG:gusto:    1 KSP no resid norm 3.792661951956e-22 true resid norm 3.792661951956e-22 ||r(i)||/||b|| 6.404046200075e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 4.082629514187e+05 true resid norm 4.082629514187e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.613042924564e-10 true resid norm 1.613042924564e-10 ||r(i)||/||b|| 3.950990210987e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.613736253633e+05 true resid norm 6.164185032864e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.069411765753e+04 true resid norm 6.344717297526e+03 ||r(i)||/||b|| 1.029287288376e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.006967024553e+02 true resid norm 2.165291243542e+02 ||r(i)||/||b|| 3.512696701994e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 4.955181626652e+00 true resid norm 5.408670055491e+00 ||r(i)||/||b|| 8.774347341385e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 4.955181626602e+00 true resid norm 5.408670055861e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.250540665393e-01 true resid norm 1.538410757549e-01 ||r(i)||/||b|| 2.844342031702e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.571771905522e-03 true resid norm 7.390528689363e-03 ||r(i)||/||b|| 1.366422542517e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.910313160878e-04 true resid norm 2.176203600240e-04 ||r(i)||/||b|| 4.023546597897e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.593703303884e-06 true resid norm 8.487476707684e-06 ||r(i)||/||b|| 1.569235434964e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.547025377642e+05 true resid norm 4.944980624256e+05 ||r(i)||/||b|| 3.196444412432e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 6.244438529493e-11 true resid norm 6.244438529493e-11 ||r(i)||/||b|| 4.036416350848e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.129613821142e+05 true resid norm 2.644949396295e+05 ||r(i)||/||b|| 2.341463380486e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 5.188974790380e-11 true resid norm 5.188974790380e-11 ||r(i)||/||b|| 4.593582951327e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "DEBUG:gusto:Leaving Semi-implicit Quasi-Newton timestep method\n",
+      "INFO     \n",
+      "INFO:gusto:\n",
+      "INFO     ========================================\n",
+      "INFO:gusto:========================================\n",
+      "INFO     at start of timestep 3, t=1800.0, dt=900.0\n",
+      "INFO:gusto:at start of timestep 3, t=1800.0, dt=900.0\n",
+      "INFO     Semi-implicit Quasi Newton: Explicit forcing\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Explicit forcing\n",
+      "DEBUG:gusto:    Residual norms for ExplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 4.918470774046e+04 true resid norm 4.603785515452e+04 ||r(i)||/||b|| 9.360196953382e-01\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.709363306585e-11 true resid norm 1.709363306585e-11 ||r(i)||/||b|| 3.475395880372e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.628475674062e+05 true resid norm 6.207601980794e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.066580831543e+04 true resid norm 6.385116455001e+03 ||r(i)||/||b|| 1.028596304137e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.011618065183e+02 true resid norm 2.163257559986e+02 ||r(i)||/||b|| 3.484852229056e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 5.009107162475e+00 true resid norm 5.422766781313e+00 ||r(i)||/||b|| 8.735686982011e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.009107162336e+00 true resid norm 5.422766781943e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.310238285189e-01 true resid norm 1.585661985134e-01 ||r(i)||/||b|| 2.924082943073e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.804915566986e-03 true resid norm 7.593813246366e-03 ||r(i)||/||b|| 1.400357705157e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.933081298615e-04 true resid norm 2.185050909137e-04 ||r(i)||/||b|| 4.029402327263e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.653093695177e-06 true resid norm 8.546658817820e-06 ||r(i)||/||b|| 1.576069774249e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 5.845374237257e-07 true resid norm 1.129613821142e+05 ||r(i)||/||b|| 1.932491873560e+11\n",
+      "DEBUG:gusto:    1 KSP no resid norm 3.760932244576e-22 true resid norm 3.760932244576e-22 ||r(i)||/||b|| 6.434031581082e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 4.068447998955e+05 true resid norm 4.068447998955e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.302395903596e-10 true resid norm 1.302395903596e-10 ||r(i)||/||b|| 3.201210643298e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.637275978130e+05 true resid norm 6.223345449248e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.068196732817e+04 true resid norm 6.400945806866e+03 ||r(i)||/||b|| 1.028537763019e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.017322034986e+02 true resid norm 2.166513722980e+02 ||r(i)||/||b|| 3.481268620949e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 5.031640760550e+00 true resid norm 5.431779556009e+00 ||r(i)||/||b|| 8.728070135757e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.031640760269e+00 true resid norm 5.431779557485e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.327606111123e-01 true resid norm 1.598490164410e-01 ||r(i)||/||b|| 2.942848006795e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.868316502526e-03 true resid norm 7.641905040180e-03 ||r(i)||/||b|| 1.406887919384e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.945661057793e-04 true resid norm 2.190233858132e-04 ||r(i)||/||b|| 4.032258369384e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.688507454708e-06 true resid norm 8.579874687314e-06 ||r(i)||/||b|| 1.579569751775e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.199022954527e+05 true resid norm 4.055284066275e+05 ||r(i)||/||b|| 3.382157156344e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 5.290535120365e-11 true resid norm 5.290535120365e-11 ||r(i)||/||b|| 4.412371840247e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 9.656092509913e+04 true resid norm 2.151199852505e+05 ||r(i)||/||b|| 2.227816117437e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 4.368178127942e-11 true resid norm 4.368178127942e-11 ||r(i)||/||b|| 4.523753395545e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "DEBUG:gusto:Leaving Semi-implicit Quasi-Newton timestep method\n",
+      "INFO     \n",
+      "INFO:gusto:\n",
+      "INFO     ========================================\n",
+      "INFO:gusto:========================================\n",
+      "INFO     at start of timestep 4, t=2700.0, dt=900.0\n",
+      "INFO:gusto:at start of timestep 4, t=2700.0, dt=900.0\n",
+      "INFO     Semi-implicit Quasi Newton: Explicit forcing\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Explicit forcing\n",
+      "DEBUG:gusto:    Residual norms for ExplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 3.623731520079e+04 true resid norm 6.791995697676e+04 ||r(i)||/||b|| 1.874309854370e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.377006062130e-11 true resid norm 1.377006062130e-11 ||r(i)||/||b|| 3.799967118149e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.655756934086e+05 true resid norm 6.263209139889e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.071506957247e+04 true resid norm 6.441284811339e+03 ||r(i)||/||b|| 1.028432017433e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.023667345222e+02 true resid norm 2.177420973958e+02 ||r(i)||/||b|| 3.476526051303e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 5.037205816478e+00 true resid norm 5.414858593667e+00 ||r(i)||/||b|| 8.645501807023e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.037205816945e+00 true resid norm 5.414858596158e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.308305953126e-01 true resid norm 1.583788999825e-01 ||r(i)||/||b|| 2.924894476374e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.800830603278e-03 true resid norm 7.535497309851e-03 ||r(i)||/||b|| 1.391633258013e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.951814065919e-04 true resid norm 2.176322724719e-04 ||r(i)||/||b|| 4.019168157527e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.684913337251e-06 true resid norm 8.591881993703e-06 ||r(i)||/||b|| 1.586723243299e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 5.930426159714e-07 true resid norm 9.656092509912e+04 ||r(i)||/||b|| 1.628229110331e+11\n",
+      "DEBUG:gusto:    1 KSP no resid norm 3.788205038540e-22 true resid norm 3.788205038540e-22 ||r(i)||/||b|| 6.387745056626e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 2.870374044092e+05 true resid norm 2.870374044092e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.282656149014e-10 true resid norm 1.282656149014e-10 ||r(i)||/||b|| 4.468602799884e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.664936934029e+05 true resid norm 6.275513396511e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.075136327118e+04 true resid norm 6.442870975745e+03 ||r(i)||/||b|| 1.026668348653e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.030257661540e+02 true resid norm 2.186361917364e+02 ||r(i)||/||b|| 3.483957055337e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 5.034072905720e+00 true resid norm 5.418192805198e+00 ||r(i)||/||b|| 8.633863817756e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.034072905562e+00 true resid norm 5.418192806697e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.298184557140e-01 true resid norm 1.575351913962e-01 ||r(i)||/||b|| 2.907522803572e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.767180921667e-03 true resid norm 7.495866514393e-03 ||r(i)||/||b|| 1.383462490506e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.951345461690e-04 true resid norm 2.173722948979e-04 ||r(i)||/||b|| 4.011896635152e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.682088429556e-06 true resid norm 8.582316804832e-06 ||r(i)||/||b|| 1.583981432743e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.346881528775e+05 true resid norm 2.754043729863e+05 ||r(i)||/||b|| 2.044755734655e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 4.386840848005e-11 true resid norm 4.386840848005e-11 ||r(i)||/||b|| 3.257035421665e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.032755124426e+05 true resid norm 2.367036499603e+05 ||r(i)||/||b|| 2.291962967425e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 3.628393263694e-11 true resid norm 3.628393263694e-11 ||r(i)||/||b|| 3.513314219293e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "DEBUG:gusto:Leaving Semi-implicit Quasi-Newton timestep method\n",
+      "INFO     \n",
+      "INFO:gusto:\n",
+      "INFO     ========================================\n",
+      "INFO:gusto:========================================\n",
+      "INFO     at start of timestep 5, t=3600.0, dt=900.0\n",
+      "INFO:gusto:at start of timestep 5, t=3600.0, dt=900.0\n",
+      "INFO     Semi-implicit Quasi Newton: Explicit forcing\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Explicit forcing\n",
+      "DEBUG:gusto:    Residual norms for ExplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 2.993856955300e+04 true resid norm 3.914845825554e+04 ||r(i)||/||b|| 1.307626210606e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 1.117982053676e-11 true resid norm 1.117982053676e-11 ||r(i)||/||b|| 3.734253407454e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.712284756826e+05 true resid norm 6.318020945328e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.092236059309e+04 true resid norm 6.475907500247e+03 ||r(i)||/||b|| 1.024989875198e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.068099444904e+02 true resid norm 2.232405514586e+02 ||r(i)||/||b|| 3.533393659033e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 5.037550826001e+00 true resid norm 5.449947978305e+00 ||r(i)||/||b|| 8.626036579279e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.037550825972e+00 true resid norm 5.449947976511e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.262184219078e-01 true resid norm 1.548699263179e-01 ||r(i)||/||b|| 2.841677149679e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.660094330537e-03 true resid norm 7.385726291473e-03 ||r(i)||/||b|| 1.355192072164e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.943513776045e-04 true resid norm 2.167003730878e-04 ||r(i)||/||b|| 3.976191589750e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.676621445170e-06 true resid norm 8.553356571439e-06 ||r(i)||/||b|| 1.569438205337e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 0: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 5.802708268023e-07 true resid norm 1.032755124426e+05 ||r(i)||/||b|| 1.779781227530e+11\n",
+      "DEBUG:gusto:    1 KSP no resid norm 3.767179856699e-22 true resid norm 3.767179856699e-22 ||r(i)||/||b|| 6.492106241941e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (0, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 2.027107463164e+05 true resid norm 2.027107463164e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 7.869384266651e-11 true resid norm 7.869384266651e-11 ||r(i)||/||b|| 3.882075523697e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (0, 1)\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: u\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: u\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.733314345061e+05 true resid norm 6.333470343240e+05 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 1.098965232923e+04 true resid norm 6.488475328056e+03 ||r(i)||/||b|| 1.024473941838e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 2.084150171367e+02 true resid norm 2.252619676010e+02 ||r(i)||/||b|| 3.556690967084e-04\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 5.035195935951e+00 true resid norm 5.465094804057e+00 ||r(i)||/||b|| 8.628910388585e-06\n",
+      "DEBUG:gusto:    Residual norms for uTrapeziumRule_ solve\n",
+      "DEBUG:gusto:    0 KSP preconditioned resid norm 5.035195935774e+00 true resid norm 5.465094802197e+00 ||r(i)||/||b|| 1.000000000000e+00\n",
+      "DEBUG:gusto:    1 KSP preconditioned resid norm 2.252681932078e-01 true resid norm 1.541017486928e-01 ||r(i)||/||b|| 2.819745206082e-02\n",
+      "DEBUG:gusto:    2 KSP preconditioned resid norm 8.634772911632e-03 true resid norm 7.362233497288e-03 ||r(i)||/||b|| 1.347137380733e-03\n",
+      "DEBUG:gusto:    3 KSP preconditioned resid norm 2.938027804860e-04 true resid norm 2.166120739369e-04 ||r(i)||/||b|| 3.963555652316e-05\n",
+      "DEBUG:gusto:    4 KSP preconditioned resid norm 9.671700592498e-06 true resid norm 8.527053629568e-06 ||r(i)||/||b|| 1.560275519125e-06\n",
+      "INFO     Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Transport 1: D\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 0)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.401792544219e+05 true resid norm 2.581791465204e+05 ||r(i)||/||b|| 1.841778568342e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 5.654885791223e-11 true resid norm 5.654885791223e-11 ||r(i)||/||b|| 4.034038998527e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 0)\n",
+      "INFO     Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Implicit forcing (1, 1)\n",
+      "DEBUG:gusto:    Residual norms for ImplicitForcingSolver_ solve\n",
+      "DEBUG:gusto:    0 KSP no resid norm 1.099362982095e+05 true resid norm 2.494150460533e+05 ||r(i)||/||b|| 2.268723343568e+00\n",
+      "DEBUG:gusto:    1 KSP no resid norm 4.166161929637e-11 true resid norm 4.166161929637e-11 ||r(i)||/||b|| 3.789614529042e-16\n",
+      "INFO     Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "INFO:gusto:Semi-implicit Quasi Newton: Mixed solve (1, 1)\n",
+      "DEBUG:gusto:Leaving Semi-implicit Quasi-Newton timestep method\n",
+      "INFO     TIMELOOP complete. t=4500.00000, tmax=4500.00000\n",
+      "INFO:gusto:TIMELOOP complete. t=4500.00000, tmax=4500.00000\n"
+     ]
+    }
+   ],
+   "source": [
+    "timesteps = 5 \n",
+    "stepper.run(0., timesteps*dt)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, the final solution of the shallow water solver is used to compute the cost functional (1) as shown in the following cell code."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 91,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "u_tf = stepper.fields('u')  # Final velocity field\n",
+    "D_tf = stepper.fields('D')  # Final depth field\n",
+    "\n",
+    "J = assemble(0.5*inner(u_tf, u_tf)*dx + 0.5*g*D_tf**2*dx)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We define the control variable: initial fluid depth $D_0$ and the velocity $\\textbf{u}_0$. Finally, we run the sensitivity computation: gradient of the cost function with respect $D_0$ and $\\textbf{u}_0$. The gradient computation is achieved with the ``compute_gradient`` function from the `firedrake.adjoint` module."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 92,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "control = [Control(D0), Control(u0)] # Control variables\n",
+    "grad = compute_gradient(J, control)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can execute a gradient verification with second order [Taylor's test](https://www.dolfinadjoint.org/en/latest/documentation/verification.html). To do so, we define the reduced functional as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 93,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "J_hat = ReducedFunctional(J, control)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We then verify the gradient computation using the ``taylor_test``. We already taped the forward model. Any additional operation as below, for  Taylor's test, does not require the tape to be enabled. Thus, we can stop annotating."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 94,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running Taylor test\n",
+      "Computed residuals: [151981167509504.0, 37995189125120.0, 9498745905152.0, 2374616748032.0]\n",
+      "Computed convergence rates: [np.float64(2.0000039015457785), np.float64(2.0000078031232102), np.float64(2.0000423626811754)]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "with stop_annotating():\n",
+    "    # Stop annotation to perform the Taylor test\n",
+    "    h0 = Function(D0.function_space())\n",
+    "    h1 = Function(u0.function_space())\n",
+    "    h0.assign(D0 * np.random.rand())\n",
+    "    h1.assign(u0 * np.random.rand())\n",
+    "    taylor_test(J_hat, [D0, u0], [h0, h1])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "That is great! We have a sensitivity computation via automated adjoint with Gusto! In addition, we have verified the gradient computation with a second-order Taylor's test, which is a good practice to ensure the correctness of the gradient computation. See the results above are close to the expected values, 2.0."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "For a notebook demonstration of adjoints, it is good practice to clear the tape. This is because if we do not clear the tape, the tape will keep growing every time we run the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 95,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tape = get_working_tape()\n",
+    "tape.clear_tape()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "firedrake",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/examples/compressible_euler/unsaturated_bubble.py b/examples/compressible_euler/unsaturated_bubble.py
index 59a911cb9..394e15f97 100644
--- a/examples/compressible_euler/unsaturated_bubble.py
+++ b/examples/compressible_euler/unsaturated_bubble.py
@@ -216,11 +216,16 @@ def unsaturated_bubble(
     R_v = eqns.parameters.R_v
     epsilon = R_d / R_v
 
+    # make expressions for determining water_v0
+    exner = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0)
+    p = thermodynamics.p(eqns.parameters, exner)
+    T = thermodynamics.T(eqns.parameters, theta0, exner, water_v0)
+    r_v_expr = thermodynamics.r_v(eqns.parameters, rel_hum, T, p)
+
     # make expressions to evaluate residual
     exner_expr = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0)
     p_expr = thermodynamics.p(eqns.parameters, exner_expr)
     T_expr = thermodynamics.T(eqns.parameters, theta0, exner_expr, water_v0)
-    water_v_expr = thermodynamics.r_v(eqns.parameters, rel_hum, T_expr, p_expr)
     rel_hum_expr = thermodynamics.RH(eqns.parameters, water_v0, T_expr, p_expr)
     rel_hum_eval = Function(Vt)
 
@@ -242,7 +247,7 @@ def unsaturated_bubble(
 
         # first solve for r_v
         for _ in range(max_inner_solve_count):
-            water_v_eval.interpolate(water_v_expr)
+            water_v_eval.interpolate(r_v_expr)
             water_v0.assign(water_v0 * (1 - delta) + delta * water_v_eval)
 
             # compute theta_vd
diff --git a/gusto/core/configuration.py b/gusto/core/configuration.py
index 40852033a..5a1cbedb7 100644
--- a/gusto/core/configuration.py
+++ b/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"
 ]
 
 
@@ -79,11 +79,7 @@ def __setattr__(self, name, value):
 
         # Almost all parameters should be Constants -- but there are some
         # specific exceptions which should be kept as integers
-        non_constants = [
-            'dumpfreq', 'pddumpfreq', 'chkptfreq',
-            'fixed_subcycles', 'max_subcycles', 'subcycle_by_courant'
-        ]
-        if type(value) in [float, int] and name not in non_constants:
+        if type(value) in [float, int] and name not in ['dumpfreq', 'pddumpfreq', 'chkptfreq']:
             object.__setattr__(self, name, Constant(value))
         else:
             object.__setattr__(self, name, value)
@@ -167,7 +163,7 @@ class ShallowWaterParameters(Configuration):
 class WrapperOptions(Configuration, metaclass=ABCMeta):
     """Base class for specifying options for a transport scheme."""
 
-    @abstractproperty
+    @abstractmethod
     def name(self):
         pass
 
@@ -269,19 +265,6 @@ class BoundaryLayerParameters(Configuration):
     mu = 100.                   # Interior penalty coefficient for vertical diffusion
 
 
-class HeldSuarezParameters(Configuration):
-    """
-    Parameters used in the default configuration for the Held Suarez test case.
-    """
-    T0stra = 200               # Stratosphere temp
-    T0surf = 315               # Surface temperature at equator
-    T0horiz = 60               # Equator to pole temperature difference
-    T0vert = 10                # Stability parameter
-    sigmab = 0.7               # Height of the boundary layer
-    tau_d = 40 * 24 * 60 * 60  # 40 day time scale
-    tau_u = 4 * 24 * 60 * 60   # 4 day timescale
-
-
 class SubcyclingOptions(Configuration):
     """
     Describes the process of subcycling a time discretisation, by dividing the
@@ -308,3 +291,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 the float type parameters attributes to 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)))
diff --git a/gusto/core/domain.py b/gusto/core/domain.py
index 797f35e90..5aaf58bff 100644
--- a/gusto/core/domain.py
+++ b/gusto/core/domain.py
@@ -6,10 +6,9 @@
 
 from gusto.core.coordinates import Coordinates
 from gusto.core.function_spaces import Spaces, check_degree_args
-from firedrake import (
-    Constant, SpatialCoordinate, sqrt, CellNormal, cross, inner, grad,
-    VectorFunctionSpace, Function, FunctionSpace, perp, curl
-)
+from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross,
+                       inner, grad, VectorFunctionSpace, Function, FunctionSpace,
+                       perp)
 import numpy as np
 
 
@@ -64,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
@@ -125,14 +125,12 @@ def __init__(self, mesh, dt, family, degree=None,
                 V = VectorFunctionSpace(mesh, "DG", sphere_degree)
                 self.outward_normals = Function(V).interpolate(CellNormal(mesh))
                 self.perp = lambda u: cross(self.outward_normals, u)
-                self.divperp = lambda u: inner(self.outward_normals, curl(u))
         else:
             kvec = [0.0]*dim
             kvec[dim-1] = 1.0
             self.k = Constant(kvec)
             if dim == 2:
                 self.perp = perp
-                self.divperp = lambda u: -u[0].dx(1) + u[1].dx(0)
 
         # -------------------------------------------------------------------- #
         # Construct information relating to height/radius
diff --git a/gusto/core/function_spaces.py b/gusto/core/function_spaces.py
index c83a15cd6..97820d454 100644
--- a/gusto/core/function_spaces.py
+++ b/gusto/core/function_spaces.py
@@ -4,8 +4,8 @@
 """
 
 from gusto.core.logging import logger
-from firedrake import (HCurl, HDiv, FunctionSpace, FiniteElement, VectorElement,
-                       TensorProductElement, BrokenElement, interval)
+from firedrake import (HCurl, HDiv, FunctionSpace, FiniteElement,
+                       TensorProductElement, interval)
 
 __all__ = ["Spaces", "check_degree_args"]
 
@@ -71,7 +71,6 @@ def __init__(self, mesh):
         self.mesh = mesh
         self.extruded_mesh = hasattr(mesh, "_base_mesh")
         self.de_rham_complex = {}
-        self.continuity = {}
 
     def __call__(self, name):
         """
@@ -90,7 +89,7 @@ def __call__(self, name):
         else:
             raise ValueError(f'The space container has no space {name}')
 
-    def add_space(self, name, space, overwrite_space=False, continuity=None):
+    def add_space(self, name, space, overwrite_space=False):
         """
         Adds a function space to the container.
 
@@ -101,10 +100,6 @@ def add_space(self, name, space, overwrite_space=False, continuity=None):
             overwrite_space (bool, optional): Logical to allow space existing in
                 container to be overwritten by an incoming space. Defaults to
                 False.
-            continuity: Whether the space is continuous or not. For spaces on
-                extruded meshes, must be a dictionary with entries for the
-                'horizontal' and 'vertical' continuity.
-                If None, defaults to value of is_cg(space).
         """
 
         if hasattr(self, name) and not overwrite_space:
@@ -114,28 +109,9 @@ def add_space(self, name, space, overwrite_space=False, continuity=None):
 
         setattr(self, name, space)
 
-        # set the continuity of the space - for extruded meshes specify both directions
-        if continuity is None:
-            continuity = is_cg(space)
-            if self.extruded_mesh:
-                self.continuity[name] = {
-                    'horizontal': continuity,
-                    'vertical': continuity
-                }
-            else:
-                self.continuity[name] = continuity
-        else:
-            if self.extruded_mesh:
-                self.continuity[name] = {
-                    'horizontal': continuity['horizontal'],
-                    'vertical': continuity['vertical']
-                }
-            else:
-                self.continuity[name] = continuity
-
-    def create_space(self, name, family, degree, overwrite_space=False, continuity=None):
+    def create_space(self, name, family, degree, overwrite_space=False):
         """
-        Creates a space and adds it to the container.
+        Creates a space and adds it to the container..
 
         Args:
             name (str): the name to give to the space.
@@ -144,18 +120,18 @@ def create_space(self, name, family, degree, overwrite_space=False, continuity=N
             overwrite_space (bool, optional): Logical to allow space existing in
                 container to be overwritten by an incoming space. Defaults to
                 False.
-            continuity: Whether the space is continuous or not. For spaces on
-                extruded meshes, must be a tuple for the (horizontal, vertical)
-                continuity. If None, defaults to value of is_cg(space).
 
         Returns:
             :class:`FunctionSpace`: the desired function space.
         """
-        space = FunctionSpace(self.mesh, family, degree, name=name)
 
-        self.add_space(name, space,
-                       overwrite_space=overwrite_space,
-                       continuity=continuity)
+        if hasattr(self, name) and not overwrite_space:
+            raise RuntimeError(f'Space {name} already exists. If you really '
+                               + 'to create it then set `overwrite_space` as '
+                               + 'to be True')
+
+        space = FunctionSpace(self.mesh, family, degree, name=name)
+        setattr(self, name, space)
         return space
 
     def build_compatible_spaces(self, family, horizontal_degree,
@@ -205,15 +181,9 @@ def build_compatible_spaces(self, family, horizontal_degree,
         setattr(self, "L2"+complex_name, de_rham_complex.L2)
         # Register L2 space as DG also
         setattr(self, "DG"+complex_name, de_rham_complex.L2)
-        if hasattr(de_rham_complex, "theta"):
+        if hasattr(de_rham_complex, "theta"+complex_name):
             setattr(self, "theta"+complex_name, de_rham_complex.theta)
 
-        # Grab the continuity information from the complex
-        for space_type in ("H1, HCurl", "HDiv", "L2", "DG", "theta"):
-            space_name = space_type + complex_name
-            if hasattr(de_rham_complex, space_type):
-                self.continuity[space_name] = de_rham_complex.continuity[space_type]
-
     def build_dg1_equispaced(self):
         """
         Builds the equispaced variant of the DG1 function space, which is used in
@@ -228,15 +198,12 @@ def build_dg1_equispaced(self):
             hori_elt = FiniteElement('DG', cell, 1, variant='equispaced')
             vert_elt = FiniteElement('DG', interval, 1, variant='equispaced')
             V_elt = TensorProductElement(hori_elt, vert_elt)
-            continuity = {'horizontal': False, 'vertical': False}
         else:
             cell = self.mesh.ufl_cell().cellname()
             V_elt = FiniteElement('DG', cell, 1, variant='equispaced')
-            continuity = False
 
         space = FunctionSpace(self.mesh, V_elt, name='DG1_equispaced')
-
-        self.add_space('DG1_equispaced', space, continuity=continuity)
+        setattr(self, 'DG1_equispaced', space)
         return space
 
 
@@ -267,7 +234,6 @@ def __init__(self, mesh, family, horizontal_degree, vertical_degree=None,
         self.extruded_mesh = hasattr(mesh, '_base_mesh')
         self.family = family
         self.complex_name = complex_name
-        self.continuity = {}
         self.build_base_spaces(family, horizontal_degree, vertical_degree)
         self.build_compatible_spaces()
 
@@ -337,24 +303,15 @@ def build_compatible_spaces(self):
         if self.extruded_mesh:
             # Horizontal and vertical degrees
             # need specifying separately. Vtheta needs returning.
-            Vcg, continuity = self.build_h1_space()
-            self.continuity["H1"] = continuity
+            Vcg = self.build_h1_space()
             setattr(self, "H1", Vcg)
-
-            Vcurl, continuity = self.build_hcurl_space()
-            self.continuity["HCurl"] = continuity
+            Vcurl = self.build_hcurl_space()
             setattr(self, "HCurl", Vcurl)
-
-            Vu, continuity = self.build_hdiv_space()
-            self.continuity["HDiv"] = continuity
+            Vu = self.build_hdiv_space()
             setattr(self, "HDiv", Vu)
-
-            Vdg, continuity = self.build_l2_space()
-            self.continuity["L2"] = continuity
+            Vdg = self.build_l2_space()
             setattr(self, "L2", Vdg)
-
-            Vth, continuity = self.build_theta_space()
-            self.continuity["theta"] = continuity
+            Vth = self.build_theta_space()
             setattr(self, "theta", Vth)
 
             return Vcg, Vcurl, Vu, Vdg, Vth
@@ -363,20 +320,13 @@ def build_compatible_spaces(self):
             # 2D: two de Rham complexes (hcurl or hdiv) with 3 spaces
             # 3D: one de Rham complexes with 4 spaces
             # either way, build all spaces
-            Vcg, continuity = self.build_h1_space()
-            self.continuity["H1"] = continuity
+            Vcg = self.build_h1_space()
             setattr(self, "H1", Vcg)
-
-            Vcurl, continuity = self.build_hcurl_space()
-            self.continuity["HCurl"] = continuity
+            Vcurl = self.build_hcurl_space()
             setattr(self, "HCurl", Vcurl)
-
-            Vu, continuity = self.build_hdiv_space()
-            self.continuity["HDiv"] = continuity
+            Vu = self.build_hdiv_space()
             setattr(self, "HDiv", Vu)
-
-            Vdg, continuity = self.build_l2_space()
-            self.continuity["L2"] = continuity
+            Vdg = self.build_l2_space()
             setattr(self, "L2", Vdg)
 
             return Vcg, Vcurl, Vu, Vdg
@@ -384,18 +334,11 @@ def build_compatible_spaces(self):
         else:
             # 1D domain, de Rham complex has 2 spaces
             # CG, hdiv and hcurl spaces should be the same
-            Vcg, continuity = self.build_h1_space()
-
-            self.continuity["H1"] = continuity
+            Vcg = self.build_h1_space()
             setattr(self, "H1", Vcg)
-
             setattr(self, "HCurl", None)
-
-            self.continuity["HDiv"] = continuity
             setattr(self, "HDiv", Vcg)
-
-            Vdg, continuity = self.build_l2_space()
-            self.continuity["L2"] = continuity
+            Vdg = self.build_l2_space()
             setattr(self, "L2", Vdg)
 
             return Vcg, Vdg
@@ -417,7 +360,7 @@ def build_hcurl_space(self):
         """
         if hdiv_hcurl_dict[self.family] is None:
             logger.warning('There is no HCurl space for this family. Not creating one')
-            return None, None
+            return None
 
         if self.extruded_mesh:
             Vh_elt = HCurl(TensorProductElement(self.base_elt_hori_hcurl,
@@ -425,13 +368,10 @@ def build_hcurl_space(self):
             Vv_elt = HCurl(TensorProductElement(self.base_elt_hori_cg,
                                                 self.base_elt_vert_dg))
             V_elt = Vh_elt + Vv_elt
-            continuity = {'horizontal': True, 'vertical': True}
         else:
             V_elt = self.base_elt_hori_hcurl
-            continuity = True
 
-        space_name = 'HCurl'+self.complex_name
-        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
+        return FunctionSpace(self.mesh, V_elt, name='HCurl'+self.complex_name)
 
     def build_hdiv_space(self):
         """
@@ -447,13 +387,9 @@ def build_hdiv_space(self):
                                           self.base_elt_vert_cg)
             Vv_elt = HDiv(Vt_elt)
             V_elt = Vh_elt + Vv_elt
-            continuity = {'horizontal': True, 'vertical': True}
         else:
             V_elt = self.base_elt_hori_hdiv
-            continuity = True
-
-        space_name = 'HDiv'+self.complex_name
-        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
+        return FunctionSpace(self.mesh, V_elt, name='HDiv'+self.complex_name)
 
     def build_l2_space(self):
         """
@@ -465,13 +401,10 @@ def build_l2_space(self):
 
         if self.extruded_mesh:
             V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_dg)
-            continuity = {'horizontal': False, 'vertical': False}
         else:
             V_elt = self.base_elt_hori_dg
-            continuity = False
 
-        space_name = 'L2'+self.complex_name
-        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
+        return FunctionSpace(self.mesh, V_elt, name='L2'+self.complex_name)
 
     def build_theta_space(self):
         """
@@ -490,10 +423,8 @@ def build_theta_space(self):
         assert self.extruded_mesh, 'Cannot create theta space if mesh is not extruded'
 
         V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_cg)
-        continuity = {'horizontal': False, 'vertical': True}
 
-        space_name = 'theta'+self.complex_name
-        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
+        return FunctionSpace(self.mesh, V_elt, name='theta'+self.complex_name)
 
     def build_h1_space(self):
         """
@@ -509,13 +440,11 @@ def build_h1_space(self):
 
         if self.extruded_mesh:
             V_elt = TensorProductElement(self.base_elt_hori_cg, self.base_elt_vert_cg)
-            continuity = {'horizontal': True, 'vertical': True}
+
         else:
             V_elt = self.base_elt_hori_cg
-            continuity = True
 
-        space_name = 'H1'+self.complex_name
-        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
+        return FunctionSpace(self.mesh, V_elt, name='H1'+self.complex_name)
 
 
 def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree):
@@ -547,24 +476,3 @@ def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree):
         raise ValueError(f'Cannot pass both "degree" and "vertical_degree" to {name}')
     if not extruded_mesh and vertical_degree is not None:
         raise ValueError(f'Cannot pass "vertical_degree" to {name} if mesh is not extruded')
-
-
-def is_cg(V):
-    """
-    Checks if a :class:`FunctionSpace` is continuous.
-
-    Function to check if a given space, V, is CG. Broken elements are always
-    discontinuous; for vector elements we check the names of the Sobolev spaces
-    of the subelements and for all other elements we just check the Sobolev
-    space name.
-
-    Args:
-        V (:class:`FunctionSpace`): the space to check.
-    """
-    ele = V.ufl_element()
-    if isinstance(ele, BrokenElement):
-        return False
-    elif type(ele) == VectorElement:
-        return all([e.sobolev_space.name == "H1" for e in ele._sub_elements])
-    else:
-        return V.ufl_element().sobolev_space.name == "H1"
diff --git a/gusto/diagnostics/compressible_euler_diagnostics.py b/gusto/diagnostics/compressible_euler_diagnostics.py
index bbacc549e..ccb3a8595 100644
--- a/gusto/diagnostics/compressible_euler_diagnostics.py
+++ b/gusto/diagnostics/compressible_euler_diagnostics.py
@@ -734,8 +734,7 @@ def setup(self, domain, state_fields):
 
         bcs = self.equations.bcs['u']
 
-        imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs,
-                                                    constant_jacobian=True)
+        imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs)
         self.imbalance_solver = LinearVariationalSolver(imbalanceproblem)
         self.expr = dot(imbalance, domain.k)
         super().setup(domain, state_fields)
@@ -790,14 +789,12 @@ def setup(self, domain, state_fields):
         eqn_rhs = domain.dt * self.phi * (rain * dot(- v, domain.k) * rho / area) * ds_b
 
         # Compute area normalisation
-        area_prob = LinearVariationalProblem(eqn_lhs, area_rhs, area,
-                                             constant_jacobian=True)
+        area_prob = LinearVariationalProblem(eqn_lhs, area_rhs, area)
         area_solver = LinearVariationalSolver(area_prob)
         area_solver.solve()
 
         # setup solver
-        rain_prob = LinearVariationalProblem(eqn_lhs, eqn_rhs, self.flux,
-                                             constant_jacobian=True)
+        rain_prob = LinearVariationalProblem(eqn_lhs, eqn_rhs, self.flux)
         self.solver = LinearVariationalSolver(rain_prob)
         self.field = state_fields(self.name, space=DG0, dump=True, pick_up=True)
         # Initialise field to zero, if picking up this will be overridden
diff --git a/gusto/diagnostics/diagnostics.py b/gusto/diagnostics/diagnostics.py
index 00272edac..b287d6d36 100644
--- a/gusto/diagnostics/diagnostics.py
+++ b/gusto/diagnostics/diagnostics.py
@@ -1,15 +1,14 @@
 """Common diagnostic fields."""
 
 
-from firedrake import (dot, dx, Function, sqrt, TestFunction,
+from firedrake import (assemble, dot, dx, Function, sqrt, TestFunction,
                        TrialFunction, Constant, grad, inner, FacetNormal,
                        LinearVariationalProblem, LinearVariationalSolver,
                        ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, pi,
                        TensorFunctionSpace, SpatialCoordinate, as_vector,
-                       Projector, assemble, FunctionSpace, FiniteElement,
+                       Projector, Interpolator, FunctionSpace, FiniteElement,
                        TensorProductElement)
 from firedrake.assign import Assigner
-from firedrake.__future__ import interpolate
 from ufl.domain import extract_unique_domain
 
 from abc import ABCMeta, abstractmethod, abstractproperty
@@ -22,8 +21,7 @@
            "XComponent", "YComponent", "ZComponent", "MeridionalComponent",
            "ZonalComponent", "RadialComponent", "Energy", "KineticEnergy",
            "Sum", "Difference", "SteadyStateError", "Perturbation",
-           "Divergence", "TracerDensity", "IterativeDiagnosticField",
-           "CumulativeSum"]
+           "Divergence", "TracerDensity", "IterativeDiagnosticField"]
 
 
 class Diagnostics(object):
@@ -194,7 +192,7 @@ def setup(self, domain, state_fields, space=None):
 
             # Solve method must be declared in diagnostic's own setup routine
             if self.method == 'interpolate':
-                self.evaluator = interpolate(self.expr, self.space)
+                self.evaluator = Interpolator(self.expr, self.field)
             elif self.method == 'project':
                 self.evaluator = Projector(self.expr, self.field)
             elif self.method == 'assign':
@@ -208,7 +206,7 @@ def compute(self):
         logger.debug(f'Computing diagnostic {self.name} with {self.method} method')
 
         if self.method == 'interpolate':
-            self.field.assign(assemble(self.evaluator))
+            self.evaluator.interpolate()
         elif self.method == 'assign':
             self.evaluator.assign()
         elif self.method == 'project':
@@ -295,7 +293,7 @@ def setup(self, domain, state_fields, space=None):
 
             # Solve method must be declared in diagnostic's own setup routine
             if self.method == 'interpolate':
-                self.evaluator = interpolate(self.expr, self.space)
+                self.evaluator = Interpolator(self.expr, self.field)
             elif self.method == 'project':
                 self.evaluator = Projector(self.expr, self.field)
             elif self.method == 'assign':
@@ -515,8 +513,7 @@ def setup(self, domain, state_fields):
             L = -inner(div(test), f)*dx
             if space.extruded:
                 L += dot(dot(test, n), f)*(ds_t + ds_b)
-            prob = LinearVariationalProblem(a, L, self.field,
-                                            constant_jacobian=True)
+            prob = LinearVariationalProblem(a, L, self.field)
             self.evaluator = LinearVariationalSolver(prob)
 
 
@@ -1043,43 +1040,3 @@ def setup(self, domain, state_fields):
 
         else:
             super().setup(domain, state_fields)
-
-
-class CumulativeSum(DiagnosticField):
-    """Base diagnostic for cumulatively summing a field in time."""
-    def __init__(self, name):
-        """
-        Args:
-            name (str): name of the field to take the cumulative sum of.
-        """
-        self.field_name = name
-        self.integral_name = name+"_cumulative"
-        super().__init__(self, method='assign', required_fields=(self.field_name,))
-
-    def setup(self, domain, state_fields):
-        """
-        Sets up the :class:`Function` for the diagnostic field.
-
-        Args:
-            domain (:class:`Domain`): the model's domain object.
-            state_fields (:class:`StateFields`): the model's field container.
-        """
-
-        # Gather the field to be summed
-        self.integrand = state_fields(self.field_name)
-        self.space = self.integrand.function_space()
-
-        # Create a new field to store the cumulative sum
-        self.field = state_fields(self.integral_name, space=self.space, dump=True, pick_up=True)
-        # Initialise the new field to zero, if picking up from a checkpoint
-        # file the original cumulative field will load and not be overwritten.
-        self.field.assign(0.0)
-
-    def compute(self):
-        """Increment the cumulative sum diagnostic."""
-        self.field.assign(self.field + self.integrand)
-
-    @property
-    def name(self):
-        """Gives the name of this diagnostic field."""
-        return self.integral_name
diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py
index 024e58ecf..6484b2afe 100644
--- a/gusto/diagnostics/shallow_water_diagnostics.py
+++ b/gusto/diagnostics/shallow_water_diagnostics.py
@@ -2,10 +2,9 @@
 
 
 from firedrake import (
-    dx, TestFunction, TrialFunction, grad, inner, curl, Function, assemble,
+    dx, TestFunction, TrialFunction, grad, inner, curl, Function, Interpolator,
     LinearVariationalProblem, LinearVariationalSolver, conditional
 )
-from firedrake.__future__ import interpolate
 from gusto.diagnostics.diagnostics import DiagnosticField, Energy
 
 __all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy",
@@ -183,18 +182,15 @@ def setup(self, domain, state_fields, vorticity_type=None):
 
             if vorticity_type == "potential":
                 a = q*gamma*D*dx
-                constant_jacobian = False
             else:
                 a = q*gamma*dx
-                constant_jacobian = True
 
             L = (- inner(domain.perp(grad(gamma)), u))*dx
             if vorticity_type != "relative":
                 f = state_fields("coriolis")
                 L += gamma*f*dx
 
-            problem = LinearVariationalProblem(a, L, self.field,
-                                               constant_jacobian=constant_jacobian)
+            problem = LinearVariationalProblem(a, L, self.field)
             self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"})
 
 
@@ -322,14 +318,14 @@ def setup(self, domain, state_fields):
 
         qsat_expr = self.equation.compute_saturation(state_fields.X(
             self.equation.field_name))
-        self.qsat_interpolate = interpolate(qsat_expr, space)
+        self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func)
         self.expr = conditional(q_t < self.qsat_func, q_t, self.qsat_func)
 
         super().setup(domain, state_fields, space=space)
 
     def compute(self):
         """Performs the computation of the diagnostic field."""
-        self.qsat_func.assign(assemble(self.qsat_interpolate))
+        self.qsat_interpolator.interpolate()
         super().compute()
 
 
@@ -372,7 +368,7 @@ def setup(self, domain, state_fields):
 
         qsat_expr = self.equation.compute_saturation(state_fields.X(
             self.equation.field_name))
-        self.qsat_interpolate = interpolate(qsat_expr, space)
+        self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func)
         vapour = conditional(q_t < self.qsat_func, q_t, self.qsat_func)
         self.expr = q_t - vapour
 
@@ -380,5 +376,5 @@ def setup(self, domain, state_fields):
 
     def compute(self):
         """Performs the computation of the diagnostic field."""
-        self.qsat_func.assign(assemble(self.qsat_interpolate))
+        self.qsat_interpolator.interpolate()
         super().compute()
diff --git a/gusto/equations/boussinesq_equations.py b/gusto/equations/boussinesq_equations.py
index b9823f110..3dcba43d3 100644
--- a/gusto/equations/boussinesq_equations.py
+++ b/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,10 @@ def __init__(self, domain, parameters,
                          active_tracers=active_tracers)
 
         self.parameters = parameters
+        # This function converts the parameters to real space.
+        # This is a preventive way to avoid adjoint issues when the parameters
+        # attribute are the control in the sensitivity computations.
+        convert_parameters_to_real_space(parameters, domain.mesh)
         self.compressible = compressible
 
         w, phi, gamma = self.tests[0:3]
@@ -168,10 +173,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
diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py
index 73bc76d45..0dbe8a01d 100644
--- a/gusto/equations/compressible_euler_equations.py
+++ b/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,10 @@ def __init__(self, domain, parameters, sponge_options=None,
                          active_tracers=active_tracers)
 
         self.parameters = parameters
+        # This function converts the parameters to real space.
+        # This is a preventive way to avoid adjoint issues when the parameters
+        # attribute are the control in the sensitivity computations.
+        convert_parameters_to_real_space(parameters, domain.mesh)
         g = parameters.g
         cp = parameters.cp
 
@@ -109,6 +113,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)
 
diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py
index fd61cedf7..aba1ed708 100644
--- a/gusto/equations/shallow_water_equations.py
+++ b/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
+        # This function converts the ``float`` and ``firedrake.Constant`` parameters
+        # attributes to a function in real space.
+        # This is a preventive way to avoid adjoint issues when the parameters
+        # attribute are the control in the sensitivity computations.
+        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)
 
@@ -417,7 +424,11 @@ def _setup_residual(self, fexpr, topog_expr, u_transport_option):
         # label these as the equivalent pressure gradient term and
         # provide linearisation
         if self.equivalent_buoyancy:
-            beta2 = self.parameters.beta2
+            try:
+                beta2 = self.parameters.beta2
+            except ValueError:
+                print("Oi")
+
             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 +658,10 @@ def __init__(self, domain, parameters,
                          active_tracers=active_tracers)
 
         self.parameters = parameters
+        # This function converts the parameters to real space.
+        # This is a preventive way to avoid adjoint issues when the parameters
+        # attribute are the control in the sensitivity computations.
+        convert_parameters_to_real_space(parameters, domain.mesh)
         g = parameters.g
         H = parameters.H
 
diff --git a/gusto/physics/__init__.py b/gusto/physics/__init__.py
index b1407a299..91f09c2ec 100644
--- a/gusto/physics/__init__.py
+++ b/gusto/physics/__init__.py
@@ -2,5 +2,4 @@
 from gusto.physics.chemistry import *                       # noqa
 from gusto.physics.boundary_and_turbulence import *         # noqa
 from gusto.physics.microphysics import *                    # noqa
-from gusto.physics.shallow_water_microphysics import *      # noqa
-from gusto.physics.held_suarez_forcing import *             # noqa
+from gusto.physics.shallow_water_microphysics import *      # noqa
\ No newline at end of file
diff --git a/gusto/physics/held_suarez_forcing.py b/gusto/physics/held_suarez_forcing.py
deleted file mode 100644
index d7573d009..000000000
--- a/gusto/physics/held_suarez_forcing.py
+++ /dev/null
@@ -1,188 +0,0 @@
-import numpy as np
-from firedrake import (Interpolator, Function, dx, pi, SpatialCoordinate,
-                       split, conditional, ge, sin, dot, ln, cos, inner, Projector)
-from firedrake.fml import subject
-from gusto.core.coord_transforms import lonlatr_from_xyz
-from gusto.recovery import Recoverer, BoundaryMethod
-from gusto.physics.physics_parametrisation import PhysicsParametrisation
-from gusto.core.labels import prognostic
-from gusto.equations import thermodynamics
-from gusto.core.configuration import HeldSuarezParameters
-from gusto.core import logger
-
-
-class Relaxation(PhysicsParametrisation):
-    """
-    Relaxation term for Held Suarez
-    """
-
-    def __init__(self, equation, variable_name, parameters, hs_parameters=None):
-        """
-        Args:
-            equation (:class:`PrognosticEquationSet`): the model's equation.
-            variable_name (str): the name of the variable
-            hs_parameters (:class'Configuration'): contains the parameters for the Held-suariez test case
-
-        """
-        label_name = f'relaxation_{variable_name}'
-        if hs_parameters is None:
-            hs_parameters = HeldSuarezParameters()
-            logger.warning('Using default Held-Suarez parameters')
-        super().__init__(equation, label_name, hs_parameters)
-
-        if equation.domain.on_sphere:
-            x, y, z = SpatialCoordinate(equation.domain.mesh)
-            _, lat, _ = lonlatr_from_xyz(x, y, z)
-        else:
-            # TODO: this could be determined some other way
-            # Take a mid-latitude
-            lat = pi / 4
-
-        self.X = Function(equation.X.function_space())
-        X = self.X
-        self.domain = equation.domain
-        theta_idx = equation.field_names.index('theta')
-        self.theta = X.subfunctions[theta_idx]
-        Vt = equation.domain.spaces('theta')
-        rho_idx = equation.field_names.index('rho')
-        rho = split(X)[rho_idx]
-
-        boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None
-        self.rho_averaged = Function(Vt)
-        self.rho_recoverer = Recoverer(rho, self.rho_averaged, boundary_method=boundary_method)
-        self.exner = Function(Vt)
-        self.exner_interpolator = Interpolator(
-            thermodynamics.exner_pressure(equation.parameters,
-                                          self.rho_averaged, self.theta), self.exner)
-        self.sigma = Function(Vt)
-        kappa = equation.parameters.kappa
-
-        T0surf = hs_parameters.T0surf
-        T0horiz = hs_parameters.T0horiz
-        T0vert = hs_parameters.T0vert
-        T0stra = hs_parameters.T0stra
-
-        sigma_b = hs_parameters.sigmab
-        tau_d = hs_parameters.tau_d
-        tau_u = hs_parameters.tau_u
-
-        theta_condition = (T0surf - T0horiz * sin(lat)**2 - (T0vert * ln(self.exner) * cos(lat)**2)/kappa)
-        Theta_eq = conditional(T0stra/self.exner >= theta_condition, T0stra/self.exner, theta_condition)
-
-        # timescale of temperature forcing
-        tau_cond = (self.sigma**(1/kappa) - sigma_b) / (1 - sigma_b)
-        newton_freq = 1 / tau_d + (1/tau_u - 1/tau_d) * conditional(0 >= tau_cond, 0, tau_cond) * cos(lat)**4
-        forcing_expr = newton_freq * (self.theta - Theta_eq)
-
-        # Create source for forcing
-        self.source_relaxation = Function(Vt)
-        self.source_interpolator = Interpolator(forcing_expr, self.source_relaxation)
-
-        # Add relaxation term to residual
-        test = equation.tests[theta_idx]
-        dx_reduced = dx(degree=equation.domain.max_quad_degree)
-        forcing_form = test * self.source_relaxation * dx_reduced
-        equation.residual += self.label(subject(prognostic(forcing_form, 'theta'), X), self.evaluate)
-
-    def evaluate(self, x_in, dt):
-        """
-        Evalutes the source term generated by the physics.
-
-        Args:
-            x_in: (:class:`Function`): the (mixed) field to be evolved.
-            dt: (:class:`Constant`): the timestep, which can be the time
-                interval for the scheme.
-        """
-        self.X.assign(x_in)
-        self.rho_recoverer.project()
-        self.exner_interpolator.interpolate()
-
-        # Determine sigma:= exner / exner_surf
-        exner_columnwise, index_data = self.domain.coords.get_column_data(self.exner, self.domain)
-        sigma_columnwise = np.zeros_like(exner_columnwise)
-        for col in range(len(exner_columnwise[:, 0])):
-            sigma_columnwise[col, :] = exner_columnwise[col, :] / exner_columnwise[col, 0]
-        self.domain.coords.set_field_from_column_data(self.sigma, sigma_columnwise, index_data)
-
-        self.source_interpolator.interpolate()
-
-
-class RayleighFriction(PhysicsParametrisation):
-    """
-    Forcing term on the velocity of the form
-    F_u = -u / a,
-    where a is some friction factor
-    """
-    def __init__(self, equation, hs_parameters=None):
-        """
-         Args:
-            equation (:class:`PrognosticEquationSet`): the model's equation.
-            hs_parameters (:class'Configuration'): contains the parameters for the Held-suariez test case
-        """
-        label_name = 'rayleigh_friction'
-        if hs_parameters is None:
-            hs_parameters = HeldSuarezParameters()
-            logger.warning('Using default Held-Suarez parameters')
-        super().__init__(equation, label_name, hs_parameters)
-
-        self.domain = equation.domain
-        self.X = Function(equation.X.function_space())
-        X = self.X
-        k = equation.domain.k
-        u_idx = equation.field_names.index('u')
-        u = split(X)[u_idx]
-        theta_idx = equation.field_names.index('theta')
-        self.theta = X.subfunctions[theta_idx]
-        rho_idx = equation.field_names.index('rho')
-        rho = split(X)[rho_idx]
-        Vt = equation.domain.spaces('theta')
-        Vu = equation.domain.spaces('HDiv')
-        u_hori = u - k*dot(u, k)
-
-        boundary_method = BoundaryMethod.extruded if self.domain == 0 else None
-        self.rho_averaged = Function(Vt)
-        self.exner = Function(Vt)
-        self.rho_recoverer = Recoverer(rho, self.rho_averaged, boundary_method=boundary_method)
-        self.exner_interpolator = Interpolator(
-            thermodynamics.exner_pressure(equation.parameters,
-                                          self.rho_averaged, self.theta), self.exner)
-
-        self.sigma = Function(Vt)
-        sigmab = hs_parameters.sigmab
-        kappa = equation.parameters.kappa
-        tau_fric = 24 * 60 * 60
-
-        tau_cond = (self.sigma**(1/kappa) - sigmab) / (1 - sigmab)
-        wind_timescale = conditional(ge(0, tau_cond), 0, tau_cond) / tau_fric
-        forcing_expr = u_hori * wind_timescale
-
-        self.source_friction = Function(Vu)
-        self.source_projector = Projector(forcing_expr, self.source_friction)
-
-        tests = equation.tests
-        test = tests[u_idx]
-        dx_reduced = dx(degree=equation.domain.max_quad_degree)
-        source_form = inner(test, self.source_friction) * dx_reduced
-        equation.residual += self.label(subject(prognostic(source_form, 'u'), X), self.evaluate)
-
-    def evaluate(self, x_in, dt):
-        """
-        Evaluates the source term generated by the physics. This does nothing if
-        the implicit formulation is not used.
-
-        Args:
-            x_in: (:class: 'Function'): the (mixed) field to be evolved.
-            dt: (:class: 'Constant'): the timestep, which can be the time
-                interval for the scheme.
-        """
-        self.X.assign(x_in)
-        self.rho_recoverer.project()
-        self.exner_interpolator.interpolate()
-        # Determine sigma:= exner / exner_surf
-        exner_columnwise, index_data = self.domain.coords.get_column_data(self.exner, self.domain)
-        sigma_columnwise = np.zeros_like(exner_columnwise)
-        for col in range(len(exner_columnwise[:, 0])):
-            sigma_columnwise[col, :] = exner_columnwise[col, :] / exner_columnwise[col, 0]
-        self.domain.coords.set_field_from_column_data(self.sigma, sigma_columnwise, index_data)
-
-        self.source_projector.project()
diff --git a/gusto/physics/microphysics.py b/gusto/physics/microphysics.py
index a5cbe3da5..e8313ab11 100644
--- a/gusto/physics/microphysics.py
+++ b/gusto/physics/microphysics.py
@@ -4,10 +4,9 @@
 """
 
 from firedrake import (
-    conditional, Function, dx, min_value, max_value, Constant, pi,
-    Projector, assemble
+    Interpolator, conditional, Function, dx, min_value, max_value, Constant, pi,
+    Projector
 )
-from firedrake.__future__ import interpolate
 from firedrake.fml import identity, Term, subject
 from gusto.equations import Phases, TracerVariableType
 from gusto.recovery import Recoverer, BoundaryMethod
@@ -171,7 +170,8 @@ def __init__(self, equation, vapour_name='water_vapour',
         # Add terms to equations and make interpolators
         # -------------------------------------------------------------------- #
         self.source = [Function(V) for factor in factors]
-        self.source_interpolate = [interpolate(sat_adj_expr*factor, V) for factor in factors]
+        self.source_interpolators = [Interpolator(sat_adj_expr*factor, source)
+                                     for factor, source in zip(factors, self.source)]
 
         tests = [equation.tests[idx] for idx in V_idxs]
 
@@ -195,8 +195,8 @@ def evaluate(self, x_in, dt):
         if isinstance(self.equation, CompressibleEulerEquations):
             self.rho_recoverer.project()
         # Evaluate the source
-        for interpolator, src in zip(self.source_interpolate, self.source):
-            src.assign(assemble(interpolator))
+        for interpolator in self.source_interpolators:
+            interpolator.interpolate()
 
 
 class AdvectedMoments(Enum):
@@ -440,7 +440,7 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain',
                                                         min_value(accu_rate, self.cloud_water / self.dt),
                                                         min_value(accr_rate + accu_rate, self.cloud_water / self.dt))))
 
-        self.source_interpolate = interpolate(rain_expr, Vt)
+        self.source_interpolator = Interpolator(rain_expr, self.source)
 
         # Add term to equation's residual
         test_cl = equation.tests[self.cloud_idx]
@@ -464,7 +464,7 @@ def evaluate(self, x_in, dt):
         self.rain.assign(x_in.subfunctions[self.rain_idx])
         self.cloud_water.assign(x_in.subfunctions[self.cloud_idx])
         # Evaluate the source
-        self.source.assign(assemble(self.source_interpolate))
+        self.source.assign(self.source_interpolator.interpolate())
 
 
 class EvaporationOfRain(PhysicsParametrisation):
@@ -609,7 +609,8 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour',
         # Add terms to equations and make interpolators
         # -------------------------------------------------------------------- #
         self.source = [Function(V) for factor in factors]
-        self.source_interpolate = [interpolate(evap_rate*factor, V) for factor in factors]
+        self.source_interpolators = [Interpolator(evap_rate*factor, source)
+                                     for factor, source in zip(factors, self.source)]
 
         tests = [equation.tests[idx] for idx in V_idxs]
 
@@ -633,5 +634,5 @@ def evaluate(self, x_in, dt):
         if isinstance(self.equation, CompressibleEulerEquations):
             self.rho_recoverer.project()
         # Evaluate the source
-        for interpolator, src in zip(self.source_interpolate, self.source):
-            src.assign(assemble(interpolator))
+        for interpolator in self.source_interpolators:
+            interpolator.interpolate()
diff --git a/gusto/physics/physics_parametrisation.py b/gusto/physics/physics_parametrisation.py
index c304c23f6..2d43dae48 100644
--- a/gusto/physics/physics_parametrisation.py
+++ b/gusto/physics/physics_parametrisation.py
@@ -8,8 +8,7 @@
 """
 
 from abc import ABCMeta, abstractmethod
-from firedrake import Function, dx, Projector, assemble
-from firedrake.__future__ import interpolate
+from firedrake import Interpolator, Function, dx, Projector
 from firedrake.fml import subject
 from gusto.core.labels import PhysicsLabel
 from gusto.core.logging import logger
@@ -118,14 +117,14 @@ def __init__(self, equation, variable_name, rate_expression,
 
         # Handle method of evaluating source/sink
         if self.method == 'interpolate':
-            self.source_interpolate = interpolate(expression, V)
+            self.source_interpolator = Interpolator(expression, V)
         else:
             self.source_projector = Projector(expression, V)
 
         # If not time-varying, evaluate for the first time here
         if not self.time_varying:
             if self.method == 'interpolate':
-                self.source.assign(assemble(self.source_interpolate))
+                self.source.assign(self.source_interpolator.interpolate())
             else:
                 self.source.assign(self.source_projector.project())
 
@@ -141,7 +140,7 @@ def evaluate(self, x_in, dt):
         if self.time_varying:
             logger.info(f'Evaluating physics parametrisation {self.label.label}')
             if self.method == 'interpolate':
-                self.source.assign(assemble(self.source_interpolate))
+                self.source.assign(self.source_interpolator.interpolate())
             else:
                 self.source.assign(self.source_projector.project())
         else:
diff --git a/gusto/physics/shallow_water_microphysics.py b/gusto/physics/shallow_water_microphysics.py
index 7f1b91227..c7d7c96cd 100644
--- a/gusto/physics/shallow_water_microphysics.py
+++ b/gusto/physics/shallow_water_microphysics.py
@@ -3,9 +3,8 @@
 """
 
 from firedrake import (
-    conditional, Function, dx, min_value, max_value, Constant, assemble
+    Interpolator, conditional, Function, dx, min_value, max_value, FunctionSpace
 )
-from firedrake.__future__ import interpolate
 from firedrake.fml import subject
 from gusto.core.logging import logger
 from gusto.physics.physics_parametrisation import PhysicsParametrisation
@@ -94,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:
@@ -135,7 +135,7 @@ def __init__(self, equation, saturation_curve,
                                             self.evaluate)
 
         # interpolator does the conversion of vapour to rain
-        self.source_interpolate = interpolate(conditional(
+        self.source_interpolator = Interpolator(conditional(
             self.water_v > self.saturation_curve,
             (1/self.tau)*gamma_r*(self.water_v - self.saturation_curve),
             0), Vv)
@@ -160,7 +160,7 @@ def evaluate(self, x_in, dt):
         if self.set_tau_to_dt:
             self.tau.assign(dt)
         self.water_v.assign(x_in.subfunctions[self.Vv_idx])
-        self.source.assign(assemble(self.source_interpolate))
+        self.source.assign(self.source_interpolator.interpolate())
 
 
 class SWSaturationAdjustment(PhysicsParametrisation):
@@ -270,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:
@@ -322,8 +323,8 @@ def __init__(self, equation, saturation_curve,
         # Add terms to equations and make interpolators
         # sources have the same order as V_idxs and factors
         self.source = [Function(Vc) for factor in factors]
-        self.source_interpolate = [interpolate(sat_adj_expr*factor, Vc)
-                                   for factor in factors]
+        self.source_interpolators = [Interpolator(sat_adj_expr*factor, source)
+                                     for factor, source in zip(factors, self.source)]
 
         # test functions have the same order as factors and sources (vapour,
         # cloud, depth, buoyancy) so that the correct test function multiplies
@@ -360,5 +361,5 @@ def evaluate(self, x_in, dt):
         self.cloud.assign(x_in.subfunctions[self.Vc_idx])
         if self.time_varying_gamma_v:
             self.gamma_v.interpolate(self.gamma_v_computation(x_in))
-        for interpolator, src in zip(self.source_interpolate, self.source):
-            src.assign(assemble(interpolator))
+        for interpolator in self.source_interpolators:
+            interpolator.interpolate()
diff --git a/gusto/recovery/recovery.py b/gusto/recovery/recovery.py
index 9d186dd48..b88c6ad6a 100644
--- a/gusto/recovery/recovery.py
+++ b/gusto/recovery/recovery.py
@@ -13,8 +13,7 @@
                        Function, FunctionSpace, Interpolator, Projector,
                        SpatialCoordinate, TensorProductElement,
                        VectorFunctionSpace, as_vector, function, interval,
-                       VectorElement, assemble)
-from firedrake.__future__ import interpolate
+                       VectorElement)
 from gusto.recovery import Averager
 from .recovery_kernels import (BoundaryRecoveryExtruded, BoundaryRecoveryHCurl,
                                BoundaryGaussianElimination)
@@ -145,14 +144,14 @@ def __init__(self, x_inout, method=BoundaryMethod.extruded, eff_coords=None):
             V_broken = FunctionSpace(mesh, BrokenElement(V_inout.ufl_element()))
             self.x_DG1_wrong = Function(V_broken)
             self.x_DG1_correct = Function(V_broken)
-            self.interpolate = interpolate(self.x_inout, V_broken)
+            self.interpolator = Interpolator(self.x_inout, self.x_DG1_wrong)
             self.averager = Averager(self.x_DG1_correct, self.x_inout)
             self.kernel = BoundaryGaussianElimination(V_broken)
 
     def apply(self):
         """Applies the boundary recovery process."""
         if self.method == BoundaryMethod.taylor:
-            self.x_DG1_wrong.assign(assemble(self.interpolate))
+            self.interpolator.interpolate()
             self.kernel.apply(self.x_DG1_wrong, self.x_DG1_correct,
                               self.act_coords, self.eff_coords, self.num_ext)
             self.averager.project()
@@ -276,7 +275,7 @@ def __init__(self, x_in, x_out, method='interpolate', boundary_method=None):
                         self.boundary_recoverers.append(BoundaryRecoverer(x_out_scalars[i],
                                                                           method=BoundaryMethod.taylor,
                                                                           eff_coords=eff_coords[i]))
-                    self.interpolate_to_vector = interpolate(as_vector(x_out_scalars), V_out)
+                    self.interpolate_to_vector = Interpolator(as_vector(x_out_scalars), self.x_out)
 
     def project(self):
         """Perform the whole recovery step."""
@@ -295,7 +294,7 @@ def project(self):
                     # Correct at boundaries
                     boundary_recoverer.apply()
                 # Combine the components to obtain the vector field
-                self.x_out.assign(assemble(self.interpolate_to_vector))
+                self.interpolate_to_vector.interpolate()
             else:
                 # Extrapolate at boundaries
                 self.boundary_recoverer.apply()
diff --git a/gusto/recovery/reversible_recovery.py b/gusto/recovery/reversible_recovery.py
index d87449817..fc8fbd332 100644
--- a/gusto/recovery/reversible_recovery.py
+++ b/gusto/recovery/reversible_recovery.py
@@ -4,8 +4,7 @@
 """
 
 from gusto.core.conservative_projection import ConservativeProjector
-from firedrake import (Projector, Function, assemble)
-from firedrake.__future__ import interpolate
+from firedrake import (Projector, Function, Interpolator)
 from .recovery import Recoverer
 
 __all__ = ["ReversibleRecoverer", "ConservativeRecoverer"]
@@ -53,7 +52,7 @@ def __init__(self, source_field, target_field, reconstruct_opts):
         elif self.opts.project_high_method == 'project':
             self.projector_high = Projector(self.q_recovered, self.q_rec_high)
         elif self.opts.project_high_method == 'interpolate':
-            self.projector_high = interpolate(self.q_recovered, target_field.function_space())
+            self.projector_high = Interpolator(self.q_recovered, self.q_rec_high)
             self.interp_high = True
         else:
             raise ValueError(f'Method {self.opts.project_high_method} '
@@ -69,7 +68,7 @@ def __init__(self, source_field, target_field, reconstruct_opts):
         elif self.opts.project_low_method == 'project':
             self.projector_low = Projector(self.q_rec_high, self.q_corr_low)
         elif self.opts.project_low_method == 'interpolate':
-            self.projector_low = interpolate(self.q_rec_high, source_field.function_space())
+            self.projector_low = Interpolator(self.q_rec_high, self.q_corr_low)
             self.interp_low = True
         else:
             raise ValueError(f'Method {self.opts.project_low_method} '
@@ -85,17 +84,17 @@ def __init__(self, source_field, target_field, reconstruct_opts):
         elif self.opts.injection_method == 'project':
             self.injector = Projector(self.q_corr_low, self.q_corr_high)
         elif self.opts.injection_method == 'interpolate':
-            self.injector = interpolate(self.q_corr_low, target_field.function_space())
+            self.injector = Interpolator(self.q_corr_low, self.q_corr_high)
             self.interp_inj = True
         else:
             raise ValueError(f'Method {self.opts.injection_method} for injection not valid')
 
     def project(self):
         self.recoverer.project()
-        self.q_rec_high.assign(assemble(self.projector_high)) if self.interp_high else self.projector_high.project()
-        self.q_corr_low.assign(assemble(self.projector_low)) if self.interp_low else self.projector_low.project()
+        self.projector_high.interpolate() if self.interp_high else self.projector_high.project()
+        self.projector_low.interpolate() if self.interp_low else self.projector_low.project()
         self.q_corr_low.assign(self.q_low - self.q_corr_low)
-        self.q_corr_high.assign(assemble(self.injector)) if self.interp_inj else self.injector.project()
+        self.injector.interpolate() if self.interp_inj else self.injector.project()
         self.q_high.assign(self.q_corr_high + self.q_rec_high)
 
 
diff --git a/gusto/solvers/__init__.py b/gusto/solvers/__init__.py
index fc4c8a7d2..efb4a5af4 100644
--- a/gusto/solvers/__init__.py
+++ b/gusto/solvers/__init__.py
@@ -1,3 +1,2 @@
-from gusto.solvers.parameters import *             # noqa
 from gusto.solvers.linear_solvers import *         # noqa
-from gusto.solvers.preconditioners import *        # noqa
+from gusto.solvers.preconditioners import *        # noqa
\ No newline at end of file
diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index ca684c6af..52141d5dc 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -11,11 +11,10 @@
     rhs, FacetNormal, div, dx, jump, avg, dS, dS_v, dS_h, ds_v, ds_t, ds_b,
     ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, cross,
     BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector,
-    assemble, conditional
+    Interpolator, conditional
 )
 from firedrake.fml import Term, drop
 from firedrake.petsc import flatten_parameters
-from firedrake.__future__ import interpolate
 from pyop2.profiling import timed_function, timed_region
 
 from gusto.equations.active_tracers import TracerVariableType
@@ -85,17 +84,6 @@ def solver_parameters(self):
     def _setup_solver(self):
         pass
 
-    @abstractmethod
-    def update_reference_profiles(self):
-        """
-        Update the solver when the reference profiles have changed.
-
-        This typically includes forcing any Jacobians that depend on
-        the reference profiles to be reassembled, and recalculating
-        any values derived from the reference values.
-        """
-        pass
-
     @abstractmethod
     def solve(self):
         pass
@@ -273,10 +261,8 @@ def L_tr(f):
         rhobar_avg = Function(Vtrace)
         exnerbar_avg = Function(Vtrace)
 
-        rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg,
-                                               constant_jacobian=True)
-        exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg,
-                                                 constant_jacobian=True)
+        rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
+        exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg)
 
         self.rho_avg_solver = LinearVariationalSolver(rho_avg_prb,
                                                       solver_parameters=cg_ilu_parameters,
@@ -332,8 +318,7 @@ def L_tr(f):
         # Function for the hybridized solutions
         self.urhol0 = Function(M)
 
-        hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0,
-                                                  constant_jacobian=True)
+        hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0)
         hybridized_solver = LinearVariationalSolver(hybridized_prb,
                                                     solver_parameters=self.solver_parameters,
                                                     options_prefix='ImplicitSolver')
@@ -359,8 +344,7 @@ def L_tr(f):
         theta_eqn = gamma*(theta - theta_in
                            + dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta_t)*dx
 
-        theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta,
-                                                 constant_jacobian=True)
+        theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta)
         self.theta_solver = LinearVariationalSolver(theta_problem,
                                                     solver_parameters=cg_ilu_parameters,
                                                     options_prefix='thetabacksubstitution')
@@ -377,6 +361,10 @@ def L_tr(f):
 
     @timed_function("Gusto:UpdateReferenceProfiles")
     def update_reference_profiles(self):
+        """
+        Updates the reference profiles.
+        """
+
         with timed_region("Gusto:HybridProjectRhobar"):
             logger.info('Compressible linear solver: rho average solve')
             self.rho_avg_solver.solve()
@@ -385,13 +373,6 @@ def update_reference_profiles(self):
             logger.info('Compressible linear solver: Exner average solve')
             self.exner_avg_solver.solve()
 
-        # Because the left hand side of the hybridised problem depends
-        # on the reference profile, the Jacobian matrix should change
-        # when the reference profiles are updated. This call will tell
-        # the hybridized_solver to reassemble the Jacobian next time
-        # `solve` is called.
-        self.hybridized_solver.invalidate_jacobian()
-
     @timed_function("Gusto:LinearSolve")
     def solve(self, xrhs, dy):
         """
@@ -530,8 +511,7 @@ def V(u):
         bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']]
 
         # Solver for u, p
-        up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs,
-                                              constant_jacobian=True)
+        up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
 
         # Provide callback for the nullspace of the trace system
         def trace_nullsp(T):
@@ -554,17 +534,12 @@ def trace_nullsp(T):
 
         b_problem = LinearVariationalProblem(lhs(b_eqn),
                                              rhs(b_eqn),
-                                             self.b,
-                                             constant_jacobian=True)
+                                             self.b)
         self.b_solver = LinearVariationalSolver(b_problem)
 
         # Log residuals on hybridized solver
         self.log_ksp_residuals(self.up_solver.snes.ksp)
 
-    @timed_function("Gusto:UpdateReferenceProfiles")
-    def update_reference_profiles(self):
-        self.up_solver.invalidate_jacobian()
-
     @timed_function("Gusto:LinearSolve")
     def solve(self, xrhs, dy):
         """
@@ -661,11 +636,11 @@ def _setup_solver(self):
             qtbar = split(equation.X_ref)[3]
 
             # set up interpolators that use the X_ref values for D and b_e
-            self.q_sat_expr_interpolate = interpolate(
-                equation.compute_saturation(equation.X_ref), VD)
-            self.q_v_interpolate = interpolate(
+            self.q_sat_expr_interpolator = Interpolator(
+                equation.compute_saturation(equation.X_ref), self.q_sat_func)
+            self.q_v_interpolator = Interpolator(
                 conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func),
-                VD)
+                self.qvbar)
 
             # bbar was be_bar and here we correct to become bbar
             bbar += equation.parameters.beta2 * self.qvbar
@@ -698,8 +673,7 @@ def _setup_solver(self):
         bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']]
 
         # Solver for u, D
-        uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs,
-                                              constant_jacobian=True)
+        uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs)
 
         # Provide callback for the nullspace of the trace system
         def trace_nullsp(T):
@@ -720,8 +694,7 @@ def trace_nullsp(T):
 
         b_problem = LinearVariationalProblem(lhs(b_eqn),
                                              rhs(b_eqn),
-                                             self.b,
-                                             constant_jacobian=True)
+                                             self.b)
         self.b_solver = LinearVariationalSolver(b_problem)
 
         # Log residuals on hybridized solver
@@ -729,9 +702,13 @@ def trace_nullsp(T):
 
     @timed_function("Gusto:UpdateReferenceProfiles")
     def update_reference_profiles(self):
+        """
+        Updates the reference profiles.
+        """
+
         if self.equations.equivalent_buoyancy:
-            self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate))
-            self.qvbar.assign(assemble(self.q_v_interpolate))
+            self.q_sat_expr_interpolator.interpolate()
+            self.q_v_interpolator.interpolate()
 
     @timed_function("Gusto:LinearSolve")
     def solve(self, xrhs, dy):
@@ -795,17 +772,13 @@ class LinearTimesteppingSolver(object):
                                         'sub_pc_type': 'ilu'}}
     }
 
-    def __init__(self, equation, alpha, reference_dependent=True):
+    def __init__(self, equation, alpha):
         """
         Args:
             equation (:class:`PrognosticEquation`): the model's equation object.
             alpha (float): the semi-implicit off-centring factor. A value of 1
                 is fully-implicit.
-            reference_dependent: this indicates that the solver Jacobian should
-                be rebuilt if the reference profiles have been updated.
         """
-        self.reference_dependent = reference_dependent
-
         residual = equation.residual.label_map(
             lambda t: t.has_label(linearisation),
             lambda t: Term(t.get(linearisation).form, t.labels),
@@ -832,18 +805,12 @@ def __init__(self, equation, alpha, reference_dependent=True):
         bcs = [DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain) for bc in equation.bcs['u']]
         problem = LinearVariationalProblem(aeqn.form,
                                            action(Leqn.form, self.xrhs),
-                                           self.dy, bcs=bcs,
-                                           constant_jacobian=True)
+                                           self.dy, bcs=bcs)
 
         self.solver = LinearVariationalSolver(problem,
                                               solver_parameters=self.solver_parameters,
                                               options_prefix='linear_solver')
 
-    @timed_function("Gusto:UpdateReferenceProfiles")
-    def update_reference_profiles(self):
-        if self.reference_dependent:
-            self.solver.invalidate_jacobian()
-
     @timed_function("Gusto:LinearSolve")
     def solve(self, xrhs, dy):
         """
@@ -928,8 +895,7 @@ def _setup_solver(self):
         bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']]
 
         # Solver for u, D
-        uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs,
-                                              constant_jacobian=True)
+        uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs)
 
         # Provide callback for the nullspace of the trace system
         def trace_nullsp(T):
@@ -943,10 +909,6 @@ def trace_nullsp(T):
         # Log residuals on hybridized solver
         self.log_ksp_residuals(self.uD_solver.snes.ksp)
 
-    @timed_function("Gusto:UpdateReferenceProfiles")
-    def update_reference_profiles(self):
-        self.uD_solver.invalidate_jacobian()
-
     @timed_function("Gusto:LinearSolve")
     def solve(self, xrhs, dy):
         """
diff --git a/gusto/solvers/parameters.py b/gusto/solvers/parameters.py
deleted file mode 100644
index a7f61edb0..000000000
--- a/gusto/solvers/parameters.py
+++ /dev/null
@@ -1,99 +0,0 @@
-"""
-This module provides some parameters sets that are good defaults
-for particular kinds of system.
-"""
-from gusto.core.function_spaces import is_cg
-
-__all__ = ['mass_parameters']
-
-
-def mass_parameters(V, spaces=None, ignore_vertical=True):
-    """
-    PETSc solver parameters for mass matrices.
-
-    Currently this sets to a monolithic CG+ILU.
-
-    TODO: implement field-by-field parameters that choose
-          preonly for discontinuous fields and CG for continuous
-          fields - see docstring below.
-
-    ================= FUTURE DOCSTRING =================
-    Any fields which are discontinuous will have block diagonal
-    mass matrices, so are solved directly using:
-        'ksp_type': 'preonly'
-        'pc_type': 'ilu'
-
-    All continuous fields are solved with CG, with the preconditioner
-    being ILU independently on each field. By solving all continuous fields
-    "monolithically", the total number of inner products is minimised, which
-    is beneficial for scaling to large core counts because it minimises the
-    total number of MPI_Allreduce calls.
-        'ksp_type': 'cg'
-        'pc_type': 'fieldsplit'
-        'pc_fieldsplit_type': 'additive'
-        'fieldsplit_ksp_type': 'preonly'
-        'fieldsplit_pc_type': 'ilu'
-
-    Args:
-        spaces: Optional `Spaces` object. If present, any subspace
-            of V that came from the `Spaces` object will use the
-            continuity information from `spaces`.
-            If not present, continuity is checked with `is_cg`.
-
-        ignore_vertical: whether to include the vertical direction when checking
-            field continuity on extruded meshes. If True, only the horizontal
-            continuity will be considered, e.g. the standard theta space will
-            be treated as discontinuous.
-    """
-    return {
-        'ksp_type': 'cg',
-        'pc_type': 'bjacobi',
-        'sub_pc_type': 'ilu'
-    }
-
-    extruded = hasattr(V.mesh, "_base_mesh")
-
-    continuous_fields = set()
-    for i, Vsub in enumerate(V.subfunctions):
-        # field = Vsub.name or str(i)
-        field = str(i)
-
-        if spaces is not None:
-            continuous = spaces.continuity.get(field, is_cg(Vsub))
-        else:
-            continuous = is_cg(Vsub)
-
-        # For extruded meshes the continuity is recorded
-        # separately for the horizontal and vertical directions.
-        if extruded and spaces is not None:
-            if ignore_vertical:
-                continuous = continuous['horizontal']
-            else:
-                continuous = (continuous['horizontal']
-                              or continuous['vertical'])
-
-        if continuous:
-            continuous_fields.add(field)
-
-    if len(V.subfunctions) == 1:
-        parameters = {
-            'ksp_type': 'cg' if all(continuous_fields) else 'preonly',
-            'pc_type': 'bjacobi',
-            'sub_pc_type': 'ilu',
-        }
-    else:
-
-        parameters = {
-            'ksp_type': 'preonly',
-            'pc_type': 'fieldsplit',
-            'pc_fieldsplit_type': 'additive',
-            'pc_fieldsplit_0_fields': ','.join(continuous_fields),
-            'fieldsplit': {
-                'ksp_type': 'preonly',
-                'pc_type': 'bjacobi',
-                'sub_pc_type': 'ilu'
-            },
-            'fieldsplit_0_ksp_type': 'cg',
-        }
-
-    return parameters
diff --git a/gusto/spatial_methods/__init__.py b/gusto/spatial_methods/__init__.py
index 349094cdf..baa1e0788 100644
--- a/gusto/spatial_methods/__init__.py
+++ b/gusto/spatial_methods/__init__.py
@@ -2,4 +2,3 @@
 from gusto.spatial_methods.diffusion_methods import *         # noqa
 from gusto.spatial_methods.transport_methods import *         # noqa
 from gusto.spatial_methods.limiters import *                  # noqa
-from gusto.spatial_methods.augmentation import *              # noqa
\ No newline at end of file
diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
deleted file mode 100644
index 42aedaa64..000000000
--- a/gusto/spatial_methods/augmentation.py
+++ /dev/null
@@ -1,238 +0,0 @@
-"""
-A module defining objects for temporarily augmenting an equation with another.
-"""
-
-
-from abc import ABCMeta, abstractmethod
-from firedrake import (
-    MixedFunctionSpace, Function, TestFunctions, split, inner, dx, grad,
-    LinearVariationalProblem, LinearVariationalSolver, lhs, rhs, dot,
-    ds_b, ds_v, ds_t, ds, FacetNormal, TestFunction, TrialFunction,
-    transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div,
-    Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC
-)
-from firedrake.fml import subject
-from gusto import (
-    time_derivative, transport, transporting_velocity, TransportEquationType,
-    logger
-)
-
-
-class Augmentation(object, metaclass=ABCMeta):
-    """
-    Augments an equation with another equation to be solved simultaneously.
-    """
-
-    @abstractmethod
-    def pre_apply(self, x_in):
-        """
-        Steps to take at the beginning of an apply method, for instance to
-        assign the input field to the internal mixed function.
-        """
-
-        pass
-
-    @abstractmethod
-    def post_apply(self, x_out):
-        """
-        Steps to take at the end of an apply method, for instance to assign the
-        internal mixed function to the output field.
-        """
-
-        pass
-
-    @abstractmethod
-    def update(self, x_in_mixed):
-        """
-        Any intermediate update steps, depending on the current mixed function.
-        """
-
-        pass
-
-
-class VorticityTransport(Augmentation):
-    """
-    Solves the transport of a vector field, simultaneously with the vorticity
-    as a mixed proble, as described in Bendall and Wimmer (2022).
-
-    Note that this is most effective with implicit time discretisations. The
-    residual-based SUPG option provides a dissipation method.
-
-    Args:
-        domain (:class:`Domain`): The domain object.
-        eqns (:class:`PrognosticEquationSet`): The overarching equation set.
-        transpose_commutator (bool, optional): Whether to include the commutator
-            of the transpose gradient terms. This is necessary for solving the
-            general vector transport equation, but is not necessary when the
-            transporting and transported fields are the same. Defaults to True.
-        supg (bool, optional): Whether to include dissipation through a
-            residual-based SUPG scheme. Defaults to False.
-    """
-
-    def __init__(
-            self, domain, eqns, transpose_commutator=True, supg=False
-    ):
-
-        V_vel = domain.spaces('HDiv')
-        V_vort = domain.spaces('H1')
-
-        self.fs = MixedFunctionSpace((V_vel, V_vort))
-        self.X = Function(self.fs)
-        self.tests = TestFunctions(self.fs)
-
-        u = Function(V_vel)
-        F, Z = split(self.X)
-        test_F, test_Z = self.tests
-
-        if hasattr(domain.mesh, "_base_mesh"):
-            self.ds = ds_b + ds_t + ds_v
-            self.dS = dS_v + dS_h
-        else:
-            self.ds = ds
-            self.dS = dS
-
-        # Add boundary conditions
-        self.bcs = []
-        if 'u' in eqns.bcs.keys():
-            for bc in eqns.bcs['u']:
-                self.bcs.append(
-                    DirichletBC(self.fs.sub(0), bc.function_arg, bc.sub_domain)
-                )
-
-        # Set up test function and the vorticity term
-        n = FacetNormal(domain.mesh)
-        sign_u = 0.5*(sign(dot(u, n)) + 1)
-        upw = lambda f: (sign_u('+')*f('+') + sign_u('-')*f('-'))
-
-        if domain.mesh.topological_dimension() == 2:
-            mix_test = test_F - domain.perp(grad(test_Z))
-            F_cross_u = Z*domain.perp(u)
-        elif domain.mesh.topological_dimension == 3:
-            mix_test = test_F - curl(test_Z)
-            F_cross_u = cross(Z, u)
-
-        time_deriv_form = inner(F, test_F)*dx + inner(Z, test_Z)*dx
-
-        # Standard vector invariant transport form -----------------------------
-        transport_form = (
-            # vorticity term
-            inner(mix_test, F_cross_u)*dx
-            + inner(n, test_Z*Z*u)*self.ds
-            # 0.5*grad(v . F)
-            - 0.5 * div(mix_test) * inner(u, F)*dx
-            + 0.5 * inner(mix_test, n) * inner(u, F)*self.ds
-        )
-
-        # Communtator of tranpose gradient terms -------------------------------
-        # This is needed for general vector transport
-        if transpose_commutator:
-            u_dot_nabla_F = dot(u, transpose(nabla_grad(F)))
-            transport_form += (
-                - inner(n, test_Z*domain.perp(u_dot_nabla_F))*self.ds
-                # + 0.5*grad(F).v
-                - 0.5 * dot(F, div(outer(u, mix_test)))*dx
-                + 0.5 * inner(mix_test('+'), n('+'))*dot(jump(u), upw(F))*self.dS
-                # - 0.5*grad(v).F
-                + 0.5 * dot(u, div(outer(F, mix_test)))*dx
-                - 0.5 * inner(mix_test('+'), n('+'))*dot(jump(F), upw(u))*self.dS
-            )
-
-        # SUPG terms -----------------------------------------------------------
-        # Add the vorticity residual to the transported vorticity,
-        # which damps enstrophy
-        if supg:
-
-            # Determine SUPG coefficient ---------------------------------------
-            tau = 0.5*domain.dt
-
-            # Find mean grid spacing to determine a Courant number
-            DG0 = FunctionSpace(domain.mesh, 'DG', 0)
-            ones = Function(DG0).interpolate(Constant(1.0))
-            area = assemble(ones*dx)
-            mean_dx = (area/DG0.dof_count)**(1/domain.mesh.geometric_dimension())
-
-            # Divide by approximately (1 + c)
-            tau /= (1.0 + sqrt(dot(u, u))*domain.dt/Constant(mean_dx))
-
-            dxqp = dx(degree=3)
-
-            if domain.mesh.topological_dimension() == 2:
-                time_deriv_form -= inner(mix_test, tau*Z*domain.perp(u)/domain.dt)*dxqp
-                transport_form -= inner(
-                    mix_test, tau*domain.perp(u)*domain.divperp(Z*domain.perp(u))
-                )*dxqp
-                if transpose_commutator:
-                    transport_form -= inner(
-                        mix_test,
-                        tau*domain.perp(u)*domain.divperp(u_dot_nabla_F)
-                    )*dxqp
-            elif domain.mesh.topological_dimension() == 3:
-                time_deriv_form -= inner(mix_test, tau*cross(Z, u)/domain.dt)*dxqp
-                transport_form -= inner(
-                    mix_test, tau*cross(curl(Z*u), u)
-                )*dxqp
-                if transpose_commutator:
-                    transport_form -= inner(
-                        mix_test,
-                        tau*cross(curl(u_dot_nabla_F), u)
-                    )*dxqp
-
-        # Assemble the residual ------------------------------------------------
-        residual = (
-            time_derivative(time_deriv_form)
-            + transport(
-                transport_form, TransportEquationType.vector_invariant
-            )
-        )
-        residual = transporting_velocity(residual, u)
-
-        self.residual = subject(residual, self.X)
-
-        self.x_in = Function(self.fs)
-        self.Z_in = Function(V_vort)
-        self.x_out = Function(self.fs)
-
-        vort_test = TestFunction(V_vort)
-        vort_trial = TrialFunction(V_vort)
-
-        F_in, _ = split(self.x_in)
-
-        eqn = (
-            inner(vort_trial, vort_test)*dx
-            + inner(domain.perp(grad(vort_test)), F_in)*dx
-            + vort_test*inner(n, domain.perp(F_in))*self.ds
-        )
-        problem = LinearVariationalProblem(
-            lhs(eqn), rhs(eqn), self.Z_in, constant_jacobian=True
-        )
-        self.solver = LinearVariationalSolver(problem)
-
-    def pre_apply(self, x_in):
-        """
-        Sets the velocity field for the local mixed function.
-
-        Args:
-            x_in (:class:`Function`): The input velocity field
-        """
-        self.x_in.subfunctions[0].assign(x_in)
-
-    def post_apply(self, x_out):
-        """
-        Sets the output velocity field from the local mixed function.
-
-        Args:
-            x_out (:class:`Function`): the output velocity field.
-        """
-        x_out.assign(self.x_out.subfunctions[0])
-
-    def update(self, x_in_mixed):
-        """
-        Performs the solve to determine the vorticity function.
-
-        Args:
-            x_in_mixed (:class:`Function`): The mixed function to update.
-        """
-        self.x_in.subfunctions[0].assign(x_in_mixed.subfunctions[0])
-        logger.debug('Vorticity solve')
-        self.solver.solve()
-        self.x_in.subfunctions[1].assign(self.Z_in)
diff --git a/gusto/spatial_methods/diffusion_methods.py b/gusto/spatial_methods/diffusion_methods.py
index c36fe1be3..a79df20d8 100644
--- a/gusto/spatial_methods/diffusion_methods.py
+++ b/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)
diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py
index 16732bd50..0055ddb89 100644
--- a/gusto/time_discretisation/explicit_runge_kutta.py
+++ b/gusto/time_discretisation/explicit_runge_kutta.py
@@ -15,8 +15,8 @@
 
 
 __all__ = [
-    "ForwardEuler", "ExplicitRungeKutta", "SSPRK2", "SSPRK3", "SSPRK4",
-    "RK4", "Heun", "RungeKuttaFormulation"
+    "ForwardEuler", "ExplicitRungeKutta", "SSPRK3", "RK4", "Heun",
+    "RungeKuttaFormulation"
 ]
 
 
@@ -89,8 +89,7 @@ class ExplicitRungeKutta(ExplicitTimeDiscretisation):
     def __init__(self, domain, butcher_matrix, field_name=None,
                  subcycling_options=None,
                  rk_formulation=RungeKuttaFormulation.increment,
-                 solver_parameters=None, limiter=None, options=None,
-                 augmentation=None):
+                 solver_parameters=None, limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -113,15 +112,11 @@ def __init__(self, domain, butcher_matrix, field_name=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
         """
         super().__init__(domain, field_name=field_name,
                          subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         limiter=limiter, options=options)
         self.butcher_matrix = butcher_matrix
         self.nbutcher = int(np.shape(self.butcher_matrix)[0])
         self.rk_formulation = rk_formulation
@@ -208,7 +203,7 @@ def lhs(self):
         if self.rk_formulation == RungeKuttaFormulation.increment:
             l = self.residual.label_map(
                 lambda t: t.has_label(time_derivative),
-                map_if_true=replace_subject(self.x_out, old_idx=self.idx),
+                map_if_true=replace_subject(self.x_out, self.idx),
                 map_if_false=drop)
 
             return l.form
@@ -218,7 +213,7 @@ def lhs(self):
             for stage in range(self.nStages):
                 l = self.residual.label_map(
                     lambda t: t.has_label(time_derivative),
-                    map_if_true=replace_subject(self.field_i[stage+1], old_idx=self.idx),
+                    map_if_true=replace_subject(self.field_i[stage+1], self.idx),
                     map_if_false=drop)
                 lhs_list.append(l)
 
@@ -227,7 +222,7 @@ def lhs(self):
         if self.rk_formulation == RungeKuttaFormulation.linear:
             l = self.residual.label_map(
                 lambda t: t.has_label(time_derivative),
-                map_if_true=replace_subject(self.x1, old_idx=self.idx),
+                map_if_true=replace_subject(self.x1, self.idx),
                 map_if_false=drop)
 
             return l.form
@@ -466,9 +461,6 @@ def apply_cycle(self, x_out, x_in):
             x_out (:class:`Function`): the output field to be computed.
         """
 
-        if self.augmentation is not None:
-            self.augmentation.update(x_in)
-
         # TODO: is this limiter application necessary?
         if self.limiter is not None:
             self.limiter.apply(x_in)
@@ -492,8 +484,7 @@ class ForwardEuler(ExplicitRungeKutta):
     def __init__(
             self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
-            solver_parameters=None, limiter=None, options=None,
-            augmentation=None
+            solver_parameters=None, limiter=None, options=None
     ):
         """
         Args:
@@ -515,9 +506,6 @@ def __init__(
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
         """
 
         butcher_matrix = np.array([1.]).reshape(1, 1)
@@ -526,136 +514,23 @@ def __init__(
                          subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
-
-
-class SSPRK2(ExplicitRungeKutta):
-    u"""
-    Implements 2nd order Strong-Stability-Preserving Runge-Kutta methods
-    for solving ∂y/∂t = F(y).                                                 \n
-
-    Spiteri & Ruuth, 2002, SIAM J. Numer. Anal.                               \n
-    "A new class of optimal high-order strong-stability-preserving time       \n
-    discretisation methods".                                                  \n
-
-    The 2-stage method (Heun's method) can be written as:                     \n
-
-    k0 = F[y^n]                                                               \n
-    k1 = F{y^n + dt*k0]                                                       \n
-    y^(n+1) = y^n + (1/2)*dt*(k0+k1)                                          \n
-
-    The 3-stage method can be written as:                                     \n
-
-    k0 = F[y^n]                                                               \n
-    k1 = F[y^n + (1/2*dt*k0]                                                  \n
-    k2 = F[y^n + (1/2)*dt*(k0+k1)]                                            \n
-    y^(n+1) = y^n + (1/3)*dt*(k0 + k1 + k2)                                   \n
-
-    The 4-stage method can be written as:                                     \n
-
-    k0 = F[y^n]                                                               \n
-    k1 = F[y^n + (1/3)*dt*k1]                                                 \n
-    k2 = F[y^n + (1/3)*dt*(k0+k1)]                                            \n
-    k3 = F[y^n + (1/3)*dt*(k0+k1+k2)]                                         \n
-    y^(n+1) = y^n + (1/4)*dt*(k0 + k1 + k2 + k3)                              \n
-    """
-    def __init__(
-            self, domain, field_name=None, subcycling_options=None,
-            rk_formulation=RungeKuttaFormulation.increment,
-            solver_parameters=None, limiter=None, options=None,
-            augmentation=None, stages=2
-    ):
-        """
-        Args:
-            domain (:class:`Domain`): the model's domain object, containing the
-                mesh and the compatible function spaces.
-            field_name (str, optional): name of the field to be evolved.
-                Defaults to None.
-            subcycling_options(:class:`SubcyclingOptions`, optional): an object
-                containing options for subcycling the time discretisation.
-                Defaults to None.
-            rk_formulation (:class:`RungeKuttaFormulation`, optional):
-                an enumerator object, describing the formulation of the Runge-
-                Kutta scheme. Defaults to the increment form.
-            solver_parameters (dict, optional): dictionary of parameters to
-                pass to the underlying solver. Defaults to None.
-            limiter (:class:`Limiter` object, optional): a limiter to apply to
-                the evolving field to enforce monotonicity. Defaults to None.
-            options (:class:`AdvectionOptions`, optional): an object containing
-                options to either be passed to the spatial discretisation, or
-                to control the "wrapper" methods, such as Embedded DG or a
-                recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
-            stages (int, optional): number of stages: (2, 3, 4). Defaults to 2.
-        """
-
-        if stages == 2:
-            butcher_matrix = np.array([
-                [1., 0.],
-                [0.5, 0.5]
-            ])
-            self.cfl_limit = 1
-
-        elif stages == 3:
-            butcher_matrix = np.array([
-                [1./2., 0., 0.],
-                [1./2., 1./2., 0.],
-                [1./3., 1./3., 1./3.]
-            ])
-            self.cfl_limit = 2
-        elif stages == 4:
-            butcher_matrix = np.array([
-                [1./3., 0., 0., 0.],
-                [1./3., 1./3., 0., 0.],
-                [1./3., 1./3., 1./3., 0.],
-                [1./4., 1./4., 1./4., 1./4.]
-            ])
-            self.cfl_limit = 3
-        else:
-            raise ValueError(f"{stages} stage 2rd order SSPRK not implemented")
-
-        super().__init__(domain, butcher_matrix, field_name=field_name,
-                         subcycling_options=subcycling_options,
-                         rk_formulation=rk_formulation,
-                         solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         limiter=limiter, options=options)
 
 
 class SSPRK3(ExplicitRungeKutta):
     u"""
-    Implements 3rd order Strong-Stability-Preserving Runge-Kutta methods
-    for solving ∂y/∂t = F(y).                                                 \n
-
-    Spiteri & Ruuth, 2002, SIAM J. Numer. Anal.                               \n
-    "A new class of optimal high-order strong-stability-preserving time       \n
-    discretisation methods".                                                  \n
-
-    The 3-stage method can be written as:                                     \n
+    Implements the 3-stage Strong-Stability-Preserving Runge-Kutta method
+    for solving ∂y/∂t = F(y). It can be written as:                           \n
 
     k0 = F[y^n]                                                               \n
     k1 = F[y^n + dt*k1]                                                       \n
     k2 = F[y^n + (1/4)*dt*(k0+k1)]                                            \n
     y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2)                                 \n
-
-    The 4-stage method can be written as:                                     \n
-
-    k0 = F[y^n]                                                               \n
-    k1 = F[y^n + (1/2)*dt*k1]                                                 \n
-    k2 = F[y^n + (1/2)*dt*(k0+k1)]                                            \n
-    k3 = F[y^n + (1/6)*dt*(k0+k1+k2)]                                         \n
-    y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + k2 + 3*k3)                            \n
-
-    The 5-stage method has numerically optimised coefficients.                \n
     """
     def __init__(
             self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
-            solver_parameters=None, limiter=None, options=None,
-            augmentation=None, stages=3
+            solver_parameters=None, limiter=None, options=None
     ):
         """
         Args:
@@ -677,63 +552,39 @@ def __init__(
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
-            stages (int, optional): number of stages: (3, 4, 5). Defaults to 3.
         """
 
-        if stages == 3:
-            butcher_matrix = np.array([
-                [1., 0., 0.],
-                [1./4., 1./4., 0.],
-                [1./6., 1./6., 2./3.]
-            ])
-            self.cfl_limit = 1
-        elif stages == 4:
-            butcher_matrix = np.array([
-                [1./2., 0., 0., 0.],
-                [1./2., 1./2., 0., 0.],
-                [1./6., 1./6., 1./6., 0.],
-                [1./6., 1./6., 1./6., 1./2.]
-            ])
-            self.cfl_limit = 2
-        elif stages == 5:
-            self.cfl_limit = 2.65062919294483
-            butcher_matrix = np.array([
-                [0.37726891511710, 0., 0., 0., 0.],
-                [0.37726891511710, 0.37726891511710, 0., 0., 0.],
-                [0.16352294089771, 0.16352294089771, 0.16352294089771, 0., 0.],
-                [0.14904059394856, 0.14831273384724, 0.14831273384724, 0.34217696850008, 0.],
-                [0.19707596384481, 0.11780316509765, 0.11709725193772, 0.27015874934251, 0.29786487010104]
-            ])
-        else:
-            raise ValueError(f"{stages} stage 3rd order SSPRK not implemented")
-
+        butcher_matrix = np.array([
+            [1., 0., 0.],
+            [1./4., 1./4., 0.],
+            [1./6., 1./6., 2./3.]
+        ])
         super().__init__(domain, butcher_matrix, field_name=field_name,
                          subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         limiter=limiter, options=options)
 
 
-class SSPRK4(ExplicitRungeKutta):
+class RK4(ExplicitRungeKutta):
     u"""
-    Implements 4th order Strong-Stability-Preserving Runge-Kutta methods
-    for solving ∂y/∂t = F(y).                                                 \n
+    Implements the classic 4-stage Runge-Kutta method.
 
-    Spiteri & Ruuth, 2002, SIAM J. Numer. Anal.                               \n
-    "A new class of optimal high-order strong-stability-preserving time       \n
-    discretisation methods".                                                  \n
+    The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be
+    written as:                                                               \n
+
+    k0 = F[y^n]                                                               \n
+    k1 = F[y^n + 1/2*dt*k1]                                                   \n
+    k2 = F[y^n + 1/2*dt*k2]                                                   \n
+    k3 = F[y^n + dt*k3]                                                       \n
+    y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3)                      \n
 
-    The 5-stage method has numerically optimised coefficients.                \n
+    where superscripts indicate the time-level.                               \n
     """
     def __init__(
             self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
-            solver_parameters=None, limiter=None, options=None,
-            augmentation=None, stages=5
+            solver_parameters=None, limiter=None, options=None
     ):
         """
         Args:
@@ -755,52 +606,37 @@ def __init__(
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
-            stages (int, optional): number of stages: (5,). Defaults to 5.
         """
-
-        if stages == 5:
-            self.cfl_limit = 1.50818004975927
-            butcher_matrix = np.array([
-                [0.39175222700392, 0., 0., 0., 0.],
-                [0.21766909633821, 0.36841059262959, 0., 0., 0.],
-                [0.08269208670950, 0.13995850206999, 0.25189177424738, 0., 0.],
-                [0.06796628370320, 0.11503469844438, 0.20703489864929, 0.54497475021237, 0.],
-                [0.14681187618661, 0.24848290924556, 0.10425883036650, 0.27443890091960, 0.22600748319395]
-            ])
-        else:
-            raise ValueError(f"{stages} stage 4rd order SSPRK not implemented")
-
+        butcher_matrix = np.array([
+            [0.5, 0., 0., 0.],
+            [0., 0.5, 0., 0.],
+            [0., 0., 1., 0.],
+            [1./6., 1./3., 1./3., 1./6.]
+        ])
         super().__init__(domain, butcher_matrix, field_name=field_name,
                          subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         limiter=limiter, options=options)
 
 
-class RK4(ExplicitRungeKutta):
+class Heun(ExplicitRungeKutta):
     u"""
-    Implements the classic 4-stage Runge-Kutta method.
+    Implements Heun's method.
 
-    The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be
-    written as:                                                               \n
+    The 2-stage Runge-Kutta scheme known as Heun's method,for solving
+    ∂y/∂t = F(y). It can be written as:                                       \n
 
-    k0 = F[y^n]                                                               \n
-    k1 = F[y^n + 1/2*dt*k1]                                                   \n
-    k2 = F[y^n + 1/2*dt*k2]                                                   \n
-    k3 = F[y^n + dt*k3]                                                       \n
-    y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3)                      \n
+    y_1 = F[y^n]                                                              \n
+    y^(n+1) = (1/2)y^n + (1/2)F[y_1]                                          \n
 
-    where superscripts indicate the time-level.                               \n
+    where superscripts indicate the time-level and subscripts indicate the stage
+    number.
     """
     def __init__(
             self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
-            solver_parameters=None, limiter=None, options=None,
-            augmentation=None
+            solver_parameters=None, limiter=None, options=None
     ):
         """
         Args:
@@ -822,40 +658,14 @@ def __init__(
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
         """
+
         butcher_matrix = np.array([
-            [0.5, 0., 0., 0.],
-            [0., 0.5, 0., 0.],
-            [0., 0., 1., 0.],
-            [1./6., 1./3., 1./3., 1./6.]
+            [1., 0.],
+            [0.5, 0.5]
         ])
         super().__init__(domain, butcher_matrix, field_name=field_name,
                          subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
-
-
-def Heun(domain, field_name=None, subcycling_options=None,
-         rk_formulation=RungeKuttaFormulation.increment,
-         solver_parameters=None, limiter=None, options=None,
-         augmentation=None):
-    u"""
-    Implements Heun's method.
-
-    The 2-stage Runge-Kutta scheme known as Heun's method,for solving
-    ∂y/∂t = F(y). It can be written as:                                       \n
-
-    y_1 = F[y^n]                                                              \n
-    y^(n+1) = (1/2)y^n + (1/2)F[y_1]                                          \n
-
-    where superscripts indicate the time-level and subscripts indicate the stage
-    number.
-    """
-    return SSPRK2(domain, field_name=field_name, subcycling_options=subcycling_options,
-                  rk_formulation=rk_formulation, solver_parameters=solver_parameters,
-                  limiter=limiter, options=options, augmentation=augmentation, stages=2)
+                         limiter=limiter, options=options)
diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py
index b220ef0be..8eb1c75ab 100644
--- a/gusto/time_discretisation/implicit_runge_kutta.py
+++ b/gusto/time_discretisation/implicit_runge_kutta.py
@@ -56,7 +56,7 @@ class ImplicitRungeKutta(TimeDiscretisation):
     # ---------------------------------------------------------------------------
 
     def __init__(self, domain, butcher_matrix, field_name=None,
-                 solver_parameters=None, options=None, augmentation=None):
+                 solver_parameters=None, options=None,):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -75,7 +75,7 @@ def __init__(self, domain, butcher_matrix, field_name=None,
         """
         super().__init__(domain, field_name=field_name,
                          solver_parameters=solver_parameters,
-                         options=options, augmentation=augmentation)
+                         options=options)
         self.butcher_matrix = butcher_matrix
         self.nStages = int(np.shape(self.butcher_matrix)[1])
 
@@ -165,7 +165,7 @@ class ImplicitMidpoint(ImplicitRungeKutta):
     y^(n+1) = y^n + dt*k0                                                     \n
     """
     def __init__(self, domain, field_name=None, solver_parameters=None,
-                 options=None, augmentation=None):
+                 options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -182,7 +182,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         butcher_matrix = np.array([[0.5], [1.]])
         super().__init__(domain, butcher_matrix, field_name,
                          solver_parameters=solver_parameters,
-                         options=options, augmentation=augmentation)
+                         options=options)
 
 
 class QinZhang(ImplicitRungeKutta):
diff --git a/gusto/time_discretisation/sdc.py b/gusto/time_discretisation/sdc.py
index 234f0c7b9..0fe0c9f29 100644
--- a/gusto/time_discretisation/sdc.py
+++ b/gusto/time_discretisation/sdc.py
@@ -46,12 +46,14 @@
 from abc import ABCMeta
 import numpy as np
 from firedrake import (
-    Function, NonlinearVariationalProblem, NonlinearVariationalSolver, Constant
+    Function, NonlinearVariationalProblem, TestFunction, TestFunctions,
+    NonlinearVariationalSolver, Constant
 )
 from firedrake.fml import (
-    replace_subject, all_terms, drop
+    replace_subject, replace_test_function, all_terms, drop
 )
 from firedrake.utils import cached_property
+from gusto.time_discretisation.wrappers import *
 from gusto.time_discretisation.time_discretisation import wrapper_apply
 from gusto.core.labels import (time_derivative, implicit, explicit)
 
@@ -116,8 +118,36 @@ def __init__(self, base_scheme, domain, M, maxk, quad_type, node_type, qdelta_im
         self.final_update = final_update
         self.formulation = formulation
         self.limiter = limiter
-        self.augmentation = self.base.augmentation
-        self.wrapper = self.base.wrapper
+
+        # Initialise wrappers
+        if options is not None:
+            self.wrapper_name = options.name
+            if self.wrapper_name == "mixed_options":
+                self.wrapper = MixedFSWrapper()
+
+                for field, suboption in options.suboptions.items():
+                    if suboption.name == 'embedded_dg':
+                        self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)})
+                    elif suboption.name == "recovered":
+                        self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)})
+                    elif suboption.name == "supg":
+                        raise RuntimeError(
+                            'SDC: suboption SUPG is currently not implemented within MixedOptions')
+                    else:
+                        raise RuntimeError(
+                            f'SDC: suboption wrapper {wrapper_name} not implemented')
+            elif self.wrapper_name == "embedded_dg":
+                self.wrapper = EmbeddedDGWrapper(self, options)
+            elif self.wrapper_name == "recovered":
+                self.wrapper = RecoveryWrapper(self, options)
+            elif self.wrapper_name == "supg":
+                self.wrapper = SUPGWrapper(self, options)
+            else:
+                raise RuntimeError(
+                    f'SDC: wrapper {self.wrapper_name} not implemented')
+        else:
+            self.wrapper = None
+            self.wrapper_name = None
 
         # Get quadrature nodes and weights
         self.nodes, self.weights, self.Q = genQCoeffs("Collocation", nNodes=M,
@@ -130,10 +160,6 @@ def __init__(self, base_scheme, domain, M, maxk, quad_type, node_type, qdelta_im
         self.dtau = np.diff(np.append(0, self.nodes))
         self.Q = float(self.dt_coarse)*self.Q
         self.Qfin = float(self.dt_coarse)*self.weights
-        self.qdelta_imp_type = qdelta_imp
-        self.formulation = formulation
-        self.node_type = node_type
-        self.quad_type = quad_type
 
         # Get Q_delta matrices
         self.Qdelta_imp = genQDeltaCoeffs(qdelta_imp, form=formulation,
@@ -187,6 +213,62 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                and (not t.has_label(time_derivative))):
                 raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit")
 
+            # Check we are not using wrappers for implicit schemes
+            if (t.has_label(implicit) and self.wrapper is not None):
+                raise NotImplementedError("Implicit terms not supported with wrappers")
+
+            # Check we are not using limiters for implicit schemes
+            if (t.has_label(implicit) and self.limiter is not None):
+                raise NotImplementedError("Implicit terms not supported with limiters")
+
+        # Set up Wrappers
+        if self.wrapper is not None:
+            if self.wrapper_name == "mixed_options":
+
+                self.wrapper.wrapper_spaces = equation.spaces
+                self.wrapper.field_names = equation.field_names
+
+                for field, subwrapper in self.wrapper.subwrappers.items():
+
+                    if field not in equation.field_names:
+                        raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
+
+                    field_idx = equation.field_names.index(field)
+                    subwrapper.setup(equation.spaces[Wrappersfield_idx])
+
+                    # Update the function space to that needed by the wrapper
+                    self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
+
+                self.wrapper.setup()
+                self.fs = self.wrapper.function_space
+                new_test_mixed = TestFunctions(self.fs)
+
+                # Replace the original test function with one from the new
+                # function space defined by the subwrappers
+                self.residual = self.residual.label_map(
+                    all_terms,
+                    map_if_true=replace_test_function(new_test_mixed))
+
+            else:
+                if self.wrapper_name == "supg":
+                    self.wrapper.setup()
+                else:
+                    self.wrapper.setup(self.fs)
+                self.fs = self.wrapper.function_space
+                if self.solver_parameters is None:
+                    self.solver_parameters = self.wrapper.solver_parameters
+                new_test = TestFunction(self.wrapper.test_space)
+                # SUPG has a special wrapper
+                if self.wrapper_name == "supg":
+                    new_test = self.wrapper.test
+
+                # Replace the original test function with the one from the wrapper
+                self.residual = self.residual.label_map(
+                    all_terms,
+                    map_if_true=replace_test_function(new_test))
+
+                self.residual = self.wrapper.label_terms(self.residual)
+
         # Set up bcs
         self.bcs = self.base.bcs
 
@@ -398,19 +480,6 @@ def apply(self, x_out, x_in):
         while k < self.maxk:
             k += 1
 
-            if self.qdelta_imp_type == "MIN-SR-FLEX":
-                # Recompute Implicit Q_delta matrix for each iteration k
-                self.Qdelta_imp = genQDeltaCoeffs(
-                    self.qdelta_imp_type,
-                    form=self.formulation,
-                    nodes=self.nodes,
-                    Q=self.Q,
-                    nNodes=self.M,
-                    nodeType=self.node_type,
-                    quadType=self.quad_type,
-                    k=k
-                )
-
             # Compute for N2N: sum(j=1,M) (s_mj*F(y_m^k) +  s_mj*S(y_m^k))
             # for Z2N: sum(j=1,M) (q_mj*F(y_m^k) +  q_mj*S(y_m^k))
             for m in range(1, self.M+1):
diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py
index d8e07bf86..d168039d2 100644
--- a/gusto/time_discretisation/time_discretisation.py
+++ b/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)
@@ -21,7 +21,6 @@
                                mass_weighted, nonlinear_time_derivative)
 from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 from gusto.time_discretisation.wrappers import *
-from gusto.solvers import mass_parameters
 
 __all__ = ["TimeDiscretisation", "ExplicitTimeDiscretisation", "BackwardEuler",
            "ThetaMethod", "TrapeziumRule", "TR_BDF2"]
@@ -31,17 +30,7 @@ def wrapper_apply(original_apply):
     """Decorator to add steps for using a wrapper around the apply method."""
     def get_apply(self, x_out, x_in):
 
-        if self.augmentation is not None:
-
-            def new_apply(self, x_out, x_in):
-
-                self.augmentation.pre_apply(x_in)
-                original_apply(self, self.augmentation.x_out, self.augmentation.x_in)
-                self.augmentation.post_apply(x_out)
-
-            return new_apply(self, x_out, x_in)
-
-        elif self.wrapper is not None:
+        if self.wrapper is not None:
 
             def new_apply(self, x_out, x_in):
 
@@ -62,17 +51,13 @@ class TimeDiscretisation(object, metaclass=ABCMeta):
     """Base class for time discretisation schemes."""
 
     def __init__(self, domain, field_name=None, subcycling_options=None,
-                 solver_parameters=None, limiter=None, options=None,
-                 augmentation=None):
+                 solver_parameters=None, limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            subcycling_options(:class:`SubcyclingOptions`, optional): an object
-                containing options for subcycling the time discretisation.
-                Defaults to None.
             solver_parameters (dict, optional): dictionary of parameters to
                 pass to the underlying solver. Defaults to None.
             limiter (:class:`Limiter` object, optional): a limiter to apply to
@@ -81,23 +66,17 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
         """
         self.domain = domain
         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
-        self.augmentation = augmentation
-        self.subcycling_options = subcycling_options
+        self.subcycling_options = None
 
         if self.subcycling_options is not None:
             self.subcycling_options.check_options()
@@ -178,11 +157,6 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             self.fs = equation.function_space
             self.idx = None
 
-        if self.augmentation is not None:
-            self.fs = self.augmentation.fs
-            self.residual = self.augmentation.residual
-            self.idx = None
-
         if len(active_labels) > 0:
             if isinstance(self.field_name, list):
                 # Multiple fields are being solved for simultaneously.
@@ -203,7 +177,6 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                         map_if_false=drop)
 
                 self.residual = residual
-
             else:
                 self.residual = self.residual.label_map(
                     lambda t: any(t.has_label(time_derivative, *active_labels)),
@@ -213,11 +186,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         if isinstance(self.field_name, list):
             self.field_name = equation.field_name
 
-        if self.augmentation is not None:
-            # Transfer BCs from appropriate function space
-            bcs = self.augmentation.bcs if hasattr(self.augmentation, "bcs") else None
-        else:
-            bcs = equation.bcs[self.field_name]
+        bcs = equation.bcs[self.field_name]
 
         self.evaluate_source = []
         self.physics_names = []
@@ -233,16 +202,8 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             for field in equation.field_names:
 
                 # Check if the mass term for this prognostic is mass-weighted
-                if len(self.residual.label_map((
-                    lambda t: t.get(prognostic) == field
-                    and t.has_label(time_derivative)
-                    and t.has_label(mass_weighted)
-                ), map_if_false=drop)) == 1:
-
-                    field_terms = self.residual.label_map(
-                        lambda t: t.get(prognostic) == field and not t.has_label(time_derivative),
-                        map_if_false=drop
-                    )
+                if len(self.residual.label_map(lambda t: t.get(prognostic) == field and t.has_label(time_derivative) and t.has_label(mass_weighted), map_if_false=drop)) == 1:
+                    field_terms = self.residual.label_map(lambda t: t.get(prognostic) == field and not t.has_label(time_derivative), map_if_false=drop)
 
                     # Check that the equation for this prognostic does not involve
                     # both mass-weighted and non-mass-weighted terms; if so, a split
@@ -444,8 +405,7 @@ class ExplicitTimeDiscretisation(TimeDiscretisation):
     """Base class for explicit time discretisations."""
 
     def __init__(self, domain, field_name=None, subcycling_options=None,
-                 solver_parameters=None, limiter=None, options=None,
-                 augmentation=None):
+                 solver_parameters=None, limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -463,15 +423,20 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
         """
         super().__init__(domain, field_name,
                          subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         limiter=limiter, options=options)
+
+        # get default solver options if none passed in
+        if solver_parameters is None:
+            self.solver_parameters = {'snes_type': 'ksponly',
+                                      'ksp_type': 'cg',
+                                      'pc_type': 'bjacobi',
+                                      'sub_pc_type': 'ilu'}
+        else:
+            self.solver_parameters = solver_parameters
 
     def setup(self, equation, apply_bcs=True, *active_labels):
         """
@@ -486,25 +451,20 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         """
         super().setup(equation, apply_bcs, *active_labels)
 
-        # get default solver options if none passed in
-        self.solver_parameters.update(mass_parameters(
-            self.fs, equation.domain.spaces))
-        self.solver_parameters['snes_type'] = 'ksponly'
-
         # if user has specified a number of fixed subcycles, then save this
         # and rescale dt accordingly; else perform just one cycle using dt
         if (self.subcycling_options is not None
                 and self.subcycling_options.fixed_subcycles is not None):
-            self.ncycles = self.subcycling_options.fixed_subcycles
-            self.dt.assign(self.dt/self.ncycles)
+            fixed_subcycles = self.subcycling_options.fixed_subcycles
+            self.dt.assign(self.dt/fixed_subcycles)
+            self.ncycles = self.fixed_subcycles
         else:
-            self.ncycles = 1
             self.dt = self.dt
+            self.ncycles = 1
         self.x0 = Function(self.fs)
         self.x1 = Function(self.fs)
 
-        # If the time_derivative term is nonlinear, we must use a nonlinear solver,
-        # but if the time_derivative term is linear, we can reuse the factorisations.
+        # If the time_derivative term is nonlinear, we must use a nonlinear solver
         if (
             len(self.residual.label_map(
                 lambda t: t.has_label(nonlinear_time_derivative),
@@ -516,11 +476,6 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                        + ' as the time derivative term is nonlinear')
             logger.warning(message)
             self.solver_parameters['snes_type'] = 'newtonls'
-        else:
-            self.solver_parameters.setdefault('snes_lag_jacobian', -2)
-            self.solver_parameters.setdefault('snes_lag_jacobian_persists', None)
-            self.solver_parameters.setdefault('snes_lag_preconditioner', -2)
-            self.solver_parameters.setdefault('snes_lag_preconditioner_persists', None)
 
     @cached_property
     def lhs(self):
@@ -579,8 +534,7 @@ class BackwardEuler(TimeDiscretisation):
     y^(n+1) = y^n + dt*F[y^(n+1)].                                               \n
     """
     def __init__(self, domain, field_name=None, subcycling_options=None,
-                 solver_parameters=None, limiter=None, options=None,
-                 augmentation=None):
+                 solver_parameters=None, limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -597,10 +551,7 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
             options (:class:`AdvectionOptions`, optional): an object containing
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
-            """
+        """
         if not solver_parameters:
             # default solver parameters
             solver_parameters = {'ksp_type': 'gmres',
@@ -609,8 +560,7 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
         super().__init__(domain=domain, field_name=field_name,
                          subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         limiter=limiter, options=options)
 
     def setup(self, equation, apply_bcs=True, *active_labels):
         """
@@ -706,7 +656,7 @@ class ThetaMethod(TimeDiscretisation):
     """
 
     def __init__(self, domain, theta, field_name=None, subcycling_options=None,
-                 solver_parameters=None, options=None, augmentation=None):
+                 solver_parameters=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -724,9 +674,6 @@ def __init__(self, domain, theta, field_name=None, subcycling_options=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
 
         Raises:
             ValueError: if theta is not provided.
@@ -745,8 +692,7 @@ def __init__(self, domain, theta, field_name=None, subcycling_options=None,
         super().__init__(domain, field_name,
                          subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
-                         options=options,
-                         augmentation=augmentation)
+                         options=options)
 
         self.theta = theta
 
@@ -823,8 +769,6 @@ def apply(self, x_out, x_in):
             x_in (:class:`Function`): the input field.
         """
         self.update_subcycling()
-        if self.augmentation is not None:
-            self.augmentation.update(x_in)
 
         self.x0.assign(x_in)
         for i in range(self.ncycles):
@@ -845,7 +789,7 @@ class TrapeziumRule(ThetaMethod):
     """
 
     def __init__(self, domain, field_name=None, subcycling_options=None,
-                 solver_parameters=None, options=None, augmentation=None):
+                 solver_parameters=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -858,14 +802,11 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
-            augmentation (:class:`Augmentation`): allows the equation solved in
-                this time discretisation to be augmented, for instances with
-                extra terms of another auxiliary variable. Defaults to None.
         """
         super().__init__(domain, 0.5, field_name,
                          subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
-                         options=options, augmentation=augmentation)
+                         options=options)
 
 
 class TR_BDF2(TimeDiscretisation):
diff --git a/gusto/time_discretisation/wrappers.py b/gusto/time_discretisation/wrappers.py
index 18381bc1f..2d388e998 100644
--- a/gusto/time_discretisation/wrappers.py
+++ b/gusto/time_discretisation/wrappers.py
@@ -6,16 +6,14 @@
 
 from abc import ABCMeta, abstractmethod
 from firedrake import (
-    FunctionSpace, Function, BrokenElement, Projector, Constant, as_ufl, dot, grad,
-    TestFunction, MixedFunctionSpace, assemble
+    FunctionSpace, Function, BrokenElement, Projector, Interpolator,
+    VectorElement, Constant, as_ufl, dot, grad, TestFunction, MixedFunctionSpace
 )
-from firedrake.__future__ import interpolate
 from firedrake.fml import Term
 from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions
 from gusto.recovery import Recoverer, ReversibleRecoverer, ConservativeRecoverer
 from gusto.core.labels import transporting_velocity
 from gusto.core.conservative_projection import ConservativeProjector
-from gusto.core.function_spaces import is_cg
 import ufl
 
 __all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedFSWrapper"]
@@ -265,7 +263,7 @@ def setup(self, original_space, post_apply_bcs):
         # Operators for projecting back
         self.interp_back = (self.options.project_low_method == 'interpolate')
         if self.options.project_low_method == 'interpolate':
-            self.x_out_projector = interpolate(self.x_out, self.original_space)
+            self.x_out_projector = Interpolator(self.x_out, self.x_projected)
         elif self.options.project_low_method == 'project':
             self.x_out_projector = Projector(self.x_out, self.x_projected,
                                              bcs=post_apply_bcs)
@@ -303,12 +301,33 @@ def post_apply(self, x_out):
         """
 
         if self.interp_back:
-            self.x_projected.assign(assemble(self.x_out_projector))
+            self.x_out_projector.interpolate()
         else:
             self.x_out_projector.project()
         x_out.assign(self.x_projected)
 
 
+def is_cg(V):
+    """
+    Checks if a :class:`FunctionSpace` is continuous.
+
+    Function to check if a given space, V, is CG. Broken elements are always
+    discontinuous; for vector elements we check the names of the Sobolev spaces
+    of the subelements and for all other elements we just check the Sobolev
+    space name.
+
+    Args:
+        V (:class:`FunctionSpace`): the space to check.
+    """
+    ele = V.ufl_element()
+    if isinstance(ele, BrokenElement):
+        return False
+    elif type(ele) == VectorElement:
+        return all([e.sobolev_space.name == "H1" for e in ele._sub_elements])
+    else:
+        return V.ufl_element().sobolev_space.name == "H1"
+
+
 class SUPGWrapper(Wrapper):
     """
     Wrapper for computing a time discretisation with SUPG, which adjusts the
diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index 19142dbaa..7c4de236b 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -4,17 +4,16 @@
 """
 
 from firedrake import (
-    Function, Constant, TrialFunctions, DirichletBC, div, assemble,
+    Function, Constant, TrialFunctions, DirichletBC, div, Interpolator,
     LinearVariationalProblem, LinearVariationalSolver
 )
 from firedrake.fml import drop, replace_subject
-from firedrake.__future__ import interpolate
 from pyop2.profiling import timed_stage
 from gusto.core import TimeLevelFields, StateFields
 from gusto.core.labels import (transport, diffusion, time_derivative,
                                linearisation, prognostic, hydrostatic,
                                physics_label, sponge, incompressible)
-from gusto.solvers import LinearTimesteppingSolver, mass_parameters
+from gusto.solvers import LinearTimesteppingSolver
 from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation
 from gusto.timestepping.timestepper import BaseTimestepper
@@ -235,8 +234,9 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
             V_DG = equation_set.domain.spaces('DG')
             self.predictor_field_in = Function(V_DG)
             div_factor = Constant(1.0) - (Constant(1.0) - self.alpha)*self.dt*div(self.x.n('u'))
-            self.predictor_interpolate = interpolate(
-                self.x.star(predictor)*div_factor, V_DG)
+            self.predictor_interpolator = Interpolator(
+                self.x.star(predictor)*div_factor, self.predictor_field_in
+            )
 
     def _apply_bcs(self):
         """
@@ -334,7 +334,7 @@ def transport_fields(self, outer, xstar, xp):
                     # Pre-multiply this variable by (1 - dt*beta*div(u))
                     V = xstar(name).function_space()
                     field_out = Function(V)
-                    self.predictor_field_in.assign(assemble(self.predictor_interpolate))
+                    self.predictor_interpolator.interpolate()
                     scheme.apply(field_out, self.predictor_field_in)
 
                     # xp is xstar plus the increment from the transported predictor
@@ -571,17 +571,13 @@ def __init__(self, equation, alpha):
             constant_jacobian=True
         )
 
-        self.solver_parameters = mass_parameters(W, equation.domain.spaces)
-
         self.solvers = {}
         self.solvers["explicit"] = LinearVariationalSolver(
             explicit_forcing_problem,
-            solver_parameters=self.solver_parameters,
             options_prefix="ExplicitForcingSolver"
         )
         self.solvers["implicit"] = LinearVariationalSolver(
             implicit_forcing_problem,
-            solver_parameters=self.solver_parameters,
             options_prefix="ImplicitForcingSolver"
         )
 
diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py
index 3c619753b..01cfb7737 100644
--- a/gusto/timestepping/timestepper.py
+++ b/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"):
diff --git a/integration-tests/adjoints/test_diffusion_sensitivity.py b/integration-tests/adjoints/test_diffusion_sensitivity.py
new file mode 100644
index 000000000..d90b3fadb
--- /dev/null
+++ b/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
diff --git a/integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py b/integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py
new file mode 100644
index 000000000..6936ec9c8
--- /dev/null
+++ b/integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py
@@ -0,0 +1,218 @@
+"""
+The moist thermal form of Test Case 5 (flow over a mountain) of Williamson et
+al, 1992:
+``A standard test set for numerical approximations to the shallow water
+equations in spherical geometry'', JCP.
+
+The initial conditions are taken from Zerroukat & Allen, 2015:
+``A moist Boussinesq shallow water equations set for testing atmospheric
+models'', JCP.
+
+Three moist variables (vapour, cloud liquid and rain) are used. This set of
+equations involves an active buoyancy field.
+
+The example here uses the icosahedral sphere mesh and degree 1 spaces. An
+explicit RK4 timestepper is used.
+"""
+import pytest
+import numpy as np
+
+from firedrake import (
+    SpatialCoordinate, as_vector, pi, sqrt, min_value, exp, cos, sin, assemble, dx, inner, Function
+)
+from firedrake.adjoint import *
+from pyadjoint import get_working_tape
+from gusto import (
+    Domain, IO, OutputParameters, Timestepper, RK4, DGUpwind,
+    ShallowWaterParameters, ThermalShallowWaterEquations, lonlatr_from_xyz,
+    InstantRain, SWSaturationAdjustment, WaterVapour, CloudWater,
+    Rain, GeneralIcosahedralSphereMesh
+)
+
+
+@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()
+
+
+def test_moist_thermal_williamson_5_sensitivity(
+        tmpdir, ncells_per_edge=8, dt=600, tmax=50.*24.*60.*60.,
+        dumpfreq=2880
+):
+    assert get_working_tape()._blocks == []
+    # ------------------------------------------------------------------------ #
+    # Parameters for test case
+    # ------------------------------------------------------------------------ #
+    radius = 6371220.           # planetary radius (m)
+    mean_depth = 5960           # reference depth (m)
+    g = 9.80616                 # acceleration due to gravity (m/s^2)
+    u_max = 20.                 # max amplitude of the zonal wind (m/s)
+    epsilon = 1/300             # linear air expansion coeff (1/K)
+    theta_SP = -40*epsilon      # value of theta at south pole (no units)
+    theta_EQ = 30*epsilon       # value of theta at equator (no units)
+    theta_NP = -20*epsilon      # value of theta at north pole (no units)
+    mu1 = 0.05                  # scaling for theta with longitude (no units)
+    mu2 = 0.98                  # proportion of qsat to make init qv (no units)
+    q0 = 135                    # qsat scaling, gives init q_v of ~0.02, (kg/kg)
+    beta2 = 10*g                # buoyancy-vaporisation factor (m/s^2)
+    nu = 20.                    # qsat factor in exponent (no units)
+    qprecip = 1e-4              # cloud to rain conversion threshold (kg/kg)
+    gamma_r = 1e-3              # rain-coalescence implicit factor
+    mountain_height = 2000.     # height of mountain (m)
+    R0 = pi/9.                  # radius of mountain (rad)
+    lamda_c = -pi/2.            # longitudinal centre of mountain (rad)
+    phi_c = pi/6.               # latitudinal centre of mountain (rad)
+
+    # ------------------------------------------------------------------------ #
+    # Our settings for this set up
+    # ------------------------------------------------------------------------ #
+
+    element_order = 1
+    u_eqn_type = 'vector_invariant_form'
+
+    # ------------------------------------------------------------------------ #
+    # Set up model objects
+    # ------------------------------------------------------------------------ #
+
+    # Domain
+    mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2)
+    domain = Domain(mesh, dt, "BDM", element_order)
+    x, y, z = SpatialCoordinate(mesh)
+    lamda, phi, _ = lonlatr_from_xyz(x, y, z)
+
+    # Equation: coriolis
+    parameters = ShallowWaterParameters(H=mean_depth, g=g)
+    Omega = parameters.Omega
+    fexpr = 2*Omega*z/radius
+
+    # Equation: topography
+    rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2)
+    r = sqrt(rsq)
+    tpexpr = mountain_height * (1 - r/R0)
+
+    # Equation: moisture
+    tracers = [
+        WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')
+    ]
+    eqns = ThermalShallowWaterEquations(
+        domain, parameters, fexpr=fexpr, topog_expr=tpexpr,
+        active_tracers=tracers, u_transport_option=u_eqn_type
+    )
+
+    # I/O
+    output = OutputParameters(
+        dirname=str(tmpdir), dumplist_latlon=['D'], dumpfreq=dumpfreq,
+        dump_vtus=True, dump_nc=False, log_courant=False,
+        dumplist=['D', 'b', 'water_vapour', 'cloud_water']
+    )
+    io = IO(domain, output)
+
+    # Physics ------------------------------------------------------------------
+    # Saturation function -- first define simple expression
+    def q_sat(b, D):
+        return (q0/(g*D + g*tpexpr)) * exp(nu*(1 - b/g))
+
+    # Function to pass to physics (takes mixed function as argument)
+    def phys_sat_func(x_in):
+        D = x_in.sub(1)
+        b = x_in.sub(2)
+        return q_sat(b, D)
+
+    # Feedback proportionality is dependent on D and b
+    def gamma_v(x_in):
+        D = x_in.sub(1)
+        b = x_in.sub(2)
+        return 1.0 / (1.0 + nu*beta2/g*q_sat(b, D))
+
+    SWSaturationAdjustment(
+        eqns, phys_sat_func, time_varying_saturation=True,
+        parameters=parameters, thermal_feedback=True,
+        beta2=beta2, gamma_v=gamma_v, time_varying_gamma_v=True
+    )
+
+    InstantRain(
+        eqns, qprecip, vapour_name="cloud_water", rain_name="rain",
+        gamma_r=gamma_r
+    )
+
+    transport_methods = [
+        DGUpwind(eqns, field_name) for field_name in eqns.field_names
+    ]
+
+    # Timestepper
+    stepper = Timestepper(
+        eqns, RK4(domain), io, spatial_methods=transport_methods
+    )
+
+    # ------------------------------------------------------------------------ #
+    # Initial conditions
+    # ------------------------------------------------------------------------ #
+
+    u0 = stepper.fields("u")
+    D0 = stepper.fields("D")
+    b0 = stepper.fields("b")
+    v0 = stepper.fields("water_vapour")
+    c0 = stepper.fields("cloud_water")
+    r0 = stepper.fields("rain")
+
+    uexpr = as_vector([-u_max*y/radius, u_max*x/radius, 0.0])
+
+    Dexpr = (
+        mean_depth - tpexpr
+        - (radius * Omega * u_max + 0.5*u_max**2)*(z/radius)**2/g
+    )
+
+    # Expression for initial buoyancy - note the bracket around 1-mu
+    theta_expr = (
+        2/(pi**2) * (
+            phi*(phi - pi/2)*theta_SP
+            - 2*(phi + pi/2) * (phi - pi/2)*(1 - mu1)*theta_EQ
+            + phi*(phi + pi/2)*theta_NP
+        )
+        + mu1*theta_EQ*cos(phi)*sin(lamda)
+    )
+    bexpr = g * (1 - theta_expr)
+
+    # Expression for initial vapour depends on initial saturation
+    vexpr = mu2 * q_sat(bexpr, Dexpr)
+
+    # Initialise (cloud and rain initially zero)
+    u0.project(uexpr)
+    D0.interpolate(Dexpr)
+    b0.interpolate(bexpr)
+    v0.interpolate(vexpr)
+    c0.assign(0.0)
+    r0.assign(0.0)
+
+    # ----------------------------------------------------------------- #
+    # Run
+    # ----------------------------------------------------------------- #
+    stepper.run(t=0, tmax=dt)
+
+    u_tf = stepper.fields('u')  # Final velocity field
+    D_tf = stepper.fields('D')  # Final depth field
+
+    J = assemble(0.5*inner(u_tf, u_tf)*dx + 0.5*g*D_tf**2*dx)
+
+    Jhat = ReducedFunctional(J, Control(D0))
+
+    assert np.allclose(Jhat(D0), J)
+    with stop_annotating():
+        # Stop annotation to perform the Taylor test
+        h0 = Function(D0.function_space())
+        h0.assign(D0 * np.random.rand())
+        assert taylor_test(Jhat, D0, h0) > 1.95
diff --git a/integration-tests/adjoints/test_shallow_water_sensitivity.py b/integration-tests/adjoints/test_shallow_water_sensitivity.py
new file mode 100644
index 000000000..fe2927707
--- /dev/null
+++ b/integration-tests/adjoints/test_shallow_water_sensitivity.py
@@ -0,0 +1,100 @@
+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()
+
+
+def test_shallow_water(tmpdir):
+    assert get_working_tape()._blocks == []
+    # setup shallow water parameters
+    R = 6371220.
+    H = 5960.
+    dt = 900.
+
+    # Domain
+    mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2)
+    x = SpatialCoordinate(mesh)
+    domain = Domain(mesh, dt, 'BDM', 1)
+    parameters = ShallowWaterParameters(H=H)
+
+    # Equation
+    Omega = parameters.Omega
+    fexpr = 2*Omega*x[2]/R
+    lamda, theta, _ = lonlatr_from_xyz(x[0], x[1], x[2])
+    R0 = pi/9.
+    R0sq = R0**2
+    lamda_c = -pi/2.
+    lsq = (lamda - lamda_c)**2
+    theta_c = pi/6.
+    thsq = (theta - theta_c)**2
+    rsq = min_value(R0sq, lsq+thsq)
+    r = sqrt(rsq)
+    bexpr = 2000 * (1 - r/R0)
+    eqn = ShallowWaterEquations(domain, parameters, fexpr=fexpr, topog_expr=bexpr)
+
+    # I/O
+    output = OutputParameters(dirname=str(tmpdir), log_courant=False)
+    io = IO(domain, output)
+
+    # Transport schemes
+    transported_fields = [TrapeziumRule(domain, "u"), SSPRK3(domain, "D")]
+    transport_methods = [DGUpwind(eqn, "u"), DGUpwind(eqn, "D")]
+
+    # Time stepper
+    stepper = SemiImplicitQuasiNewton(
+        eqn, io, transported_fields, transport_methods
+    )
+
+    u0 = stepper.fields('u')
+    D0 = stepper.fields('D')
+    u_max = 20.   # Maximum amplitude of the zonal wind (m/s)
+    uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
+    g = parameters.g
+    Rsq = R**2
+    Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*x[2]**2/Rsq)/g - bexpr
+
+    u0.project(uexpr)
+    D0.interpolate(Dexpr)
+
+    Dbar = Function(D0.function_space()).assign(H)
+    stepper.set_reference_profiles([('D', Dbar)])
+
+    stepper.run(0., 5*dt)
+
+    u_tf = stepper.fields('u')  # Final velocity field
+    D_tf = stepper.fields('D')  # Final depth field
+
+    J = assemble(0.5*inner(u_tf, u_tf)*dx + 0.5*g*D_tf**2*dx)
+
+    control = [Control(D0), Control(u0)]  # Control variables
+    J_hat = ReducedFunctional(J, control)
+    assert np.isclose(J_hat([D0, u0]), J, rtol=1e-10)
+    with stop_annotating():
+        # Stop annotation to perform the Taylor test
+        h0 = Function(D0.function_space())
+        h1 = Function(u0.function_space())
+        h0.assign(D0 * np.random.rand())
+        h1.assign(u0 * np.random.rand())
+        assert taylor_test(J_hat, [D0, u0], [h0, h1]) > 1.95
diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
index 02b1addfc..b0508aae3 100644
--- a/integration-tests/conftest.py
+++ b/integration-tests/conftest.py
@@ -25,7 +25,7 @@ def tracer_sphere(tmpdir, degree, small_dt):
     # Parameters chosen so that dt != 1
     # Gaussian is translated from (lon=pi/2, lat=0) to (lon=0, lat=0)
     # to demonstrate that transport is working correctly
-    if small_dt:
+    if (small_dt):
         dt = pi/3. * 0.005
     else:
         dt = pi/3. * 0.02
@@ -47,7 +47,7 @@ def tracer_sphere(tmpdir, degree, small_dt):
                        uexpr, umax, radius, tol)
 
 
-def tracer_slice(tmpdir, degree, small_dt):
+def tracer_slice(tmpdir, degree):
     n = 30 if degree == 0 else 15
     m = PeriodicIntervalMesh(n, 1.)
     mesh = ExtrudedMesh(m, layers=n, layer_height=1./n)
@@ -56,10 +56,7 @@ def tracer_slice(tmpdir, degree, small_dt):
     # Gaussian is translated by 1.5 times width of domain to demonstrate
     # that transport is working correctly
 
-    if small_dt:
-        dt = 0.002
-    else:
-        dt = 0.01
+    dt = 0.01
     tmax = 0.75
     output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
     domain = Domain(mesh, dt, family="CG", degree=degree)
@@ -83,11 +80,8 @@ def tracer_slice(tmpdir, degree, small_dt):
     return TracerSetup(domain, tmax, io, f_init, f_end, degree, uexpr, tol=tol)
 
 
-def tracer_blob_slice(tmpdir, degree, small_dt):
-    if small_dt:
-        dt = 0.002
-    else:
-        dt = 0.01
+def tracer_blob_slice(tmpdir, degree):
+    dt = 0.01
     L = 10.
     m = PeriodicIntervalMesh(10, L)
     mesh = ExtrudedMesh(m, layers=10, layer_height=1.)
@@ -111,9 +105,10 @@ def _tracer_setup(tmpdir, geometry, blob=False, degree=1, small_dt=False):
             assert not blob
             return tracer_sphere(tmpdir, degree, small_dt)
         elif geometry == "slice":
+            assert not small_dt
             if blob:
-                return tracer_blob_slice(tmpdir, degree, small_dt)
+                return tracer_blob_slice(tmpdir, degree)
             else:
-                return tracer_slice(tmpdir, degree, small_dt)
+                return tracer_slice(tmpdir, degree)
 
     return _tracer_setup
diff --git a/integration-tests/data/boussinesq_compressible_chkpt.h5 b/integration-tests/data/boussinesq_compressible_chkpt.h5
index 2ad2cab9d..a466f1fbe 100644
Binary files a/integration-tests/data/boussinesq_compressible_chkpt.h5 and b/integration-tests/data/boussinesq_compressible_chkpt.h5 differ
diff --git a/integration-tests/data/boussinesq_incompressible_chkpt.h5 b/integration-tests/data/boussinesq_incompressible_chkpt.h5
index c2aadff89..7719ee798 100644
Binary files a/integration-tests/data/boussinesq_incompressible_chkpt.h5 and b/integration-tests/data/boussinesq_incompressible_chkpt.h5 differ
diff --git a/integration-tests/data/dry_compressible_chkpt.h5 b/integration-tests/data/dry_compressible_chkpt.h5
index 523aef482..9d084a148 100644
Binary files a/integration-tests/data/dry_compressible_chkpt.h5 and b/integration-tests/data/dry_compressible_chkpt.h5 differ
diff --git a/integration-tests/data/linear_sw_wave_chkpt.h5 b/integration-tests/data/linear_sw_wave_chkpt.h5
index 313448078..65d96abdc 100644
Binary files a/integration-tests/data/linear_sw_wave_chkpt.h5 and b/integration-tests/data/linear_sw_wave_chkpt.h5 differ
diff --git a/integration-tests/data/moist_compressible_chkpt.h5 b/integration-tests/data/moist_compressible_chkpt.h5
index bd43cdadc..f4065e241 100644
Binary files a/integration-tests/data/moist_compressible_chkpt.h5 and b/integration-tests/data/moist_compressible_chkpt.h5 differ
diff --git a/integration-tests/data/simult_SIQN_order0_chkpt.h5 b/integration-tests/data/simult_SIQN_order0_chkpt.h5
index ae942022f..4e49fb596 100644
Binary files a/integration-tests/data/simult_SIQN_order0_chkpt.h5 and b/integration-tests/data/simult_SIQN_order0_chkpt.h5 differ
diff --git a/integration-tests/data/simult_SIQN_order1_chkpt.h5 b/integration-tests/data/simult_SIQN_order1_chkpt.h5
index d8210cdd3..96108b658 100644
Binary files a/integration-tests/data/simult_SIQN_order1_chkpt.h5 and b/integration-tests/data/simult_SIQN_order1_chkpt.h5 differ
diff --git a/integration-tests/data/sw_fplane_chkpt.h5 b/integration-tests/data/sw_fplane_chkpt.h5
index 03cde643b..b752f0d70 100644
Binary files a/integration-tests/data/sw_fplane_chkpt.h5 and b/integration-tests/data/sw_fplane_chkpt.h5 differ
diff --git a/integration-tests/equations/test_sw_fplane.py b/integration-tests/equations/test_sw_fplane.py
index 3d0acff5f..8e0779c2b 100644
--- a/integration-tests/equations/test_sw_fplane.py
+++ b/integration-tests/equations/test_sw_fplane.py
@@ -38,12 +38,10 @@ def run_sw_fplane(tmpdir):
     io = IO(domain, output, diagnostic_fields=[CourantNumber()])
 
     # Transport schemes
-    vorticity_transport = VorticityTransport(domain, eqns, supg=True)
-    transported_fields = [
-        TrapeziumRule(domain, "u", augmentation=vorticity_transport),
-        SSPRK3(domain, "D")
-    ]
-    transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
+    transported_fields = [TrapeziumRule(domain, "u"),
+                          SSPRK3(domain, "D")]
+    transport_methods = [DGUpwind(eqns, "u"),
+                         DGUpwind(eqns, "D")]
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields,
diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py
index 0c259e221..6d484bfd1 100644
--- a/integration-tests/model/test_time_discretisation.py
+++ b/integration-tests/model/test_time_discretisation.py
@@ -10,19 +10,9 @@ def run(timestepper, tmax, f_end):
 
 @pytest.mark.parametrize(
     "scheme", [
-        "ssprk2_increment_2", "ssprk2_predictor_2", "ssprk2_linear_2",
-        "ssprk2_increment_3", "ssprk2_predictor_3", "ssprk2_linear_3",
-        "ssprk2_increment_4", "ssprk2_predictor_4", "ssprk2_linear_4",
-
-        "ssprk3_increment_3", "ssprk3_predictor_3", "ssprk3_linear_3",
-        "ssprk3_increment_4", "ssprk3_predictor_4", "ssprk3_linear_4",
-        "ssprk3_increment_5", "ssprk3_predictor_5", "ssprk3_linear_5",
-
-        "ssprk4_increment_5", "ssprk4_predictor_5", "ssprk4_linear_5",
-
-        "TrapeziumRule", "ImplicitMidpoint", "QinZhang",
+        "ssprk3_increment", "TrapeziumRule", "ImplicitMidpoint", "QinZhang",
         "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog",
-        "AdamsMoulton", "AdamsMoulton"
+        "AdamsMoulton", "AdamsMoulton", "ssprk3_predictor", "ssprk3_linear"
     ]
 )
 def test_time_discretisation(tmpdir, scheme, tracer_setup):
@@ -40,52 +30,12 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup):
         V = domain.spaces("DG")
         eqn = AdvectionEquation(domain, V, "f")
 
-    if scheme == "ssprk2_increment_2":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment)
-    elif scheme == "ssprk2_predictor_2":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor)
-    elif scheme == "ssprk2_linear_2":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear)
-    elif scheme == "ssprk2_increment_3":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment, stages=3)
-    elif scheme == "ssprk2_predictor_3":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=3)
-    elif scheme == "ssprk2_linear_3":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear, stages=3)
-    elif scheme == "ssprk2_increment_4":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment, stages=4)
-    elif scheme == "ssprk2_predictor_4":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=4)
-    elif scheme == "ssprk2_linear_4":
-        transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear, stages=4)
-
-    elif scheme == "ssprk3_increment_3":
+    if scheme == "ssprk3_increment":
         transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment)
-    elif scheme == "ssprk3_predictor_3":
+    elif scheme == "ssprk3_predictor":
         transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor)
-    elif scheme == "ssprk3_linear_3":
+    elif scheme == "ssprk3_linear":
         transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear)
-    elif scheme == "ssprk3_increment_4":
-        transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment, stages=4)
-    elif scheme == "ssprk3_predictor_4":
-        transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=4)
-    elif scheme == "ssprk3_linear_4":
-        transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear, stages=4)
-
-    elif scheme == "ssprk3_increment_5":
-        transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment, stages=5)
-    elif scheme == "ssprk3_predictor_5":
-        transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=5)
-    elif scheme == "ssprk3_linear_5":
-        transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear, stages=5)
-
-    elif scheme == "ssprk4_increment_5":
-        transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.increment)
-    elif scheme == "ssprk4_predictor_5":
-        transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.predictor)
-    elif scheme == "ssprk4_linear_5":
-        transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.linear)
-
     elif scheme == "TrapeziumRule":
         transport_scheme = TrapeziumRule(domain)
     elif scheme == "ImplicitMidpoint":
diff --git a/integration-tests/physics/test_held_suarez_friction.py b/integration-tests/physics/test_held_suarez_friction.py
deleted file mode 100644
index dfe7e34f9..000000000
--- a/integration-tests/physics/test_held_suarez_friction.py
+++ /dev/null
@@ -1,114 +0,0 @@
-"""
-This tests the Rayleigh friction term used in the Held Suarez test case.
-"""
-
-from gusto import *
-import gusto.equations.thermodynamics as td
-from gusto.core.labels import physics_label
-from firedrake import (Constant, PeriodicIntervalMesh, as_vector, norm,
-                       ExtrudedMesh, Function, dot)
-from firedrake.fml import identity, drop
-
-
-def run_apply_rayleigh_friction(dirname):
-    # ------------------------------------------------------------------------ #
-    # Set up model objects
-    # ------------------------------------------------------------------------ #
-
-    dt = 3600.0
-
-    # declare grid shape, with length L and height H
-    L = 500.
-    H = 500.
-    nlayers = int(H / 5.)
-    ncolumns = int(L / 5.)
-
-    # make mesh and domain
-    m = PeriodicIntervalMesh(ncolumns, L)
-    mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
-    domain = Domain(mesh, dt, "CG", 0)
-
-    # Set up equation
-    parameters = CompressibleParameters()
-    eqn = CompressibleEulerEquations(domain, parameters)
-
-    # I/O
-    output = OutputParameters(dirname=dirname+"/held_suarez_friction",
-                              dumpfreq=1,
-                              dumplist=['u'])
-    io = IO(domain, output)
-
-    # Physics scheme
-    physics_parametrisation = RayleighFriction(eqn)
-
-    time_discretisation = BackwardEuler(domain)
-
-    # time_discretisation = ForwardEuler(domain)
-    physics_schemes = [(physics_parametrisation, time_discretisation)]
-
-    # Only want time derivatives and physics terms in equation, so drop the rest
-    eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)),
-                                          map_if_true=identity, map_if_false=drop)
-
-    # Time stepper
-    scheme = ForwardEuler(domain)
-    stepper = SplitPhysicsTimestepper(eqn, scheme, io,
-                                      physics_schemes=physics_schemes)
-
-    # ------------------------------------------------------------------------ #
-    # Initial conditions
-    # ------------------------------------------------------------------------ #
-
-    Vu = domain.spaces("HDiv")
-    Vt = domain.spaces("theta")
-    Vr = domain.spaces("DG")
-
-    # Declare prognostic fields
-    u0 = stepper.fields("u")
-    rho0 = stepper.fields("rho")
-    theta0 = stepper.fields("theta")
-
-    # Set a background state with constant pressure and temperature
-    pressure = Function(Vr).interpolate(Constant(100000.))
-    temperature = Function(Vt).interpolate(Constant(295.))
-    theta_d = td.theta(parameters, temperature, pressure)
-
-    theta0.project(theta_d)
-    rho0.interpolate(pressure / (temperature*parameters.R_d))
-
-    # Constant horizontal wind
-    u0.project(as_vector([864, 0.0]))
-
-    # Answer: slower winds than initially
-    u_true = Function(Vu)
-    u_true.project(as_vector([828, 0.0]))
-
-    # ------------------------------------------------------------------------ #
-    # Run
-    # ------------------------------------------------------------------------ #
-
-    stepper.run(t=0, tmax=dt)
-
-    return mesh, stepper, u_true
-
-
-def test_rayleigh_friction(tmpdir):
-
-    dirname = str(tmpdir)
-    mesh, stepper, u_true = run_apply_rayleigh_friction(dirname)
-
-    u_final = stepper.fields('u')
-
-    # Project into CG1 to get sensible values
-    e_x = as_vector([1.0, 0.0])
-    e_z = as_vector([0.0, 1.0])
-
-    DG0 = FunctionSpace(mesh, "DG", 0)
-    u_x_final = Function(DG0).project(dot(u_final, e_x))
-    u_x_true = Function(DG0).project(dot(u_true, e_x))
-    u_z_final = Function(DG0).project(dot(u_final, e_z))
-    u_z_true = Function(DG0).project(dot(u_true, e_z))
-
-    denom = norm(u_x_true)
-    assert norm(u_x_final - u_x_true) / denom < 0.0001, 'Final horizontal wind is incorrect'
-    assert norm(u_z_final - u_z_true) < 1e-12, 'Final vertical wind is incorrect'
diff --git a/integration-tests/physics/test_held_suarez_relaxation.py b/integration-tests/physics/test_held_suarez_relaxation.py
deleted file mode 100644
index 615ba8d29..000000000
--- a/integration-tests/physics/test_held_suarez_relaxation.py
+++ /dev/null
@@ -1,107 +0,0 @@
-"""
-This tests the Held-Suarez physics routine to apply Rayleigh friction.
-"""
-from gusto import *
-import gusto.equations.thermodynamics as td
-from gusto.core.labels import physics_label
-from firedrake import (Constant, PeriodicIntervalMesh, as_vector,
-                       ExtrudedMesh, Function)
-from firedrake.fml import identity, drop
-import pytest
-
-
-def run_held_suarez_relaxation(dirname, temp):
-
-    # ------------------------------------------------------------------------ #
-    # Set up model objects
-    # ------------------------------------------------------------------------ #
-
-    dt = 3600.0
-
-    # declare grid shape, with length L and height H
-    L = 500.
-    H = 500.
-    nlayers = int(H / 5.)
-    ncolumns = int(L / 5.)
-
-    # make mesh and domain
-    m = PeriodicIntervalMesh(ncolumns, L)
-    mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
-    domain = Domain(mesh, dt, "CG", 0)
-
-    # Set up equation
-    parameters = CompressibleParameters()
-    eqn = CompressibleEulerEquations(domain, parameters)
-
-    # I/O
-    output = OutputParameters(dirname=dirname+"/held_suarez_friction",
-                              dumpfreq=1,
-                              dumplist=['u'])
-    io = IO(domain, output)
-
-    # Physics scheme
-    physics_parametrisation = Relaxation(eqn, 'theta', parameters)
-
-    time_discretisation = ForwardEuler(domain)
-
-    # time_discretisation = ForwardEuler(domain)
-    physics_schemes = [(physics_parametrisation, time_discretisation)]
-
-    # Only want time derivatives and physics terms in equation, so drop the rest
-    eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)),
-                                          map_if_true=identity, map_if_false=drop)
-
-    # Time stepper
-    scheme = ForwardEuler(domain)
-    stepper = SplitPhysicsTimestepper(eqn, scheme, io,
-                                      physics_schemes=physics_schemes)
-
-    # ------------------------------------------------------------------------ #
-    # Initial conditions
-    # ------------------------------------------------------------------------ #
-
-    Vt = domain.spaces("theta")
-    Vr = domain.spaces("DG")
-
-    # Declare prognostic fields
-    u0 = stepper.fields("u")
-    rho0 = stepper.fields("rho")
-    theta0 = stepper.fields("theta")
-
-    # Set a background state with constant pressure and temperature
-    p0 = 100000.
-    pressure = Function(Vr).interpolate(Constant(p0))
-    temperature = Function(Vt).interpolate(Constant(temp))
-    theta_d = td.theta(parameters, temperature, pressure)
-
-    theta0.project(theta_d)
-    rho0.interpolate(p0 / (theta_d*parameters.R_d))  # This ensures that exner = 1
-
-    # Constant horizontal wind
-    u0.project(as_vector([1.0, 1.0]))
-    theta_initial = Function(Vt)
-    theta_initial.interpolate(theta_d)
-
-    # ------------------------------------------------------------------------ #
-    # Run
-    # ------------------------------------------------------------------------ #
-
-    stepper.run(t=0, tmax=dt)
-
-    return stepper, theta_initial
-
-
-@pytest.mark.parametrize('temp', [280, 290])
-def test_held_suarez_relaxation(tmpdir, temp):
-    # By configuring the fields we have set the equilibrium temperature to 285K
-    # We test a temperature value eith side to check it moves in the right direction
-    dirname = str(tmpdir)
-    stepper, theta_initial = run_held_suarez_relaxation(dirname, temp)
-
-    theta_final = stepper.fields('theta')
-    final_data = theta_final.dat.data
-    initial_data = theta_initial.dat.data
-    if temp == 280:
-        assert np.mean(final_data) > np.mean(initial_data)
-    if temp == 290:
-        assert np.mean(final_data) < np.mean(initial_data)
diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py
index 9f6dd773f..54c951cd7 100644
--- a/integration-tests/transport/test_vector_recovered_space.py
+++ b/integration-tests/transport/test_vector_recovered_space.py
@@ -15,9 +15,6 @@ def run(timestepper, tmax, f_end):
     return norm(timestepper.fields("f") - f_end) / norm(f_end)
 
 
-# NB: The default vector transport test is not valid on the sphere as it is
-# designed for a 3-component vector function space, and not the space of tangent
-# vectors on the sphere
 @pytest.mark.parametrize("geometry", ["slice"])
 def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup):
 
diff --git a/integration-tests/transport/test_vorticity_transport.py b/integration-tests/transport/test_vorticity_transport.py
deleted file mode 100644
index 1cd9bdcbb..000000000
--- a/integration-tests/transport/test_vorticity_transport.py
+++ /dev/null
@@ -1,111 +0,0 @@
-"""
-This tests the transport of a vector-valued field using vorticity augmentation.
-The computed solution is compared with a true one to check that the transport
-is working correctly.
-"""
-
-from gusto import *
-from firedrake import (
-    as_vector, norm, exp, PeriodicRectangleMesh, SpatialCoordinate, min_value
-)
-import pytest
-
-
-def run(timestepper, tmax, f_end):
-    timestepper.run(0, tmax)
-
-    return norm(timestepper.fields("f") - f_end) / norm(f_end)
-
-
-@pytest.mark.parametrize("supg", [False, True])
-def test_vorticity_transport_setup(tmpdir, supg):
-
-    # ------------------------------------------------------------------------ #
-    # Parameters for test case
-    # ------------------------------------------------------------------------ #
-
-    Lx = 2000.       # length of domain in x direction, in m
-    Ly = 2000.       # width of domain in y direction, in m
-    ncells_1d = 20   # number of points in x and y directions
-    tmax = 500.
-    degree = 1
-
-    if supg:
-        # Smaller time steps for RK scheme
-        dt = 5.0
-        dumpfreq = 50
-    else:
-        dt = 25.0
-        dumpfreq = 10
-
-    # ------------------------------------------------------------------------ #
-    # Set up model objects
-    # ------------------------------------------------------------------------ #
-
-    mesh = PeriodicRectangleMesh(ncells_1d, ncells_1d, Lx, Ly, quadrilateral=True)
-    domain = Domain(mesh, dt, "RTCF", degree)
-    x, y = SpatialCoordinate(mesh)
-
-    Vu = domain.spaces("HDiv")
-
-    # Equation
-    eqn = AdvectionEquation(domain, Vu, "f")
-
-    # I/O
-    dirname = f'{tmpdir}/vorticity_plane'
-    output = OutputParameters(
-        dirname=dirname, dumpfreq=dumpfreq, dump_nc=False, dump_vtus=True
-    )
-    io = IO(domain, output)
-
-    augmentation = VorticityTransport(domain, eqn, supg=supg)
-
-    # Make equation
-    if supg:
-        transport_scheme = SSPRK3(
-            domain, augmentation=augmentation,
-            rk_formulation=RungeKuttaFormulation.predictor
-        )
-    else:
-        transport_scheme = TrapeziumRule(domain, augmentation=augmentation)
-    transport_method = DGUpwind(eqn, "f")
-
-    time_varying_velocity = False
-    timestepper = PrescribedTransport(
-        eqn, transport_scheme, io, time_varying_velocity, transport_method
-    )
-
-    # ------------------------------------------------------------------------ #
-    # Initial conditions
-    # ------------------------------------------------------------------------ #
-
-    # Specify locations of the two Gaussians
-    xc = Lx/2.
-    xend = 3.*Lx/4.
-    yc = Ly/2.
-
-    def l2_dist(xc, yc):
-        return min_value(abs(x - xc), Lx - abs(x - xc))**2 + (y - yc)**2
-
-    f0 = 1.
-    lc = 4.*Lx/25.
-
-    init_scalar_expr = f0*exp(-l2_dist(xc, yc)/lc**2)
-    uexpr = as_vector([1.0, 0.0])
-
-    # Set fields
-    f = timestepper.fields("f")
-    f.project(as_vector([init_scalar_expr, init_scalar_expr]))
-
-    u0 = timestepper.fields("u")
-    u0.project(uexpr)
-
-    final_scalar_expr = f0*exp(-l2_dist(xend, yc)/lc**2)
-    final_vector_expr = as_vector([final_scalar_expr, final_scalar_expr])
-
-    # Run and check error
-    error = run(timestepper, tmax, final_vector_expr)
-
-    tol = 1e-1
-    assert error < tol, \
-        'The transport error is greater than the permitted tolerance'
diff --git a/unit-tests/diagnostic_tests/test_cumulative_sum.py b/unit-tests/diagnostic_tests/test_cumulative_sum.py
deleted file mode 100644
index 295e522ca..000000000
--- a/unit-tests/diagnostic_tests/test_cumulative_sum.py
+++ /dev/null
@@ -1,48 +0,0 @@
-
-from gusto.diagnostics import CumulativeSum
-from gusto.core.fields import StateFields, PrescribedFields, TimeLevelFields
-from gusto import (Domain, CompressibleParameters, CompressibleEulerEquations,
-                   WaterVapour)
-from firedrake import PeriodicIntervalMesh, ExtrudedMesh
-import numpy as np
-
-
-def test_dewpoint():
-
-    L = 10
-    H = 10
-    ncol = 3
-    nlayers = 3
-
-    m = PeriodicIntervalMesh(ncol, L)
-    mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
-
-    domain = Domain(mesh, 0.1, 'CG', 1)
-    params = CompressibleParameters()
-    active_tracers = [WaterVapour()]
-    eqn = CompressibleEulerEquations(domain, params, active_tracers=active_tracers)
-    prog_fields = TimeLevelFields(eqn)
-
-    Vtheta = domain.spaces('theta')
-
-    # Setting up prognostic fields for the diagnostic to use
-    prescribed_fields = PrescribedFields()
-    state_fields = StateFields(prog_fields, prescribed_fields)
-
-    theta = state_fields('theta', Vtheta)
-
-    # Initial conditions
-    theta.interpolate(300.0)
-
-    diagnostic = CumulativeSum(name='theta')
-    diagnostic.setup(domain, state_fields)
-    diagnostic.compute()
-
-    assert np.allclose(diagnostic.field.dat.data, 300.0, atol=0.0), \
-        'The cumulative sum diagnostic does not seem to be correct after 1 iteration'
-
-    theta.interpolate(250.0)
-    diagnostic.compute()
-
-    assert np.allclose(diagnostic.field.dat.data, 550.0, atol=0.0), \
-        'The cumulative sum diagnostic does not seem to be correct after 2 iterations'
diff --git a/unit-tests/test_function_spaces.py b/unit-tests/test_function_spaces.py
index bcf105e35..b39aa8cc5 100644
--- a/unit-tests/test_function_spaces.py
+++ b/unit-tests/test_function_spaces.py
@@ -100,27 +100,6 @@ def test_de_rham_spaces(domain, family):
     assert elt.degree() == degree, '"L2" space does not seem to be degree ' \
         + f'{degree}. Found degree {elt.degree()}'
 
-    # Check that continuities have been recorded correctly
-    if hasattr(mesh, "_base_mesh"):
-        expected_continuity = {
-            "H1": {'horizontal': True, 'vertical': True},
-            "L2": {'horizontal': False, 'vertical': False},
-            "HDiv": {'horizontal': True, 'vertical': True},
-            "HCurl": {'horizontal': True, 'vertical': True},
-            "theta": {'horizontal': False, 'vertical': True}
-        }
-    else:
-        expected_continuity = {
-            "H1": True,
-            "L2": False,
-            "HDiv": True,
-            "HCurl": True
-        }
-
-    for space, continuity in expected_continuity.items():
-        if space in spaces.continuity:
-            assert spaces.continuity[space] == continuity
-
 
 # ---------------------------------------------------------------------------- #
 # Test creation of DG1 equispaced
@@ -139,13 +118,6 @@ def test_dg_equispaced(domain, family):
     assert elt.variant() == "equispaced", '"DG1 equispaced" does not seem to ' \
         + f'be equispaced variant. Found variant {elt.variant()}'
 
-    if hasattr(mesh, "_base_mesh"):
-        expected_continuity = {'horizontal': False, 'vertical': False}
-    else:
-        expected_continuity = False
-
-    assert spaces.continuity['DG1_equispaced'] == expected_continuity, "DG is discontinuous"
-
 
 # ---------------------------------------------------------------------------- #
 # Test creation of DG0 space
@@ -161,13 +133,6 @@ def test_dg0(domain, family):
     assert elt.degree() in [0, (0, 0)], '"DG0" space does not seem to be' \
         + f'degree 0. Found degree {elt.degree()}'
 
-    if hasattr(mesh, "_base_mesh"):
-        expected_continuity = {'horizontal': False, 'vertical': False}
-    else:
-        expected_continuity = False
-
-    assert spaces.continuity['DG0'] == expected_continuity, "DG is discontinuous"
-
 
 # ---------------------------------------------------------------------------- #
 # Test creation of a general CG space
@@ -185,10 +150,3 @@ def test_cg(domain, family):
     assert elt.degree() == degree or elt.degree() == (degree, degree), \
         (f'"CG" space does not seem to be degree {degree}. '
          + f'Found degree {elt.degree()}')
-
-    if hasattr(mesh, "_base_mesh"):
-        expected_continuity = {'horizontal': True, 'vertical': True}
-    else:
-        expected_continuity = True
-
-    assert spaces.continuity['CG'] == expected_continuity, "CG is continuous"