diff --git a/autofit/non_linear/initializer.py b/autofit/non_linear/initializer.py index 5b5b57eed..e18ebaad4 100644 --- a/autofit/non_linear/initializer.py +++ b/autofit/non_linear/initializer.py @@ -378,6 +378,11 @@ def __init__( super().__init__(parameter_dict=parameter_dict_new) + # Set only by `from_result` when `n_points > 1` (or `jitter != 0.0`). + # `None` preserves the exact single-point behaviour above unchanged. + self._point_dicts: Optional[List[Dict[Prior, float]]] = None + self._point_index: int = 0 + def info_value_from(self, value : Tuple[float, float]) -> float: """ Returns the value that is used to display the starting point of the parameters in the initializer. @@ -393,6 +398,225 @@ def info_value_from(self, value : Tuple[float, float]) -> float: """ return (value[1] + value[0]) / 2.0 + @classmethod + def from_result( + cls, + result, + model: Optional[AbstractPriorModel] = None, + point: str = "max_log_likelihood", + n_points: int = 1, + jitter: float = 0.0, + seed: int = 0, + ) -> "InitializerParamStartPoints": + """ + Build an initializer whose starting point(s) are taken from a previous `Result`, for example to + warm-start a search (e.g. `BlackJAXNUTS`) from an earlier fit's best-fit point(s). + + This maps the previous result's inferred parameter values onto the (possibly different) target + model **by prior path**, never by list index or by touching prior objects directly. This is + deliberately distinct from `Result.model_centred*`, which *rewrite* priors (e.g. to `GaussianPrior`s + centred on the previous best-fit); `from_result` only ever reads values off `result.samples` and + writes a `{Prior: float}` starting-point dictionary, so the scientific priors of `model` (or, if not + given, `result.model`) are left completely untouched. + + Parameters + ---------- + result + The previous `Result` (or any object exposing `.samples` and `.model`) start points are drawn + from. + model + The model the start point(s) are mapped onto. If `None`, `result.model` is used (the common + case: continuing to fit the same model with a different / warm-started search). + point + Which point of `result.samples` to draw the starting values from: `"max_log_likelihood"` (the + default) or `"median_pdf"`. + n_points + The number of starting points to generate (e.g. one per NUTS chain). `n_points=1` (the default) + reproduces the exact single-point `InitializerParamStartPoints` behaviour (a fixed +-1e-8 + unit-space jitter). + jitter + For `n_points > 1` (or when explicitly non-zero), each of the `n_points` starting values is + offset in *physical* parameter space by `jitter * sigma * N(0, 1)`, where `sigma` is the + corresponding diagonal entry of `result.samples.covariance_matrix` (its square root) when that + covariance is available and finite; otherwise `jitter` is interpreted as an absolute physical + offset (`jitter * N(0, 1)`). `jitter=0.0` (the default) still applies a fixed, tiny (1e-8) + offset so that `n_points > 1` starting points are not bit-identical. + seed + Seed for the `numpy.random.RandomState` used to draw the (reproducible) jitter offsets. + + Returns + ------- + InitializerParamStartPoints + An initializer that draws its `n_points` starting vectors, in order, from the previous result. + + Raises + ------ + InitializerException + If `result.samples` is unavailable, if `point` is not a recognised option, or if the source + and target models cannot be matched by path (dimension mismatch or an unmatched path). + """ + samples = result.samples + + if samples is None: + raise exc.InitializerException( + "InitializerParamStartPoints.from_result: `result.samples` is None " + "(no samples available in memory or on disk) -- cannot build start points." + ) + + if point == "max_log_likelihood": + vector = samples.max_log_likelihood(as_instance=False) + elif point == "median_pdf": + vector = samples.median_pdf(as_instance=False) + else: + raise exc.InitializerException( + "InitializerParamStartPoints.from_result: `point` must be " + f"'max_log_likelihood' or 'median_pdf', got {point!r}." + ) + + # The source model is `samples.model` (not `result.model`, which for a `Result` is a *freshly + # constructed* mapper via `mapper_via_defaults_from()` and therefore has priors with different + # `id`s / ordering). `vector` above is ordered consistently with `samples.model.priors_ordered_by_id` + # (both derive from the same `Samples` object), so path-matching must start from that same model. + source_model = samples.model + source_priors = source_model.priors_ordered_by_id + + if len(source_priors) != len(vector): + raise exc.InitializerException( + "InitializerParamStartPoints.from_result: the source result's samples vector has " + f"{len(vector)} entries but its model has {len(source_priors)} priors -- cannot map by path." + ) + + value_by_source_path = {} + for prior, value in zip(source_priors, vector): + path = source_model.path_for_prior(prior) + if path is None: + raise exc.InitializerException( + "InitializerParamStartPoints.from_result: no path found for a prior in the source " + "result's model; cannot map its value by path." + ) + value_by_source_path[path] = float(value) + + target_model = model if model is not None else result.model + target_priors = target_model.priors_ordered_by_id + + base_dict = {} + missing_paths = [] + for prior in target_priors: + path = target_model.path_for_prior(prior) + if path is None or path not in value_by_source_path: + missing_paths.append(path) + continue + base_dict[prior] = value_by_source_path[path] + + if missing_paths: + raise exc.InitializerException( + "InitializerParamStartPoints.from_result: could not match the following target-model " + f"prior paths to the source result's ({point}) values: {missing_paths}. The target " + "model's free parameters must live at the same paths as the source result's model." + ) + + if n_points < 1: + raise exc.InitializerException( + f"InitializerParamStartPoints.from_result: n_points must be >= 1, got {n_points}." + ) + + if n_points == 1 and jitter == 0.0: + return cls(parameter_dict=base_dict) + + # Multi-point and/or explicitly-jittered path: draw `n_points` starting vectors, each offset in + # physical space, deterministically from `seed`. Priors are only ever read (`.id`, `path_for_prior`) + # -- never mutated. + sigma_by_source_path = {} + covariance = None + try: + covariance = np.asarray(samples.covariance_matrix) + except Exception: + covariance = None + + if ( + covariance is not None + and covariance.ndim == 2 + and covariance.shape[0] == covariance.shape[1] == len(source_priors) + ): + diagonal = np.diag(covariance) + for index, prior in enumerate(source_priors): + path = source_model.path_for_prior(prior) + sigma_value = diagonal[index] + sigma_by_source_path[path] = ( + float(np.sqrt(sigma_value)) + if np.isfinite(sigma_value) and sigma_value > 0 + else None + ) + + # `jitter == 0.0` still uses a tiny fixed magnitude so `n_points > 1` points differ from one + # another (matching the single-point class's own +-1e-8 default jitter). + magnitude = jitter if jitter != 0.0 else 1.0e-8 + random_state = np.random.RandomState(seed) + + point_dicts: List[Dict[Prior, float]] = [] + for _ in range(n_points): + point_dict = {} + for prior in target_priors: + path = target_model.path_for_prior(prior) + value = base_dict[prior] + sigma = sigma_by_source_path.get(path) + if sigma is not None: + value = value + magnitude * sigma * random_state.normal() + else: + value = value + magnitude * random_state.normal() + point_dict[prior] = value + point_dicts.append(point_dict) + + initializer = cls(parameter_dict=base_dict) + initializer._point_dicts = point_dicts + initializer._point_index = 0 + return initializer + + def _generate_unit_parameter_list(self, model: AbstractPriorModel) -> List[float]: + """ + Generate a unit vector for the model. + + When this initializer was built via `from_result(..., n_points > 1)` (or a non-zero `jitter`), + successive calls cycle through the `n_points` pre-computed starting vectors, one per call (i.e. one + per requested starting point / chain), wrapping around if more points are requested than were + generated. Otherwise this defers to the base `InitializerParamBounds` behaviour (a fixed +-1e-8 + unit-space jitter around the single stored `parameter_dict`). + """ + if self._point_dicts is None: + return super()._generate_unit_parameter_list(model) + + point_dict = self._point_dicts[self._point_index % len(self._point_dicts)] + self._point_index += 1 + + unit_parameter_list = [] + for prior in model.priors_ordered_by_id: + try: + value = point_dict[prior] + except KeyError: + key = ".".join(model.path_for_prior(prior)) + if key not in self._generated_warnings: + logger.warning( + f"Range for {key} not set in the InitializerParamStartPoints. " + f"Using defaults." + ) + self._generated_warnings.add(key) + value = prior.random(self.lower_limit, self.upper_limit) + + unit_parameter_list.append(prior.unit_value_for(value)) + + return unit_parameter_list + + def samples_from_model(self, *args, **kwargs): + """ + As `AbstractInitializer.samples_from_model`, but first resets the per-point cycling counter used by + `_generate_unit_parameter_list` (when this initializer was built via `from_result(n_points > 1)`) so + repeated fits/resumes start again from point 0. + """ + if self._point_dicts is not None: + self._point_index = 0 + + return super().samples_from_model(*args, **kwargs) + class Initializer(AbstractInitializer): def __init__(self, lower_limit: float, upper_limit: float): diff --git a/autofit/non_linear/result.py b/autofit/non_linear/result.py index d9418b844..042a954e0 100644 --- a/autofit/non_linear/result.py +++ b/autofit/non_linear/result.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from autofit.non_linear.analysis.analysis import Analysis + from autofit.non_linear.initializer import InitializerParamStartPoints from autofit import exc from autofit.mapper.prior_model.abstract import AbstractPriorModel @@ -123,6 +124,76 @@ def instance(self): def max_log_likelihood_instance(self): return self.instance + @property + def start_point(self) -> "InitializerParamStartPoints": + """ + A convenience wrapper around `InitializerParamStartPoints.from_result(self)`, using the defaults + (`point="max_log_likelihood"`, `n_points=1`, `jitter=0.0`). + + Unlike `model_centred` (and its siblings below), this does **not** modify any priors: it reads the + previous result's inferred parameter values and maps them, by path, onto a `{Prior: float}` starting + point for a search's `initializer`. The model this result was fit with (`self.model`) keeps its + original priors completely untouched -- only a separate initializer object is built. + + For the multi-chain / jittered / warm-started-onto-a-different-model case, use + `start_point_from(...)` instead. + + Returns + ------- + An initializer whose single starting point is this result's maximum log likelihood sample. + """ + from autofit.non_linear.initializer import InitializerParamStartPoints + + return InitializerParamStartPoints.from_result(self) + + def start_point_from( + self, + point: str = "max_log_likelihood", + n_points: int = 1, + jitter: float = 0.0, + seed: int = 0, + model: Optional[AbstractPriorModel] = None, + ) -> "InitializerParamStartPoints": + """ + As `start_point`, but exposing the full `InitializerParamStartPoints.from_result` interface: multiple + starting points (e.g. one per chain of a multi-chain search), physical-space jitter around each, and/or + mapping onto a different (but path-compatible) target `model`. + + Like `start_point`, this never modifies priors -- it only ever reads values off `self.samples` and + writes a `{Prior: float}` (or per-point list thereof) starting-point dictionary. Contrast with + `model_centred` / `model_centred_absolute` / `model_centred_relative` / `model_centred_max_lh_bounded` + above, which construct a *new model* with rewritten (e.g. `GaussianPrior`) priors. + + Parameters + ---------- + point + Which point of `self.samples` to draw the starting values from: `"max_log_likelihood"` (the + default) or `"median_pdf"`. + n_points + The number of starting points to generate (e.g. one per NUTS chain). + jitter + Physical-space jitter magnitude applied around each starting point (see + `InitializerParamStartPoints.from_result` for the exact definition). + seed + Seed for the deterministic jitter draws. + model + The model the start point(s) are mapped onto. If `None`, `self.model` is used. + + Returns + ------- + An initializer that draws its `n_points` starting vectors, in order, from this result. + """ + from autofit.non_linear.initializer import InitializerParamStartPoints + + return InitializerParamStartPoints.from_result( + self, + model=model, + point=point, + n_points=n_points, + jitter=jitter, + seed=seed, + ) + @property def model_centred(self) -> AbstractPriorModel: """ diff --git a/autofit/non_linear/samples/mcmc.py b/autofit/non_linear/samples/mcmc.py index c8e32e188..4c5521008 100644 --- a/autofit/non_linear/samples/mcmc.py +++ b/autofit/non_linear/samples/mcmc.py @@ -203,6 +203,51 @@ def total_steps(self) -> int: def total_walkers(self) -> int: return self.samples_info["total_walkers"] + @property + def ess_bulk(self) -> Optional[List[float]]: + """ + Bulk effective sample size per free parameter (rank-normalized, split-chain), where available (e.g. + `BlackJAXNUTS`). `None` for searches that do not populate this diagnostic. + """ + return self.samples_info.get("ess_bulk_per_param") + + @property + def ess_tail(self) -> Optional[List[float]]: + """ + Tail effective sample size per free parameter (rank-normalized, split-chain), where available (e.g. + `BlackJAXNUTS`). `None` for searches that do not populate this diagnostic. + """ + return self.samples_info.get("ess_tail_per_param") + + @property + def rhat(self) -> Optional[List[float]]: + """ + Rank-normalized split-R-hat per free parameter, where available (e.g. `BlackJAXNUTS`). Values close + to 1.0 indicate convergence; values above ~1.01 suggest the chains have not converged. + + With a single chain (`num_chains == 1`) this is a *split-chain* R-hat only (each chain is split in + half internally before comparison) -- a useful non-stationarity check, but not a genuine + multi-chain convergence diagnostic. `None` for searches that do not populate this diagnostic. + """ + return self.samples_info.get("rhat_per_param") + + @property + def n_divergent(self) -> Optional[int]: + """ + The total number of divergent transitions across all chains and samples, where available (e.g. + `BlackJAXNUTS`). `None` for searches that do not populate this diagnostic. + """ + return self.samples_info.get("n_divergent") + + @property + def tree_depths(self) -> Optional[dict]: + """ + A histogram of the NUTS trajectory tree depths reached across all chains and samples, as a dict + mapping tree depth to count, where available (e.g. ``BlackJAXNUTS``). ``None`` for searches that + do not populate this diagnostic. + """ + return self.samples_info.get("tree_depth_histogram") + @property def log_evidence(self): return None diff --git a/autofit/non_linear/search/mcmc/blackjax/chains.py b/autofit/non_linear/search/mcmc/blackjax/chains.py new file mode 100644 index 000000000..705c8e0b6 --- /dev/null +++ b/autofit/non_linear/search/mcmc/blackjax/chains.py @@ -0,0 +1,194 @@ +""" +Sampler-agnostic helpers for multi-chain BlackJAX-family searches. + +Kept separate from ``blackjax/nuts/search.py`` so a future second BlackJAX +sampler (e.g. HMC) can share the same warm-start / multi-chain / inverse +mass matrix plumbing without importing NUTS-specific code. Following the +existing module convention in this package (see ``_ess_per_param_from`` in +``nuts/search.py``), JAX/blackjax imports live *inside* the functions that +need them so this module -- and anything that merely imports it for the +pure-numpy helpers -- stays importable without the ``[optional]`` extras +installed. +""" + +from __future__ import annotations + +from typing import Optional, Tuple, Union + +import numpy as np + +# The `inverse_mass_matrix` kwarg on `BlackJAXNUTS` accepts any of these. +InverseMassMatrixSpec = Union[None, str, np.ndarray, object] + +_VALID_STRING_KINDS = ("diagonal", "dense") + + +def inverse_mass_matrix_kind_from(spec: InverseMassMatrixSpec) -> str: + """ + Classify an ``inverse_mass_matrix`` specification into the small set of string kinds recorded in + ``BlackJAXNUTS.__identifier_fields__`` and persisted onto ``search_internal`` / ``samples_info``. + + Parameters + ---------- + spec + ``None``, ``"diagonal"``, ``"dense"``, a numpy array (seed matrix), or a `Result`/`Samples`-like + object (anything exposing ``.covariance_matrix`` directly, or ``.samples.covariance_matrix``). + + Returns + ------- + One of ``"none"``, ``"diagonal"``, ``"dense"``, ``"array"``, ``"result"``. + """ + if spec is None: + return "none" + + if isinstance(spec, str): + if spec not in _VALID_STRING_KINDS: + raise ValueError( + f"inverse_mass_matrix string must be one of {_VALID_STRING_KINDS}, got {spec!r}" + ) + return spec + + if isinstance(spec, np.ndarray): + return "array" + + if hasattr(spec, "covariance_matrix") or hasattr(spec, "samples"): + return "result" + + raise ValueError( + "Unsupported inverse_mass_matrix specification: expected None, 'diagonal', 'dense', " + f"a numpy array, or a Result/Samples object, got {spec!r}" + ) + + +def resolve_inverse_mass_matrix( + spec: InverseMassMatrixSpec, n_dim: int +) -> Tuple[bool, Optional[np.ndarray]]: + """ + Resolve an ``inverse_mass_matrix`` specification into the ``(is_mass_matrix_diagonal, + initial_inverse_mass_matrix)`` pair consumed by ``blackjax.window_adaptation`` / + ``blackjax.adaptation.staged_adaptation``. + + Parameters + ---------- + spec + See `inverse_mass_matrix_kind_from`. + n_dim + The number of free model parameters. Used to validate array/covariance shapes. + + Returns + ------- + is_mass_matrix_diagonal + Whether ``blackjax.window_adaptation`` should adapt a diagonal (``True``) or dense (``False``) + inverse mass matrix. + seed + ``None`` for a fresh (identity-seeded) adaptation, a 1-D array of shape ``(n_dim,)`` for a diagonal + seed, or a 2-D array of shape ``(n_dim, n_dim)`` for a dense seed -- passed straight through as + ``initial_inverse_mass_matrix``. + + Raises + ------ + ValueError + If a string is not ``"diagonal"``/``"dense"``, an array has the wrong shape, or a `Result`/`Samples` + source's covariance is unusable (MLE-only: too few samples, or non-finite/degenerate) -- in which + case an explicit seed array (e.g. from a Laplace approximation) should be passed instead. + """ + if spec is None: + return True, None + + if isinstance(spec, str): + if spec == "diagonal": + return True, None + if spec == "dense": + return False, None + raise ValueError( + f"inverse_mass_matrix string must be one of {_VALID_STRING_KINDS}, got {spec!r}" + ) + + if isinstance(spec, np.ndarray): + if spec.ndim == 1: + if spec.shape[0] != n_dim: + raise ValueError( + "inverse_mass_matrix: a 1-D (diagonal) seed array must have shape " + f"({n_dim},), got {spec.shape}" + ) + return True, np.asarray(spec, dtype=float) + if spec.ndim == 2: + if spec.shape != (n_dim, n_dim): + raise ValueError( + "inverse_mass_matrix: a 2-D (dense) seed array must have shape " + f"({n_dim}, {n_dim}), got {spec.shape}" + ) + return False, np.asarray(spec, dtype=float) + raise ValueError( + f"inverse_mass_matrix: array must be 1-D or 2-D, got ndim={spec.ndim}" + ) + + # Result / Samples: duck-typed, in preference order, to avoid a circular import on + # `autofit.non_linear.result.Result` / `autofit.non_linear.samples.Samples`. + if hasattr(spec, "covariance_matrix"): + samples = spec + elif hasattr(spec, "samples"): + samples = spec.samples + else: + raise ValueError( + "Unsupported inverse_mass_matrix specification: expected None, 'diagonal', 'dense', " + f"a numpy array, or a Result/Samples object, got {spec!r}" + ) + + if samples is None: + raise ValueError( + "inverse_mass_matrix: the given Result has no samples (`result.samples` is None) -- " + "cannot derive a covariance seed from it." + ) + + n_samples = len(samples) + if n_samples < 2 * n_dim: + raise ValueError( + f"inverse_mass_matrix: the given Result/Samples has only {n_samples} samples for " + f"{n_dim} free parameters (< 2 * n_dim = {2 * n_dim}), so its covariance is unreliable " + "(this typically means the previous fit was an MLE-style optimizer, not a sampler that " + "explores the posterior). Pass an explicit inverse_mass_matrix array instead (e.g. from " + "a Laplace approximation)." + ) + + covariance = np.asarray(samples.covariance_matrix, dtype=float) + + if covariance.shape != (n_dim, n_dim): + raise ValueError( + "inverse_mass_matrix: the given Result/Samples covariance_matrix has shape " + f"{covariance.shape}, expected ({n_dim}, {n_dim})" + ) + + if not np.all(np.isfinite(covariance)): + raise ValueError( + "inverse_mass_matrix: the given Result/Samples covariance_matrix contains non-finite " + "values -- cannot use it as a warm-start seed. Pass an explicit array instead." + ) + + if np.allclose(covariance, np.eye(n_dim)): + raise ValueError( + "inverse_mass_matrix: the given Result/Samples covariance_matrix is the identity matrix, " + "which `Samples.covariance_matrix` falls back to when it has too few samples to compute a " + "real covariance (this typically means the previous fit was an MLE-style optimizer). Pass " + "an explicit inverse_mass_matrix array instead (e.g. from a Laplace approximation)." + ) + + return False, covariance + + +def stack_initial_positions(parameter_lists) -> np.ndarray: + """ + Stack the per-chain starting vectors generated by an initializer (a length-``num_chains`` list of + length-``n_dim`` lists/vectors) into a single ``(num_chains, n_dim)`` array, the shape + ``BlackJAXNUTS._fit`` vmaps warmup and sampling over. + """ + return np.asarray(parameter_lists, dtype=float) + + +def split_chain_diagnostics(positions: np.ndarray) -> np.ndarray: + """ + Reshape ``(n_samples, n_chains, n_dim)`` sampler-order positions (the persisted ``search_internal`` + layout) into ``(n_chains, n_samples, n_dim)`` -- the ``chain_axis=0``, ``sample_axis=1`` convention + required by every function in ``blackjax.diagnostics``. + """ + return np.moveaxis(positions, 0, 1) diff --git a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py index 997507cd7..4b4eda976 100644 --- a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py @@ -16,6 +16,13 @@ from autofit.non_linear.search.mcmc.abstract_mcmc import AbstractMCMC from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelations from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelationsSettings +from autofit.non_linear.search.mcmc.blackjax.chains import ( + InverseMassMatrixSpec, + inverse_mass_matrix_kind_from, + resolve_inverse_mass_matrix, + split_chain_diagnostics, + stack_initial_positions, +) from autofit.non_linear.test_mode import is_test_mode from autofit.non_linear.samples.mcmc import SamplesMCMC from autofit.non_linear.samples.sample import Sample @@ -27,7 +34,12 @@ class BlackJAXNUTS(AbstractMCMC): - __identifier_fields__ = ("num_warmup", "num_samples", "num_chains") + __identifier_fields__ = ( + "num_warmup", + "num_samples", + "num_chains", + "inverse_mass_matrix", + ) def __init__( self, @@ -41,6 +53,9 @@ def __init__( max_num_doublings: int = 10, seed: int = 42, initializer: Optional[Initializer] = None, + inverse_mass_matrix: InverseMassMatrixSpec = None, + mass_matrix_shrinkage: float = 0.0, + share_adaptation: bool = False, auto_correlation_settings: AutoCorrelationsSettings = AutoCorrelationsSettings( check_for_convergence=False ), @@ -65,19 +80,40 @@ def __init__( https://github.com/blackjax-devs/blackjax - The fit runs in two phases: + The fit runs in two phases, both vmapped over ``num_chains`` independent + chains: 1) ``blackjax.window_adaptation`` warmup, which tunes the leapfrog step - size (dual averaging) and a diagonal inverse mass matrix. - 2) NUTS sampling with the tuned kernel, run inside a ``jax.lax.scan`` - so the inner step is fully JIT-compiled. The scan is broken into - ``iterations_per_full_update``-sized chunks so partial state is - persisted and ``perform_update`` runs periodically (samples.csv - flush, plotting), mirroring the Emcee chunking pattern. + size (dual averaging) and an inverse mass matrix (diagonal or dense, + see ``inverse_mass_matrix`` below) per chain. + 2) NUTS sampling with each chain's tuned kernel, run inside a + ``jax.lax.scan`` so the inner step is fully JIT-compiled. The scan is + broken into ``iterations_per_full_update``-sized chunks so partial + state is persisted and ``perform_update`` runs periodically + (samples.csv flush, plotting), mirroring the Emcee chunking pattern. + + **Warm starting.** Both the starting point(s) and the mass matrix can be + seeded from a previous `Result`, without ever modifying that result's + model's priors (contrast with ``Result.model_centred*``, which rewrite + priors to new `GaussianPrior`s). For example, to warm-start a 4-chain + run from an earlier (e.g. single-chain, or MLE) fit's ``result``:: + + search = af.BlackJAXNUTS( + num_chains=4, + initializer=result.start_point_from(n_points=4, jitter=0.05), + inverse_mass_matrix=result, + ) - The single-chain default is the natural fit for the autofit MCMC - plumbing (``SamplesMCMC``, ``AutoCorrelations``). Resume support is - not implemented in v1; the on-disk pickle layout leaves room for it. + ``result.start_point_from(...)`` (equivalently + ``InitializerParamStartPoints.from_result(result, n_points=4, jitter=0.05)``) + maps the previous result's best-fit point onto the (possibly different) + target model by path, then jitters ``n_points`` physical-space starting + vectors around it — one per chain. Passing ``inverse_mass_matrix=result`` + (or ``result.samples``) seeds warmup's dense inverse mass matrix from + that result's ``samples.covariance_matrix``; this raises a clear error + if the result looks MLE-only (too few samples for a reliable + covariance), in which case an explicit array (e.g. from a Laplace + approximation) should be passed instead. Parameters ---------- @@ -92,12 +128,12 @@ def __init__( path prefix and the search name and as the SQLite identifier. num_warmup Number of warmup steps used by ``blackjax.window_adaptation`` - to tune the leapfrog step size and diagonal inverse mass matrix. + to tune the leapfrog step size and inverse mass matrix. num_samples - Number of post-warmup samples drawn from the tuned NUTS kernel. + Number of post-warmup samples drawn from the tuned NUTS kernel, + per chain. num_chains - Number of independent chains. v1 uses 1; values >1 raise an - error until multi-chain support lands. + Number of independent chains sampled in parallel via ``vmap``. target_accept Target Metropolis acceptance rate for window adaptation (default 0.8 — the standard Stan setting). @@ -107,9 +143,40 @@ def __init__( seed Integer seed passed to ``jax.random.PRNGKey``. initializer - Generates the starting point. Defaults to + Generates the starting point(s), one per chain. Defaults to ``InitializerBall(0.49, 0.51)`` from ``AbstractMCMC`` — small - ball around the prior median in unit-cube coordinates. + ball around the prior median in unit-cube coordinates. Pass + ``result.start_point_from(n_points=num_chains, jitter=...)`` (or + an ``InitializerParamStartPoints.from_result(...)``) to warm-start + from a previous result instead. + inverse_mass_matrix + Controls the metric adapted by ``blackjax.window_adaptation``: + + - ``None`` (default): adapt a fresh diagonal inverse mass matrix, + seeded at the identity (standard behaviour). + - ``"diagonal"`` / ``"dense"``: adapt a diagonal / dense inverse + mass matrix, still seeded at the identity. + - a ``numpy`` array: seed the adaptation with this matrix -- a 1-D + array of shape ``(n_dim,)`` seeds a diagonal metric, a 2-D array + of shape ``(n_dim, n_dim)`` seeds a dense metric. + - a `Result` or `Samples` object: seed a dense metric from its + ``samples.covariance_matrix``. Raises ``ValueError`` if that + covariance looks MLE-only (too few samples, or non-finite / + identity) -- pass an explicit array in that case. + mass_matrix_shrinkage + Forwarded to ``blackjax.window_adaptation``'s + ``imm_shrinkage_to_previous``: shrinkage of each warmup window's + adapted inverse mass matrix toward the *previous* window's (in + addition to Stan's existing shrinkage toward the identity). + ``0.0`` (default) reproduces Stan's standard per-window-reset + behaviour; a positive value lets a high-confidence + ``inverse_mass_matrix`` seed persist further into warmup. + share_adaptation + If ``True``, after per-chain warmup the ``num_chains`` tuned + inverse mass matrices are averaged and the tuned step sizes + combined via their median, and every chain's sampling phase uses + this single shared, shared kernel. If ``False`` (default) each + chain samples with its own independently-tuned kernel. auto_correlation_settings Configures the per-parameter ESS-derived integrated auto-correlation diagnostics. ``check_for_convergence`` defaults @@ -119,7 +186,7 @@ def __init__( Sample chunk size between ``perform_update`` calls. Inherited from the autonerves config when ``None``. number_of_cores - Currently unused — single chain runs on a single device. Kept + Currently unused — chains run vmapped on a single device. Kept for API parity with the other MCMC searches. silence If True, the default print output of the non-linear search is @@ -142,12 +209,6 @@ def __init__( **kwargs, ) - if num_chains != 1: - raise NotImplementedError( - "BlackJAXNUTS currently supports num_chains=1 only. " - "Multi-chain support will be added in a future revision." - ) - self.num_warmup = num_warmup self.num_samples = num_samples self.num_chains = num_chains @@ -155,6 +216,17 @@ def __init__( self.max_num_doublings = max_num_doublings self.seed = seed + # The raw specification (used at fit time to resolve an actual seed + # array) is kept private; the public `inverse_mass_matrix` attribute + # is the small descriptive string used by `__identifier_fields__` and + # persisted onto `search_internal` / `samples_info`. Validated eagerly + # here (`inverse_mass_matrix_kind_from` raises on an unrecognised + # spec) so a bad value fails at construction, not mid-fit. + self._inverse_mass_matrix_spec = inverse_mass_matrix + self.inverse_mass_matrix = inverse_mass_matrix_kind_from(inverse_mass_matrix) + self.mass_matrix_shrinkage = mass_matrix_shrinkage + self.share_adaptation = share_adaptation + if is_test_mode(): self.apply_test_mode() @@ -165,10 +237,11 @@ def __init__( def apply_test_mode(self): logger.warning( "TEST MODE 1 (reduced iterations): BlackJAXNUTS will run with " - "num_warmup=20, num_samples=20 for faster completion." + "num_warmup=20, num_samples=20, num_chains<=2 for faster completion." ) self.num_warmup = 20 self.num_samples = 20 + self.num_chains = min(self.num_chains, 2) def _fit(self, model: AbstractPriorModel, analysis): """ @@ -210,11 +283,12 @@ def _fit(self, model: AbstractPriorModel, analysis): live_visual_update=self.live_visual_update, ) - # Initial position: borrow the standard initializer machinery so users - # can substitute their own (InitializerBall by default for MCMC). - # Single chain → ask for one starting point. + # 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=1, + total_points=self.num_chains, model=model, fitness=fitness, paths=self.paths, @@ -227,7 +301,10 @@ def _fit(self, model: AbstractPriorModel, analysis): analysis=analysis, ) - initial_position = jnp.asarray(parameter_lists[0]) + n_dim = model.prior_count + + # (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 @@ -239,39 +316,91 @@ def log_density(params): return fitness.call(params) # One-shot trace + compile so warmup timing is honest. - _ = float(log_density(initial_position)) + _ = float(log_density(initial_positions[0])) rng_key = jax.random.PRNGKey(self.seed) - # ---- Warmup ---------------------------------------------------- + # ---- Inverse mass matrix ----------------------------------------- + is_mass_matrix_diagonal, imm_seed = resolve_inverse_mass_matrix( + self._inverse_mass_matrix_spec, n_dim=n_dim + ) + + # ---- Warmup (vmapped over chains) -------------------------------- self.logger.info( - f"BlackJAXNUTS: window adaptation ({self.num_warmup} steps, " - f"target_accept={self.target_accept})" + f"BlackJAXNUTS: window adaptation ({self.num_warmup} steps x " + f"{self.num_chains} chains, target_accept={self.target_accept}, " + f"inverse_mass_matrix={self.inverse_mass_matrix})" ) warmup = blackjax.window_adaptation( blackjax.nuts, log_density, + is_mass_matrix_diagonal=is_mass_matrix_diagonal, + initial_inverse_mass_matrix=imm_seed, + imm_shrinkage_to_previous=self.mass_matrix_shrinkage, target_acceptance_rate=self.target_accept, max_num_doublings=self.max_num_doublings, ) rng_key, warmup_key = jax.random.split(rng_key) - (last_state, tuned_params), _ = warmup.run( - warmup_key, initial_position, num_steps=self.num_warmup + warmup_keys = jax.random.split(warmup_key, self.num_chains) + + def run_warmup(key, position): + return warmup.run(key, position, num_steps=self.num_warmup) + + (last_state, tuned_params), _ = jax.vmap(run_warmup)( + warmup_keys, initial_positions ) jax.block_until_ready(last_state.position) - # ---- Sampling -------------------------------------------------- + if self.share_adaptation: + shared_step_size = jnp.median(tuned_params["step_size"]) + shared_inverse_mass_matrix = jnp.mean( + tuned_params["inverse_mass_matrix"], axis=0 + ) + tuned_params = { + "step_size": jnp.full_like( + tuned_params["step_size"], shared_step_size + ), + "inverse_mass_matrix": jnp.broadcast_to( + shared_inverse_mass_matrix, + tuned_params["inverse_mass_matrix"].shape, + ), + } + + # ---- Sampling (vmapped over chains) ------------------------------ self.logger.info( - f"BlackJAXNUTS: sampling ({self.num_samples} steps, " - f"chunked {self.iterations_per_full_update} per perform_update)" + f"BlackJAXNUTS: sampling ({self.num_samples} steps x " + f"{self.num_chains} chains, chunked " + f"{self.iterations_per_full_update} per perform_update)" ) - nuts_kernel = blackjax.nuts(log_density, **tuned_params) + # `blackjax.nuts(...)` fixes a single (step_size, inverse_mass_matrix) + # pair; building the kernel directly and vmapping it over the + # per-chain tuned params lets each chain keep (or share, see + # `share_adaptation` above) its own tuned metric. + nuts_kernel = blackjax.nuts.build_kernel() + + def step_fn(key, state, step_size, inverse_mass_matrix): + return nuts_kernel( + key, + state, + log_density, + step_size, + inverse_mass_matrix, + self.max_num_doublings, + ) + + vmapped_step = jax.vmap(step_fn, in_axes=(0, 0, 0, 0)) def one_step(state, rng_key): - new_state, info = nuts_kernel.step(rng_key, state) + step_keys = jax.random.split(rng_key, self.num_chains) + new_state, info = vmapped_step( + step_keys, + state, + tuned_params["step_size"], + tuned_params["inverse_mass_matrix"], + ) return new_state, (new_state, info) def run_chunk(rng_key, initial_state, n_steps): @@ -288,6 +417,7 @@ def run_chunk(rng_key, initial_state, n_steps): "acceptance_rate": [], "num_integration_steps": [], "is_divergent": [], + "num_trajectory_expansions": [], } state = last_state @@ -306,8 +436,12 @@ def run_chunk(rng_key, initial_state, n_steps): # Per-sample log-likelihood for the chunk (NUTS only stores the # log-density inside its kernel state — we recompute the # log-likelihood explicitly so the resulting SamplesMCMC has a - # clean log_likelihood / log_prior split). - chunk_log_l = jax.vmap(_log_likelihood_only(fitness))(states.position) + # clean log_likelihood / log_prior split). ``states.position`` has + # shape (chunk_n, num_chains, n_dim); vmap the per-sample + # log-likelihood over both leading axes. + chunk_log_l = jax.vmap(jax.vmap(_log_likelihood_only(fitness)))( + states.position + ) positions_chunks.append(np.asarray(states.position)) log_likelihood_chunks.append(np.asarray(chunk_log_l)) @@ -316,6 +450,9 @@ def run_chunk(rng_key, initial_state, n_steps): np.asarray(infos.num_integration_steps) ) info_chunks["is_divergent"].append(np.asarray(infos.is_divergent)) + info_chunks["num_trajectory_expansions"].append( + np.asarray(infos.num_trajectory_expansions) + ) # Carry forward the last state position for the next chunk. state = jax.tree_util.tree_map(lambda x: x[-1], states) @@ -333,6 +470,8 @@ def run_chunk(rng_key, initial_state, n_steps): num_samples_completed=total_done, num_samples_total=self.num_samples, num_chains=self.num_chains, + warm_start_source=type(self.initializer).__name__, + inverse_mass_matrix_kind=self.inverse_mass_matrix, ) self.output_search_internal(search_internal=search_internal) @@ -394,23 +533,52 @@ def _test_mode_samples_info(self) -> dict: "num_chains": int(self.num_chains), "ess_min": float("nan"), "ess_per_param": [], + "ess_bulk_per_param": [], + "ess_tail_per_param": [], + "ess_bulk_min": float("nan"), + "ess_tail_min": float("nan"), + "rhat_per_param": [], + "rhat_max": float("nan"), "mean_acceptance": float("nan"), "n_divergent": 0, + "divergent_indices": [], + "tree_depth_histogram": {}, "n_logl_evals": 0, "total_walkers": int(self.num_chains), "total_steps": 0, + "warm_start_source": type(self.initializer).__name__, + "inverse_mass_matrix_kind": self.inverse_mass_matrix, } def samples_info_from(self, search_internal=None): search_internal = search_internal if search_internal is not None else self.backend - positions = search_internal["positions"] + positions = search_internal["positions"] # (n_samples, n_chains, n_dim) info = search_internal["infos"] ess_per_param = _ess_per_param_from(positions) n_logl_evals = int(info["num_integration_steps"].sum()) mean_acceptance = float(info["acceptance_rate"].mean()) - n_divergent = int(info["is_divergent"].sum()) + + is_divergent = info["is_divergent"] + n_divergent = int(is_divergent.sum()) + divergent_indices = [ + [int(sample_index), int(chain_index)] + for sample_index, chain_index in np.argwhere(is_divergent) + ] + + tree_depths = np.asarray(info["num_trajectory_expansions"]) + tree_depth_values, tree_depth_counts = np.unique( + tree_depths, return_counts=True + ) + tree_depth_histogram = { + int(depth): int(count) + for depth, count in zip(tree_depth_values, tree_depth_counts) + } + + ess_bulk_per_param, ess_tail_per_param, rhat_per_param = _chain_diagnostics_from( + positions + ) return { "num_warmup": int(search_internal["num_warmup"]), @@ -418,14 +586,26 @@ def samples_info_from(self, search_internal=None): "num_chains": int(search_internal["num_chains"]), "ess_min": float(ess_per_param.min()), "ess_per_param": ess_per_param.tolist(), + "ess_bulk_per_param": ess_bulk_per_param.tolist(), + "ess_tail_per_param": ess_tail_per_param.tolist(), + "ess_bulk_min": float(ess_bulk_per_param.min()), + "ess_tail_min": float(ess_tail_per_param.min()), + "rhat_per_param": rhat_per_param.tolist(), + "rhat_max": float(rhat_per_param.max()), "mean_acceptance": mean_acceptance, "n_divergent": n_divergent, + "divergent_indices": divergent_indices, + "tree_depth_histogram": tree_depth_histogram, "n_logl_evals": n_logl_evals, "check_size": self.auto_correlation_settings.check_size, "required_length": self.auto_correlation_settings.required_length, "change_threshold": self.auto_correlation_settings.change_threshold, "total_walkers": int(search_internal["num_chains"]), "total_steps": int(search_internal["num_samples_completed"]), + "warm_start_source": search_internal.get("warm_start_source", "none"), + "inverse_mass_matrix_kind": search_internal.get( + "inverse_mass_matrix_kind", "none" + ), "time": self.timer.time if self.timer else None, } @@ -437,11 +617,21 @@ def samples_via_internal_from(self, model, search_internal=None): """ search_internal = search_internal if search_internal is not None else self.backend - positions = search_internal["positions"] # (num_samples, n_dim) - log_likelihood_array = search_internal["log_likelihood_history"] + positions = search_internal["positions"] # (n_samples, n_chains, n_dim) + log_likelihood_array = search_internal["log_likelihood_history"] # (n_samples, n_chains) - parameter_lists = positions.tolist() - log_likelihood_list = [float(x) for x in log_likelihood_array] + # Chain-major flatten: all of chain 0's samples (in draw order), then + # chain 1's, etc. -- matches `total_walkers = num_chains` below. + n_samples, n_chains, n_dim = positions.shape + positions_chain_major = np.moveaxis(positions, 0, 1).reshape( + n_chains * n_samples, n_dim + ) + log_likelihood_chain_major = np.moveaxis(log_likelihood_array, 0, 1).reshape( + n_chains * n_samples + ) + + parameter_lists = positions_chain_major.tolist() + log_likelihood_list = [float(x) for x in log_likelihood_chain_major] log_prior_list = model.log_prior_list_from(parameter_lists=parameter_lists) weight_list = [1.0] * len(parameter_lists) @@ -471,7 +661,7 @@ def auto_correlations_from(self, search_internal=None): relative-change convergence metric stays meaningful. """ search_internal = search_internal if search_internal is not None else self.backend - positions = search_internal["positions"] + positions = search_internal["positions"] # (n_samples, n_chains, n_dim) n_samples = positions.shape[0] check_size = self.auto_correlation_settings.check_size @@ -527,15 +717,44 @@ def log_l(params): def _ess_per_param_from(positions: np.ndarray) -> np.ndarray: """ Per-parameter effective sample size via BlackJAX's Geyer-style - monotone variance estimator. Single-chain shape contract: positions is - ``(num_samples, n_dim)``; we add a length-1 chain axis so the - diagnostic sees ``(n_chains=1, num_samples, n_dim)``. + monotone variance estimator. + + ``positions`` has shape ``(num_samples, num_chains, n_dim)`` (the + ``search_internal`` layout, for any ``num_chains >= 1``); this is + reshaped to BlackJAX's ``(chain_axis=0, sample_axis=1)`` convention via + ``split_chain_diagnostics`` before the diagnostic is applied, so a + single-chain run (``num_chains == 1``) sees the same + ``(1, num_samples, n_dim)`` shape the original single-chain + implementation used. """ import jax.numpy as jnp from blackjax.diagnostics import effective_sample_size - ess = effective_sample_size(jnp.asarray(positions)[None, ...]) - return np.asarray(ess) + chain_major = split_chain_diagnostics(positions) + ess = effective_sample_size(jnp.asarray(chain_major)) + return np.atleast_1d(np.asarray(ess)) + + +def _chain_diagnostics_from(positions: np.ndarray): + """ + Bulk/tail ESS and rank-normalised split-R-hat per parameter, via + ``blackjax.diagnostics``. See ``_ess_per_param_from`` for the + ``(num_samples, num_chains, n_dim)`` -> ``(chain_axis, sample_axis)`` + reshape. R-hat with a single chain is split-chain only (BlackJAX splits + each chain in half internally before comparing) -- it is still a useful + non-stationarity check, just not a genuine multi-chain convergence + diagnostic in that case. + """ + import jax.numpy as jnp + from blackjax.diagnostics import ess_bulk, ess_tail, rhat + + chain_major = jnp.asarray(split_chain_diagnostics(positions)) + + ess_bulk_per_param = np.atleast_1d(np.asarray(ess_bulk(chain_major))) + ess_tail_per_param = np.atleast_1d(np.asarray(ess_tail(chain_major))) + rhat_per_param = np.atleast_1d(np.asarray(rhat(chain_major))) + + return ess_bulk_per_param, ess_tail_per_param, rhat_per_param def _times_from_positions(positions: np.ndarray) -> np.ndarray: @@ -545,8 +764,13 @@ def _times_from_positions(positions: np.ndarray) -> np.ndarray: e.g. mostly-divergent chain), we floor ``ESS = 1`` so ``times`` does not go past ``num_samples`` and ``AutoCorrelations.check_if_converged`` stays well-defined. + + ``positions`` has shape ``(num_samples, num_chains, n_dim)``; ``N`` in + the identity above is the *total* sample count across all chains + (``num_samples * num_chains``), matching what ``effective_sample_size`` + itself pools over. """ - n_samples = positions.shape[0] + n_samples = positions.shape[0] * positions.shape[1] ess = _ess_per_param_from(positions) ess = np.clip(ess, a_min=1.0, a_max=None) return n_samples / ess @@ -562,6 +786,8 @@ def _build_search_internal( num_samples_completed, num_samples_total, num_chains, + warm_start_source="none", + inverse_mass_matrix_kind="none", ): """ Glue chunked sampling output into the persistence dict pickled under @@ -576,6 +802,9 @@ def _build_search_internal( info_chunks["num_integration_steps"] ), "is_divergent": np.concatenate(info_chunks["is_divergent"]), + "num_trajectory_expansions": np.concatenate( + info_chunks["num_trajectory_expansions"] + ), }, # tuned_params often contains JAX arrays; convert to numpy so the # pickle is portable across JAX versions. @@ -588,4 +817,6 @@ def _build_search_internal( "num_samples_completed": num_samples_completed, "num_samples": num_samples_total, "num_chains": num_chains, + "warm_start_source": warm_start_source, + "inverse_mass_matrix_kind": inverse_mass_matrix_kind, } diff --git a/test_autofit/non_linear/result/test_result.py b/test_autofit/non_linear/result/test_result.py index dc89cfefa..0a25b7584 100644 --- a/test_autofit/non_linear/result/test_result.py +++ b/test_autofit/non_linear/result/test_result.py @@ -76,6 +76,23 @@ def test_raises(self, result): result.samples.prior_means, a=2.0, r=1.0 ) + def test_start_point(self, result): + start_point = result.start_point + + assert isinstance(start_point, af.InitializerParamStartPoints) + + expected = result.samples.max_log_likelihood(as_instance=False) + parameter_dict = { + ".".join(result.model.path_for_prior(prior)): value + for prior, value in start_point.parameter_dict.items() + } + for path, value in zip( + [".".join(path) for path in result.model.paths], expected + ): + assert parameter_dict[path] == pytest.approx( + (value - 1e-8, value + 1e-8) + ) + @pytest.fixture(name="results") def make_results_collection(): diff --git a/test_autofit/non_linear/search/mcmc/test_blackjax_chains.py b/test_autofit/non_linear/search/mcmc/test_blackjax_chains.py new file mode 100644 index 000000000..473448448 --- /dev/null +++ b/test_autofit/non_linear/search/mcmc/test_blackjax_chains.py @@ -0,0 +1,170 @@ +import numpy as np +import pytest + +# `autofit.non_linear.search.mcmc.blackjax.chains` is pure numpy -- no JAX or +# blackjax import at module scope (following the repo convention of no JAX in +# unit tests) -- so these tests run unconditionally, unlike +# `test_blackjax_nuts.py`'s `@requires_blackjax`-guarded diagnostics tests. +from autofit.non_linear.search.mcmc.blackjax.chains import ( + inverse_mass_matrix_kind_from, + resolve_inverse_mass_matrix, + split_chain_diagnostics, + stack_initial_positions, +) + + +class _FakeSamples: + def __init__(self, covariance_matrix, n_samples): + self.covariance_matrix = covariance_matrix + self._n_samples = n_samples + + def __len__(self): + return self._n_samples + + +class _FakeResult: + def __init__(self, samples): + self.samples = samples + + +class TestInverseMassMatrixKindFrom: + def test__none(self): + assert inverse_mass_matrix_kind_from(None) == "none" + + def test__diagonal_and_dense_strings(self): + assert inverse_mass_matrix_kind_from("diagonal") == "diagonal" + assert inverse_mass_matrix_kind_from("dense") == "dense" + + def test__bad_string_raises(self): + with pytest.raises(ValueError): + inverse_mass_matrix_kind_from("not_a_kind") + + def test__array_1d_and_2d(self): + assert inverse_mass_matrix_kind_from(np.ones(3)) == "array" + assert inverse_mass_matrix_kind_from(np.eye(3)) == "array" + + def test__result_like(self): + samples = _FakeSamples(covariance_matrix=np.eye(3), n_samples=20) + assert inverse_mass_matrix_kind_from(samples) == "result" + assert inverse_mass_matrix_kind_from(_FakeResult(samples)) == "result" + + def test__unsupported_type_raises(self): + with pytest.raises(ValueError): + inverse_mass_matrix_kind_from(object()) + + +class TestResolveInverseMassMatrix: + def test__none_gives_diagonal_no_seed(self): + is_diagonal, seed = resolve_inverse_mass_matrix(None, n_dim=3) + assert is_diagonal is True + assert seed is None + + def test__diagonal_string(self): + is_diagonal, seed = resolve_inverse_mass_matrix("diagonal", n_dim=3) + assert is_diagonal is True + assert seed is None + + def test__dense_string(self): + is_diagonal, seed = resolve_inverse_mass_matrix("dense", n_dim=3) + assert is_diagonal is False + assert seed is None + + def test__bad_string_raises(self): + with pytest.raises(ValueError): + resolve_inverse_mass_matrix("not_a_kind", n_dim=3) + + def test__1d_array_seeds_diagonal(self): + diag = np.array([1.0, 2.0, 3.0]) + is_diagonal, seed = resolve_inverse_mass_matrix(diag, n_dim=3) + assert is_diagonal is True + assert seed is diag or np.array_equal(seed, diag) + + def test__1d_array_wrong_shape_raises(self): + with pytest.raises(ValueError): + resolve_inverse_mass_matrix(np.ones(2), n_dim=3) + + def test__2d_array_seeds_dense(self): + dense = np.eye(3) * 2.0 + is_diagonal, seed = resolve_inverse_mass_matrix(dense, n_dim=3) + assert is_diagonal is False + assert np.array_equal(seed, dense) + + def test__2d_array_wrong_shape_raises(self): + with pytest.raises(ValueError): + resolve_inverse_mass_matrix(np.eye(2), n_dim=3) + + def test__3d_array_raises(self): + with pytest.raises(ValueError): + resolve_inverse_mass_matrix(np.ones((2, 2, 2)), n_dim=3) + + def test__result_covariance_used_as_dense_seed(self): + cov = np.array([[2.0, 0.1, 0.0], [0.1, 3.0, 0.2], [0.0, 0.2, 1.5]]) + samples = _FakeSamples(covariance_matrix=cov, n_samples=20) + + is_diagonal, seed = resolve_inverse_mass_matrix(_FakeResult(samples), n_dim=3) + + assert is_diagonal is False + assert np.array_equal(seed, cov) + + # Also works passed directly as a `Samples`-like object (not wrapped + # in a `Result`-like object). + is_diagonal, seed = resolve_inverse_mass_matrix(samples, n_dim=3) + assert is_diagonal is False + assert np.array_equal(seed, cov) + + def test__mle_only_too_few_samples_raises(self): + # < 2 * n_dim samples -- an MLE-style optimizer's "samples", not a + # posterior-exploring sampler's. + samples = _FakeSamples(covariance_matrix=np.eye(3), n_samples=5) + + with pytest.raises(ValueError, match="MLE-style"): + resolve_inverse_mass_matrix(_FakeResult(samples), n_dim=3) + + def test__degenerate_identity_covariance_raises(self): + # `Samples.covariance_matrix` itself falls back to `np.eye(1)` (or + # similar) when it can't compute a real covariance -- treat an + # identity covariance from enough samples as still-unusable. + samples = _FakeSamples(covariance_matrix=np.eye(3), n_samples=20) + + with pytest.raises(ValueError, match="identity"): + resolve_inverse_mass_matrix(_FakeResult(samples), n_dim=3) + + def test__non_finite_covariance_raises(self): + cov = np.array([[1.0, np.nan, 0.0], [np.nan, 1.0, 0.0], [0.0, 0.0, 1.0]]) + samples = _FakeSamples(covariance_matrix=cov, n_samples=20) + + with pytest.raises(ValueError, match="non-finite"): + resolve_inverse_mass_matrix(_FakeResult(samples), n_dim=3) + + def test__result_with_no_samples_raises(self): + with pytest.raises(ValueError): + resolve_inverse_mass_matrix(_FakeResult(samples=None), n_dim=3) + + +class TestStackInitialPositions: + def test__list_of_lists_stacked(self): + stacked = stack_initial_positions([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + + assert stacked.shape == (3, 2) + assert np.array_equal(stacked, np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])) + + +class TestSplitChainDiagnostics: + def test__moves_sample_axis_behind_chain_axis(self): + n_samples, n_chains, n_dim = 200, 2, 3 + rng = np.random.default_rng(0) + positions = rng.normal(size=(n_samples, n_chains, n_dim)) + + chain_major = split_chain_diagnostics(positions) + + assert chain_major.shape == (n_chains, n_samples, n_dim) + assert np.array_equal(chain_major[0], positions[:, 0, :]) + assert np.array_equal(chain_major[1], positions[:, 1, :]) + + def test__single_chain(self): + n_samples, n_dim = 50, 4 + positions = np.zeros((n_samples, 1, n_dim)) + + chain_major = split_chain_diagnostics(positions) + + assert chain_major.shape == (1, n_samples, n_dim) diff --git a/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py b/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py index fbc55bd28..a6f4d730e 100644 --- a/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py +++ b/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py @@ -5,6 +5,7 @@ import autofit as af from autofit.non_linear.search.mcmc.blackjax.nuts.search import ( + _build_search_internal, _times_from_positions, ) @@ -19,6 +20,15 @@ pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") +class _FakeSamples: + """A minimal `Samples`-like object -- just enough for + `inverse_mass_matrix_kind_from` (construction-time classification only; + no resolution) to recognise it as a "result" kind.""" + + def __init__(self, covariance_matrix): + self.covariance_matrix = covariance_matrix + + def test__explicit_params(): search = af.BlackJAXNUTS( num_warmup=321, @@ -62,11 +72,50 @@ def test__defaults(): # Convergence checking is intentionally OFF by default for NUTS — the # sampler runs to a fixed budget after warmup tunes the kernel. assert search.auto_correlation_settings.check_for_convergence is False + assert search.inverse_mass_matrix == "none" + assert search.mass_matrix_shrinkage == 0.0 + assert search.share_adaptation is False + + +def test__multi_chain_construction(): + search = af.BlackJAXNUTS(num_chains=4) + + assert search.num_chains == 4 + + +class TestInverseMassMatrix: + def test__none(self): + search = af.BlackJAXNUTS(inverse_mass_matrix=None) + assert search.inverse_mass_matrix == "none" + + def test__diagonal_string(self): + search = af.BlackJAXNUTS(inverse_mass_matrix="diagonal") + assert search.inverse_mass_matrix == "diagonal" + + def test__dense_string(self): + search = af.BlackJAXNUTS(inverse_mass_matrix="dense") + assert search.inverse_mass_matrix == "dense" + + def test__1d_array(self): + search = af.BlackJAXNUTS(inverse_mass_matrix=np.array([1.0, 2.0, 3.0])) + assert search.inverse_mass_matrix == "array" + + def test__2d_array(self): + search = af.BlackJAXNUTS(inverse_mass_matrix=np.eye(3)) + assert search.inverse_mass_matrix == "array" + + def test__result_like(self): + fake_samples = _FakeSamples(covariance_matrix=np.eye(3)) + search = af.BlackJAXNUTS(inverse_mass_matrix=fake_samples) + assert search.inverse_mass_matrix == "result" + def test__bad_string_raises(self): + with pytest.raises(ValueError): + af.BlackJAXNUTS(inverse_mass_matrix="not_a_valid_kind") -def test__multi_chain_not_implemented(): - with pytest.raises(NotImplementedError): - af.BlackJAXNUTS(num_chains=2) + def test__bad_type_raises(self): + with pytest.raises(ValueError): + af.BlackJAXNUTS(inverse_mass_matrix=object()) def test__test_mode_reduces_iterations(monkeypatch): @@ -75,40 +124,63 @@ def test__test_mode_reduces_iterations(monkeypatch): # the env var (not after-the-fact mutate the search instance). monkeypatch.setenv("PYAUTO_TEST_MODE", "1") - search = af.BlackJAXNUTS(num_warmup=10000, num_samples=10000) + search = af.BlackJAXNUTS(num_warmup=10000, num_samples=10000, num_chains=8) assert search.num_warmup == 20 assert search.num_samples == 20 + # `apply_test_mode` clamps `num_chains` to at most 2 so test-mode fits + # stay fast even when a warm-started multi-chain config is requested. + assert search.num_chains == 2 + + +def test__test_mode_leaves_single_chain_alone(monkeypatch): + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + search = af.BlackJAXNUTS(num_chains=1) + + assert search.num_chains == 1 def test__identifier_fields_distinguish_run_shape(): a = af.BlackJAXNUTS(num_warmup=100, num_samples=500, num_chains=1) b = af.BlackJAXNUTS(num_warmup=200, num_samples=500, num_chains=1) c = af.BlackJAXNUTS(num_warmup=100, num_samples=999, num_chains=1) + d = af.BlackJAXNUTS( + num_warmup=100, num_samples=500, num_chains=1, inverse_mass_matrix="dense" + ) - # __identifier_fields__ is a class attribute; make sure it covers all - # three knobs. The autofit identifier hash uses these to keep distinct - # runs from colliding on the same output path. + # __identifier_fields__ is a class attribute; make sure it covers every + # knob that changes the run's shape. The autofit identifier hash uses + # these to keep distinct runs from colliding on the same output path. assert af.BlackJAXNUTS.__identifier_fields__ == ( "num_warmup", "num_samples", "num_chains", + "inverse_mass_matrix", ) # Instance attributes should round-trip exactly. assert (a.num_warmup, a.num_samples, a.num_chains) == (100, 500, 1) assert (b.num_warmup, b.num_samples, b.num_chains) == (200, 500, 1) assert (c.num_warmup, c.num_samples, c.num_chains) == (100, 999, 1) + assert a.inverse_mass_matrix == "none" + assert d.inverse_mass_matrix == "dense" + + from autofit.mapper.identifier import Identifier + + assert str(Identifier(a)) != str(Identifier(d)) @requires_blackjax def test__times_from_positions_clamps_low_ess(): # Construct a degenerate "chain" — every sample is identical → ESS # collapses; the clamp must keep ``times`` finite and equal to N. + # Positions carry the `(n_samples, n_chains, n_dim)` search_internal + # layout, so a single chain still gets an explicit size-1 chain axis. rng = np.random.default_rng(0) n_dim = 3 n_samples = 200 - positions = np.tile(rng.normal(size=(1, n_dim)), (n_samples, 1)) + positions = np.tile(rng.normal(size=(1, 1, n_dim)), (n_samples, 1, 1)) times = _times_from_positions(positions) @@ -124,10 +196,131 @@ def test__times_from_positions_independent_chain(): rng = np.random.default_rng(1) n_samples = 2000 n_dim = 3 - positions = rng.normal(size=(n_samples, n_dim)) + positions = rng.normal(size=(n_samples, 1, n_dim)) times = _times_from_positions(positions) # Expect each per-param ``τ`` to be O(1) — well under N. We use a loose # bound (10) so the test isn't flaky on small-sample variance. assert np.all(times < 10.0) + + +@requires_blackjax +def test__times_from_positions_multi_chain_pools_samples(): + # `N` in the tau = N / ESS identity should be the *total* sample count + # across all chains, not just one chain's. + rng = np.random.default_rng(2) + n_samples = 500 + n_chains = 4 + n_dim = 2 + positions = rng.normal(size=(n_samples, n_chains, n_dim)) + + times = _times_from_positions(positions) + + assert times.shape == (n_dim,) + assert np.all(np.isfinite(times)) + + +def test__build_search_internal_shape_contract(): + # Pure numpy: exercises the persistence-dict glue directly, without + # touching blackjax/jax. + n_samples, n_chains, n_dim = 5, 3, 2 + + rng = np.random.default_rng(0) + positions_chunks = [rng.normal(size=(n_samples, n_chains, n_dim))] + log_likelihood_chunks = [rng.normal(size=(n_samples, n_chains))] + info_chunks = { + "acceptance_rate": [np.full((n_samples, n_chains), 0.85)], + "num_integration_steps": [np.full((n_samples, n_chains), 7, dtype=int)], + "is_divergent": [np.zeros((n_samples, n_chains), dtype=bool)], + "num_trajectory_expansions": [np.full((n_samples, n_chains), 3, dtype=int)], + } + + search_internal = _build_search_internal( + positions_chunks=positions_chunks, + log_likelihood_chunks=log_likelihood_chunks, + info_chunks=info_chunks, + tuned_params={"step_size": np.ones(n_chains), "inverse_mass_matrix": np.ones((n_chains, n_dim))}, + last_state_position=positions_chunks[0][-1], + num_warmup=50, + num_samples_completed=n_samples, + num_samples_total=n_samples, + num_chains=n_chains, + warm_start_source="InitializerParamStartPoints", + inverse_mass_matrix_kind="dense", + ) + + assert search_internal["positions"].shape == (n_samples, n_chains, n_dim) + assert search_internal["log_likelihood_history"].shape == (n_samples, n_chains) + assert search_internal["infos"]["is_divergent"].shape == (n_samples, n_chains) + assert search_internal["infos"]["num_trajectory_expansions"].shape == ( + n_samples, + n_chains, + ) + assert search_internal["num_chains"] == n_chains + assert search_internal["warm_start_source"] == "InitializerParamStartPoints" + assert search_internal["inverse_mass_matrix_kind"] == "dense" + + +@requires_blackjax +def test__samples_via_internal_from_shape_contract(): + n_samples, n_chains = 6, 3 + + model = af.Model(af.m.MockClassx2) + model.one = af.UniformPrior(lower_limit=-10.0, upper_limit=10.0) + model.two = af.UniformPrior(lower_limit=-10.0, upper_limit=10.0) + n_dim = model.prior_count + + rng = np.random.default_rng(3) + positions = rng.uniform(-1.0, 1.0, size=(n_samples, n_chains, n_dim)) + log_likelihood_history = rng.normal(size=(n_samples, n_chains)) + + search_internal = { + "positions": positions, + "log_likelihood_history": log_likelihood_history, + "infos": { + "acceptance_rate": np.full((n_samples, n_chains), 0.8), + "num_integration_steps": np.full((n_samples, n_chains), 5, dtype=int), + "is_divergent": np.zeros((n_samples, n_chains), dtype=bool), + "num_trajectory_expansions": np.full((n_samples, n_chains), 2, dtype=int), + }, + "tuned_params": { + "step_size": np.ones(n_chains), + "inverse_mass_matrix": np.ones((n_chains, n_dim)), + }, + "last_state_position": positions[-1], + "num_warmup": 10, + "num_samples_completed": n_samples, + "num_samples": n_samples, + "num_chains": n_chains, + "warm_start_source": "InitializerBall", + "inverse_mass_matrix_kind": "none", + } + + search = af.BlackJAXNUTS(num_chains=n_chains) + samples = search.samples_via_internal_from(model=model, search_internal=search_internal) + + assert len(samples.sample_list) == n_samples * n_chains + assert samples.total_walkers == n_chains + assert samples.samples_info["num_chains"] == n_chains + assert samples.samples_info["warm_start_source"] == "InitializerBall" + assert samples.samples_info["inverse_mass_matrix_kind"] == "none" + assert len(samples.samples_info["ess_bulk_per_param"]) == n_dim + assert len(samples.samples_info["ess_tail_per_param"]) == n_dim + assert len(samples.samples_info["rhat_per_param"]) == n_dim + assert samples.samples_info["n_divergent"] == 0 + assert samples.samples_info["divergent_indices"] == [] + assert samples.samples_info["tree_depth_histogram"] == {2: n_samples * n_chains} + + # `SamplesMCMC` read-only convenience properties (autofit/non_linear/samples/mcmc.py). + assert samples.ess_bulk == samples.samples_info["ess_bulk_per_param"] + assert samples.ess_tail == samples.samples_info["ess_tail_per_param"] + assert samples.rhat == samples.samples_info["rhat_per_param"] + assert samples.n_divergent == 0 + assert samples.tree_depths == {2: n_samples * n_chains} + + # chain-major flatten: the first n_samples rows are chain 0. + chain_0_first_row = samples.sample_list[0].parameter_lists_for_paths( + samples.paths + ) + assert chain_0_first_row == pytest.approx(positions[0, 0].tolist()) diff --git a/test_autofit/non_linear/test_initializer_from_result.py b/test_autofit/non_linear/test_initializer_from_result.py new file mode 100644 index 000000000..735447757 --- /dev/null +++ b/test_autofit/non_linear/test_initializer_from_result.py @@ -0,0 +1,186 @@ +import pytest + +import autofit as af +from autofit import Sample +from autofit.non_linear.mock.mock_result import MockResult +from autofit.non_linear.mock.mock_samples import MockSamples +from autofit.non_linear.mock.mock_samples_summary import MockSamplesSummary + +# Pure numpy: `InitializerParamStartPoints.from_result` never imports JAX/blackjax, +# so these tests run unconditionally (no `skipif`/`importorskip` needed). + + +@pytest.fixture(name="model") +def make_model(): + model = af.Model(af.m.MockClassx4) + model.one = af.UniformPrior(lower_limit=0.0, upper_limit=1.0) + model.two = af.UniformPrior(lower_limit=0.0, upper_limit=1.0) + model.three = af.UniformPrior(lower_limit=0.0, upper_limit=1.0) + model.four = af.UniformPrior(lower_limit=0.0, upper_limit=1.0) + return model + + +@pytest.fixture(name="result") +def make_result(model): + sample_list = [ + Sample( + log_likelihood=1.0, + log_prior=0.0, + weight=0.2, + kwargs={"one": 0.1, "two": 0.2, "three": 0.3, "four": 0.4}, + ), + Sample( + log_likelihood=2.0, + log_prior=0.0, + weight=0.5, + kwargs={"one": 0.11, "two": 0.22, "three": 0.33, "four": 0.44}, + ), + Sample( + log_likelihood=1.5, + log_prior=0.0, + weight=0.3, + kwargs={"one": 0.12, "two": 0.19, "three": 0.28, "four": 0.41}, + ), + ] + samples = MockSamples(sample_list=sample_list, model=model) + return MockResult( + samples=samples, + samples_summary=MockSamplesSummary( + model=model, + max_log_likelihood_instance=[0.11, 0.22, 0.33, 0.44], + median_pdf_sample=sample_list[1], + ), + model=model, + ) + + +def parameter_dict_by_path(model, initializer): + return { + ".".join(model.path_for_prior(prior)): value + for prior, value in initializer.parameter_dict.items() + } + + +class TestFromResult: + def test__default_point_is_max_log_likelihood(self, result, model): + vector = result.samples.max_log_likelihood(as_instance=False) + assert vector == [0.11, 0.22, 0.33, 0.44] + + initializer = af.InitializerParamStartPoints.from_result(result) + + parameter_dict = parameter_dict_by_path(model, initializer) + assert parameter_dict["one"] == pytest.approx((0.11 - 1e-8, 0.11 + 1e-8)) + assert parameter_dict["two"] == pytest.approx((0.22 - 1e-8, 0.22 + 1e-8)) + assert parameter_dict["three"] == pytest.approx((0.33 - 1e-8, 0.33 + 1e-8)) + assert parameter_dict["four"] == pytest.approx((0.44 - 1e-8, 0.44 + 1e-8)) + + def test__median_pdf_point(self, result, model): + expected = result.samples.median_pdf(as_instance=False) + + initializer = af.InitializerParamStartPoints.from_result( + result, point="median_pdf" + ) + + parameter_dict = parameter_dict_by_path(model, initializer) + for path, value in zip(["one", "two", "three", "four"], expected): + assert parameter_dict[path] == pytest.approx((value - 1e-8, value + 1e-8)) + + def test__unrecognised_point_raises(self, result): + with pytest.raises(af.exc.InitializerException): + af.InitializerParamStartPoints.from_result(result, point="not_a_point") + + def test__priors_are_never_mutated(self, result, model): + before = [(p.id, p.lower_limit, p.upper_limit) for p in model.priors_ordered_by_id] + + af.InitializerParamStartPoints.from_result(result) + af.InitializerParamStartPoints.from_result(result, point="median_pdf") + af.InitializerParamStartPoints.from_result(result, n_points=3, jitter=0.05, seed=0) + + after = [(p.id, p.lower_limit, p.upper_limit) for p in model.priors_ordered_by_id] + assert before == after + + def test__n_points_shapes_and_distinct_points(self, result, model): + initializer = af.InitializerParamStartPoints.from_result( + result, n_points=5, jitter=0.05, seed=1 + ) + + assert len(initializer._point_dicts) == 5 + for point_dict in initializer._point_dicts: + assert len(point_dict) == model.prior_count + + values_for_one = { + point_dict[prior] + for point_dict in initializer._point_dicts + for prior in point_dict + if model.path_for_prior(prior) == ("one",) + } + assert len(values_for_one) == 5 # every point jittered to a distinct value + + def test__n_points_generation_reproducible_with_seed(self, result): + initializer_a = af.InitializerParamStartPoints.from_result( + result, n_points=4, jitter=0.05, seed=7 + ) + initializer_b = af.InitializerParamStartPoints.from_result( + result, n_points=4, jitter=0.05, seed=7 + ) + + for point_a, point_b in zip(initializer_a._point_dicts, initializer_b._point_dicts): + for prior in point_a: + assert point_a[prior] == point_b[prior] + + def test__n_points_used_by_samples_from_model(self, result, model): + initializer = af.InitializerParamStartPoints.from_result( + result, n_points=4, jitter=0.05, seed=1 + ) + + # Non-constant fitness so the identical-figure-of-merit guard in + # `samples_from_model` does not trip (mirrors the pattern used + # elsewhere in test_initializer.py for `MockFitness`). + from random import random + + def fitness(parameters): + return random() + + _, parameter_lists, _ = initializer.samples_from_model( + total_points=4, + model=model, + fitness=fitness, + paths=af.DirectoryPaths(), + ) + + assert len(parameter_lists) == 4 + # Each of the 4 generated points should be distinct (one per chain). + assert len({tuple(p) for p in parameter_lists}) == 4 + + def test__n_points_one_matches_legacy_behaviour(self, result, model): + # n_points=1, jitter=0.0 (defaults) must reproduce the exact classic + # `InitializerParamStartPoints` single-point behaviour: no `_point_dicts`. + initializer = af.InitializerParamStartPoints.from_result(result) + assert initializer._point_dicts is None + + def test__dimension_mismatch_raises(self): + model_a = af.Model(af.m.MockClassx2) + + sample_list = [ + Sample( + log_likelihood=1.0, + log_prior=0.0, + weight=0.0, + kwargs={"one": 0.1, "two": 0.2}, + ), + ] + samples = MockSamples(sample_list=sample_list, model=model_a) + result = MockResult( + samples=samples, + samples_summary=MockSamplesSummary( + model=model_a, + max_log_likelihood_instance=[0.1, 0.2], + median_pdf_sample=sample_list[0], + ), + model=model_a, + ) + + model_b = af.Model(af.m.MockClassx4) + + with pytest.raises(af.exc.InitializerException): + af.InitializerParamStartPoints.from_result(result, model=model_b)