Integrate checkpoint_schedules library

.github/workflows/test-base.yml

@@ -29,12 +29,14 @@ jobs:
           sudo apt-get update
           sudo apt-get install flake8 python3-h5py python3-numpy python3-pytest python3-pytest-timeout python3-pytest-xdist python3-scipy python3-sympy
           python3 -m pip install ruff
+          git clone --depth 1 https://github.com/firedrakeproject/checkpoint_schedules.git
       - name: Lint
         run: |
           cd tlm_adjoint
           flake8
           ruff check
       - name: Run tests
         run: |
+          export PYTHONPATH=$PWD/checkpoint_schedules:$PYTHONPATH
           cd tlm_adjoint
           pytest-3 -v -n 2 --timeout=300 --timeout-method=thread

.github/workflows/test-fenics.yml

@@ -29,12 +29,14 @@ jobs:
           sudo apt-get update
           sudo apt-get install flake8 python3-dolfin python3-h5py python3-numpy python3-pytest python3-pytest-timeout python3-pytest-xdist python3-scipy python3-sympy
           python3 -m pip install ruff
+          git clone --depth 1 https://github.com/firedrakeproject/checkpoint_schedules.git
       - name: Lint
         run: |
           cd tlm_adjoint
           flake8
           ruff check
       - name: Run tests
         run: |
+          export PYTHONPATH=$PWD/checkpoint_schedules:$PYTHONPATH
           cd tlm_adjoint
           pytest-3 -v tests/base tests/checkpoint_schedules tests/fenics -n 2 --timeout=300 --timeout-method=thread

.github/workflows/test-firedrake.yml

@@ -31,6 +31,7 @@ jobs:
         run: |
           . /home/firedrake/firedrake/bin/activate
           python3 -m pip install jax jaxlib ruff pytest-timeout pytest-xdist
+          git clone --depth 1 https://github.com/firedrakeproject/checkpoint_schedules.git
       - name: Lint
         run: |
           . /home/firedrake/firedrake/bin/activate
@@ -40,6 +41,7 @@ jobs:
       - name: Run tests
         run: |
           . /home/firedrake/firedrake/bin/activate
+          export PYTHONPATH=$PWD/checkpoint_schedules:$PYTHONPATH
           cd tlm_adjoint
           pytest -v tests/base tests/checkpoint_schedules tests/firedrake -n 2 --timeout=300 --timeout-method=thread
   test-complex:
@@ -60,6 +62,7 @@ jobs:
         run: |
           . /home/firedrake/firedrake/bin/activate
           python3 -m pip install jax jaxlib ruff pytest-timeout pytest-xdist
+          git clone --depth 1 https://github.com/firedrakeproject/checkpoint_schedules.git
       - name: Lint
         run: |
           . /home/firedrake/firedrake/bin/activate
@@ -69,5 +72,6 @@ jobs:
       - name: Run tests
         run: |
           . /home/firedrake/firedrake/bin/activate
+          export PYTHONPATH=$PWD/checkpoint_schedules:$PYTHONPATH
           cd tlm_adjoint
           pytest -v tests/base tests/checkpoint_schedules tests/firedrake -n 2 --timeout=300 --timeout-method=thread

docs/source/dependencies.rst

@@ -8,6 +8,7 @@ tlm_adjoint requires:
 
   - `NumPy <https://numpy.org/>`_
   - `SymPy <https://www.sympy.org>`_
+  - `checkpoint_schedules <https://www.firedrakeproject.org/checkpoint_schedules/>`_
 
 Backend dependencies
 --------------------

tests/checkpoint_schedules/test_validity.py

@@ -11,10 +11,6 @@
 import functools
 import pytest
 
-try:
-    import hrevolve
-except ImportError:
-    hrevolve = None
 try:
     import mpi4py.MPI as MPI
 except ImportError:
@@ -72,10 +68,7 @@ def mixed(n, s):
      (two_level, {"period": 2}),
      (two_level, {"period": 7}),
      (two_level, {"period": 10}),
