fix: three cadence/update bugs from the PR#1421 review (+2 found in review)

autofit/non_linear/search/abstract_search.py

@@ -1037,6 +1037,82 @@ def output_search_internal(self, search_internal):
             obj=search_internal,
         )
 
+    def _steps_until_full_update(self, iterations_remaining: int) -> int:
+        """
+        How many iterations to run before the next ``perform_update``
+        checkpoint, given how much of the budget is left.
+
+        Searches that chunk their run around ``iterations_per_full_update`` all
+        need this same number, and they all need it as an ``int`` — most feed it
+        to something that ultimately does ``range(...)``. But
+        ``iterations_per_full_update`` is stored as a **float** by
+        ``__init__``, because the packaged default is the inf-like ``1e99``
+        meaning "one chunk, checkpoint only at the end". That float is kept
+        deliberately: ``int(1e99)`` is a 99-digit integer, and writing *that*
+        into every saved ``search.json`` in place of a readable ``1e99``
+        sentinel would be a poor trade for a conversion each caller can do at
+        the point of use.
+
+        So the conversion lives here, once, instead of being re-derived (and
+        occasionally forgotten) per search. Forgetting it is not hypothetical:
+        it crashed the MultiStart gradient search for every user who set a real
+        cadence (PyAutoFit#1420), and left the same latent crash in Emcee and
+        BlackJAX NUTS (PyAutoFit#1422).
+
+        There is no second "no intermediate checkpoint" sentinel: the config
+        default ``1e99`` already means that, and it flows through the ``min``
+        below without a special case. A stored ``0`` is therefore a
+        misconfiguration — most plausibly an HPC override — and is rejected
+        rather than silently reinterpreted as "checkpoint never".
+
+        Parameters
+        ----------
+        iterations_remaining
+            Iterations left in this search's budget. Validated, because the
+            returned chunk can only be a usable positive whole number if this
+            is one — searches store their budget under different names
+            (``n_steps``, ``nsteps``, ``num_samples``, ``maxiter``) and none of
+            them validates it.
+
+        Raises
+        ------
+        ValueError
+            If ``iterations_remaining``, or an explicitly supplied
+            ``iterations_per_full_update``, is not a whole number of at
+            least 1.
+        """
+        self._check_step_count(iterations_remaining, "iterations_remaining")
+        self._check_step_count(
+            self.iterations_per_full_update, "iterations_per_full_update"
+        )
+        return int(min(self.iterations_per_full_update, iterations_remaining))
+
+    def _check_step_count(self, value, name: str):
+        """
+        Reject an iteration count that cannot describe a whole number of
+        iterations. ``float`` values are allowed — the packaged config default
+        ``1e99`` is one — provided they are integral.
+
+        Rejecting rather than clamping is deliberate: a value below 1 truncates
+        to a zero-length chunk, so the enclosing ``while`` loop makes no
+        progress and spins forever re-running ``perform_update``. Clamping it to
+        1 would hide the mistake and silently run a schedule the user never
+        asked for, and a silent hang on a cluster is far more expensive than an
+        error at the first chunk boundary.
+        """
+        try:
+            is_whole = value == int(value)
+        except (TypeError, ValueError, OverflowError):
+            is_whole = False
+
+        if not is_whole or value < 1:
+            raise ValueError(
+                f"{type(self).__name__}: `{name}` must be a whole number of "
+                f"iterations and at least 1, but was {value!r}. A fractional "
+                "or sub-1 value gives a zero-length chunk, which never advances "
+                "the search and would hang it rather than fail."
+            )
+
     @property
     def _updater(self):
         # The cached ``SearchUpdater`` must be invalidated whenever

autofit/non_linear/search/mcmc/blackjax/nuts/search.py

@@ -288,7 +288,9 @@ def run_chunk(rng_key, initial_state, n_steps):
         iterations_remaining = self.num_samples
 
         while iterations_remaining > 0:
-            chunk_n = min(self.iterations_per_full_update, iterations_remaining)
+            # ``run_chunk`` scans ``chunk_n`` steps and the key split below sizes
+            # itself from it, so this must be an ``int`` (PyAutoFit#1422).
+            chunk_n = self._steps_until_full_update(iterations_remaining)
 
             rng_key, sample_key = jax.random.split(rng_key)
             states, infos = run_chunk(sample_key, state, chunk_n)

autofit/non_linear/search/mcmc/emcee/search.py

@@ -200,10 +200,9 @@ def _fit(self, model: AbstractPriorModel, analysis):
             iterations_remaining = self.nsteps
 
         while iterations_remaining > 0:
