diff --git a/autofit/non_linear/analysis/latent.py b/autofit/non_linear/analysis/latent.py index 562a0d79a..ab7f98b6f 100644 --- a/autofit/non_linear/analysis/latent.py +++ b/autofit/non_linear/analysis/latent.py @@ -23,6 +23,7 @@ import functools import logging import time +import traceback from typing import Optional import numpy as np @@ -142,6 +143,19 @@ def latent_samples_from( latent.variables, analysis, model=samples.model ) + # A latent function that raises becomes a NaN row (below) so one bad + # sample cannot abort the batch. That must never be *silent*: a raise + # on every sample means the latent function is broken for this model + # (e.g. a `jax.jit` trace failure), and the only symptom was "no + # finite latent samples remained" with the cause lost. Count the + # failures and keep the first traceback for the warnings below. + failures = {"count": 0, "first_traceback": None} + + def _record_failure(): + failures["count"] += 1 + if failures["first_traceback"] is None: + failures["first_traceback"] = traceback.format_exc() + if analysis._use_jax: import jax import jax.numpy as jnp @@ -167,9 +181,11 @@ def _safe_jitted(p): # A latent that raises (any exception, not just # FitException) becomes a NaN row, which the global mask # below drops — one bad sample must not abort the batch. + # The raise is recorded, not swallowed: see `failures`. try: return jitted_compute_latent(p) except Exception: + _record_failure() return nan_tuple def batched_compute_latent(parameters_batch): @@ -197,9 +213,11 @@ def _safe_compute(xx): # Any exception (not just FitException) becomes a NaN row, # which the global mask below drops — a single failing latent # evaluation must not abort the whole post-fit latent pass. + # The raise is recorded, not swallowed: see `failures`. try: return compute_latent_for_model(xx) except Exception: + _record_failure() return nan_row def batched_compute_latent(x): @@ -278,11 +296,35 @@ def batched_compute_latent(x): print(f"Time to compute latent variables: {time.time() - start_latent} seconds for {len(samples)} samples.") - if not kept_idx: + n_samples = len(parameter_array) + + if failures["count"]: logger.warning( - "compute_latent_samples: no finite latent samples remained " - "after masking; skipping latent output." + "compute_latent_samples: the latent function raised on %d of " + "%d samples; each becomes a NaN row and is dropped. First " + "failure:\n%s", + failures["count"], + n_samples, + failures["first_traceback"], ) + + def _warn_no_latents(): + if n_samples and failures["count"] == n_samples: + logger.warning( + "compute_latent_samples: no finite latent samples remained " + "because the latent function raised on EVERY sample -- " + "the latent function is broken for this model (see the " + "traceback above), the samples are not at fault; skipping " + "latent output." + ) + else: + logger.warning( + "compute_latent_samples: no finite latent samples remained " + "after masking; skipping latent output." + ) + + if not kept_idx: + _warn_no_latents() return None kept_keys = [keys[i] for i in kept_idx] @@ -292,10 +334,7 @@ def batched_compute_latent(x): kept_samples = [s for s, keep in zip(all_samples, row_mask) if keep] if len(kept_samples) == 0: - logger.warning( - "compute_latent_samples: no finite latent samples remained " - "after masking; skipping latent output." - ) + _warn_no_latents() return None latent_samples = [ diff --git a/test_autofit/analysis/test_latent.py b/test_autofit/analysis/test_latent.py index 819e62bba..ef0092f5d 100644 --- a/test_autofit/analysis/test_latent.py +++ b/test_autofit/analysis/test_latent.py @@ -7,6 +7,8 @@ (``compute_latent_variables`` / ``LATENT_KEYS``) is covered by ``test_latent_variables.py``, which also exercises the back-compat shim. """ +import logging + import numpy as np import pytest @@ -129,13 +131,39 @@ def log_likelihood_function(self, instance): return 1.0 -def test_new_path_skips_arbitrary_exception_samples(): - latent = RaisingAnalysis().compute_latent_samples( - _samples([_sample(centre=1.0), _sample(centre=-1.0, weight=0.0)]) - ) +def test_new_path_skips_arbitrary_exception_samples(caplog): + with caplog.at_level(logging.WARNING, logger="autofit.non_linear.analysis.latent"): + latent = RaisingAnalysis().compute_latent_samples( + _samples([_sample(centre=1.0), _sample(centre=-1.0, weight=0.0)]) + ) assert len(latent.sample_list) == 1 assert latent.sample_list[0].kwargs == {("fwhm",): FWHM_SIGMA_3} + # The swallowed raise is reported, with its count and traceback, rather + # than vanishing into a NaN row. + assert "raised on 1 of 2 samples" in caplog.text + assert "boom from latent function" in caplog.text + assert "ValueError" in caplog.text + + +def test_every_sample_raising_names_the_cause_instead_of_a_blank_block(caplog): + """ + The Euclid ``vis_lp`` failure mode: the latent function raised on every + sample (a ``jax.jit`` trace failure), the engine turned each into a NaN + row, and the only trace was "no finite latent samples remained" -- the + cause was lost. Every-sample failure must say the function is broken and + show the first traceback. + """ + with caplog.at_level(logging.WARNING, logger="autofit.non_linear.analysis.latent"): + latent = RaisingAnalysis().compute_latent_samples( + _samples([_sample(centre=-1.0), _sample(centre=-2.0)]) + ) + assert latent is None + assert "raised on 2 of 2 samples" in caplog.text + assert "boom from latent function" in caplog.text + assert "raised on EVERY sample" in caplog.text + assert "latent function is broken" in caplog.text + class AntiCorrelatedNaNLatent(af.Latent): @staticmethod