fix: report the real quick-update cadence and JAX compile status

autofit/non_linear/analysis/latent.py

@@ -28,6 +28,7 @@
 import numpy as np
 
 from autofit import exc
+from autofit.non_linear.jax_compile import log_on_first_compile
 from autofit.non_linear.samples.sample import Sample
 from autofit.non_linear.samples.samples import Samples
 from autofit.non_linear.samples.util import simple_model_for_kwargs
@@ -144,18 +145,21 @@ def latent_samples_from(
         if analysis._use_jax:
             import jax
             import jax.numpy as jnp
-            start = time.time()
             if batch_mode == "vmap":
-                logger.info("JAX: Applying vmap and jit to likelihood function for latent variables -- may take a few seconds.")
                 # vmap traces `variables` once for the whole batch, so a
                 # per-sample try/except is not possible here — latent functions
                 # on the vmap path must express failures as NaN (e.g.
                 # `jnp.where`), never by raising. The `jit` and numpy paths
                 # below do guard per sample.
-                batched_compute_latent = jax.jit(jax.vmap(compute_latent_for_model))
+                batched_compute_latent = log_on_first_compile(
+                    jax.jit(jax.vmap(compute_latent_for_model)),
+                    "latent variable function (vmap)",
+                )
             elif batch_mode == "jit":
-                logger.info("JAX: Applying per-sample jit to latent variables (LATENT_BATCH_MODE='jit') -- may take a few seconds on first sample.")
-                jitted_compute_latent = jax.jit(compute_latent_for_model)
+                jitted_compute_latent = log_on_first_compile(
+                    jax.jit(compute_latent_for_model),
+                    "latent variable function (per-sample jit)",
+                )
                 n_latents = len(keys)
                 nan_tuple = tuple(jnp.nan for _ in range(n_latents))
 
@@ -185,7 +189,6 @@ def batched_compute_latent(parameters_batch):
                 raise ValueError(
                     f"Unknown LATENT_BATCH_MODE={batch_mode!r}; expected 'vmap' or 'jit'."
                 )
-            logger.info(f"JAX: {batch_mode} dispatch applied in {time.time() - start} seconds.")
         else:
             n_latents = len(keys)
             nan_row = np.full(n_latents, np.nan)

autofit/non_linear/fitness.py

@@ -16,6 +16,7 @@
 
 
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
+from autofit.non_linear.jax_compile import log_on_first_compile
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.non_linear.analysis import Analysis
 
@@ -425,8 +426,15 @@ def manage_quick_update(self, parameters, log_likelihood):
                                 "Live display update raised an exception (ignored)."
                             )
 
+            # Searches hand their parameters over in whatever type they hold them:
+            # ndarray (Nautilus), JAX array, or a plain Python list (Dynesty's
+            # initializer). `np.asarray` normalizes all three -- calling `.tolist()`
+            # directly assumed the array case, which held only while Nautilus was
+            # the sole search wired up to quick updates (PyAutoFit#1434).
             result_info = text_util.result_max_lh_info_from(
-                max_log_likelihood_sample=self.quick_update_max_lh_parameters.tolist(),
+                max_log_likelihood_sample=np.asarray(
+                    self.quick_update_max_lh_parameters
+                ).tolist(),
                 max_log_likelihood=self.quick_update_max_lh,
                 model=self.model,
             )
@@ -504,11 +512,11 @@ def _vmap(self):
         after its first creation, avoiding repeated JIT compilation overhead.
         """
         import jax
-        start = time.time()
-        logger.info("JAX: Applying vmap and jit to likelihood function -- may take a few seconds.")
-        func = jax.vmap(jax.jit(self.call))
-        logger.info(f"JAX: vmap and jit applied in {time.time() - start} seconds.")
-        return func
+
+        return log_on_first_compile(
+            jax.vmap(jax.jit(self.call)),
+            "vectorized (vmap) likelihood function",
+        )
 
     @cached_property
     def _jit(self):
@@ -524,11 +532,11 @@ def _jit(self):
         first use, so JIT compilation only occurs once.
         """
         import jax
-        start = time.time()
-        logger.info("JAX: Applying jit to likelihood function -- may take a few seconds.")
-        func = jax.jit(self.call)
-        logger.info(f"JAX: jit applied in {time.time() - start} seconds.")
-        return func
+
+        return log_on_first_compile(
+            jax.jit(self.call),
+            "likelihood function",
+        )
 
     @cached_property
     def _grad(self):
