No tape-able operations at module level in adjoint tests (#4867)

Author: JHopeCollins

tests/firedrake/adjoint/test_assemble.py

@@ -6,28 +6,16 @@
 from firedrake.adjoint import *
 
 
+@pytest.fixture(autouse=True)
+def autouse_set_test_tape(set_test_tape):
+    pass
+
+
 @pytest.fixture
 def rg():
     return RandomGenerator(PCG64(seed=1234))
 
 
-@pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
-
-
 @pytest.mark.skipcomplex
 def test_assemble_0_forms():
     mesh = IntervalMesh(10, 0, 1)

tests/firedrake/adjoint/test_assignment.py

@@ -7,26 +7,14 @@
 from numpy.testing import assert_approx_equal, assert_allclose
 
 
-@pytest.fixture
-def rg():
-    return RandomGenerator(PCG64(seed=1234))
-
-
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
+@pytest.fixture
+def rg():
+    return RandomGenerator(PCG64(seed=1234))
 
 
 @pytest.mark.skipcomplex

tests/firedrake/adjoint/test_burgers_newton.py

@@ -13,20 +13,8 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @pytest.fixture
@@ -119,7 +107,7 @@ def J(ic, nu, solve_type, timestep, total_steps, V, nu_time_dependent=False):
     # The comment below and the others like it are used to generate the
     # documentation for the firedrake/docs/source/chekpointing.rst file.
     # [test_disk_checkpointing 10]
-    for step in tape.timestepper(range(total_steps)):
+    for step in tape.timestepper(iter(range(total_steps))):
         # Advance the forward model
         # [test_disk_checkpointing 11]
         if nu_time_dependent and step > 4:

tests/firedrake/adjoint/test_checkpointing_multistep.py

@@ -10,29 +10,20 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 total_steps = 20
 dt = 0.01
-mesh = UnitIntervalMesh(1)
-V = FunctionSpace(mesh, "DG", 0)
 
 
-def J(displacement_0):
+@pytest.fixture
+def V():
+    return FunctionSpace(UnitIntervalMesh(1), "DG", 0)
+
+
+def J(displacement_0, V):
     stiff = Constant(2.5)
     damping = Constant(0.3)
     rho = Constant(1.0)
@@ -59,12 +50,12 @@ def J(displacement_0):
 
 
 @pytest.mark.skipcomplex
-def test_multisteps():
+def test_multisteps(V):
     tape = get_working_tape()
     tape.progress_bar = ProgressBar
     tape.enable_checkpointing(MixedCheckpointSchedule(total_steps, 2, storage=StorageType.RAM))
     displacement_0 = Function(V).assign(1.0)
-    val = J(displacement_0)
+    val = J(displacement_0, V)
     _check_forward(tape)
     c = Control(displacement_0)
     J_hat = ReducedFunctional(val, c)
@@ -82,20 +73,20 @@ def test_multisteps():
 
 
 @pytest.mark.skipcomplex
-def test_validity():
+def test_validity(V):
     tape = get_working_tape()
     tape.progress_bar = ProgressBar
     displacement_0 = Function(V).assign(1.0)
     # Without checkpointing.
-    val0 = J(displacement_0)
+    val0 = J(displacement_0, V)
     J_hat0 = ReducedFunctional(val0, Control(displacement_0))
     dJ0 = J_hat0.derivative()
-    val_recomputed0 = J(displacement_0)
+    val_recomputed0 = J(displacement_0, V)
     tape.clear_tape()
 
     # With checkpointing.
     tape.enable_checkpointing(MixedCheckpointSchedule(total_steps, 2, storage=StorageType.RAM))
-    val = J(displacement_0)
+    val = J(displacement_0, V)
     J_hat = ReducedFunctional(val, Control(displacement_0))
     dJ = J_hat.derivative()
     val_recomputed = J_hat(displacement_0)

tests/firedrake/adjoint/test_disk_checkpointing.py

@@ -9,22 +9,8 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
+def autouse_test_taping(set_test_tape):
     yield
-    tape = get_working_tape()
-    tape.clear_tape()
-    tape._package_data = {}
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
-
     if disk_checkpointing():
         pause_disk_checkpointing()
 

tests/firedrake/adjoint/test_dynamic_meshes.py

@@ -6,25 +6,13 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @pytest.mark.skipcomplex
-@pytest.mark.parametrize("mesh", [UnitSquareMesh(10, 10)])
-def test_dynamic_meshes_2D(mesh):
+def test_dynamic_meshes_2D():
+    mesh = UnitSquareMesh(10, 10)
     S = mesh.coordinates.function_space()
     s = [Function(S), Function(S), Function(S)]
     mesh.coordinates.assign(mesh.coordinates + s[0])
@@ -71,13 +59,26 @@ def test_dynamic_meshes_2D(mesh):
 
 
 @pytest.mark.skipcomplex
-@pytest.mark.parametrize("mesh", [UnitCubeMesh(4, 4, 5),
-                                  UnitOctahedralSphereMesh(3),
-                                  UnitIcosahedralSphereMesh(3),
-                                  UnitCubedSphereMesh(3),
-                                  TorusMesh(25, 10, 1, 0.5),
-                                  CylinderMesh(10, 25, radius=0.5, depth=0.8)])
-def test_dynamic_meshes_3D(mesh):
+@pytest.mark.parametrize("mesh_type", ["UnitCubeMesh",
+                                       "UnitOctahedralSphereMesh",
+                                       "UnitIcosahedralSphereMesh",
+                                       "UnitCubedSphereMesh",
+                                       "TorusMesh",
+                                       "CylinderMesh"])
+def test_dynamic_meshes_3D(mesh_type):
+    if mesh_type == "UnitCubeMesh":
+        mesh = UnitCubeMesh(4, 4, 5)
+    if mesh_type == "UnitOctahedralSphereMesh":
+        mesh = UnitOctahedralSphereMesh(3)
+    if mesh_type == "UnitIcosahedralSphereMesh":
+        mesh = UnitIcosahedralSphereMesh(3)
+    if mesh_type == "UnitCubedSphereMesh":
+        mesh = UnitCubedSphereMesh(3)
+    if mesh_type == "TorusMesh":
+        mesh = TorusMesh(25, 10, 1, 0.5)
+    if mesh_type == "CylinderMesh":
+        mesh = CylinderMesh(10, 25, radius=0.5, depth=0.8)
+
     S = mesh.coordinates.function_space()
     s = [Function(S), Function(S), Function(S)]
     mesh.coordinates.assign(mesh.coordinates + s[0])

tests/firedrake/adjoint/test_ensemble_reduced_functional.py

@@ -6,20 +6,8 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @pytest.mark.parallel(nprocs=4)

tests/firedrake/adjoint/test_external_modification.py

@@ -6,20 +6,8 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @pytest.mark.skipcomplex

tests/firedrake/adjoint/test_hessian.py

@@ -3,32 +3,17 @@
 from firedrake import *
 from firedrake.adjoint import *
 
-from numpy.random import default_rng
-rng = default_rng()
+
+@pytest.fixture(autouse=True)
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @pytest.fixture
 def rg():
     return RandomGenerator(PCG64(seed=1234))
 
 
-@pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
-
-
 @pytest.mark.skipcomplex
 def test_simple_solve(rg):
     tape = Tape()

tests/firedrake/adjoint/test_optimisation.py

@@ -12,20 +12,8 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    if annotate_tape():
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @pytest.mark.skipcomplex