From 71ed05666eb049cfa79bed06daa7570ae2823f1f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 10:35:52 +0100 Subject: [PATCH 01/14] feat: declare static capability attributes on all 15 searches Every search now declares jax_use, gradient, batched, honours_gradient_mode, posterior_kind, produces_evidence, resumable, warm_start, install_extra, upstream_url, citation_keys, status, test_mode_budget, objective_target and invalid_value as plain class attributes (str enums in the new autofit.non_linear.search.capabilities module), with defaults on NonLinearSearch. None is an identifier field; the A0 golden identifier table is unchanged. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- autofit/non_linear/search/abstract_search.py | 21 ++ autofit/non_linear/search/capabilities.py | 260 ++++++++++++++++++ .../search/mcmc/blackjax/nuts/search.py | 19 ++ .../search/mcmc/blackjax/smc/search.py | 19 ++ .../non_linear/search/mcmc/emcee/search.py | 19 ++ autofit/non_linear/search/mcmc/zeus/search.py | 19 ++ autofit/non_linear/search/mle/bfgs/search.py | 19 ++ .../non_linear/search/mle/drawer/search.py | 19 ++ .../search/mle/multi_start_gradient/search.py | 23 ++ .../search/nest/dynesty/search/abstract.py | 19 ++ .../non_linear/search/nest/nautilus/search.py | 19 ++ autofit/non_linear/search/nest/nss/search.py | 19 ++ .../non_linear/search/test_capabilities.py | 164 +++++++++++ 13 files changed, 639 insertions(+) create mode 100644 autofit/non_linear/search/capabilities.py create mode 100644 test_autofit/non_linear/search/test_capabilities.py diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 837d145f0..7963e61de 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -52,6 +52,7 @@ from autofit.graphical.expectation_propagation import AbstractFactorOptimiser from autofit.non_linear.fitness import get_timeout_seconds +from autofit.non_linear.search import capabilities as cap from autofit.non_linear.test_mode import ( test_mode_level, test_mode_samples, @@ -335,6 +336,26 @@ def __init__( __identifier_fields__ = tuple() + # Static capability declarations (``autofit.non_linear.search.capabilities``). + # These are the defaults for a search that declares nothing; every public search + # overrides them and the registry (``search/registry.py``) mirrors its values. + # None of them is ever an identifier field. + jax_use = cap.JaxUse.NONE + gradient = cap.Gradient.NONE + batched = False + honours_gradient_mode = False + posterior_kind = None + produces_evidence = False + resumable = False + warm_start = cap.WarmStart.NEUTRAL + install_extra = "" + upstream_url = "" + citation_keys = () + status = cap.Status.EXPERIMENTAL + test_mode_budget = {} + objective_target = None + invalid_value = -float("inf") + def optimise( self, factor_approx: FactorApproximation, diff --git a/autofit/non_linear/search/capabilities.py b/autofit/non_linear/search/capabilities.py new file mode 100644 index 000000000..c5eb81078 --- /dev/null +++ b/autofit/non_linear/search/capabilities.py @@ -0,0 +1,260 @@ +""" +Static capability declarations of the non-linear searches. + +Every search declares what it is as plain class attributes, defined on +``NonLinearSearch`` (the defaults) and overridden by each concrete search: + +======================== ============================================================== +Attribute Meaning +======================== ============================================================== +``jax_use`` ``none`` (the backend never needs JAX), ``optional`` (it has a + jitted or batched fast path when the analysis is JAX) or + ``required`` (the backend cannot run without a traceable JAX + likelihood). +``gradient`` ``none`` or ``uses``: whether the backend differentiates the + likelihood. +``batched`` Whether the backend evaluates a vectorised (``jax.vmap``) + likelihood over many points at once. +``honours_gradient_mode`` Whether the backend reads ``Analysis.gradient_mode``. +``posterior_kind`` ``chain`` (equally weighted MCMC draws), ``weighted`` (weighted + samples, e.g. nested sampling) or ``point`` (a point estimate). +``produces_evidence`` Whether the backend estimates the Bayesian evidence. +``resumable`` Whether an interrupted run resumes from its own checkpoint. +``warm_start`` ``provider`` (starts cold and can seed another search), + ``consumer`` (wants a warm start) or ``neutral``. +``install_extra`` The pip extra that installs the backend (``""`` = base). +``upstream_url`` The backend project's home page. +``citation_keys`` Keys into ``files/citations.bib``. +``status`` ``stable``, ``experimental`` or ``archived``. +``test_mode_budget`` The reduced budget ``PYAUTO_TEST_MODE=1`` applies. +``objective_target`` What the backend's objective returns and in which coordinates. +``invalid_value`` The sentinel the objective returns for an invalid model. +======================== ============================================================== + +These are **static** capabilities, a property of the class. The capability a given run +actually has can be narrower: ``SMC`` estimates the evidence only when it starts from a +prior-sampling initializer (a warm start skips the tempering from the prior). Effective +run capabilities are not declared here. + +None of these attributes is ever in a search's ``__identifier_fields__``, so declaring +or changing one never changes an output path. The declarative registry +(``autofit.non_linear.search.registry``) mirrors them for every public search, and a +test asserts the two agree. +""" + +import enum +import math +from typing import Any, Dict, NamedTuple + + +class JaxUse(str, enum.Enum): + """ + How a search's backend relates to JAX. + """ + + NONE = "none" + OPTIONAL = "optional" + REQUIRED = "required" + + def __str__(self): + return self.value + + +class Gradient(str, enum.Enum): + """ + Whether a search's backend differentiates the likelihood. + """ + + NONE = "none" + USES = "uses" + + def __str__(self): + return self.value + + +class PosteriorKind(str, enum.Enum): + """ + The representation of a search's posterior samples. + """ + + CHAIN = "chain" + WEIGHTED = "weighted" + POINT = "point" + + def __str__(self): + return self.value + + +class WarmStart(str, enum.Enum): + """ + A search's role when searches are chained. + """ + + PROVIDER = "provider" + CONSUMER = "consumer" + NEUTRAL = "neutral" + + def __str__(self): + return self.value + + +class Status(str, enum.Enum): + """ + The maturity of a search. + """ + + STABLE = "stable" + EXPERIMENTAL = "experimental" + ARCHIVED = "archived" + + def __str__(self): + return self.value + + +class ObjectiveQuantity(str, enum.Enum): + """ + The statistical quantity a search's objective returns. + """ + + LOG_LIKELIHOOD = "log_likelihood" + LOG_POSTERIOR = "log_posterior" + NEG2_LOG_POSTERIOR = "neg2_log_posterior" + + def __str__(self): + return self.value + + +class CoordinateSpace(str, enum.Enum): + """ + The coordinates a search's backend proposes points in. ``Fitness`` always receives + physical parameters; a ``unit_cube`` backend maps through the prior transform first. + """ + + PHYSICAL = "physical" + UNIT_CUBE = "unit_cube" + + def __str__(self): + return self.value + + +class ObjectiveTarget(NamedTuple): + """ + What a search's objective returns (``quantity``) and the coordinates its backend + proposes points in (``space``). + """ + + quantity: ObjectiveQuantity + space: CoordinateSpace + + def to_dict(self) -> Dict[str, str]: + return {"quantity": str(self.quantity), "space": str(self.space)} + + +CAPABILITY_ATTRIBUTES = ( + "jax_use", + "gradient", + "batched", + "honours_gradient_mode", + "posterior_kind", + "produces_evidence", + "resumable", + "warm_start", + "install_extra", + "upstream_url", + "citation_keys", + "status", + "test_mode_budget", + "objective_target", + "invalid_value", +) +""" +The names of the static capability attributes, in manifest order. +""" + + +def invalid_value_to_str(value: float) -> str: + """ + The JSON-safe spelling of an ``invalid_value`` sentinel (JSON has no ``-inf``). + """ + value = float(value) + if math.isinf(value): + return "-inf" if value < 0 else "inf" + return format(value, "g") + + +def capability_value(name: str, value: Any) -> Any: + """ + Normalise one capability attribute value into plain JSON data. + """ + if name == "objective_target": + return value.to_dict() if value is not None else None + if name == "invalid_value": + return invalid_value_to_str(value) + if name == "citation_keys": + return list(value) + if name == "test_mode_budget": + return dict(value) + if isinstance(value, enum.Enum): + return value.value + return value + + +def capabilities_from(cls) -> Dict[str, Any]: + """ + The static capabilities of a search class, as plain JSON data in manifest order. + """ + return { + name: capability_value(name, getattr(cls, name)) + for name in CAPABILITY_ATTRIBUTES + } + + +JAX_REQUIRED_MESSAGE = ( + "{search} requires a JAX analysis: its backend is JAX-native (jax_use='required') " + "and cannot run on a numpy likelihood, but the analysis passed to fit has " + "is_jax=False. Construct the analysis with use_jax=True (and do not set " + "PYAUTO_DISABLE_JAX=1 outside test mode), or choose a search whose jax_use is " + "'none' or 'optional' (see the capability matrix in the PyAutoFit docs)." +) +""" +The one message raised when a ``jax_use='required'`` search is given a numpy analysis. +""" + + +def check_jax_required(search, analysis): + """ + Raise ``SearchException`` when ``search`` requires JAX and ``analysis`` is not JAX. + + This is the single fail-fast gate every ``jax_use='required'`` search goes through: + ``NonLinearSearch.start_resume_fit`` calls it after the test-mode bypass return, and + the JAX-native searches call it again at the top of ``_fit`` for direct callers. + """ + from autofit import exc + + if search.jax_use == JaxUse.REQUIRED and not analysis.is_jax: + raise exc.SearchException( + JAX_REQUIRED_MESSAGE.format(search=type(search).__name__) + ) + + +def capability_summary_from(search) -> str: + """ + A short human-readable block naming a search and its main capabilities, written as + the header of ``model.info`` and as a section of ``search.summary``. + """ + cls = type(search) + target = cls.objective_target + target_text = ( + f"{target.quantity} ({target.space} coordinates)" + if target is not None + else "undeclared" + ) + return ( + f"Search: {cls.__name__} ({cls.__module__}.{cls.__qualname__})\n" + f"JAX use = {cls.jax_use}, gradient = {cls.gradient}, batched = {cls.batched}\n" + f"Posterior kind = {cls.posterior_kind}, produces evidence = " + f"{cls.produces_evidence}, resumable = {cls.resumable}\n" + f"Objective = {target_text}, invalid value = " + f"{invalid_value_to_str(cls.invalid_value)}\n" + f"Status = {cls.status}\n" + ) diff --git a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py index b4f10ba67..148591d4f 100644 --- a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py @@ -27,6 +27,7 @@ from autofit.non_linear.test_mode import is_test_mode from autofit.non_linear.samples.mcmc import SamplesMCMC from autofit.non_linear.samples.sample import Sample +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa @@ -35,6 +36,24 @@ class BlackJAXNUTS(AbstractMCMC): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.REQUIRED + gradient = cap.Gradient.USES + batched = True + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.CHAIN + produces_evidence = False + resumable = False + warm_start = cap.WarmStart.CONSUMER + install_extra = "optional" + upstream_url = "https://github.com/blackjax-devs/blackjax" + citation_keys = ("blackjax",) + status = cap.Status.STABLE + test_mode_budget = {"num_warmup": 20, "num_samples": 20, "num_chains": 2} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) + invalid_value = -float("inf") + __identifier_fields__ = ( "num_warmup", "num_samples", diff --git a/autofit/non_linear/search/mcmc/blackjax/smc/search.py b/autofit/non_linear/search/mcmc/blackjax/smc/search.py index d6f18659a..979e4c77c 100644 --- a/autofit/non_linear/search/mcmc/blackjax/smc/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/smc/search.py @@ -27,6 +27,7 @@ from autofit.non_linear.search.mcmc.blackjax.smc.samples import SamplesSMC from autofit.non_linear.test_mode import is_test_mode from autofit.non_linear.samples.sample import Sample +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa @@ -51,6 +52,24 @@ class SMC(AbstractMCMC): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.REQUIRED + gradient = cap.Gradient.USES + batched = True + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.WEIGHTED + produces_evidence = True + resumable = False + warm_start = cap.WarmStart.PROVIDER + install_extra = "optional" + upstream_url = "https://github.com/blackjax-devs/blackjax" + citation_keys = ("blackjax",) + status = cap.Status.EXPERIMENTAL + test_mode_budget = {"num_particles": 16, "num_mcmc_steps": 2, "max_smc_steps": 5} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.PHYSICAL) + invalid_value = -1.0e99 + __identifier_fields__ = ( "num_particles", "kernel", diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index 912eb207e..5bc0bdc5b 100644 --- a/autofit/non_linear/search/mcmc/emcee/search.py +++ b/autofit/non_linear/search/mcmc/emcee/search.py @@ -20,6 +20,7 @@ from autofit.non_linear.test_mode import is_test_mode from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.mcmc import SamplesMCMC +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa @@ -28,6 +29,24 @@ class Emcee(AbstractMCMC): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.NONE + gradient = cap.Gradient.NONE + batched = False + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.CHAIN + produces_evidence = False + resumable = False + warm_start = cap.WarmStart.CONSUMER + install_extra = "" + upstream_url = "https://github.com/dfm/emcee" + citation_keys = ("emcee",) + status = cap.Status.STABLE + test_mode_budget = {"nwalkers": 20, "nsteps": 10} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) + invalid_value = -float("inf") + __identifier_fields__ = ("nwalkers",) def __init__( diff --git a/autofit/non_linear/search/mcmc/zeus/search.py b/autofit/non_linear/search/mcmc/zeus/search.py index 378987d37..57abd4099 100644 --- a/autofit/non_linear/search/mcmc/zeus/search.py +++ b/autofit/non_linear/search/mcmc/zeus/search.py @@ -17,6 +17,7 @@ from autofit.non_linear.samples.sample import Sample from autofit.non_linear.test_mode import is_test_mode from autofit.non_linear.samples.mcmc import SamplesMCMC +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa @@ -25,6 +26,24 @@ class Zeus(AbstractMCMC): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.NONE + gradient = cap.Gradient.NONE + batched = False + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.CHAIN + produces_evidence = False + resumable = False + warm_start = cap.WarmStart.CONSUMER + install_extra = "optional" + upstream_url = "https://github.com/minaskar/zeus" + citation_keys = ("zeus1", "zeus2") + status = cap.Status.STABLE + test_mode_budget = {"nwalkers": 20, "nsteps": 10} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) + invalid_value = -float("inf") + __identifier_fields__ = ( "nwalkers", "tune", diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index 705a63d6a..d79ee2c93 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -16,6 +16,7 @@ from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.samples import Samples +from autofit.non_linear.search import capabilities as cap import numpy as np @@ -31,6 +32,24 @@ class AbstractBFGS(AbstractMLE): # ``Drawer`` deliberately does NOT declare this: it inherits the attribute # from ``AbstractMLE`` but never consumes it, and a setting that cannot # affect the result must not re-key stored results. + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.OPTIONAL + gradient = cap.Gradient.NONE + batched = False + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.POINT + produces_evidence = False + resumable = False + warm_start = cap.WarmStart.PROVIDER + install_extra = "" + upstream_url = "https://github.com/scipy/scipy" + citation_keys = ("scipy",) + status = cap.Status.STABLE + test_mode_budget = {} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.NEG2_LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) + invalid_value = -float("inf") + __identifier_fields__ = ("clipper",) method = None diff --git a/autofit/non_linear/search/mle/drawer/search.py b/autofit/non_linear/search/mle/drawer/search.py index 8d1753537..3355a86e3 100644 --- a/autofit/non_linear/search/mle/drawer/search.py +++ b/autofit/non_linear/search/mle/drawer/search.py @@ -9,12 +9,31 @@ from autofit.non_linear.search.mle.abstract_mle import AbstractMLE from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples import Samples, Sample +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa class Drawer(AbstractMLE): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.NONE + gradient = cap.Gradient.NONE + batched = False + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.POINT + produces_evidence = False + resumable = False + warm_start = cap.WarmStart.NEUTRAL + install_extra = "" + upstream_url = "https://github.com/PyAutoLabs/PyAutoFit" + citation_keys = () + status = cap.Status.STABLE + test_mode_budget = {} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) + invalid_value = -float("inf") + __identifier_fields__ = ("total_draws",) def __init__( diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 0d2792ecf..3a4cc5bef 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -27,6 +27,7 @@ from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.samples import Samples +from autofit.non_linear.search import capabilities as cap from autofit.non_linear.search.mle.multi_start_gradient.convergence import ( MultiStartGradientConvergence, ) @@ -43,6 +44,24 @@ class AbstractMultiStartGradient(AbstractMLE): # on disk, not deleted). Scoped to the clipper-consuming searches only: the # nested samplers and MCMC searches never touch the clipper and their # identifiers must stay byte-identical (PyAutoFit#1493). + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.REQUIRED + gradient = cap.Gradient.USES + batched = True + honours_gradient_mode = True + posterior_kind = cap.PosteriorKind.POINT + produces_evidence = False + resumable = True + warm_start = cap.WarmStart.PROVIDER + install_extra = "" + upstream_url = "https://github.com/google-deepmind/optax" + citation_keys = ("optax",) + status = cap.Status.STABLE + test_mode_budget = {"convergence.window": 1, "convergence.min_steps": 1} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.NEG2_LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) + invalid_value = -float("inf") + __identifier_fields__ = ("clipper",) # Name of the optax update rule, resolved lazily at fit time from ``optax`` @@ -2227,6 +2246,7 @@ class MultiStartADABelief(AbstractMultiStartGradient): optax_method = "adabelief" _default_learning_rate = 1.0e-2 + status = cap.Status.EXPERIMENTAL class MultiStartLion(AbstractMultiStartGradient): @@ -2237,6 +2257,7 @@ class MultiStartLion(AbstractMultiStartGradient): optax_method = "lion" _default_learning_rate = 1.0e-3 + status = cap.Status.EXPERIMENTAL class MultiStartProdigy(AbstractMultiStartGradient): @@ -2256,3 +2277,5 @@ class MultiStartProdigy(AbstractMultiStartGradient): optax_method = "prodigy" _default_learning_rate = None + upstream_url = "https://github.com/konstmish/prodigy" + citation_keys = ("optax", "prodigy") diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index 484316e79..5fd941a0a 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -18,6 +18,7 @@ from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.nest import SamplesNest from autofit.non_linear.test_mode import is_test_mode +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa @@ -68,6 +69,24 @@ def prior_transform(cube, model): class AbstractDynesty(AbstractNest, ABC): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.OPTIONAL + gradient = cap.Gradient.NONE + batched = False + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.WEIGHTED + produces_evidence = True + resumable = True + warm_start = cap.WarmStart.PROVIDER + install_extra = "" + upstream_url = "https://github.com/joshspeagle/dynesty" + citation_keys = ("dynesty",) + status = cap.Status.STABLE + test_mode_budget = {"maxcall": 1} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.UNIT_CUBE) + invalid_value = -1.0e99 + def __init__( self, name: Optional[str] = None, diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 13320851a..fe2d344b0 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -20,6 +20,7 @@ from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.nest import SamplesNest from autofit.non_linear.test_mode import is_test_mode +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa @@ -155,6 +156,24 @@ def __exit__(self, exc_type, exc_val, exc_tb): class Nautilus(abstract_nest.AbstractNest): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.OPTIONAL + gradient = cap.Gradient.NONE + batched = True + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.WEIGHTED + produces_evidence = True + resumable = True + warm_start = cap.WarmStart.PROVIDER + install_extra = "optional" + upstream_url = "https://github.com/johannesulf/nautilus" + citation_keys = ("nautilus",) + status = cap.Status.STABLE + test_mode_budget = {"n_like_max": 1} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.UNIT_CUBE) + invalid_value = -1.0e99 + __identifier_fields__ = ( "n_live", "n_update", diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index 7c358d696..5790046d3 100644 --- a/autofit/non_linear/search/nest/nss/search.py +++ b/autofit/non_linear/search/nest/nss/search.py @@ -15,6 +15,7 @@ from .samples import NSSamples from autofit.non_linear.samples.sample import Sample from autofit.non_linear.test_mode import is_test_mode +from autofit.non_linear.search import capabilities as cap if TYPE_CHECKING: from autofit.database.sqlalchemy_ import sa @@ -192,6 +193,24 @@ def __init__( class NSS(abstract_nest.AbstractNest): + # Static capabilities (see ``autofit.non_linear.search.capabilities``); mirrored + # by ``search/registry.py``. Never identifier fields. + jax_use = cap.JaxUse.REQUIRED + gradient = cap.Gradient.NONE + batched = True + honours_gradient_mode = False + posterior_kind = cap.PosteriorKind.WEIGHTED + produces_evidence = True + resumable = True + warm_start = cap.WarmStart.PROVIDER + install_extra = "optional" + upstream_url = "https://github.com/blackjax-devs/blackjax" + citation_keys = ("blackjax",) + status = cap.Status.EXPERIMENTAL + test_mode_budget = {"termination": -1.0} + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.UNIT_CUBE) + invalid_value = NSS_INVALID_LOG_LIKELIHOOD + __identifier_fields__ = ( "n_live", "num_mcmc_steps", diff --git a/test_autofit/non_linear/search/test_capabilities.py b/test_autofit/non_linear/search/test_capabilities.py new file mode 100644 index 000000000..8ff1cd7fd --- /dev/null +++ b/test_autofit/non_linear/search/test_capabilities.py @@ -0,0 +1,164 @@ +""" +The static capability attributes every public search declares +(``autofit.non_linear.search.capabilities``), parametrised per class over the +conformance roster. + +Classes are resolved inside each test, so the module collects on every CI leg. The +values themselves are mirrored by ``search/registry.py`` and compared there +(``test_registry.py``); these tests check that each declaration is well formed, that +none of them is an identifier field and that ``test_mode_budget`` is what +``PYAUTO_TEST_MODE=1`` actually applies. +""" + +import json +import math +import re +from pathlib import Path + +import pytest + +from autofit.non_linear.search import capabilities as cap +from autofit.non_linear.search.abstract_search import NonLinearSearch + +from test_autofit.non_linear.search.conformance_roster import searches_under_test + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +ROSTER = searches_under_test() + +CITATIONS_BIB = Path(__file__).parents[3] / "files" / "citations.bib" + + +def _params(): + return [pytest.param(entry, id=entry.name) for entry in ROSTER] + + +def _bib_keys(): + return set(re.findall(r"^@\w+\{([^,]+),", CITATIONS_BIB.read_text(), re.M)) + + +def _search_class(entry): + missing = entry.missing_requirements() + if missing: + pytest.skip(f"{entry.name} cannot be constructed without {missing}") + return entry.resolve() + + +def test_defaults_on_non_linear_search(): + assert NonLinearSearch.jax_use == cap.JaxUse.NONE + assert NonLinearSearch.gradient == cap.Gradient.NONE + assert NonLinearSearch.batched is False + assert NonLinearSearch.posterior_kind is None + assert NonLinearSearch.status == cap.Status.EXPERIMENTAL + assert NonLinearSearch.objective_target is None + assert NonLinearSearch.invalid_value == -math.inf + + +@pytest.mark.parametrize("entry", _params()) +def test_capabilities_are_declared_and_well_formed(entry): + cls = entry.resolve() + + assert isinstance(cls.jax_use, cap.JaxUse) + assert isinstance(cls.gradient, cap.Gradient) + assert isinstance(cls.batched, bool) + assert isinstance(cls.honours_gradient_mode, bool) + assert isinstance(cls.posterior_kind, cap.PosteriorKind) + assert isinstance(cls.produces_evidence, bool) + assert isinstance(cls.resumable, bool) + assert isinstance(cls.warm_start, cap.WarmStart) + assert cls.install_extra in ("", "optional") + assert cls.upstream_url.startswith("https://") + assert isinstance(cls.citation_keys, tuple) + assert isinstance(cls.status, cap.Status) + assert isinstance(cls.test_mode_budget, dict) + assert isinstance(cls.objective_target, cap.ObjectiveTarget) + assert isinstance(cls.objective_target.quantity, cap.ObjectiveQuantity) + assert isinstance(cls.objective_target.space, cap.CoordinateSpace) + assert isinstance(cls.invalid_value, float) + + +@pytest.mark.parametrize("entry", _params()) +def test_capability_rules(entry): + """ + Rules that hold between attributes: a gradient backend is JAX-native, only a + gradient backend can honour ``gradient_mode``, evidence comes with weighted + samples and point estimators never claim evidence. + """ + cls = entry.resolve() + + if cls.gradient == cap.Gradient.USES: + assert cls.jax_use == cap.JaxUse.REQUIRED + if cls.honours_gradient_mode: + assert cls.gradient == cap.Gradient.USES + if cls.produces_evidence: + assert cls.posterior_kind == cap.PosteriorKind.WEIGHTED + if cls.posterior_kind == cap.PosteriorKind.POINT: + assert not cls.produces_evidence + + +@pytest.mark.parametrize("entry", _params()) +def test_roster_jax_use_matches_the_class(entry): + assert entry.jax_use == str(entry.resolve().jax_use) + + +@pytest.mark.parametrize("entry", _params()) +def test_capabilities_are_never_identifier_fields(entry): + cls = entry.resolve() + + assert not set(cap.CAPABILITY_ATTRIBUTES) & set(cls.__identifier_fields__) + + +@pytest.mark.parametrize("entry", _params()) +def test_citation_keys_are_in_citations_bib(entry): + assert set(entry.resolve().citation_keys) <= _bib_keys() + + +@pytest.mark.parametrize("entry", _params()) +def test_capabilities_serialise_to_json(entry): + capabilities = cap.capabilities_from(entry.resolve()) + + assert list(capabilities) == list(cap.CAPABILITY_ATTRIBUTES) + assert json.loads(json.dumps(capabilities)) == capabilities + + +def _attribute(search, dotted): + value = search + for part in dotted.split("."): + value = getattr(value, part) + return value + + +@pytest.mark.parametrize("entry", _params()) +def test_test_mode_budget_is_what_test_mode_applies(entry, monkeypatch): + """ + Each budget entry is the value ``PYAUTO_TEST_MODE=1`` sets, or the default when the + default is already below it (``num_chains`` and ``num_particles`` are capped with + ``min``). + """ + cls = _search_class(entry) + + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + default = cls() + + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + reduced = cls() + + assert reduced is not default + for key, budget in cls.test_mode_budget.items(): + value = _attribute(reduced, key) + default_value = _attribute(default, key) + assert value == budget or value == default_value <= budget, key + + +def test_jax_required_message_names_the_search_and_the_fix(): + message = cap.JAX_REQUIRED_MESSAGE.format(search="BlackJAXNUTS") + + assert "BlackJAXNUTS" in message + assert "use_jax=True" in message + assert "PYAUTO_DISABLE_JAX" in message + + +def test_invalid_value_spelling(): + assert cap.invalid_value_to_str(-math.inf) == "-inf" + assert cap.invalid_value_to_str(-1.0e99) == "-1e+99" + assert cap.invalid_value_to_str(-1.0e30) == "-1e+30" From 196abf39ac51f97dedbd4f26f9dfbde206e673b3 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 10:52:50 +0100 Subject: [PATCH 02/14] feat: Analysis.is_jax as the single JAX probe, with D12 wrapper rules Analysis.is_jax replaces every getattr(analysis, "_use_jax", False) / analysis._use_jax probe in the searches, Fitness and the latent machinery. ModelAnalysis and AnalysisFactor forward is_jax and gradient_mode from the analysis they wrap. FactorGraphModel(use_jax=None) derives its backend from its factors and its gradient_mode from their common mode (reverse when mixed); whole-graph fitting raises SearchException naming any factor that disagrees, while per-factor EP still allows mixed factors. HierarchicalFactor honours PYAUTO_DISABLE_JAX and the graph flag survives pytree flattening and pickling. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- autofit/graphical/declarative/collection.py | 99 ++++++- .../graphical/declarative/factor/analysis.py | 21 ++ .../declarative/factor/hierarchical.py | 8 + autofit/non_linear/analysis/analysis.py | 25 +- autofit/non_linear/analysis/latent.py | 4 +- autofit/non_linear/analysis/model_analysis.py | 24 +- autofit/non_linear/fitness.py | 4 +- autofit/non_linear/search/abstract_search.py | 11 +- autofit/non_linear/search/mle/bfgs/search.py | 2 +- .../search/nest/dynesty/search/abstract.py | 6 +- .../non_linear/search/nest/nautilus/search.py | 4 +- test_autofit/analysis/test_is_jax.py | 241 ++++++++++++++++++ 12 files changed, 422 insertions(+), 27 deletions(-) create mode 100644 test_autofit/analysis/test_is_jax.py diff --git a/autofit/graphical/declarative/collection.py b/autofit/graphical/declarative/collection.py index c194c9f2d..dbe3a4077 100644 --- a/autofit/graphical/declarative/collection.py +++ b/autofit/graphical/declarative/collection.py @@ -1,4 +1,4 @@ -from typing import Union, Optional +from typing import List, Union, Optional from autofit.graphical.declarative.factor.hierarchical import HierarchicalFactor @@ -19,7 +19,7 @@ def __init__( *model_factors: Union[AbstractDeclarativeFactor, HierarchicalFactor], name=None, include_prior_factors=True, - use_jax : bool = False + use_jax: Optional[bool] = None, ): """ A collection of factors that describe models, which can be @@ -31,6 +31,15 @@ def __init__( ---------- model_factors Factors which are hierarchical or associated with a specific analysis + use_jax + Whether the whole graph is evaluated on JAX when it is fitted as one + analysis (``search.fit(model=graph.global_prior_model, analysis=graph)``). + ``None`` (the default) derives it from the factors: the graph is JAX when + every factor is. An explicit bool is honoured, but whole-graph fitting + requires every factor to agree with it and raises ``SearchException`` + naming the factors that do not (``check_backend_agreement``). Per-factor + expectation propagation (``graph.optimise``) fits each factor with its own + analysis, so mixed numpy/JAX factors are allowed there. """ super().__init__( include_prior_factors=include_prior_factors, @@ -42,7 +51,7 @@ def __init__( def tree_flatten(self): return ( (self._model_factors,), - (self._name, self._include_prior_factors), + (self._name, self._include_prior_factors, self._explicit_use_jax), ) @classmethod @@ -51,8 +60,84 @@ def tree_unflatten(cls, aux_data, children): *children[0], name=aux_data[0], include_prior_factors=aux_data[1], + use_jax=aux_data[2], ) + @property + def _use_jax(self) -> bool: + """ + The graph's backend: the explicit ``use_jax`` when one was given, otherwise + ``True`` exactly when every factor is JAX (``is_jax`` reads this). + """ + explicit = self.__dict__.get("_explicit_use_jax") + if explicit is not None: + return bool(explicit) + flags = [factor.is_jax for factor in self.__dict__.get("_model_factors", [])] + return bool(flags) and all(flags) + + @_use_jax.setter + def _use_jax(self, value: Optional[bool]): + # Set by `Analysis.__init__`: `None` means "derive from the factors". + self._explicit_use_jax = value + + def _flat_factors(self): + model_factors = list() + for model_factor in self.__dict__.get("_model_factors", []): + if isinstance(model_factor, HierarchicalFactor): + model_factors.extend(model_factor.factors) + else: + model_factors.append(model_factor) + return model_factors + + def factors_disagreeing_on_backend(self) -> List[str]: + """ + The names of the factors whose ``is_jax`` differs from the graph's. + """ + is_jax = self.is_jax + return [ + factor.name for factor in self._flat_factors() if factor.is_jax != is_jax + ] + + def check_backend_agreement(self): + """ + Raise ``SearchException`` unless every factor agrees with the graph's backend. + + Whole-graph fitting evaluates every factor inside one objective, so a JAX graph + traces numpy factors (which fail on tracers) and a numpy graph never compiles a + JAX factor. ``NonLinearSearch`` calls this before fitting a graph as one + analysis; per-factor expectation propagation never does. + """ + disagreeing = self.factors_disagreeing_on_backend() + if disagreeing: + from autofit import exc + + raise exc.SearchException( + f"The FactorGraphModel {self.name} is fitted as one analysis with " + f"is_jax={self.is_jax}, but these factors have " + f"is_jax={not self.is_jax}: {', '.join(disagreeing)}. Whole-graph " + f"fitting requires every factor to use the same backend: build every " + f"factor's analysis with the same use_jax (and HierarchicalFactor with " + f"the same use_jax), or omit FactorGraphModel(use_jax=...) so the graph " + f"derives it from its factors. Per-factor expectation propagation " + f"(graph.optimise) allows mixed factors." + ) + + @property + def gradient_mode(self) -> str: + """ + The factors' common ``gradient_mode``; ``"reverse"`` (the ``Analysis`` default) + when the factors declare different modes. Both modes compute the same gradient, + so a mixed graph falls back to the mode whose memory does not scale with the + number of free parameters. + """ + modes = { + getattr(factor, "gradient_mode", "reverse") + for factor in self._flat_factors() + } + if len(modes) == 1: + return modes.pop() + return "reverse" + @property def prior_model(self): """ @@ -144,13 +229,7 @@ def _shared_state_from(self, instance: ModelInstance): @property def model_factors(self): - model_factors = list() - for model_factor in self._model_factors: - if isinstance(model_factor, HierarchicalFactor): - model_factors.extend(model_factor.factors) - else: - model_factors.append(model_factor) - return model_factors + return self._flat_factors() def make_result( self, diff --git a/autofit/graphical/declarative/factor/analysis.py b/autofit/graphical/declarative/factor/analysis.py index 55dda7d8d..061895932 100644 --- a/autofit/graphical/declarative/factor/analysis.py +++ b/autofit/graphical/declarative/factor/analysis.py @@ -86,6 +86,27 @@ def __init__( name=name, ) + @property + def is_jax(self) -> bool: + """ + Whether this factor's analysis is JAX (``Analysis.is_jax`` of the wrapped + analysis). Per-factor expectation propagation fits each factor with its own + analysis, so mixed numpy/JAX factors are allowed there; whole-graph fitting of a + ``FactorGraphModel`` requires every factor to agree. + """ + return self.analysis.is_jax + + @property + def _use_jax(self) -> bool: + return self.analysis.is_jax + + @property + def gradient_mode(self) -> str: + """ + The wrapped analysis's declared gradient mode. + """ + return self.analysis.gradient_mode + def tree_flatten(self): return ( (self.prior_model,), diff --git a/autofit/graphical/declarative/factor/hierarchical.py b/autofit/graphical/declarative/factor/hierarchical.py index 276325dc8..a956cf950 100644 --- a/autofit/graphical/declarative/factor/hierarchical.py +++ b/autofit/graphical/declarative/factor/hierarchical.py @@ -78,6 +78,14 @@ def __init__( self._name = name or namer(self.__class__.__name__) self._factors = list() self.optimiser = optimiser + + # The same switch `Analysis.__init__` honours, so a hierarchical factor in a + # graph agrees with its analysis factors when JAX is disabled for a smoke run. + import os + + if os.environ.get("PYAUTO_DISABLE_JAX") == "1": + use_jax = False + self._use_jax = use_jax @property diff --git a/autofit/non_linear/analysis/analysis.py b/autofit/non_linear/analysis/analysis.py index 47046d179..f586cd50d 100644 --- a/autofit/non_linear/analysis/analysis.py +++ b/autofit/non_linear/analysis/analysis.py @@ -151,9 +151,26 @@ def method(*args, **kwargs): return method + @property + def is_jax(self) -> bool: + """ + Whether this analysis evaluates its likelihood on the JAX backend. + + The single read-only probe every search, ``Fitness`` and the latent machinery use + to ask "is this analysis JAX?". ``use_jax`` passed to ``__init__`` is the user's + choice; ``is_jax`` is the resolved value after ``PYAUTO_DISABLE_JAX=1`` and the + no-JAX-installed fallback have been applied. An analysis that never called + ``Analysis.__init__`` is numpy. + + Wrappers forward it: ``ModelAnalysis`` takes its wrapped analysis's value, + ``AnalysisFactor`` reads its analysis and ``FactorGraphModel`` derives it from its + factors. + """ + return bool(getattr(self, "_use_jax", False)) + @property def _xp(self): - if self._use_jax: + if self.is_jax: import jax.numpy as jnp return jnp return np @@ -356,7 +373,7 @@ def supports_background_update(self) -> bool: @property def supports_jax_visualization(self) -> bool: """Whether the visualizer can work directly with JAX arrays.""" - return self._use_jax + return self.is_jax def perform_quick_update(self, paths, instance): raise NotImplementedError @@ -388,7 +405,7 @@ def batched_memory_bytes(self, model, batch_size: int, gradient: bool = False): The projected bytes, or ``None`` when this analysis is not on the JAX path (nothing to measure). """ - if not self._use_jax: + if not self.is_jax: return None import jax @@ -441,7 +458,7 @@ def print_vram_use(self, model, batch_size : int, gradient: bool = False) -> str if skip_fit_output(): return - if not self._use_jax: + if not self.is_jax: print("use_jax=False for this analysis, therefore does not use GPU and VRAM use cannot be profiled.") return diff --git a/autofit/non_linear/analysis/latent.py b/autofit/non_linear/analysis/latent.py index ab7f98b6f..dfecd9d7a 100644 --- a/autofit/non_linear/analysis/latent.py +++ b/autofit/non_linear/analysis/latent.py @@ -156,7 +156,7 @@ def _record_failure(): if failures["first_traceback"] is None: failures["first_traceback"] = traceback.format_exc() - if analysis._use_jax: + if analysis.is_jax: import jax import jax.numpy as jnp if batch_mode == "vmap": @@ -248,7 +248,7 @@ def batched_compute_latent(x): # batched JAX call on this chunk latent_values_batch = batched_compute_latent(batch) - if analysis._use_jax: + if analysis.is_jax: import jax.numpy as jnp latent_values_batch = jnp.stack(latent_values_batch, axis=-1) # (batch, n_latents) diff --git a/autofit/non_linear/analysis/model_analysis.py b/autofit/non_linear/analysis/model_analysis.py index 8b1a92b59..379fa72a1 100644 --- a/autofit/non_linear/analysis/model_analysis.py +++ b/autofit/non_linear/analysis/model_analysis.py @@ -6,20 +6,42 @@ class ModelAnalysis(Analysis): - def __init__(self, analysis: Analysis, model: AbstractPriorModel, use_jax : bool = False): + def __init__( + self, + analysis: Analysis, + model: AbstractPriorModel, + use_jax: Optional[bool] = None, + ): """ Comprises a model and an analysis that can be applied to instances of that model. Parameters ---------- analysis + The wrapped analysis. model + The model whose instances the wrapped analysis is evaluated on. + use_jax + ``None`` (the default) forwards the wrapped analysis's ``is_jax``, so the + wrapper and the analysis it wraps can never disagree about the backend. An + explicit bool is honoured as before. """ + if use_jax is None: + use_jax = analysis.is_jax + super().__init__(use_jax=use_jax) self.analysis = analysis self.model = model + @property + def gradient_mode(self) -> str: + """ + The wrapped analysis's declared gradient mode (``Analysis.gradient_mode``), so a + gradient search fitting the wrapper differentiates the way the analysis asked. + """ + return self.analysis.gradient_mode + def __getattr__(self, item): if item in ("__getstate__", "__setstate__"): raise AttributeError(item) diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index b448e6618..848d5c80e 100644 --- a/autofit/non_linear/fitness.py +++ b/autofit/non_linear/fitness.py @@ -272,7 +272,7 @@ def __init__( self.use_jax_jit = use_jax_jit self.gradient_mode = validate_gradient_mode(gradient_mode) - if getattr(self.analysis, "_use_jax", False): + if self.analysis.is_jax: from autofit.jax.pytrees import enable_pytrees, register_model enable_pytrees() @@ -299,7 +299,7 @@ def __init__( from autofit.non_linear.quick_update import BackgroundQuickUpdate convert_jax = ( - getattr(self.analysis, "_use_jax", False) + self.analysis.is_jax and not getattr(self.analysis, "supports_jax_visualization", False) ) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 7963e61de..40a3ee616 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -447,7 +447,7 @@ def optimise( analysis = factor.analysis - uses_jax = getattr(analysis, "_use_jax", False) + uses_jax = analysis.is_jax self.logger.info( f"EP factor step [{factor.name}]: running the factor search " @@ -674,7 +674,7 @@ class represented by model M and gives a score for their fitness. """ self.check_model(model=model) - if getattr(analysis, "_use_jax", False): + if analysis.is_jax: try: import jax devices = jax.devices() @@ -856,6 +856,13 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res call_likelihood=(mode == 2), ) + # Runtime gates, after the test-mode bypass return so `PYAUTO_TEST_MODE>=2` + # (which never runs a backend) is not affected by them. + from autofit.graphical.declarative.collection import FactorGraphModel + + if isinstance(analysis, FactorGraphModel): + analysis.check_backend_agreement() + model.freeze() search_internal, fitness = self._fit( model=model, diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index d79ee2c93..e5b00db49 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -314,7 +314,7 @@ def _fit( # box makes the objective non-finite on either branch. Bounds are # passed on both — they are correct on both, and having them # diverge by branch would be a trap of its own. - if analysis._use_jax: + if analysis.is_jax: search_internal = optimize.minimize( fun=fitness._jit, diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index 5fd941a0a..acd6ae760 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -244,7 +244,7 @@ def _fit( 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, + use_jax_jit=analysis.is_jax and self.use_jax_jit, ) if not isinstance(self.paths, NullPaths): @@ -265,7 +265,7 @@ def _fit( while not finished: try: - if self.number_of_cores <= 1 or self.force_x1_cpu or analysis._use_jax: + if self.number_of_cores <= 1 or self.force_x1_cpu or analysis.is_jax: raise RuntimeError Pool = _fork_pool_cls() @@ -291,7 +291,7 @@ def _fit( except RuntimeError as e: if not checkpoint_exists: - if getattr(analysis, "_use_jax", False): + if analysis.is_jax: self.logger.info( "Running Dynesty with JAX-jitted likelihood (single CPU, no pool)." ) diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index fe2d344b0..9b55a651f 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -346,7 +346,7 @@ def _fit(self, model: AbstractPriorModel, analysis): "Starting new Nautilus non-linear search (no previous samples found)." ) - if self.force_x1_cpu or analysis._use_jax: + if self.force_x1_cpu or analysis.is_jax: fitness = Fitness( model=model, @@ -435,7 +435,7 @@ def fit_x1_cpu(self, fitness, model, analysis): the log likelihood the search maximizes. """ - if analysis._use_jax: + if analysis.is_jax: self.logger.info( "Running search with JAX vectorization (parallelization handled by JAX)." ) diff --git a/test_autofit/analysis/test_is_jax.py b/test_autofit/analysis/test_is_jax.py new file mode 100644 index 000000000..6477ce56a --- /dev/null +++ b/test_autofit/analysis/test_is_jax.py @@ -0,0 +1,241 @@ +""" +``Analysis.is_jax``, the single "is this analysis JAX?" probe, and how the wrappers +forward it (search-extensibility decision D12): + +- ``ModelAnalysis`` and ``AnalysisFactor`` take their wrapped analysis's ``is_jax`` and + ``gradient_mode``; +- ``FactorGraphModel`` derives ``is_jax`` from its factors, and whole-graph fitting + requires every factor to agree; per-factor expectation propagation allows mixed + factors; +- hierarchical factors carry the flag and honour ``PYAUTO_DISABLE_JAX``; +- the flags survive pytree flattening and pickling. +""" + +import importlib.util +import pickle + +import pytest + +import autofit as af +import autofit.graphical as g +from autofit import exc +from autofit.non_linear.analysis.model_analysis import ModelAnalysis + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +requires_jax = pytest.mark.skipif( + importlib.util.find_spec("jax") is None, reason="jax is not installed" +) + + +class NumpyAnalysis(af.Analysis): + def __init__(self, use_jax=False): + super().__init__(use_jax=use_jax) + + def log_likelihood_function(self, instance): + return -0.5 * float((instance.centre - 50.0) ** 2) + + +class ForwardAnalysis(NumpyAnalysis): + gradient_mode = "forward" + + +class NoInitAnalysis(af.Analysis): + """An analysis that never calls ``Analysis.__init__``.""" + + def __init__(self): + pass + + def log_likelihood_function(self, instance): + return 0.0 + + +def _factor(analysis, name): + return g.AnalysisFactor(af.Model(af.ex.Gaussian), analysis=analysis, name=name) + + +def test_is_jax_false_by_default(): + assert NumpyAnalysis().is_jax is False + + +def test_is_jax_false_without_analysis_init(): + assert NoInitAnalysis().is_jax is False + + +@requires_jax +def test_is_jax_true_when_use_jax(): + assert NumpyAnalysis(use_jax=True).is_jax is True + + +@requires_jax +def test_disable_jax_switch(monkeypatch): + monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") + + assert NumpyAnalysis(use_jax=True).is_jax is False + + +@requires_jax +def test_model_analysis_forwards_is_jax_and_gradient_mode(): + analysis = ForwardAnalysis(use_jax=True) + + wrapped = analysis.with_model(af.Model(af.ex.Gaussian)) + + assert isinstance(wrapped, ModelAnalysis) + assert wrapped.is_jax is True + assert wrapped.gradient_mode == "forward" + assert pickle.loads(pickle.dumps(wrapped)).is_jax is True + + +def test_model_analysis_explicit_use_jax_is_honoured(): + wrapped = ModelAnalysis( + analysis=NumpyAnalysis(), model=af.Model(af.ex.Gaussian), use_jax=False + ) + + assert wrapped.is_jax is False + assert wrapped.gradient_mode == "reverse" + + +@requires_jax +def test_analysis_factor_forwards_is_jax_and_gradient_mode(): + factor = _factor(ForwardAnalysis(use_jax=True), "forward") + + assert factor.is_jax is True + assert factor.gradient_mode == "forward" + assert _factor(NumpyAnalysis(), "numpy").is_jax is False + + +@requires_jax +def test_factor_graph_model_inherits_use_jax_from_its_factors(): + graph = g.FactorGraphModel( + _factor(NumpyAnalysis(use_jax=True), "a"), + _factor(NumpyAnalysis(use_jax=True), "b"), + ) + + assert graph.is_jax is True + assert graph.factors_disagreeing_on_backend() == [] + graph.check_backend_agreement() + + +def test_numpy_factor_graph_model_is_numpy(): + graph = g.FactorGraphModel( + _factor(NumpyAnalysis(), "a"), _factor(NumpyAnalysis(), "b") + ) + + assert graph.is_jax is False + graph.check_backend_agreement() + + +@requires_jax +def test_mixed_factor_graph_whole_graph_check_names_the_disagreeing_factors(): + graph = g.FactorGraphModel( + _factor(NumpyAnalysis(), "numpy_factor"), + _factor(NumpyAnalysis(use_jax=True), "jax_factor"), + ) + + assert graph.is_jax is False + assert graph.factors_disagreeing_on_backend() == ["jax_factor"] + + with pytest.raises(exc.SearchException, match="jax_factor"): + graph.check_backend_agreement() + + +@requires_jax +def test_explicit_graph_use_jax_must_agree_with_the_factors(): + graph = g.FactorGraphModel(_factor(NumpyAnalysis(), "numpy_factor"), use_jax=True) + + assert graph.is_jax is True + + with pytest.raises(exc.SearchException, match="numpy_factor"): + graph.check_backend_agreement() + + +@requires_jax +def test_whole_graph_fit_of_a_mixed_graph_raises(monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + graph = g.FactorGraphModel( + _factor(NumpyAnalysis(), "numpy_factor"), + _factor(NumpyAnalysis(use_jax=True), "jax_factor"), + ) + + with pytest.raises(exc.SearchException, match="jax_factor"): + af.Drawer(total_draws=2).fit(model=graph.global_prior_model, analysis=graph) + + +@requires_jax +@pytest.mark.filterwarnings("ignore::RuntimeWarning") +def test_mixed_factor_graph_still_runs_per_factor_ep(monkeypatch): + """ + Per-factor EP fits each factor with its own analysis, so a numpy factor and a JAX + factor in one graph each run on their own backend. + """ + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + graph = g.FactorGraphModel( + g.AnalysisFactor( + af.Model(af.ex.Gaussian), + analysis=NumpyAnalysis(), + optimiser=af.Drawer(total_draws=3), + name="numpy_factor", + ), + g.AnalysisFactor( + af.Model(af.ex.Gaussian), + analysis=NumpyAnalysis(use_jax=True), + optimiser=af.Drawer(total_draws=3), + name="jax_factor", + ), + ) + + result = graph.optimise( + af.LaplaceOptimiser(), + paths=False, + ep_history=af.EPHistory(kl_tol=None), + max_steps=1, + ) + + assert result is not None + + +def test_factor_graph_gradient_mode_is_the_common_mode_else_reverse(): + forward = g.FactorGraphModel( + _factor(ForwardAnalysis(), "a"), _factor(ForwardAnalysis(), "b") + ) + mixed = g.FactorGraphModel( + _factor(ForwardAnalysis(), "a"), _factor(NumpyAnalysis(), "b") + ) + + assert forward.gradient_mode == "forward" + assert mixed.gradient_mode == "reverse" + + +@requires_jax +def test_factor_graph_flags_survive_pytree_flattening_and_pickling(): + graph = g.FactorGraphModel(_factor(NumpyAnalysis(), "numpy_factor"), use_jax=True) + + children, aux_data = graph.tree_flatten() + rebuilt = g.FactorGraphModel.tree_unflatten(aux_data, children) + + assert rebuilt.is_jax is True + assert pickle.loads(pickle.dumps(graph)).is_jax is True + + +@requires_jax +def test_hierarchical_factor_carries_the_flag(monkeypatch): + def hierarchical(use_jax): + factor = g.HierarchicalFactor( + af.GaussianPrior, + mean=af.GaussianPrior(mean=0.0, sigma=1.0), + sigma=1.0, + use_jax=use_jax, + ) + factor.add_drawn_variable(af.GaussianPrior(mean=0.0, sigma=1.0)) + return factor + + assert all(factor.is_jax for factor in hierarchical(True).factors) + assert not any(factor.is_jax for factor in hierarchical(False).factors) + + monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") + + assert not any(factor.is_jax for factor in hierarchical(True).factors) + + +def test_supports_jax_visualization_reads_is_jax(): + assert NumpyAnalysis().supports_jax_visualization is False From 9a835bed64f8e7b0ddd64a8f9ccd7e08eb4acdbb Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 10:55:48 +0100 Subject: [PATCH 03/14] feat: fail-fast SearchException gate for jax_use='required' searches NonLinearSearch.start_resume_fit raises one shared SearchException (capabilities.JAX_REQUIRED_MESSAGE) when a JAX-required search gets a numpy analysis, before any backend state exists. The gate sits after the test-mode bypass return, so PYAUTO_DISABLE_JAX=1 + PYAUTO_TEST_MODE=2 smoke runs still complete. BlackJAXNUTS, SMC and the MultiStart searches replace their own ValueError checks with the same gate (error type changes to SearchException); NSS gains it. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- autofit/non_linear/search/abstract_search.py | 5 + .../search/mcmc/blackjax/nuts/search.py | 15 +-- .../search/mcmc/blackjax/smc/search.py | 14 +-- .../search/mle/multi_start_gradient/search.py | 9 +- autofit/non_linear/search/nest/nss/search.py | 4 + .../search/mcmc/test_blackjax_smc.py | 2 +- .../search/nest/nss/test_checkpoint.py | 1 + .../search/test_jax_required_gate.py | 107 ++++++++++++++++++ 8 files changed, 126 insertions(+), 31 deletions(-) create mode 100644 test_autofit/non_linear/search/test_jax_required_gate.py diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 40a3ee616..e352097d5 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -863,6 +863,11 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res if isinstance(analysis, FactorGraphModel): analysis.check_backend_agreement() + # Fail fast: a jax_use='required' search given a numpy analysis raises the one + # shared SearchException before any backend state exists. `PYAUTO_DISABLE_JAX=1` + # smoke runs use `PYAUTO_TEST_MODE>=2`, which returned above. + cap.check_jax_required(self, analysis) + model.freeze() search_internal, fitness = self._fit( model=model, diff --git a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py index 148591d4f..832defa0a 100644 --- a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py @@ -289,18 +289,9 @@ def _fit(self, model: AbstractPriorModel, analysis): import jax.numpy as jnp import blackjax - # JAX is mandatory: NUTS needs gradients of the log-density. Refuse - # cleanly if the analysis was built without ``use_jax=True``. - xp = getattr(analysis, "_xp", None) - if xp is None or not xp.__name__.startswith("jax"): - raise ValueError( - "BlackJAXNUTS requires an Analysis built with use_jax=True. " - "NUTS is a gradient-based sampler and the log-likelihood must " - "flow through jax.grad. Construct Analysis(..., use_jax=True) " - "and call enable_pytrees() / register_model(model) before fit. " - "See autofit_workspace_test/scripts/searches/BlackJAXNUTS.py " - "for a worked example." - ) + # JAX is mandatory (jax_use='required'): NUTS needs gradients of the + # log-density. `fit` already ran this gate; repeated for direct callers. + cap.check_jax_required(self, analysis) fitness = Fitness( model=model, diff --git a/autofit/non_linear/search/mcmc/blackjax/smc/search.py b/autofit/non_linear/search/mcmc/blackjax/smc/search.py index 979e4c77c..16e800896 100644 --- a/autofit/non_linear/search/mcmc/blackjax/smc/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/smc/search.py @@ -347,17 +347,9 @@ def _fit(self, model: AbstractPriorModel, analysis): import blackjax import blackjax.smc.resampling as resampling - # JAX is mandatory: MALA/HMC need gradients of the log-density. Refuse cleanly if the analysis was built - # without ``use_jax=True``. - xp = getattr(analysis, "_xp", None) - if xp is None or not xp.__name__.startswith("jax"): - raise ValueError( - "SMC requires an Analysis built with use_jax=True. Its inner kernel (MALA/HMC) is " - "gradient-based and the log-likelihood must flow through jax.grad. Construct " - "Analysis(..., use_jax=True) and call enable_pytrees() / register_model(model) before fit. " - "See autofit_workspace_test/scripts/searches/BlackJAXNUTS.py for a worked example of the " - "sibling BlackJAX search." - ) + # JAX is mandatory (jax_use='required'): MALA/HMC need gradients of the log-density. `fit` already ran + # this gate; repeated for direct callers. + cap.check_jax_required(self, analysis) n_dim = model.prior_count diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 3a4cc5bef..205dd2dbb 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -974,13 +974,8 @@ def _fit( "`pip install jax optax`." ) from e - if not getattr(analysis, "_use_jax", False): - raise ValueError( - f"{type(self).__name__} is a JAX-native gradient search and " - "requires a JAX-traceable Analysis (e.g. `AnalysisImaging(..., " - "use_jax=True)`). The supplied analysis is not running on the JAX " - "backend." - ) + # jax_use='required': `fit` already ran this gate; repeated for direct callers. + cap.check_jax_required(self, analysis) # ``fitness.call`` is differentiated inside ``jax.jit(jax.vmap(...))`` # below, so the per-evaluation quick-update counting that fires from diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index 5790046d3..b0d7c24df 100644 --- a/autofit/non_linear/search/nest/nss/search.py +++ b/autofit/non_linear/search/nest/nss/search.py @@ -416,6 +416,10 @@ def _fit(self, model: AbstractPriorModel, analysis): like latent-sample generation, which calls ``fitness.batch_size``. """ + # jax_use='required': NSS traces its likelihood inside blackjax and has no numpy + # path. `fit` already ran this gate; repeated for direct callers. + cap.check_jax_required(self, analysis) + import jax import jax.numpy as jnp import time diff --git a/test_autofit/non_linear/search/mcmc/test_blackjax_smc.py b/test_autofit/non_linear/search/mcmc/test_blackjax_smc.py index b44939147..cd7630c5d 100644 --- a/test_autofit/non_linear/search/mcmc/test_blackjax_smc.py +++ b/test_autofit/non_linear/search/mcmc/test_blackjax_smc.py @@ -540,7 +540,7 @@ def test__hmc_kernel_runs_and_holds_acceptance(): def test__non_jax_analysis_raises_a_clear_error(): search = af.SMC(num_particles=8, num_mcmc_steps=1, max_smc_steps=1) - with pytest.raises(ValueError, match="use_jax=True"): + with pytest.raises(af.exc.SearchException, match="use_jax=True"): search.fit( model=gaussian_target_model(), analysis=GaussianTargetAnalysis(), diff --git a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py index 5f32d72ea..0188df1d1 100644 --- a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py +++ b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py @@ -190,6 +190,7 @@ def test__load_checkpoint_called_when_file_exists(tmp_path): ) mock_analysis = SimpleNamespace( log_likelihood_function=lambda instance: 0.0, + is_jax=True, ) search = af.NSS(n_live=4, num_mcmc_steps=1, num_delete=1, termination=-3.0) diff --git a/test_autofit/non_linear/search/test_jax_required_gate.py b/test_autofit/non_linear/search/test_jax_required_gate.py new file mode 100644 index 000000000..e662f8d72 --- /dev/null +++ b/test_autofit/non_linear/search/test_jax_required_gate.py @@ -0,0 +1,107 @@ +""" +The fail-fast gate for ``jax_use='required'`` searches (search-extensibility D2). + +``NonLinearSearch.start_resume_fit`` raises one shared ``SearchException`` +(``capabilities.JAX_REQUIRED_MESSAGE``) when a JAX-required search is given a numpy +analysis, before the backend runs. The gate sits after the test-mode bypass return, so +``PYAUTO_DISABLE_JAX=1`` + ``PYAUTO_TEST_MODE=2`` smoke runs still complete. + +No backend executes here: the gate raises before ``_fit``, and test mode 2 bypasses it, +so these cases run on every CI leg (``NSS`` alone needs blackjax to be constructed). +""" + +import pytest + +import autofit as af +from autofit import exc +from autofit.non_linear.search import capabilities as cap + +from test_autofit.non_linear.search.conformance_roster import searches_under_test + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +REQUIRED = [entry for entry in searches_under_test() if entry.jax_use == "required"] + + +class GaussianAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + return -0.5 * float( + (instance.centre - 50.0) ** 2 + + (instance.normalization - 25.0) ** 2 + + (instance.sigma - 10.0) ** 2 + ) + + +def _search(entry): + missing = entry.missing_requirements() + if missing: + pytest.skip(f"{entry.name} cannot be constructed without {missing}") + return entry.resolve()() + + +def test_seven_searches_require_jax(): + assert sorted(entry.name for entry in REQUIRED) == sorted( + [ + "BlackJAXNUTS", + "SMC", + "NSS", + "MultiStartAdam", + "MultiStartADABelief", + "MultiStartLion", + "MultiStartProdigy", + ] + ) + + +@pytest.mark.parametrize("entry", REQUIRED, ids=[entry.name for entry in REQUIRED]) +def test_required_search_with_numpy_analysis_raises_the_shared_message( + entry, monkeypatch +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + search = _search(entry) + + def backend_must_not_run(*args, **kwargs): + raise AssertionError("the backend ran before the gate") + + monkeypatch.setattr(search, "_fit", backend_must_not_run) + + with pytest.raises(exc.SearchException) as error: + search.fit(model=af.Model(af.ex.Gaussian), analysis=GaussianAnalysis()) + + assert str(error.value) == cap.JAX_REQUIRED_MESSAGE.format(search=entry.name) + + +@pytest.mark.parametrize("entry", REQUIRED, ids=[entry.name for entry in REQUIRED]) +def test_required_search_completes_with_disable_jax_and_test_mode_2( + entry, monkeypatch +): + monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") + monkeypatch.setenv("PYAUTO_TEST_MODE", "2") + search = _search(entry) + analysis = GaussianAnalysis(use_jax=True) + + assert analysis.is_jax is False + + result = search.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert isinstance(result, af.Result) + + +@pytest.mark.parametrize("entry", REQUIRED, ids=[entry.name for entry in REQUIRED]) +def test_direct_fit_call_raises_the_shared_message(entry, monkeypatch): + """ + The JAX-native searches repeat the gate at the top of ``_fit``, so a caller that + bypasses ``fit`` gets the same error rather than a tracer error deep in the backend. + """ + for module in entry.jax_modules: + pytest.importorskip(module) + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + search = _search(entry) + + with pytest.raises(exc.SearchException, match=entry.name): + search._fit(model=af.Model(af.ex.Gaussian), analysis=GaussianAnalysis()) + + +def test_optional_search_with_numpy_analysis_passes_the_gate(): + cap.check_jax_required(af.Drawer(), GaussianAnalysis()) + cap.check_jax_required(af.DynestyStatic(), GaussianAnalysis()) From baf2becee56c57e8e21bda0bda6b41f1e639bd86 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 10:58:21 +0100 Subject: [PATCH 04/14] fix: Nautilus(force_x1_cpu=True) runs with a numpy analysis fit_x1_cpu built its Fitness with use_jax_vmap=self.use_jax_vmap (default True) whatever the analysis, so a numpy likelihood was traced through jax.vmap and raised TracerArrayConversionError. The vectorised path is now requested only when analysis.is_jax. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- .../non_linear/search/nest/nautilus/search.py | 4 ++- .../search/nest/test_nautilus_force_x1_cpu.py | 34 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) create mode 100644 test_autofit/non_linear/search/nest/test_nautilus_force_x1_cpu.py diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 9b55a651f..018d7683b 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -357,7 +357,9 @@ def _fit(self, model: AbstractPriorModel, analysis): 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_vmap=self.use_jax_vmap, + # The vectorised (jax.vmap) likelihood needs a JAX analysis; + # `force_x1_cpu` with a numpy analysis evaluates point by point. + use_jax_vmap=self.use_jax_vmap and analysis.is_jax, batch_size=self.n_batch, ) diff --git a/test_autofit/non_linear/search/nest/test_nautilus_force_x1_cpu.py b/test_autofit/non_linear/search/nest/test_nautilus_force_x1_cpu.py new file mode 100644 index 000000000..57fdabd44 --- /dev/null +++ b/test_autofit/non_linear/search/nest/test_nautilus_force_x1_cpu.py @@ -0,0 +1,34 @@ +""" +``Nautilus(force_x1_cpu=True)`` with a numpy analysis runs on one CPU without the JAX +vectorised likelihood (search-extensibility phase A1). + +``force_x1_cpu`` routes the fit through ``fit_x1_cpu``, whose ``Fitness`` was built +with ``use_jax_vmap=self.use_jax_vmap`` (default ``True``) whatever the analysis, so a +numpy analysis was batched through ``jax.vmap`` and crashed. The vectorised path is +now requested only when the analysis is JAX. +""" + +import numpy as np +import pytest + +import autofit as af + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +class NumpyAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + vector = np.array([instance.centre, instance.normalization, instance.sigma]) + return float(-0.5 * np.sum((vector - np.array([50.0, 25.0, 10.0])) ** 2)) + + +def test_force_x1_cpu_runs_with_a_numpy_analysis(monkeypatch): + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + search = af.Nautilus(force_x1_cpu=True) + + assert search.use_jax_vmap is True + + result = search.fit(model=af.Model(af.ex.Gaussian), analysis=NumpyAnalysis()) + + assert isinstance(result, af.Result) + assert len(result.samples.sample_list) > 0 From f4404933ecfe27d83c48d3b6689e9cf04060778c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 10:59:00 +0100 Subject: [PATCH 05/14] feat: declarative lazy search registry and search-manifest@1 autofit/non_linear/search/registry.py lists the 15 public searches as data (class_path string, family, lazy, requires, the mirrored capability attributes, example and integration_test anchors) and never imports a search module. python -m autofit search-manifest --json publishes it as the versioned search-manifest@1, the only cross-repo format. Tests: completeness against the autofit exports (deleting an entry fails), entry == class attributes (skipped per entry when its backend is absent), manifest round trip, and no search module or optional backend imported. The conformance roster now derives jax_use/jax_modules from the registry; its golden table is unchanged. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- autofit/__main__.py | 58 ++ autofit/non_linear/search/registry.py | 518 ++++++++++++++++++ .../non_linear/search/conformance_roster.py | 90 ++- .../non_linear/search/test_registry.py | 208 +++++++ 4 files changed, 855 insertions(+), 19 deletions(-) create mode 100644 autofit/__main__.py create mode 100644 autofit/non_linear/search/registry.py create mode 100644 test_autofit/non_linear/search/test_registry.py diff --git a/autofit/__main__.py b/autofit/__main__.py new file mode 100644 index 000000000..976cc7fe6 --- /dev/null +++ b/autofit/__main__.py @@ -0,0 +1,58 @@ +""" +Command-line entry point: ``python -m autofit ``. + +Commands +-------- +search-manifest [--json] + Print the declarative search registry. ``--json`` emits the versioned manifest + (``search-manifest@1``), the only cross-repo format for search facts; without it a + plain-text table is printed. +""" + +import argparse +import json +import sys + + +def _search_manifest(arguments) -> int: + from autofit.non_linear.search.registry import manifest + + data = manifest() + + if arguments.json: + json.dump(data, sys.stdout, indent=2) + sys.stdout.write("\n") + return 0 + + print(f"{data['schema']} (autofit {data['autofit_version']})") + for search in data["searches"]: + capabilities = search["capabilities"] + print( + f"{search['name']:<20} {search['family']:<5} " + f"jax={capabilities['jax_use']:<9} gradient={capabilities['gradient']:<5} " + f"posterior={capabilities['posterior_kind']:<9} " + f"evidence={capabilities['produces_evidence']!s:<5} " + f"status={capabilities['status']}" + ) + return 0 + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(prog="python -m autofit") + commands = parser.add_subparsers(dest="command", required=True) + + search_manifest = commands.add_parser( + "search-manifest", + help="Print the declarative search registry (search-manifest@1).", + ) + search_manifest.add_argument( + "--json", action="store_true", help="Emit the versioned JSON manifest." + ) + search_manifest.set_defaults(handler=_search_manifest) + + arguments = parser.parse_args(argv) + return arguments.handler(arguments) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/autofit/non_linear/search/registry.py b/autofit/non_linear/search/registry.py new file mode 100644 index 000000000..dd563a05b --- /dev/null +++ b/autofit/non_linear/search/registry.py @@ -0,0 +1,518 @@ +""" +The declarative registry of PyAutoFit's public non-linear searches. + +One entry per search, written as **data**: the class is named by a ``class_path`` +string and is never imported by this module, so ``import autofit`` never imports a +search backend through it (``autofit.__init__._LAZY_ATTRS`` and the no-JAX CI leg keep +holding). Each entry mirrors the search's static capability class attributes +(``autofit.non_linear.search.capabilities``) in their JSON form; a unit test resolves +every class whose dependencies are installed and asserts entry == class attributes, and +a completeness test asserts every search ``autofit`` exports has an entry. + +The registry is published as a versioned JSON manifest, + + python -m autofit search-manifest --json + +whose schema is ``search-manifest@1`` (``MANIFEST_SCHEMA``). The manifest is the only +cross-repo format: the generated docs (``docs/_generate_searches.py``), workspaces and +assistants read it, never this module's Python objects. + +The capabilities are **static** (a property of the class). What a particular run can +deliver may be narrower: for example ``SMC`` estimates the evidence only when it starts +from a prior-sampling initializer. Effective run capabilities are documented, not +recorded here. + +Adding a search costs one entry here plus its ``autofit`` export; the completeness +test fails until both exist. +""" + +import importlib +import importlib.util +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Optional, Tuple + +from autofit.non_linear.search.capabilities import CAPABILITY_ATTRIBUTES + +MANIFEST_SCHEMA = "search-manifest@1" +""" +The manifest schema identifier. Bump the suffix on any incompatible change to the +manifest layout (a removed or renamed key, a changed value type); adding a key is +compatible. +""" + +FAMILIES = { + "mcmc": "Markov chain Monte Carlo", + "nest": "Nested sampling", + "mle": "Maximum likelihood / maximum a posteriori", +} +""" +The search families, in documentation order, with their display titles. +""" + +_WORKSPACE = "autofit_workspace:scripts/searches" +_WORKSPACE_TEST = "autofit_workspace_test:scripts/searches" + + +@dataclass(frozen=True) +class RegistryEntry: + """ + One public search. + + Parameters + ---------- + name + The public short name, as exported by ``autofit``. + class_path + The dotted ``module.Class`` path of the search class. + family + ``mcmc``, ``nest`` or ``mle`` (a key of ``FAMILIES``). + lazy + Whether ``autofit`` exports the class lazily (``_LAZY_ATTRS``), because its + module imports a heavy optional backend. + requires + The top-level modules the search's backend needs to run. + capabilities + The search's static capability attributes in JSON form, keyed and ordered by + ``CAPABILITY_ATTRIBUTES``. + example + A runnable example anchor, ``repo:path#section``, or ``None``. + integration_test + The integration-test script anchor, ``repo:path``, or ``None``. + """ + + name: str + class_path: str + family: str + lazy: bool + requires: Tuple[str, ...] + capabilities: Mapping[str, Any] + example: Optional[str] = None + integration_test: Optional[str] = None + + def missing_requirements(self) -> List[str]: + """ + The modules in ``requires`` that are not installed. + """ + return [ + module + for module in self.requires + if importlib.util.find_spec(module) is None + ] + + def resolve(self): + """ + Import and return the search class (imports its module). + """ + module_name, _, attribute = self.class_path.rpartition(".") + return getattr(importlib.import_module(module_name), attribute) + + def to_dict(self) -> Dict[str, Any]: + """ + The entry as it appears in the manifest. + """ + return { + "name": self.name, + "class_path": self.class_path, + "family": self.family, + "lazy": self.lazy, + "requires": list(self.requires), + "example": self.example, + "integration_test": self.integration_test, + "capabilities": { + name: _plain(self.capabilities[name]) for name in CAPABILITY_ATTRIBUTES + }, + } + + +def _plain(value): + if isinstance(value, Mapping): + return {key: _plain(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain(item) for item in value] + return value + + +def _capabilities( + *, + jax_use: str, + gradient: str, + batched: bool, + honours_gradient_mode: bool, + posterior_kind: str, + produces_evidence: bool, + resumable: bool, + warm_start: str, + install_extra: str, + upstream_url: str, + citation_keys: Tuple[str, ...], + status: str, + test_mode_budget: Dict[str, Any], + objective_target: Tuple[str, str], + invalid_value: str, +) -> Dict[str, Any]: + quantity, space = objective_target + values = dict( + jax_use=jax_use, + gradient=gradient, + batched=batched, + honours_gradient_mode=honours_gradient_mode, + posterior_kind=posterior_kind, + produces_evidence=produces_evidence, + resumable=resumable, + warm_start=warm_start, + install_extra=install_extra, + upstream_url=upstream_url, + citation_keys=list(citation_keys), + status=status, + test_mode_budget=dict(test_mode_budget), + objective_target={"quantity": quantity, "space": space}, + invalid_value=invalid_value, + ) + return {name: values[name] for name in CAPABILITY_ATTRIBUTES} + + +_SEARCH = "autofit.non_linear.search" +_MULTI_START = f"{_SEARCH}.mle.multi_start_gradient.search" + +_BLACKJAX = "https://github.com/blackjax-devs/blackjax" +_DYNESTY = "https://github.com/joshspeagle/dynesty" +_SCIPY = "https://github.com/scipy/scipy" +_OPTAX = "https://github.com/google-deepmind/optax" + +_EMCEE_ZEUS_BUDGET = {"nwalkers": 20, "nsteps": 10} +_MULTI_START_BUDGET = {"convergence.window": 1, "convergence.min_steps": 1} + + +def _dynesty(name: str, module: str) -> RegistryEntry: + return RegistryEntry( + name=name, + class_path=f"{_SEARCH}.nest.dynesty.search.{module}.{name}", + family="nest", + lazy=False, + requires=("dynesty",), + capabilities=_capabilities( + jax_use="optional", + gradient="none", + batched=False, + honours_gradient_mode=False, + posterior_kind="weighted", + produces_evidence=True, + resumable=True, + warm_start="provider", + install_extra="", + upstream_url=_DYNESTY, + citation_keys=("dynesty",), + status="stable", + test_mode_budget={"maxcall": 1}, + objective_target=("log_likelihood", "unit_cube"), + invalid_value="-1e+99", + ), + example=f"{_WORKSPACE}/nest.py#Search: {name}", + integration_test=f"{_WORKSPACE_TEST}/{name}.py", + ) + + +def _bfgs(name: str) -> RegistryEntry: + return RegistryEntry( + name=name, + class_path=f"{_SEARCH}.mle.bfgs.search.{name}", + family="mle", + lazy=False, + requires=("scipy",), + capabilities=_capabilities( + jax_use="optional", + gradient="none", + batched=False, + honours_gradient_mode=False, + posterior_kind="point", + produces_evidence=False, + resumable=False, + warm_start="provider", + install_extra="", + upstream_url=_SCIPY, + citation_keys=("scipy",), + status="stable", + test_mode_budget={}, + objective_target=("neg2_log_posterior", "physical"), + invalid_value="-inf", + ), + example=f"{_WORKSPACE}/mle.py#Search: LBFGS" if name == "LBFGS" else None, + integration_test=f"{_WORKSPACE_TEST}/{name}.py", + ) + + +def _multi_start( + name: str, + status: str = "stable", + upstream_url: str = _OPTAX, + citation_keys: Tuple[str, ...] = ("optax",), + example: Optional[str] = None, +) -> RegistryEntry: + return RegistryEntry( + name=name, + class_path=f"{_MULTI_START}.{name}", + family="mle", + lazy=False, + requires=("jax", "optax"), + capabilities=_capabilities( + jax_use="required", + gradient="uses", + batched=True, + honours_gradient_mode=True, + posterior_kind="point", + produces_evidence=False, + resumable=True, + warm_start="provider", + install_extra="", + upstream_url=upstream_url, + citation_keys=citation_keys, + status=status, + test_mode_budget=_MULTI_START_BUDGET, + objective_target=("neg2_log_posterior", "physical"), + invalid_value="-inf", + ), + example=example, + integration_test=f"{_WORKSPACE_TEST}/{name}.py", + ) + + +SEARCHES: Tuple[RegistryEntry, ...] = ( + RegistryEntry( + name="Emcee", + class_path=f"{_SEARCH}.mcmc.emcee.search.Emcee", + family="mcmc", + lazy=False, + requires=("emcee",), + capabilities=_capabilities( + jax_use="none", + gradient="none", + batched=False, + honours_gradient_mode=False, + posterior_kind="chain", + produces_evidence=False, + resumable=False, + warm_start="consumer", + install_extra="", + upstream_url="https://github.com/dfm/emcee", + citation_keys=("emcee",), + status="stable", + test_mode_budget=_EMCEE_ZEUS_BUDGET, + objective_target=("log_posterior", "physical"), + invalid_value="-inf", + ), + example=f"{_WORKSPACE}/mcmc.py#Search: Emcee", + integration_test=f"{_WORKSPACE_TEST}/Emcee.py", + ), + RegistryEntry( + name="Zeus", + class_path=f"{_SEARCH}.mcmc.zeus.search.Zeus", + family="mcmc", + lazy=False, + requires=("zeus",), + capabilities=_capabilities( + jax_use="none", + gradient="none", + batched=False, + honours_gradient_mode=False, + posterior_kind="chain", + produces_evidence=False, + resumable=False, + warm_start="consumer", + install_extra="optional", + upstream_url="https://github.com/minaskar/zeus", + citation_keys=("zeus1", "zeus2"), + status="stable", + test_mode_budget=_EMCEE_ZEUS_BUDGET, + objective_target=("log_posterior", "physical"), + invalid_value="-inf", + ), + example=f"{_WORKSPACE}/mcmc.py#Search: Zeus", + integration_test=f"{_WORKSPACE_TEST}/Zeus.py", + ), + RegistryEntry( + name="BlackJAXNUTS", + class_path=f"{_SEARCH}.mcmc.blackjax.nuts.search.BlackJAXNUTS", + family="mcmc", + lazy=False, + requires=("jax", "blackjax"), + capabilities=_capabilities( + jax_use="required", + gradient="uses", + batched=True, + honours_gradient_mode=False, + posterior_kind="chain", + produces_evidence=False, + resumable=False, + warm_start="consumer", + install_extra="optional", + upstream_url=_BLACKJAX, + citation_keys=("blackjax",), + status="stable", + test_mode_budget={"num_warmup": 20, "num_samples": 20, "num_chains": 2}, + objective_target=("log_posterior", "physical"), + invalid_value="-inf", + ), + example=f"{_WORKSPACE}/mcmc.py#Search: BlackJAXNUTS", + integration_test=f"{_WORKSPACE_TEST}/BlackJAXNUTS.py", + ), + RegistryEntry( + name="SMC", + class_path=f"{_SEARCH}.mcmc.blackjax.smc.search.SMC", + family="mcmc", + lazy=True, + requires=("jax", "blackjax"), + capabilities=_capabilities( + jax_use="required", + gradient="uses", + batched=True, + honours_gradient_mode=False, + posterior_kind="weighted", + produces_evidence=True, + resumable=False, + warm_start="provider", + install_extra="optional", + upstream_url=_BLACKJAX, + citation_keys=("blackjax",), + status="experimental", + test_mode_budget={ + "num_particles": 16, + "num_mcmc_steps": 2, + "max_smc_steps": 5, + }, + objective_target=("log_likelihood", "physical"), + invalid_value="-1e+99", + ), + example=None, + integration_test=f"{_WORKSPACE_TEST}/SMC.py", + ), + _dynesty("DynestyStatic", "static"), + _dynesty("DynestyDynamic", "dynamic"), + RegistryEntry( + name="Nautilus", + class_path=f"{_SEARCH}.nest.nautilus.search.Nautilus", + family="nest", + lazy=False, + requires=("nautilus",), + capabilities=_capabilities( + jax_use="optional", + gradient="none", + batched=True, + honours_gradient_mode=False, + posterior_kind="weighted", + produces_evidence=True, + resumable=True, + warm_start="provider", + install_extra="optional", + upstream_url="https://github.com/johannesulf/nautilus", + citation_keys=("nautilus",), + status="stable", + test_mode_budget={"n_like_max": 1}, + objective_target=("log_likelihood", "unit_cube"), + invalid_value="-1e+99", + ), + example=f"{_WORKSPACE}/nest.py#Search: Nautilus", + integration_test=f"{_WORKSPACE_TEST}/Nautilus.py", + ), + RegistryEntry( + name="NSS", + class_path=f"{_SEARCH}.nest.nss.search.NSS", + family="nest", + lazy=True, + requires=("jax", "blackjax"), + capabilities=_capabilities( + jax_use="required", + gradient="none", + batched=True, + honours_gradient_mode=False, + posterior_kind="weighted", + produces_evidence=True, + resumable=True, + warm_start="provider", + install_extra="optional", + upstream_url=_BLACKJAX, + citation_keys=("blackjax",), + status="experimental", + test_mode_budget={"termination": -1.0}, + objective_target=("log_likelihood", "unit_cube"), + invalid_value="-1e+30", + ), + example=None, + integration_test=f"{_WORKSPACE_TEST}/NSS.py", + ), + RegistryEntry( + name="Drawer", + class_path=f"{_SEARCH}.mle.drawer.search.Drawer", + family="mle", + lazy=False, + requires=(), + capabilities=_capabilities( + jax_use="none", + gradient="none", + batched=False, + honours_gradient_mode=False, + posterior_kind="point", + produces_evidence=False, + resumable=False, + warm_start="neutral", + install_extra="", + upstream_url="https://github.com/PyAutoLabs/PyAutoFit", + citation_keys=(), + status="stable", + test_mode_budget={}, + objective_target=("log_posterior", "physical"), + invalid_value="-inf", + ), + example=f"{_WORKSPACE}/mle.py#Search: Drawer", + integration_test=f"{_WORKSPACE_TEST}/Drawer.py", + ), + _bfgs("BFGS"), + _bfgs("LBFGS"), + _multi_start("MultiStartAdam", example=f"{_WORKSPACE}/mle.py#Search: MultiStartAdam"), + _multi_start("MultiStartADABelief", status="experimental"), + _multi_start("MultiStartLion", status="experimental"), + _multi_start( + "MultiStartProdigy", + upstream_url="https://github.com/konstmish/prodigy", + citation_keys=("optax", "prodigy"), + ), +) +""" +Every public search, in documentation order (by family, then by name order above). +""" + + +def entries() -> Tuple[RegistryEntry, ...]: + """ + Every registered search. + """ + return SEARCHES + + +def entry(name: str) -> RegistryEntry: + """ + The registry entry of the search exported as ``autofit.``. + """ + for search_entry in SEARCHES: + if search_entry.name == name: + return search_entry + raise KeyError(f"No search named {name!r} is registered") + + +def manifest() -> Dict[str, Any]: + """ + The versioned search manifest (``search-manifest@1``), the only cross-repo format. + + Keys: ``schema``, ``autofit_version``, ``families`` (family key -> title), + ``capability_attributes`` (the capability keys, in order) and ``searches`` (one + object per search: ``name``, ``class_path``, ``family``, ``lazy``, ``requires``, + ``example``, ``integration_test`` and ``capabilities``). + """ + import autofit + + return { + "schema": MANIFEST_SCHEMA, + "autofit_version": autofit.__version__, + "families": dict(FAMILIES), + "capability_attributes": list(CAPABILITY_ATTRIBUTES), + "searches": [search_entry.to_dict() for search_entry in SEARCHES], + } diff --git a/test_autofit/non_linear/search/conformance_roster.py b/test_autofit/non_linear/search/conformance_roster.py index 04819b35b..0e6d90687 100644 --- a/test_autofit/non_linear/search/conformance_roster.py +++ b/test_autofit/non_linear/search/conformance_roster.py @@ -29,7 +29,10 @@ capability the roster declares for it: ``jax_use`` is ``"none"`` (the search never touches JAX), ``"optional"`` (the search runs on a numpy likelihood and has a JAX path) or ``"required"`` (the backend is JAX-native), and ``jax_modules`` names the modules a -``"required"`` backend cannot run without. +``"required"`` backend cannot run without. Both are derived from the declarative search +registry (``autofit/non_linear/search/registry.py``), which mirrors each class's +``jax_use`` attribute and is tested equal to it; reading the registry imports no search +module, so the roster still collects on every leg. Golden table ------------ @@ -130,9 +133,18 @@ def resolve(self): _MULTI_START = f"{_SEARCH}.mle.multi_start_gradient.search" -_BLACKJAX = dict(jax_use="required", jax_modules=("jax", "blackjax")) -_OPTAX = dict(jax_use="required", jax_modules=("jax", "optax")) -_OPTIONAL = dict(jax_use="optional") +def _declared(name: str) -> dict: + """ + The ``jax_use`` / ``jax_modules`` of a search, derived from its registry entry. + """ + from autofit.non_linear.search.registry import entry + + registered = entry(name) + jax_use = registered.capabilities["jax_use"] + return dict( + jax_use=jax_use, + jax_modules=registered.requires if jax_use == "required" else (), + ) def searches_under_test() -> List[SearchEntry]: @@ -140,56 +152,96 @@ def searches_under_test() -> List[SearchEntry]: The 15 public searches of PyAutoFit, in a fixed order. """ return [ - SearchEntry("Emcee", f"{_SEARCH}.mcmc.emcee.search.Emcee", "mcmc"), - SearchEntry("Zeus", f"{_SEARCH}.mcmc.zeus.search.Zeus", "mcmc"), + SearchEntry( + "Emcee", + f"{_SEARCH}.mcmc.emcee.search.Emcee", + "mcmc", + **_declared("Emcee"), + ), + SearchEntry( + "Zeus", + f"{_SEARCH}.mcmc.zeus.search.Zeus", + "mcmc", + **_declared("Zeus"), + ), SearchEntry( "BlackJAXNUTS", f"{_SEARCH}.mcmc.blackjax.nuts.search.BlackJAXNUTS", "mcmc", - **_BLACKJAX, + **_declared("BlackJAXNUTS"), ), SearchEntry( - "SMC", f"{_SEARCH}.mcmc.blackjax.smc.search.SMC", "mcmc", **_BLACKJAX + "SMC", + f"{_SEARCH}.mcmc.blackjax.smc.search.SMC", + "mcmc", + **_declared("SMC"), ), SearchEntry( "DynestyStatic", f"{_SEARCH}.nest.dynesty.search.static.DynestyStatic", "nest", - **_OPTIONAL, + **_declared("DynestyStatic"), ), SearchEntry( "DynestyDynamic", f"{_SEARCH}.nest.dynesty.search.dynamic.DynestyDynamic", "nest", - **_OPTIONAL, + **_declared("DynestyDynamic"), ), SearchEntry( - "Nautilus", f"{_SEARCH}.nest.nautilus.search.Nautilus", "nest", **_OPTIONAL + "Nautilus", + f"{_SEARCH}.nest.nautilus.search.Nautilus", + "nest", + **_declared("Nautilus"), ), SearchEntry( "NSS", f"{_SEARCH}.nest.nss.search.NSS", "nest", requires=["blackjax"], - **_BLACKJAX, + **_declared("NSS"), ), - SearchEntry("Drawer", f"{_SEARCH}.mle.drawer.search.Drawer", "mle"), - SearchEntry("BFGS", f"{_SEARCH}.mle.bfgs.search.BFGS", "mle", **_OPTIONAL), - SearchEntry("LBFGS", f"{_SEARCH}.mle.bfgs.search.LBFGS", "mle", **_OPTIONAL), SearchEntry( - "MultiStartAdam", f"{_MULTI_START}.MultiStartAdam", "mle", **_OPTAX + "Drawer", + f"{_SEARCH}.mle.drawer.search.Drawer", + "mle", + **_declared("Drawer"), + ), + SearchEntry( + "BFGS", + f"{_SEARCH}.mle.bfgs.search.BFGS", + "mle", + **_declared("BFGS"), + ), + SearchEntry( + "LBFGS", + f"{_SEARCH}.mle.bfgs.search.LBFGS", + "mle", + **_declared("LBFGS"), + ), + SearchEntry( + "MultiStartAdam", + f"{_MULTI_START}.MultiStartAdam", + "mle", + **_declared("MultiStartAdam"), ), SearchEntry( "MultiStartADABelief", f"{_MULTI_START}.MultiStartADABelief", "mle", - **_OPTAX, + **_declared("MultiStartADABelief"), ), SearchEntry( - "MultiStartLion", f"{_MULTI_START}.MultiStartLion", "mle", **_OPTAX + "MultiStartLion", + f"{_MULTI_START}.MultiStartLion", + "mle", + **_declared("MultiStartLion"), ), SearchEntry( - "MultiStartProdigy", f"{_MULTI_START}.MultiStartProdigy", "mle", **_OPTAX + "MultiStartProdigy", + f"{_MULTI_START}.MultiStartProdigy", + "mle", + **_declared("MultiStartProdigy"), ), ] diff --git a/test_autofit/non_linear/search/test_registry.py b/test_autofit/non_linear/search/test_registry.py new file mode 100644 index 000000000..da997192f --- /dev/null +++ b/test_autofit/non_linear/search/test_registry.py @@ -0,0 +1,208 @@ +""" +The declarative search registry (``autofit/non_linear/search/registry.py``) and the +versioned ``search-manifest@1``. + +- completeness: every search ``autofit`` exports is registered, and every entry is + exported, so deleting an entry (or adding an export without one) fails; +- entry == class attributes, skipped per entry only when a module in its ``requires`` + is not installed; +- the registry is data: building the manifest imports no search module and no optional + backend; +- ``python -m autofit search-manifest --json`` emits the manifest. + +None of these tests needs a sibling repository; the ``example`` / ``integration_test`` +anchors are only checked for their format here. +""" + +import inspect +import json +import re +import subprocess +import sys + +import pytest + +import autofit as af +from autofit.non_linear.search import registry +from autofit.non_linear.search.abstract_search import NonLinearSearch +from autofit.non_linear.search.capabilities import CAPABILITY_ATTRIBUTES, capabilities_from + +from test_autofit.non_linear.search.conformance_roster import searches_under_test + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +_SEARCH_PACKAGE = "autofit.non_linear.search." + +OPTIONAL_BACKENDS = ( + "jax", + "jaxlib", + "blackjax", + "optax", + "prodigyopt", + "nautilus", + "dynesty", + "emcee", + "zeus", +) + + +def exported_search_names(): + """ + The names ``autofit`` exports for concrete non-linear searches. Lazy exports are + recognised by their module path, so this never imports a lazy search module. + """ + names = { + name + for name, (module, _) in af._LAZY_ATTRS.items() + if module.startswith(_SEARCH_PACKAGE) + } + for name, value in vars(af).items(): + if ( + inspect.isclass(value) + and issubclass(value, NonLinearSearch) + and not inspect.isabstract(value) + and value.__module__.startswith(_SEARCH_PACKAGE) + ): + names.add(name) + return names + + +def _params(): + return [pytest.param(entry, id=entry.name) for entry in registry.entries()] + + +def test_every_exported_search_is_registered(): + registered = {entry.name for entry in registry.entries()} + + assert exported_search_names() == registered + assert len(registered) == 15 + + +def test_deleting_an_entry_fails_completeness(monkeypatch): + monkeypatch.setattr(registry, "SEARCHES", registry.SEARCHES[1:]) + + assert exported_search_names() != {entry.name for entry in registry.entries()} + + +def test_registry_and_conformance_roster_agree(): + roster = {entry.name: entry for entry in searches_under_test()} + + assert set(roster) == {entry.name for entry in registry.entries()} + for entry in registry.entries(): + assert roster[entry.name].class_path == entry.class_path + assert roster[entry.name].family == entry.family + assert roster[entry.name].jax_use == entry.capabilities["jax_use"] + + +@pytest.mark.parametrize("entry", _params()) +def test_entry_names_its_export(entry): + if entry.lazy: + module, attribute = af._LAZY_ATTRS[entry.name] + assert f"{module}.{attribute}" == entry.class_path + else: + assert entry.name not in af._LAZY_ATTRS + cls = vars(af)[entry.name] + assert f"{cls.__module__}.{cls.__qualname__}" == entry.class_path + + +@pytest.mark.parametrize("entry", _params()) +def test_entry_equals_class_attributes(entry): + missing = entry.missing_requirements() + if missing: + pytest.skip(f"{entry.name} needs {missing}, which is not installed") + + cls = entry.resolve() + + assert cls.__name__ == entry.name + assert capabilities_from(cls) == dict(entry.to_dict()["capabilities"]) + + +@pytest.mark.parametrize("entry", _params()) +def test_entry_is_well_formed(entry): + assert entry.family in registry.FAMILIES + assert list(entry.capabilities) == list(CAPABILITY_ATTRIBUTES) + if entry.example is not None: + assert re.fullmatch(r"\w+:scripts/[\w/]+\.py#Search: \w+", entry.example) + if entry.integration_test is not None: + assert re.fullmatch(r"\w+:scripts/[\w/]+\.py", entry.integration_test) + + +def test_manifest_schema_and_round_trip(): + manifest = registry.manifest() + + assert manifest["schema"] == registry.MANIFEST_SCHEMA == "search-manifest@1" + assert manifest["autofit_version"] == af.__version__ + assert manifest["capability_attributes"] == list(CAPABILITY_ATTRIBUTES) + assert list(manifest["families"]) == ["mcmc", "nest", "mle"] + assert [search["name"] for search in manifest["searches"]] == [ + entry.name for entry in registry.entries() + ] + assert json.loads(json.dumps(manifest)) == manifest + + for search in manifest["searches"]: + assert set(search) == { + "name", + "class_path", + "family", + "lazy", + "requires", + "example", + "integration_test", + "capabilities", + } + + +def _run_python(code): + return subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + check=True, + ).stdout + + +def test_manifest_imports_no_search_module_and_no_backend(): + """ + ``import autofit`` plus building the manifest imports neither the lazy search + modules nor any optional sampler backend: the registry is data. + """ + output = _run_python( + "import sys, json\n" + "import autofit\n" + "from autofit.non_linear.search import registry\n" + "registry.manifest()\n" + "print(json.dumps(sorted(sys.modules)))\n" + ) + modules = set(json.loads(output.strip().splitlines()[-1])) + + for entry in registry.entries(): + if entry.lazy: + assert entry.class_path.rpartition(".")[0] not in modules + for backend in OPTIONAL_BACKENDS: + assert backend not in modules, backend + + +def test_python_m_autofit_search_manifest_json(): + output = subprocess.run( + [sys.executable, "-m", "autofit", "search-manifest", "--json"], + capture_output=True, + text=True, + check=True, + ).stdout + manifest = json.loads(output) + + assert manifest["schema"] == "search-manifest@1" + assert len(manifest["searches"]) == 15 + assert manifest == json.loads(json.dumps(registry.manifest())) + + +def test_python_m_autofit_search_manifest_table(): + output = subprocess.run( + [sys.executable, "-m", "autofit", "search-manifest"], + capture_output=True, + text=True, + check=True, + ).stdout + + assert output.startswith("search-manifest@1") + assert "Nautilus" in output From 8096bde11209c6e501817150a3f492a5262a1684 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 11:14:55 +0100 Subject: [PATCH 06/14] feat: surface search capabilities in search.summary and model.info Both files now open with the search's class path and its static capabilities (JAX use, gradient, batched, posterior kind, evidence, resumable, objective, invalid value, status), from capabilities.capability_summary_from. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- autofit/non_linear/paths/abstract.py | 8 ++++ autofit/non_linear/paths/directory.py | 9 +++++ autofit/text/text_util.py | 6 +++ .../non_linear/search/test_capabilities.py | 40 +++++++++++++++++++ 4 files changed, 63 insertions(+) diff --git a/autofit/non_linear/paths/abstract.py b/autofit/non_linear/paths/abstract.py index 958842d82..b0faf6d43 100644 --- a/autofit/non_linear/paths/abstract.py +++ b/autofit/non_linear/paths/abstract.py @@ -637,11 +637,19 @@ def save_summary( with open_(filename, "w") as f: f.write(result_info) + search = getattr(self, "search", None) + search_capabilities = None + if search is not None and getattr(type(search), "objective_target", None): + from autofit.non_linear.search.capabilities import capability_summary_from + + search_capabilities = capability_summary_from(search) + text_util.search_summary_to_file( samples=samples, log_likelihood_function_time=log_likelihood_function_time, visualization_time=visualization_time, filename=self.output_path / "search.summary", + search_capabilities=search_capabilities, ) @property diff --git a/autofit/non_linear/paths/directory.py b/autofit/non_linear/paths/directory.py index 2323fd250..78110cb37 100644 --- a/autofit/non_linear/paths/directory.py +++ b/autofit/non_linear/paths/directory.py @@ -497,6 +497,15 @@ def _save_model_info(self, model): """ info = model.info + # Header: the search and its static capabilities (jax_use, gradient, + # posterior kind, ...), so the file says what produced the fit. Guarded + # for paths without a search (or a search double with no declarations). + search = getattr(self, "search", None) + if search is not None and getattr(type(search), "objective_target", None): + from autofit.non_linear.search.capabilities import capability_summary_from + + info = f"{capability_summary_from(search)}\n{info}" + try: scaler_info = self.search.scaler.info_from_model(model=model) except (NotImplementedError, AttributeError): diff --git a/autofit/text/text_util.py b/autofit/text/text_util.py index 899e1d5a7..983aacf9d 100644 --- a/autofit/text/text_util.py +++ b/autofit/text/text_util.py @@ -318,8 +318,14 @@ def search_summary_to_file( log_likelihood_function_time, filename, visualization_time=None, + search_capabilities=None, ): summary = search_summary_from_samples(samples=samples) + + # The search and its static capabilities open the file, one blank line above + # the sample counts (`capabilities.capability_summary_from`). + if search_capabilities: + summary = [search_capabilities, "\n"] + summary summary.append( f"Log Likelihood Function Evaluation Time (seconds) = {log_likelihood_function_time}\n" ) diff --git a/test_autofit/non_linear/search/test_capabilities.py b/test_autofit/non_linear/search/test_capabilities.py index 8ff1cd7fd..3edeec14e 100644 --- a/test_autofit/non_linear/search/test_capabilities.py +++ b/test_autofit/non_linear/search/test_capabilities.py @@ -13,11 +13,14 @@ import json import math import re +import zipfile from pathlib import Path import pytest +import autofit as af from autofit.non_linear.search import capabilities as cap +from autonerves import conf from autofit.non_linear.search.abstract_search import NonLinearSearch from test_autofit.non_linear.search.conformance_roster import searches_under_test @@ -28,6 +31,8 @@ CITATIONS_BIB = Path(__file__).parents[3] / "files" / "citations.bib" +CONFIG_PATH = Path(__file__).parents[2] / "config" + def _params(): return [pytest.param(entry, id=entry.name) for entry in ROSTER] @@ -162,3 +167,38 @@ def test_invalid_value_spelling(): assert cap.invalid_value_to_str(-math.inf) == "-inf" assert cap.invalid_value_to_str(-1.0e99) == "-1e+99" assert cap.invalid_value_to_str(-1.0e30) == "-1e+30" + + +class _NumpyAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + return -0.5 * float((instance.centre - 50.0) ** 2) + + +def _read_output(search, filename): + zip_path = Path(search.paths._zip_path) + if zip_path.exists(): + with zipfile.ZipFile(zip_path) as archive: + return archive.read(filename).decode() + return (Path(search.paths.output_path) / filename).read_text() + + +def test_capabilities_surface_in_model_info_and_search_summary(tmp_path, monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + original_configs = list(conf.instance.configs) + conf.instance.push(new_path=str(CONFIG_PATH), output_path=str(tmp_path)) + try: + search = af.Drawer(name="capabilities_surface", total_draws=3) + search.fit(model=af.Model(af.ex.Gaussian), analysis=_NumpyAnalysis()) + + model_info = _read_output(search, "model.info") + summary = _read_output(search, "search.summary") + finally: + conf.instance.configs = original_configs + + header = cap.capability_summary_from(search) + + assert model_info.startswith(header) + assert summary.startswith(header) + assert "JAX use = none" in header + assert "Posterior kind = point" in header + assert "Total Samples" in summary From ee7e8b3c2cd70300dcb610e10d15836e13a7766a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 11:14:55 +0100 Subject: [PATCH 07/14] docs: generate the search API page, capability matrix and citations docs/_generate_searches.py writes docs/api/searches.rst, the new docs/searches/index.rst capability matrix (15 rows) and the canonical docs/searches/citations.rst from python -m autofit search-manifest --json and files/citations.bib; --check fails on a stale page and runs in a new Docs workflow job and in test_autofit. docs/conf.py mocks the optional sampler backends that are not installed (autodoc_mock_imports), so the minimal [docs] environment builds. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/docs.yml | 15 + docs/_generate_searches.py | 409 +++++++++++++++++++++ docs/api/searches.rst | 40 +- docs/conf.py | 23 ++ docs/general/citations.md | 3 + docs/index.md | 10 + docs/searches/citations.rst | 222 +++++++++++ docs/searches/index.rst | 321 ++++++++++++++++ test_autofit/test_generated_search_docs.py | 40 ++ 9 files changed, 1066 insertions(+), 17 deletions(-) create mode 100644 docs/_generate_searches.py create mode 100644 docs/searches/citations.rst create mode 100644 docs/searches/index.rst create mode 100644 test_autofit/test_generated_search_docs.py diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 7c356e7a1..71e001b6c 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -29,3 +29,18 @@ jobs: with: package: autofit secrets: inherit + + # The search pages (docs/api/searches.rst, docs/searches/index.rst and + # docs/searches/citations.rst) are generated from the search manifest + # (python -m autofit search-manifest --json); fail when they are stale. + generated-search-docs: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install autofit + run: pip install -e . + - name: Check the generated search docs + run: python docs/_generate_searches.py --check diff --git a/docs/_generate_searches.py b/docs/_generate_searches.py new file mode 100644 index 000000000..1d92393af --- /dev/null +++ b/docs/_generate_searches.py @@ -0,0 +1,409 @@ +""" +Generate the search documentation pages from the search manifest. + + python docs/_generate_searches.py # (re)write the generated pages + python docs/_generate_searches.py --check # exit 1 if any page is out of date + +The pages are generated from ``python -m autofit search-manifest --json`` +(``search-manifest@1``, built by ``autofit/non_linear/search/registry.py``) and +``files/citations.bib``; they are never edited by hand: + +- ``docs/api/searches.rst``: the API reference, one autosummary block per family; +- ``docs/searches/index.rst``: the capability matrix ("which search"); +- ``docs/searches/citations.rst``: the canonical citations page for the searches. + +The pages deliberately omit ``autofit_version``, so a version bump never makes them +stale; only a change to the registry or the bibliography does. + +Run from anywhere: paths resolve relative to this file. ``--check`` is run by the Docs +workflow (``.github/workflows/docs.yml``) and by ``test_autofit``. +""" + +import argparse +import json +import re +import subprocess +import sys +from pathlib import Path +from typing import Dict, List + +DOCS = Path(__file__).resolve().parent +REPO = DOCS.parent +BIB = REPO / "files" / "citations.bib" + +HEADER = ( + ".. This file is generated by docs/_generate_searches.py from the search manifest\n" + " (python -m autofit search-manifest --json). Do not edit it by hand: change\n" + " autofit/non_linear/search/registry.py (or files/citations.bib) and re-run the\n" + " generator.\n\n" +) + +# Bibliography fields left out of the rendered entries: long, or local to an author's +# machine. The full entries stay in files/citations.bib. +_DROPPED_BIB_FIELDS = ("abstract", "file", "keywords", "mendeley-groups") + + +def load_manifest() -> dict: + """ + The search manifest, read through the public CLI exactly as other repos read it. + """ + output = subprocess.run( + [sys.executable, "-m", "autofit", "search-manifest", "--json"], + capture_output=True, + text=True, + check=True, + cwd=REPO, + ).stdout + return json.loads(output) + + +def _searches_by_family(manifest: dict) -> Dict[str, List[dict]]: + grouped = {family: [] for family in manifest["families"]} + for search in manifest["searches"]: + grouped[search["family"]].append(search) + return grouped + + +def _yes(value: bool) -> str: + return "yes" if value else "no" + + +def _link(search: dict) -> str: + return f":class:`~autofit.{search['name']}`" + + +def api_searches_rst(manifest: dict) -> str: + lines = [ + HEADER, + ".. _api_searches:\n\n", + "===================\n", + "Non-Linear Searches\n", + "===================\n\n", + "A non-linear search is an algorithm which fits a model to data.\n\n", + f"**PyAutoFit** supports {len(manifest['searches'])} non-linear searches in " + "three families: nested samplers (nest), Markov chain Monte Carlo (MCMC) and " + "maximum likelihood / maximum a posteriori estimators (MLE). The " + ":doc:`capability matrix <../searches/index>` compares them (JAX use, " + "gradients, evidence, posterior kind, install) and the " + ":doc:`search citations <../searches/citations>` page lists what to cite " + "for each.\n\n", + "**Examples / Tutorials:**\n\n", + "- `readthedocs: example using non-linear searches " + "`_.\n", + "- `autofit_workspace: simple tutorial `_\n", + "- `autofit_workspace: complex tutorial `_\n", + "- `HowToFit: tutorial lectures (detailed step-by-step examples) " + "`_\n\n", + ] + for family, searches in _searches_by_family(manifest).items(): + title = manifest["families"][family] + lines += [ + f"{title}\n", + f"{'-' * len(title)}\n\n", + ".. currentmodule:: autofit\n\n", + ".. autosummary::\n", + " :toctree: _autosummary\n", + " :template: custom-class-template.rst\n", + " :recursive:\n\n", + ] + lines += [f" {search['name']}\n" for search in searches] + lines.append("\n") + lines += [ + "There are also a number of tools which are used to customize the behaviour of " + "non-linear searches in **PyAutoFit**,\n", + "including directory output structure, parameter sample initialization and MCMC " + "auto correlation analysis.\n\n", + "Tools\n", + "-----\n\n", + ".. currentmodule:: autofit\n\n", + ".. autosummary::\n", + " :toctree: _autosummary\n", + " :template: custom-class-template.rst\n", + " :recursive:\n\n", + " DirectoryPaths\n", + " DatabasePaths\n", + " Result\n", + " InitializerBall\n", + " InitializerPrior\n", + " ClipperNone\n", + " ClipperPriorBox\n", + " AutoCorrelationsSettings\n\n", + "**PyAutoFit** can perform a parallelized grid-search of non-linear searches, " + "where a subset of parameters in the\n", + "model are fitted in over a discrete grid.\n\n", + "**Examples / Tutorials:**\n\n", + "- `readthedocs: example using a non-linear search grid search " + "`_.\n", + "- `autofit_workspace: example using a non-linear search grid search " + "`_\n\n", + "GridSearch\n", + "----------\n\n", + ".. autosummary::\n", + " :toctree: _autosummary\n", + " :template: custom-class-template.rst\n", + " :recursive:\n\n", + " SearchGridSearch\n", + " GridSearchResult\n", + ] + return "".join(lines) + + +def _list_table(header: List[str], rows: List[List[str]], widths: str) -> List[str]: + lines = [ + ".. list-table::\n", + " :header-rows: 1\n", + f" :widths: {widths}\n\n", + ] + for row in [header] + rows: + lines.append(f" * - {row[0]}\n") + lines += [f" - {cell}\n" for cell in row[1:]] + lines.append("\n") + return lines + + +def _example(search: dict) -> str: + example = search["example"] + if example is None: + return "--" + repo, _, rest = example.partition(":") + path, _, section = rest.partition("#") + url = f"https://github.com/PyAutoLabs/{repo}/blob/main/{path}" + return f"`{Path(path).name} <{url}>`__" + + +def _budget(budget: dict) -> str: + if not budget: + return "--" + return ", ".join(f"``{key}={value}``" for key, value in budget.items()) + + +def searches_index_rst(manifest: dict) -> str: + title = "Choosing a search: the capability matrix" + lines = [ + HEADER, + ".. _search_capability_matrix:\n\n", + f"{'=' * len(title)}\n", + f"{title}\n", + f"{'=' * len(title)}\n\n", + f"Every one of **PyAutoFit**'s {len(manifest['searches'])} non-linear searches " + "declares its capabilities as class attributes (``autofit.Nautilus.jax_use``, " + "``autofit.Emcee.posterior_kind``, ...). The tables below are generated from " + "those declarations, through the versioned search manifest " + f"(``{manifest['schema']}``)::\n\n", + " python -m autofit search-manifest --json\n\n", + "The same flags are enforced at run time: a search whose ``jax_use`` is " + "``required`` raises ``SearchException`` before fitting when the analysis is " + "not JAX (``Analysis(use_jax=True)``).\n\n", + "These are **static** capabilities, a property of the search class. What one " + "run delivers can be narrower: ``SMC`` estimates the evidence only when it " + "starts from a prior-sampling initializer, not from a warm start.\n\n", + "Capabilities\n", + "------------\n\n", + ] + header = [ + "Search", + "Family", + "JAX", + "Gradients", + "Batched", + "Posterior", + "Evidence", + "Resumable", + "Warm start", + "Install", + "Status", + "Example", + ] + rows = [] + for search in manifest["searches"]: + capabilities = search["capabilities"] + rows.append( + [ + _link(search), + search["family"], + capabilities["jax_use"], + capabilities["gradient"], + _yes(capabilities["batched"]), + capabilities["posterior_kind"], + _yes(capabilities["produces_evidence"]), + _yes(capabilities["resumable"]), + capabilities["warm_start"], + f"``autofit[{capabilities['install_extra']}]``" + if capabilities["install_extra"] + else "base", + capabilities["status"], + _example(search), + ] + ) + lines += _list_table(header, rows, "14 6 8 8 7 8 7 8 8 10 9 9") + lines += [ + "Objectives\n", + "----------\n\n", + "What each search's objective returns, the coordinates its backend proposes " + "points in (``unit_cube`` backends map through the prior transform; the " + "likelihood always receives physical parameters), the value it returns for an " + "invalid model, and the reduced budget ``PYAUTO_TEST_MODE=1`` applies.\n\n", + ] + header = ["Search", "Objective", "Coordinates", "Invalid value", "Test-mode budget"] + rows = [ + [ + _link(search), + f"``{search['capabilities']['objective_target']['quantity']}``", + f"``{search['capabilities']['objective_target']['space']}``", + f"``{search['capabilities']['invalid_value']}``", + _budget(search["capabilities"]["test_mode_budget"]), + ] + for search in manifest["searches"] + ] + lines += _list_table(header, rows, "16 18 12 12 42") + lines += [ + "Column meanings\n", + "---------------\n\n", + "- **JAX**: ``none`` (the backend never needs JAX), ``optional`` (a jitted or " + "batched fast path when the analysis is JAX) or ``required`` (the backend " + "cannot run without a traceable JAX likelihood).\n", + "- **Gradients**: ``uses`` when the backend differentiates the likelihood.\n", + "- **Batched**: the backend evaluates a vectorised (``jax.vmap``) likelihood " + "over many points at once.\n", + "- **Posterior**: ``chain`` (equally weighted MCMC draws), ``weighted`` " + "(weighted samples) or ``point`` (a point estimate).\n", + "- **Resumable**: an interrupted run resumes from its own checkpoint.\n", + "- **Warm start**: ``provider`` (starts cold and can seed another search), " + "``consumer`` (wants a warm start, e.g. from a provider's posterior) or " + "``neutral``.\n", + "- **Install**: ``base`` ships with ``pip install autofit``; " + "``autofit[optional]`` needs the optional extra.\n\n", + "See also the :doc:`search citations ` page and the " + ":doc:`API reference <../api/searches>`.\n", + ] + return "".join(lines) + + +def bib_entries(text: str) -> Dict[str, str]: + """ + The entries of a ``.bib`` file keyed by citation key, each as its source text. + """ + entries = {} + for match in re.finditer(r"^@\w+\{([^,]+),", text, re.M): + depth = 0 + for index in range(match.start(), len(text)): + if text[index] == "{": + depth += 1 + elif text[index] == "}": + depth -= 1 + if depth == 0: + entries[match.group(1).strip()] = text[match.start() : index + 1] + break + return entries + + +def _rendered_entry(source: str) -> str: + kept = [ + line + for line in source.splitlines() + if not re.match( + rf"\s*({'|'.join(_DROPPED_BIB_FIELDS)})\s*=", line, re.IGNORECASE + ) + ] + return "\n".join(kept) + + +def searches_citations_rst(manifest: dict, bib_text: str) -> str: + entries = bib_entries(bib_text) + title = "Citing the searches" + lines = [ + HEADER, + ".. _search_citations:\n\n", + f"{'=' * len(title)}\n", + f"{title}\n", + f"{'=' * len(title)}\n\n", + "If you used a non-linear search, please cite the software behind it as well " + "as **PyAutoFit** (key ``pyautofit``). This is the one canonical list: it is " + "generated from each search's declared ``citation_keys`` and the bibliography " + "in `files/citations.bib `__ (long fields such as abstracts are omitted here; the " + "file has the full entries).\n\n", + ] + header = ["Search", "Cite", "Upstream"] + rows = [] + for search in manifest["searches"]: + keys = search["capabilities"]["citation_keys"] + rows.append( + [ + _link(search), + ", ".join(f"``{key}``" for key in keys) if keys else "``pyautofit`` only", + f"`{search['capabilities']['upstream_url']} " + f"<{search['capabilities']['upstream_url']}>`__", + ] + ) + lines += _list_table(header, rows, "20 30 50") + lines += ["BibTeX entries\n", "--------------\n\n"] + keys = [] + for search in manifest["searches"]: + for key in search["capabilities"]["citation_keys"]: + if key not in keys: + keys.append(key) + missing = [key for key in keys if key not in entries] + if missing: + raise KeyError(f"citation keys missing from {BIB.name}: {missing}") + for key in sorted(keys): + lines += [ + f"``{key}``\n\n", + ".. code-block:: bibtex\n\n", + ] + lines += [f" {line}\n" if line else "\n" for line in _rendered_entry(entries[key]).splitlines()] + lines.append("\n") + return "".join(lines) + + +def generated_pages(manifest: dict) -> Dict[Path, str]: + return { + DOCS / "api" / "searches.rst": api_searches_rst(manifest), + DOCS / "searches" / "index.rst": searches_index_rst(manifest), + DOCS / "searches" / "citations.rst": searches_citations_rst( + manifest, BIB.read_text() + ), + } + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[1]) + parser.add_argument( + "--check", + action="store_true", + help="Do not write; exit 1 if any generated page differs from the manifest.", + ) + arguments = parser.parse_args(argv) + + pages = generated_pages(load_manifest()) + + if arguments.check: + stale = [ + path + for path, text in pages.items() + if not path.exists() or path.read_text() != text + ] + for path in stale: + print(f"out of date: {path.relative_to(REPO)}", file=sys.stderr) + if stale: + print( + "Run `python docs/_generate_searches.py` and commit the result.", + file=sys.stderr, + ) + return 1 + print(f"{len(pages)} generated search pages are up to date.") + return 0 + + for path, text in pages.items(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + print(f"wrote {path.relative_to(REPO)}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docs/api/searches.rst b/docs/api/searches.rst index 60e3840ba..94e713749 100644 --- a/docs/api/searches.rst +++ b/docs/api/searches.rst @@ -1,11 +1,17 @@ +.. This file is generated by docs/_generate_searches.py from the search manifest + (python -m autofit search-manifest --json). Do not edit it by hand: change + autofit/non_linear/search/registry.py (or files/citations.bib) and re-run the + generator. + +.. _api_searches: + =================== Non-Linear Searches =================== A non-linear search is an algorithm which fits a model to data. -**PyAutoFit** currently supports three types of non-linear search algorithms: nested samplers (nest), -Markov Chain Monte Carlo (MCMC) and Maximum Likelihood Estimators (MLE). +**PyAutoFit** supports 15 non-linear searches in three families: nested samplers (nest), Markov chain Monte Carlo (MCMC) and maximum likelihood / maximum a posteriori estimators (MLE). The :doc:`capability matrix <../searches/index>` compares them (JAX use, gradients, evidence, posterior kind, install) and the :doc:`search citations <../searches/citations>` page lists what to cite for each. **Examples / Tutorials:** @@ -14,8 +20,8 @@ Markov Chain Monte Carlo (MCMC) and Maximum Likelihood Estimators (MLE). - `autofit_workspace: complex tutorial `_ - `HowToFit: tutorial lectures (detailed step-by-step examples) `_ -Nested Samplers ---------------- +Markov chain Monte Carlo +------------------------ .. currentmodule:: autofit @@ -24,13 +30,13 @@ Nested Samplers :template: custom-class-template.rst :recursive: - Nautilus - DynestyDynamic - DynestyStatic - NSS + Emcee + Zeus + BlackJAXNUTS + SMC -MCMC ----- +Nested sampling +--------------- .. currentmodule:: autofit @@ -39,13 +45,13 @@ MCMC :template: custom-class-template.rst :recursive: - Emcee - Zeus - BlackJAXNUTS - SMC + DynestyStatic + DynestyDynamic + Nautilus + NSS -Maximum Likelihood Estimators ------------------------------ +Maximum likelihood / maximum a posteriori +----------------------------------------- .. currentmodule:: autofit @@ -54,13 +60,13 @@ Maximum Likelihood Estimators :template: custom-class-template.rst :recursive: + Drawer BFGS LBFGS MultiStartAdam MultiStartADABelief MultiStartLion MultiStartProdigy - Drawer There are also a number of tools which are used to customize the behaviour of non-linear searches in **PyAutoFit**, including directory output structure, parameter sample initialization and MCMC auto correlation analysis. diff --git a/docs/conf.py b/docs/conf.py index 4e2178b1b..7a8439bd1 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -100,6 +100,29 @@ "show-inheritance": True, } autodoc_class_signature = "separated" + +# Optional sampler backends are mocked when they are not installed, so the API pages +# (autosummary of every search, including NSS / SMC / Zeus / Nautilus, whose modules +# import their backend) build in the minimal `pip install autofit[docs]` environment +# that ReadTheDocs uses. In the full-extras environment nothing is mocked. The list +# mirrors the `requires` of autofit/non_linear/search/registry.py. +import importlib.util + +_OPTIONAL_BACKENDS = [ + "jax", + "jaxlib", + "blackjax", + "optax", + "prodigyopt", + "nautilus", + "dynesty", + "emcee", + "zeus", + "getdist", +] +autodoc_mock_imports = [ + module for module in _OPTIONAL_BACKENDS if importlib.util.find_spec(module) is None +] autoclass_content = "init" numpydoc_show_class_members = False diff --git a/docs/general/citations.md b/docs/general/citations.md index 29342bf6c..62d9b1ce4 100644 --- a/docs/general/citations.md +++ b/docs/general/citations.md @@ -11,6 +11,9 @@ would greatly appreciate it if you mention **PyAutoFit** by name and include a l **PyAutoFit** is published in the [Journal of Open Source Software](https://joss.theoj.org/papers/10.21105/joss.02550#) and its entry in the above .bib file is under the key `pyautofit`. +What to cite for each non-linear search is listed on the generated +[search citations](../searches/citations.rst) page, the one canonical list. + ## Dynesty If you used the nested sampling algorithm Dynesty, please follow the citation instructions [on the dynesty readthedocs](https://dynesty.readthedocs.io/en/latest/references.html). diff --git a/docs/index.md b/docs/index.md index 891b68978..723aa00fb 100644 --- a/docs/index.md +++ b/docs/index.md @@ -97,6 +97,16 @@ cookbooks/multiple_datasets cookbooks/multi_level_model ``` +```{toctree} +:caption: 'Searches:' +:hidden: true +:maxdepth: 1 + +searches/index +searches/citations +design/run_ctx +``` + ```{toctree} :caption: 'Features:' :hidden: true diff --git a/docs/searches/citations.rst b/docs/searches/citations.rst new file mode 100644 index 000000000..8dc0f1f54 --- /dev/null +++ b/docs/searches/citations.rst @@ -0,0 +1,222 @@ +.. This file is generated by docs/_generate_searches.py from the search manifest + (python -m autofit search-manifest --json). Do not edit it by hand: change + autofit/non_linear/search/registry.py (or files/citations.bib) and re-run the + generator. + +.. _search_citations: + +=================== +Citing the searches +=================== + +If you used a non-linear search, please cite the software behind it as well as **PyAutoFit** (key ``pyautofit``). This is the one canonical list: it is generated from each search's declared ``citation_keys`` and the bibliography in `files/citations.bib `__ (long fields such as abstracts are omitted here; the file has the full entries). + +.. list-table:: + :header-rows: 1 + :widths: 20 30 50 + + * - Search + - Cite + - Upstream + * - :class:`~autofit.Emcee` + - ``emcee`` + - `https://github.com/dfm/emcee `__ + * - :class:`~autofit.Zeus` + - ``zeus1``, ``zeus2`` + - `https://github.com/minaskar/zeus `__ + * - :class:`~autofit.BlackJAXNUTS` + - ``blackjax`` + - `https://github.com/blackjax-devs/blackjax `__ + * - :class:`~autofit.SMC` + - ``blackjax`` + - `https://github.com/blackjax-devs/blackjax `__ + * - :class:`~autofit.DynestyStatic` + - ``dynesty`` + - `https://github.com/joshspeagle/dynesty `__ + * - :class:`~autofit.DynestyDynamic` + - ``dynesty`` + - `https://github.com/joshspeagle/dynesty `__ + * - :class:`~autofit.Nautilus` + - ``nautilus`` + - `https://github.com/johannesulf/nautilus `__ + * - :class:`~autofit.NSS` + - ``blackjax`` + - `https://github.com/blackjax-devs/blackjax `__ + * - :class:`~autofit.Drawer` + - ``pyautofit`` only + - `https://github.com/PyAutoLabs/PyAutoFit `__ + * - :class:`~autofit.BFGS` + - ``scipy`` + - `https://github.com/scipy/scipy `__ + * - :class:`~autofit.LBFGS` + - ``scipy`` + - `https://github.com/scipy/scipy `__ + * - :class:`~autofit.MultiStartAdam` + - ``optax`` + - `https://github.com/google-deepmind/optax `__ + * - :class:`~autofit.MultiStartADABelief` + - ``optax`` + - `https://github.com/google-deepmind/optax `__ + * - :class:`~autofit.MultiStartLion` + - ``optax`` + - `https://github.com/google-deepmind/optax `__ + * - :class:`~autofit.MultiStartProdigy` + - ``optax``, ``prodigy`` + - `https://github.com/konstmish/prodigy `__ + +BibTeX entries +-------------- + +``blackjax`` + +.. code-block:: bibtex + + @article{blackjax, + author = {Cabezas, Alberto and Corenflos, Adrien and Lao, Junpeng and Louf, R{\'e}mi and others}, + title = {{BlackJAX}: Composable {B}ayesian inference in {JAX}}, + journal = {arXiv e-prints}, + eprint = {2402.10797}, + archivePrefix = {arXiv}, + year = {2024} + } + +``dynesty`` + +.. code-block:: bibtex + + @article{dynesty, + archivePrefix = {arXiv}, + arxivId = {1904.02180}, + author = {Speagle, Joshua S}, + doi = {10.1093/mnras/staa278}, + eprint = {1904.02180}, + issn = {0035-8711}, + journal = {Monthly Notices of the Royal Astronomical Society}, + number = {3}, + pages = {3132--3158}, + title = {{dynesty: a dynamic nested sampling package for estimating Bayesian posteriors and evidences}}, + volume = {493}, + year = {2020} + } + +``emcee`` + +.. code-block:: bibtex + + @article{emcee, + archivePrefix = {arXiv}, + arxivId = {1202.3665}, + author = {Foreman-Mackey, Daniel and Hogg, David W. and Lang, Dustin and Goodman, Jonathan}, + doi = {10.1086/670067}, + eprint = {1202.3665}, + issn = {00046280}, + journal = {Publications of the Astronomical Society of the Pacific}, + number = {925}, + pages = {306--312}, + title = {{emcee : The MCMC Hammer }}, + volume = {125}, + year = {2013} + } + +``nautilus`` + +.. code-block:: bibtex + + @ARTICLE{nautilus, + author = {{Lange}, Johannes U.}, + title = "{NAUTILUS: boosting Bayesian importance nested sampling with deep learning}", + journal = {\mnras}, + year = 2023, + month = oct, + volume = {525}, + number = {2}, + pages = {3181-3194}, + doi = {10.1093/mnras/stad2441}, + archivePrefix = {arXiv}, + eprint = {2306.16923}, + primaryClass = {astro-ph.IM}, + adsurl = {https://ui.adsabs.harvard.edu/abs/2023MNRAS.525.3181L}, + adsnote = {Provided by the SAO/NASA Astrophysics Data System} + } + +``optax`` + +.. code-block:: bibtex + + @software{optax, + author = {{DeepMind} and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and others}, + title = {The {D}eep{M}ind {JAX} Ecosystem}, + url = {http://github.com/google-deepmind}, + year = {2020} + } + +``prodigy`` + +.. code-block:: bibtex + + @inproceedings{prodigy, + author = {Mishchenko, Konstantin and Defazio, Aaron}, + title = {{Prodigy: An Expeditiously Adaptive Parameter-Free Learner}}, + booktitle = {Proceedings of the 41st International Conference on Machine Learning}, + series = {Proceedings of Machine Learning Research}, + volume = {235}, + pages = {35779--35804}, + publisher = {PMLR}, + url = {https://proceedings.mlr.press/v235/mishchenko24a.html}, + eprint = {2306.06101}, + archivePrefix = {arXiv}, + year = {2024} + } + +``scipy`` + +.. code-block:: bibtex + + @article{scipy, + author = {{Virtanen}, Pauli and {Gommers}, Ralf and {Oliphant}, + Travis E. and {Haberland}, Matt and {Reddy}, Tyler and + {Cournapeau}, David and {Burovski}, Evgeni and {Peterson}, Pearu + and {Weckesser}, Warren and {Bright}, Jonathan and {van der Walt}, + St{\'e}fan J. and {Brett}, Matthew and {Wilson}, Joshua and + {Jarrod Millman}, K. and {Mayorov}, Nikolay and {Nelson}, Andrew + R.~J. and {Jones}, Eric and {Kern}, Robert and {Larson}, Eric and + {Carey}, CJ and {Polat}, {\.I}lhan and {Feng}, Yu and {Moore}, + Eric W. and {Vand erPlas}, Jake and {Laxalde}, Denis and + {Perktold}, Josef and {Cimrman}, Robert and {Henriksen}, Ian and + {Quintero}, E.~A. and {Harris}, Charles R and {Archibald}, Anne M. + and {Ribeiro}, Ant{\^o}nio H. and {Pedregosa}, Fabian and + {van Mulbregt}, Paul and {Contributors}, SciPy 1. 0}, + title = "{SciPy 1.0: Fundamental Algorithms for Scientific + Computing in Python}", + journal = {Nature Methods}, + year = "2020", + volume={17}, + pages={261--272}, + adsurl = {https://rdcu.be/b08Wh}, + doi = {10.1038/s41592-019-0686-2}, + } + +``zeus1`` + +.. code-block:: bibtex + + @article{zeus1, + title={zeus: A Python Implementation of the Ensemble Slice Sampling method}, + author={Minas Karamanis and Florian Beutler}, + year={2021}, + note={in prep} + } + +``zeus2`` + +.. code-block:: bibtex + + @article{zeus2, + title={Ensemble Slice Sampling}, + author={Minas Karamanis and Florian Beutler}, + year={2020}, + eprint={2002.06212}, + archivePrefix={arXiv}, + primaryClass={stat.ML} + } + diff --git a/docs/searches/index.rst b/docs/searches/index.rst new file mode 100644 index 000000000..2419057a0 --- /dev/null +++ b/docs/searches/index.rst @@ -0,0 +1,321 @@ +.. This file is generated by docs/_generate_searches.py from the search manifest + (python -m autofit search-manifest --json). Do not edit it by hand: change + autofit/non_linear/search/registry.py (or files/citations.bib) and re-run the + generator. + +.. _search_capability_matrix: + +======================================== +Choosing a search: the capability matrix +======================================== + +Every one of **PyAutoFit**'s 15 non-linear searches declares its capabilities as class attributes (``autofit.Nautilus.jax_use``, ``autofit.Emcee.posterior_kind``, ...). The tables below are generated from those declarations, through the versioned search manifest (``search-manifest@1``):: + + python -m autofit search-manifest --json + +The same flags are enforced at run time: a search whose ``jax_use`` is ``required`` raises ``SearchException`` before fitting when the analysis is not JAX (``Analysis(use_jax=True)``). + +These are **static** capabilities, a property of the search class. What one run delivers can be narrower: ``SMC`` estimates the evidence only when it starts from a prior-sampling initializer, not from a warm start. + +Capabilities +------------ + +.. list-table:: + :header-rows: 1 + :widths: 14 6 8 8 7 8 7 8 8 10 9 9 + + * - Search + - Family + - JAX + - Gradients + - Batched + - Posterior + - Evidence + - Resumable + - Warm start + - Install + - Status + - Example + * - :class:`~autofit.Emcee` + - mcmc + - none + - none + - no + - chain + - no + - no + - consumer + - base + - stable + - `mcmc.py `__ + * - :class:`~autofit.Zeus` + - mcmc + - none + - none + - no + - chain + - no + - no + - consumer + - ``autofit[optional]`` + - stable + - `mcmc.py `__ + * - :class:`~autofit.BlackJAXNUTS` + - mcmc + - required + - uses + - yes + - chain + - no + - no + - consumer + - ``autofit[optional]`` + - stable + - `mcmc.py `__ + * - :class:`~autofit.SMC` + - mcmc + - required + - uses + - yes + - weighted + - yes + - no + - provider + - ``autofit[optional]`` + - experimental + - -- + * - :class:`~autofit.DynestyStatic` + - nest + - optional + - none + - no + - weighted + - yes + - yes + - provider + - base + - stable + - `nest.py `__ + * - :class:`~autofit.DynestyDynamic` + - nest + - optional + - none + - no + - weighted + - yes + - yes + - provider + - base + - stable + - `nest.py `__ + * - :class:`~autofit.Nautilus` + - nest + - optional + - none + - yes + - weighted + - yes + - yes + - provider + - ``autofit[optional]`` + - stable + - `nest.py `__ + * - :class:`~autofit.NSS` + - nest + - required + - none + - yes + - weighted + - yes + - yes + - provider + - ``autofit[optional]`` + - experimental + - -- + * - :class:`~autofit.Drawer` + - mle + - none + - none + - no + - point + - no + - no + - neutral + - base + - stable + - `mle.py `__ + * - :class:`~autofit.BFGS` + - mle + - optional + - none + - no + - point + - no + - no + - provider + - base + - stable + - -- + * - :class:`~autofit.LBFGS` + - mle + - optional + - none + - no + - point + - no + - no + - provider + - base + - stable + - `mle.py `__ + * - :class:`~autofit.MultiStartAdam` + - mle + - required + - uses + - yes + - point + - no + - yes + - provider + - base + - stable + - `mle.py `__ + * - :class:`~autofit.MultiStartADABelief` + - mle + - required + - uses + - yes + - point + - no + - yes + - provider + - base + - experimental + - -- + * - :class:`~autofit.MultiStartLion` + - mle + - required + - uses + - yes + - point + - no + - yes + - provider + - base + - experimental + - -- + * - :class:`~autofit.MultiStartProdigy` + - mle + - required + - uses + - yes + - point + - no + - yes + - provider + - base + - stable + - -- + +Objectives +---------- + +What each search's objective returns, the coordinates its backend proposes points in (``unit_cube`` backends map through the prior transform; the likelihood always receives physical parameters), the value it returns for an invalid model, and the reduced budget ``PYAUTO_TEST_MODE=1`` applies. + +.. list-table:: + :header-rows: 1 + :widths: 16 18 12 12 42 + + * - Search + - Objective + - Coordinates + - Invalid value + - Test-mode budget + * - :class:`~autofit.Emcee` + - ``log_posterior`` + - ``physical`` + - ``-inf`` + - ``nwalkers=20``, ``nsteps=10`` + * - :class:`~autofit.Zeus` + - ``log_posterior`` + - ``physical`` + - ``-inf`` + - ``nwalkers=20``, ``nsteps=10`` + * - :class:`~autofit.BlackJAXNUTS` + - ``log_posterior`` + - ``physical`` + - ``-inf`` + - ``num_warmup=20``, ``num_samples=20``, ``num_chains=2`` + * - :class:`~autofit.SMC` + - ``log_likelihood`` + - ``physical`` + - ``-1e+99`` + - ``num_particles=16``, ``num_mcmc_steps=2``, ``max_smc_steps=5`` + * - :class:`~autofit.DynestyStatic` + - ``log_likelihood`` + - ``unit_cube`` + - ``-1e+99`` + - ``maxcall=1`` + * - :class:`~autofit.DynestyDynamic` + - ``log_likelihood`` + - ``unit_cube`` + - ``-1e+99`` + - ``maxcall=1`` + * - :class:`~autofit.Nautilus` + - ``log_likelihood`` + - ``unit_cube`` + - ``-1e+99`` + - ``n_like_max=1`` + * - :class:`~autofit.NSS` + - ``log_likelihood`` + - ``unit_cube`` + - ``-1e+30`` + - ``termination=-1.0`` + * - :class:`~autofit.Drawer` + - ``log_posterior`` + - ``physical`` + - ``-inf`` + - -- + * - :class:`~autofit.BFGS` + - ``neg2_log_posterior`` + - ``physical`` + - ``-inf`` + - -- + * - :class:`~autofit.LBFGS` + - ``neg2_log_posterior`` + - ``physical`` + - ``-inf`` + - -- + * - :class:`~autofit.MultiStartAdam` + - ``neg2_log_posterior`` + - ``physical`` + - ``-inf`` + - ``convergence.window=1``, ``convergence.min_steps=1`` + * - :class:`~autofit.MultiStartADABelief` + - ``neg2_log_posterior`` + - ``physical`` + - ``-inf`` + - ``convergence.window=1``, ``convergence.min_steps=1`` + * - :class:`~autofit.MultiStartLion` + - ``neg2_log_posterior`` + - ``physical`` + - ``-inf`` + - ``convergence.window=1``, ``convergence.min_steps=1`` + * - :class:`~autofit.MultiStartProdigy` + - ``neg2_log_posterior`` + - ``physical`` + - ``-inf`` + - ``convergence.window=1``, ``convergence.min_steps=1`` + +Column meanings +--------------- + +- **JAX**: ``none`` (the backend never needs JAX), ``optional`` (a jitted or batched fast path when the analysis is JAX) or ``required`` (the backend cannot run without a traceable JAX likelihood). +- **Gradients**: ``uses`` when the backend differentiates the likelihood. +- **Batched**: the backend evaluates a vectorised (``jax.vmap``) likelihood over many points at once. +- **Posterior**: ``chain`` (equally weighted MCMC draws), ``weighted`` (weighted samples) or ``point`` (a point estimate). +- **Resumable**: an interrupted run resumes from its own checkpoint. +- **Warm start**: ``provider`` (starts cold and can seed another search), ``consumer`` (wants a warm start, e.g. from a provider's posterior) or ``neutral``. +- **Install**: ``base`` ships with ``pip install autofit``; ``autofit[optional]`` needs the optional extra. + +See also the :doc:`search citations ` page and the :doc:`API reference <../api/searches>`. diff --git a/test_autofit/test_generated_search_docs.py b/test_autofit/test_generated_search_docs.py new file mode 100644 index 000000000..7954203ce --- /dev/null +++ b/test_autofit/test_generated_search_docs.py @@ -0,0 +1,40 @@ +""" +The generated search docs (``docs/_generate_searches.py``) are up to date with the +search manifest, and the capability matrix has one row per registered search. + +Skipped when the ``docs`` directory is absent (an installed-package test run). +""" + +import subprocess +import sys +from pathlib import Path + +import pytest + +from autofit.non_linear.search import registry + +DOCS = Path(__file__).resolve().parents[1] / "docs" +GENERATOR = DOCS / "_generate_searches.py" + +pytestmark = pytest.mark.skipif( + not GENERATOR.exists(), reason="docs/ is not part of this checkout" +) + + +def test_generated_search_pages_are_up_to_date(): + completed = subprocess.run( + [sys.executable, str(GENERATOR), "--check"], + capture_output=True, + text=True, + ) + + assert completed.returncode == 0, completed.stderr + + +def test_capability_matrix_renders_one_row_per_search(): + text = (DOCS / "searches" / "index.rst").read_text() + capabilities = text.split("Capabilities\n------------")[1].split("Objectives\n")[0] + + rows = [line for line in capabilities.splitlines() if line.startswith(" * - :class:")] + + assert len(rows) == len(registry.entries()) == 15 From c03f027a1a2249e4d6436d633ddd82332673d65c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 11:14:55 +0100 Subject: [PATCH 08/14] docs: design note freezing the run(ctx) signature and FitContext docs/design/run_ctx.md (RTD-visible, under Searches) freezes run(self, ctx) -> internal and the FitContext members (objective factory, model, paths, test-mode level, pool factory, start points, RNG, resume, checkpointer slot, schedule, update callback), its lifecycle rules and what the context does not hold, for phase A2 to implement. Search-extensibility phase A1, PyAutoFit#1674. Co-Authored-By: Claude Opus 5.5 --- docs/design/run_ctx.md | 117 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 docs/design/run_ctx.md diff --git a/docs/design/run_ctx.md b/docs/design/run_ctx.md new file mode 100644 index 000000000..9fefdff24 --- /dev/null +++ b/docs/design/run_ctx.md @@ -0,0 +1,117 @@ +# Design note: the `run(ctx)` search hook and `FitContext` + +Status: **frozen** (search-extensibility phase A1, decision D7). Phase A2 ships the +bridge and migrates `Drawer` and `Nautilus` onto it; phase A5 migrates the remaining +searches one PR each. This note fixes the signature and the members of `FitContext` +so that samplers written during the epic can target it before it exists. A change to +anything below is a design change: it needs a new revision of this note, not an +in-passing edit in an implementation PR. + +## Why + +Today every search implements `_fit(model, analysis) -> (search_internal, fitness)` +and re-implements, inside it, the same chores: building a `Fitness` (11 sites, 4 +builders), choosing a pool (4 idioms), drawing start points (8 initializer calls), +resuming (7 idioms), chunking the run between updates, and calling +`perform_update(during_analysis=True)`. A new search therefore costs 400-700 lines, +most of them copied. `run(ctx)` moves the chores into one per-fit context object that +`NonLinearSearch.fit` builds, so a search implements only its backend loop. + +## The contract a new search implements + +```python +class MySearch(NonLinearSearch): + # 1. static capabilities (autofit.non_linear.search.capabilities), declared as + # class attributes and mirrored in autofit/non_linear/search/registry.py + jax_use = JaxUse.OPTIONAL + gradient = Gradient.NONE + posterior_kind = PosteriorKind.WEIGHTED + ... + + # 2. the backend loop + def run(self, ctx: FitContext) -> Any: + """Run the backend and return its internal state (`search_internal`).""" + + # 3. the samples adapter + def raw_samples_from(self, model, internal) -> RawSamples: + """Map the internal state onto parameters, log likelihoods and weights.""" + + # 4. optional: extra samples_info entries + def info_from(self, internal) -> dict: + ... +``` + +The frozen signature is: + +```python +def run(self, ctx: "FitContext") -> Any +``` + +- `run` receives exactly one argument and returns the backend's internal state, the + object `paths.save_search_internal` archives and `samples_via_internal_from` reads. +- `run` never builds a `Fitness`, a pool or an initializer itself, never reads + `PYAUTO_TEST_MODE`, never calls `perform_update` directly and never writes + `.completed`: it asks `ctx`. +- `run` may raise; `fit` cleans the context up (pools closed, quick-update threads + shut down) on every exit path. +- `RawSamples(parameters, log_likelihood | log_posterior, weights, info)` and + `samples_from_raw` are phase A3's samples adapter; `raw_samples_from` is named here + so a search written now has the right shape. + +## `FitContext` members + +`FitContext` is a plain object created by `NonLinearSearch.fit` once per call, after the +test-mode bypass and the fail-fast gates have run. + +| Member | Type / shape | What it gives `run` | Replaces | +|---|---|---|---| +| `ctx.model` | `AbstractPriorModel` | The (frozen) model being fitted. | the `model` argument of `_fit` | +| `ctx.paths` | `AbstractPaths` | Output paths; capability flags (`has_timer`, `supports_checkpointing`) rather than `isinstance(NullPaths)` checks. | `self.paths` reads inside `_fit` | +| `ctx.objective(kind)` | `kind in {"scalar", "batched", "value_and_grad", "batched_value_and_grad"}` -> callable | The objective factory handle (phase A2's `Fitness.objective`). `kind` is execution only; what the objective returns, its coordinates and its invalid value are the search's declared `objective_target` and `invalid_value`. On numpy the grad kinds raise; on JAX the callables are lazily jitted. | 11 `Fitness(...)` sites, 4 builders | +| `ctx.fitness` | `Fitness` | The underlying `Fitness`, for searches that need its history or quick-update hooks. | direct `Fitness` construction | +| `ctx.test_mode_level` | `int` | The resolved test-mode level (always `< 2` inside `run`: levels 2 and 3 bypass the backend). Level 1 budgets are already applied from `test_mode_budget`. | `is_test_mode()` / `test_mode_level()` calls | +| `ctx.pool` | `PoolFactory` | `ctx.pool(n_cores)` returns a pool or `None`, applying the one fork rule (JAX analyses and EP factor searches never fork) and recording the effective worker count in `search.summary`. | 4 pool idioms | +| `ctx.start_points(n)` | `(n, n_dim)` arrays of parameters and figures of merit | Initial points from the search's initializer, seeded from `ctx.rng`, including the start-point plot. | 8 initializer calls + manual `plot_start_point` | +| `ctx.rng` | `numpy.random.SeedSequence` fan-out | Stream 0 feeds the initializer, streams 1+ the backend (`ctx.rng.spawn(k)`). One search-level `seed` (phase A4, D4). | per-search seed handling | +| `ctx.resume` | `Optional[resume_state]` | The interrupted run's resume state when one exists, else `None` (phase A3's `resume_state`, separate from the final archive). | 7 resume idioms | +| `ctx.checkpointer` | `Checkpointer` slot | The archive strategy chosen by the search's class attribute (`DillCheckpointer`, `PickleCheckpointer`, `NativeFileCheckpointer`), with `save/load/exists/finalize`. A1 freezes the slot; A3 fills it. | `conf.instance[...]` mutations, ad-hoc pickling | +| `ctx.schedule` | `chunks()` / `next_budget(done)` | How many iterations to run before the next update. | `_steps_until_full_update`, Dynesty/Nautilus `iterations_from` | +| `ctx.update(internal)` | callback | Converts the internal state once, writes samples and visuals, saves the checkpoint (`perform_update(during_analysis=True)` + checkpoint save). Hands the updater a snapshot, never the live object. | `perform_update` calls inside `_fit` | + +## Lifecycle rules (D6) + +- One `FitContext` per `fit()` call. It is **never** stored on `self`, never pickled or + dill-dumped with the search, and never copied by `copy_with_paths` (a shallow copy, + used by grid-search children): each child fit builds its own. +- It is created after the test-mode bypass return and after the fail-fast gates + (`jax_use='required'` with a numpy analysis, whole-graph backend agreement), so + `run` only ever sees a runnable configuration. +- It is cleaned up on every exit, including exceptions: pools closed, quick-update + threads shut down, compiled objectives released. +- Updates receive snapshots, so a background visualizer never races the backend. + +## What is not in the context + +- **The analysis.** `run` sees the likelihood only through `ctx.objective(kind)` / + `ctx.fitness`; it never calls `analysis.log_likelihood_function` directly. + Visualization and result construction stay in `NonLinearSearch`. +- **Configuration.** No `conf.instance` access through the context; search settings + are constructor arguments, and the test-mode reduction is already applied. +- **Static capabilities.** They are class attributes on the search (`type(self)`), not + context members, so docs and the registry read them without a fit. +- **Identifier inputs.** Nothing in the context enters the search identifier; the + identifier is computed from `__identifier_fields__` before the context exists. +- **Grid-search / EP orchestration.** Outer loops (grid search, sensitivity mapping, + expectation propagation) build one context per inner fit; the context does not know + it is inside one. + +## Migration + +- **A2** ships `FitContext`, `PoolFactory`, `start_points`, the objective factory and a + bridge in `NonLinearSearch.fit`: a search that defines `run` is driven through the + context, a search that only defines `_fit` keeps working unchanged. `Drawer` and + `Nautilus` migrate as the proofs. +- **A5** migrates the rest, one PR each: Emcee, Zeus, BFGS; Dynesty, NSS, NUTS, SMC; + MultiStart last (its `_fit` is split internally first). +- External subclasses that override `_fit` (the developer-tier `nss`, `ultranest` and + `pyswarms` searches) keep working for at least one release after A5. From 2c07ebea698e090ee89ddbfda35cfa7db0521e47 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 13:46:51 +0100 Subject: [PATCH 09/14] fix: migrate legacy FactorGraphModel pickles to _explicit_use_jax Graphs pickled before the backend was derived from the factors stored a plain `_use_jax` attribute, which the new `_use_jax` property shadows, so a legacy graph built with use_jax=False and JAX children re-derived is_jax=True and `tree_flatten()` raised AttributeError on the missing `_explicit_use_jax`. `__setstate__` now moves the old key over. Co-Authored-By: Claude Fable 5.1 --- autofit/graphical/declarative/collection.py | 15 +++++++++++ test_autofit/analysis/test_is_jax.py | 28 +++++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/autofit/graphical/declarative/collection.py b/autofit/graphical/declarative/collection.py index dbe3a4077..e14a9b743 100644 --- a/autofit/graphical/declarative/collection.py +++ b/autofit/graphical/declarative/collection.py @@ -63,6 +63,21 @@ def tree_unflatten(cls, aux_data, children): use_jax=aux_data[2], ) + def __setstate__(self, state): + """ + Restore a pickled graph, migrating pickles written before ``use_jax`` could be + derived from the factors. + + Those stored the graph's backend as a plain ``_use_jax`` attribute, which the + ``_use_jax`` property below now shadows; it is moved to ``_explicit_use_jax`` + so the graph keeps the backend it was built with (then always an explicit + bool, ``False`` by default) and ``tree_flatten`` finds the attribute. + """ + state = dict(state) + if "_explicit_use_jax" not in state: + state["_explicit_use_jax"] = state.pop("_use_jax", None) + self.__dict__.update(state) + @property def _use_jax(self) -> bool: """ diff --git a/test_autofit/analysis/test_is_jax.py b/test_autofit/analysis/test_is_jax.py index 6477ce56a..eff2389d2 100644 --- a/test_autofit/analysis/test_is_jax.py +++ b/test_autofit/analysis/test_is_jax.py @@ -217,6 +217,34 @@ def test_factor_graph_flags_survive_pytree_flattening_and_pickling(): assert pickle.loads(pickle.dumps(graph)).is_jax is True +@requires_jax +def test_legacy_factor_graph_pickle_keeps_its_explicit_backend(): + """ + Graphs pickled before the backend was derived from the factors stored an explicit + ``_use_jax`` attribute (``False`` by default). Unpickling migrates it, so the graph + keeps that choice rather than re-deriving JAX from its children, and + ``tree_flatten`` still works. + """ + graph = g.FactorGraphModel( + _factor(NumpyAnalysis(use_jax=True), "a"), + _factor(NumpyAnalysis(use_jax=True), "b"), + use_jax=False, + ) + + legacy_state = dict(graph.__dict__) + del legacy_state["_explicit_use_jax"] + legacy_state["_use_jax"] = False + legacy = g.FactorGraphModel.__new__(g.FactorGraphModel) + legacy.__dict__.update(legacy_state) + + restored = pickle.loads(pickle.dumps(legacy)) + + assert restored.is_jax is False + children, aux_data = restored.tree_flatten() + assert aux_data[2] is False + assert g.FactorGraphModel.tree_unflatten(aux_data, children).is_jax is False + + @requires_jax def test_hierarchical_factor_carries_the_flag(monkeypatch): def hierarchical(use_jax): From 29ec7f0412cdf30ce6c6c71b8c7cb557dc0636e5 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 13:48:37 +0100 Subject: [PATCH 10/14] fix: derive FactorGraphModel backend from the flattened factors `FactorGraphModel(h)` with `h = HierarchicalFactor(..., use_jax=True)` reported is_jax=False and then rejected its own JAX children: inference read the unflattened `_model_factors` (the HierarchicalFactor container, which has no `is_jax`) while the agreement check read the flattened ones. Both now read `_flat_factors()`. Co-Authored-By: Claude Fable 5.1 --- autofit/graphical/declarative/collection.py | 6 +++++- test_autofit/analysis/test_is_jax.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/autofit/graphical/declarative/collection.py b/autofit/graphical/declarative/collection.py index e14a9b743..98ccc1f5e 100644 --- a/autofit/graphical/declarative/collection.py +++ b/autofit/graphical/declarative/collection.py @@ -83,11 +83,15 @@ def _use_jax(self) -> bool: """ The graph's backend: the explicit ``use_jax`` when one was given, otherwise ``True`` exactly when every factor is JAX (``is_jax`` reads this). + + The factors are the flattened ones (a ``HierarchicalFactor`` contributes the + factors it generates), the same set ``check_backend_agreement`` reads, so a + derived backend always agrees with its own factors. """ explicit = self.__dict__.get("_explicit_use_jax") if explicit is not None: return bool(explicit) - flags = [factor.is_jax for factor in self.__dict__.get("_model_factors", [])] + flags = [factor.is_jax for factor in self._flat_factors()] return bool(flags) and all(flags) @_use_jax.setter diff --git a/test_autofit/analysis/test_is_jax.py b/test_autofit/analysis/test_is_jax.py index eff2389d2..ee45cbd35 100644 --- a/test_autofit/analysis/test_is_jax.py +++ b/test_autofit/analysis/test_is_jax.py @@ -265,5 +265,21 @@ def hierarchical(use_jax): assert not any(factor.is_jax for factor in hierarchical(True).factors) +@requires_jax +def test_factor_graph_derives_jax_through_a_hierarchical_factor(): + h = af.HierarchicalFactor( + af.GaussianPrior, + mean=af.GaussianPrior(mean=0, sigma=1), + sigma=1.0, + use_jax=True, + ) + h.add_drawn_variable(af.GaussianPrior(mean=0, sigma=1)) + graph = af.FactorGraphModel(h) + + assert all(factor.is_jax for factor in h.factors) + assert graph.is_jax is True + graph.check_backend_agreement() + + def test_supports_jax_visualization_reads_is_jax(): assert NumpyAnalysis().supports_jax_visualization is False From c43a87be7fc6c80b74995fe746c63ba2acf30a35 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 13:49:12 +0100 Subject: [PATCH 11/14] fix: run the whole-graph backend check through analysis wrappers `search.fit(analysis=ModelAnalysis(graph, ...))` on a mixed numpy/JAX graph skipped `check_backend_agreement` and reached the backend, because the check only recognised a bare `FactorGraphModel`. The fit gate now unwraps `ModelAnalysis` and `AnalysisFactor` before the check. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/search/abstract_search.py | 16 ++++++++++--- test_autofit/analysis/test_is_jax.py | 25 ++++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index e352097d5..57adf7b4a 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -859,9 +859,19 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res # Runtime gates, after the test-mode bypass return so `PYAUTO_TEST_MODE>=2` # (which never runs a backend) is not affected by them. from autofit.graphical.declarative.collection import FactorGraphModel - - if isinstance(analysis, FactorGraphModel): - analysis.check_backend_agreement() + from autofit.graphical.declarative.factor.analysis import AnalysisFactor + from autofit.non_linear.analysis.model_analysis import ModelAnalysis + + # Unwrap the analysis wrappers that forward their wrapped analysis's + # likelihood, so a graph wrapped in `ModelAnalysis` (or an `AnalysisFactor`) + # is still checked as one whole graph. + inner = analysis + while isinstance(inner, (ModelAnalysis, AnalysisFactor)) and not isinstance( + inner, FactorGraphModel + ): + inner = inner.__dict__.get("analysis") + if isinstance(inner, FactorGraphModel): + inner.check_backend_agreement() # Fail fast: a jax_use='required' search given a numpy analysis raises the one # shared SearchException before any backend state exists. `PYAUTO_DISABLE_JAX=1` diff --git a/test_autofit/analysis/test_is_jax.py b/test_autofit/analysis/test_is_jax.py index ee45cbd35..80ea93d3f 100644 --- a/test_autofit/analysis/test_is_jax.py +++ b/test_autofit/analysis/test_is_jax.py @@ -161,6 +161,31 @@ def test_whole_graph_fit_of_a_mixed_graph_raises(monkeypatch): af.Drawer(total_draws=2).fit(model=graph.global_prior_model, analysis=graph) +class _BackendReached(Exception): + pass + + +@requires_jax +def test_whole_graph_check_survives_a_model_analysis_wrapper(monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + graph = g.FactorGraphModel( + _factor(NumpyAnalysis(), "numpy_factor"), + _factor(NumpyAnalysis(use_jax=True), "jax_factor"), + ) + search = af.Drawer(total_draws=2) + + def reached(*args, **kwargs): + raise _BackendReached() + + monkeypatch.setattr(search, "_fit", reached) + + with pytest.raises(exc.SearchException, match="jax_factor"): + search.fit( + model=graph.global_prior_model, + analysis=ModelAnalysis(graph, graph.global_prior_model), + ) + + @requires_jax @pytest.mark.filterwarnings("ignore::RuntimeWarning") def test_mixed_factor_graph_still_runs_per_factor_ep(monkeypatch): From 1999efa8223b19fd8b0dc969a05c622feaead609 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 13:50:00 +0100 Subject: [PATCH 12/14] fix: declare NSS objective coordinates as physical NSS draws its initial live points in the unit cube but transforms them before `algo.init`; the sampler then proposes physical vectors evaluated through `instance_from_vector` against a physical-space prior density. Declare `physical` in the class, the registry and the generated matrix. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/search/nest/nss/search.py | 5 ++++- autofit/non_linear/search/registry.py | 2 +- docs/searches/index.rst | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index b0d7c24df..d90ff2261 100644 --- a/autofit/non_linear/search/nest/nss/search.py +++ b/autofit/non_linear/search/nest/nss/search.py @@ -208,7 +208,10 @@ class NSS(abstract_nest.AbstractNest): citation_keys = ("blackjax",) status = cap.Status.EXPERIMENTAL test_mode_budget = {"termination": -1.0} - objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.UNIT_CUBE) + # Physical: the unit-cube initial draws are transformed before ``algo.init``, and + # the sampler then proposes physical vectors (``instance_from_vector``) against a + # physical-space prior density. + objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.PHYSICAL) invalid_value = NSS_INVALID_LOG_LIKELIHOOD __identifier_fields__ = ( diff --git a/autofit/non_linear/search/registry.py b/autofit/non_linear/search/registry.py index dd563a05b..92badd492 100644 --- a/autofit/non_linear/search/registry.py +++ b/autofit/non_linear/search/registry.py @@ -433,7 +433,7 @@ def _multi_start( citation_keys=("blackjax",), status="experimental", test_mode_budget={"termination": -1.0}, - objective_target=("log_likelihood", "unit_cube"), + objective_target=("log_likelihood", "physical"), invalid_value="-1e+30", ), example=None, diff --git a/docs/searches/index.rst b/docs/searches/index.rst index 2419057a0..d44cd0be0 100644 --- a/docs/searches/index.rst +++ b/docs/searches/index.rst @@ -268,7 +268,7 @@ What each search's objective returns, the coordinates its backend proposes point - ``n_like_max=1`` * - :class:`~autofit.NSS` - ``log_likelihood`` - - ``unit_cube`` + - ``physical`` - ``-1e+30`` - ``termination=-1.0`` * - :class:`~autofit.Drawer` From b183f9e6c4c2502892fef70a49c1c7b767106545 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 13:51:16 +0100 Subject: [PATCH 13/14] fix: declare the minimizers' observed invalid value (+inf) and its condition BFGS, LBFGS and the MultiStart family declared invalid_value=-inf, but a NaN likelihood reaches their backend as +inf: Fitness replaces it with -inf and the chi-squared conversion multiplies by -2. A FitException early return (and a failed traced assertion) still returns -inf. Declare the NaN path's +inf, record the conditional behaviour in the capability docstring and a footnote under the generated objectives table, and pin both observed values in a test. Normalising the sentinel is left to A3's adapter. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/search/capabilities.py | 6 +++ autofit/non_linear/search/mle/bfgs/search.py | 7 +++- .../search/mle/multi_start_gradient/search.py | 7 +++- autofit/non_linear/search/registry.py | 4 +- docs/_generate_searches.py | 7 ++++ docs/searches/index.rst | 14 ++++--- .../non_linear/search/test_capabilities.py | 40 +++++++++++++++++++ 7 files changed, 75 insertions(+), 10 deletions(-) diff --git a/autofit/non_linear/search/capabilities.py b/autofit/non_linear/search/capabilities.py index c5eb81078..e3626bbdd 100644 --- a/autofit/non_linear/search/capabilities.py +++ b/autofit/non_linear/search/capabilities.py @@ -31,6 +31,12 @@ ``invalid_value`` The sentinel the objective returns for an invalid model. ======================== ============================================================== +``invalid_value`` records what the backend actually sees and is not yet normalised. For +the ``neg2_log_posterior`` minimizers (``BFGS``, ``LBFGS`` and the ``MultiStart*`` +family) it is conditional: a NaN or infinite likelihood is replaced by ``-inf`` and then +multiplied by -2 into ``+inf`` (the declared value), while a ``FitException`` early +return and a failed traced assertion return ``-inf``. + These are **static** capabilities, a property of the class. The capability a given run actually has can be narrower: ``SMC`` estimates the evidence only when it starts from a prior-sampling initializer (a warm start skips the tempering from the prior). Effective diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index e5b00db49..478409be0 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -48,7 +48,12 @@ class AbstractBFGS(AbstractMLE): status = cap.Status.STABLE test_mode_budget = {} objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.NEG2_LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) - invalid_value = -float("inf") + # The value a NaN or infinite likelihood reaches the backend as: ``Fitness`` + # replaces it with ``resample_figure_of_merit=-inf`` and the chi-squared + # conversion (x -2) turns that into ``+inf``. The sentinel is conditional: a + # ``FitException`` early return (numpy) and a failed traced assertion still return + # ``-inf``. Normalising it is the A3 objective adapter's job, not A1's. + invalid_value = float("inf") __identifier_fields__ = ("clipper",) diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 205dd2dbb..dc8c7ca07 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -60,7 +60,12 @@ class AbstractMultiStartGradient(AbstractMLE): status = cap.Status.STABLE test_mode_budget = {"convergence.window": 1, "convergence.min_steps": 1} objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.NEG2_LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) - invalid_value = -float("inf") + # The value a NaN or infinite likelihood reaches the backend as: ``Fitness`` + # replaces it with ``resample_figure_of_merit=-inf`` and the chi-squared + # conversion (x -2) turns that into ``+inf``. The sentinel is conditional: a + # ``FitException`` early return (numpy) and a failed traced assertion still return + # ``-inf``. Normalising it is the A3 objective adapter's job, not A1's. + invalid_value = float("inf") __identifier_fields__ = ("clipper",) diff --git a/autofit/non_linear/search/registry.py b/autofit/non_linear/search/registry.py index 92badd492..288a65727 100644 --- a/autofit/non_linear/search/registry.py +++ b/autofit/non_linear/search/registry.py @@ -234,7 +234,7 @@ def _bfgs(name: str) -> RegistryEntry: status="stable", test_mode_budget={}, objective_target=("neg2_log_posterior", "physical"), - invalid_value="-inf", + invalid_value="inf", ), example=f"{_WORKSPACE}/mle.py#Search: LBFGS" if name == "LBFGS" else None, integration_test=f"{_WORKSPACE_TEST}/{name}.py", @@ -269,7 +269,7 @@ def _multi_start( status=status, test_mode_budget=_MULTI_START_BUDGET, objective_target=("neg2_log_posterior", "physical"), - invalid_value="-inf", + invalid_value="inf", ), example=example, integration_test=f"{_WORKSPACE_TEST}/{name}.py", diff --git a/docs/_generate_searches.py b/docs/_generate_searches.py index 1d92393af..2daef4e9e 100644 --- a/docs/_generate_searches.py +++ b/docs/_generate_searches.py @@ -260,6 +260,13 @@ def searches_index_rst(manifest: dict) -> str: for search in manifest["searches"] ] lines += _list_table(header, rows, "16 18 12 12 42") + lines += [ + "The invalid value is what the backend actually sees, not yet normalised. For " + "the ``neg2_log_posterior`` minimizers it is conditional: a NaN or infinite " + "likelihood is replaced by ``-inf`` and then multiplied by -2 into ``inf`` " + "(the value listed), while a model rejected by a ``FitException`` or a failed " + "assertion returns ``-inf``.\n\n", + ] lines += [ "Column meanings\n", "---------------\n\n", diff --git a/docs/searches/index.rst b/docs/searches/index.rst index d44cd0be0..cf0bc1467 100644 --- a/docs/searches/index.rst +++ b/docs/searches/index.rst @@ -279,34 +279,36 @@ What each search's objective returns, the coordinates its backend proposes point * - :class:`~autofit.BFGS` - ``neg2_log_posterior`` - ``physical`` - - ``-inf`` + - ``inf`` - -- * - :class:`~autofit.LBFGS` - ``neg2_log_posterior`` - ``physical`` - - ``-inf`` + - ``inf`` - -- * - :class:`~autofit.MultiStartAdam` - ``neg2_log_posterior`` - ``physical`` - - ``-inf`` + - ``inf`` - ``convergence.window=1``, ``convergence.min_steps=1`` * - :class:`~autofit.MultiStartADABelief` - ``neg2_log_posterior`` - ``physical`` - - ``-inf`` + - ``inf`` - ``convergence.window=1``, ``convergence.min_steps=1`` * - :class:`~autofit.MultiStartLion` - ``neg2_log_posterior`` - ``physical`` - - ``-inf`` + - ``inf`` - ``convergence.window=1``, ``convergence.min_steps=1`` * - :class:`~autofit.MultiStartProdigy` - ``neg2_log_posterior`` - ``physical`` - - ``-inf`` + - ``inf`` - ``convergence.window=1``, ``convergence.min_steps=1`` +The invalid value is what the backend actually sees, not yet normalised. For the ``neg2_log_posterior`` minimizers it is conditional: a NaN or infinite likelihood is replaced by ``-inf`` and then multiplied by -2 into ``inf`` (the value listed), while a model rejected by a ``FitException`` or a failed assertion returns ``-inf``. + Column meanings --------------- diff --git a/test_autofit/non_linear/search/test_capabilities.py b/test_autofit/non_linear/search/test_capabilities.py index 3edeec14e..7d3245a62 100644 --- a/test_autofit/non_linear/search/test_capabilities.py +++ b/test_autofit/non_linear/search/test_capabilities.py @@ -169,6 +169,46 @@ def test_invalid_value_spelling(): assert cap.invalid_value_to_str(-1.0e30) == "-1e+30" +class _NaNAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + return float("nan") + + +class _FitExceptionAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + raise af.exc.FitException() + + +@pytest.mark.parametrize("name", ["BFGS", "LBFGS", "MultiStartAdam"]) +def test_minimizer_invalid_value_is_what_the_backend_sees(name): + """ + Pins the observed, not yet normalised, sentinel of the ``neg2_log_posterior`` + minimizers, built with the ``Fitness`` settings their ``_fit`` uses: a NaN + likelihood reaches the backend as ``+inf`` (the declared value) and a + ``FitException`` as ``-inf``. A3's objective adapter normalises this and flips the + test deliberately. + """ + import numpy as np + from autofit.non_linear.fitness import Fitness + + cls = getattr(af, name) + model = af.Model(af.ex.Gaussian) + vector = model.physical_values_from_prior_medians + + def fom(analysis): + return Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=False, + resample_figure_of_merit=-np.inf, + convert_to_chi_squared=True, + )(vector) + + assert cls.invalid_value == math.inf + assert fom(_NaNAnalysis()) == cls.invalid_value + assert fom(_FitExceptionAnalysis()) == -math.inf + + class _NumpyAnalysis(af.Analysis): def log_likelihood_function(self, instance): return -0.5 * float((instance.centre - 50.0) ** 2) From 53e39c9c03bf79de390385ccd8dfa2ec7f0359cf Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 13:53:49 +0100 Subject: [PATCH 14/14] fix: declare Emcee resumable Emcee's `_fit` reopens its HDF backend, loads the last sample and iteration count and samples only the remaining steps, so it resumes from its own checkpoint under the published definition. The report's section 3.1 table listed it as not resumable; corrected in the class (with a note), the registry and the generated matrix. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/search/mcmc/emcee/search.py | 6 +++++- autofit/non_linear/search/registry.py | 2 +- docs/searches/index.rst | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index 5bc0bdc5b..b07202396 100644 --- a/autofit/non_linear/search/mcmc/emcee/search.py +++ b/autofit/non_linear/search/mcmc/emcee/search.py @@ -37,7 +37,11 @@ class Emcee(AbstractMCMC): honours_gradient_mode = False posterior_kind = cap.PosteriorKind.CHAIN produces_evidence = False - resumable = False + # Resumable: ``_fit`` reopens the HDF backend, loads the last sample and iteration + # count and samples only the remaining steps. The search-extensibility report's + # section 3.1 table listed Emcee as not resumable; that was corrected here (A1 + # review) against the published meaning "resumes from its own checkpoint". + resumable = True warm_start = cap.WarmStart.CONSUMER install_extra = "" upstream_url = "https://github.com/dfm/emcee" diff --git a/autofit/non_linear/search/registry.py b/autofit/non_linear/search/registry.py index 288a65727..18bcf0dca 100644 --- a/autofit/non_linear/search/registry.py +++ b/autofit/non_linear/search/registry.py @@ -290,7 +290,7 @@ def _multi_start( honours_gradient_mode=False, posterior_kind="chain", produces_evidence=False, - resumable=False, + resumable=True, warm_start="consumer", install_extra="", upstream_url="https://github.com/dfm/emcee", diff --git a/docs/searches/index.rst b/docs/searches/index.rst index cf0bc1467..a8288e623 100644 --- a/docs/searches/index.rst +++ b/docs/searches/index.rst @@ -43,7 +43,7 @@ Capabilities - no - chain - no - - no + - yes - consumer - base - stable