diff --git a/gusto/solvers/preconditioners.py b/gusto/solvers/preconditioners.py
index dccfa9b2b..e523054b6 100644
--- a/gusto/solvers/preconditioners.py
+++ b/gusto/solvers/preconditioners.py
@@ -346,7 +346,7 @@ def apply(self, pc, x, y):
 
             unbroken_res_hdiv = self.unbroken_residual.subfunctions[self.vidx]
             broken_res_hdiv = self.broken_residual.subfunctions[self.vidx]
-            broken_res_hdiv.assign(0)
+            broken_res_hdiv.zero()
             self.average_kernel.apply(broken_res_hdiv, self.weight, unbroken_res_hdiv)
 
             # Compute the rhs for the multiplier system
@@ -374,7 +374,7 @@ def apply(self, pc, x, y):
             # Compute the hdiv projection of the broken hdiv solution
             broken_hdiv = self.broken_solution.subfunctions[self.vidx]
             unbroken_hdiv = self.unbroken_solution.subfunctions[self.vidx]
-            unbroken_hdiv.assign(0)
+            unbroken_hdiv.zero()
 
             self.average_kernel.apply(unbroken_hdiv, self.weight, broken_hdiv)
 
@@ -735,7 +735,7 @@ def apply(self, pc, x, y):
 
         # Recover broken u and rho
         u_broken, rho, l = self.y_hybrid.subfunctions
-        self.u_hdiv.assign(0)
+        self.u_hdiv.zero()
         self._average_kernel.apply(self.u_hdiv, self._weight, u_broken)
         for bc in self.bcs:
             bc.zero(self.u_hdiv)
@@ -745,7 +745,7 @@ def apply(self, pc, x, y):
         self.y.subfunctions[1].assign(rho)
 
         # Recover theta
-        self.theta.assign(0)
+        self.theta.zero()
         self.theta_solver.solve()
         self.y.subfunctions[2].assign(self.theta)
 
diff --git a/gusto/solvers/solver_presets.py b/gusto/solvers/solver_presets.py
index ff718d44a..76392ffa0 100644
--- a/gusto/solvers/solver_presets.py
+++ b/gusto/solvers/solver_presets.py
@@ -40,7 +40,7 @@ def hybridised_solver_parameters(equation, solver_prognostics, alpha=0.5, tau_va
 
     # Callback function for the nullspace of the trace system
     def trace_nullsp(T):
-        return VectorSpaceBasis(constant=True)
+        return VectorSpaceBasis(constant=True, comm=equation.function_space.mesh().comm)
 
     # Prepare some generic variables to help with logic below ------------------
 
diff --git a/integration-tests/adjoints/test_diffusion_sensitivity.py b/integration-tests/adjoints/test_diffusion_sensitivity.py
index 7c353eef9..89297e1c9 100644
--- a/integration-tests/adjoints/test_diffusion_sensitivity.py
+++ b/integration-tests/adjoints/test_diffusion_sensitivity.py
@@ -3,32 +3,16 @@
 
 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 autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @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))
@@ -41,38 +25,59 @@ def test_diffusion_sensitivity(nu_is_control, tmpdir):
 
     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)
+    nu = Function(R, val=0.2)
     diffusion_params = DiffusionParameters(mesh, kappa=nu)
     eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params)
 
-    diffusion_scheme = BackwardEuler(domain)
+    # Don't let an inexact solve get in the way of a good Taylor test
+    solver_parameters = {
+        'snes_type': 'ksponly',
+        'ksp_type': 'preonly',
+        'pc_type': 'lu',
+        'pc_factor_mat_solver_type': 'mumps',
+    }
+    diffusion_scheme = BackwardEuler(domain, solver_parameters=solver_parameters)
     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)
+    ic = Function(V).interpolate(fexpr)
 
-    end = 0.1
-    timestepper.run(0., end)
+    # These are the only operations we are interested in rerunning with the tape.
+    with set_working_tape() as tape:
+        # initialise the solution
+        timestepper.fields("f").assign(ic)
 
-    u = timestepper.fields("f")
-    J = assemble(inner(u, u)*dx)
+        # propagate forward
+        end = 0.1
+        timestepper.run(0., end)
 
-    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())
+        # Make sure we have more than a quadratic nonlinearity so Jhat.hessian isn't exact.
+        u = timestepper.fields("f")
+        J = assemble(inner(u, u)*dx)**2
 
