Inference¶

Package¶

inference ¶

inference¶

Inference engines for WPPM.

This subpackage provides different strategies for fitting model parameters to data and returning posterior objects.

Implementations

MAPOptimizer : maximum a posteriori fit with Optax optimizers.
NUTSSampler : full posterior sampling via NUTS (requires blackjax).
LaplaceApproximation : approximate posterior covariance around MAP (stub).
LangevinSampler : skeleton for Langevin-based sampling (stub).

Optional dependencies

NUTSSampler requires blackjax: pip install 'psyphy[sampling]'

Future extensions

adjusted MC samplers, e.g., MALA (for Bayesian posterior inference).

Classes:

Name	Description
`InferenceEngine`	Abstract interface for inference engines.
`LangevinSampler`	Langevin sampler (stub).
`LaplaceApproximation`	Laplace approximation around MAP estimate.
`MAPOptimizer`	MAP (Maximum A Posteriori) optimizer.
`NUTSSampler`	NUTS (No-U-Turn Sampler) using BlackJAX.

Attributes:

Name	Type	Description
`INFERENCE_ENGINES`

INFERENCE_ENGINES ¶

INFERENCE_ENGINES = {'map': MAPOptimizer, 'nuts': NUTSSampler, 'laplace': LaplaceApproximation, 'langevin': LangevinSampler}

InferenceEngine ¶

Bases: ABC

Abstract interface for inference engines.

Methods:

Name	Description
`fit`	Fit model parameters to data and return a Posterior object.

fit ¶

fit(model: Any, data: Any) -> Any

Fit model parameters to data.

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Psychophysical model to fit.	required
`data`	`ResponseData`	Observed trials.	required

Returns:

Type	Description
`Posterior`	Posterior object wrapping fitted params and model reference.

Source code in src/psyphy/inference/base.py

@abstractmethod
def fit(self, model: Any, data: Any) -> Any:
    """
    Fit model parameters to data.

    Parameters
    ----------
    model : WPPM
        Psychophysical model to fit.
    data : ResponseData
        Observed trials.

    Returns
    -------
    Posterior
        Posterior object wrapping fitted params and model reference.
    """
    ...

LangevinSampler ¶

LangevinSampler(steps: int = 1000, step_size: float = 0.001, temperature: float = 1.0)

Bases: InferenceEngine

Langevin sampler (stub).

Parameters:

Name	Type	Description	Default
`steps`	`int`	Number of Langevin steps.	`1000`
`step_size`	`float`	Integration step size.	`1e-3`
`temperature`	`float`	Noise scale (temperature).	`1.0`

Methods:

Name	Description
`fit`	Fit model parameters with Langevin dynamics (stub).

Attributes:

Name	Type	Description
`step_size`
`steps`
`temperature`

Source code in src/psyphy/inference/langevin.py

def __init__(
    self, steps: int = 1000, step_size: float = 1e-3, temperature: float = 1.0
):
    self.steps = steps
    self.step_size = step_size
    self.temperature = temperature

step_size ¶

step_size = step_size

steps ¶

steps = steps

temperature ¶

temperature = temperature

fit ¶

fit(model, data) -> MAPPosterior

Fit model parameters with Langevin dynamics (stub).

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Model instance.	required
`data`	`ResponseData`	Observed trials.	required

Source code in src/psyphy/inference/langevin.py

def fit(self, model, data) -> MAPPosterior:
    """
    Fit model parameters with Langevin dynamics (stub).

    Parameters
    ----------
    model : WPPM
        Model instance.
    data : ResponseData
        Observed trials.

    """

    raise NotImplementedError("Langevin Sampler not implemented yet.")

LaplaceApproximation ¶

Bases: InferenceEngine

Laplace approximation around MAP estimate.

Methods:

Name	Description
`from_map`	Construct a Gaussian approximation centered at MAP.

fit ¶

fit(model: Any, data: Any) -> Any

Fit model parameters to data.

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Psychophysical model to fit.	required
`data`	`ResponseData`	Observed trials.	required

Returns:

Type	Description
`Posterior`	Posterior object wrapping fitted params and model reference.

Source code in src/psyphy/inference/base.py

@abstractmethod
def fit(self, model: Any, data: Any) -> Any:
    """
    Fit model parameters to data.

    Parameters
    ----------
    model : WPPM
        Psychophysical model to fit.
    data : ResponseData
        Observed trials.

    Returns
    -------
    Posterior
        Posterior object wrapping fitted params and model reference.
    """
    ...

from_map ¶

from_map(map_posterior: MAPPosterior) -> MAPPosterior

Return posterior approximation from MAP.

Parameters:

Name	Type	Description	Default
`map_posterior`	`Posterior`	Posterior object from MAP optimization.	required

Returns:

Type	Description
`MAPPosterior`	Posterior distribution containing Laplace approximation.

Source code in src/psyphy/inference/laplace.py

def from_map(self, map_posterior: MAPPosterior) -> MAPPosterior:
    """
    Return posterior approximation from MAP.

    Parameters
    ----------
    map_posterior : Posterior
        Posterior object from MAP optimization.

    Returns
    -------
    MAPPosterior
        Posterior distribution containing Laplace approximation.
    """
    # 1. First find MAP estimate
    # 2. Compute Hessian at MAP
    # 3. Invert Hessian to get covariance
    # 4. Return Gaussian posterior with MAP mean and covariance

    return map_posterior

MAPOptimizer ¶

MAPOptimizer(steps: int = 500, learning_rate: float = 5e-05, momentum: float = 0.9, optimizer: GradientTransformation | None = None, *, reduction: Literal['mean', 'sum'] = 'mean', track_history: bool = True, log_every: int = 1, progress_every: int = 10, show_progress: bool = False, max_grad_norm: float | None = None)

Bases: InferenceEngine

MAP (Maximum A Posteriori) optimizer.

Parameters:

Name	Type	Description	Default
`steps`	`int`	Number of optimization steps.	`500`
`optimizer`	`GradientTransformation`	Optax optimizer to use. Default: SGD with momentum.	`None`

Notes

Loss function = negative log posterior.
Gradients computed with jax.grad.
By default the objective is scaled to a per-trial quantity (reduction="mean"); see the reduction parameter.

Create a MAP optimizer.

Parameters:

