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/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/graphical/declarative/collection.py b/autofit/graphical/declarative/collection.py index c194c9f2d..98ccc1f5e 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,103 @@ def tree_unflatten(cls, aux_data, children): *children[0], name=aux_data[0], include_prior_factors=aux_data[1], + 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: + """ + 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._flat_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 +248,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/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/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 837d145f0..57adf7b4a 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, @@ -426,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 " @@ -653,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() @@ -835,6 +856,28 @@ 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 + 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` + # 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/capabilities.py b/autofit/non_linear/search/capabilities.py new file mode 100644 index 000000000..e3626bbdd --- /dev/null +++ b/autofit/non_linear/search/capabilities.py @@ -0,0 +1,266 @@ +""" +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. +======================== ============================================================== + +``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 +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..832defa0a 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", @@ -270,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 d6f18659a..16e800896 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", @@ -328,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/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index 912eb207e..b07202396 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,28 @@ 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: ``_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" + 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..478409be0 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,29 @@ 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) + # 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",) method = None @@ -295,7 +319,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/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..dc8c7ca07 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,29 @@ 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) + # 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",) # Name of the optax update rule, resolved lazily at fit time from ``optax`` @@ -955,13 +979,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 @@ -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..acd6ae760 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, @@ -225,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): @@ -246,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() @@ -272,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 13320851a..018d7683b 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", @@ -327,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, @@ -338,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, ) @@ -416,7 +437,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/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index 7c358d696..d90ff2261 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,27 @@ 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} + # 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__ = ( "n_live", "num_mcmc_steps", @@ -397,6 +419,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/autofit/non_linear/search/registry.py b/autofit/non_linear/search/registry.py new file mode 100644 index 000000000..18bcf0dca --- /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=True, + 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", "physical"), + 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/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/docs/_generate_searches.py b/docs/_generate_searches.py new file mode 100644 index 000000000..2daef4e9e --- /dev/null +++ b/docs/_generate_searches.py @@ -0,0 +1,416 @@ +""" +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 += [ + "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", + "- **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/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. 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..a8288e623 --- /dev/null +++ b/docs/searches/index.rst @@ -0,0 +1,323 @@ +.. 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 + - yes + - 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`` + - ``physical`` + - ``-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`` + +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 +--------------- + +- **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/analysis/test_is_jax.py b/test_autofit/analysis/test_is_jax.py new file mode 100644 index 000000000..80ea93d3f --- /dev/null +++ b/test_autofit/analysis/test_is_jax.py @@ -0,0 +1,310 @@ +""" +``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) + + +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): + """ + 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_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): + 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) + + +@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 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/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/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 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..7d3245a62 --- /dev/null +++ b/test_autofit/non_linear/search/test_capabilities.py @@ -0,0 +1,244 @@ +""" +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 +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 + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +ROSTER = searches_under_test() + +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] + + +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" + + +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) + + +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 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()) 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 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