-    # the functional as a pure function of nu
-    Jhat = ReducedFunctional(J, control)
+        if nu_is_control:
+            m = nu
+        else:
+            m = ic
 
+        # the functional as a pure function of the control
+        Jhat = ReducedFunctional(J, Control(m), tape=tape)
+
+    assert np.allclose(J, Jhat(m)), "Functional re-evaluation does not match original evaluation."
+
+    # perturbation direction for taylor test
     if nu_is_control:
-        assert np.allclose(J, Jhat(nu))
-        assert taylor_test(Jhat, nu, h) > 1.95
+        h = Function(R, val=0.1)
     else:
-        assert np.allclose(J, Jhat(u))
-        assert taylor_test(Jhat, u, h) > 1.95
+        h = Function(V).interpolate(as_vector([cos(4*pi*x[1]), sin(4*pi*x[0])]))
+
+    # Check the TLM explicitly before checking the Hessian (which relies on the tlm)
+    assert taylor_test(Jhat, m, h, dJdm=Jhat.tlm(h)) > 1.95, "TLM is not second order accurate."
+
+    # Check the re-evaluation, derivative, and Hessian all converge at the expected rates.
+    taylor = taylor_to_dict(Jhat, m, h)
+    assert min(taylor['R0']['Rate']) > 0.95, taylor['R0']
+    assert min(taylor['R1']['Rate']) > 1.95, taylor['R1']
+    assert min(taylor['R2']['Rate']) > 2.95, taylor['R2']
diff --git a/integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py b/integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py
index 3f076bf35..810451e77 100644
--- a/integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py
+++ b/integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py
@@ -21,7 +21,6 @@
     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,
@@ -31,29 +30,14 @@
 
 
 @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 autouse_set_test_tape(set_test_tape):
+    pass
 
 
 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
     # ------------------------------------------------------------------------ #
@@ -150,9 +134,18 @@ def gamma_v(x_in):
         DGUpwind(eqns, field_name) for field_name in eqns.field_names
     ]
 
+    # Don't let an inexact solve get in the way of a good Taylor test
+    solver_parameters = {
+        'snes_type': 'ksponly',
+        'ksp_type': 'preonly',
+        'pc_type': 'lu',
+        'pc_factor_mat_solver_type': 'mumps',
+    }
+
     # Timestepper
     stepper = Timestepper(
-        eqns, RK4(domain), io, spatial_methods=transport_methods
+        eqns, RK4(domain, solver_parameters=solver_parameters),
+        io, spatial_methods=transport_methods
     )
 
     # ------------------------------------------------------------------------ #
@@ -188,29 +181,41 @@ def gamma_v(x_in):
     # 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)
+    # Control
+    m_D = Function(D0.function_space()).interpolate(Dexpr)
+
+    with set_working_tape() as tape:
+        # Initialise (cloud and rain initially zero)
+        u0.project(uexpr)
+        D0.interpolate(m_D)
+        b0.interpolate(bexpr)
+        v0.interpolate(vexpr)
+        c0.zero()
+        r0.zero()
+
+        # ----------------------------------------------------------------- #
+        # Run
+        # ----------------------------------------------------------------- #
+        # two timesteps so we hit the codepath for reshuffling data between steps
+        stepper.run(t=0., tmax=2*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(m_D), tape=tape)
 
-    # ----------------------------------------------------------------- #
-    # Run
-    # ----------------------------------------------------------------- #
-    stepper.run(t=0, tmax=dt)
+    assert np.allclose(Jhat(m_D), J), "Functional re-evaluation does not match original evaluation."
 
-    u_tf = stepper.fields('u')  # Final velocity field
-    D_tf = stepper.fields('D')  # Final depth field
+    # perturbation direction for taylor test
+    h_D = Function(D0.function_space())
+    h_D.interpolate(0.1*(Dexpr - (mean_depth - tpexpr)))
 
-    J = assemble(0.5*inner(u_tf, u_tf)*dx + 0.5*g*D_tf**2*dx)
+    # Check the TLM explicitly before checking the Hessian (which relies on the tlm)
+    assert taylor_test(Jhat, m_D, h_D, dJdm=Jhat.tlm(h_D)) > 1.95, "TLM is not second order accurate."
 