@@ -545,11 +553,11 @@ def _grad(self):
         only once.
         """
         import jax
-        start = time.time()
-        logger.info("JAX: Applying grad to likelihood function -- may take a few seconds.")
-        func = jax.grad(self.call)
-        logger.info(f"JAX: grad applied in {time.time() - start} seconds.")
-        return func
+
+        return log_on_first_compile(
+            jax.grad(self.call),
+            "likelihood function gradient",
+        )
 
     def grad(self, *args, **kwargs):
         return self._grad(*args, **kwargs)

autofit/non_linear/jax_compile.py

@@ -0,0 +1,70 @@
+import logging
+import time
+
+logger = logging.getLogger(__name__)
+
+
+def log_on_first_compile(func, description):
+    """
+    Wrap `func` so that its first invocation reports the JAX compilation it
+    triggers.
+
+    `jax.jit`, `jax.vmap` and `jax.grad` all return immediately -- they only
+    build a wrapper. Tracing, lowering and XLA compilation happen on the *first
+    call* to the function they return, and that call is where the user waits:
+    seconds for a small model, minutes for a large one. Logging at
+    wrapper-construction time therefore announced a wait that had not started
+    yet and reported a duration of roughly zero, leaving the real wait that
+    followed unexplained.
+
+    The one-shot flag lives in a closure rather than on an object because the
+    callers cache these wrappers on attributes that are stripped for pickling; a
+    fresh process rebuilds the wrapper and genuinely does recompile, so it
+    should log again.
+
+    Parameters
+    ----------
+    func
+        The JAX-transformed callable whose first call triggers compilation.
+    description
+        What is being compiled, as it should read mid-sentence in the log.
+
+    Returns
+    -------
+    A callable with the same signature as `func`.
+    """
+    state = {"compiled": False}
+
+    def wrapper(*args, **kwargs):
+        if state["compiled"]:
+            return func(*args, **kwargs)
+
+        logger.info(
+            f"JAX jit compiling {description}, could take seconds or minutes..."
+        )
+        start = time.time()
+
+        try:
+            result = func(*args, **kwargs)
+
+            # JAX dispatches asynchronously, so `result` may be a future that is
+            # not yet materialized. Block once, on this first call only, so the
+            # duration logged below is the wait the user actually sat through
+            # rather than the dispatch latency.
+            try:
+                import jax
+
+                jax.block_until_ready(result)
+            except Exception:
+                pass
+        finally:
+            state["compiled"] = True
+
+        logger.info(
+            f"JAX jit compilation of {description} complete in "
+            f"{time.time() - start:.1f} seconds."
+        )
+
+        return result
+
+    return wrapper

autofit/non_linear/search/abstract_search.py

@@ -57,6 +57,14 @@
 
 logger = logging.getLogger(__name__)
 
+#: Iteration cadences at or above this mean "never". The packaged config default
+#: for both ``iterations_per_quick_update`` and ``iterations_per_full_update`` is
+#: the inf-like ``1e99`` sentinel documented on ``_steps_until_full_update`` --
+#: kept as a float so ``search.json`` stores a readable ``1e99`` rather than a
+#: 99-digit integer. Compared against a threshold rather than ``1e99`` exactly so
+#: a hand-set ``1e100`` in a workspace config reads as "never" too.
+ITERATIONS_NEVER = 1e90
+
 
 def check_cores(func):
     """
@@ -444,6 +452,35 @@ def paths(self, paths: Optional[AbstractPaths]):
             paths.search = self
         self._paths = paths
 
+    @property
+    def quick_update_message(self) -> str:
+        """
+        One line, logged at the start of every search, telling the user the real
+        cadence of the on-the-fly maximum-likelihood updates.
+
+        The cadence is worth stating because it is the only thing that explains
+        the terminal's behaviour during a long fit: either updates appear every
+        N iterations, or nothing appears at all until the search finishes. The
+        packaged default is the ``ITERATIONS_NEVER`` sentinel, so "nothing at
+        all" is what most users get -- and a message that reported that as
+        ``1e+99 iterations`` would be technically true and practically useless.
+        Hence the two-branch wording, and hence naming the config key: the
+        disabled branch is the one a user is most likely to want to act on.
+        """
+        iterations = self.iterations_per_quick_update
+
+        if not np.isfinite(iterations) or iterations >= ITERATIONS_NEVER:
+            return (
+                "On-the-fly updates of the maximum likelihood model are disabled. "
+                "Set `updates: iterations_per_quick_update` in config/general.yaml "
+                "to a finite number of iterations to enable them."
+            )
+
+        return (
+            "On-the-fly updates of the maximum likelihood model every "
+            f"{int(iterations)} iterations."
+        )
+
     def copy_with_paths(self, paths):
         self.logger.debug(f"Creating a copy of {self._paths.name}")
         search_instance = copy.copy(self)
@@ -507,6 +544,7 @@ class represented by model M and gives a score for their fitness.
                 logger.info("Starting non-linear search with JAX.")
         else:
             logger.info(f"Starting non-linear search with {self.number_of_cores} cores.")
+        logger.info(self.quick_update_message)
         self._log_process_state()
 
         model = analysis.modify_model(model)

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

@@ -201,6 +201,9 @@ def _fit(self, model: AbstractPriorModel, analysis):
             paths=self.paths,
             fom_is_log_likelihood=False,  # log-posterior target for NUTS
             resample_figure_of_merit=-jnp.inf,
+            iterations_per_quick_update=self.iterations_per_quick_update,
+            background_quick_update=self.quick_update_background,
+            live_visual_update=self.live_visual_update,
         )
 
         # Initial position: borrow the standard initializer machinery so users

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

@@ -136,6 +136,9 @@ def _fit(self, model: AbstractPriorModel, analysis):
             paths=self.paths,
             fom_is_log_likelihood=False,
             resample_figure_of_merit=-np.inf,