Name	Type	Description	Default
`steps`	`int`	Number of optimization steps.	`500`
`optimizer`	`GradientTransformation \| None`	Optax optimizer to use.	`None`
`learning_rate`	`float`	Learning rate for the default optimizer (SGD with momentum).	`5e-05`
`momentum`	`float`	Momentum for the default optimizer (SGD with momentum).	`0.9`
`reduction`	`('mean', 'sum')`	How to scale the objective before differentiating. `"mean"` (default): divide the negative log posterior by the number of trials N, giving a per-trial objective. (Hong et al. 2025, elife, used this and no gradient clipping) `"sum"`: use the negative log posterior as-is. The two objectives differ by the positive constant N, so they have the same minimizer, and with no gradient clipping they trace an identical path under `lr_sum = lr_mean / N`. The choice matters for two practical reasons: Learning-rate portability. Under `"sum"`, gradient magnitude grows with N, so `learning_rate` must be retuned whenever the dataset size changes. Under `"mean"` the gradient is an average of per-trial gradients, so a working learning rate transfers across dataset sizes. Gradient clipping. `max_grad_norm` is a fixed threshold. Under `"sum"` the raw gradient norm scales with N, so the clip saturates on essentially every step for realistic N — which discards gradient magnitude and silently turns SGD into normalized fixed-step descent. To disable it entirely (e.g. to match a reference implementation that does no clipping) pass `max_grad_norm=None`. Note that the model is unaffected: `WPPM.log_posterior_from_data` still returns the true (summed) log posterior. Scaling is step-size conditioning and lives here, in the optimizer, so that density-based consumers (e.g. a Laplace approximation taking the Hessian at the mode) keep seeing the unnormalized log posterior. Because the recorded loss is per-trial under `"mean"`, learning curves are not comparable to those produced with `"sum"`: the values differ by a factor of N.	`"mean"`
`track_history`	`bool`	When True, record loss history during fitting for plotting.	`True`
`log_every`	`int`	Record every N steps (also records the last step).	`1`
`progress_every`	`int`	Update the progress-bar loss display every N steps (and the last step) when show_progress=True. This is kept separate from log_every so you can record loss at high frequency for plotting (e.g. log_every=1) without forcing a device->host sync for the progress UI every step.	`10`
`show_progress`	`bool`	When True, display a tqdm progress bar during fitting. This is a UI feature: `tqdm` ships in the `progress` extra (`pip install 'psyphy[progress]'`), not as a core dependency, and if it is not installed fitting proceeds without a progress bar.	`False`
`max_grad_norm`	`float \| None`	If set, clip gradients by global norm to this value before applying optimizer updates: `g <- g * min(1, max_grad_norm / \|\|g\|\|)`. Off by default. Clipping is a guard rail against rare gradient spikes, and only behaves as one when the threshold sits above the typical gradient norm so that it binds on a small fraction of steps. A threshold below the typical norm binds on every step, which fixes the step length at `learning_rate * max_grad_norm` regardless of gradient magnitude -- that is normalized gradient descent, a different algorithm, and it makes `learning_rate` non-transferable (only the product of the two matters). The typical norm is model- and data-dependent, so there is no safe universal default; for the WPPM it is O(10^2)-O(10^3). Choose a threshold from the observed distribution rather than a convention borrowed from deep-learning recipes, and check :attr:`clip_rate` after :meth:`fit` to confirm which regime you are in. Note this is orthogonal to `reduction`: under `"sum"` the norm additionally scales with the trial count N, but a mis-set threshold saturates under `"mean"` too.	`None`

Attributes:

Name	Type	Description
`clip_rate`	`float \| None`	Fraction of steps on which gradient clipping bound, populated after :meth:`fit`. `None` when `max_grad_norm` is `None`. A value near 1.0 means clipping is not acting as a guard rail but as a change of optimizer.

Methods:

Name	Description
`fit`	Fit model parameters with MAP optimization.
`get_history`	Return (steps, losses) recorded during the last fit when tracking was enabled.

Attributes:

Name	Type	Description
`clip_rate`	`float \| None`
`log_every`
`loss_history`	`list[float]`
`loss_steps`	`list[int]`
`max_grad_norm`
`n_clipped_steps`	`int`
`optimizer`
`progress_every`
`reduction`
`show_progress`
`steps`
`track_history`

Source code in src/psyphy/inference/map_optimizer.py

def __init__(
    self,
    steps: int = 500,
    learning_rate: float = 5e-5,
    momentum: float = 0.9,
    optimizer: optax.GradientTransformation | None = None,
    *,
    reduction: Literal["mean", "sum"] = "mean",
    track_history: bool = True,
    log_every: int = 1,
    progress_every: int = 10,
    show_progress: bool = False,
    max_grad_norm: float | None = None,
):
    """Create a MAP optimizer.

    Parameters
    ----------
    steps : int
        Number of optimization steps.
    optimizer : optax.GradientTransformation | None
        Optax optimizer to use.
    learning_rate : float, optional
        Learning rate for the default optimizer (SGD with momentum).
    momentum : float, optional
        Momentum for the default optimizer (SGD with momentum).
    reduction : {"mean", "sum"}, optional
        How to scale the objective before differentiating.

        - ``"mean"`` (default): divide the negative log posterior by the
          number of trials N, giving a *per-trial* objective.
          (Hong et al. 2025, elife, used this and no gradient clipping)
        - ``"sum"``: use the negative log posterior as-is.

        The two objectives differ by the positive constant N, so they have
        the **same minimizer**, and with no gradient clipping they trace an
        identical path under ``lr_sum = lr_mean / N``. The choice matters
        for two practical reasons:

        1. **Learning-rate portability.** Under ``"sum"``, gradient
           magnitude grows with N, so ``learning_rate`` must be retuned
           whenever the dataset size changes. Under ``"mean"`` the gradient
           is an average of per-trial gradients, so a working learning rate
           transfers across dataset sizes.
        2. **Gradient clipping.** ``max_grad_norm`` is a fixed threshold.
           Under ``"sum"`` the raw gradient norm scales with N, so the clip
           saturates on essentially every step for realistic N — which
           discards gradient *magnitude* and silently turns SGD into
           normalized fixed-step descent.

           To disable it entirely (e.g. to match a reference implementation
           that does no clipping) pass ``max_grad_norm=None``.

        Note that the *model* is unaffected: ``WPPM.log_posterior_from_data``
        still returns the true (summed) log posterior. Scaling is
        step-size conditioning and lives here, in the optimizer, so that
        density-based consumers (e.g. a Laplace approximation taking the
        Hessian at the mode) keep seeing the unnormalized log posterior.

        Because the recorded loss is per-trial under ``"mean"``, learning
        curves are not comparable to those produced with ``"sum"``: the
        values differ by a factor of N.
    track_history : bool, optional
        When True, record loss history during fitting for plotting.
    log_every : int, optional
        Record every N steps (also records the last step).
    progress_every : int, optional
        Update the progress-bar loss display every N steps (and the last step)
        when show_progress=True.
        This is kept separate from log_every so you can record loss at high
        frequency for plotting (e.g. log_every=1) without forcing a device->host
        sync for the progress UI every step.
    show_progress : bool, optional
        When True, display a tqdm progress bar during fitting.
        This is a UI feature: `tqdm` ships in the ``progress`` extra
        (``pip install 'psyphy[progress]'``), not as a core dependency, and
        if it is not installed fitting proceeds without a progress bar.
    max_grad_norm : float | None, optional
        If set, clip gradients by global norm to this value before applying
        optimizer updates: ``g <- g * min(1, max_grad_norm / ||g||)``.

        **Off by default.** Clipping is a guard rail against rare gradient
        spikes, and only behaves as one when the threshold sits *above* the
        typical gradient norm so that it binds on a small fraction of steps.
        A threshold below the typical norm binds on every step, which fixes
        the step length at ``learning_rate * max_grad_norm`` regardless of
        gradient magnitude -- that is normalized gradient descent, a
        different algorithm, and it makes ``learning_rate`` non-transferable
        (only the product of the two matters).

        The typical norm is model- and data-dependent, so there is no safe
        universal default; for the WPPM it is O(10^2)-O(10^3). Choose a
        threshold from the observed distribution rather than a convention
        borrowed from deep-learning recipes, and check
        :attr:`clip_rate` after :meth:`fit` to confirm which regime you are
        in. Note this is orthogonal to ``reduction``: under ``"sum"`` the
        norm additionally scales with the trial count N, but a mis-set
        threshold saturates under ``"mean"`` too.

    Attributes
    ----------
    clip_rate : float | None
        Fraction of steps on which gradient clipping bound, populated after
        :meth:`fit`. ``None`` when ``max_grad_norm`` is ``None``. A value
        near 1.0 means clipping is not acting as a guard rail but as a
        change of optimizer.
    """
    self.steps = steps
    if reduction not in ("mean", "sum"):
        raise ValueError(f'reduction must be "mean" or "sum", got {reduction!r}.')
    self.reduction = reduction
    base_optimizer = optimizer or optax.sgd(
        learning_rate=learning_rate, momentum=momentum
    )
    if max_grad_norm is None:
        self.optimizer = base_optimizer
    else:
        self.optimizer = optax.chain(
            optax.clip_by_global_norm(float(max_grad_norm)),
            base_optimizer,
        )

    self.track_history = track_history
    self.log_every = max(1, int(log_every))
    self.progress_every = max(1, int(progress_every))
    self.show_progress = bool(show_progress)
    self.max_grad_norm = max_grad_norm
    # Exposed after fit() when tracking is enabled
    self.loss_steps: list[int] = []
    self.loss_history: list[float] = []
    # Exposed after fit(); None when clipping is off. See `max_grad_norm`.
    self.clip_rate: float | None = None
    self.n_clipped_steps: int = 0

