From 2ddc1b6ca2aa6c0a20eda5f9e26ef83ab60c604f Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 21 Oct 2024 08:46:28 -0300
Subject: [PATCH 01/19] Disk checkpointing managing start, pause and continue
 checkpointing.

---
 firedrake/adjoint_utils/checkpointing.py | 40 ++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index 583d7155eb..91395f2770 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -1,6 +1,7 @@
 """A module providing support for disk checkpointing of the adjoint tape."""
 from pyadjoint import get_working_tape, OverloadedType
 from pyadjoint.tape import TapePackageData
+from pyadjoint.checkpointing import ManageDiskCheckpointing
 from pyop2.mpi import COMM_WORLD
 import tempfile
 import os
@@ -13,7 +14,8 @@
 
 __all__ = ["enable_disk_checkpointing", "disk_checkpointing",
            "pause_disk_checkpointing", "continue_disk_checkpointing",
-           "stop_disk_checkpointing", "checkpointable_mesh"]
+           "stop_disk_checkpointing", "checkpointable_mesh",
+           "DiskCheckpointing"]
 
 
 def current_checkpoint_file(init=None):
@@ -251,7 +253,7 @@ def restore(self):
         pass
 
 
-class CheckpointFunction(CheckpointBase):
+class CheckpointFunction(CheckpointBase, OverloadedType):
     """Metadata for a Function checkpointed to disk.
 
     An object of this class replaces the :class:`~firedrake.Function` stored as
@@ -304,6 +306,9 @@ def restore(self):
         return type(function)(function.function_space(),
                               function.dat, name=self.name(), count=self.count)
 
+    def _ad_restore_at_checkpoint(self, checkpoint):
+        return checkpoint.restore()
+
 
 def maybe_disk_checkpoint(function):
     """Checkpoint a Function to disk if disk checkpointing is active."""
@@ -341,3 +346,34 @@ def _ad_restore_at_checkpoint(self, checkpoint):
             raise ValueError("We must not have output and checkpoint as "
                              "DelegatedFunctionCheckpoint objects.")
         return checkpoint
+
+
+class DiskCheckpointing(ManageDiskCheckpointing):
+    """Manager for the disk checkpointing process.
+
+    Parameters
+    ----------
+    dirname : str
+        The directory in which the disk checkpoints should be stored. If not
+        specified then the current working directory is used. Checkpoints are
+        stored in a temporary subdirectory of this directory.
+    comm : mpi4py.MPI.Intracomm
+        The MPI communicator over which the computation to be disk checkpointed
+        is defined. This will usually match the communicator on which the
+        mesh(es) are defined.
+    cleanup : bool
+        If set to False, checkpoint files will not be deleted when no longer
+        required. This is usually only useful for debugging.
+    """
+
+    def __init__(self, dirname=None, comm=COMM_WORLD, cleanup=True):
+        super().__init__(dirname, comm=comm, cleanup=cleanup)
+
+    def start_checkpointing(self):
+        enable_disk_checkpointing(self.dirname, self.comm, self.cleanup)
+
+    def continue_checkpointing(self):
+        continue_disk_checkpointing()
+
+    def pause_checkpointing(self):
+        pause_disk_checkpointing()
\ No newline at end of file

From 21455505a3237ff7e2c46ed033e1fc0bbb47907a Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 21 Oct 2024 15:24:13 -0300
Subject: [PATCH 02/19] AdjointDiskCheckpoint

---
 firedrake/adjoint/__init__.py            | 2 +-
 firedrake/adjoint_utils/checkpointing.py | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py
index 8373a4285a..cabcb33ebf 100644
--- a/firedrake/adjoint/__init__.py
+++ b/firedrake/adjoint/__init__.py
@@ -24,7 +24,7 @@
 from firedrake.adjoint_utils.checkpointing import \
     enable_disk_checkpointing, pause_disk_checkpointing, \
     continue_disk_checkpointing, stop_disk_checkpointing, \
-    checkpointable_mesh   # noqa F401
+    checkpointable_mesh, AdjointDiskCheckpointing   # noqa F401
 from firedrake.adjoint_utils import get_solve_blocks  # noqa F401
 
 from pyadjoint.verification import taylor_test, taylor_to_dict  # noqa F401
diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index 91395f2770..f60d5c358e 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -15,7 +15,7 @@
 __all__ = ["enable_disk_checkpointing", "disk_checkpointing",
            "pause_disk_checkpointing", "continue_disk_checkpointing",
            "stop_disk_checkpointing", "checkpointable_mesh",
-           "DiskCheckpointing"]
+           "AdjointDiskCheckpointing"]
 
 
 def current_checkpoint_file(init=None):
@@ -348,7 +348,7 @@ def _ad_restore_at_checkpoint(self, checkpoint):
         return checkpoint
 
 
-class DiskCheckpointing(ManageDiskCheckpointing):
+class AdjointDiskCheckpointing(ManageDiskCheckpointing):
     """Manager for the disk checkpointing process.
 
     Parameters