+            iterations_per_quick_update=self.iterations_per_quick_update,
+            background_quick_update=self.quick_update_background,
+            live_visual_update=self.live_visual_update,
         )
 
         pool = self.make_sneaky_pool(fitness)

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

@@ -173,6 +173,9 @@ def _fit(self, model: AbstractPriorModel, analysis):
             paths=self.paths,
             fom_is_log_likelihood=False,
             resample_figure_of_merit=-np.inf,
+            iterations_per_quick_update=self.iterations_per_quick_update,
+            background_quick_update=self.quick_update_background,
+            live_visual_update=self.live_visual_update,
         )
 
         try:

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

@@ -124,7 +124,10 @@ def _fit(
             fom_is_log_likelihood=False,
             resample_figure_of_merit=-np.inf,
             convert_to_chi_squared=True,
-            store_history=self.should_plot_start_point
+            store_history=self.should_plot_start_point,
+            iterations_per_quick_update=self.iterations_per_quick_update,
+            background_quick_update=self.quick_update_background,
+            live_visual_update=self.live_visual_update,
         )
 
         try:

autofit/non_linear/search/mle/drawer/search.py

@@ -114,6 +114,9 @@ def _fit(self, model: AbstractPriorModel, analysis):
             fom_is_log_likelihood=False,
             resample_figure_of_merit=-np.inf,
             convert_to_chi_squared=False,
+            iterations_per_quick_update=self.iterations_per_quick_update,
+            background_quick_update=self.quick_update_background,
+            live_visual_update=self.live_visual_update,
         )
 
         total_draws = self.total_draws

autofit/non_linear/search/nest/dynesty/search/abstract.py

@@ -182,6 +182,9 @@ def _fit(
             paths=self.paths,
             fom_is_log_likelihood=True,
             resample_figure_of_merit=-1.0e99,
+            iterations_per_quick_update=self.iterations_per_quick_update,
+            background_quick_update=self.quick_update_background,
+            live_visual_update=self.live_visual_update,
             use_jax_jit=getattr(analysis, "_use_jax", False) and self.use_jax_jit,
         )
 

test_autofit/non_linear/search/test_quick_update_message.py

@@ -0,0 +1,70 @@
+import pytest
+
+import autofit as af
+
+# ``quick_update_message`` is logged once at the start of every search. It is the
+# only thing that explains the terminal's behaviour during a long fit: either
+# on-the-fly maximum-likelihood updates appear every N iterations, or nothing
+# appears until the search finishes.
+#
+# It exists because the workspace example scripts used to print the *name* of the
+# knob rather than its value ("On-the-fly updates every
+# iterations_per_quick_update are printed to the notebook"). Interpolating the
+# value alone would not have fixed it: the packaged default is the inf-like
+# ``1e99`` sentinel, so the honest message for most users is not a number at all.
+#
+# NumPy-only, like the rest of the library suite.
+
+pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning")
+
+
+def test__real_cadence_renders_as_a_plain_integer():
+    # The knob is stored as a float (so ``search.json`` keeps a readable
+    # ``1e99`` rather than a 99-digit integer), so the message has to cast --
+    # otherwise a user configuring 250000 is told "every 250000.0 iterations".
+    search = af.MultiStartAdam(n_steps=3000, iterations_per_quick_update=250000)
+
+    assert isinstance(search.iterations_per_quick_update, float)
+
+    message = search.quick_update_message
+
+    assert "every 250000 iterations" in message
+    assert "250000.0" not in message
+    assert "disabled" not in message
+
+
+def test__default_sentinel_says_disabled_and_names_the_config_key():
+    # The packaged default is 1e99 — updates never fire. Reporting that as
+    # "every 1e+99 iterations" would be true and useless, so the disabled branch
+    # says so plainly and names the key that turns them on, because that is the
+    # branch a user is most likely to want to act on.
+    search = af.MultiStartAdam(n_steps=300)
+
+    assert search.iterations_per_quick_update == pytest.approx(1e99)
+
+    message = search.quick_update_message
+
+    assert "disabled" in message
+    assert "iterations_per_quick_update" in message
+    assert "config/general.yaml" in message
+    assert "1e+99" not in message
+
+
+@pytest.mark.parametrize("never", [1e99, 1e100, float("inf")])
+def test__any_never_sized_cadence_reads_as_disabled(never):
+    # The threshold is compared against, not equality-tested on 1e99, so a
+    # hand-set 1e100 in a workspace config reads as "never" too. ``inf`` is
+    # covered separately because it is not comparable-then-castable: ``int(inf)``
+    # raises, so the finite check has to come first.
+    search = af.MultiStartAdam(n_steps=300)
+    search.iterations_per_quick_update = never
+
+    assert "disabled" in search.quick_update_message
+
+
+def test__a_cadence_of_one_still_reads_as_a_number():
+    # Guards the boundary from the other side: the disabled branch must not
+    # swallow small real cadences.
+    search = af.MultiStartAdam(n_steps=300, iterations_per_quick_update=1)
+
+    assert "every 1 iterations" in search.quick_update_message