From 42c0cc702eff4450eb23dfe98af738ac86a78255 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 8 Jan 2024 11:31:03 -0500
Subject: [PATCH 01/10] Initial add checkpoint and maxiter feature

---
 src/aspire/reconstruction/estimator.py | 53 +++++++++++++++++++++++++-
 src/aspire/reconstruction/mean.py      | 27 +++++++++++--
 2 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index dec6a901da..2d65df5fff 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -1,4 +1,6 @@
 import logging
+import os
+from pathlib import Path
 
 from aspire.basis import Coef
 from aspire.reconstruction.kernel import FourierKernel
@@ -7,7 +9,16 @@
 
 
 class Estimator:
-    def __init__(self, src, basis, batch_size=512, preconditioner="circulant"):
+    def __init__(
+        self,
+        src,
+        basis,
+        batch_size=512,
+        preconditioner="circulant",
+        checkpoint_iterations=10,
+        checkpoint_prefix="volume_checkpoint",
+        maxiter=100,
+    ):
         """
         An object representing a 2*L-by-2*L-by-2*L array containing the non-centered Fourier transform of the mean
         least-squares estimator kernel.
@@ -15,6 +26,24 @@ def __init__(self, src, basis, batch_size=512, preconditioner="circulant"):
         projection directions (with the appropriate amplitude multipliers and CTFs) and averaging over the whole
         dataset.
         Note that this is a non-centered Fourier transform, so the zero frequency is found at index 1.
+
+
+        :param src: `ImageSource` to be used for estimation.
+        :param basis: 3D Basis to be used during estimation.
+        :param batch_size: Optional batch size of images drawn from
+            `src` during back projection and kernel estimation steps.
+        :param preconditioner: Optional kernel preconditioner (`string`).
+        :param checkpoint_iterations: Optionally save `cg` estimated `Volume`
+            instance periodically each `checkpoint_iterations`.
+            Setting to None disables, otherwise checks for positive integer.
+        :param checkpoint_prefix: Optional path prefix for `cg`
+            checkpoint files.  If the parent directory does not exist,
+            creation is attempted.  `_iter{N}` will be appended to the
+            prefix.
+        :param maxiter: Optional max number of `cg` iterations
+            before returning.  This should be used in conjunction with
+            `checkpoint_iterations` to prevent excessive disk usage.
+            `None` disables.
         """
 
         self.src = src
@@ -23,6 +52,7 @@ def __init__(self, src, basis, batch_size=512, preconditioner="circulant"):
         self.batch_size = batch_size
         self.preconditioner = preconditioner
 
+        # dtype configuration
         if not self.dtype == self.basis.dtype:
             logger.warning(
                 f"Inconsistent types in {self.dtype} Estimator."
@@ -35,6 +65,27 @@ def __init__(self, src, basis, batch_size=512, preconditioner="circulant"):
                 f" Given src.L={src.L} != {basis.nres}"
             )
 
+        # Checkpoint configuration
+        if checkpoint_iterations is not None:
+            checkpoint_iterations = int(checkpoint_iterations)
+            if not checkpoint_iterations > 0:
+                raise ValueError(
+                    "`checkpoint_iterations` should be a positive integer or `None`."
+                )
+        self.checkpoint_iterations = checkpoint_iterations
+        # Create checkpointing dirs as needed
+        parent = Path(checkpoint_prefix).parent
+        if not os.path.exists(parent):
+            os.mkdirs(parents)
+        self.checkpoint_prefix = checkpoint_prefix
+
+        # Maximum iteration configuration
+        if maxiter is not None:
+            maxiter = int(maxiter)
+            if not maxiter > 0:
+                raise ValueError("`maxiter` should be a positive integer or `None`.")
+        self.maxiter = maxiter
+
     def __getattr__(self, name):
         """Lazy attributes instantiated on first-access"""
 
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index fda258cecd..e7f4e7b42d 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -200,12 +200,33 @@ def conj_grad(self, b_coef, tol=1e-5, regularizer=0):
         tol = tol or config.mean.cg_tol
         target_residual = tol * norm(b_coef)
 
-        def cb(xk):
+        # callback setup
+        i = 0  # iteration counter
+
+        def cb(xk, i=i):
+            i += 1  # increment iteration count
+
             logger.info(
-                f"Delta {norm(b_coef - self.apply_kernel(xk))} (target {target_residual})"
+                f"[Iter {i}]: Delta {norm(b_coef - self.apply_kernel(xk))} (target {target_residual})"
             )
 