-    Jhat = ReducedFunctional(J, Control(D0))
+    assert taylor_test(Jhat, m_D, h_D) > 1.95, "Adjoint derivative is not second order accurate."
 
-    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
+    # NOTE: Currently computing the Hessian will raise an exception
+    # see https://github.com/FEniCS/ufl/issues/477
diff --git a/integration-tests/adjoints/test_shallow_water_sensitivity.py b/integration-tests/adjoints/test_shallow_water_sensitivity.py
index e598ce5f6..02a134192 100644
--- a/integration-tests/adjoints/test_shallow_water_sensitivity.py
+++ b/integration-tests/adjoints/test_shallow_water_sensitivity.py
@@ -1,45 +1,32 @@
 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 == []
+def autouse_set_test_tape(set_test_tape):
+    pass
+
+
+@pytest.mark.parametrize("control", ["u", "D"])
+@pytest.mark.parametrize("stepper_type", ["BackwardEuler", "RK4", "SemiImplicitQuasiNewton"])
+def test_shallow_water(tmpdir, control, stepper_type):
     # setup shallow water parameters
     R = 6371220.
     H = 5960.
-    dt = 900.
+    mountain = 2000.
+    dt = 300.
+    u_max = 20.
 
     # Domain
-    mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2)
-    x = SpatialCoordinate(mesh)
+    mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=1)
+    x, y, z = SpatialCoordinate(mesh)
     domain = Domain(mesh, dt, 'BDM', 1)
 
     # Equation
-    lamda, theta, _ = lonlatr_from_xyz(x[0], x[1], x[2])
+    lamda, theta, _ = lonlatr_from_xyz(x, y, z)
     R0 = pi/9.
     R0sq = R0**2
     lamda_c = -pi/2.
@@ -48,52 +35,122 @@ def test_shallow_water(tmpdir):
     thsq = (theta - theta_c)**2
     rsq = min_value(R0sq, lsq+thsq)
     r = sqrt(rsq)
-    bexpr = 2000 * (1 - r/R0)
-    parameters = ShallowWaterParameters(mesh, H=H, topog_expr=bexpr)
+    bexpr = mountain * (1 - r/R0)
+    b = Function(domain.spaces("DG")).interpolate(bexpr)
+    parameters = ShallowWaterParameters(mesh, H=H, topog_expr=b)
     eqn = ShallowWaterEquations(domain, parameters)
 
     # I/O
     output = OutputParameters(dirname=str(tmpdir), log_courant=False)
     io = IO(domain, output)
 
+    # Don't let an inexact solve get in the way of a good Taylor test
+    lu_parameters = {
+        'ksp_type': 'preonly',
+        'pc_type': 'lu',
+        'pc_factor_mat_solver_type': 'mumps',
+    }
+    mass_parameters = {
+        'snes_type': 'ksponly',
+        **lu_parameters,
+        'ksp_reuse_preconditioner': None,
+    }
+    snes_parameters = {
+        'snes_monitor': None,
+        'snes_converged_reason': None,
+        'snes_type': 'newtonls',
+        'snes_stol': 0,
+        'snes_rtol': 1e-10,
+        **lu_parameters,
+    }
+    linear_solver_parameters = {
+        'mat_type': 'matfree',
+        'ksp_type': 'preonly',
+        'pc_type': 'python',
+        'pc_python_type': 'firedrake.HybridizationPC',
+        'hybridization': lu_parameters,
+    }
+
     # 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
-    )
+    if stepper_type == "BackwardEuler":
+        stepper = Timestepper(
+            eqn, BackwardEuler(domain, solver_parameters=snes_parameters),
+            io, spatial_methods=transport_methods
+        )
+    elif stepper_type == "RK4":
+        stepper = Timestepper(
+            eqn, RK4(domain, solver_parameters=mass_parameters),
+            io, spatial_methods=transport_methods
+        )
+    else:
+        assert stepper_type == "SemiImplicitQuasiNewton"
+        transported_fields = [
+            TrapeziumRule(domain, "u", solver_parameters=snes_parameters),
+            SSPRK3(domain, "D", solver_parameters=mass_parameters),
+        ]
+        stepper = SemiImplicitQuasiNewton(
+            eqn, io, transported_fields, transport_methods,
+            linear_solver_parameters=linear_solver_parameters
+        )
 
     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])
