Docs for checkpointing in adjoint simulations (#4094)

* Docs for checkpointing in adjoint simulations

---------

Co-authored-by: David A. Ham <[email protected]>

Author: Ig-dolci

docs/source/checkpointing.rst

@@ -230,24 +230,137 @@ with `idx` parameter always unset, and the same :class:`~.Function` can only be
 loaded using the same mode.
 
 
-Using disk checkpointing in adjoint simulations
+Using checkpointing in adjoint simulations
 ===============================================
 
 When adjoint annotation is active, the result of every Firedrake operation is
-stored in memory. For some simulations, this can result in a very large memory
-footprint. As an alternative, it is possible to specify that those intermediate
-results in forward evaluations of the tape which have type
-:class:`~firedrake.function.Function` be written to disk. This is usually the
-bulk of the data stored on the tape so this largely alleviates the memory
-problem, at the cost of the time taken to read to and write from disk.
+stored in memory. For some time-dependent simulations, this can lead to a
+large memory footprint. To alleviate this, we can use checkpointing strategies to store only some intermediate forward data in memory or on disk.
 
-Having imported `firedrake.adjoint`, there are two steps required to enable
-disk checkpointing of the forward tape state.
+Checkpointing for time-dependent adjoint simulations in Firedrake employs
+schedules, which determine how forward data is stored. These schedules are
+implemented in the `checkpoint_schedules package
+<https://www.firedrakeproject.org/checkpoint_schedules/>`_.
 
-1. Call :func:`~firedrake.adjoint_utils.checkpointing.enable_disk_checkpointing`.
-2. Wrap all mesh constructors in :func:`~firedrake.adjoint_utils.checkpointing.checkpointable_mesh`.
 
-See the documentation of those functions for more detail.
+To store every time step of the forward data required for adjoint-based gradient
+computation **in memory**, first import the schedule from the
+``checkpoint_schedules`` package, start adjoint annotation with ``continue_annotation()``,
+get the working tape with ``get_working_tape()``:
+
+.. code-block:: python3
+
+    from firedrake import *
+    from firedrake.adjoint import *
+    from checkpoint_schedules import SingleMemoryStorageSchedule
+    continue_annotation()
+    tape = get_working_tape()
+
+Define the schedule:
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_burgers_newton.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 4]
+    :end-before: [test_disk_checkpointing 5]
+
+and enable checkpointing:
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_burgers_newton.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 6]
+    :end-before: [test_disk_checkpointing 7]
+
+**For any checkpointing approach, it is essential to call the time loop as
+follows when advancing the solver in time:**
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_burgers_newton.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 10]
+    :end-before: [test_disk_checkpointing 11]
+
+
+``SingleMemoryStorageSchedule`` stores only the adjoint variables from the last adjoint
+time step, which corresponds to the zero forward time step due to the time-reversed nature
+of the adjoint solver.
+
+
+To store every time step of the forward data required for adjoint-based gradient
+computation **on disk**, write the necessary imports and start adjoint annotation:
+
+.. code-block:: python3
+
+    from firedrake import *
+    from firedrake.adjoint import *
+    from checkpoint_schedules import SingleDiskStorageSchedule
+
+    continue_annotation()
+    tape = get_working_tape()
+
+Then, enable disk checkpointing following the code below:
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_disk_checkpointing.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 1]
+    :end-before: [test_disk_checkpointing 2]
+
+For disk checkpointing, all mesh constructors must be wrapped using
+:func:`~.checkpointing.checkpointable_mesh`. For example:
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_disk_checkpointing.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 2]
+    :end-before: [test_disk_checkpointing 3]
+
+``SingleDiskStorageSchedule`` stores only the adjoint variables from the last adjoint
+time step (equivalent to the zero forward time step) in memory. This checkpointing strategy
+reduces memory usage at the cost of reading and writing data from disk.
+
+The ``checkpoint_schedules`` package provides other checkpointing
+strategies, such as Revolve, Mixed Schedule, and HRevolve. These methods
+store only a limited set of time steps, reducing memory usage at the cost of
+increased computational effort due to repeated forward calculations.
+
+For example, to use the **Revolve** schedule:
+
+.. code-block:: python3
+
+    from firedrake import *
+    from firedrake.adjoint import *
+    from checkpoint_schedules import Revolve
+
+    continue_annotation()
+    tape = get_working_tape()
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_burgers_newton.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 8]
+    :end-before: [test_disk_checkpointing 9]
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_burgers_newton.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 6]
+    :end-before: [test_disk_checkpointing 7]
+
+Then, advance the solver in time as follows:
+
+.. literalinclude:: ../../tests/firedrake/adjoint/test_burgers_newton.py
+    :language: python3
+    :dedent:
+    :start-after: [test_disk_checkpointing 10]
+    :end-before: [test_disk_checkpointing 11]
+
+``steps_to_store`` is the number of time steps stored in memory.
+
+For more details on available checkpointing strategies, refer to the
+`checkpoint_schedules package documentation
+<https://www.firedrakeproject.org/checkpoint_schedules/>`_.
 
 
 Checkpointing with DumbCheckpoint