@@ -376,4 +376,4 @@ def continue_checkpointing(self):
         continue_disk_checkpointing()
 
     def pause_checkpointing(self):
-        pause_disk_checkpointing()
\ No newline at end of file
+        pause_disk_checkpointing()

From 2d4743af1c1832c679911ad7473279c65b0feaec Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 22 Oct 2024 15:25:40 -0300
Subject: [PATCH 03/19] Minor change

---
 firedrake/adjoint_utils/checkpointing.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index f60d5c358e..d3f38b36c9 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -307,7 +307,10 @@ def restore(self):
                               function.dat, name=self.name(), count=self.count)
 
     def _ad_restore_at_checkpoint(self, checkpoint):
-        return checkpoint.restore()
+        if isinstance(checkpoint, CheckpointFunction):
+            return checkpoint.restore()
+        else:
+            return checkpoint
 
 
 def maybe_disk_checkpoint(function):

From c197eff70149fda4073aa9d975e21c78c66690ea Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 22 Oct 2024 15:29:37 -0300
Subject: [PATCH 04/19] Testing branch

---
 .github/workflows/build.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 0fff428f43..b21335899e 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -83,6 +83,7 @@ jobs:
             --install defcon \
             --install gadopt \
             --install asQ \
+            --package-branch pyadjoint dolci/disk_checkpointing \
             || (cat firedrake-install.log && /bin/false)
       - name: Install test dependencies
         run: |

From 89535155b33fca8f1080f0bcab7b7e4fa2b9a258 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 22 Oct 2024 16:21:20 -0300
Subject: [PATCH 05/19] Minor change