+
+    uexpr = as_vector([-u_max*y/R, u_max*x/R, 0.0])
     g = parameters.g
     Omega = parameters.Omega
     Rsq = R**2
-    Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*x[2]**2/Rsq)/g - bexpr
+    Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*z**2/Rsq)/g - bexpr
 
-    u0.project(uexpr)
-    D0.interpolate(Dexpr)
+    # controls m for the initial conditions
+    m_u = Function(u0.function_space()).project(uexpr)
+    m_D = Function(D0.function_space()).interpolate(Dexpr)
+
+    if control == "u":
+        m = m_u
+    else:
+        m = m_D
 
     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
+    # These are the only operations we are interested in rerunning with the tape.
+    with set_working_tape() as tape:
+        # initialise the solution
+        u0.assign(m_u)
+        D0.assign(m_D)
+
+        # propagate forwards
+        stepper.run(0., 2*dt)
+
+        if control == "u":
+            u_tf = stepper.fields('u')
+            J = assemble(inner(u_tf, u_tf)*dx)**2
+        else:
+            D_tf = stepper.fields('D')
+            dD = D_tf - (H - b)
+            J = assemble(0.5*g*inner(dD, dD)*dx)**2
+
+        Jhat = ReducedFunctional(J, Control(m), tape=tape)
+
+    # Perturbation directions for taylor test
+    # pyadjoint will multiply h by 1e-2, 1e-4 etc so we pre-multiply by 10
+    if control == "u":
+        h = Function(u0.function_space()).interpolate(
+            10.0*u_max*as_vector([cos(2*pi*y/R), 0.0, 1 + cos(pi*z/R)*sin(4*pi*x/R)]))
+    else:
+        h = Function(D0.function_space()).interpolate(1.0*(Dexpr - (H - b))*cos(pi*z/R))
+
+    assert abs(float(J) - float(Jhat(m))) < 1e-10, "Functional re-evaluation does not match original evaluation."
+
+    # Check the TLM explicitly before checking the Hessian (which relies on the tlm)
+    assert taylor_test(Jhat, m, h, dJdm=Jhat.tlm(h)) > 1.95, "TLM is not second order accurate."
+
+    # Check the re-evaluation, derivative, and Hessian all converge at the expected rates.
+    taylor = taylor_to_dict(Jhat, m, h)
+    assert min(taylor['R0']['Rate']) > 0.95, taylor['R0']
+    assert min(taylor['R1']['Rate']) > 1.95, taylor['R1']
+    assert min(taylor['R2']['Rate']) > 2.95, taylor['R2']
diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
index 7fce74377..939ab8bb9 100644
--- a/integration-tests/conftest.py
+++ b/integration-tests/conftest.py
@@ -8,6 +8,10 @@
 from gusto import *
 from collections import namedtuple
 import pytest
+from pyadjoint.tape import (
+    annotate_tape, get_working_tape, set_working_tape,
+    continue_annotation, pause_annotation
+)
 
 opts = ('domain', 'tmax', 'io', 'f_init', 'f_end', 'degree',
         'uexpr', 'umax', 'radius', 'tol')
@@ -124,3 +128,29 @@ def _tracer_setup(tmpdir, geometry, blob=False, degree=1, small_dt=False):
                 return tracer_slice(tmpdir, degree, small_dt)
 
     return _tracer_setup
+
+
+@pytest.fixture(scope="module", autouse=True)
+def check_empty_tape(request):
+    """Check that the tape is empty at the end of each module.
+    """
+    def finalizer():
+        # make sure taping is switched off
+        assert not annotate_tape()
+
+        # make sure the tape is empty
+        tape = get_working_tape()
+        if tape is not None:
+            assert len(tape.get_blocks()) == 0
+
+    request.addfinalizer(finalizer)
+
+
+@pytest.fixture
+def set_test_tape():
+    """Set a new working tape specifically for this test.
+    """
+    continue_annotation()
+    with set_working_tape():
+        yield
+    pause_annotation()
diff --git a/integration-tests/data/linear_sw_wave_rexi_chkpt.h5 b/integration-tests/data/linear_sw_wave_rexi_chkpt.h5
deleted file mode 100644
index da948670c..000000000
Binary files a/integration-tests/data/linear_sw_wave_rexi_chkpt.h5 and /dev/null differ