-        x, info = cg(operator, b_coef.flatten(), M=M, callback=cb, tol=tol, atol=0)
+            # Optional checkpoint
+            if self.checkpoint_iterations and (i % self.checkpoint_iterations) == 0:
+                # Construct checkpoint path
+                path = f"{self.checkpoint_prefix}_iter{i}"
+                # Write out the current solution
+                np.save(path, xk.reshape(self.r, self.basis.count))
+                logger.info("Checkpoint saved to `{path}`")
+
+        x, info = cg(
+            operator,
+            b_coef.flatten(),
+            M=M,
+            callback=cb,
+            tol=tol,
+            atol=0,
+            maxiter=self.maxiter,
+        )
 
         if info != 0:
             raise RuntimeError("Unable to converge!")

From 5664823cb1d840b4d7b1c10d4c07759f7aa44c1f Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 8 Jan 2024 13:31:13 -0500
Subject: [PATCH 02/10] initial test add

---
 tests/test_mean_estimator.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index 2c7563dbb8..367c37e97a 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -1,4 +1,5 @@
 import os.path
+import tempfile
 from unittest import TestCase
 
 import numpy as np
@@ -26,12 +27,12 @@ def setUp(self):
             ],
             dtype=self.dtype,
         )
-        basis = FBBasis3D((self.resolution,) * 3, dtype=self.dtype)
+        self.basis = FBBasis3D((self.resolution,) * 3, dtype=self.dtype)
 
-        self.estimator = MeanEstimator(self.sim, basis, preconditioner="none")
+        self.estimator = MeanEstimator(self.sim, self.basis, preconditioner="none")
 
         self.estimator_with_preconditioner = MeanEstimator(
-            self.sim, basis, preconditioner="circulant"
+            self.sim, self.basis, preconditioner="circulant"
         )
 
     def tearDown(self):
@@ -675,3 +676,17 @@ def testOptimize2(self):
                 atol=1e-4,
             )
         )
+
+    def testCheckpoint(self):
+        """Exercise the checkpointing and max iterations branches."""
+        with tempfile.TemporaryDirectory() as tmp_input_dir:
+            prefix = os.path.join(tmp_input_dir, "chk")
+            estimator = MeanEstimator(self.sim, self.basis, preconditioner="none",
+                                      checkpoint_iterations=2, maxiter=2, checkpoint_prefix=prefix)
+
+            _ = self.estimator.estimate()
+            # Load the checkpoint coefficients
+            b_chk = np.load(f"{prefix}_iter2")
+
+            # Estimate, starting from checkpoint
+            estimate = self.estimator.estimate(b_coef=b_chk)

From 2a3317267944e57966cb39755f884dc3ccc853cf Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 8 Jan 2024 13:55:53 -0500
Subject: [PATCH 03/10] Fix bugs and implement test for checkpointing.