---
 firedrake/adjoint_utils/checkpointing.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index d3f38b36c9..e6ec1c6816 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -370,7 +370,9 @@ class AdjointDiskCheckpointing(ManageDiskCheckpointing):
     """
 
     def __init__(self, dirname=None, comm=COMM_WORLD, cleanup=True):
-        super().__init__(dirname, comm=comm, cleanup=cleanup)
+        self.dirname = dirname
+        self.comm = comm
+        self.cleanup = cleanup
 
     def start_checkpointing(self):
         enable_disk_checkpointing(self.dirname, self.comm, self.cleanup)

From 91f3e5f0d4f87f1e6783fa398ed0c029078e9c3d Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 12 Nov 2024 21:25:15 +0000
Subject: [PATCH 06/19] testing docs

---
 .github/workflows/docs.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 34086ead44..9047a8af4d 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -66,7 +66,7 @@ jobs:
   deploy:
     name: Deploy Github pages
     needs: build_docs
-    if: ${{ github.ref == 'refs/heads/master' }}
+    if: ${{ github.ref == 'refs/heads/dolci/disk_checkpointing' }}
     permissions:
       pages: write
       id-token: write

From 3af88b5cff722db9485f60de0a9ba8a8dd1dfee6 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 12 Nov 2024 21:36:49 +0000
Subject: [PATCH 07/19] testing docs

---
 .github/workflows/docs.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 9047a8af4d..5e0f2b70bd 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -4,7 +4,7 @@ on:
   # Run on pushes to master
   push:
     branches:
-      - master
+      - dolci/disk_checkpointing
   # And all pull requests
   pull_request:
 
@@ -41,7 +41,7 @@ jobs:
           cd docs
           make validate-bibtex
       - name: Check documentation links
-        if: ${{ github.ref == 'refs/heads/master' }}
+        if: ${{ github.ref == 'refs/heads/dolci/disk_checkpointing' }}
         run: |
           . /home/firedrake/firedrake/bin/activate
           cd docs

From 041ae42690ac6b87da76a88c81d95209100057e8 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 12 Nov 2024 21:45:59 +0000
Subject: [PATCH 08/19] testing docs

---
 .github/workflows/docs.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 5e0f2b70bd..b44223fb78 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -41,7 +41,7 @@ jobs:
           cd docs
           make validate-bibtex
       - name: Check documentation links
-        if: ${{ github.ref == 'refs/heads/dolci/disk_checkpointing' }}
+        if: ${{ github.ref == 'refs/heads/master' }}
         run: |
           . /home/firedrake/firedrake/bin/activate
           cd docs
@@ -66,7 +66,7 @@ jobs:
   deploy:
     name: Deploy Github pages
     needs: build_docs
-    if: ${{ github.ref == 'refs/heads/dolci/disk_checkpointing' }}
+    if: ${{ github.ref == 'refs/heads/master' }}
     permissions:
       pages: write
       id-token: write

From 96060d81726f806bf37b92502fb9af6f46364ac6 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 12 Nov 2024 22:00:01 +0000
Subject: [PATCH 09/19] testing docs

---
 .github/workflows/docs.yml | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index b44223fb78..a7c1e93a5c 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -4,6 +4,7 @@ on:
   # Run on pushes to master
   push:
     branches:
+      - master
       - dolci/disk_checkpointing
   # And all pull requests
   pull_request:
@@ -41,7 +42,7 @@ jobs:
           cd docs
           make validate-bibtex
       - name: Check documentation links
-        if: ${{ github.ref == 'refs/heads/master' }}
+        if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
         run: |
           . /home/firedrake/firedrake/bin/activate
           cd docs
@@ -66,7 +67,7 @@ jobs:
   deploy:
     name: Deploy Github pages
     needs: build_docs
-    if: ${{ github.ref == 'refs/heads/master' }}
+    if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
     permissions:
       pages: write
       id-token: write

From d6f960d4874e8ffe2c305786610472aac7f0eebd Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 12 Nov 2024 22:08:52 +0000
Subject: [PATCH 10/19] dd

---
 .github/workflows/docs.yml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index a7c1e93a5c..34086ead44 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -5,7 +5,6 @@ on:
   push:
     branches:
       - master
-      - dolci/disk_checkpointing
   # And all pull requests
   pull_request:
 
@@ -42,7 +41,7 @@ jobs:
           cd docs
           make validate-bibtex
       - name: Check documentation links
-        if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
+        if: ${{ github.ref == 'refs/heads/master' }}
         run: |
           . /home/firedrake/firedrake/bin/activate
           cd docs
@@ -67,7 +66,7 @@ jobs:
   deploy:
     name: Deploy Github pages
     needs: build_docs
-    if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
+    if: ${{ github.ref == 'refs/heads/master' }}
     permissions:
       pages: write
       id-token: write

From b73712a3d5e5fcf4960eb602f7f043b975ae256a Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 12 Nov 2024 22:31:55 +0000
Subject: [PATCH 11/19] DiskCheckpointingManager

---
 firedrake/adjoint_utils/checkpointing.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index e6ec1c6816..39ec4cc351 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -1,7 +1,7 @@
 """A module providing support for disk checkpointing of the adjoint tape."""
 from pyadjoint import get_working_tape, OverloadedType
 from pyadjoint.tape import TapePackageData
-from pyadjoint.checkpointing import ManageDiskCheckpointing
+from pyadjoint.checkpointing import DiskCheckpointingManager
 from pyop2.mpi import COMM_WORLD
 import tempfile
 import os
@@ -351,7 +351,7 @@ def _ad_restore_at_checkpoint(self, checkpoint):
         return checkpoint
 
 
-class AdjointDiskCheckpointing(ManageDiskCheckpointing):
+class AdjointDiskCheckpointing(DiskCheckpointingManager):
     """Manager for the disk checkpointing process.
 
     Parameters