clip_rate ¶

clip_rate: float | None = None

log_every ¶

log_every = max(1, int(log_every))

loss_history ¶

loss_history: list[float] = []

loss_steps ¶

loss_steps: list[int] = []

max_grad_norm ¶

max_grad_norm = max_grad_norm

n_clipped_steps ¶

n_clipped_steps: int = 0

optimizer ¶

optimizer = base_optimizer

progress_every ¶

progress_every = max(1, int(progress_every))

reduction ¶

reduction = reduction

show_progress ¶

show_progress = bool(show_progress)

steps ¶

steps = steps

track_history ¶

track_history = track_history

fit ¶

fit(model, data, init_params: dict | None = None, seed: int | None = None) -> MAPPosterior

Fit model parameters with MAP optimization.

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Model instance.	required
`data`	`ResponseData`	Observed trials.	required
`init_params`	`dict \| None`	Initial parameter PyTree to start optimization from. If provided, this takes precedence over the seed.	`None`
`seed`	`int \| None`	PRNG seed used to draw initial parameters from the model's prior when init_params is not provided, and as the base key for the MC likelihood random stream during optimization. If None, defaults to 0.	`None`

Returns:

Type	Description
`MAPPosterior`	Posterior wrapper around MAP params and model.

Source code in src/psyphy/inference/map_optimizer.py

def fit(
    self,
    model,
    data,
    init_params: dict | None = None,
    seed: int | None = None,
) -> MAPPosterior:
    """
    Fit model parameters with MAP optimization.

    Parameters
    ----------
    model : WPPM
        Model instance.
    data : ResponseData
        Observed trials.
    init_params : dict | None, optional
        Initial parameter PyTree to start optimization from. If provided,
        this takes precedence over the seed.
    seed : int | None, optional
        PRNG seed used to draw initial parameters from the model's prior
        when init_params is not provided, and as the base key for the MC
        likelihood random stream during optimization. If None, defaults to 0.

    Returns
    -------
    MAPPosterior
        Posterior wrapper around MAP params and model.
    """

    rng_seed = 0 if seed is None else int(seed)
    # Master key: split into init key and optimization key stream.
    master_key = jax.random.PRNGKey(rng_seed)
    init_key, opt_key = jax.random.split(master_key)

    # Initialize parameters
    params = init_params if init_params is not None else model.init_params(init_key)
    opt_state = self.optimizer.init(params)

    # Objective scale. Resolved here, outside the jitted step, so it is a
    # Python float baked in at trace time rather than a traced value.
    # `model.log_posterior_from_data` always returns the true (summed) log
    # posterior; "mean" turns it into a per-trial quantity for the gradient
    # step. See the `reduction` docstring for why this lives here and not
    # in the model.
    scale = 1.0
    if self.reduction == "mean":
        n_trials = int(jax.numpy.asarray(data.responses).shape[0])
        if n_trials == 0:
            raise ValueError('reduction="mean" requires at least one trial.')
        scale = 1.0 / n_trials

    # key is now an explicit argument so each JIT-compiled call receives a
    # distinct random key — a fresh MC noise realization per gradient step.
    clip_at = self.max_grad_norm

    @jax.jit
    def step(params, opt_state, key):
        loss, grads = jax.value_and_grad(
            lambda p: -scale * model.log_posterior_from_data(p, data, key=key)
        )(params)
        # Measured before the update, so it is the norm the clip actually saw.
        # Accumulated on device and read once after the loop, to avoid adding
        # a host sync per step.
        clipped = (
            jax.numpy.zeros((), dtype=jax.numpy.int32)
            if clip_at is None
            else (_global_norm(grads) > clip_at).astype(jax.numpy.int32)
        )
        updates, opt_state = self.optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, loss, clipped

    # clear any previous history
    if self.track_history:
        self.loss_steps.clear()
        self.loss_history.clear()
    self.clip_rate = None
    self.n_clipped_steps = 0
    n_clipped = jax.numpy.zeros((), dtype=jax.numpy.int32)
    steps_run = 0

    # Optional progress bar.
    #
    # Why we *manually* advance the bar:
    # - When JAX runs on GPU, the first `step(...)` call can spend a long time in
    #   compilation, and tqdm may not visibly advance if the underlying iterator
    #   doesn't get a chance to redraw.
    # - By keeping a normal `range(self.steps)` loop and calling `pbar.update(1)`
    #   ourselves, we ensure the bar advances exactly once per iteration.
    #
    # Performance note: *displaying the loss* requires transferring `loss` from
    # device -> host, which can add sync overhead. We therefore only attach a
    # loss postfix every `progress_every` steps.
    pbar = None
    if self.show_progress:
        try:
            from tqdm.auto import tqdm

            pbar = tqdm(total=self.steps, desc="MAP fit", leave=False)
        except Exception:
            # Soft dependency: tqdm not available (or terminal unsuitable).
            pbar = None

    for i in range(self.steps):
        # Split a fresh subkey for each step so the MC likelihood sees a
        # different noise realization on every gradient evaluation.
        opt_key, subkey = jax.random.split(opt_key)
        params, opt_state, loss, clipped = step(params, opt_state, subkey)
        n_clipped = n_clipped + clipped
        steps_run = i + 1

        # Non-finite guard: if loss becomes NaN/Inf, optimization has diverged.
        # Stop early so downstream plots don’t look “truncated” due to NaNs.
        if not bool(jax.numpy.isfinite(loss)):
            if self.track_history:
                try:
                    self.loss_steps.append(i)
                    self.loss_history.append(float(loss))
                except Exception:
                    pass
            print(
                f"[MAPOptimizer] Non-finite loss at step {i}: {loss}. "
                "Stopping early."
            )
            if pbar is not None:
                with contextlib.suppress(Exception):
                    pbar.update(1)
            break

        if self.track_history and (
            (i % self.log_every == 0) or (i == self.steps - 1)
        ):
            # Pull scalar to host and record
            try:
                self.loss_steps.append(i)
                self.loss_history.append(float(loss))
            except Exception:
                #  do not break fitting if logging fails
                pass

        #  progress bar loss display (avoid host sync every step)
        if pbar is not None and (
            (i % self.progress_every == 0) or (i == self.steps - 1)
        ):
            with contextlib.suppress(Exception):
                pbar.set_postfix(loss=float(loss))

        if pbar is not None:
            with contextlib.suppress(Exception):
                pbar.update(1)
                # Encourage a redraw occasionally in environments with buffered/stale
                # TTY updates.
                if (i % self.progress_every == 0) or (i == self.steps - 1):
                    pbar.refresh()

    if pbar is not None:
        with contextlib.suppress(Exception):
            pbar.close()

    # Clipping regime. Read once, after the loop, so the counter costs no
    # per-step host sync. A rate near 1.0 means the threshold sits below the
    # model's typical gradient norm, so every step is rescaled to
    # `max_grad_norm` and the effective algorithm is no longer plain SGD.
    if self.max_grad_norm is not None and steps_run > 0:
        self.n_clipped_steps = int(n_clipped)
        self.clip_rate = self.n_clipped_steps / steps_run
        if self.clip_rate > 0.5:
            warnings.warn(
                f"Gradient clipping bound on {self.clip_rate:.0%} of steps "
                f"({self.n_clipped_steps}/{steps_run}) at "
                f"max_grad_norm={self.max_grad_norm}. At this rate clipping "
                "is not a guard rail: the step length is fixed at "
                "learning_rate * max_grad_norm regardless of gradient "
                "magnitude, so `learning_rate` no longer means what it would "
                "without clipping. Raise max_grad_norm above the typical "
                "gradient norm, or pass max_grad_norm=None.",
                stacklevel=2,
            )

    return MAPPosterior(params=params, model=model)

