From 1dd789d642c6c59d6f2fdb22495c63d8d00f7875 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 02:46:48 +0000 Subject: [PATCH] fix: compare log likelihoods, not figures of merit, in the resume sanity check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Fitness.check_log_likelihood` compared the log likelihood stored in a previous run's samples summary against `fitness(parameters)`, which returns the figure of merit in the search's *own* convention. Those agree only when `fom_is_log_likelihood=True` and `convert_to_chi_squared=False`. For `MultiStartAdam` / `MultiStartProdigy` (`fom_is_log_likelihood=False`, `convert_to_chi_squared=True`) the fresh value is `-2 * log_posterior`, so resuming a search killed mid-run raised `SearchException` at roughly `-2x` the stored value on a completely unchanged likelihood function, making the resume path unusable. `LBFGS`, `Emcee`, `Zeus`, `NUTS` and `Drawer` carry the same mismatch. The persisted side is correct: `Sample.log_likelihood` is a true log likelihood (`MultiStartGradient.samples_via_internal_from` explicitly stores `-0.5 * best_fom - log_prior`), and every downstream consumer reads it as one. So the comparison side is converted instead, via a new `log_likelihood_from` helper that applies the exact inverse of the figure-of-merit mapping. `call_wrap` already needed that same inverse for its quick-update / history bookkeeping and now shares it. The guard still fires on a genuinely changed likelihood function, which is what it exists for — that is asserted for all three figure-of-merit conventions. Also makes the log-prior subtraction out-of-place. As an in-place `-=` it aliased and mutated `figure_of_merit` whenever `convert_to_chi_squared` was `False`, so a `call_wrap` over a NumPy array returned the log likelihood in place of the figure of merit the search asked for. Verified end-to-end by killing a `MultiStartAdam` mid-run and resuming: before, `Old = -56.3875 / New = 112.7750` (exactly -2x) and exit 1; after, the resume runs to completion. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JNbe6eLQbxUtY52EGqej5o --- autofit/non_linear/fitness.py | 69 +++++- .../test_fitness_check_log_likelihood.py | 210 ++++++++++++++++++ 2 files changed, 267 insertions(+), 12 deletions(-) create mode 100644 test_autofit/non_linear/test_fitness_check_log_likelihood.py diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index 4cb4307d2..796af8bce 100644 --- a/autofit/non_linear/fitness.py +++ b/autofit/non_linear/fitness.py @@ -277,6 +277,45 @@ def call(self, parameters): return figure_of_merit + def log_likelihood_from(self, figure_of_merit, parameters): + """ + Invert the figure-of-merit convention to recover the log likelihood. + + `call` maps a log likelihood to the figure of merit (FoM) the search consumes: it adds the summed log prior + when `fom_is_log_likelihood` is `False` (giving a log posterior) and multiplies by `-2.0` when + `convert_to_chi_squared` is `True` (giving a chi-squared). This method applies the exact inverse, so any + code holding a FoM can get back the log likelihood on the scale that `Samples` persist. + + Both callers need that inverse. `call_wrap` uses it for the quick-update / history bookkeeping, which is + defined on log likelihoods. `check_log_likelihood` uses it to compare a freshly computed value against the + log likelihood stored in a previous run's samples summary; without it, that check compares a stored log + likelihood against a value in the search's own FoM convention and every resume of a non-log-likelihood + search (e.g. `MultiStartAdam`, `LBFGS`, `Emcee`) fails its sanity check on an unchanged likelihood + function. + + Parameters + ---------- + figure_of_merit + The figure of merit returned by `call`, in this fitness's own convention. + parameters + The parameter vector the figure of merit was computed for, needed to evaluate the log priors that + `fom_is_log_likelihood=False` folded in. + + Returns + ------- + The log likelihood, on the same scale as `Sample.log_likelihood`. + """ + if self.convert_to_chi_squared: + log_likelihood = -0.5 * figure_of_merit + else: + log_likelihood = figure_of_merit + + if not self.fom_is_log_likelihood: + log_prior_list = np.array(self.model.log_prior_list_from_vector(vector=parameters, xp=np)) + log_likelihood = log_likelihood - np.sum(log_prior_list) + + return log_likelihood + def call_wrap(self, parameters): """ Wrapper around a JAX-jitted likelihood function that optionally stores @@ -313,14 +352,9 @@ def call_wrap(self, parameters): if self.use_jax_jit: figure_of_merit = float(figure_of_merit) - if self.convert_to_chi_squared: - log_likelihood = -0.5 * figure_of_merit - else: - log_likelihood = figure_of_merit - - if not self.fom_is_log_likelihood: - log_prior_list = np.array(self.model.log_prior_list_from_vector(vector=parameters, xp=np)) - log_likelihood -= np.sum(log_prior_list) + log_likelihood = self.log_likelihood_from( + figure_of_merit=figure_of_merit, parameters=parameters + ) self.manage_quick_update(parameters=parameters, log_likelihood=log_likelihood) @@ -608,16 +642,27 @@ def check_log_likelihood(self, fitness): parameters = max_log_likelihood_sample.parameter_lists_for_model(model=self.model) - log_likelihood_new = fitness(parameters=parameters) + # `fitness(...)` returns the figure of merit in this search's own convention, which is only the log + # likelihood when `fom_is_log_likelihood=True` and `convert_to_chi_squared=False`. The stored value is + # always a log likelihood (`Sample.log_likelihood`), so the fresh value is converted back onto that scale + # before comparison -- otherwise every resume of a log-posterior / chi-squared search fails this check on + # an unchanged likelihood function. + log_likelihood_new = self.log_likelihood_from( + figure_of_merit=fitness(parameters=parameters), parameters=parameters + ) if not np.isclose(log_likelihood_old, log_likelihood_new): raise exc.SearchException( f""" - Figure of merit sanity check failed. + Log likelihood sanity check failed. This means that the existing results of a model fit used a different likelihood function compared to the one implemented now. - Old Figure of Merit = {log_likelihood_old} - New Figure of Merit = {log_likelihood_new} + + Both values below are log likelihoods, converted out of this search's + figure-of-merit convention, so they are directly comparable. + + Old Log Likelihood = {log_likelihood_old} + New Log Likelihood = {log_likelihood_new} """ ) \ No newline at end of file diff --git a/test_autofit/non_linear/test_fitness_check_log_likelihood.py b/test_autofit/non_linear/test_fitness_check_log_likelihood.py new file mode 100644 index 000000000..de5d8551f --- /dev/null +++ b/test_autofit/non_linear/test_fitness_check_log_likelihood.py @@ -0,0 +1,210 @@ +import numpy as np +import pytest + +from autonerves import conf + +import autofit as af +from autofit import exc +from autofit.non_linear.fitness import Fitness +from autofit.non_linear.samples.sample import Sample +from autofit.non_linear.samples.summary import SamplesSummary + + +@pytest.fixture(name="check_likelihood_function_on") +def make_check_likelihood_function_on(): + """ + `test_autofit/config/general.yaml` disables the resume sanity check for the + suite at large. These tests are about that check, so switch it on for their + duration and restore it afterwards. + """ + original = conf.instance["general"]["test"]["check_likelihood_function"] + conf.instance["general"]["test"]["check_likelihood_function"] = True + yield + conf.instance["general"]["test"]["check_likelihood_function"] = original + + +@pytest.fixture(name="model") +def make_model(): + """ + A model whose priors give a *non-zero* summed log prior, so the log-prior leg + of the figure-of-merit conversion is actually exercised. With the default + uniform priors the log prior is 0.0 and a log-posterior FoM is numerically + indistinguishable from a log-likelihood FoM. + """ + model = af.Model(af.ex.Gaussian) + model.centre = af.GaussianPrior(mean=50.0, sigma=20.0) + model.normalization = af.GaussianPrior(mean=1.0, sigma=0.5) + model.sigma = af.GaussianPrior(mean=5.0, sigma=2.0) + return model + + +@pytest.fixture(name="analysis") +def make_analysis(): + return af.ex.Analysis(data=np.ones(20), noise_map=np.ones(20) * 0.1) + + +# Deliberately off the prior means: a `GaussianPrior` has log prior 0.0 at its mean, +# which would collapse the log-posterior conventions onto the log-likelihood one. +PARAMETERS = [10.0, 2.0, 3.0] + + +def _paths_with_stored_summary(model, analysis, tmp_path, name): + """ + Write out the `samples_summary` a previous run would have checkpointed, holding + the *true* log likelihood of `PARAMETERS` (which is what `Sample.log_likelihood` + always stores, whatever figure-of-merit convention the search itself uses). + """ + paths = af.DirectoryPaths(name=name, path_prefix=str(tmp_path)) + paths.model = model + + log_likelihood = float(Fitness(model=model, analysis=analysis).call(PARAMETERS)) + log_prior = float( + np.sum(model.log_prior_list_from_vector(vector=PARAMETERS, xp=np)) + ) + assert not np.isclose(log_prior, 0.0), "fixture must exercise the log-prior leg" + + sample = Sample.from_lists( + model=model, + parameter_lists=[PARAMETERS], + log_likelihood_list=[log_likelihood], + log_prior_list=[log_prior], + weight_list=[1.0], + )[0] + paths.save_samples_summary( + samples_summary=SamplesSummary(max_log_likelihood_sample=sample, model=model) + ) + + return paths, log_likelihood + + +@pytest.mark.parametrize( + "fom_is_log_likelihood, convert_to_chi_squared", + [ + (True, False), # nested samplers (Dynesty, Nautilus) + (False, False), # log-posterior searches (Emcee, Zeus, NUTS, Drawer) + (False, True), # chi-squared searches (MultiStartAdam/Prodigy, LBFGS) + ], +) +def test_check_log_likelihood_passes_on_resume_for_every_fom_convention( + model, + analysis, + tmp_path, + check_likelihood_function_on, + fom_is_log_likelihood, + convert_to_chi_squared, +): + """ + Resuming a search whose likelihood function has *not* changed must not raise, + whatever figure-of-merit convention that search uses. + + Regression for the bug where `check_log_likelihood` compared the stored log + likelihood against `fitness(...)`, i.e. against the search's own FoM. For a + `MultiStartGradient` (`-2 * log_posterior`) that made the fresh value roughly + `-2x` the stored one, so every resume of a killed mid-run search died with a + `SearchException` on an unchanged likelihood function. + """ + paths, _ = _paths_with_stored_summary( + model=model, + analysis=analysis, + tmp_path=tmp_path, + name=f"resume_{fom_is_log_likelihood}_{convert_to_chi_squared}", + ) + + # `check_log_likelihood` runs from `Fitness.__init__` whenever `paths` is set. + Fitness( + model=model, + analysis=analysis, + paths=paths, + fom_is_log_likelihood=fom_is_log_likelihood, + resample_figure_of_merit=-np.inf, + convert_to_chi_squared=convert_to_chi_squared, + ) + + +@pytest.mark.parametrize( + "fom_is_log_likelihood, convert_to_chi_squared", + [(True, False), (False, False), (False, True)], +) +def test_check_log_likelihood_still_raises_when_likelihood_function_changes( + model, + analysis, + tmp_path, + check_likelihood_function_on, + fom_is_log_likelihood, + convert_to_chi_squared, +): + """ + The guard exists to catch a genuinely changed likelihood function between runs. + Converting the fresh value out of the FoM convention must not defeat that: a + resume against a shifted likelihood still has to raise, in every convention. + """ + paths, _ = _paths_with_stored_summary( + model=model, + analysis=analysis, + tmp_path=tmp_path, + name=f"changed_{fom_is_log_likelihood}_{convert_to_chi_squared}", + ) + + # Stand in for "the source changed underneath the resume" by shifting the data + # the analysis fits, which moves the log likelihood and nothing else. + changed_analysis = af.ex.Analysis( + data=np.ones(20) * 2.0, noise_map=np.ones(20) * 0.1 + ) + + with pytest.raises(exc.SearchException): + Fitness( + model=model, + analysis=changed_analysis, + paths=paths, + fom_is_log_likelihood=fom_is_log_likelihood, + resample_figure_of_merit=-np.inf, + convert_to_chi_squared=convert_to_chi_squared, + ) + + +@pytest.mark.parametrize( + "fom_is_log_likelihood, convert_to_chi_squared", + [(True, False), (False, False), (False, True)], +) +def test_log_likelihood_from_inverts_the_figure_of_merit_convention( + model, analysis, fom_is_log_likelihood, convert_to_chi_squared +): + """ + `log_likelihood_from` is the exact inverse of the FoM mapping applied by `call`, + and is the single conversion both `call_wrap` and `check_log_likelihood` use. + """ + fitness = Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=fom_is_log_likelihood, + resample_figure_of_merit=-np.inf, + convert_to_chi_squared=convert_to_chi_squared, + ) + + expected = float(Fitness(model=model, analysis=analysis).call(PARAMETERS)) + + log_likelihood = fitness.log_likelihood_from( + figure_of_merit=fitness.call(PARAMETERS), parameters=PARAMETERS + ) + + assert log_likelihood == pytest.approx(expected) + + +def test_call_wrap_returns_the_figure_of_merit_not_the_log_likelihood(model, analysis): + """ + `call_wrap` derives the log likelihood for its quick-update / history bookkeeping + but must still return the FoM the search consumes. Guards against the conversion + leaking into the return value (the log-prior subtraction used to be an in-place + `-=`, which mutated the returned array when the FoM was not chi-squared). + """ + fitness = Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=False, + resample_figure_of_merit=-np.inf, + convert_to_chi_squared=False, + ) + + assert float(fitness.call_wrap(np.array(PARAMETERS))) == pytest.approx( + float(fitness.call(PARAMETERS)) + )