diff --git a/autofit/non_linear/analysis/analysis.py b/autofit/non_linear/analysis/analysis.py index f586cd50d..b92747306 100644 --- a/autofit/non_linear/analysis/analysis.py +++ b/autofit/non_linear/analysis/analysis.py @@ -417,7 +417,7 @@ def batched_memory_bytes(self, model, batch_size: int, gradient: bool = False): model=model, analysis=self, fom_is_log_likelihood=True, - use_jax_vmap=True, + batched=True, batch_size=batch_size, ) diff --git a/autofit/non_linear/analysis/multiprocessing.py b/autofit/non_linear/analysis/multiprocessing.py index a024ebef0..24d5b6848 100644 --- a/autofit/non_linear/analysis/multiprocessing.py +++ b/autofit/non_linear/analysis/multiprocessing.py @@ -82,6 +82,19 @@ def __init__(self, analyses, n_cores: int): n_cores The number of cores available to perform analyses over """ + self.processes = list() + + # The one JAX fork rule (`autofit.non_linear.parallel.pool`). This pool always + # forks (at least one process), so it refuses a JAX analysis outright. + if any(getattr(analysis, "is_jax", False) for analysis in analyses): + from autofit import exc + from autofit.non_linear.parallel.pool import JAX_FORK_REASON + + raise exc.SearchException( + f"AnalysisPool cannot evaluate a JAX analysis: {JAX_FORK_REASON}. " + "Evaluate the analyses in this process instead." + ) + self.analyses = analyses self.n_cores = n_cores diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index 848d5c80e..8b8cfc313 100644 --- a/autofit/non_linear/fitness.py +++ b/autofit/non_linear/fitness.py @@ -2,6 +2,7 @@ import numpy as np import os import time +import warnings from timeout_decorator import timeout from typing import Optional @@ -17,6 +18,14 @@ from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.jax.gradient import validate_gradient_mode from autofit.non_linear.jax_compile import log_on_first_compile +from autofit.non_linear.objective import ( + BATCHED, + GRAD_KINDS, + SCALAR, + VALUE_AND_GRAD, + check_kind, + jax_objective, +) from autofit.non_linear.paths.abstract import AbstractPaths from autofit.non_linear.analysis import Analysis @@ -29,6 +38,20 @@ ITERATIONS_NEVER = 1e90 +def _warn_deprecated_fitness_kwarg(name: str, instead: str): + """ + Warn, as a ``FutureWarning`` so end users see it, that a deprecated ``Fitness`` + keyword argument was passed. + """ + warnings.warn( + f"Fitness({name}=...) is deprecated and will be removed after one release: " + f"{instead}. Every search now selects its objective through " + f"`Fitness.objective(kind)` (search-extensibility phase A2).", + FutureWarning, + stacklevel=3, + ) + + def get_timeout_seconds(): try: @@ -177,13 +200,15 @@ def __init__( resample_figure_of_merit: float = None, convert_to_chi_squared: bool = False, store_history: bool = False, - use_jax_vmap : bool = False, - use_jax_jit : bool = False, + use_jax_vmap : Optional[bool] = None, + use_jax_jit : Optional[bool] = None, batch_size : Optional[int] = None, iterations_per_quick_update: Optional[int] = None, background_quick_update: bool = False, live_visual_update: bool = False, gradient_mode: Optional[str] = None, + batched: bool = False, + compile: bool = True, ): """ Interfaces with any non-linear search to fit the model to the data and return a log likelihood via @@ -246,6 +271,18 @@ def __init__( How `grad` differentiates the likelihood: `"reverse"` (`jax.grad`) or `"forward"` (`jax.jacfwd` over the flat parameter vector). `None` (default) uses the analysis's declared `Analysis.gradient_mode`. See `autofit.jax.gradient`. + batched + If `True`, `__call__` / `call_wrap` evaluate a batch of parameter vectors at once (the `"batched"` + objective); a single vector is promoted to a batch of one. Otherwise they evaluate one vector (the + `"scalar"` objective). + compile + On a JAX analysis the objectives are lazily jitted (compiled on their first call). `False` is the + debugging escape hatch: every objective then runs eagerly, op by op. Ignored on numpy. + use_jax_vmap + Deprecated alias of `batched`. + use_jax_jit + Deprecated and without effect: a JAX analysis is always jitted unless `compile=False`. `False` + maps to `compile=False` for one release. """ self.analysis = analysis @@ -268,8 +305,23 @@ def __init__( self.parameters_history_list = [] self.log_likelihood_history_list = [] - self.use_jax_vmap = use_jax_vmap - self.use_jax_jit = use_jax_jit + if use_jax_vmap is not None: + _warn_deprecated_fitness_kwarg( + "use_jax_vmap", "pass `batched=` instead (the same meaning)" + ) + batched = bool(use_jax_vmap) + + if use_jax_jit is not None: + _warn_deprecated_fitness_kwarg( + "use_jax_jit", + "a JAX analysis is always jitted lazily; pass `compile=False` only to " + "debug eagerly", + ) + if not use_jax_jit: + compile = False + + self.batched = bool(batched) + self.compile = bool(compile) self.gradient_mode = validate_gradient_mode(gradient_mode) if self.analysis.is_jax: @@ -278,12 +330,8 @@ def __init__( enable_pytrees() register_model(self.model) - self._call = self.call - - if self.use_jax_vmap: - self._call = self._vmap - elif self.use_jax_jit: - self._call = self._jit + self._objectives = {} + self._call = self.objective(BATCHED if self.batched else SCALAR) self.batch_size = batch_size self.iterations_per_quick_update = iterations_per_quick_update @@ -590,15 +638,71 @@ def call_wrap(self, parameters): depending on configuration. """ - if self.use_jax_vmap: + if self.batched: if len(np.array(parameters).shape) == 1: parameters = np.array(parameters)[None, :] + if self._is_jax: + # One dtype and one container for every caller: a jitted objective is cached + # on the pytree structure and dtype of its input, so a Python list from the + # initializer and an ndarray from the sampler would otherwise be two compiles. + parameters = np.asarray(parameters, dtype=float) + figure_of_merit = self._call(parameters) - if self.use_jax_jit: + if self._is_jax and not self.batched: figure_of_merit = float(figure_of_merit) + self._record(parameters=parameters, figure_of_merit=figure_of_merit) + + return figure_of_merit + + def call_wrap_value_and_grad(self, parameters): + """ + The host-side `(figure_of_merit, gradient)` wrapper of the `"value_and_grad"` objective. + + The counterpart of `call_wrap` for a gradient-based backend that drives a scalar objective + from Python (scipy's `minimize(jac=True)` in `BFGS` / `LBFGS`): the value and the exact + gradient come from one jitted call, and the value is booked exactly as `call_wrap` books it + (quick updates, `store_history`), so a gradient-driven run keeps its history and quick + updates instead of bypassing them. + + Parameters + ---------- + parameters + One parameter vector. + + Returns + ------- + The figure of merit as a Python float and its gradient as a float64 numpy array. + """ + parameters = np.asarray(parameters, dtype=float) + + figure_of_merit, gradient = self.objective(VALUE_AND_GRAD)(parameters) + + figure_of_merit = float(figure_of_merit) + gradient = np.asarray(gradient, dtype=float) + + self._record(parameters=parameters, figure_of_merit=figure_of_merit) + + return figure_of_merit, gradient + + def _record(self, parameters, figure_of_merit): + """ + The per-evaluation bookkeeping of `call_wrap`: convert the figure of merit back to a log + likelihood, feed the quick update and append to the history. + + Both consumers are optional, and with neither active (the default: quick updates + off, no history) nothing here is needed, so the conversion -- which re-evaluates + the log priors on the host for a log-posterior search -- is skipped. That keeps + `call_wrap` on a jitted objective close to the jitted call's own cost. + """ + if not self.store_history and ( + self.iterations_per_quick_update is None + or self.iterations_per_quick_update >= ITERATIONS_NEVER + ): + return + log_likelihood = self.log_likelihood_from( figure_of_merit=figure_of_merit, parameters=parameters ) @@ -610,8 +714,6 @@ def call_wrap(self, parameters): self.parameters_history_list.append(np.array(parameters)) self.log_likelihood_history_list.append(np.array(log_likelihood)) - return figure_of_merit - def manage_quick_update(self, parameters, log_likelihood): """ Manage quick updates during the non-linear search. @@ -836,8 +938,8 @@ def __getstate__(self): state = self.__dict__.copy() # Strip JAX-compiled callables: jax.jit / jax.vmap / jax.grad return # functions tied to C++ XLA state that can't roundtrip through pickle. - # cached_property values lazily recompile on first access after unpickle. - for attr in ("_call", "_jit", "_vmap", "_grad"): + # The objective cache and `_grad` are rebuilt lazily on first use after unpickle. + for attr in ("_call", "_objectives", "_jit", "_vmap", "_grad"): state.pop(attr, None) return state @@ -856,71 +958,134 @@ def __setstate__(self, state): self._set_traced_assertions() # `Fitness` objects pickled before the gradient-mode override existed defer to the analysis. self.__dict__.setdefault("gradient_mode", None) - self._call = self.call - if getattr(self, "use_jax_vmap", False): - self._call = self._vmap - elif getattr(self, "use_jax_jit", False): - self._call = self._jit + # `Fitness` objects pickled before the objective factory (search-extensibility A2) + # carry `use_jax_vmap` / `use_jax_jit` instead of `batched` / `compile`. Their + # dispatch maps onto the factory: vmap -> the batched objective; jit or not, a JAX + # scalar objective is now jitted lazily. + legacy_vmap = self.__dict__.pop("use_jax_vmap", None) + self.__dict__.pop("use_jax_jit", None) + self.__dict__.setdefault("batched", bool(legacy_vmap)) + self.__dict__.setdefault("compile", True) + self._objectives = {} + self._call = self.objective(BATCHED if self.batched else SCALAR) - @cached_property - def _vmap(self): + @property + def use_jax_vmap(self) -> bool: """ - Vectorized and JIT-compiled likelihood function. - - This wraps the base likelihood function (`self.call`) so that it evaluates a whole batch of - parameter vectors at once, as `jax.jit(jax.vmap(self.call))`. The order is jit **of** vmap: - the conventional composition, and the one `analysis/latent.py` already uses. - - That order means the batch is one XLA program. The outer jit traces `vmap(call)` once, keyed - on the shape of the batch it was given, and caches the compiled executable; a second call at - the same batch shape is a cache hit served from jit's C++ fast path, with no Python batching - trace in between. The previous composition, `jax.vmap(jax.jit(self.call))`, re-entered - Python's batching machinery on every call and dispatched the inner `pjit` eagerly underneath - the batching trace, which is where the `autogalaxy_workspace_test#118` hangs are parked. - - The inner jit is dropped rather than kept: a jit nested inside a jit is inlined, so it adds a - call boundary and nothing else. - - This is **not** a fix for the XLA CPU Eigen-pool FFT deadlock (`PyAutoFit#1530`). That fires - at execution time, and execution still happens once per call whatever the tracing order is. - The only measured effect of the ordering is the A/B run of 2026-08-23 on `rectangular_mge.py` - (control 8/10 stalls against 3/10 for jit-of-vmap, Fisher exact p=0.070), recorded there as - contributory rather than causal. - - The compiled executable is still specialised per batch *length*, so a search that varies its - batch size recompiles per distinct length (see the PyAutoMind draft - `vmap_jit_recompiles_per_nautilus_batch_length.md`). This change neither fixes nor worsens - that. - - Because this is a `cached_property`, the compiled function is stored after its first - creation, avoiding repeated JIT compilation overhead. + Deprecated: whether `call_wrap` evaluates batches. Read `batched` instead. """ - import jax + return self.batched - return log_on_first_compile( - jax.jit(jax.vmap(self.call)), - "vectorized (vmap) likelihood function", - ) + @property + def use_jax_jit(self) -> bool: + """ + Deprecated: whether the scalar objective `call_wrap` dispatches to is jitted. + """ + return self._is_jax and self.compile and not self.batched - @cached_property - def _jit(self): + def objective(self, kind: str = SCALAR, compile: Optional[bool] = None): + """ + The objective of `kind`, the one factory every search builds its likelihood + callable from. + + `kind` is execution only, one of `"scalar"`, `"batched"`, `"value_and_grad"` and + `"batched_value_and_grad"` (see `autofit.non_linear.objective`). What it returns, + in which coordinates and for an invalid model is this `Fitness`'s figure-of-merit + convention, which `NonLinearSearch.make_fitness` derives from the search's declared + `objective_target` and `invalid_value`. + + - **numpy**: `"scalar"` is `call` itself and `"batched"` a Python loop over it; + the grad kinds raise `SearchException`, since numpy cannot differentiate. + - **JAX**: every kind is built on the one unjitted, composable `call` and wrapped + in one lazy `jax.jit` (compiled on its first call, with the compile logged), + cached per kind (so each kind compiles once per input shape) and stripped on + pickle. The grad kinds honour `gradient_mode` (else the analysis's declared + `Analysis.gradient_mode`). + + These are the raw objectives: they bypass the quick-update and history + bookkeeping of `call_wrap`, which a backend that evaluates from Python should + call instead (`__call__`, `call_wrap`, `call_wrap_value_and_grad`). + + Parameters + ---------- + kind + The objective kind. + compile + `False` returns the JAX objective unjitted and uncached (the debugging escape + hatch). `None` (default) uses this `Fitness`'s `compile` setting. + """ + check_kind(kind) + + if not self._is_jax: + if kind in GRAD_KINDS: + raise exc.SearchException( + f"The {kind!r} objective differentiates the likelihood, which needs a JAX " + f"analysis, but {type(self.analysis).__name__} has is_jax=False. Construct " + "the analysis with use_jax=True, or use the 'scalar' or 'batched' objective." + ) + if kind == SCALAR: + return self.call + return self._batched_numpy + + compile = self.compile if compile is None else compile + + if not compile: + return jax_objective( + self.call, + kind, + gradient_mode=self._resolved_gradient_mode(kind), + compile=False, + ) + + objectives = self.__dict__.setdefault("_objectives", {}) + + if kind not in objectives: + objectives[kind] = jax_objective( + self.call, + kind, + gradient_mode=self._resolved_gradient_mode(kind), + ) + + return objectives[kind] + + def _resolved_gradient_mode(self, kind: str) -> str: + """ + The gradient mode a grad kind is built with (`"reverse"` for the others, unused). """ - JIT-compiled likelihood function. + if kind not in GRAD_KINDS: + return "reverse" - This wraps the base likelihood function (`self.call`) with `jax.jit`, - producing a compiled version optimized for repeated evaluation on a - single set of parameters. The first call triggers compilation, while - later calls benefit from the compiled execution. + from autofit.jax import gradient + + return gradient.resolve_gradient_mode(self.analysis, self.gradient_mode) - As a `cached_property`, the compiled function is cached after its - first use, so JIT compilation only occurs once. + def _batched_numpy(self, parameters): + """ + The numpy `"batched"` objective: `call` over each row of the batch. """ - import jax + return np.array([self.call(vector) for vector in parameters]) - return log_on_first_compile( - jax.jit(self.call), - "likelihood function", - ) + @property + def _vmap(self): + """ + The `"batched"` objective, `jax.jit(jax.vmap(call))` on JAX (`objective("batched")`). + + The order is jit **of** vmap, so a batch is one XLA program, traced once per batch shape + and served from jit's C++ fast path afterwards (PyAutoFit#1636; the previous + `jax.vmap(jax.jit(call))` re-entered Python's batching machinery on every call). The + executable is still specialised per batch *length*: a backend that varies its batch size + recompiles per distinct length unless it pads (`autofit.non_linear.objective. + evaluate_in_chunks`). Kept as a property for existing callers. + """ + return self.objective(BATCHED) + + @property + def _jit(self): + """ + The `"scalar"` objective, `jax.jit(call)` on JAX (`objective("scalar")`). Kept as a + property for existing callers. + """ + return self.objective(SCALAR) @cached_property def _grad(self): diff --git a/autofit/non_linear/grid/grid_search/__init__.py b/autofit/non_linear/grid/grid_search/__init__.py index 443703449..e9caa67b1 100644 --- a/autofit/non_linear/grid/grid_search/__init__.py +++ b/autofit/non_linear/grid/grid_search/__init__.py @@ -7,7 +7,7 @@ from autonerves.dictable import to_dict from autofit import exc from autofit.mapper import prior as p -from autofit.non_linear.parallel import Process +from autofit.non_linear.parallel import Process, effective_number_of_cores from autofit.text.formatter import write_table from .job import Job from .result import GridSearchResult @@ -182,7 +182,15 @@ def fit( self.logger.info("Running grid search...") - process_class = Process if self.parallel else Sequential + # The one JAX fork rule (`autofit.non_linear.parallel.pool`): the grid's job + # pool never forks a JAX analysis, it runs the cells sequentially instead. + parallel = self.parallel and effective_number_of_cores( + requested=self.number_of_cores, + is_jax=getattr(analysis, "is_jax", False), + where="GridSearch", + ) > 1 + + process_class = Process if parallel else Sequential # noinspection PyArgumentList return self._fit( model=model, diff --git a/autofit/non_linear/grid/sensitivity/__init__.py b/autofit/non_linear/grid/sensitivity/__init__.py index 0eb8207e3..554c810f7 100644 --- a/autofit/non_linear/grid/sensitivity/__init__.py +++ b/autofit/non_linear/grid/sensitivity/__init__.py @@ -13,7 +13,11 @@ from autofit.non_linear.grid.sensitivity.job import Job, MaskedJobResult from autofit.non_linear.grid.sensitivity.job import JobResult from autofit.non_linear.grid.sensitivity.result import SensitivityResult -from autofit.non_linear.parallel import Process +from autofit.non_linear.parallel import ( + Process, + effective_number_of_cores, + jax_backend_initialized, +) from autofit.text.formatter import write_table @@ -147,7 +151,16 @@ def run(self) -> SensitivityResult: ] physical_values = list(self._physical_values) - process_class = Process if self.number_of_cores > 1 else Sequential + # The one JAX fork rule (`autofit.non_linear.parallel.pool`). The jobs build + # their analyses after the fork, so the signal is whether this process has + # already initialised JAX, which is what makes forking it unsafe. + number_of_cores = effective_number_of_cores( + requested=self.number_of_cores, + is_jax=jax_backend_initialized(), + where="Sensitivity", + ) + + process_class = Process if number_of_cores > 1 else Sequential results = [] jobs = [] diff --git a/autofit/non_linear/objective.py b/autofit/non_linear/objective.py new file mode 100644 index 000000000..fadfad7b2 --- /dev/null +++ b/autofit/non_linear/objective.py @@ -0,0 +1,202 @@ +""" +The objective factory shared by every non-linear search. + +A search evaluates its likelihood through one of four objective *kinds*: + +========================== ============================================================= +Kind What the callable takes and returns +========================== ============================================================= +``scalar`` One parameter vector ``(n_dim,)`` -> one figure of merit. +``batched`` A batch ``(n, n_dim)`` -> ``(n,)`` figures of merit. +``value_and_grad`` One vector -> ``(figure_of_merit, gradient)``. +``batched_value_and_grad`` A batch -> ``((n,), (n, n_dim))``. +========================== ============================================================= + +``kind`` is execution only. What the objective returns (log likelihood, log posterior or +-2 x log posterior), its coordinates and the value it returns for an invalid model are +the search's declared ``objective_target`` and ``invalid_value`` +(``autofit.non_linear.search.capabilities``), which ``NonLinearSearch.make_fitness`` +turns into the ``Fitness`` the objective is built from. + +``Fitness.objective(kind)`` is the entry point; this module holds the pieces it is built +from so a search that composes its own coordinate transform on top of the shared +objective (``MultiStartGradient``, ``SMC``) builds it the same way: + +- ``jax_objective(func, kind, ...)`` turns one unjitted, composable scalar function into + the callable for ``kind``: ``jax.vmap`` for the batched kinds, the declared gradient + mode's ``value_and_grad`` for the grad kinds, and, unless ``compile=False``, one outer + ``jax.jit`` wrapped in ``log_on_first_compile``. Building it compiles nothing: tracing + and XLA compilation happen on the first call (lazy jit). +- ``chunk_slices`` / ``evaluate_in_chunks`` sweep a batched callable over a batch in + fixed-width chunks, padding the ragged final chunk, so the compiled program is keyed + on one ``(chunk, n_dim)`` shape however many rows are evaluated (lifted from + ``MultiStartGradient``). + +JAX is imported lazily inside the builders, so importing this module (and ``autofit``) +loads no optional backend. +""" +from typing import Callable, List, Optional, Tuple + +SCALAR = "scalar" +BATCHED = "batched" +VALUE_AND_GRAD = "value_and_grad" +BATCHED_VALUE_AND_GRAD = "batched_value_and_grad" + +#: Every objective kind, in documentation order. +OBJECTIVE_KINDS = (SCALAR, BATCHED, VALUE_AND_GRAD, BATCHED_VALUE_AND_GRAD) + +#: The kinds that differentiate the likelihood, and therefore need JAX. +GRAD_KINDS = (VALUE_AND_GRAD, BATCHED_VALUE_AND_GRAD) + +#: The kinds that evaluate a batch of parameter vectors at once. +BATCHED_KINDS = (BATCHED, BATCHED_VALUE_AND_GRAD) + +_DESCRIPTIONS = { + SCALAR: "likelihood function", + BATCHED: "vectorized (vmap) likelihood function", + VALUE_AND_GRAD: "likelihood function value and gradient", + BATCHED_VALUE_AND_GRAD: "vectorized (vmap) likelihood function value and gradient", +} + + +def check_kind(kind: str) -> str: + """ + Return ``kind`` unchanged if it is an objective kind, else raise ``ValueError``. + """ + if kind not in OBJECTIVE_KINDS: + raise ValueError( + f"Unknown objective kind {kind!r}. Valid kinds are " + f"{', '.join(repr(k) for k in OBJECTIVE_KINDS)}." + ) + return kind + + +def jax_objective( + func: Callable, + kind: str, + gradient_mode: str = "reverse", + compile: bool = True, + description: Optional[str] = None, +) -> Callable: + """ + The JAX callable of objective ``kind`` built on the unjitted scalar ``func``. + + The transforms compose as ``jit(vmap(value_and_grad(func)))``: one outer jit, so a + batch is a single XLA program cached on its shape (the ``Fitness._vmap`` ordering of + PyAutoFit#1636), and a jit nested inside a jit would only add a call boundary. + + Parameters + ---------- + func + A scalar function of one flat parameter vector, traceable by JAX. + kind + One of ``OBJECTIVE_KINDS``. + gradient_mode + ``"reverse"`` or ``"forward"`` (``autofit.jax.gradient``); read by the grad kinds + only. + compile + ``False`` returns the transformed callable unjitted (the debugging escape hatch: + every call then runs eagerly, op by op). + description + How the compile notice names the callable; defaults to one per kind. + + Returns + ------- + The callable. With ``compile=True`` its first call traces, lowers and compiles (and + logs that it did), and the jitted function is reachable on ``__wrapped__``, whose + ``_cache_size()`` is the compile-count probe. + """ + check_kind(kind) + + import jax + + transformed = func + + if kind in GRAD_KINDS: + from autofit.jax.gradient import value_and_grad_from + + transformed = value_and_grad_from(transformed, gradient_mode) + + if kind in BATCHED_KINDS: + transformed = jax.vmap(transformed) + + if not compile: + return transformed + + from autofit.non_linear.jax_compile import log_on_first_compile + + return log_on_first_compile( + jax.jit(transformed), + description or _DESCRIPTIONS[kind], + ) + + +def chunk_slices(n_rows: int, batch_size: int) -> List[Tuple[int, int, int]]: + """ + The ``(lo, hi, pad)`` chunk bounds of a sweep over ``n_rows`` rows in + ``batch_size``-wide chunks. + + Rows ``lo:hi`` of the batch form each chunk, padded by ``pad`` repeats of its final + row so every chunk presents the same ``(batch_size, n_dim)`` shape to the compiled + function (one XLA compile for the whole sweep). Only the final chunk can be ragged + (``pad > 0``). + """ + return [ + (lo, min(lo + batch_size, n_rows), max(0, lo + batch_size - n_rows)) + for lo in range(0, n_rows, batch_size) + ] + + +def evaluate_in_chunks(func: Callable, parameters, batch_size: int, xp=None): + """ + Evaluate the batched callable ``func`` over ``parameters`` in fixed-width chunks. + + The ragged final chunk is padded with repeats of its last row and the padded rows are + discarded from every output, so the result is numerically identical to one call on + the whole batch; only the allocation and the number of compiled shapes change. + + Parameters + ---------- + func + A batched callable returning one array, or a tuple of arrays, each with a leading + batch axis. + parameters + The ``(n, n_dim)`` batch. + batch_size + The chunk width. + xp + The array module the chunks are concatenated with (``jax.numpy`` by default). + + Returns + ------- + The outputs of ``func`` over the whole batch, in the same structure ``func`` returns. + """ + if xp is None: + import jax.numpy as xp + + # One array type for every chunk, so the padded final chunk and the unpadded ones + # present identical inputs to the compiled function. + parameters = xp.asarray(parameters) + + chunks = [] + + for lo, hi, pad in chunk_slices(parameters.shape[0], batch_size): + chunk = parameters[lo:hi] + if pad: + chunk = xp.concatenate([chunk, xp.tile(chunk[-1:], (pad, 1))]) + + outputs = func(chunk) + is_tuple = isinstance(outputs, tuple) + outputs = outputs if is_tuple else (outputs,) + + if pad: + outputs = tuple(output[:-pad] for output in outputs) + + chunks.append(outputs) + + combined = tuple( + xp.concatenate([chunk[index] for chunk in chunks]) + for index in range(len(chunks[0])) + ) + + return combined if is_tuple else combined[0] diff --git a/autofit/non_linear/parallel/__init__.py b/autofit/non_linear/parallel/__init__.py index d222b0c35..b8ed99830 100644 --- a/autofit/non_linear/parallel/__init__.py +++ b/autofit/non_linear/parallel/__init__.py @@ -4,3 +4,9 @@ from .process import Process from .sneaky import SneakyJob from .sneaky import SneakyPool +from .pool import ( + PoolFactory, + check_factor_search_cores, + effective_number_of_cores, + jax_backend_initialized, +) diff --git a/autofit/non_linear/parallel/pool.py b/autofit/non_linear/parallel/pool.py new file mode 100644 index 000000000..c9ad6cd65 --- /dev/null +++ b/autofit/non_linear/parallel/pool.py @@ -0,0 +1,238 @@ +""" +The pools PyAutoFit forks, and the one rule that decides whether it may fork at all. + +**The JAX fork rule.** A process that has initialised JAX (XLA's thread pools, its +compilation cache, a GPU context) must not ``fork``: a forked child inherits locks held by +threads that do not exist in it, and an XLA compile or execution in the child can deadlock +(the same class of hang as the 27-hour EP run guarded against by +``check_factor_search_cores``). It also defeats the point of JAX: every forked worker +recompiles the likelihood, so ``N`` cores cost ``N`` compiles and run slower than one jitted +process. So: + +- a JAX analysis is never evaluated through a multiprocessing pool. A search given one + with ``number_of_cores > 1`` runs with **one** core, logs exactly one INFO line saying + so, and records the requested and effective counts in ``search.summary``. Parallelism + comes from the batched (``jax.vmap``) objective instead; +- the grid-search and sensitivity-mapping job pools apply the same rule to their jobs; +- ``AnalysisPool`` (which always forks) refuses a JAX analysis outright; +- expectation propagation never forks a factor search, JAX or not + (``check_factor_search_cores``, a deliberate refusal). + +``effective_number_of_cores`` is the rule; ``PoolFactory`` applies it to a search's fit and +builds the pools the searches use, with the same pool classes, start method and arguments +as before (``fork_context``), so pool semantics are unchanged wherever forking is allowed. +""" +import logging +import sys +from typing import Optional + +from autofit import exc + +from .context import fork_context + +logger = logging.getLogger(__name__) + +JAX_FORK_REASON = ( + "JAX analyses never fork a multiprocessing pool: a forked child of a " + "JAX-initialised process can deadlock in XLA and recompiles the likelihood per " + "worker; parallelism comes from the vectorised (jax.vmap) objective instead" +) + + +def jax_backend_initialized() -> bool: + """ + Whether this process has already initialised a JAX (XLA) backend. + + Importing ``jax`` initialises nothing; the first array operation does. Used where the + analysis that will run is not known before the fork (sensitivity mapping builds its + analyses inside each job), so the parent's own JAX state is the only signal. + """ + if "jax" not in sys.modules: + return False + try: + from jax._src import xla_bridge + + return bool(getattr(xla_bridge, "_backends", None)) + except Exception: + return False + + +def effective_number_of_cores(requested: int, is_jax: bool, where: str) -> int: + """ + The JAX fork rule: the number of worker processes ``where`` may fork. + + Parameters + ---------- + requested + The ``number_of_cores`` the user asked for. + is_jax + Whether the work being parallelised is JAX (``Analysis.is_jax``, or a JAX backend + already initialised in this process). + where + What is asking, as it should read in the one INFO line. + + Returns + ------- + ``requested``, or 1 when ``is_jax`` and ``requested > 1``, in which case one INFO line + names the downgrade and why. + """ + if requested is None: + return 1 + + if is_jax and requested > 1: + logger.info( + f"{where}: number_of_cores={requested} requested, running with 1 core. " + f"{JAX_FORK_REASON}." + ) + return 1 + + return requested + + +def check_factor_search_cores(search): + """ + Refuse an expectation-propagation factor search built with ``number_of_cores > 1``. + + A deliberate refusal, not a silent downgrade. ``number_of_cores`` is the user's stated + intent, and the *same* search instance is re-entered once per factor per EP step — so + quietly rewriting it to 1 would mutate shared state the caller still owns and would + hide the misconfiguration rather than fix it. Raising makes the operator change the + search they built. + + ``AbstractSearch.optimise`` is the single door every factor optimisation passes + through, so every search type is covered by this one check. + + ``getattr`` with a default of 1: a handful of search doubles subclass + ``NonLinearSearch`` without running its ``__init__`` (e.g. the regression suite's + ``StaticSearch``), and they run nothing in parallel anyway. + """ + number_of_cores = getattr(search, "number_of_cores", 1) + + if number_of_cores > 1: + raise exc.SearchException( + f"Expectation propagation never runs a factor search through a " + f"Python multiprocessing pool (human ruling 2026-09-09), but the " + f"factor optimiser {search.__class__.__name__} was built with " + f"number_of_cores={number_of_cores}.\n\n" + f"Why: a forked likelihood worker that dies (a segfault, an OOM " + f"kill) is silently replaced by `multiprocessing.Pool`, but the " + f"task it was running is never re-issued, so the `Pool.map` " + f"driving the fit blocks forever and the EP run hangs to the wall " + f"clock rather than failing. RAL job 342351_0 burned 27 hours " + f"exactly this way.\n\n" + f"Fix, either of:\n" + f" - build the factor search with number_of_cores=1;\n" + f" - get parallelism from a vectorised JAX likelihood instead, " + f"`Analysis(use_jax=True)` — Nautilus then takes `fit_x1_cpu` " + f"with vectorized=True and no pool at all." + ) + + +class PoolFactory: + def __init__(self, search, analysis=None): + """ + The pools one fit of ``search`` may use, with the JAX fork rule applied once. + + Built by ``NonLinearSearch`` once per fit, after the test-mode bypass and the + fail-fast gates, and reached by a ``run(ctx)`` search as ``ctx.pool`` and by a + legacy ``_fit`` through ``make_pool`` / ``make_sneaky_pool`` / ``number_of_cores`` + of the fit's factory. + + Parameters + ---------- + search + The search being fitted; its ``number_of_cores`` is the request and its + ``paths`` go to ``SneakyPool``. + analysis + The analysis being fitted, whose ``is_jax`` drives the fork rule. ``None`` + (no fit context) applies no JAX rule. + """ + self.search = search + self.requested = getattr(search, "number_of_cores", 1) or 1 + self.is_jax = bool(getattr(analysis, "is_jax", False)) + + self.number_of_cores = effective_number_of_cores( + requested=self.requested, + is_jax=self.is_jax, + where=type(search).__name__, + ) + + self._created = [] + + @property + def summary(self) -> str: + """ + The ``search.summary`` line recording the requested and effective worker counts. + """ + line = ( + f"Number of cores = {self.number_of_cores} " + f"(requested {self.requested})" + ) + if self.number_of_cores < self.requested: + line += f": {JAX_FORK_REASON}" + return line + "\n" + + def __call__(self, n_cores: Optional[int] = None, **pool_kwargs): + """ + A ``fork_context().Pool`` of ``n_cores`` workers, or ``None`` when the fit runs + on one core (no pool at all: a ``Pool(1)`` would push every likelihood call + through a forked worker). + + Parameters + ---------- + n_cores + The worker count; defaults to the fit's effective ``number_of_cores``. A + larger count than the effective one is clamped to it, so the fork rule + cannot be bypassed. + pool_kwargs + Passed to ``Pool`` (e.g. ``initializer`` / ``initargs``). + """ + n_cores = self.number_of_cores if n_cores is None else min( + n_cores, self.number_of_cores + ) + + if n_cores <= 1: + return None + + getattr(self.search, "logger", logger).info("...using pool") + pool = fork_context().Pool(processes=n_cores, **pool_kwargs) + self._created.append(pool) + return pool + + def sneaky(self, fitness): + """ + A ``SneakyPool`` over ``fitness`` (copied to each worker once, at fork time), + or ``None`` on one core. + """ + if self.number_of_cores <= 1: + return None + + from .sneaky import SneakyPool + + getattr(self.search, "logger", logger).warning( + "...using SneakyPool. This copies the likelihood function " + "to each process on instantiation to avoid copying multiple " + "times." + ) + pool = SneakyPool( + processes=self.number_of_cores, + paths=self.search.paths, + fitness=fitness, + ) + self._created.append(pool) + return pool + + def close(self): + """ + Terminate every pool this factory built, after a failed fit. A successful fit's + pools are closed by the search that used them, exactly as before. + """ + for pool in self._created: + terminate = getattr(pool, "terminate", None) + if terminate is None: + continue + try: + terminate() + except Exception: + logger.debug("Pool termination after a failed fit raised", exc_info=True) + self._created = [] diff --git a/autofit/non_linear/paths/abstract.py b/autofit/non_linear/paths/abstract.py index b0faf6d43..6fb6e2713 100644 --- a/autofit/non_linear/paths/abstract.py +++ b/autofit/non_linear/paths/abstract.py @@ -640,9 +640,14 @@ def save_summary( 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 + from autofit.non_linear.search.capabilities import ( + capability_summary_from, + run_summary_from, + ) - search_capabilities = capability_summary_from(search) + search_capabilities = capability_summary_from(search) + run_summary_from( + search + ) text_util.search_summary_to_file( samples=samples, diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 57adf7b4a..b0d7715ac 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -7,7 +7,7 @@ import os import time import warnings -from abc import ABC, abstractmethod +from abc import ABC, ABCMeta, abstractmethod from collections import Counter from functools import wraps from pathlib import Path @@ -37,7 +37,12 @@ # `ITERATIONS_NEVER` is defined in `fitness.py` (so `Fitness.manage_quick_update` can # use it) and re-exported here for existing importers, e.g. `multi_start_gradient`. from autofit.non_linear.fitness import ITERATIONS_NEVER # noqa: F401 -from autofit.non_linear.parallel import SneakyPool, fork_context +from autofit.non_linear.parallel import ( + PoolFactory, + SneakyPool, + check_factor_search_cores, + fork_context, +) from autofit.non_linear.search.updater import log_process_state from autofit.non_linear.paths.abstract import AbstractPaths from autofit.non_linear.paths.database import DatabasePaths @@ -68,6 +73,24 @@ TEST_MODE_REPRESENTATIVE_MAX_ATTEMPTS = 100 +def warn_deprecated_jax_knob(owner: str, name: str, instead: str): + """ + Warn, as a ``FutureWarning`` so end users see it, that the deprecated JAX knob + ``name`` was passed to ``owner``. + + ``use_jax_jit`` and ``use_jax_vmap`` stay accepted for one release (search- + extensibility phase A2): every search now selects its objective through + ``Fitness.objective(kind)``, jitting a JAX analysis lazily and batching it where + the backend has a batched fast path. + """ + warnings.warn( + f"{owner}({name}=...) is deprecated and will be removed after one release: " + f"{instead}", + FutureWarning, + stacklevel=3, + ) + + def check_cores(func): """ Checks how many cores the search has been configured to @@ -136,7 +159,30 @@ def decorated(self, *args, **kwargs): return decorated -class NonLinearSearch(AbstractFactorOptimiser, ABC): +class _SearchMeta(ABCMeta): + """ + Keeps ``NonLinearSearch`` abstract until a subclass implements its backend through + either hook: overriding ``_fit`` (the legacy contract) or ``run(ctx)`` (the bridge, + ``docs/design/run_ctx.md``). ``_fit`` stays an abstract method; a class that + defines ``run`` is concrete because the inherited ``_fit`` is the bridge to it. + """ + + def __new__(mcls, name, bases, namespace, **kwargs): + cls = super().__new__(mcls, name, bases, namespace, **kwargs) + + run = getattr(cls, "run", None) + + if ( + "_fit" in cls.__abstractmethods__ + and run is not None + and not getattr(run, "_is_bridge_default", False) + ): + cls.__abstractmethods__ = frozenset(cls.__abstractmethods__ - {"_fit"}) + + return cls + + +class NonLinearSearch(AbstractFactorOptimiser, ABC, metaclass=_SearchMeta): # Visualization switches toggled per factor search by ``optimise`` (EP), # governed by ``general.yaml -> output -> visualize_ep_factor_searches``. # Class-level (not set in ``__init__``) so search doubles that skip @@ -402,41 +448,9 @@ def optimise( f" AnalysisFactors, HierarchicalFactors and PriorFactors" ) - # A deliberate refusal, not a silent downgrade. `number_of_cores` is the - # user's stated intent, and the *same* search instance is re-entered once - # per factor per EP step — so quietly rewriting it to 1 here would mutate - # shared state the caller still owns and would hide the misconfiguration - # rather than fix it. Raising makes the operator change the search they - # built. - # - # The guard lives here, in `AbstractSearch.optimise`, rather than in any - # one search: `optimise` is the single door every factor optimisation - # passes through, so every search type (Nautilus, Dynesty, Emcee, the MLE - # searches) is covered by one check instead of N. - # - # `getattr` with a default of 1: a handful of search doubles subclass - # `NonLinearSearch` without running its `__init__` (e.g. the regression - # suite's `StaticSearch`), and they run nothing in parallel anyway. - number_of_cores = getattr(self, "number_of_cores", 1) - - if number_of_cores > 1: - raise exc.SearchException( - f"Expectation propagation never runs a factor search through a " - f"Python multiprocessing pool (human ruling 2026-09-09), but the " - f"factor optimiser {self.__class__.__name__} was built with " - f"number_of_cores={number_of_cores}.\n\n" - f"Why: a forked likelihood worker that dies (a segfault, an OOM " - f"kill) is silently replaced by `multiprocessing.Pool`, but the " - f"task it was running is never re-issued, so the `Pool.map` " - f"driving the fit blocks forever and the EP run hangs to the wall " - f"clock rather than failing. RAL job 342351_0 burned 27 hours " - f"exactly this way.\n\n" - f"Fix, either of:\n" - f" - build the factor search with number_of_cores=1;\n" - f" - get parallelism from a vectorised JAX likelihood instead, " - f"`Analysis(use_jax=True)` — Nautilus then takes `fit_x1_cpu` " - f"with vectorized=True and no pool at all." - ) + # A deliberate refusal, not a silent downgrade: EP never forks a factor + # search (`autofit.non_linear.parallel.pool.check_factor_search_cores`). + check_factor_search_cores(self) model = factor.prior_model.mapper_from_prior_arguments( { @@ -557,6 +571,8 @@ def __getstate__(self): del state["_logger"] if "paths" in state: del state["paths"] + # Per-fit state, never serialized with the search. + state.pop("_fit_pools", None) return state @property @@ -674,6 +690,10 @@ class represented by model M and gives a score for their fitness. """ self.check_model(model=model) + # Per-fit: set by `start_resume_fit` when the fit's pools are resolved, and read + # by `search.summary`; cleared so a reused search never reports a stale count. + self._parallel_summary = None + if analysis.is_jax: try: import jax @@ -878,11 +898,19 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res # smoke runs use `PYAUTO_TEST_MODE>=2`, which returned above. cap.check_jax_required(self, analysis) + # The fit's pools, with the one JAX fork rule applied once (one INFO line when it + # downgrades) and the effective worker count recorded for `search.summary`. + self._fit_pools = PoolFactory(self, analysis) + self._parallel_summary = self._fit_pools.summary + model.freeze() - search_internal, fitness = self._fit( - model=model, - analysis=analysis, - ) + try: + search_internal, fitness = self._fit( + model=model, + analysis=analysis, + ) + finally: + self._fit_pools = None if hasattr(fitness, "shutdown_quick_update"): fitness.shutdown_quick_update() @@ -1407,9 +1435,144 @@ def _build_fake_samples(model, parameter_vector, log_likelihood): weight_list=(weights / weights.sum()).tolist(), ) + # Whether `ctx.start_points` makes the start-point plot. `Drawer` draws its whole + # result through the initializer, so it has no starting point to plot. + _plots_start_point = True + + # The `Samples` class `samples_from_raw` builds for a `run(ctx)` search. + samples_cls = Samples + @abstractmethod def _fit(self, model: AbstractPriorModel, analysis: Analysis): - pass + """ + Run the backend and return ``(search_internal, fitness)``. + + A search either overrides this (every search not yet migrated) or implements + ``run(ctx)``, in which case this is the bridge (``docs/design/run_ctx.md``): it + builds the fit's ``Fitness`` and ``FitContext`` -- after the test-mode bypass and + the fail-fast gates, which ``start_resume_fit`` has already run -- calls + ``run(ctx)`` and cleans the context up on every exit. + """ + if type(self).run is NonLinearSearch.run: + raise NotImplementedError( + f"{type(self).__name__} implements neither run(ctx) nor _fit." + ) + + from autofit.non_linear.search.fit_context import FitContext + + fitness = self.make_fitness( + analysis=analysis, model=model, **self.fitness_overrides(analysis) + ) + + ctx = FitContext( + search=self, + model=model, + analysis=analysis, + fitness=fitness, + pool=self._pools(), + test_mode_level=test_mode_level(), + ) + + try: + search_internal = self.run(ctx) + except BaseException: + ctx.close(failed=True) + raise + + ctx.close(failed=False) + + return search_internal, fitness + + def run(self, ctx) -> Any: + """ + The backend loop of a search on the ``run(ctx)`` hook: run the backend through + ``ctx`` and return its internal state (``search_internal``). See + ``docs/design/run_ctx.md``. + """ + raise NotImplementedError + + run._is_bridge_default = True + + def raw_samples_from(self, model: AbstractPriorModel, internal): + """ + Map a ``run(ctx)`` search's internal state onto a ``RawSamples`` + (``autofit.non_linear.search.fit_context``). + """ + raise NotImplementedError + + def info_from(self, internal) -> Optional[dict]: + """ + The ``samples_info`` entries of a ``run(ctx)`` search's internal state. + """ + return None + + def fitness_overrides(self, analysis: Analysis) -> dict: + """ + ``Fitness`` keyword arguments a ``run(ctx)`` search adds to ``make_fitness`` for + this analysis (e.g. ``batched=True`` for a batched backend). Empty by default. + """ + return {} + + def make_fitness( + self, analysis: Analysis, model: AbstractPriorModel, **overrides + ) -> Fitness: + """ + The `Fitness` this search evaluates its likelihood through, the one construction + site every search shares. + + Its figure-of-merit convention comes from the search's declared capabilities + (`autofit.non_linear.search.capabilities`), never from per-site literals: + + - `objective_target.quantity` picks what the objective returns: the log + likelihood (`fom_is_log_likelihood=True`), the log posterior, or -2 x the log + posterior (`convert_to_chi_squared=True`, for the minimizers); + - `invalid_value` is the sentinel the backend sees for an invalid model. For a + `neg2_log_posterior` search that is the value *after* the -2 multiply, so the + `resample_figure_of_merit` substituted before it is `invalid_value / -2` + (`+inf` -> `-inf`). + + The quick-update cadence, background worker and live-visual settings are always + forwarded, so no search can silently drop quick updates (what + `test_quick_update_wiring.py` used to scan for). + + Parameters + ---------- + analysis + The analysis whose likelihood is wrapped. + model + The model being fitted. + overrides + Any other `Fitness` keyword argument, which wins over the derived and + forwarded ones (e.g. `batched=True`, `store_history=True`). + """ + target = type(self).objective_target + + kwargs = dict( + model=model, + analysis=analysis, + paths=self.paths, + iterations_per_quick_update=self.iterations_per_quick_update, + background_quick_update=self.quick_update_background, + live_visual_update=self.live_visual_update, + ) + + if target is not None: + quantity = target.quantity + neg2 = quantity == cap.ObjectiveQuantity.NEG2_LOG_POSTERIOR + + kwargs.update( + fom_is_log_likelihood=quantity == cap.ObjectiveQuantity.LOG_LIKELIHOOD, + convert_to_chi_squared=neg2, + resample_figure_of_merit=( + type(self).invalid_value / -2.0 + if neg2 + else type(self).invalid_value + ), + ) + + kwargs.update(overrides) + + return Fitness(**kwargs) def check_model(self, model: AbstractPriorModel): if model is not None and model.prior_count == 0: @@ -1639,6 +1802,61 @@ def plot_start_point( plot_search=False, ) + def start_points( + self, + model: AbstractPriorModel, + fitness: Fitness, + n: int, + plot: bool = True, + **initializer_kwargs, + ): + """ + The search's ``n`` initial points, drawn by its initializer, with the start-point + plot made once, here, from the first of them. + + The one place the searches draw start points: the initializer is evaluated through + ``fitness`` (so a JAX analysis runs its lazily jitted scalar objective) with the + fit's effective core count (the JAX fork rule of ``PoolFactory``). + + Parameters + ---------- + model + The model being fitted. + fitness + The fit's ``Fitness``; its analysis visualizes the start point. + n + The number of points. + plot + ``False`` skips the start-point plot, for searches whose initial points are + not a starting point (``Drawer`` draws its whole result this way, Dynesty its + initial live points). + initializer_kwargs + Passed to ``initializer.samples_from_model`` (e.g. ``test_mode_samples``). + + Returns + ------- + The initializer's ``(unit_parameter_lists, parameter_lists, figure_of_merit_list)``. + """ + unit_parameter_lists, parameter_lists, figure_of_merit_list = ( + self.initializer.samples_from_model( + total_points=n, + model=model, + fitness=fitness, + paths=self.paths, + n_cores=self._pools().number_of_cores, + **initializer_kwargs, + ) + ) + + if plot: + self.plot_start_point( + parameter_vector=parameter_lists[0], + model=model, + analysis=fitness.analysis, + ) + + return unit_parameter_lists, parameter_lists, figure_of_merit_list + def samples_from(self, model: AbstractPriorModel, search_internal=None) -> Samples: """ Loads the samples of a non-linear search from its output files. @@ -1671,9 +1889,32 @@ def samples_from(self, model: AbstractPriorModel, search_internal=None) -> Sampl def samples_via_internal_from( self, model: AbstractPriorModel, search_internal=None ): - raise NotImplementedError + """ + The samples of the internal state. A ``run(ctx)`` search gets them from its + ``raw_samples_from``; any other search overrides this. + """ + if type(self).raw_samples_from is NonLinearSearch.raw_samples_from: + raise NotImplementedError + + from autofit.non_linear.search.fit_context import samples_from_raw + + if search_internal is None: + search_internal = self.paths.load_search_internal() + + return samples_from_raw( + model=model, + raw=self.raw_samples_from(model, search_internal), + samples_cls=self.samples_cls, + ) + + def _pools(self) -> PoolFactory: + """ + The current fit's `PoolFactory` (the JAX fork rule applied to its analysis), or, + outside a fit, one built from `number_of_cores` alone. + """ + pools = self.__dict__.get("_fit_pools") + return pools if pools is not None else PoolFactory(self) - @check_cores def make_pool(self): """Make the pool instance used to parallelize a `NonLinearSearch` alongside a set of unique ids for every process in the pool. If the specified number of cores is 1, a pool instance is not made and None is returned. @@ -1684,10 +1925,8 @@ def make_pool(self): The pool instance is also set up with a list of unique pool ids, which are used during model-fitting to identify a 'master core' (the one whose id value is lowest) which handles model result output, visualization, etc.""" - self.logger.info("...using pool") - return fork_context().Pool(processes=self.number_of_cores) + return self._pools()() - @check_cores def make_sneaky_pool(self, fitness: Fitness) -> Optional[SneakyPool]: """ Create a pool for multiprocessing that uses slight-of-hand @@ -1705,14 +1944,7 @@ def make_sneaky_pool(self, fitness: Fitness) -> Optional[SneakyPool]: An implementation of a multiprocessing pool """ - self.logger.warning( - "...using SneakyPool. This copies the likelihood function " - "to each process on instantiation to avoid copying multiple " - "times." - ) - return SneakyPool( - processes=self.number_of_cores, paths=self.paths, fitness=fitness - ) + return self._pools().sneaky(fitness) def __eq__(self, other): return isinstance(other, NonLinearSearch) and self.__dict__ == other.__dict__ diff --git a/autofit/non_linear/search/capabilities.py b/autofit/non_linear/search/capabilities.py index e3626bbdd..a9b891950 100644 --- a/autofit/non_linear/search/capabilities.py +++ b/autofit/non_linear/search/capabilities.py @@ -264,3 +264,12 @@ def capability_summary_from(search) -> str: f"{invalid_value_to_str(cls.invalid_value)}\n" f"Status = {cls.status}\n" ) + + +def run_summary_from(search) -> str: + """ + The run-time lines ``search.summary`` adds below the static capabilities: the + requested and effective worker counts of the last fit (the JAX fork rule of + ``autofit.non_linear.parallel.pool``). Empty before a fit has resolved its pools. + """ + return getattr(search, "_parallel_summary", None) or "" diff --git a/autofit/non_linear/search/fit_context.py b/autofit/non_linear/search/fit_context.py new file mode 100644 index 000000000..7d7677918 --- /dev/null +++ b/autofit/non_linear/search/fit_context.py @@ -0,0 +1,200 @@ +""" +The minimal ``run(ctx)`` bridge (search-extensibility phase A2, decision D7). + +``docs/design/run_ctx.md`` (frozen) fixes the ``run(ctx)`` signature and the members of +``FitContext``. A search that implements ``run(ctx)`` (and ``raw_samples_from``) instead +of ``_fit`` is driven through a ``FitContext`` by ``NonLinearSearch._fit``, the bridge; +a search that overrides ``_fit`` keeps working unchanged. ``Drawer`` and ``Nautilus`` are +the first two searches on the bridge. + +Members filled by later phases are present with their A2 value: + +- ``ctx.resume`` is ``None`` and ``ctx.checkpointer`` is ``None`` (phase A3 fills both; + until then a backend resumes from its own native checkpoint, as ``Nautilus`` does); +- ``ctx.rng`` is a ``numpy.random.SeedSequence`` from the search's ``seed`` attribute + when it has one (phase A4 adds one search-level seed and feeds the initializer). + +``RawSamples`` and ``samples_from_raw`` are a provisional, minimal stand-in for phase +A3's samples adapter (``autofit/non_linear/samples/adapter.py``), so a ``run(ctx)`` search +has the right shape now; A3's adapter replaces them. +""" +from typing import Any, Iterator, List, NamedTuple, Optional + +import numpy as np + + +class RawSamples(NamedTuple): + """ + A backend's samples, before conversion to ``Samples``: parameter vectors, their log + likelihoods **or** log posteriors, their weights and the ``samples_info`` entries. + + Provisional: phase A3's ``samples/adapter.py`` owns the final type. + """ + + parameters: List[List[float]] + log_likelihood: Optional[List[float]] = None + log_posterior: Optional[List[float]] = None + weights: Optional[List[float]] = None + info: Optional[dict] = None + + +def samples_from_raw(model, raw: RawSamples, samples_cls): + """ + Convert ``raw`` into a ``samples_cls`` instance: the log prior of every vector, the log + likelihood from the log posterior when only that was given, and unit weights when none + were. Weights are never renormalised. + + Provisional: phase A3's adapter replaces it. + """ + from autofit.non_linear.samples.sample import Sample + + parameter_lists = raw.parameters + + log_prior_list = [ + sum(model.log_prior_list_from_vector(vector=vector)) + for vector in parameter_lists + ] + + if raw.log_likelihood is not None: + log_likelihood_list = raw.log_likelihood + else: + log_likelihood_list = [ + log_posterior - log_prior + for log_posterior, log_prior in zip(raw.log_posterior, log_prior_list) + ] + + weight_list = ( + raw.weights if raw.weights is not None else len(log_likelihood_list) * [1.0] + ) + + sample_list = Sample.from_lists( + model=model, + parameter_lists=parameter_lists, + log_likelihood_list=log_likelihood_list, + log_prior_list=log_prior_list, + weight_list=weight_list, + ) + + return samples_cls( + model=model, + sample_list=sample_list, + samples_info=raw.info, + ) + + +class UpdateSchedule: + def __init__(self, search): + """ + How many iterations a ``run(ctx)`` backend runs before its next ``ctx.update``, + from the search's ``iterations_per_full_update``. + + The backend's total budget is a search setting (``nsteps``, ``n_like_max``, ...), + so it is passed in: ``next_budget(done, total)`` and ``chunks(total)``. + """ + self._search = search + + def next_budget(self, done: int, total: int) -> int: + """ + The next chunk's length, given ``done`` of ``total`` iterations already run. + """ + return self._search._steps_until_full_update(total - done) + + def chunks(self, total: int) -> Iterator[int]: + """ + The chunk lengths that run ``total`` iterations, one update between chunks. + """ + done = 0 + while done < total: + budget = self.next_budget(done, total) + yield budget + done += budget + + +class FitContext: + def __init__( + self, + search, + model, + analysis, + fitness, + pool, + test_mode_level: int, + ): + """ + Everything one ``fit`` hands a ``run(ctx)`` backend (``docs/design/run_ctx.md``). + + Created by the bridge once per fit, after the test-mode bypass and the fail-fast + gates; never stored on the search, never pickled, never copied with it; closed on + every exit. + """ + self._search = search + self._analysis = analysis + + self.model = model + self.paths = search.paths + self.fitness = fitness + self.test_mode_level = test_mode_level + self.pool = pool + self.rng = np.random.SeedSequence(getattr(search, "seed", None)) + self.resume = None + self.checkpointer = None + self.schedule = UpdateSchedule(search) + + def objective(self, kind: str): + """ + The objective of ``kind`` (``Fitness.objective``): what it returns, in which + coordinates and for an invalid model are the search's declared + ``objective_target`` and ``invalid_value``. + """ + return self.fitness.objective(kind) + + def start_points(self, n: int): + """ + ``n`` initial points from the search's initializer, with the start-point plot. + + Returns + ------- + The ``(n, n_dim)`` parameters and the ``(n,)`` figures of merit, as arrays. + """ + _, parameter_lists, figure_of_merit_list = self._search.start_points( + model=self.model, + fitness=self.fitness, + n=n, + plot=self._search._plots_start_point, + ) + return np.asarray(parameter_lists), np.asarray(figure_of_merit_list) + + def update(self, internal: Any): + """ + Write the samples, visuals and summary of the run so far + (``perform_update(during_analysis=True)``). + + The backend is paused for the duration of the call (the update is synchronous), + so ``internal`` is read as it stands at this point of the run. + """ + return self._search.perform_update( + model=self.model, + analysis=self._analysis, + during_analysis=True, + fitness=self.fitness, + search_internal=internal, + ) + + def close(self, failed: bool): + """ + Clean up after ``run``. On failure: close the pools the context built, shut down + the quick-update threads and release the compiled objectives. On success the + fitness is still needed for the final update, which shuts its threads down. + """ + if not failed: + return + + self.pool.close() + + shutdown = getattr(self.fitness, "shutdown_quick_update", None) + if shutdown is not None: + shutdown() + + objectives = getattr(self.fitness, "_objectives", None) + if objectives is not None: + objectives.clear() diff --git a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py index 832defa0a..a0876de53 100644 --- a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py @@ -293,33 +293,16 @@ def _fit(self, model: AbstractPriorModel, analysis): # log-density. `fit` already ran this gate; repeated for direct callers. cap.check_jax_required(self, analysis) - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=False, # log-posterior target for NUTS - resample_figure_of_merit=-jnp.inf, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - ) + fitness = self.make_fitness(analysis=analysis, model=model) # Initial position(s): borrow the standard initializer machinery so # users can substitute their own (InitializerBall by default for # MCMC; InitializerParamStartPoints.from_result(...) to warm-start). # One starting point per chain. - unit_lists, parameter_lists, _ = self.initializer.samples_from_model( - total_points=self.num_chains, + unit_lists, parameter_lists, _ = self.start_points( model=model, fitness=fitness, - paths=self.paths, - n_cores=self.number_of_cores, - ) - - self.plot_start_point( - parameter_vector=parameter_lists[0], - model=model, - analysis=analysis, + n=self.num_chains, ) n_dim = model.prior_count @@ -327,14 +310,14 @@ def _fit(self, model: AbstractPriorModel, analysis): # (num_chains, n_dim) initial_positions = jnp.asarray(stack_initial_positions(parameter_lists)) - # Build the JIT'd log-density target. ``fitness.call`` is the pure - # JAX-traceable path (it routes through model.instance_from_vector and - # analysis.log_likelihood_function with xp=jnp) — distinct from - # ``call_wrap``/``__call__``, which add Python-side history tracking - # and a ``float()`` conversion that would break NUTS gradients. - @jax.jit - def log_density(params): - return fitness.call(params) + # The JIT'd log-density target: the shared "scalar" objective, the + # lazily jitted ``fitness.call`` (the pure JAX-traceable path through + # model.instance_from_vector and analysis.log_likelihood_function with + # xp=jnp) — distinct from ``call_wrap``/``__call__``, which add + # Python-side history tracking and a ``float()`` conversion that would + # break NUTS gradients. blackjax builds its own ``value_and_grad`` of + # it; honouring ``gradient_mode`` here is a separate spike. + log_density = fitness.objective("scalar") # One-shot trace + compile so warmup timing is honest. _ = float(log_density(initial_positions[0])) diff --git a/autofit/non_linear/search/mcmc/blackjax/smc/search.py b/autofit/non_linear/search/mcmc/blackjax/smc/search.py index 16e800896..0b3f5c48e 100644 --- a/autofit/non_linear/search/mcmc/blackjax/smc/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/smc/search.py @@ -12,7 +12,7 @@ from autonerves import conf from autofit.mapper.prior_model.abstract import AbstractPriorModel -from autofit.non_linear.fitness import Fitness +from autofit.non_linear.objective import jax_objective from autofit.non_linear.initializer import Initializer, InitializerPrior from autofit.non_linear.search.mcmc.abstract_mcmc import AbstractMCMC from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelationsSettings @@ -360,16 +360,7 @@ def _fit(self, model: AbstractPriorModel, analysis): # The resample sentinel is the library-standard `-1e99` rather than `-inf`: blackjax's ESS root solver # multiplies the log-likelihood by a candidate `delta` that can be exactly 0.0, and `0 * -inf` is NaN. # A large finite negative value gives the same zero weight without that hazard. - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=True, - resample_figure_of_merit=-1.0e99, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - ) + fitness = self.make_fitness(analysis=analysis, model=model) # ---- Whitening --------------------------------------------------- # @@ -404,18 +395,10 @@ def _fit(self, model: AbstractPriorModel, analysis): "Pass initializer=result.start_point_from() alongside inverse_mass_matrix=result." ) - _, parameter_lists, _ = self.initializer.samples_from_model( - total_points=n_start_points, + _, parameter_lists, _ = self.start_points( model=model, fitness=fitness, - paths=self.paths, - n_cores=self.number_of_cores, - ) - - self.plot_start_point( - parameter_vector=parameter_lists[0], - model=model, - analysis=analysis, + n=n_start_points, ) start_points = stack_initial_positions(parameter_lists) @@ -442,6 +425,10 @@ def _fit(self, model: AbstractPriorModel, analysis): lower_jax = jnp.asarray(lower_limits) upper_jax = jnp.asarray(upper_limits) + # The shared unjitted scalar objective (``fitness.call``), composed below into the + # whitened coordinates SMC steps in. + objective = fitness.objective("scalar", compile=False) + def physical_from_z(z): return shift_jax + whitening_jax @ z @@ -460,7 +447,7 @@ def log_likelihood_z(z): # rejects it. Without this the likelihood may be evaluated on a nonsensical instance and return NaN, # whose derivative is already on the autodiff tape by the time any guard sees it. parameters = jnp.clip(physical_from_z(z), lower_jax, upper_jax) - return fitness.call(parameters) + return objective(parameters) # ---- The tempering path ------------------------------------------ if self.is_warm_start: @@ -527,7 +514,13 @@ def smc_loglikelihood_fn(z): state = smc.init(initial_particles) - vmapped_log_likelihood = jax.jit(jax.vmap(log_likelihood_z)) + # The "batched" objective of the whitened log likelihood, built by the shared + # factory: one lazy jit of the vmap, with its first compile logged. + vmapped_log_likelihood = jax_objective( + log_likelihood_z, + "batched", + description="SMC whitened log likelihood (vmap)", + ) self.logger.info( f"SMC: adaptive tempered SMC ({self.num_particles} particles, kernel={self.kernel}, " diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index b07202396..3b158e375 100644 --- a/autofit/non_linear/search/mcmc/emcee/search.py +++ b/autofit/non_linear/search/mcmc/emcee/search.py @@ -12,7 +12,6 @@ from autofit import exc from autofit.mapper.model_mapper import ModelMapper from autofit.mapper.prior_model.abstract import AbstractPriorModel -from autofit.non_linear.fitness import Fitness from autofit.non_linear.initializer import Initializer from autofit.non_linear.search.mcmc.abstract_mcmc import AbstractMCMC from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelationsSettings @@ -158,16 +157,7 @@ def _fit(self, model: AbstractPriorModel, analysis): """ import emcee - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=False, - resample_figure_of_merit=-np.inf, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - ) + fitness = self.make_fitness(analysis=analysis, model=model) pool = self.make_sneaky_pool(fitness) @@ -204,18 +194,10 @@ def _fit(self, model: AbstractPriorModel, analysis): unit_parameter_lists, parameter_lists, log_posterior_list, - ) = self.initializer.samples_from_model( - total_points=search_internal.nwalkers, + ) = self.start_points( model=model, fitness=fitness, - paths=self.paths, - n_cores=self.number_of_cores, - ) - - self.plot_start_point( - parameter_vector=parameter_lists[0], - model=model, - analysis=analysis, + n=search_internal.nwalkers, ) state = np.zeros(shape=(search_internal.nwalkers, model.prior_count)) diff --git a/autofit/non_linear/search/mcmc/zeus/search.py b/autofit/non_linear/search/mcmc/zeus/search.py index 57abd4099..9fbd0a62b 100644 --- a/autofit/non_linear/search/mcmc/zeus/search.py +++ b/autofit/non_linear/search/mcmc/zeus/search.py @@ -9,7 +9,6 @@ from autofit import exc from autofit.mapper.model_mapper import ModelMapper from autofit.mapper.prior_model.abstract import AbstractPriorModel -from autofit.non_linear.fitness import Fitness from autofit.non_linear.initializer import Initializer from autofit.non_linear.search.mcmc.abstract_mcmc import AbstractMCMC from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelationsSettings @@ -191,16 +190,7 @@ def _fit(self, model: AbstractPriorModel, analysis): pool = self.make_pool() - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=False, - resample_figure_of_merit=-np.inf, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - ) + fitness = self.make_fitness(analysis=analysis, model=model) try: search_internal = self.paths.load_search_internal() @@ -235,19 +225,11 @@ def _fit(self, model: AbstractPriorModel, analysis): unit_parameter_lists, parameter_lists, log_posterior_list, - ) = self.initializer.samples_from_model( - total_points=search_internal.nwalkers, + ) = self.start_points( model=model, fitness=fitness, + n=search_internal.nwalkers, test_mode_samples=False, - paths=self.paths, - n_cores=self.number_of_cores, - ) - - self.plot_start_point( - parameter_vector=parameter_lists[0], - model=model, - analysis=analysis, ) state = np.zeros(shape=(search_internal.nwalkers, model.prior_count)) diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index 478409be0..6dfa30563 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -7,7 +7,6 @@ from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.search.mle.abstract_mle import AbstractMLE from autofit.non_linear.analysis import Analysis -from autofit.non_linear.fitness import Fitness from autofit.non_linear.clipper import ( AbstractClipper, ClipperNone, @@ -245,17 +244,10 @@ def _fit( # immediately rather than after the first chunk of iterations. bounds = self._bounds_from(model=model) - fitness = Fitness( - model=model, + fitness = self.make_fitness( analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=False, - resample_figure_of_merit=-np.inf, - convert_to_chi_squared=True, + model=model, store_history=self.should_plot_start_point, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, ) try: @@ -274,12 +266,10 @@ def _fit( unit_parameter_lists, parameter_lists, log_posterior_list, - ) = self.initializer.samples_from_model( - total_points=1, + ) = self.start_points( model=model, fitness=fitness, - paths=self.paths, - n_cores=self.number_of_cores, + n=1, ) x0 = np.asarray(parameter_lists[0]) @@ -290,12 +280,6 @@ def _fit( f"Starting new {self.method} non-linear search (no previous samples found)." ) - self.plot_start_point( - parameter_vector=x0, - model=model, - analysis=analysis, - ) - while total_iterations < self.maxiter: iterations_remaining = self.maxiter - total_iterations @@ -321,8 +305,13 @@ def _fit( # diverge by branch would be a trap of its own. if analysis.is_jax: + # The exact gradient from the "value_and_grad" objective, one jitted + # call per step: no finite differences (which cost n_params + 1 + # likelihood calls per gradient). The host-side wrapper books every + # call as `call_wrap` does, so history and quick updates still fire. search_internal = optimize.minimize( - fun=fitness._jit, + fun=fitness.call_wrap_value_and_grad, + jac=True, x0=x0, method=self.method, options=options, diff --git a/autofit/non_linear/search/mle/drawer/search.py b/autofit/non_linear/search/mle/drawer/search.py index 3355a86e3..3491776bb 100644 --- a/autofit/non_linear/search/mle/drawer/search.py +++ b/autofit/non_linear/search/mle/drawer/search.py @@ -5,7 +5,6 @@ from autofit.mapper.prior_model.abstract import AbstractPriorModel -from autofit.non_linear.fitness import Fitness 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 @@ -34,6 +33,9 @@ class Drawer(AbstractMLE): objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_POSTERIOR, cap.CoordinateSpace.PHYSICAL) invalid_value = -float("inf") + # The draws are the result, not a starting point: nothing to plot. + _plots_start_point = False + __identifier_fields__ = ("total_draws",) def __init__( @@ -111,112 +113,55 @@ def __init__( self.logger.debug("Creating Drawer Search") - def _fit(self, model: AbstractPriorModel, analysis): + def run(self, ctx): """ - Fit a model using Drawer and the Analysis class which contains the data and returns the log likelihood from - instances of the model, which the `NonLinearSearch` seeks to maximize. + Draw ``total_draws`` points from the initializer and return them (the ``run(ctx)`` + hook; see ``docs/design/run_ctx.md``). + + Every log likelihood evaluation happens in the initializer, through the fit's + objective (lazily jitted on a JAX analysis). Parameters ---------- - model : ModelMapper - The model which generates instances for different points in parameter space. - analysis : Analysis - Contains the data and the log likelihood function which fits an instance of the model to the data, returning - the log likelihood the `NonLinearSearch` maximizes. + ctx + The fit's ``FitContext``. Returns ------- - A result object comprising the Samples object that inclues the maximum log likelihood instance and full - chains used by the fit. + The internal state: the drawn parameter lists, their log posteriors and the run time. """ - - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=False, - resample_figure_of_merit=-np.inf, - convert_to_chi_squared=False, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - ) - - total_draws = self.total_draws - self.logger.info( - f"Performing DrawerSearch for a total of {total_draws} points." + f"Performing DrawerSearch for a total of {self.total_draws} points." ) - ( - unit_parameter_lists, - parameter_lists, - log_posterior_list, - ) = self.initializer.samples_from_model( - total_points=self.total_draws, - model=model, - fitness=fitness, - paths=self.paths, - n_cores=self.number_of_cores, - ) + parameters, log_posteriors = ctx.start_points(self.total_draws) search_internal = { - "parameter_lists": parameter_lists, - "log_posterior_list": log_posterior_list, + "parameter_lists": parameters.tolist(), + "log_posterior_list": log_posteriors.tolist(), "time": self.timer.time if self.timer else None, } - self.paths.save_search_internal( + ctx.paths.save_search_internal( obj=search_internal, ) self.logger.info("Drawer complete") - return search_internal, fitness + return search_internal - def samples_via_internal_from(self, model, search_internal=None): + def raw_samples_from(self, model, internal): """ - Returns a `Samples` object from the Drawer internal results. - - The `search_internal` dictionary returned by `_fit` is used when passed, so a fit with `NullPaths` - (which writes nothing to disk) can still build its samples; otherwise it is loaded from the output folder. - - Parameters - ---------- - model - Maps input vectors of unit parameter values to physical values and model instances via priors. - search_internal - The dictionary of parameter lists, log posteriors and run time that `_fit` returns. + The drawn points as ``RawSamples``: log posteriors, unit weights, and the internal + dictionary itself as ``samples_info``. """ - search_internal_dict = ( - search_internal - if search_internal is not None - else self.paths.load_search_internal() - ) + from autofit.non_linear.search.fit_context import RawSamples - parameter_lists = search_internal_dict["parameter_lists"] - log_posterior_list = search_internal_dict["log_posterior_list"] - - log_prior_list = [ - sum(model.log_prior_list_from_vector(vector=vector)) - for vector in parameter_lists - ] - log_likelihood_list = [ - lp - prior for lp, prior in zip(log_posterior_list, log_prior_list) - ] - - weight_list = len(log_likelihood_list) * [1.0] - - sample_list = Sample.from_lists( - model=model, - parameter_lists=parameter_lists, - log_likelihood_list=log_likelihood_list, - log_prior_list=log_prior_list, - weight_list=weight_list, + return RawSamples( + parameters=internal["parameter_lists"], + log_posterior=internal["log_posterior_list"], + info=self.info_from(internal), ) - return Samples( - model=model, - sample_list=sample_list, - samples_info=search_internal_dict, - ) + def info_from(self, internal): + return internal 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 dc8c7ca07..9103e176b 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -12,6 +12,7 @@ from autofit.non_linear.search.mle.abstract_mle import AbstractMLE from autofit.non_linear.analysis import Analysis from autofit.non_linear.fitness import Fitness +from autofit.non_linear.objective import chunk_slices, evaluate_in_chunks from autofit.jax.gradient import ( resolve_gradient_mode, validate_gradient_mode, @@ -999,16 +1000,10 @@ def _fit( else self.iterations_per_quick_update ) - fitness = Fitness( - model=model, + fitness = self.make_fitness( analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=False, - resample_figure_of_merit=-np.inf, - convert_to_chi_squared=True, + model=model, iterations_per_quick_update=iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, ) # -2 * log_posterior, to MINIMIZE. value_and_grad over every start. @@ -1044,7 +1039,11 @@ def _fit( f"MultiStartGradient gradient mode: {gradient_mode} ({source})." ) - _value_and_grad = value_and_grad_from(fitness.call, gradient_mode) + # Built on the shared unjitted scalar objective (``fitness.call``): this + # search keeps its own transformed-coordinate builder on top of it. + objective = fitness.objective("scalar", compile=False) + + _value_and_grad = value_and_grad_from(objective, gradient_mode) # Per-start gradient finiteness is reduced *inside* the jitted call, as a # third output, rather than by a separate op on its result. The reduction @@ -1143,7 +1142,7 @@ def _to_physical(vector): # draws, and those draws are, and stay, physical. _value_and_grad_stepped = ( value_and_grad_from( - lambda phi: fitness.call(_to_physical(phi)), gradient_mode + lambda phi: objective(_to_physical(phi)), gradient_mode ) if (has_scaler or has_bijector) else _value_and_grad @@ -1170,29 +1169,11 @@ def _value_and_grad_finite(vector): batch_size = self.batch_size def batched_value_and_grad(params): - foms_chunks = [] - grads_chunks = [] - grad_finite_chunks = [] - violation_chunks = [] - for lo, hi, pad in _chunk_slices(params.shape[0], batch_size): - chunk = params[lo:hi] - if pad: - chunk = jnp.concatenate([chunk, jnp.tile(chunk[-1:], (pad, 1))]) - foms, grads, grad_finite, violation = _vmapped(chunk) - if pad: - foms = foms[:-pad] - grads = grads[:-pad] - grad_finite = grad_finite[:-pad] - violation = violation[:-pad] - foms_chunks.append(foms) - grads_chunks.append(grads) - grad_finite_chunks.append(grad_finite) - violation_chunks.append(violation) - return ( - jnp.concatenate(foms_chunks), - jnp.concatenate(grads_chunks), - jnp.concatenate(grad_finite_chunks), - jnp.concatenate(violation_chunks), + # The shared batching helper: chunks of ``batch_size`` rows, + # the ragged final chunk padded to that shape and its padded + # rows discarded, so one chunk shape is ever compiled. + return evaluate_in_chunks( + _vmapped, params, batch_size=batch_size, xp=jnp ) # The optax rule (resolved from optax / optax.contrib), guarded by @@ -2211,18 +2192,10 @@ def projected(n): def _chunk_slices(n_rows, batch_size): """ - The ``(lo, hi, pad)`` chunk bounds the batched ``value_and_grad`` sweep - iterates over: rows ``lo:hi`` of the params array, padded by ``pad`` repeats - of the final row so every chunk presents the same ``(batch_size, ndim)`` - shape to the compiled function (one XLA compile for the whole sweep). Only - the final chunk can be ragged (``pad > 0``) — the broad-start collection may - return fewer than ``n_starts`` rows, so raggedness is not restricted to - ``n_starts % batch_size``. + The ``(lo, hi, pad)`` chunk bounds of the batched ``value_and_grad`` sweep; + see ``autofit.non_linear.objective.chunk_slices``, where it now lives. """ - return [ - (lo, min(lo + batch_size, n_rows), max(0, lo + batch_size - n_rows)) - for lo in range(0, n_rows, batch_size) - ] + return chunk_slices(n_rows, batch_size) class MultiStartAdam(AbstractMultiStartGradient): diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index acd6ae760..d22bc8d29 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -11,7 +11,6 @@ import warnings from autofit import exc -from autofit.non_linear.fitness import Fitness from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.paths.null import NullPaths from autofit.non_linear.search.nest.abstract_nest import AbstractNest @@ -109,7 +108,7 @@ def __init__( number_of_cores: int = 1, silence: bool = False, force_x1_cpu: bool = False, - use_jax_jit: bool = True, + use_jax_jit: Optional[bool] = None, session: Optional[sa.orm.Session] = None, **kwargs, ): @@ -175,6 +174,18 @@ def __init__( self.maxcall = maxcall self.force_x1_cpu = force_x1_cpu + if use_jax_jit is not None: + from autofit.non_linear.search.abstract_search import ( + warn_deprecated_jax_knob, + ) + + warn_deprecated_jax_knob( + type(self).__name__, + "use_jax_jit", + "a JAX analysis is always jitted lazily (use_jax_jit=False still " + "evaluates it eagerly until the argument is removed).", + ) + self.use_jax_jit = use_jax_jit self.logger.debug(f"Creating {self.__class__.__name__} Search") @@ -235,16 +246,13 @@ def _fit( set of accepted samples of the fit. """ - fitness = Fitness( - model=model, + fitness = self.make_fitness( analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=True, - resample_figure_of_merit=-1.0e99, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - use_jax_jit=analysis.is_jax and self.use_jax_jit, + model=model, + # A JAX analysis is jitted lazily by the scalar objective `__call__` + # dispatches to; the deprecated `use_jax_jit=False` keeps its old meaning + # (eager JAX) for one release. + compile=self.use_jax_jit is not False, ) if not isinstance(self.paths, NullPaths): @@ -261,17 +269,21 @@ def _fit( "Starting new Dynesty non-linear search (no previous samples found)." ) + # The fit's effective core count: the JAX fork rule makes it 1 for a JAX + # analysis (`autofit.non_linear.parallel.pool`). + number_of_cores = self._pools().number_of_cores + finished = False while not finished: try: - if self.number_of_cores <= 1 or self.force_x1_cpu or analysis.is_jax: + if number_of_cores <= 1 or self.force_x1_cpu or analysis.is_jax: raise RuntimeError Pool = _fork_pool_cls() with Pool( - njobs=self.number_of_cores, + njobs=number_of_cores, loglike=fitness, prior_transform=prior_transform, logl_args=(model, fitness), @@ -282,7 +294,7 @@ def _fit( fitness=fitness, checkpoint_exists=checkpoint_exists, pool=pool, - queue_size=self.number_of_cores, + queue_size=number_of_cores, ) finished = self.run_search_internal(search_internal=search_internal) @@ -299,7 +311,7 @@ def _fit( self.logger.info( "Running Dynesty single-CPU per `force_x1_cpu=True` (no pool)." ) - elif self.number_of_cores <= 1: + elif number_of_cores <= 1: self.logger.info( "Running Dynesty single-CPU (number_of_cores=1, no pool)." ) @@ -567,12 +579,11 @@ def live_points_init_from(self, model, fitness): unit_parameters, parameters, log_likelihood_list, - ) = self.initializer.samples_from_model( - total_points=self.number_live_points, + ) = self.start_points( model=model, fitness=fitness, - paths=self.paths, - n_cores=self.number_of_cores, + n=self.number_live_points, + plot=False, ) init_unit_parameters = np.zeros( diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 018d7683b..c5802bfa3 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -13,7 +13,6 @@ from autofit import exc from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.mapper.prior.vectorized import PriorVectorized -from autofit.non_linear.fitness import Fitness from autofit.non_linear.parallel import fork_context from autofit.non_linear.paths.null import NullPaths from autofit.non_linear.search.nest import abstract_nest @@ -214,7 +213,7 @@ def __init__( silence: bool = False, force_x1_cpu: bool = False, session: Optional[sa.orm.Session] = None, - use_jax_vmap: bool = True, + use_jax_vmap: Optional[bool] = None, **kwargs, ): """ @@ -298,6 +297,19 @@ def __init__( self.verbose = verbose self.force_x1_cpu = force_x1_cpu + if use_jax_vmap is not None: + from autofit.non_linear.search.abstract_search import ( + warn_deprecated_jax_knob, + ) + + warn_deprecated_jax_knob( + type(self).__name__, + "use_jax_vmap", + "a JAX analysis is evaluated through the batched objective " + "automatically (use_jax_vmap=False still evaluates it point by point " + "until the argument is removed).", + ) + self.use_jax_vmap = use_jax_vmap if is_test_mode(): @@ -312,26 +324,35 @@ def apply_test_mode(self): ) self.n_like_max = 1 - def _fit(self, model: AbstractPriorModel, analysis): + def fitness_overrides(self, analysis) -> dict: + """ + The single-process path evaluates batches: a JAX analysis through the batched + objective, `n_batch` points per call (`force_x1_cpu` with a numpy analysis + evaluates point by point). + """ + if self.force_x1_cpu or analysis.is_jax: + return { + "batched": self.use_jax_vmap is not False and analysis.is_jax, + "batch_size": self.n_batch, + } + return {} + + def run(self, ctx): """ - Fit a model using the search and the Analysis class which contains the data and returns the log likelihood from - instances of the model, which the `NonLinearSearch` seeks to maximize. + Run Nautilus (the ``run(ctx)`` hook; see ``docs/design/run_ctx.md``) and return + the sampler. + + A JAX analysis (or ``force_x1_cpu``) runs in this process through the fit's + objective, batched on JAX; otherwise likelihoods are spread over the fit's pool. + Nautilus resumes from its own native checkpoint (``checkpoint.hdf5``). Parameters ---------- - model : ModelMapper - The model which generates instances for different points in parameter space. - analysis : Analysis - Contains the data and the log likelihood function which fits an instance of the model to the data, returning - the log likelihood the `NonLinearSearch` maximizes. - - Returns - ------- - A result object comprising the Samples object that includes the maximum log likelihood instance and full - set of accepted ssamples of the fit. + ctx + The fit's ``FitContext``. """ - if not isinstance(self.paths, NullPaths): + if not isinstance(ctx.paths, NullPaths): checkpoint_exists = Path(self.checkpoint_file).exists() else: checkpoint_exists = False @@ -346,49 +367,19 @@ def _fit(self, model: AbstractPriorModel, analysis): "Starting new Nautilus non-linear search (no previous samples found)." ) - if self.force_x1_cpu or analysis.is_jax: - - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=True, - resample_figure_of_merit=-1.0e99, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - # 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, + if self.force_x1_cpu or ctx.pool.is_jax: + return self.fit_x1_cpu( + fitness=ctx.fitness, + model=ctx.model, + update=ctx.update, ) - search_internal = self.fit_x1_cpu( - fitness=fitness, - model=model, - analysis=analysis, - ) - - else: - - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=True, - resample_figure_of_merit=-1.0e99, - iterations_per_quick_update=self.iterations_per_quick_update, - background_quick_update=self.quick_update_background, - live_visual_update=self.live_visual_update, - ) - - search_internal = self.fit_multiprocessing( - fitness=fitness, - model=model, - analysis=analysis, - ) - - return search_internal, fitness + return self.fit_multiprocessing( + fitness=ctx.fitness, + model=ctx.model, + pools=ctx.pool, + update=ctx.update, + ) @property def sampler_cls(self): @@ -418,7 +409,7 @@ def checkpoint_file(self): except TypeError: pass - def fit_x1_cpu(self, fitness, model, analysis): + def fit_x1_cpu(self, fitness, model, analysis=None, update=None): """ Perform the non-linear search, using one CPU core. @@ -437,7 +428,7 @@ def fit_x1_cpu(self, fitness, model, analysis): the log likelihood the search maximizes. """ - if analysis.is_jax: + if fitness.analysis.is_jax: self.logger.info( "Running search with JAX vectorization (parallelization handled by JAX)." ) @@ -452,7 +443,7 @@ def fit_x1_cpu(self, fitness, model, analysis): n_dim=model.prior_count, filepath=self.checkpoint_file, pool=None, - vectorized=fitness.use_jax_vmap, + vectorized=fitness.batched, n_live=self.n_live, n_update=self.n_update, enlarge_per_dim=self.enlarge_per_dim, @@ -464,9 +455,15 @@ def fit_x1_cpu(self, fitness, model, analysis): seed=self.seed, ) - return self.call_search(search_internal=search_internal, model=model, analysis=analysis, fitness=fitness) + return self.call_search( + search_internal=search_internal, + model=model, + analysis=analysis, + fitness=fitness, + update=update, + ) - def fit_multiprocessing(self, fitness, model, analysis): + def fit_multiprocessing(self, fitness, model, analysis=None, pools=None, update=None): """ Perform the non-linear search, using multiple CPU cores parallelized via Python's multiprocessing module. @@ -520,14 +517,16 @@ def fit_multiprocessing(self, fitness, model, analysis): # unpickles a fresh copy for every likelihood call. Worker memory then # grows by a few MB per call (8-core RAL runs reached the 96 GB cgroup # limit) while the parent starves the pool serialising it (#1547). - if self.number_of_cores <= 1: + pools = pools if pools is not None else self._pools() + + if pools.number_of_cores <= 1: pool_context = nullcontext(None) else: from nautilus.pool import initialize_worker pool_context = _LikelihoodWorkerPool( fork_context().Pool( - self.number_of_cores, + pools.number_of_cores, initializer=initialize_worker, initargs=(fitness.call_wrap,), ) @@ -556,7 +555,8 @@ def fit_multiprocessing(self, fitness, model, analysis): search_internal=search_internal, model=model, analysis=analysis, - fitness=fitness + fitness=fitness, + update=update, ) # Drop the pool references so their finalizers don't fire at interpreter @@ -574,7 +574,7 @@ def fit_multiprocessing(self, fitness, model, analysis): return search_internal - def call_search(self, search_internal, model, analysis, fitness): + def call_search(self, search_internal, model, analysis, fitness, update=None): """ The x1 CPU and multiprocessing searches both call this function to perform the non-linear search. @@ -652,13 +652,16 @@ def call_search(self, search_internal, model, analysis, fitness): if not finished: - self.perform_update( - model=model, - analysis=analysis, - during_analysis=True, - fitness=fitness, - search_internal=search_internal - ) + if update is not None: + update(search_internal) + else: + self.perform_update( + model=model, + analysis=analysis, + during_analysis=True, + fitness=fitness, + search_internal=search_internal + ) return search_internal @@ -743,51 +746,26 @@ def samples_info_from(self, search_internal=None): "number_live_points": int(search_internal.n_live), } - def samples_via_internal_from( - self, model: AbstractPriorModel, search_internal=None - ): - """ - Returns a `Samples` object from the nautilus internal results. - - The samples contain all information on the parameter space sampling (e.g. the parameters, - log likelihoods, etc.). + samples_cls = SamplesNest - The internal search results are converted from the native format used by the search to lists of values - (e.g. `parameter_lists`, `log_likelihood_list`). - - Parameters - ---------- - model - Maps input vectors of unit parameter values to physical values and model instances via priors. + def raw_samples_from(self, model, internal): """ + Nautilus's posterior as ``RawSamples``: parameters, log likelihoods and the + exponentiated log weights. + """ + from autofit.non_linear.search.fit_context import RawSamples - if search_internal is None: - search_internal = self.paths.load_search_internal() - - parameters, log_weights, log_likelihoods = search_internal.posterior() - - parameter_lists = parameters.tolist() - log_likelihood_list = log_likelihoods.tolist() - weight_list = np.exp(log_weights).tolist() - - log_prior_list = [ - sum(model.log_prior_list_from_vector(vector=vector)) - for vector in parameter_lists - ] + parameters, log_weights, log_likelihoods = internal.posterior() - sample_list = Sample.from_lists( - model=model, - parameter_lists=parameter_lists, - log_likelihood_list=log_likelihood_list, - log_prior_list=log_prior_list, - weight_list=weight_list, + return RawSamples( + parameters=parameters.tolist(), + log_likelihood=log_likelihoods.tolist(), + weights=np.exp(log_weights).tolist(), + info=self.info_from(internal), ) - return SamplesNest( - model=model, - sample_list=sample_list, - samples_info=self.samples_info_from(search_internal=search_internal), - ) + def info_from(self, internal): + return self.samples_info_from(search_internal=internal) @property def batch_size(self): diff --git a/autofit/non_linear/settings.py b/autofit/non_linear/settings.py index 04d2537f0..61c570747 100644 --- a/autofit/non_linear/settings.py +++ b/autofit/non_linear/settings.py @@ -16,7 +16,7 @@ def __init__( number_of_cores: Optional[int] = 1, session: Optional[sa.orm.Session] = None, info: Optional[dict] = None, - use_jax_vmap: bool = True, + use_jax_vmap: Optional[bool] = None, ): """ Stores all the input settings that are used in search's and their `fit functions. @@ -40,7 +40,21 @@ def __init__( info Optional dictionary containing information about the model-fit that is stored in the database and can be loaded by the aggregator after the model-fit is complete. + use_jax_vmap + Deprecated: a JAX analysis is evaluated through the batched objective of a search that has one + automatically. Still forwarded (to the searches that accept it) for one release. """ + if use_jax_vmap is not None: + from autofit.non_linear.search.abstract_search import ( + warn_deprecated_jax_knob, + ) + + warn_deprecated_jax_knob( + "SettingsSearch", + "use_jax_vmap", + "a JAX analysis is evaluated through the batched objective of a search " + "that has one automatically.", + ) self.path_prefix = path_prefix self.unique_tag = unique_tag diff --git a/docs/design/run_ctx.md b/docs/design/run_ctx.md index 9fefdff24..91c4cb18d 100644 --- a/docs/design/run_ctx.md +++ b/docs/design/run_ctx.md @@ -115,3 +115,31 @@ test-mode bypass and the fail-fast gates have run. 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. + +## Revision 1 (phase A2, 2026-10-08): what the bridge shipped + +The signature `run(self, ctx) -> Any` and the member list above are unchanged. A2 +pins the following details the table left open; each is flagged for review in the A2 PR. + +- **`ctx.schedule`** takes the backend's total budget, which is a search setting the + context cannot know: `next_budget(done, total)` and `chunks(total)`. +- **`ctx.start_points(n)`** returns `(parameters, figures_of_merit)` as `(n, n_dim)` and + `(n,)` arrays. A search with no starting point to plot (`Drawer`) sets the class + attribute `_plots_start_point = False`. +- **`ctx.update(internal)`** is synchronous: the backend is paused while it runs, so the + internal state is read as it stands (no copy). The snapshot hand-off arrives with A3's + samples adapter, which converts the samples once. +- **`ctx.resume`, `ctx.checkpointer`** are `None` until A3 fills them; **`ctx.rng`** is a + `SeedSequence` of the search's `seed` attribute, if any, until A4's search-level seed. +- **`ctx.close(failed)`** is the cleanup the lifecycle rules require; on failure it closes + the pools the context built, shuts the quick-update threads down and releases the + compiled objectives. +- **The bridge is `NonLinearSearch._fit`**: when a search defines `run`, the base `_fit` + builds the `Fitness` and the context (after the gates `start_resume_fit` runs) and calls + `run(ctx)`, so every caller of `_fit` (and the conformance suite's call counter) sees a + migrated search exactly as before. +- **Search-side hooks** (not context members): `fitness_overrides(analysis)` lets a search + add `Fitness` arguments (`Nautilus`: `batched`, `batch_size`), and `samples_cls` names the + `Samples` class `raw_samples_from`'s result is converted into. `RawSamples` and + `samples_from_raw` live provisionally in `autofit/non_linear/search/fit_context.py` until + A3's `samples/adapter.py` replaces them. diff --git a/test_autofit/analysis/test_is_jax.py b/test_autofit/analysis/test_is_jax.py index 80ea93d3f..ab9311b77 100644 --- a/test_autofit/analysis/test_is_jax.py +++ b/test_autofit/analysis/test_is_jax.py @@ -33,7 +33,9 @@ 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) + # Traceable on both backends: a `use_jax=True` analysis is jitted by every search + # (search-extensibility phase A2), so a host `float(...)` would not trace. + return -0.5 * (instance.centre - 50.0) ** 2 class ForwardAnalysis(NumpyAnalysis): diff --git a/test_autofit/graphical/test_ep_no_multiprocessing.py b/test_autofit/graphical/test_ep_no_multiprocessing.py index a54c6b13c..7d5709f93 100644 --- a/test_autofit/graphical/test_ep_no_multiprocessing.py +++ b/test_autofit/graphical/test_ep_no_multiprocessing.py @@ -52,6 +52,7 @@ def no_fork_context(*args, **kwargs): module_paths = ( "autofit.non_linear.parallel", "autofit.non_linear.parallel.context", + "autofit.non_linear.parallel.pool", "autofit.non_linear.search.abstract_search", "autofit.non_linear.search.nest.dynesty.search.abstract", "autofit.non_linear.search.nest.nautilus.search", diff --git a/test_autofit/non_linear/search/mcmc/test_emcee_jax_timing.py b/test_autofit/non_linear/search/mcmc/test_emcee_jax_timing.py new file mode 100644 index 000000000..888ee2e1b --- /dev/null +++ b/test_autofit/non_linear/search/mcmc/test_emcee_jax_timing.py @@ -0,0 +1,56 @@ +""" +Eager-versus-jit timing guard for Emcee on a JAX analysis (search-extensibility A2). + +Emcee declares `jax_use='none'` and evaluates one walker at a time through +`Fitness.call_wrap`. Before A2 that ran the JAX likelihood eagerly, op by op (7.80 ms +eager vs 0.29 ms jitted per call on `af.ex.Gaussian`, survey 02 §3.1). Now `call_wrap` +dispatches to the lazily jitted scalar objective, so the per-call cost Emcee pays must +stay within 2x the jitted call itself. +""" +import time + +import numpy as np +import pytest + +import autofit as af + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +def _per_call(func, vector, n=200, repeats=5, sync=False): + func(vector) + best = float("inf") + for _ in range(repeats): + start = time.perf_counter() + for _ in range(n): + result = func(vector) + if sync: + float(result) + best = min(best, (time.perf_counter() - start) / n) + return best + + +def test_emcee_call_path_runs_at_jitted_speed(monkeypatch): + pytest.importorskip("jax") + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + x = np.arange(100.0) + data = 3.0 * np.exp(-0.5 * ((x - 50.0) / 8.0) ** 2) + analysis = af.ex.Analysis(data=data, noise_map=np.full(100, 0.1), use_jax=True) + model = af.Model(af.ex.Gaussian) + + fitness = af.Emcee().make_fitness(analysis=analysis, model=model) + vector = np.array(model.physical_values_from_prior_medians) + + jitted = fitness.objective("scalar") + + # Emcee's log_prob_fn is `fitness.call_wrap`, which dispatches to the jitted objective. + assert fitness._call is jitted + assert hasattr(jitted, "__wrapped__") + + jit_per_call = _per_call(jitted, vector, sync=True) + emcee_per_call = _per_call(fitness.call_wrap, vector) + + assert emcee_per_call <= 2.0 * jit_per_call, ( + f"call_wrap {emcee_per_call * 1e3:.3f} ms vs jit {jit_per_call * 1e3:.3f} ms" + ) diff --git a/test_autofit/non_linear/search/mle/test_bfgs_jax_gradient.py b/test_autofit/non_linear/search/mle/test_bfgs_jax_gradient.py new file mode 100644 index 000000000..c20d8499b --- /dev/null +++ b/test_autofit/non_linear/search/mle/test_bfgs_jax_gradient.py @@ -0,0 +1,74 @@ +""" +BFGS / LBFGS on a JAX analysis minimise with the exact gradient of the "value_and_grad" +objective (`jac=True`), not scipy's finite differences (search-extensibility phase A2). + +With `jac=True` every scipy function evaluation is one call of the jitted value-and-gradient +objective, which returns the gradient too: `nfev` equals the number of those calls and no +scalar (finite-difference) call happens inside `minimize`. Finite differences cost +`n_params` extra function evaluations per gradient, so on the numpy path +`nfev >= (n_params + 1) * njev`. +""" +import numpy as np +import pytest + +import autofit as af + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +def _data(): + x = np.arange(50) + signal = 3.0 * np.exp(-0.5 * ((x - 25.0) / 5.0) ** 2) + return signal, np.full(50, 0.1) + + +def _model(): + model = af.Model(af.ex.Gaussian) + model.centre = af.UniformPrior(lower_limit=10.0, upper_limit=40.0) + model.normalization = af.UniformPrior(lower_limit=0.5, upper_limit=6.0) + model.sigma = af.UniformPrior(lower_limit=1.0, upper_limit=10.0) + return model + + +@pytest.mark.parametrize("search_cls", [af.BFGS, af.LBFGS]) +def test_jax_analysis_uses_the_exact_gradient(search_cls, monkeypatch): + pytest.importorskip("jax") + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + data, noise_map = _data() + analysis = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=True) + + from autofit.non_linear.fitness import Fitness + + calls = {"value_and_grad": 0} + original = Fitness.call_wrap_value_and_grad + + def counted(self, parameters): + calls["value_and_grad"] += 1 + return original(self, parameters) + + monkeypatch.setattr(Fitness, "call_wrap_value_and_grad", counted) + + search = search_cls(maxiter=50) + result = search.fit(model=_model(), analysis=analysis) + + internal = result.search_internal + n_params = 3 + + assert internal.njev > 0 + assert calls["value_and_grad"] == internal.nfev + assert internal.nfev < (n_params + 1) * internal.njev + + +def test_numpy_analysis_still_differences_finitely(monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + data, noise_map = _data() + analysis = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=False) + + search = af.BFGS(maxiter=20) + result = search.fit(model=_model(), analysis=analysis) + + internal = result.search_internal + + assert internal.nfev >= 4 * internal.njev 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 index 57fdabd44..2f504af47 100644 --- 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 @@ -26,7 +26,9 @@ 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 + # Default: the deprecated knob is unset, so the batched objective is chosen by + # the analysis alone (numpy -> point by point). + assert search.use_jax_vmap is None result = search.fit(model=af.Model(af.ex.Gaussian), analysis=NumpyAnalysis()) diff --git a/test_autofit/non_linear/search/test_abstract_search.py b/test_autofit/non_linear/search/test_abstract_search.py index 9823854ee..f5eb0b0c1 100644 --- a/test_autofit/non_linear/search/test_abstract_search.py +++ b/test_autofit/non_linear/search/test_abstract_search.py @@ -912,10 +912,29 @@ def test_settings_search_dict_for_filters_use_jax_vmap(caplog): settings_search = af.SettingsSearch(path_prefix="settings_search_test") assert "use_jax_vmap" in settings_search.search_dict - assert settings_search.search_dict_for(af.Nautilus)["use_jax_vmap"] is True + # The deprecated knob keeps its key (search_dict keeps its shape); its default is + # None, which a search reads as "batched when the analysis is JAX" and never warns. + assert settings_search.search_dict_for(af.Nautilus)["use_jax_vmap"] is None assert "use_jax_vmap" not in settings_search.search_dict_for(af.DynestyStatic) with caplog.at_level(logging.WARNING): af.DynestyStatic(**settings_search.search_dict_for(af.DynestyStatic)) assert _unknown_kwarg_messages(caplog) == [] + + +def test_deprecated_jax_knobs_warn_only_when_passed(): + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + af.SettingsSearch(path_prefix="settings_search_test") + af.Nautilus() + af.DynestyStatic() + + with pytest.warns(FutureWarning, match="SettingsSearch\\(use_jax_vmap"): + af.SettingsSearch(path_prefix="settings_search_test", use_jax_vmap=True) + with pytest.warns(FutureWarning, match="Nautilus\\(use_jax_vmap"): + af.Nautilus(use_jax_vmap=False) + with pytest.warns(FutureWarning, match="DynestyStatic\\(use_jax_jit"): + af.DynestyStatic(use_jax_jit=True) diff --git a/test_autofit/non_linear/search/test_fit_context.py b/test_autofit/non_linear/search/test_fit_context.py new file mode 100644 index 000000000..1cbd88c19 --- /dev/null +++ b/test_autofit/non_linear/search/test_fit_context.py @@ -0,0 +1,157 @@ +""" +The minimal `run(ctx)` bridge (search-extensibility phase A2, `docs/design/run_ctx.md`). + +`Drawer` and `Nautilus` run through `run(ctx)`; a search that overrides `_fit` keeps +working; a toy search written only against the frozen contract fits end to end. +""" +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.search import capabilities as cap +from autofit.non_linear.search.abstract_search import NonLinearSearch +from autofit.non_linear.search.fit_context import FitContext, RawSamples + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +FROZEN_MEMBERS = ( + "model", + "paths", + "objective", + "fitness", + "test_mode_level", + "pool", + "start_points", + "rng", + "resume", + "checkpointer", + "schedule", + "update", +) + + +def _analysis(): + return af.ex.Analysis( + data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=False + ) + + +@pytest.mark.parametrize("search_cls", [af.Drawer, af.Nautilus]) +def test_proofs_are_on_the_bridge(search_cls): + assert search_cls.run is not NonLinearSearch.run + assert search_cls.raw_samples_from is not NonLinearSearch.raw_samples_from + assert search_cls._fit is NonLinearSearch._fit + + +def test_drawer_fit_runs_through_run_ctx(monkeypatch): + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + seen = [] + original = af.Drawer.run + + def spy(self, ctx): + seen.append(ctx) + return original(self, ctx) + + monkeypatch.setattr(af.Drawer, "run", spy) + + search = af.Drawer(total_draws=5) + result = search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert len(seen) == 1 + ctx = seen[0] + assert isinstance(ctx, FitContext) + for member in FROZEN_MEMBERS: + assert hasattr(ctx, member), member + assert ctx.test_mode_level == 1 + assert ctx.resume is None and ctx.checkpointer is None + + # Never stored on the search. + assert not any(isinstance(value, FitContext) for value in vars(search).values()) + + assert len(result.samples.sample_list) == 5 + + +class ToySearch(NonLinearSearch): + """ + A search written only against the frozen `run(ctx)` contract: draw points, evaluate + them through the batched objective, keep them all. + """ + + posterior_kind = cap.PosteriorKind.POINT + objective_target = cap.ObjectiveTarget( + cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.PHYSICAL + ) + invalid_value = -float("inf") + + def run(self, ctx): + parameters, _ = ctx.start_points(6) + log_likelihoods = np.asarray(ctx.objective("batched")(parameters)) + return {"parameters": parameters, "log_likelihoods": log_likelihoods} + + def raw_samples_from(self, model, internal): + return RawSamples( + parameters=internal["parameters"].tolist(), + log_likelihood=internal["log_likelihoods"].tolist(), + info={"total_samples": len(internal["log_likelihoods"])}, + ) + + def plot_results(self, samples): + pass + + +def test_a_new_search_needs_only_run_and_raw_samples_from(monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + search = ToySearch(initializer=af.InitializerPrior()) + model = af.Model(af.ex.Gaussian) + result = search.fit(model=model, analysis=_analysis()) + + assert len(result.samples.sample_list) == 6 + assert result.samples.max_log_likelihood() is not None + + +def test_failure_in_run_closes_the_context(monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + closed = [] + + def failing_run(self, ctx): + ctx.objective("scalar") + original_close = ctx.close + + def close(failed): + closed.append(failed) + return original_close(failed) + + ctx.close = close + raise RuntimeError("backend failed") + + monkeypatch.setattr(ToySearch, "run", failing_run) + + with pytest.raises(RuntimeError, match="backend failed"): + ToySearch().fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert closed == [True] + + +def test_a_search_with_neither_hook_is_abstract(): + import inspect + + class Empty(NonLinearSearch): + pass + + assert inspect.isabstract(NonLinearSearch) + assert inspect.isabstract(Empty) + assert not inspect.isabstract(ToySearch) + assert not inspect.isabstract(af.Drawer) + + with pytest.raises(TypeError): + Empty() + + +def test_schedule_chunks_cover_the_budget(): + search = af.Drawer(iterations_per_full_update=4) + from autofit.non_linear.search.fit_context import UpdateSchedule + + assert list(UpdateSchedule(search).chunks(10)) == [4, 4, 2] diff --git a/test_autofit/non_linear/search/test_make_fitness.py b/test_autofit/non_linear/search/test_make_fitness.py new file mode 100644 index 000000000..44c1d3d8f --- /dev/null +++ b/test_autofit/non_linear/search/test_make_fitness.py @@ -0,0 +1,96 @@ +""" +`NonLinearSearch.make_fitness`, the one `Fitness` construction site the searches share +(search-extensibility phase A2). + +It replaces `test_quick_update_wiring.py`, which AST-scanned every `Fitness(...)` call +under the search tree for a forwarded `iterations_per_quick_update`: a search that +forgot it silently never fired quick updates (PyAutoFit#1434). With every search +building its `Fitness` through `make_fitness`, which always forwards the cadence, that +scan has nothing left to guard; what remains is that `make_fitness` forwards it, and that +no search builds a `Fitness` any other way (NSS keeps its own path until phase A3b). +""" +import ast +from pathlib import Path + +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.search import capabilities as cap + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +SEARCH_ROOT = Path(__file__).parents[3] / "autofit" / "non_linear" / "search" + + +def _analysis(): + return af.ex.Analysis(data=np.ones(10), noise_map=np.ones(10), use_jax=False) + + +def _fitness_for(search_cls, **overrides): + search = search_cls(iterations_per_quick_update=7) + return search, search.make_fitness( + analysis=_analysis(), model=af.Model(af.ex.Gaussian), **overrides + ) + + +@pytest.mark.parametrize( + "search_cls, fom_is_log_likelihood, convert_to_chi_squared, resample", + [ + (af.Emcee, False, False, -np.inf), + (af.DynestyStatic, True, False, -1.0e99), + (af.Nautilus, True, False, -1.0e99), + (af.Drawer, False, False, -np.inf), + (af.BFGS, False, True, -np.inf), + (af.LBFGS, False, True, -np.inf), + ], +) +def test_convention_follows_the_declared_target( + search_cls, fom_is_log_likelihood, convert_to_chi_squared, resample +): + search, fitness = _fitness_for(search_cls) + + assert fitness.fom_is_log_likelihood is fom_is_log_likelihood + assert fitness.convert_to_chi_squared is convert_to_chi_squared + assert fitness.resample_figure_of_merit == resample + + +def test_neg2_resample_reaches_the_backend_as_the_declared_invalid_value(): + search, fitness = _fitness_for(af.BFGS) + + assert search.objective_target.quantity == cap.ObjectiveQuantity.NEG2_LOG_POSTERIOR + assert fitness.resample_figure_of_merit * -2.0 == type(search).invalid_value + + +def test_quick_update_settings_are_always_forwarded(): + search, fitness = _fitness_for(af.Emcee) + + assert fitness.iterations_per_quick_update == search.iterations_per_quick_update == 7 + assert fitness.live_visual_update == search.live_visual_update + assert fitness.paths is search.paths + + +def test_overrides_win(): + _, fitness = _fitness_for(af.Nautilus, batched=True, store_history=True) + + assert fitness.batched is True + assert fitness.store_history is True + + +def _fitness_constructions(): + for path in sorted(SEARCH_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + name = func.id if isinstance(func, ast.Name) else getattr(func, "attr", None) + if name == "Fitness": + yield path.relative_to(SEARCH_ROOT).as_posix() + + +def test_only_make_fitness_and_nss_construct_a_fitness(): + assert sorted(_fitness_constructions()) == [ + "abstract_search.py", + "nest/nss/search.py", + ] diff --git a/test_autofit/non_linear/search/test_quick_update_wiring.py b/test_autofit/non_linear/search/test_quick_update_wiring.py deleted file mode 100644 index 58204d8e4..000000000 --- a/test_autofit/non_linear/search/test_quick_update_wiring.py +++ /dev/null @@ -1,87 +0,0 @@ -import ast -from pathlib import Path - -import pytest - -# Every search builds its own `Fitness`, and the on-the-fly quick-update machinery -# only runs if that construction forwards `iterations_per_quick_update`. Forgetting -# it fails *silently*: `Fitness.manage_quick_update` returns at its `is None` guard, -# so the search runs perfectly and simply never updates. -# -# That is not hypothetical. Before PyAutoFit#1434 only Nautilus forwarded it, so -# setting a real cadence did nothing under Dynesty, Emcee, Zeus, BlackJAX NUTS, -# BFGS or Drawer — and once the cadence is announced in the startup log, a missing -# forward turns a dead feature into a false claim in the CLI. -# -# Checked structurally rather than by running each search: the samplers are -# optional dependencies and their `_fit` bodies need real data, while the defect -# is entirely visible in the call site. - -SEARCH_ROOT = Path(__file__).parents[3] / "autofit" / "non_linear" / "search" - -REQUIRED_KWARG = "iterations_per_quick_update" - -# path suffix -> why this construction site legitimately omits the kwarg. -EXEMPT = { - "nest/nss/search.py": ( - "af.NSS samples through inline JAX closures, never through Fitness -- " - "its Fitness exists only for AbstractNest.perform_update post-fit " - "work (latent samples, batch_size), so forwarding the cadence there " - "would wire it to a dead path. Quick updates fire through the " - "search's own outer loop instead: `iterations_per_quick_update` is " - "consumed directly in `_fit` via `_fire_quick_update`." - ), - # mle/multi_start_gradient/search.py was exempt here from PyAutoFit#1433 - # (its `fitness.call` is differentiated inside jit/vmap, so the counter in - # `call_wrap` cannot run) until PyAutoFit#1552 wired the cadence through - # the host-side step loop instead: the search now forwards the kwarg to its - # `Fitness` and feeds the counter one gradient step at a time via - # `_manage_quick_update_step`. -} - - -def _fitness_call_sites(): - """Every `Fitness(...)` construction under the search tree, as - (path-relative-to-search-root, set-of-keyword-names).""" - for path in sorted(SEARCH_ROOT.rglob("*.py")): - tree = ast.parse(path.read_text(encoding="utf-8")) - - for node in ast.walk(tree): - if not isinstance(node, ast.Call): - continue - - func = node.func - name = func.id if isinstance(func, ast.Name) else getattr(func, "attr", None) - - if name != "Fitness": - continue - - keywords = {kw.arg for kw in node.keywords if kw.arg is not None} - yield path.relative_to(SEARCH_ROOT).as_posix(), keywords - - -def test__the_scan_finds_the_call_sites_it_is_meant_to_guard(): - # Without this, a refactor that renamed or re-exported `Fitness` would make - # every assertion below vacuously pass. - sites = list(_fitness_call_sites()) - - assert len(sites) >= 8, sites - assert any("nautilus" in path for path, _ in sites) - assert any("dynesty" in path for path, _ in sites) - - -@pytest.mark.parametrize( - "relative_path, keywords", - list(_fitness_call_sites()), - ids=lambda value: value if isinstance(value, str) else "", -) -def test__every_search_forwards_the_quick_update_cadence(relative_path, keywords): - if relative_path in EXEMPT: - pytest.skip(EXEMPT[relative_path]) - - assert REQUIRED_KWARG in keywords, ( - f"{relative_path} builds a Fitness without `{REQUIRED_KWARG}`, so " - "quick updates silently never fire for this search while the startup " - "log still announces a cadence. Forward it, or add the site to EXEMPT " - "with the reason." - ) diff --git a/test_autofit/non_linear/search/test_start_points.py b/test_autofit/non_linear/search/test_start_points.py new file mode 100644 index 000000000..df9406617 --- /dev/null +++ b/test_autofit/non_linear/search/test_start_points.py @@ -0,0 +1,73 @@ +""" +`NonLinearSearch.start_points(model, fitness, n)`: the one place the searches draw their +initial points, which makes the start-point plot exactly once (search-extensibility A2). +""" +import numpy as np +import pytest + +import autofit as af + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +def _search_and_fitness(monkeypatch, use_jax=False, number_of_cores=1): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + search = af.Emcee(number_of_cores=number_of_cores) + model = af.Model(af.ex.Gaussian) + analysis = af.ex.Analysis(data=np.ones(10), noise_map=np.ones(10), use_jax=use_jax) + return search, model, search.make_fitness(analysis=analysis, model=model) + + +def _count_plots(monkeypatch, search): + plotted = [] + monkeypatch.setattr( + search, + "plot_start_point", + lambda parameter_vector, model, analysis: plotted.append(parameter_vector), + ) + return plotted + + +def test_plots_the_first_point_once(monkeypatch): + search, model, fitness = _search_and_fitness(monkeypatch) + plotted = _count_plots(monkeypatch, search) + + unit, parameters, figures_of_merit = search.start_points( + model=model, fitness=fitness, n=4 + ) + + assert len(parameters) == len(unit) == len(figures_of_merit) == 4 + assert plotted == [parameters[0]] + + +def test_plot_false_skips_the_plot(monkeypatch): + search, model, fitness = _search_and_fitness(monkeypatch) + plotted = _count_plots(monkeypatch, search) + + search.start_points(model=model, fitness=fitness, n=2, plot=False) + + assert plotted == [] + + +def test_initializer_gets_the_effective_core_count(monkeypatch): + pytest.importorskip("jax") + from autofit.non_linear.parallel import PoolFactory + + search, model, fitness = _search_and_fitness( + monkeypatch, use_jax=True, number_of_cores=2 + ) + _count_plots(monkeypatch, search) + + seen = {} + original = search.initializer.samples_from_model + + def spy(**kwargs): + seen.update(kwargs) + return original(**kwargs) + + monkeypatch.setattr(search.initializer, "samples_from_model", spy) + + search._fit_pools = PoolFactory(search, fitness.analysis) + search.start_points(model=model, fitness=fitness, n=2) + + assert seen["n_cores"] == 1 diff --git a/test_autofit/non_linear/test_dict.py b/test_autofit/non_linear/test_dict.py index eff43521b..89495c64f 100644 --- a/test_autofit/non_linear/test_dict.py +++ b/test_autofit/non_linear/test_dict.py @@ -44,7 +44,7 @@ def make_dynesty_dict(): "slices": 5, "unique_tag": None, "update_interval": None, - "use_jax_jit": True, + "use_jax_jit": None, "walks": 5, }, "class_path": "autofit.non_linear.search.nest.dynesty.search.static.DynestyStatic", diff --git a/test_autofit/non_linear/test_fitness_objective.py b/test_autofit/non_linear/test_fitness_objective.py new file mode 100644 index 000000000..d01abca41 --- /dev/null +++ b/test_autofit/non_linear/test_fitness_objective.py @@ -0,0 +1,227 @@ +""" +`Fitness.objective(kind)`, the one objective factory (search-extensibility phase A2). + +The compile-count probe follows `test_fitness_vmap_cache.py`: under `jax.jit` the +likelihood's Python body runs only while tracing, so a counting analysis counts traces, and +`objective(kind).__wrapped__._cache_size()` counts compiled executables. Each kind must +compile exactly once for a fixed input shape, however often it is called. +""" +import numpy as np +import pytest + +import autofit as af +from autofit import exc +from autofit.non_linear.fitness import Fitness +from autofit.non_linear.objective import ( + OBJECTIVE_KINDS, + chunk_slices, + evaluate_in_chunks, +) + + +class CountingAnalysis(af.ex.Analysis): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.trace_count = 0 + + def log_likelihood_function(self, instance, **kwargs): + self.trace_count += 1 + return super().log_likelihood_function(instance=instance, **kwargs) + + +def _fitness(use_jax, **kwargs): + model = af.Model(af.ex.Gaussian) + analysis = CountingAnalysis( + data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=use_jax + ) + return Fitness(model=model, analysis=analysis, **kwargs) + + +def _vector(fitness): + return np.array(fitness.model.physical_values_from_prior_medians) + + +def _input_for(kind, fitness): + vector = _vector(fitness) + if kind.startswith("batched"): + return np.tile(vector, (4, 1)) * np.linspace(0.9, 1.1, 4)[:, None] + return vector + + +def test_unknown_kind_raises(): + fitness = _fitness(use_jax=False) + with pytest.raises(ValueError, match="Unknown objective kind"): + fitness.objective("hessian") + + +def test_numpy_scalar_is_the_plain_call(): + fitness = _fitness(use_jax=False) + assert fitness.objective("scalar").__func__ is Fitness.call + assert fitness._call.__func__ is Fitness.call + + +def test_numpy_batched_loops_over_call(): + fitness = _fitness(use_jax=False) + batch = _input_for("batched", fitness) + + values = fitness.objective("batched")(batch) + + assert values.shape == (4,) + assert np.allclose(values, [fitness.call(row) for row in batch]) + + +@pytest.mark.parametrize("kind", ["value_and_grad", "batched_value_and_grad"]) +def test_numpy_grad_kinds_raise(kind): + fitness = _fitness(use_jax=False) + with pytest.raises(exc.SearchException, match="needs a JAX analysis"): + fitness.objective(kind) + + +@pytest.mark.parametrize("kind", OBJECTIVE_KINDS) +def test_one_compile_per_kind(kind): + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True) + objective = fitness.objective(kind) + parameters = _input_for(kind, fitness) + + # Building the objective compiles nothing: jit is lazy. + assert fitness.analysis.trace_count == 0 + assert objective.__wrapped__._cache_size() == 0 + + first = objective(parameters) + second = objective(parameters) + + assert fitness.analysis.trace_count == 1 + assert objective.__wrapped__._cache_size() == 1 + # Cached per kind: asking again returns the same compiled callable. + assert fitness.objective(kind) is objective + + first = first if isinstance(first, tuple) else (first,) + second = second if isinstance(second, tuple) else (second,) + for a, b in zip(first, second): + assert np.array_equal(np.asarray(a), np.asarray(b)) + + +def test_kinds_agree_with_the_eager_call(): + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True) + vector = _vector(fitness) + batch = _input_for("batched", fitness) + + expected = float(fitness.call(vector)) + + assert float(fitness.objective("scalar")(vector)) == pytest.approx(expected) + + value, grad = fitness.objective("value_and_grad")(vector) + assert float(value) == pytest.approx(expected) + assert np.asarray(grad).shape == vector.shape + + values, grads = fitness.objective("batched_value_and_grad")(batch) + assert np.allclose(np.asarray(values), np.asarray(fitness.objective("batched")(batch))) + assert np.asarray(grads).shape == batch.shape + + +def test_compile_false_is_the_eager_escape_hatch(): + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True) + vector = _vector(fitness) + + eager = fitness.objective("scalar", compile=False) + eager(vector) + eager(vector) + + # Unjitted: the Python body runs on every call, and nothing is cached. + assert fitness.analysis.trace_count == 2 + assert not hasattr(eager, "__wrapped__") + assert eager is not fitness._objectives["scalar"] + + hatch = _fitness(use_jax=True, compile=False) + hatch(vector) + hatch(vector) + assert hatch.analysis.trace_count == 2 + + +def test_call_wrap_on_jax_is_jitted_and_compiles_once_across_input_types(): + """ + The scalar path of `call_wrap` (Emcee, Zeus, Drawer, the initializer) is jitted on a JAX + analysis, and a Python list and an ndarray share one compiled executable. + """ + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True) + vector = _vector(fitness) + + a = fitness(vector) + b = fitness(list(vector)) + + assert isinstance(a, float) and a == b + assert fitness.analysis.trace_count == 1 + assert fitness.objective("scalar").__wrapped__._cache_size() == 1 + + +def test_call_wrap_value_and_grad_books_history(): + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True, store_history=True) + vector = _vector(fitness) + + value, grad = fitness.call_wrap_value_and_grad(vector) + + assert isinstance(value, float) + assert grad.dtype == np.float64 and grad.shape == vector.shape + assert len(fitness.parameters_history_list) == 1 + assert len(fitness.log_likelihood_history_list) == 1 + + +def test_objective_cache_is_stripped_on_pickle(): + pytest.importorskip("jax") + import pickle + + fitness = _fitness(use_jax=True) + vector = _vector(fitness) + fitness.objective("scalar")(vector) + + assert "_objectives" not in fitness.__getstate__() + + restored = pickle.loads(pickle.dumps(fitness)) + + # Rebuilt lazily: nothing compiled until the restored objective is first called. + assert restored.objective("scalar").__wrapped__._cache_size() == 0 + assert float(restored.objective("scalar")(vector)) == pytest.approx( + float(fitness.objective("scalar")(vector)) + ) + + +def test_deprecated_kwargs_warn_and_alias(): + with pytest.warns(FutureWarning, match="use_jax_vmap"): + fitness = _fitness(use_jax=False, use_jax_vmap=True) + assert fitness.batched is True + assert fitness.use_jax_vmap is True + + with pytest.warns(FutureWarning, match="use_jax_jit"): + fitness = _fitness(use_jax=False, use_jax_jit=False) + assert fitness.compile is False + + +def test_chunk_slices_pad_only_the_final_chunk(): + assert chunk_slices(5, 2) == [(0, 2, 0), (2, 4, 0), (4, 5, 1)] + assert chunk_slices(4, 2) == [(0, 2, 0), (2, 4, 0)] + + +def test_evaluate_in_chunks_matches_one_call_and_compiles_one_shape(): + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True) + batch = np.tile(_vector(fitness), (5, 1)) * np.linspace(0.9, 1.1, 5)[:, None] + + batched = fitness.objective("batched") + + chunked = evaluate_in_chunks(batched, batch, batch_size=2) + whole = batched(batch) + + assert np.allclose(np.asarray(chunked), np.asarray(whole)) + # Shapes (2, d) and (5, d): the padded ragged chunk reuses the (2, d) executable. + assert batched.__wrapped__._cache_size() == 2 diff --git a/test_autofit/non_linear/test_fork_rule.py b/test_autofit/non_linear/test_fork_rule.py new file mode 100644 index 000000000..5ba302a18 --- /dev/null +++ b/test_autofit/non_linear/test_fork_rule.py @@ -0,0 +1,112 @@ +""" +The one JAX fork rule and `PoolFactory` (`autofit.non_linear.parallel.pool`, +search-extensibility phase A2, decision D11). + +A JAX analysis is never evaluated through a forked pool: a search given one with +`number_of_cores > 1` runs on one core, says so in exactly one INFO line, and records the +requested and effective counts in `search.summary`. Forking is unchanged everywhere else. +""" +import logging +import zipfile + +import numpy as np +import pytest + +import autofit as af +from autofit import exc +from autofit.non_linear.parallel import pool as pool_module +from autofit.non_linear.parallel import PoolFactory, effective_number_of_cores + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +class _Analysis: + def __init__(self, is_jax): + self.is_jax = is_jax + + +def _rule_records(caplog): + return [ + record + for record in caplog.records + if record.name == pool_module.logger.name and "running with 1 core" in record.message + ] + + +def test_rule_downgrades_jax_only(caplog): + with caplog.at_level(logging.INFO): + assert effective_number_of_cores(4, is_jax=False, where="X") == 4 + assert effective_number_of_cores(1, is_jax=True, where="X") == 1 + assert _rule_records(caplog) == [] + + assert effective_number_of_cores(4, is_jax=True, where="X") == 1 + + assert len(_rule_records(caplog)) == 1 + + +def test_pool_factory_under_jax_builds_no_pool(monkeypatch): + def no_fork(): + raise AssertionError("must not fork") + + monkeypatch.setattr(pool_module, "fork_context", no_fork) + + search = af.Emcee(number_of_cores=2) + factory = PoolFactory(search, _Analysis(is_jax=True)) + + assert factory.requested == 2 + assert factory.number_of_cores == 1 + assert factory() is None + assert factory(n_cores=8) is None + assert factory.sneaky(fitness=None) is None + assert "Number of cores = 1 (requested 2)" in factory.summary + assert "JAX analyses never fork" in factory.summary + + +def test_pool_factory_on_numpy_keeps_the_request(): + search = af.Emcee(number_of_cores=3) + factory = PoolFactory(search, _Analysis(is_jax=False)) + + assert factory.number_of_cores == 3 + assert factory.summary == "Number of cores = 3 (requested 3)\n" + + +def test_emcee_jax_two_cores_runs_single_core_and_records_it(monkeypatch, caplog, tmp_path): + pytest.importorskip("jax") + from autonerves import conf + + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + def no_fork(*args, **kwargs): + raise AssertionError("must not fork") + + monkeypatch.setattr(pool_module, "fork_context", no_fork) + monkeypatch.setattr("autofit.non_linear.parallel.sneaky.SneakyPool", no_fork) + + analysis = af.ex.Analysis(data=np.ones(20), noise_map=np.ones(20), use_jax=True) + + search = af.Emcee( + name="fork_rule_emcee", + path_prefix=str(tmp_path), + number_of_cores=2, + nwalkers=8, + nsteps=4, + auto_correlation_settings=af.AutoCorrelationsSettings(check_for_convergence=False), + ) + + with caplog.at_level(logging.INFO): + search.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert len(_rule_records(caplog)) == 1 + assert search.number_of_cores == 2 # the user's request is never rewritten + + with zipfile.ZipFile(search.paths._zip_path) as archive: + summary = archive.read("search.summary").decode() + + assert "Number of cores = 1 (requested 2)" in summary + + +def test_analysis_pool_refuses_a_jax_analysis(): + from autofit.non_linear.analysis.multiprocessing import AnalysisPool + + with pytest.raises(exc.SearchException, match="JAX analyses never fork"): + AnalysisPool([_Analysis(is_jax=True)], n_cores=2)