---
 src/aspire/reconstruction/estimator.py |  2 +-
 src/aspire/reconstruction/mean.py      | 17 ++++++++++-------
 tests/test_mean_estimator.py           | 26 ++++++++++++++++++--------
 3 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index 2d65df5fff..08e64d5e74 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -76,7 +76,7 @@ def __init__(
         # Create checkpointing dirs as needed
         parent = Path(checkpoint_prefix).parent
         if not os.path.exists(parent):
-            os.mkdirs(parents)
+            os.makedirs(parent)
         self.checkpoint_prefix = checkpoint_prefix
 
         # Maximum iteration configuration
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index e7f4e7b42d..e7bce3fff9 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -201,22 +201,25 @@ def conj_grad(self, b_coef, tol=1e-5, regularizer=0):
         target_residual = tol * norm(b_coef)
 
         # callback setup
-        i = 0  # iteration counter
+        self.i = 0  # iteration counter
 
-        def cb(xk, i=i):
-            i += 1  # increment iteration count
+        def cb(xk):
+            self.i += 1  # increment iteration count
 
             logger.info(
-                f"[Iter {i}]: Delta {norm(b_coef - self.apply_kernel(xk))} (target {target_residual})"
+                f"[Iter {self.i}]: Delta {norm(b_coef - self.apply_kernel(xk))} (target {target_residual})"
             )
 
             # Optional checkpoint
-            if self.checkpoint_iterations and (i % self.checkpoint_iterations) == 0:
+            if (
+                self.checkpoint_iterations
+                and (self.i % self.checkpoint_iterations) == 0
+            ):
                 # Construct checkpoint path
-                path = f"{self.checkpoint_prefix}_iter{i}"
+                path = f"{self.checkpoint_prefix}_iter{self.i}.npy"
                 # Write out the current solution
                 np.save(path, xk.reshape(self.r, self.basis.count))
-                logger.info("Checkpoint saved to `{path}`")
+                logger.info(f"Checkpoint saved to `{path}`")
 
         x, info = cg(
             operator,
diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index 367c37e97a..1a64356109 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -679,14 +679,24 @@ def testOptimize2(self):
 
     def testCheckpoint(self):
         """Exercise the checkpointing and max iterations branches."""
+        test_iter = 2
         with tempfile.TemporaryDirectory() as tmp_input_dir:
-            prefix = os.path.join(tmp_input_dir, "chk")
-            estimator = MeanEstimator(self.sim, self.basis, preconditioner="none",
-                                      checkpoint_iterations=2, maxiter=2, checkpoint_prefix=prefix)
+            prefix = os.path.join(tmp_input_dir, "new", "dirs", "chk")
+            estimator = MeanEstimator(
+                self.sim,
+                self.basis,
+                preconditioner="none",
+                checkpoint_iterations=test_iter,
+                maxiter=test_iter + 1,
+                checkpoint_prefix=prefix,
+            )
+
+            # Assert we raise when reading `maxiter`.
+            with raises(RuntimeError, match="Unable to converge!"):
+                _ = estimator.estimate()
 
-            _ = self.estimator.estimate()
-            # Load the checkpoint coefficients
-            b_chk = np.load(f"{prefix}_iter2")
+            # Load the checkpoint coefficients while tmp_input_dir exists.
+            b_chk = np.load(f"{prefix}_iter{test_iter}.npy")
 
-            # Estimate, starting from checkpoint
-            estimate = self.estimator.estimate(b_coef=b_chk)
+        # Estimate, starting from checkpoint
+        estimate = self.estimator.estimate(b_coef=b_chk)

From ea3e5cf839f9de6f7b989a30c9c0e7cc5cdae69e Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 8 Jan 2024 15:15:13 -0500
Subject: [PATCH 04/10] Fixup tox check

---
 src/aspire/reconstruction/estimator.py | 1 -
 tests/test_mean_estimator.py           | 4 ++--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index 08e64d5e74..b3b8eed7f8 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -27,7 +27,6 @@ def __init__(
         dataset.
         Note that this is a non-centered Fourier transform, so the zero frequency is found at index 1.
 
-
         :param src: `ImageSource` to be used for estimation.
         :param basis: 3D Basis to be used during estimation.
         :param batch_size: Optional batch size of images drawn from
diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index 1a64356109..fb1020162e 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -698,5 +698,5 @@ def testCheckpoint(self):
             # Load the checkpoint coefficients while tmp_input_dir exists.
             b_chk = np.load(f"{prefix}_iter{test_iter}.npy")
 
-        # Estimate, starting from checkpoint
-        estimate = self.estimator.estimate(b_coef=b_chk)
+        # Restart estimate from checkpoint
+        _ = self.estimator.estimate(b_coef=b_chk)

From e581adb982b8f4222465b834872520eafb328958 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 9 Jan 2024 10:20:53 -0500
Subject: [PATCH 05/10] Cover new checkpoint arg handling lines.

---
 src/aspire/reconstruction/estimator.py | 14 +++++++++++--
 tests/test_mean_estimator.py           | 28 ++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index b3b8eed7f8..e3d39037f7 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -66,12 +66,18 @@ def __init__(
 
         # Checkpoint configuration
         if checkpoint_iterations is not None:
-            checkpoint_iterations = int(checkpoint_iterations)
+            try:
+                checkpoint_iterations = int(checkpoint_iterations)
+            except ValueError:
+                # Sentinel value to emit a more descriptive message below.
+                checkpoint_iterations = -1
+
             if not checkpoint_iterations > 0:
                 raise ValueError(
                     "`checkpoint_iterations` should be a positive integer or `None`."
                 )
         self.checkpoint_iterations = checkpoint_iterations
+
         # Create checkpointing dirs as needed
         parent = Path(checkpoint_prefix).parent
         if not os.path.exists(parent):
@@ -80,7 +86,11 @@ def __init__(
 
         # Maximum iteration configuration
         if maxiter is not None:
-            maxiter = int(maxiter)
+            try:
+                maxiter = int(maxiter)
+            except ValueError:
+                # Sentinel value to emit a more descriptive message below.
+                maxiter = -1
             if not maxiter > 0:
                 raise ValueError("`maxiter` should be a positive integer or `None`.")
         self.maxiter = maxiter
diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index fb1020162e..557b7590a4 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -700,3 +700,31 @@ def testCheckpoint(self):
 
         # Restart estimate from checkpoint
         _ = self.estimator.estimate(b_coef=b_chk)
+
+    def testCheckpointArgs(self):
+        with tempfile.TemporaryDirectory() as tmp_input_dir:
+            prefix = os.path.join(tmp_input_dir, "chk")
+
+            for junk in [-1, 0, "abc"]:
+                # Junk `checkpoint_iterations` values
+                with raises(
+                    ValueError, match=r".*iterations.*should be a positive integer.*"
+                ):
+                    _ = MeanEstimator(
+                        self.sim,
+                        self.basis,
+                        preconditioner="none",
+                        checkpoint_iterations=junk,
+                        checkpoint_prefix=prefix,
+                    )
+                # Junk `maxiter` values
+                with raises(
+                    ValueError, match=r".*maxiter.*should be a positive integer.*"
+                ):
+                    _ = MeanEstimator(
+                        self.sim,
+                        self.basis,
+                        preconditioner="none",
+                        maxiter=-1,
+                        checkpoint_prefix=prefix,
+                    )

From 95da0ca10a35cd5437d63a25e01da0d607cdfeb0 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 26 Jan 2024 10:15:30 -0500
Subject: [PATCH 06/10] use the junk values in test

---
 tests/test_mean_estimator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index 557b7590a4..46b4690fd1 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -725,6 +725,6 @@ def testCheckpointArgs(self):
                         self.sim,
                         self.basis,
                         preconditioner="none",
-                        maxiter=-1,
+                        maxiter=junk,
                         checkpoint_prefix=prefix,
                     )

From c78e5473f2156968a03b5bcccc6bfd202328452a Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 29 Jan 2024 11:06:09 -0500
Subject: [PATCH 07/10] Fix bug when prefix path is None

---
 src/aspire/reconstruction/estimator.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index e3d39037f7..68e3b80186 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -79,9 +79,10 @@ def __init__(
         self.checkpoint_iterations = checkpoint_iterations
 
         # Create checkpointing dirs as needed
-        parent = Path(checkpoint_prefix).parent
-        if not os.path.exists(parent):
-            os.makedirs(parent)
+        if checkpoint_prefix:
+            parent = Path(checkpoint_prefix).parent
+            if not os.path.exists(parent):
+                os.makedirs(parent)
         self.checkpoint_prefix = checkpoint_prefix
 
         # Maximum iteration configuration

From 058c508d2cdfe69e4fdcb191c46fa12dfbad1750 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 29 Jan 2024 11:28:14 -0500
Subject: [PATCH 08/10] Save recon attempt on iteration before maxiter

---
 src/aspire/reconstruction/mean.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index e7bce3fff9..0bce470856 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -210,11 +210,17 @@ def cb(xk):
                 f"[Iter {self.i}]: Delta {norm(b_coef - self.apply_kernel(xk))} (target {target_residual})"
             )
 
-            # Optional checkpoint
-            if (
+            # Do checkpoint at `checkpoint_iterations`,
+            _do_checkpoint = (
                 self.checkpoint_iterations
                 and (self.i % self.checkpoint_iterations) == 0
-            ):
+            )
+            # or the last iteration when `maxiter` provided.
+            if self.maxiter:
+                _do_checkpoint |= self.i == (self.maxiter - 1)
+
+            # Optional checkpoint
+            if _do_checkpoint:
                 # Construct checkpoint path
                 path = f"{self.checkpoint_prefix}_iter{self.i}.npy"
                 # Write out the current solution

From 2f841b12f4f194082baf357a740451fd444a17f9 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 7 Feb 2024 14:12:00 -0500
Subject: [PATCH 09/10] Add support preconditioner options to docstring

---
 src/aspire/reconstruction/estimator.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index 68e3b80186..9e0e561229 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -32,6 +32,7 @@ def __init__(
         :param batch_size: Optional batch size of images drawn from
             `src` during back projection and kernel estimation steps.
         :param preconditioner: Optional kernel preconditioner (`string`).
+            Currently supported options are "circulant" or None.
         :param checkpoint_iterations: Optionally save `cg` estimated `Volume`
             instance periodically each `checkpoint_iterations`.
             Setting to None disables, otherwise checks for positive integer.

From b69fe38561b1fc35c121b29d2b1d3a52a02ac276 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 7 Feb 2024 14:19:21 -0500
Subject: [PATCH 10/10] Add zero prefix checkpoint iteration

---
 src/aspire/reconstruction/mean.py | 2 +-
 tests/test_mean_estimator.py      | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 0bce470856..d25d914276 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -222,7 +222,7 @@ def cb(xk):
             # Optional checkpoint
             if _do_checkpoint:
                 # Construct checkpoint path
-                path = f"{self.checkpoint_prefix}_iter{self.i}.npy"
+                path = f"{self.checkpoint_prefix}_iter{self.i:04d}.npy"
                 # Write out the current solution
                 np.save(path, xk.reshape(self.r, self.basis.count))
                 logger.info(f"Checkpoint saved to `{path}`")
diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index 46b4690fd1..616c258fcc 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -696,7 +696,7 @@ def testCheckpoint(self):
                 _ = estimator.estimate()
 
             # Load the checkpoint coefficients while tmp_input_dir exists.
-            b_chk = np.load(f"{prefix}_iter{test_iter}.npy")
+            b_chk = np.load(f"{prefix}_iter{test_iter:04d}.npy")
 
         # Restart estimate from checkpoint
         _ = self.estimator.estimate(b_coef=b_chk)