Disk checkpointing

[Ig-dolci]

Am I employing enable_disk_checkpointing correctly in the following code?

from firedrake import *
from firedrake.adjoint import *

continue_annotation()

enable_disk_checkpointing()
n = 30
mesh = UnitIntervalMesh(n)
V = FunctionSpace(mesh, "CG", 2)
end = 0.3
checkpointable_mesh(mesh)

def Dt(u, u_, timestep):
    return (u - u_) / timestep

def J(ic):
    u_ = Function(V)
    u = Function(V)
    v = TestFunction(V)
    u_.assign(ic)
    nu = Constant(0.0001)
    F = (Dt(u, u_, timestep) * v
         + u * u.dx(0) * v + nu * u.dx(0) * v.dx(0)) * dx
    bc = DirichletBC(V, 0.0, "on_boundary")
    t = 0.0
    while t < end:
        solve(F == 0, u, bc)
        u_.assign(u)
        t += float(timestep)

    return assemble(u_ * u_ * dx + ic * ic * dx)


timestep = Constant(1.0 / n)
tape = get_working_tape()
tape.progress_bar = ProgressBar
x, = SpatialCoordinate(mesh)
ic = project(sin(2. * pi * x), V)

val = J(ic)
pause_disk_checkpointing()
J_hat = ReducedFunctional(val, Control(ic))
J_hat(ic)

Why have I the error below?

int/tests/disk_checkpointing.py
firedrake:WARNING OMP_NUM_THREADS is not set or is set to a value greater than 1, we suggest setting OMP_NUM_THREADS=1 to improve performance
Evaluating functional ▣▣▣▣▣▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢▢ 3/30 [0:00:00]
Traceback (most recent call last):
  File "/Users/ddolci/tes_fire_install/firedrake/src/pyadjoint/tests/disk_checkpointing.py", line 41, in <module>
    J_hat(ic)
  File "/Users/ddolci/tes_fire_install/firedrake/src/pyadjoint/pyadjoint/tape.py", line 109, in wrapper
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/pyadjoint/pyadjoint/reduced_functional.py", line 212, in __call__
    blocks[i].recompute()
  File "/Users/ddolci/tes_fire_install/firedrake/src/pyadjoint/pyadjoint/block.py", line 341, in recompute
    inputs = [bv.saved_output for bv in self.get_dependencies()]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/pyadjoint/pyadjoint/block.py", line 341, in <listcomp>
    inputs = [bv.saved_output for bv in self.get_dependencies()]
              ^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/pyadjoint/pyadjoint/block_variable.py", line 56, in saved_output
    return self.output._ad_restore_at_checkpoint(self.checkpoint)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/adjoint_utils/function.py", line 277, in _ad_restore_at_checkpoint
    return checkpoint.restore()
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/adjoint_utils/checkpointing.py", line 298, in restore
    function = infile.load_function(self.mesh, self.stored_name,
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/checkpointing.py", line 1172, in load_function
    V = self._load_function_space(mesh, V_name)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/checkpointing.py", line 1100, in _load_function_space
    tV = self._load_function_space_topology(tmesh, element)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/Users/ddolci/tes_fire_install/firedrake/src/firedrake/firedrake/checkpointing.py", line 1131, in _load_function_space_topology
    sfXC = base_tmesh.sfXC
           ^^^^^^^^^^^^^^^
AttributeError: 'MeshTopology' object has no attribute 'sfXC'. Did you mean: 'sfBC'?

[ksagiyam]

The error message indicates that the mesh is not the loaded one (as currently required due to some restriction of CheckpointFile). checkpiontable_mesh is to return the loaded mesh that can be used in adjoint checkpointing, so you need to do:

mesh = UnitIntervalMesh(n)
mesh = checkpointable_mesh(mesh)
V = FunctionSpace(mesh, "CG", 2)
end = 0.3
...

[Ig-dolci]

Thanks, @ksagiyam !