From 33155dcd142808c896f79781bc004e7dfdad4fb5 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Wed, 13 Nov 2024 08:42:46 +0000
Subject: [PATCH 12/19] testing docs

---
 .github/workflows/docs.yml | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 34086ead44..a7c1e93a5c 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - master
+      - dolci/disk_checkpointing
   # And all pull requests
   pull_request:
 
@@ -41,7 +42,7 @@ jobs:
           cd docs
           make validate-bibtex
       - name: Check documentation links
-        if: ${{ github.ref == 'refs/heads/master' }}
+        if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
         run: |
           . /home/firedrake/firedrake/bin/activate
           cd docs
@@ -66,7 +67,7 @@ jobs:
   deploy:
     name: Deploy Github pages
     needs: build_docs
-    if: ${{ github.ref == 'refs/heads/master' }}
+    if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
     permissions:
       pages: write
       id-token: write

From 1c3f625bc552b9d806b2fba62ef8dbc98d6eafd9 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Thu, 14 Nov 2024 10:06:49 +0000
Subject: [PATCH 13/19] Use package data to manage disk checkpointing

---
 .github/workflows/docs.yml    | 5 ++---
 firedrake/adjoint/__init__.py | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index a7c1e93a5c..34086ead44 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -5,7 +5,6 @@ on:
   push:
     branches:
       - master
-      - dolci/disk_checkpointing
   # And all pull requests
   pull_request:
 
@@ -42,7 +41,7 @@ jobs:
           cd docs
           make validate-bibtex
       - name: Check documentation links
-        if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
+        if: ${{ github.ref == 'refs/heads/master' }}
         run: |
           . /home/firedrake/firedrake/bin/activate
           cd docs
@@ -67,7 +66,7 @@ jobs:
   deploy:
     name: Deploy Github pages
     needs: build_docs