-            if self.iterations_per_full_update > iterations_remaining:
-                iterations = iterations_remaining
-            else:
-                iterations = self.iterations_per_full_update
+            # ``emcee``'s ``sample`` does ``range(iterations)`` internally and
+            # never casts, so this must be an ``int`` (PyAutoFit#1422).
+            iterations = self._steps_until_full_update(iterations_remaining)
 
             for sample in search_internal.sample(
                 initial_state=state,

autofit/non_linear/search/mcmc/zeus/search.py

@@ -236,10 +236,12 @@ def _fit(self, model: AbstractPriorModel, analysis):
             iterations_remaining = self.nsteps
 
         while iterations_remaining > 0:
-            if self.iterations_per_full_update > iterations_remaining:
-                iterations = iterations_remaining
-            else:
-                iterations = self.iterations_per_full_update
+            # zeus casts internally (``self.nsteps = int(iterations)``), so a
+            # float never crashes it — but ``total_iterations`` below is
+            # incremented by this value, so an uncast float would drift the
+            # bookkeeping away from the samples zeus actually drew
+            # (PyAutoFit#1422).
+            iterations = self._steps_until_full_update(iterations_remaining)
 
             for sample in search_internal.sample(
                 start=state,

autofit/non_linear/search/mle/bfgs/search.py

@@ -168,7 +168,10 @@ def _fit(
         while total_iterations < self.maxiter:
 
             iterations_remaining = self.maxiter - total_iterations
-            iterations = min(self.iterations_per_full_update, iterations_remaining)
+            # SciPy tolerates an integral float ``maxiter``, but a fractional one
+            # would quietly acquire ceiling semantics instead of being rejected
+            # (PyAutoFit#1422).
+            iterations = self._steps_until_full_update(iterations_remaining)
 
             if iterations > 0:
                 options = dict(self.options)

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -156,78 +156,39 @@ def __init__(
 
         self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search")
 
-    def _steps_in_chunk(self, steps_remaining: int) -> int:
+    def _is_final_boundary(self, converged: bool, total_steps: int) -> bool:
         """
-        The number of gradient steps to run before the next ``perform_update``
-        checkpoint boundary, given how much of the ``n_steps`` budget is left.
-
-        ``AbstractSearch.__init__`` stores ``iterations_per_full_update`` as a
-        **float**, because the packaged default is the inf-like ``1e99`` (a
-        single chunk, i.e. checkpoint only at the end). ``range`` requires an
-        ``int``, so the chunk size is cast here at the consumer rather than in
-        the shared float coercion, which every other search relies on.
-
-        Without the cast the loop only survives on the default: ``min(1e99,
-        steps_remaining)`` returns the ``int`` operand, so the float never
-        reaches ``range``. A user-supplied cadence *below* the remaining budget
-        (e.g. ``iterations_per_full_update=50`` with ``n_steps=3000``) does
-        reach it, and raises ``TypeError: 'float' object cannot be interpreted
-        as an integer``.
-
-        Both the cadence and ``n_steps`` are validated as whole numbers of at
-        least one step, rather than being coerced into something plausible. A
-        value below 1 (or a fractional one) truncates to ``range(0)``: no step
-        runs, ``total_steps`` never advances, and the enclosing ``while`` loop
-        spins forever re-running ``perform_update`` — a silent hang that on a
-        cluster burns the whole allocation. Clamping such a value to 1 would
-        hide the mistake instead of reporting it, so it raises here. Once both
-        are validated, ``steps_remaining`` is an ``int`` of at least 1 whenever
-        the loop is entered, so the returned chunk is at least 1 and never
-        overshoots the remaining budget.
-
-        Parameters
-        ----------
-        steps_remaining
-            The number of steps left in the ``n_steps`` budget.
-
-        Raises
-        ------
-        ValueError
-            If ``n_steps``, or an explicitly supplied
-            ``iterations_per_full_update``, is not a whole number of at least 1.
+        Whether the chunk boundary just reached ends the search — either the
+        auto-convergence check fired, or the ``n_steps`` ceiling was hit.
+
+        ``_fit`` must not emit a ``perform_update`` at a terminal boundary:
+        ``start_resume_fit`` performs the final update unconditionally the
+        moment ``_fit`` returns, so one here is duplicated work. Flipping
+        ``during_analysis`` does not avoid that — ``SearchUpdater.update``
+        rebuilds the samples, recomputes the summary and re-runs likelihood
+        profiling on every call regardless of the flag.
+
+        A named predicate rather than an inline expression so the rule can be
+        tested directly; ``_fit`` itself needs jax + optax + a JAX-traceable
+        ``Analysis`` and cannot be driven from the NumPy-only library suite.
         """
-        self._check_step_count(self.n_steps, "n_steps")
-
-        # A falsy cadence means "no intermediate checkpoint": one chunk covering
-        # the whole budget. Only a value the user actually supplied is validated.
-        if self.iterations_per_full_update:
-            self._check_step_count(
-                self.iterations_per_full_update, "iterations_per_full_update"
-            )
-            cadence = self.iterations_per_full_update
-        else:
-            cadence = self.n_steps
+        return bool(converged) or total_steps >= self.n_steps
 
-        return int(min(cadence, steps_remaining))
-
-    def _check_step_count(self, value, name: str):
+    @staticmethod
+    def _stop_reason_on_resume(stop_reason):
         """
-        Reject a step count that cannot describe a whole number of gradient
-        steps. ``float`` values are allowed — the packaged config default
-        ``1e99`` is one — provided they are integral.
+        The ``stop_reason`` a resumed run should start from, given the one
+        restored from its previous ``search_internal``.
+
+        ``"converged"`` survives: the step loop's guard uses it to refuse
+        resuming a converged search into further steps. Every other reason
+        describes a run that is now over — keeping it would stamp a stale
+        ``"max_steps"`` on every intermediate checkpoint of a search that is in
+        fact still running, which is exactly what raising ``n_steps`` (the
+        documented way to extend a budget) does. It is re-derived at each chunk
+        boundary, so it starts clear.
         """
-        try:
-            is_whole = value == int(value)
-        except (TypeError, ValueError, OverflowError):
-            is_whole = False
-
-        if not is_whole or value < 1:
-            raise ValueError(
-                f"{type(self).__name__}: `{name}` must be a whole number of "
-                f"gradient steps and at least 1, but was {value!r}. A fractional "
-                "or sub-1 value gives a zero-length step chunk, which never "
-                "advances the search and would hang it rather than fail."
-            )
+        return stop_reason if stop_reason == "converged" else None
 
     def _fit(
         self,
@@ -348,13 +309,15 @@ def batched_value_and_grad(params):
         # ``resurrect`` is on); seeded independently of the broad-start draw.
         resurrect_rng = np.random.default_rng(1)
 
+        stop_reason = self._stop_reason_on_resume(stop_reason)
+
         # ``n_steps`` is the hard ceiling / max budget; ``stop_reason`` becomes
         # ``"converged"`` if the auto-convergence check stops the search early
         # (this also short-circuits the loop on a resumed, already-converged run).
         while total_steps < self.n_steps and stop_reason != "converged":
 
             steps_remaining = self.n_steps - total_steps
-            iterations = self._steps_in_chunk(steps_remaining)
+            iterations = self._steps_until_full_update(steps_remaining)
 
             # Convergence is assessed every step (``fom_history`` updates every
             # step) rather than only at the ``iterations_per_full_update`` boundary,
@@ -427,16 +390,29 @@ def batched_value_and_grad(params):
             }
             self.paths.save_search_internal(obj=search_internal)
 
-            # A converged (or ceiling-reached) boundary is the final update, so it
-            # runs with ``during_analysis=False``; intermediate boundaries do not.
-            is_final = converged or total_steps >= self.n_steps
-            self.perform_update(
-                model=model,
-                analysis=analysis,
-                during_analysis=not is_final,
-                fitness=fitness,
-                search_internal=search_internal,
-            )
+            # Update only at *intermediate* boundaries. ``start_resume_fit``
+            # performs the final update unconditionally the moment ``_fit``
+            # returns, so any update emitted here at a terminal boundary is pure
+            # duplicated work — and flipping ``during_analysis`` does not avoid
+            # it: ``SearchUpdater.update`` rebuilds the samples, recomputes the
+            # summary and re-runs likelihood profiling on *every* call,
+            # regardless of the flag. The previous ``during_analysis=not
+            # is_final`` was worse still, running the whole final pass twice
+            # including final visualization and latents.
+            #
+            # Nothing is lost by skipping: the checkpoint is
+            # ``save_search_internal`` above, and the ``search_internal`` dict
+            # this loop just built (carrying ``stop_reason``, ``converged`` and
+            # ``fom_history``) is what ``_fit`` returns and what that final
+            # update is computed from.
+            if not self._is_final_boundary(converged=converged, total_steps=total_steps):
+                self.perform_update(
+                    model=model,
+                    analysis=analysis,
+                    during_analysis=True,
+                    fitness=fitness,
+                    search_internal=search_internal,
+                )
 
             if converged:
                 self.logger.info(