-     pytest.param(
-         h_revolve, {},
-         marks=pytest.mark.skipif(hrevolve is None,
-                                  reason="H-Revolve not available")),
+     (h_revolve, {}),
      (mixed, {})])
 @pytest.mark.parametrize("n, S", [(1, (0,)),
                                   (2, (1,)),

tests/fenics/test_equations.py

@@ -792,10 +792,8 @@ def adjoint_jacobian_solve(self, adj_x, nl_deps, b):
     x_0, _, adj_x_0, J = forward(y)
     stop_manager()
 
-    assert len(manager()._cp._refs) == 3
-    assert tuple(manager()._cp._refs.keys()) == (var_id(y),
-                                                 var_id(zero),
-                                                 var_id(adj_x_0))
+    assert len(manager()._cp._refs) == 1
+    assert tuple(manager()._cp._refs.keys()) == (var_id(adj_x_0),)
     assert len(manager()._cp._cp) == 0
     if test_adj_ic:
         assert len(manager()._cp._data) == 8
@@ -805,7 +803,7 @@ def adjoint_jacobian_solve(self, adj_x, nl_deps, b):
         assert len(manager()._cp._data) == 9
         assert tuple(map(len, manager()._cp._data.values())) \
             == (0, 0, 0, 1, 0, 0, 2, 0, 0)
-    assert len(manager()._cp._storage) == 5
+    assert len(manager()._cp._storage) == 3
 
     dJdx_0, dJdy = compute_gradient(J, [x_0, y])
     assert var_linf_norm(dJdx_0) == 0.0

tests/fenics/test_models.py

@@ -6,11 +6,6 @@
 import numpy as np
 import pytest
 
-try:
-    import hrevolve
-except ImportError:
-    hrevolve = None
-
 pytestmark = pytest.mark.skipif(
     DEFAULT_COMM.size not in {1, 4},
     reason="tests must be run in serial, or with 4 processes")
@@ -86,10 +81,7 @@ def diffusion_ref():
                      "snaps_in_ram": 2}),
      ("multistage", {"format": "hdf5", "snaps_on_disk": 1,
                      "snaps_in_ram": 2}),
-     pytest.param(
-         "H-Revolve", {"snapshots_on_disk": 1, "snapshots_in_ram": 2},
-         marks=pytest.mark.skipif(hrevolve is None,
-                                  reason="H-Revolve not available")),
+     ("H-Revolve", {"snapshots_on_disk": 1, "snapshots_in_ram": 2}),
      ("mixed", {"snapshots": 2, "storage": "disk"})])
 @seed_test
 def test_oscillator(setup_test, test_leaks,

tests/firedrake/test_equations.py

@@ -1043,10 +1043,8 @@ def adjoint_jacobian_solve(self, adj_x, nl_deps, b):
     x_0, _, adj_x_0, J = forward(y)
     stop_manager()
 
-    assert len(manager()._cp._refs) == 3
-    assert tuple(manager()._cp._refs.keys()) == (var_id(y),
-                                                 var_id(zero),
-                                                 var_id(adj_x_0))
+    assert len(manager()._cp._refs) == 1
+    assert tuple(manager()._cp._refs.keys()) == (var_id(adj_x_0),)
     assert len(manager()._cp._cp) == 0
     if test_adj_ic:
         assert len(manager()._cp._data) == 9
@@ -1056,7 +1054,7 @@ def adjoint_jacobian_solve(self, adj_x, nl_deps, b):
         assert len(manager()._cp._data) == 10
         assert tuple(map(len, manager()._cp._data.values())) \
             == (0, 0, 0, 0, 1, 0, 0, 2, 0, 0)
-    assert len(manager()._cp._storage) == 5
+    assert len(manager()._cp._storage) == 3
 
     dJdx_0, dJdy = compute_gradient(J, [x_0, y])
     assert var_linf_norm(dJdx_0) == 0.0

tests/firedrake/test_models.py

@@ -6,11 +6,6 @@
 import numpy as np
 import pytest
 
-try:
-    import hrevolve
-except ImportError:
-    hrevolve = None
-
 pytestmark = pytest.mark.skipif(
     DEFAULT_COMM.size not in {1, 4},
     reason="tests must be run in serial, or with 4 processes")
@@ -84,10 +79,7 @@ def diffusion_ref():
                      "snaps_in_ram": 2}),
      ("multistage", {"format": "hdf5", "snaps_on_disk": 1,
                      "snaps_in_ram": 2}),
-     pytest.param(
-         "H-Revolve", {"snapshots_on_disk": 1, "snapshots_in_ram": 2},
-         marks=pytest.mark.skipif(hrevolve is None,
-                                  reason="H-Revolve not available")),
+     ("H-Revolve", {"snapshots_on_disk": 1, "snapshots_in_ram": 2}),
      ("mixed", {"snapshots": 2, "storage": "disk"})])
 @seed_test
 def test_oscillator(setup_test, test_leaks,