tests/firedrake/adjoint/test_burgers_newton.py

@@ -99,7 +99,7 @@ def _check_reverse(tape):
                     assert out._checkpoint is None
 
 
-def J(ic, nu, solve_type, timestep, steps, V, nu_time_dependent=False):
+def J(ic, nu, solve_type, timestep, total_steps, V, nu_time_dependent=False):
     """Burgers equation solver."""
     u_ = Function(V, name="u_")
     u = Function(V, name="u")
@@ -115,9 +115,15 @@ def J(ic, nu, solve_type, timestep, steps, V, nu_time_dependent=False):
 
     tape = get_working_tape()
     J = 0.0
-    for j in tape.timestepper(range(steps)):
-        if nu_time_dependent and j > 4:
-            nu.assign(nu*(1.0 + j/1000))
+
+    # 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)):
+        # Advance the forward model
+        # [test_disk_checkpointing 11]
+        if nu_time_dependent and step > 4:
+            nu.assign(nu*(1.0 + step/1000))
         if solve_type == "NLVS":
             solver.solve()
         else:
@@ -135,26 +141,33 @@ def test_burgers_newton(solve_type, checkpointing, basics):
     """
     tape = get_working_tape()
     tape.progress_bar = ProgressBar
-    mesh, timestep, steps = basics
+    mesh, timestep, total_steps = basics
     if checkpointing:
+        steps_to_store = total_steps//3
         if checkpointing == "Revolve":
-            schedule = Revolve(steps, steps//3)
+            # [test_disk_checkpointing 8]
+            schedule = Revolve(total_steps, steps_to_store)
+            # [test_disk_checkpointing 9]
         if checkpointing == "SingleMemory":
+            # [test_disk_checkpointing 4]
             schedule = SingleMemoryStorageSchedule()
+            # [test_disk_checkpointing 5]
         if checkpointing == "Mixed":
             enable_disk_checkpointing()
-            schedule = MixedCheckpointSchedule(steps, steps//3, storage=StorageType.DISK)
+            schedule = MixedCheckpointSchedule(total_steps, steps_to_store, storage=StorageType.DISK)
         if checkpointing == "NoneAdjoint":
             schedule = NoneCheckpointSchedule()
+        # [test_disk_checkpointing 6]
         tape.enable_checkpointing(schedule)
+        # [test_disk_checkpointing 7]
 
     if checkpointing and schedule.uses_storage_type(StorageType.DISK):
         mesh = checkpointable_mesh(mesh)
 
     V, ic, nu = setup_test(mesh)
-    val = J(ic, nu, solve_type, timestep, steps, V)
+    val = J(ic, nu, solve_type, timestep, total_steps, V)
     if checkpointing:
-        assert len(tape.timesteps) == steps
+        assert len(tape.timesteps) == total_steps
         if checkpointing == "Revolve" or checkpointing == "Mixed":
             _check_forward(tape)
 

tests/firedrake/adjoint/test_disk_checkpointing.py

@@ -108,9 +108,14 @@ def test_disk_checkpointing_parallel():
     tape = get_working_tape()
     tape.clear_tape()
     continue_annotation()
+    # 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 1]
     enable_disk_checkpointing()
     tape.enable_checkpointing(SingleDiskStorageSchedule())
+    # [test_disk_checkpointing 2]
     mesh = checkpointable_mesh(UnitSquareMesh(10, 10))
+    # [test_disk_checkpointing 3]
     J_disk, grad_J_disk = adjoint_example(mesh)
 
     assert disk_checkpointing() is False