get_history ¶

get_history() -> tuple[list[int], list[float]]

Return (steps, losses) recorded during the last fit when tracking was enabled.

Source code in src/psyphy/inference/map_optimizer.py

def get_history(self) -> tuple[list[int], list[float]]:
    """Return (steps, losses) recorded during the last fit when tracking was enabled."""
    return self.loss_steps, self.loss_history

NUTSSampler ¶

NUTSSampler(num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 4, *, logdensity_key_seed: int = 0, target_acceptance_rate: float = 0.8, step_size: float | None = None, max_num_doublings: int = 8, prior_preconditioning: bool = True, seed: int = 0, show_progress: bool = True)

Bases: InferenceEngine

NUTS (No-U-Turn Sampler) using BlackJAX.

Draws samples from the full posterior p(θ | data) via gradient-based MCMC. Returns an MCMCPosterior with shape (num_chains, num_samples, *param_shape).

Parameters:

Name	Type	Description	Default
`num_warmup`	`int`	Number of warmup/burn-in steps per chain.	`500`
`num_samples`	`int`	Number of posterior draws per chain (after warmup).	`1000`
`num_chains`	`int`	Number of independent chains to run.	`4`
`logdensity_key_seed`	`int`	Seed for the fixed JAX key baked into the logdensity closure. Controls which MC noise realization is used throughout sampling. See module docstring for rationale.	`0`
`target_acceptance_rate`	`float`	Target Metropolis acceptance rate for window adaptation (Mode A only). Typical values: 0.-0.9.	`0.8`
`step_size`	`float \| None`	If None (default), use blackjax.window_adaptation to automatically tune step_size and mass matrix (Mode A). If a float, skip adaptation and use this fixed step_size for all chains — warmup becomes plain burn-in and all chains are vmapped in parallel (Mode B). The user is responsible for choosing a suitable value.	`None`
`max_num_doublings`	`int`	Maximum NUTS tree depth. One draw costs at most 2**max_num_doublings - 1 gradient evaluations (255 at the default; BlackJAX's own default of 10 allows 1023). Applies to both modes, warmup and sampling. If `sampler_stats["n_steps"]` sits at the cap for most draws, the trajectories are being truncated: mixing gets worse (lower ESS) but the chain is still a valid MCMC sampler. Raise the cap, or improve the parameterization.	`8`
`prior_preconditioning`	`bool`	Mode A only. If True and the model exposes `model.prior.param_scales()` (WPPM does), sample z = W / prior_std instead of W, so adaptation starts from a mass matrix matched to the prior scales rather than the identity. Results are returned in W coordinates either way. Set False to reproduce the old identity-start behaviour.	`True`
`seed`	`int`	Master PRNG seed for chain initialization and sampling keys.	`0`
`show_progress`	`bool`	Display a tqdm progress bar during warmup and sampling. Requires `tqdm`, which ships in the `progress` extra (`pip install 'psyphy[progress]'`) and is NOT a core dependency. Silently skipped if tqdm is not installed, so on a clean install this default is a no-op and a long fit prints nothing.	`True`

Notes

Two warmup modes:

Mode A (step_size=None, default): blackjax.window_adaptation adapts step_size and inverse_mass_matrix independently per chain. The adaptation state is all JAX arrays, so warmup is vmapped over chains and JIT-compiled once; sampling is likewise one vmapped, jitted lax.scan using each chain's own adapted parameters.

Mode B (step_size=): No adaptation. All chains share the same fixed step_size and a diagonal mass matrix of ones. Warmup is pure burn-in. All chains are fully vmapped -> single compiled call. Faster for many chains but requires manual step_size tuning.

Examples:

>>> from psyphy.inference import NUTSSampler
>>> sampler = NUTSSampler(num_warmup=200, num_samples=500, num_chains=4)
>>> nuts_posterior = sampler.fit(model, data, init_params=map_posterior.params)
>>> print(nuts_posterior.n_chains, nuts_posterior.n_draws)
4 500

Methods:

Name	Description
`fit`	Run NUTS sampling and return a posterior over model parameters.

Attributes:

Name	Type	Description
`logdensity_key_seed`
`max_num_doublings`
`num_chains`
`num_samples`
`num_warmup`
`prior_preconditioning`
`seed`
`show_progress`
`step_size`
`target_acceptance_rate`

Source code in src/psyphy/inference/nuts.py

def __init__(
    self,
    num_warmup: int = 500,
    num_samples: int = 1000,
    num_chains: int = 4,
    *,
    logdensity_key_seed: int = 0,
    target_acceptance_rate: float = 0.8,
    step_size: float | None = None,
    max_num_doublings: int = 8,
    prior_preconditioning: bool = True,
    seed: int = 0,
    show_progress: bool = True,
):
    self.num_warmup = int(num_warmup)
    self.num_samples = int(num_samples)
    self.num_chains = int(num_chains)
    self.logdensity_key_seed = int(logdensity_key_seed)
    self.target_acceptance_rate = float(target_acceptance_rate)
    self.step_size = float(step_size) if step_size is not None else None
    if int(max_num_doublings) < 1:
        raise ValueError(f"max_num_doublings must be >= 1, got {max_num_doublings}")
    self.max_num_doublings = int(max_num_doublings)
    self.prior_preconditioning = bool(prior_preconditioning)
    self.seed = int(seed)
    self.show_progress = bool(show_progress)

logdensity_key_seed ¶

logdensity_key_seed = int(logdensity_key_seed)

max_num_doublings ¶

max_num_doublings = int(max_num_doublings)

num_chains ¶

num_chains = int(num_chains)

num_samples ¶

num_samples = int(num_samples)

num_warmup ¶

num_warmup = int(num_warmup)

prior_preconditioning ¶

prior_preconditioning = bool(prior_preconditioning)

seed ¶

seed = int(seed)

show_progress ¶

show_progress = bool(show_progress)

step_size ¶

step_size = float(step_size) if step_size is not None else None

target_acceptance_rate ¶

target_acceptance_rate = float(target_acceptance_rate)

fit ¶

fit(model, data, init_params: dict | None = None, seed: int | None = None) -> MCMCPosterior

Run NUTS sampling and return a posterior over model parameters.

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Model instance providing log_posterior_from_data.	required
`data`	`TrialData`	Observed trials.	required
`init_params`	`dict \| None`	Initial parameter PyTree. If None, samples from the model prior using the master seed. Providing MAP params as init_params dramatically reduces the warmup needed.	`None`
`seed`	`int \| None`	Override the master PRNG seed set at construction.	`None`

Returns:

Type	Description
`MCMCPosterior`	Samples shaped (num_chains, num_samples, *param_shape), sampler_stats with acceptance_rate.

Source code in src/psyphy/inference/nuts.py

def fit(
    self,
    model,
    data,
    init_params: dict | None = None,
    seed: int | None = None,
) -> MCMCPosterior:
    """
    Run NUTS sampling and return a posterior over model parameters.

    Parameters
    ----------
    model : WPPM
        Model instance providing log_posterior_from_data.
    data : TrialData
        Observed trials.
    init_params : dict | None, optional
        Initial parameter PyTree. If None, samples from the model prior
        using the master seed. Providing MAP params as init_params
        dramatically reduces the warmup needed.
    seed : int | None, optional
        Override the master PRNG seed set at construction.

    Returns
    -------
    MCMCPosterior
        Samples shaped (num_chains, num_samples, *param_shape),
        sampler_stats with acceptance_rate.
    """
    try:
        import blackjax
    except ImportError as e:
        raise ImportError(
            "BlackJAX is required for NUTSSampler. "
            "Install it with: pip install 'psyphy[sampling]'\n"
            f"Note: blackjax >= 1.4 requires Python >= 3.11 (this is "
            f"{sys.version_info.major}.{sys.version_info.minor}), so on an "
            "older interpreter the 'sampling' extra resolves to nothing and "
            "this error appears even after installing it."
        ) from e

    rng_seed = self.seed if seed is None else int(seed)
    master_key = jax.random.PRNGKey(rng_seed)
    init_key, chain_key = jax.random.split(master_key)

    # init_position: dict matching the model's parameter PyTree structure.
    # For WPPM: {"W": jnp.ndarray of shape (degree+1, degree+1, input_dim, embedding_dim)}
    # e.g. degree=1, input_dim=2, extra_dims=1 -> {"W": shape (2, 2, 2, 3)}
    # BlackJAX requires this as the starting point for Hamiltonian dynamics.
    init_position = (
        init_params if init_params is not None else model.init_params(init_key)
    )

    # Fixed MC key baked into logdensity — makes the stochastic
    # likelihood deterministic for NUTS gradient evaluations.
    fixed_mc_key = jax.random.PRNGKey(self.logdensity_key_seed)

    # -- BlackJAX API boundary --------------------------------
    # logdensity_fn is the ONLY thing BlackJAX needs from psyphy.
    # Contract:
    #   input:  params  — dict, same structure as init_position
    #                      e.g. {"W": (2, 2, 2, 3)}
    #   output: scalar  — log p(θ | data), a JAX float (shape ())
    #   must be: JAX-differentiable (jax.grad must work on it)
    #            (nearly) deterministic — hence the fixed_mc_key
    # -----------------------------------------
    def logdensity_fn(params: dict) -> jnp.ndarray:
        return model.log_posterior_from_data(params, data, key=fixed_mc_key)

    # chain_keys: shape (num_chains, 2) — one PRNGKey per chain,
    # split further inside each mode into warmup_key and sample_key.
    chain_keys = jax.random.split(chain_key, self.num_chains)

    # blackjax is passed as an argument rather than imported again in each
    # private method. Because it is a soft dependency, the single import
    # (with its error-handling) lives here in fit(). In Python, modules are
    # first-class objects — 'import blackjax' just binds the module to a
    # local variable, which can be passed around like any other value.
    if self.step_size is None:
        # Prior preconditioning (Mode A only -- in Mode B the user picked
        # step_size for W coordinates, and silently rescaling would change
        # what that number means). See "Notes on warmup cost" above.
        scales = (
            self._prior_scales(model, init_position)
            if self.prior_preconditioning
            else None
        )
        if scales is None:
            positions, sampler_stats = self._run_adaptive(
                blackjax, logdensity_fn, init_position, chain_keys
            )
        else:
            # Change of variables W = scales * z. The Jacobian of a fixed
            # linear map is a constant, so it is dropped: it does not
            # change the posterior or any accept/reject decision.
            def z_logdensity_fn(z: dict) -> jnp.ndarray:
                return logdensity_fn(jax.tree.map(jnp.multiply, z, scales))

            init_z = jax.tree.map(jnp.divide, init_position, scales)
            z_positions, sampler_stats = self._run_adaptive(
                blackjax, z_logdensity_fn, init_z, chain_keys
            )
            # Map draws back to W. z leaves are (n_chains, n_draws, *shape)
            # and scales leaves are (*shape), so they broadcast over the
            # two leading axes. Every downstream consumer (MCMCPosterior,
            # ArviZ, plots) therefore only ever sees W.
            positions = jax.tree.map(jnp.multiply, z_positions, scales)
    else:
        positions, sampler_stats = self._run_fixed_stepsize(
            blackjax, logdensity_fn, init_position, chain_keys
        )

    # positions: dict, each value shaped (n_chains, n_draws, *param_shape)
    #   e.g. {"W": (num_chains, num_samples, 2, 2, 2, 3)}
    # sampler_stats: each value shaped (num_chains, num_samples) — one scalar
    # per draw per chain. Keys follow ArviZ's sample_stats convention.
    return MCMCPosterior(positions, model, sampler_stats=sampler_stats)

Base¶

base ¶

base.py

Abstract base class for inference engines.

All inference engines must implement a fit(model, data) method that returns a Posterior object.

All inference engines (MAPOptimizer, LangevinSampler, LaplaceApproximation) subclass from this base.

Classes:

Name	Description
`InferenceEngine`	Abstract interface for inference engines.

InferenceEngine ¶

Bases: ABC

Abstract interface for inference engines.

Methods:

Name	Description
`fit`	Fit model parameters to data and return a Posterior object.

fit ¶

fit(model: Any, data: Any) -> Any

Fit model parameters to data.

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Psychophysical model to fit.	required
`data`	`ResponseData`	Observed trials.	required

Returns:

Type	Description
`Posterior`	Posterior object wrapping fitted params and model reference.

Source code in src/psyphy/inference/base.py

@abstractmethod
def fit(self, model: Any, data: Any) -> Any:
    """
    Fit model parameters to data.

    Parameters
    ----------
    model : WPPM
        Psychophysical model to fit.
    data : ResponseData
        Observed trials.

    Returns
    -------
    Posterior
        Posterior object wrapping fitted params and model reference.
    """
    ...

MAP Optimizer¶

map_optimizer ¶

map_optimizer.py

MAP (Maximum A Posteriori) optimizer using Optax.

MVP implementation: - Uses gradient ascent on log posterior. - Defaults to SGD with momentum, but any Optax optimizer can be passed in.

Connections

Calls WPPM.log_posterior_from_data(params, data) as the objective.
Returns a Posterior object wrapping the MAP estimate.

Classes:

Name	Description
`MAPOptimizer`	MAP (Maximum A Posteriori) optimizer.

MAPOptimizer ¶

MAPOptimizer(steps: int = 500, learning_rate: float = 5e-05, momentum: float = 0.9, optimizer: GradientTransformation | None = None, *, reduction: Literal['mean', 'sum'] = 'mean', track_history: bool = True, log_every: int = 1, progress_every: int = 10, show_progress: bool = False, max_grad_norm: float | None = None)

Bases: InferenceEngine

MAP (Maximum A Posteriori) optimizer.

Parameters:

Name	Type	Description	Default
`steps`	`int`	Number of optimization steps.	`500`
`optimizer`	`GradientTransformation`	Optax optimizer to use. Default: SGD with momentum.	`None`

Notes

Loss function = negative log posterior.
Gradients computed with jax.grad.
By default the objective is scaled to a per-trial quantity (reduction="mean"); see the reduction parameter.

Create a MAP optimizer.

Parameters:

Name	Type	Description	Default
`steps`	`int`	Number of optimization steps.	`500`
`optimizer`	`GradientTransformation \| None`	Optax optimizer to use.	`None`
`learning_rate`	`float`	Learning rate for the default optimizer (SGD with momentum).	`5e-05`
`momentum`	`float`	Momentum for the default optimizer (SGD with momentum).	`0.9`
`reduction`	`('mean', 'sum')`	How to scale the objective before differentiating. `"mean"` (default): divide the negative log posterior by the number of trials N, giving a per-trial objective. (Hong et al. 2025, elife, used this and no gradient clipping) `"sum"`: use the negative log posterior as-is. The two objectives differ by the positive constant N, so they have the same minimizer, and with no gradient clipping they trace an identical path under `lr_sum = lr_mean / N`. The choice matters for two practical reasons: Learning-rate portability. Under `"sum"`, gradient magnitude grows with N, so `learning_rate` must be retuned whenever the dataset size changes. Under `"mean"` the gradient is an average of per-trial gradients, so a working learning rate transfers across dataset sizes. Gradient clipping. `max_grad_norm` is a fixed threshold. Under `"sum"` the raw gradient norm scales with N, so the clip saturates on essentially every step for realistic N — which discards gradient magnitude and silently turns SGD into normalized fixed-step descent. To disable it entirely (e.g. to match a reference implementation that does no clipping) pass `max_grad_norm=None`. Note that the model is unaffected: `WPPM.log_posterior_from_data` still returns the true (summed) log posterior. Scaling is step-size conditioning and lives here, in the optimizer, so that density-based consumers (e.g. a Laplace approximation taking the Hessian at the mode) keep seeing the unnormalized log posterior. Because the recorded loss is per-trial under `"mean"`, learning curves are not comparable to those produced with `"sum"`: the values differ by a factor of N.	`"mean"`
`track_history`	`bool`	When True, record loss history during fitting for plotting.	`True`
`log_every`	`int`	Record every N steps (also records the last step).	`1`
`progress_every`	`int`	Update the progress-bar loss display every N steps (and the last step) when show_progress=True. This is kept separate from log_every so you can record loss at high frequency for plotting (e.g. log_every=1) without forcing a device->host sync for the progress UI every step.	`10`
`show_progress`	`bool`	When True, display a tqdm progress bar during fitting. This is a UI feature: `tqdm` ships in the `progress` extra (`pip install 'psyphy[progress]'`), not as a core dependency, and if it is not installed fitting proceeds without a progress bar.	`False`
`max_grad_norm`	`float \| None`	If set, clip gradients by global norm to this value before applying optimizer updates: `g <- g * min(1, max_grad_norm / \|\|g\|\|)`. Off by default. Clipping is a guard rail against rare gradient spikes, and only behaves as one when the threshold sits above the typical gradient norm so that it binds on a small fraction of steps. A threshold below the typical norm binds on every step, which fixes the step length at `learning_rate * max_grad_norm` regardless of gradient magnitude -- that is normalized gradient descent, a different algorithm, and it makes `learning_rate` non-transferable (only the product of the two matters). The typical norm is model- and data-dependent, so there is no safe universal default; for the WPPM it is O(10^2)-O(10^3). Choose a threshold from the observed distribution rather than a convention borrowed from deep-learning recipes, and check :attr:`clip_rate` after :meth:`fit` to confirm which regime you are in. Note this is orthogonal to `reduction`: under `"sum"` the norm additionally scales with the trial count N, but a mis-set threshold saturates under `"mean"` too.	`None`

Attributes:

Name	Type	Description
`clip_rate`	`float \| None`	Fraction of steps on which gradient clipping bound, populated after :meth:`fit`. `None` when `max_grad_norm` is `None`. A value near 1.0 means clipping is not acting as a guard rail but as a change of optimizer.

Methods:

Name	Description
`fit`	Fit model parameters with MAP optimization.
`get_history`	Return (steps, losses) recorded during the last fit when tracking was enabled.

Attributes:

Name	Type	Description
`clip_rate`	`float \| None`
`log_every`
`loss_history`	`list[float]`
`loss_steps`	`list[int]`
`max_grad_norm`
`n_clipped_steps`	`int`
`optimizer`
`progress_every`
`reduction`
`show_progress`
`steps`
`track_history`

Source code in src/psyphy/inference/map_optimizer.py

def __init__(
    self,
    steps: int = 500,
    learning_rate: float = 5e-5,
    momentum: float = 0.9,
    optimizer: optax.GradientTransformation | None = None,
    *,
    reduction: Literal["mean", "sum"] = "mean",
    track_history: bool = True,
    log_every: int = 1,
    progress_every: int = 10,
    show_progress: bool = False,
    max_grad_norm: float | None = None,
):
    """Create a MAP optimizer.

    Parameters
    ----------
    steps : int
        Number of optimization steps.
    optimizer : optax.GradientTransformation | None
        Optax optimizer to use.
    learning_rate : float, optional
        Learning rate for the default optimizer (SGD with momentum).
    momentum : float, optional
        Momentum for the default optimizer (SGD with momentum).
    reduction : {"mean", "sum"}, optional
        How to scale the objective before differentiating.

        - ``"mean"`` (default): divide the negative log posterior by the
          number of trials N, giving a *per-trial* objective.
          (Hong et al. 2025, elife, used this and no gradient clipping)
        - ``"sum"``: use the negative log posterior as-is.

        The two objectives differ by the positive constant N, so they have
        the **same minimizer**, and with no gradient clipping they trace an
        identical path under ``lr_sum = lr_mean / N``. The choice matters
        for two practical reasons:

        1. **Learning-rate portability.** Under ``"sum"``, gradient
           magnitude grows with N, so ``learning_rate`` must be retuned
           whenever the dataset size changes. Under ``"mean"`` the gradient
           is an average of per-trial gradients, so a working learning rate
           transfers across dataset sizes.
        2. **Gradient clipping.** ``max_grad_norm`` is a fixed threshold.
           Under ``"sum"`` the raw gradient norm scales with N, so the clip
           saturates on essentially every step for realistic N — which
           discards gradient *magnitude* and silently turns SGD into
           normalized fixed-step descent.

           To disable it entirely (e.g. to match a reference implementation
           that does no clipping) pass ``max_grad_norm=None``.

        Note that the *model* is unaffected: ``WPPM.log_posterior_from_data``
        still returns the true (summed) log posterior. Scaling is
        step-size conditioning and lives here, in the optimizer, so that
        density-based consumers (e.g. a Laplace approximation taking the
        Hessian at the mode) keep seeing the unnormalized log posterior.

        Because the recorded loss is per-trial under ``"mean"``, learning
        curves are not comparable to those produced with ``"sum"``: the
        values differ by a factor of N.
    track_history : bool, optional
        When True, record loss history during fitting for plotting.
    log_every : int, optional
        Record every N steps (also records the last step).
    progress_every : int, optional
        Update the progress-bar loss display every N steps (and the last step)
        when show_progress=True.
        This is kept separate from log_every so you can record loss at high
        frequency for plotting (e.g. log_every=1) without forcing a device->host
        sync for the progress UI every step.
    show_progress : bool, optional
        When True, display a tqdm progress bar during fitting.
        This is a UI feature: `tqdm` ships in the ``progress`` extra
        (``pip install 'psyphy[progress]'``), not as a core dependency, and
        if it is not installed fitting proceeds without a progress bar.
    max_grad_norm : float | None, optional
        If set, clip gradients by global norm to this value before applying
        optimizer updates: ``g <- g * min(1, max_grad_norm / ||g||)``.

        **Off by default.** Clipping is a guard rail against rare gradient
        spikes, and only behaves as one when the threshold sits *above* the
        typical gradient norm so that it binds on a small fraction of steps.
        A threshold below the typical norm binds on every step, which fixes
        the step length at ``learning_rate * max_grad_norm`` regardless of
        gradient magnitude -- that is normalized gradient descent, a
        different algorithm, and it makes ``learning_rate`` non-transferable
        (only the product of the two matters).

        The typical norm is model- and data-dependent, so there is no safe
        universal default; for the WPPM it is O(10^2)-O(10^3). Choose a
        threshold from the observed distribution rather than a convention
        borrowed from deep-learning recipes, and check
        :attr:`clip_rate` after :meth:`fit` to confirm which regime you are
        in. Note this is orthogonal to ``reduction``: under ``"sum"`` the
        norm additionally scales with the trial count N, but a mis-set
        threshold saturates under ``"mean"`` too.

    Attributes
    ----------
    clip_rate : float | None
        Fraction of steps on which gradient clipping bound, populated after
        :meth:`fit`. ``None`` when ``max_grad_norm`` is ``None``. A value
        near 1.0 means clipping is not acting as a guard rail but as a
        change of optimizer.
    """
    self.steps = steps
    if reduction not in ("mean", "sum"):
        raise ValueError(f'reduction must be "mean" or "sum", got {reduction!r}.')
    self.reduction = reduction
    base_optimizer = optimizer or optax.sgd(
        learning_rate=learning_rate, momentum=momentum
    )
    if max_grad_norm is None:
        self.optimizer = base_optimizer
    else:
        self.optimizer = optax.chain(
            optax.clip_by_global_norm(float(max_grad_norm)),
            base_optimizer,
        )

    self.track_history = track_history
    self.log_every = max(1, int(log_every))
    self.progress_every = max(1, int(progress_every))
    self.show_progress = bool(show_progress)
    self.max_grad_norm = max_grad_norm
    # Exposed after fit() when tracking is enabled
    self.loss_steps: list[int] = []
    self.loss_history: list[float] = []
    # Exposed after fit(); None when clipping is off. See `max_grad_norm`.
    self.clip_rate: float | None = None
    self.n_clipped_steps: int = 0

clip_rate ¶

clip_rate: float | None = None

log_every ¶

log_every = max(1, int(log_every))

loss_history ¶

loss_history: list[float] = []

loss_steps ¶

loss_steps: list[int] = []

max_grad_norm ¶

max_grad_norm = max_grad_norm

n_clipped_steps ¶

n_clipped_steps: int = 0

optimizer ¶

optimizer = base_optimizer

progress_every ¶

progress_every = max(1, int(progress_every))

reduction ¶

reduction = reduction

show_progress ¶

show_progress = bool(show_progress)

steps ¶

steps = steps

track_history ¶

track_history = track_history

fit ¶

fit(model, data, init_params: dict | None = None, seed: int | None = None) -> MAPPosterior

Fit model parameters with MAP optimization.

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Model instance.	required
`data`	`ResponseData`	Observed trials.	required
`init_params`	`dict \| None`	Initial parameter PyTree to start optimization from. If provided, this takes precedence over the seed.	`None`
`seed`	`int \| None`	PRNG seed used to draw initial parameters from the model's prior when init_params is not provided, and as the base key for the MC likelihood random stream during optimization. If None, defaults to 0.	`None`

Returns:

Type	Description
`MAPPosterior`	Posterior wrapper around MAP params and model.

Source code in src/psyphy/inference/map_optimizer.py

def fit(
    self,
    model,
    data,
    init_params: dict | None = None,
    seed: int | None = None,
) -> MAPPosterior:
    """
    Fit model parameters with MAP optimization.

    Parameters
    ----------
    model : WPPM
        Model instance.
    data : ResponseData
        Observed trials.
    init_params : dict | None, optional
        Initial parameter PyTree to start optimization from. If provided,
        this takes precedence over the seed.
    seed : int | None, optional
        PRNG seed used to draw initial parameters from the model's prior
        when init_params is not provided, and as the base key for the MC
        likelihood random stream during optimization. If None, defaults to 0.

    Returns
    -------
    MAPPosterior
        Posterior wrapper around MAP params and model.
    """

    rng_seed = 0 if seed is None else int(seed)
    # Master key: split into init key and optimization key stream.
    master_key = jax.random.PRNGKey(rng_seed)
    init_key, opt_key = jax.random.split(master_key)

    # Initialize parameters
    params = init_params if init_params is not None else model.init_params(init_key)
    opt_state = self.optimizer.init(params)

    # Objective scale. Resolved here, outside the jitted step, so it is a
    # Python float baked in at trace time rather than a traced value.
    # `model.log_posterior_from_data` always returns the true (summed) log
    # posterior; "mean" turns it into a per-trial quantity for the gradient
    # step. See the `reduction` docstring for why this lives here and not
    # in the model.
    scale = 1.0
    if self.reduction == "mean":
        n_trials = int(jax.numpy.asarray(data.responses).shape[0])
        if n_trials == 0:
            raise ValueError('reduction="mean" requires at least one trial.')
        scale = 1.0 / n_trials

    # key is now an explicit argument so each JIT-compiled call receives a
    # distinct random key — a fresh MC noise realization per gradient step.
    clip_at = self.max_grad_norm

    @jax.jit
    def step(params, opt_state, key):
        loss, grads = jax.value_and_grad(
            lambda p: -scale * model.log_posterior_from_data(p, data, key=key)
        )(params)
        # Measured before the update, so it is the norm the clip actually saw.
        # Accumulated on device and read once after the loop, to avoid adding
        # a host sync per step.
        clipped = (
            jax.numpy.zeros((), dtype=jax.numpy.int32)
            if clip_at is None
            else (_global_norm(grads) > clip_at).astype(jax.numpy.int32)
        )
        updates, opt_state = self.optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, loss, clipped

    # clear any previous history
    if self.track_history:
        self.loss_steps.clear()
        self.loss_history.clear()
    self.clip_rate = None
    self.n_clipped_steps = 0
    n_clipped = jax.numpy.zeros((), dtype=jax.numpy.int32)
    steps_run = 0

    # Optional progress bar.
    #
    # Why we *manually* advance the bar:
    # - When JAX runs on GPU, the first `step(...)` call can spend a long time in
    #   compilation, and tqdm may not visibly advance if the underlying iterator
    #   doesn't get a chance to redraw.
    # - By keeping a normal `range(self.steps)` loop and calling `pbar.update(1)`
    #   ourselves, we ensure the bar advances exactly once per iteration.
    #
    # Performance note: *displaying the loss* requires transferring `loss` from
    # device -> host, which can add sync overhead. We therefore only attach a
    # loss postfix every `progress_every` steps.
    pbar = None
    if self.show_progress:
        try:
            from tqdm.auto import tqdm

            pbar = tqdm(total=self.steps, desc="MAP fit", leave=False)
        except Exception:
            # Soft dependency: tqdm not available (or terminal unsuitable).
            pbar = None

    for i in range(self.steps):
        # Split a fresh subkey for each step so the MC likelihood sees a
        # different noise realization on every gradient evaluation.
        opt_key, subkey = jax.random.split(opt_key)
        params, opt_state, loss, clipped = step(params, opt_state, subkey)
        n_clipped = n_clipped + clipped
        steps_run = i + 1

        # Non-finite guard: if loss becomes NaN/Inf, optimization has diverged.
        # Stop early so downstream plots don’t look “truncated” due to NaNs.
        if not bool(jax.numpy.isfinite(loss)):
            if self.track_history:
                try:
                    self.loss_steps.append(i)
                    self.loss_history.append(float(loss))
                except Exception:
                    pass
            print(
                f"[MAPOptimizer] Non-finite loss at step {i}: {loss}. "
                "Stopping early."
            )
            if pbar is not None:
                with contextlib.suppress(Exception):
                    pbar.update(1)
            break

        if self.track_history and (
            (i % self.log_every == 0) or (i == self.steps - 1)
        ):
            # Pull scalar to host and record
            try:
                self.loss_steps.append(i)
                self.loss_history.append(float(loss))
            except Exception:
                #  do not break fitting if logging fails
                pass

        #  progress bar loss display (avoid host sync every step)
        if pbar is not None and (
            (i % self.progress_every == 0) or (i == self.steps - 1)
        ):
            with contextlib.suppress(Exception):
                pbar.set_postfix(loss=float(loss))

        if pbar is not None:
            with contextlib.suppress(Exception):
                pbar.update(1)
                # Encourage a redraw occasionally in environments with buffered/stale
                # TTY updates.
                if (i % self.progress_every == 0) or (i == self.steps - 1):
                    pbar.refresh()

    if pbar is not None:
        with contextlib.suppress(Exception):
            pbar.close()

    # Clipping regime. Read once, after the loop, so the counter costs no
    # per-step host sync. A rate near 1.0 means the threshold sits below the
    # model's typical gradient norm, so every step is rescaled to
    # `max_grad_norm` and the effective algorithm is no longer plain SGD.
    if self.max_grad_norm is not None and steps_run > 0:
        self.n_clipped_steps = int(n_clipped)
        self.clip_rate = self.n_clipped_steps / steps_run
        if self.clip_rate > 0.5:
            warnings.warn(
                f"Gradient clipping bound on {self.clip_rate:.0%} of steps "
                f"({self.n_clipped_steps}/{steps_run}) at "
                f"max_grad_norm={self.max_grad_norm}. At this rate clipping "
                "is not a guard rail: the step length is fixed at "
                "learning_rate * max_grad_norm regardless of gradient "
                "magnitude, so `learning_rate` no longer means what it would "
                "without clipping. Raise max_grad_norm above the typical "
                "gradient norm, or pass max_grad_norm=None.",
                stacklevel=2,
            )

    return MAPPosterior(params=params, model=model)

get_history ¶

get_history() -> tuple[list[int], list[float]]

Return (steps, losses) recorded during the last fit when tracking was enabled.

Source code in src/psyphy/inference/map_optimizer.py

def get_history(self) -> tuple[list[int], list[float]]:
    """Return (steps, losses) recorded during the last fit when tracking was enabled."""
    return self.loss_steps, self.loss_history

Langevin Samplers¶

langevin ¶

langevin.py

Langevin samplers for posterior inference.

Will lkely implement: - Overdamped (unadjusted) Langevin Algorithm (ULA) (exists in BlackJax) - Underdamped Langevin (with BAOAB splitting scheme?)

Used for posterior-aware trial placement (InfoGain).

MVP implementation: - Stub that returns an initial Posterior. - Future: implement underdamped Langevin dynamics (e.g. BAOAB integrator).

Classes:

Name	Description
`LangevinSampler`	Langevin sampler (stub).
`NumpyroSampler`	Numpyro-based sampler (stub).

LangevinSampler ¶

LangevinSampler(steps: int = 1000, step_size: float = 0.001, temperature: float = 1.0)

Bases: InferenceEngine

Langevin sampler (stub).

Parameters:

Name	Type	Description	Default
`steps`	`int`	Number of Langevin steps.	`1000`
`step_size`	`float`	Integration step size.	`1e-3`
`temperature`	`float`	Noise scale (temperature).	`1.0`

Methods:

Name	Description
`fit`	Fit model parameters with Langevin dynamics (stub).

Attributes:

Name	Type	Description
`step_size`
`steps`
`temperature`

Source code in src/psyphy/inference/langevin.py

def __init__(
    self, steps: int = 1000, step_size: float = 1e-3, temperature: float = 1.0
):
    self.steps = steps
    self.step_size = step_size
    self.temperature = temperature

step_size ¶

step_size = step_size

steps ¶

steps = steps

temperature ¶

temperature = temperature

fit ¶

fit(model, data) -> MAPPosterior

Fit model parameters with Langevin dynamics (stub).

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Model instance.	required
`data`	`ResponseData`	Observed trials.	required

Source code in src/psyphy/inference/langevin.py

def fit(self, model, data) -> MAPPosterior:
    """
    Fit model parameters with Langevin dynamics (stub).

    Parameters
    ----------
    model : WPPM
        Model instance.
    data : ResponseData
        Observed trials.

    """

    raise NotImplementedError("Langevin Sampler not implemented yet.")

NumpyroSampler ¶

NumpyroSampler(steps: int = 1000)

Bases: InferenceEngine

Numpyro-based sampler (stub).

Parameters:

Name	Type	Description	Default
`steps`	`int`	Number of sampling steps.	`1000`

Methods:

Name	Description
`fit`	Fit model parameters using Numpyro (stub).

Attributes:

Name	Type	Description
`steps`

Source code in src/psyphy/inference/langevin.py

def __init__(self, steps: int = 1000):
    self.steps = steps

steps ¶

steps = steps

fit ¶

fit(model, data) -> MAPPosterior

Fit model parameters using Numpyro (stub).

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Model instance.	required
`data`	`ResponseData`	Observed trials.	required

Returns:

Type	Description
`MAPPosterior`	Posterior distribution containing sample history.

Source code in src/psyphy/inference/langevin.py

def fit(self, model, data) -> MAPPosterior:
    """
    Fit model parameters using Numpyro (stub).

    Parameters
    ----------
    model : WPPM
        Model instance.
    data : ResponseData
        Observed trials.

    Returns
    -------
    MAPPosterior
        Posterior distribution containing sample history.
    """
    raise NotImplementedError("NumpyroSampler.fit not implemented yet.")

Laplace Approximation¶

laplace ¶

laplace.py

Laplace approximation to posterior.

Approximates posterior with a Gaussian: N(mean = MAP, covariance = H^-1 at MAP)

Provides posterior.sample() cheaply. Useful for InfoGainPlacement when only MAP fit is available.

MVP implementation: - Stub that just returns the MAP posterior. - Future: compute covariance from Hessian at MAP params.

Classes:

Name	Description
`LaplaceApproximation`	Laplace approximation around MAP estimate.

LaplaceApproximation ¶

Bases: InferenceEngine

Laplace approximation around MAP estimate.

Methods:

Name	Description
`from_map`	Construct a Gaussian approximation centered at MAP.

fit ¶

fit(model: Any, data: Any) -> Any

Fit model parameters to data.

Parameters:

Name	Type	Description	Default
`model`	`WPPM`	Psychophysical model to fit.	required
`data`	`ResponseData`	Observed trials.	required

Returns:

Type	Description
`Posterior`	Posterior object wrapping fitted params and model reference.

Source code in src/psyphy/inference/base.py

@abstractmethod
def fit(self, model: Any, data: Any) -> Any:
    """
    Fit model parameters to data.

    Parameters
    ----------
    model : WPPM
        Psychophysical model to fit.
    data : ResponseData
        Observed trials.

    Returns
    -------
    Posterior
        Posterior object wrapping fitted params and model reference.
    """
    ...

from_map ¶

from_map(map_posterior: MAPPosterior) -> MAPPosterior

Return posterior approximation from MAP.

Parameters:

Name	Type	Description	Default
`map_posterior`	`Posterior`	Posterior object from MAP optimization.	required

Returns:

Type	Description
`MAPPosterior`	Posterior distribution containing Laplace approximation.

Source code in src/psyphy/inference/laplace.py

def from_map(self, map_posterior: MAPPosterior) -> MAPPosterior:
    """
    Return posterior approximation from MAP.

    Parameters
    ----------
    map_posterior : Posterior
        Posterior object from MAP optimization.

    Returns
    -------
    MAPPosterior
        Posterior distribution containing Laplace approximation.
    """
    # 1. First find MAP estimate
    # 2. Compute Hessian at MAP
    # 3. Invert Hessian to get covariance
    # 4. Return Gaussian posterior with MAP mean and covariance

    return map_posterior