"""Snooker (DE-MCzs) move primitives for the GUIBRUSHR sampler.
This module is a dependency-free (numpy-only) reference implementation of the
snooker update of ter Braak and Vrugt (2008), "Differential Evolution Markov
Chain with snooker updater and fewer chains", Stat. Comput. 18(4), 435-446.
It is deliberately isolated from the rest of the retrieval pipeline so it can be
unit-tested in isolation and reused verbatim by the production worker
(`ModelData.parallel_chain`). It implements only the snooker-specific maths:
- `snooker_propose` : the line-constrained proposal x_prop = x_i + gamma_s * (zp1 - zp2)
- `snooker_log_jacobian` : the radial Jacobian (d-1)*(log||x_prop-z|| - log||x_i-z||)
The 90% "parallel" DE-MCz move (x_i + gamma*(zR1 - zR2) + e) reuses the existing
GUIBRUSHR machinery and is not duplicated here; only its archive draw differs.
Conventions (validated at Gate G1 against the project guide snooker_demc.pdf
section 3 Eq. 3, ter Braak and Vrugt 2008, mc3 chain.py and the first author's
demc_zs.r reference):
- The archive ``Z`` is a 2D array of shape ``(M, d)``: M past+present states,
each a parameter vector of length ``d``. Draws are uniform over the M rows.
- ``gamma_s`` is dimension-independent and drawn from ``Uniform(1.2, 2.2)`` per
step (mc3 parity). It is NOT divided by sqrt(d) (that is the parallel-move
convention only).
- There is no additive noise term ``e`` on the snooker move.
- The Jacobian is always combined in log-space as ``(d-1) * dlog||.||``; the raw
power ``||.||**(d-1)`` is never formed (it overflows/underflows for large d).
"""
from __future__ import annotations
import numpy as np
# Probability of choosing a snooker move per step in DE-MCzs (the remaining
# 1 - PSNOOKER_DEFAULT are parallel DE moves). mc3 / ter Braak and Vrugt default.
PSNOOKER_DEFAULT: float = 0.1
# gamma_s ~ Uniform(GAMMA_S_LOW, GAMMA_S_HIGH); mean 1.7 = 2.38/sqrt(2).
GAMMA_S_LOW: float = 1.2
GAMMA_S_HIGH: float = 2.2
# Floor used to guard a degenerate line (anchor coincident with the current
# state) and a proposal landing exactly on the anchor. Matches the demc_zs.r
# practice of flooring the squared norm rather than crashing.
_NORM_EPS: float = 1e-12
[docs]
def draw_gamma_s(rng: np.random.Generator) -> float:
"""Draw the snooker step scale gamma_s ~ Uniform(1.2, 2.2)."""
return float(rng.uniform(GAMMA_S_LOW, GAMMA_S_HIGH))
[docs]
def sample_distinct_indices(
rng: np.random.Generator, n: int, k: int
) -> tuple[int, ...]:
"""Draw ``k`` distinct integer indices uniformly from ``[0, n)``.
Args:
rng: numpy random Generator.
n: archive size (number of rows to choose from).
k: number of distinct indices required.
Returns:
A tuple of ``k`` distinct ints.
Raises:
ValueError: if ``n < k`` (not enough states to draw distinct indices).
"""
if n < k:
raise ValueError(
f"cannot draw {k} distinct indices from an archive of size {n}"
)
return tuple(int(i) for i in rng.choice(n, size=k, replace=False))
[docs]
def project_onto_line(
point: np.ndarray, anchor: np.ndarray, unit_dir: np.ndarray
) -> np.ndarray:
"""Orthogonal projection of ``point`` onto the line through ``anchor``.
The line has unit direction ``unit_dir``. Returns the nearest point on the
line, i.e. ``anchor + <point - anchor, unit_dir> * unit_dir``.
"""
return anchor + np.dot(point - anchor, unit_dir) * unit_dir
[docs]
def snooker_propose(
x_i: np.ndarray,
archive: np.ndarray,
rng: np.random.Generator,
gamma_s: float | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Generate a snooker proposal for the current state ``x_i``.
Picks an anchor ``z`` and two further states ``z1, z2`` uniformly at random
(distinct rows) from the archive, projects the difference ``z1 - z2`` onto
the line through ``x_i`` and ``z``, and steps along that line. No additive
noise term is used.
Args:
x_i: current state, shape ``(d,)``.
archive: archive ``Z`` of past+present states, shape ``(M, d)`` with
``M >= 3``.
rng: numpy random Generator.
gamma_s: optional fixed step scale. If None, drawn from Uniform(1.2, 2.2).
Returns:
``(x_prop, z_anchor)`` where ``x_prop`` is the proposed state, shape
``(d,)``, and ``z_anchor`` is the anchor ``z`` (shape ``(d,)``) needed
by `snooker_log_jacobian` for the acceptance correction.
Raises:
ValueError: if the archive has fewer than 3 states.
Note:
Mathematically ``zp1 - zp2 == <z1 - z2, u> * u`` (the compact form used
here), so the explicit orthogonal projections are not materialised.
"""
n_states = archive.shape[0]
if n_states < 3:
raise ValueError(
"snooker move requires at least 3 archive states "
f"(anchor + 2 differences); got {n_states}"
)
idx_anchor, idx_1, idx_2 = sample_distinct_indices(rng, n_states, 3)
z_anchor = archive[idx_anchor]
direction = z_anchor - x_i
norm = float(np.linalg.norm(direction))
if norm < _NORM_EPS:
# Degenerate line: anchor coincident with current state. Floor the norm
# (demc_zs.r style) so the unit direction is well-defined instead of
# dividing by zero. Extremely rare once chains have spread.
norm = _NORM_EPS
unit_dir = direction / norm
if gamma_s is None:
gamma_s = draw_gamma_s(rng)
# Compact form of gamma_s * (zp1 - zp2): the projected difference lies along
# the line, so it equals gamma_s * <z1 - z2, u> * u.
projected_diff = float(np.dot(archive[idx_1] - archive[idx_2], unit_dir))
x_prop = x_i + gamma_s * projected_diff * unit_dir
return x_prop, z_anchor
[docs]
def snooker_log_jacobian(
x_i: np.ndarray,
x_prop: np.ndarray,
z_anchor: np.ndarray,
d: int,
) -> float:
"""Radial-Jacobian term of the snooker log-acceptance ratio.
Returns ``(d - 1) * (log||x_prop - z|| - log||x_i - z||)``, computed entirely
in log-space so that the raw power ``||.||**(d-1)`` is never formed.
For ``d == 1`` the term is exactly 0 (snooker reduces to plain Metropolis
along the line). If the proposal lands on the anchor (or the current state
coincides with it) the radius underflows and the move is rejected by
returning ``-inf``, matching the fact that the spherical shell has zero
volume at ``z``.
Args:
x_i: current state, shape ``(d,)``.
x_prop: proposed state, shape ``(d,)``.
z_anchor: anchor ``z`` returned by `snooker_propose`, shape ``(d,)``.
d: number of free parameters.
Returns:
The log-Jacobian contribution to ``log_alpha``; ``-inf`` to force reject.
"""
if d == 1:
return 0.0
r_i = float(np.linalg.norm(x_i - z_anchor))
r_prop = float(np.linalg.norm(x_prop - z_anchor))
if r_prop <= _NORM_EPS or r_i <= _NORM_EPS:
# Cannot jump onto the anchor (r = 0 => acceptance ratio 0), and a
# degenerate current radius is equally pathological: reject.
return -np.inf
return (d - 1) * (np.log(r_prop) - np.log(r_i))