-    if: ${{ github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dolci/disk_checkpointing' }}
+    if: ${{ github.ref == 'refs/heads/master' }}
     permissions:
       pages: write
       id-token: write
diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py
index cabcb33ebf..2c50c88853 100644
--- a/firedrake/adjoint/__init__.py
+++ b/firedrake/adjoint/__init__.py
@@ -24,7 +24,7 @@
 from firedrake.adjoint_utils.checkpointing import \
     enable_disk_checkpointing, pause_disk_checkpointing, \
     continue_disk_checkpointing, stop_disk_checkpointing, \
-    checkpointable_mesh, AdjointDiskCheckpointing   # noqa F401
+    checkpointable_mesh  # noqa F401
 from firedrake.adjoint_utils import get_solve_blocks  # noqa F401
 
 from pyadjoint.verification import taylor_test, taylor_to_dict  # noqa F401

From b1c99d9a07bd966eccc1aa28c8ea7e935c5dd116 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Thu, 14 Nov 2024 10:07:46 +0000
Subject: [PATCH 14/19] dd

---
 firedrake/adjoint_utils/checkpointing.py | 52 ++++++------------------
 1 file changed, 12 insertions(+), 40 deletions(-)

diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index 39ec4cc351..03ea60cf6e 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -1,7 +1,6 @@
 """A module providing support for disk checkpointing of the adjoint tape."""
 from pyadjoint import get_working_tape, OverloadedType
 from pyadjoint.tape import TapePackageData
-from pyadjoint.checkpointing import DiskCheckpointingManager
 from pyop2.mpi import COMM_WORLD
 import tempfile
 import os
@@ -14,8 +13,7 @@
 
 __all__ = ["enable_disk_checkpointing", "disk_checkpointing",
            "pause_disk_checkpointing", "continue_disk_checkpointing",
-           "stop_disk_checkpointing", "checkpointable_mesh",
-           "AdjointDiskCheckpointing"]
+           "stop_disk_checkpointing", "checkpointable_mesh"]
 
 
 def current_checkpoint_file(init=None):
@@ -206,6 +204,16 @@ def restore_from_checkpoint(self, state):
         self.init_checkpoint_file = state["init"]
         self.current_checkpoint_file = state["current"]
 
+    def start_checkpointing(self):
+        enable_disk_checkpointing(self.dirname, self.comm, self.cleanup)
+
+    def continue_checkpointing(self):
+        continue_disk_checkpointing()
+
+    def pause_checkpointing(self):
+        pause_disk_checkpointing()
+
+
 
 def checkpointable_mesh(mesh):
     """Write a mesh to disk and read it back.
@@ -307,10 +315,7 @@ def restore(self):
                               function.dat, name=self.name(), count=self.count)
 
     def _ad_restore_at_checkpoint(self, checkpoint):
-        if isinstance(checkpoint, CheckpointFunction):
-            return checkpoint.restore()
-        else:
-            return checkpoint
+        return checkpoint.restore()
 
 
 def maybe_disk_checkpoint(function):
@@ -349,36 +354,3 @@ def _ad_restore_at_checkpoint(self, checkpoint):
             raise ValueError("We must not have output and checkpoint as "
                              "DelegatedFunctionCheckpoint objects.")
         return checkpoint
-
-
-class AdjointDiskCheckpointing(DiskCheckpointingManager):
-    """Manager for the disk checkpointing process.
-
-    Parameters
-    ----------
-    dirname : str
-        The directory in which the disk checkpoints should be stored. If not
-        specified then the current working directory is used. Checkpoints are
-        stored in a temporary subdirectory of this directory.
-    comm : mpi4py.MPI.Intracomm
-        The MPI communicator over which the computation to be disk checkpointed
-        is defined. This will usually match the communicator on which the
-        mesh(es) are defined.
-    cleanup : bool
-        If set to False, checkpoint files will not be deleted when no longer
-        required. This is usually only useful for debugging.
-    """
-
-    def __init__(self, dirname=None, comm=COMM_WORLD, cleanup=True):
-        self.dirname = dirname
-        self.comm = comm
-        self.cleanup = cleanup
-
-    def start_checkpointing(self):
-        enable_disk_checkpointing(self.dirname, self.comm, self.cleanup)
-
-    def continue_checkpointing(self):
-        continue_disk_checkpointing()
-
-    def pause_checkpointing(self):
-        pause_disk_checkpointing()

From 0b503160c59b3ac6e93fe86d178f82a43f0bcd5e Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Thu, 14 Nov 2024 10:14:41 +0000
Subject: [PATCH 15/19] flake8

---
 firedrake/adjoint_utils/checkpointing.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index 03ea60cf6e..16cdd3f197 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -214,7 +214,6 @@ def pause_checkpointing(self):
         pause_disk_checkpointing()
 
 
-
 def checkpointable_mesh(mesh):
     """Write a mesh to disk and read it back.
 

From 92ac3ecb774925a4bd2eb2c8f0c1976fb1c8bec4 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Wed, 20 Nov 2024 09:26:51 +0000
Subject: [PATCH 16/19] remove start_checkpointing and add a pyadjoint callback
 interface for checkpointinf on disk

---
 firedrake/adjoint/__init__.py            | 3 +++
 firedrake/adjoint_utils/checkpointing.py | 3 ---
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py
index 2c50c88853..f3afa186b6 100644
--- a/firedrake/adjoint/__init__.py
+++ b/firedrake/adjoint/__init__.py
@@ -21,6 +21,7 @@
                             pause_annotation, continue_annotation, \
                             stop_annotating, annotate_tape  # noqa F401
 from pyadjoint.reduced_functional import ReducedFunctional  # noqa F401
+from pyadjoint.checkpointing import disk_checkpointing_callback  # noqa F401
 from firedrake.adjoint_utils.checkpointing import \
     enable_disk_checkpointing, pause_disk_checkpointing, \
     continue_disk_checkpointing, stop_disk_checkpointing, \
@@ -54,3 +55,5 @@ def __call__(self, form):
 sys.modules[__name__].__class__ = _AdjointModule
 
 set_working_tape(Tape())
+disk_checkpointing_callback["error"] = "Please call enable_disk_checkpointing() "\
+    "before checkpointing on the disk."
diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index 16cdd3f197..27ff9fe58c 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -204,9 +204,6 @@ def restore_from_checkpoint(self, state):
         self.init_checkpoint_file = state["init"]
         self.current_checkpoint_file = state["current"]
 
-    def start_checkpointing(self):
-        enable_disk_checkpointing(self.dirname, self.comm, self.cleanup)
-
     def continue_checkpointing(self):
         continue_disk_checkpointing()
 

From 8e253b792f3c7413fbddeda7fdf0849fb06dec1d Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Wed, 20 Nov 2024 10:04:31 +0000
Subject: [PATCH 17/19] callback interface in checkpointing file

---
 firedrake/adjoint/__init__.py            | 2 --
 firedrake/adjoint_utils/checkpointing.py | 4 +++-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py
index f3afa186b6..c48b990420 100644
--- a/firedrake/adjoint/__init__.py
+++ b/firedrake/adjoint/__init__.py
@@ -55,5 +55,3 @@ def __call__(self, form):
 sys.modules[__name__].__class__ = _AdjointModule
 
 set_working_tape(Tape())
-disk_checkpointing_callback["error"] = "Please call enable_disk_checkpointing() "\
-    "before checkpointing on the disk."
diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index 27ff9fe58c..7d22d8692d 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -1,5 +1,5 @@
 """A module providing support for disk checkpointing of the adjoint tape."""
-from pyadjoint import get_working_tape, OverloadedType
+from pyadjoint import get_working_tape, OverloadedType, disk_checkpointing_callback
 from pyadjoint.tape import TapePackageData
 from pyop2.mpi import COMM_WORLD
 import tempfile
@@ -10,6 +10,8 @@
 from numbers import Number
 _enable_disk_checkpoint = False
 _checkpoint_init_data = False
+disk_checkpointing_callback["error"] = "Please call enable_disk_checkpointing() "\
+    "before checkpointing on the disk."
 
 __all__ = ["enable_disk_checkpointing", "disk_checkpointing",
            "pause_disk_checkpointing", "continue_disk_checkpointing",

From 042148511a8079d9653211786ee3b1cbd5909ea0 Mon Sep 17 00:00:00 2001
From: Daiane Iglesia Dolci <63597005+Ig-dolci@users.noreply.github.com>
Date: Thu, 21 Nov 2024 10:18:57 +0000
Subject: [PATCH 18/19] Update firedrake/adjoint_utils/checkpointing.py

Co-authored-by: David A. Ham <david.ham@imperial.ac.uk>
---
 firedrake/adjoint_utils/checkpointing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/checkpointing.py b/firedrake/adjoint_utils/checkpointing.py
index 7d22d8692d..0ab337e17b 100644
--- a/firedrake/adjoint_utils/checkpointing.py
+++ b/firedrake/adjoint_utils/checkpointing.py
@@ -10,7 +10,7 @@
 from numbers import Number
 _enable_disk_checkpoint = False
 _checkpoint_init_data = False
-disk_checkpointing_callback["error"] = "Please call enable_disk_checkpointing() "\
+disk_checkpointing_callback["firedrake"] = "Please call enable_disk_checkpointing() "\
     "before checkpointing on the disk."
 
 __all__ = ["enable_disk_checkpointing", "disk_checkpointing",

From 1bd8a2943bfc2c53fa3adb24ca73ffad8984ddc2 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Thu, 21 Nov 2024 12:59:41 +0000
Subject: [PATCH 19/19] removing pyadjoint package for the tests.

---
 .github/workflows/build.yml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index b21335899e..0fff428f43 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -83,7 +83,6 @@ jobs:
             --install defcon \
             --install gadopt \
             --install asQ \
-            --package-branch pyadjoint dolci/disk_checkpointing \
             || (cat firedrake-install.log && /bin/false)
       - name: Install test dependencies
         run: |