Posterior¶

Package Overview¶

posterior ¶

posterior¶

Posterior representations.

This subpackage provides: - ParameterPosterior: protocol for posteriors over model parameters p(θ | data) - MAPPosterior: delta distribution at θ_MAP (point estimate)

Two-tier design

ParameterPosterior: represents p(θ | data)
PredictivePosterior: represents p(f(X*) | data)

Future extensions

LaplacePosterior: Gaussian approximation N(θ_MAP, Σ)

Classes:

Name	Description
`MAPPosterior`	MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.
`MCMCPosterior`	Backend-agnostic container for MCMC posterior samples.
`ParameterPosterior`	Protocol for posterior distributions over model parameters p(θ \| data).
`PredictivePosterior`	Protocol for predictive distributions p(f(X*) \| data) at test stimuli.
`ThresholdConfig`	Configuration for :class:`WPPMPredictivePosterior`'s `threshold_pred` mode.
`WPPMPredictivePosterior`	Predictive posterior for WPPM models.

MAPPosterior ¶

MAPPosterior(params, model)

MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.

Represents a point estimate with no uncertainty.

Parameters:

Name	Type	Description	Default
`params`	`dict`	MAP parameter dictionary (θ_MAP)	required
`model`	`WPPM`	Model instance used for predictions	required

Notes

This implements the ParameterPosterior protocol.

Methods:

Name	Description
`sample`	Sample from delta distribution (returns repeated θ_MAP).

Attributes:

Name	Type	Description
`model`		Return the associated model.
`params`		Return the MAP parameters (θ_MAP).

Source code in src/psyphy/posterior/posterior.py

def __init__(self, params, model):
    self._params = params
    self._model = model

model ¶

model

Return the associated model.

params ¶

params

Return the MAP parameters (θ_MAP).

sample ¶

sample(n: int = 1, *, key=None)

Sample from delta distribution (returns repeated θ_MAP).

Parameters:

Name	Type	Description	Default
`n`	`int`	Number of samples	`1`
`key`	`Array`	PRNG key (unused for delta distribution)	`None`

Returns:

Type	Description
`dict`	Parameter PyTree with leading dimension n. Each array has shape (n, ...) with identical values.

Notes

Delta distribution has no randomness - returns repeated MAP estimate.

Source code in src/psyphy/posterior/posterior.py

def sample(self, n: int = 1, *, key=None):
    """
    Sample from delta distribution (returns repeated θ_MAP).

    Parameters
    ----------
    n : int, default=1
        Number of samples
    key : jax.Array, optional
        PRNG key (unused for delta distribution)

    Returns
    -------
    dict
        Parameter PyTree with leading dimension n.
        Each array has shape (n, ...) with identical values.

    Notes
    -----
    Delta distribution has no randomness - returns repeated MAP estimate.
    """
    return jax.tree.map(
        lambda x: jnp.tile(x[None, ...], (n,) + (1,) * x.ndim), self._params
    )

MCMCPosterior ¶

MCMCPosterior(samples: dict, model, *, sampler_stats: dict | None = None)

Backend-agnostic container for MCMC posterior samples.

Satisfies the ParameterPosterior protocol: - params: posterior mean over all chains and draws - model: the fitted WPPM instance - sample(n, key): draw n parameter vectors from the chain pool

Parameters:

Name	Type	Description	Default
`samples`	`dict`	Parameter samples. Each value has shape `(n_chains, n_draws, *param_shape)`. For WPPM this is `{"W": ndarray of shape (n_chains, n_draws, d0, d1, d2, d3)}`.	required
`model`	`WPPM`	The model instance used during inference.	required
`sampler_stats`	`dict \| None`	Per-chain diagnostics returned by the sampler, each of shape `(n_chains, n_draws)`. NUTSSampler supplies ArviZ's `sample_stats` names: `acceptance_rate`, `diverging`, `energy`, `n_steps`.	`None`

Notes

Implements the ParameterPosterior protocol (i.e., structural, not inherited).

Methods:

Name	Description
`sample`	Draw n parameter vectors from the pooled chain samples.
`to_arviz`	Convert MCMC samples to an ArviZ InferenceData object.

Attributes:

Name	Type	Description
`model`		The WPPM instance used during inference.
`n_chains`	`int`	Number of chains.
`n_draws`	`int`	Number of draws per chain.
`params`	`dict`	Posterior mean over all chains and draws.
`sampler_stats`

Source code in src/psyphy/posterior/mcmc_posterior.py

def __init__(
    self,
    samples: dict,
    model,
    *,
    sampler_stats: dict | None = None,
):
    # samples: dict of JAX arrays, each shaped (n_chains, n_draws, *param_shape)
    # For WPPM: {"W": (n_chains, n_draws, degree+1, degree+1, input_dim, embedding_dim)}
    # e.g. 4 chains × 1000 draws, degree=1, input_dim=2, extra_dims=1:
    #      {"W": (4, 1000, 2, 2, 2, 3)}
    # This is the canonical shape that flows from NUTSSampler -> here -> ArviZ.
    self._samples = samples
    self._model = model
    # sampler_stats: optional dict of per-chain diagnostics from the sampler.
    # Keys from NUTSSampler: acceptance_rate, diverging, energy, n_steps --
    # each (n_chains, n_draws), named per ArviZ's sample_stats convention so
    # to_arviz() feeds az.plot_energy and divergence reporting directly.
    self.sampler_stats = sampler_stats or {}

model ¶

model

The WPPM instance used during inference.

n_chains ¶

n_chains: int

Number of chains.

n_draws ¶

n_draws: int

Number of draws per chain.

params ¶

params: dict

Posterior mean over all chains and draws.

Returns:

Type	Description
`dict`	Parameter PyTree with the same structure as a single draw. For WPPM: `{"W": ndarray of shape (*W_shape)}`. e.g. {"W": (2, 2, 2, 3)} — mean over axes 0 (chains) and 1 (draws).

sampler_stats ¶

sampler_stats = sampler_stats or {}

sample ¶

sample(n: int, *, key=None) -> dict

Draw n parameter vectors from the pooled chain samples.

Pools all chains and draws into a flat collection of n_chains * n_draws samples, then randomly subsamples n of them (without replacement if n <= total, with replacement otherwise).

Parameters:

Name	Type	Description	Default
`n`	`int`	Number of samples to return.	required
`key`	`Array \| None`	PRNG key. Defaults to `jax.random.PRNGKey(0)` when None.	`None`

Returns:

Type	Description
`dict`	Parameter PyTree with leading dimension `n`. For WPPM: `{"W": ndarray of shape (n, *W_shape)}`.

Source code in src/psyphy/posterior/mcmc_posterior.py

def sample(self, n: int, *, key=None) -> dict:
    """Draw n parameter vectors from the pooled chain samples.

    Pools all chains and draws into a flat collection of
    ``n_chains * n_draws`` samples, then randomly subsamples ``n``
    of them (without replacement if n <= total, with replacement otherwise).

    Parameters
    ----------
    n : int
        Number of samples to return.
    key : jax.Array | None
        PRNG key. Defaults to ``jax.random.PRNGKey(0)`` when None.

    Returns
    -------
    dict
        Parameter PyTree with leading dimension ``n``.
        For WPPM: ``{"W": ndarray of shape (n, *W_shape)}``.
    """
    if key is None:
        key = jr.PRNGKey(0)

    def _subsample(x):
        # x     : (n_chains, n_draws, *param_shape)
        #           e.g. (4, 1000, 2, 2, 2, 3)
        # pool  : (n_chains * n_draws, *param_shape)
        #           e.g. (4000, 2, 2, 2, 3)  — all draws from all chains flattened
        # output: (n, *param_shape)
        #           e.g. (100, 2, 2, 2, 3)   — n draws subsampled from pool
        n_chains, n_draws = x.shape[:2]
        event_shape = x.shape[2:]
        pool = x.reshape(n_chains * n_draws, *event_shape)
        total = n_chains * n_draws
        replace = n > total
        indices = jr.choice(key, total, shape=(n,), replace=replace)
        return pool[indices]

    return jax.tree.map(_subsample, self._samples)

to_arviz ¶

to_arviz()

Convert MCMC samples to an ArviZ InferenceData object.

Requires arviz to be installed::

1	`pip install 'psyphy[diagnostics]'`

Returns:

Type Description

DataTree

ArviZ >= 1.0 returns xarray's DataTree; arviz.InferenceData was the 0.x type and is gone (it is still importable on 1.x but emits a MigrationWarning). Measured: on arviz 0.23.4 every access pattern documented below raises, which is why the extra pins arviz>=1.0. Contains a posterior group with variables shaped (n_chains, n_draws, *param_shape). For WPPM: variable W of shape (n_chains, n_draws, d0, d1, d2, d3).

Notes

ArviZ handles arbitrary-dimensional variables natively — no flattening required. R-hat and ESS are computed per scalar element.

If sampler_stats were provided at construction (e.g. acceptance_rate), they are included in the sample_stats group.

Examples:

>>> idata = nuts_posterior.to_arviz()
>>> import arviz as az
>>> print(az.summary(idata))

Source code in src/psyphy/posterior/mcmc_posterior.py

def to_arviz(self):
    """Convert MCMC samples to an ArviZ InferenceData object.

    Requires `arviz` to be installed::

        pip install 'psyphy[diagnostics]'

    Returns
    -------
    xarray.DataTree
        ArviZ >= 1.0 returns xarray's DataTree; ``arviz.InferenceData`` was
        the 0.x type and is gone (it is still importable on 1.x but emits a
        MigrationWarning). Measured: on arviz 0.23.4 every access pattern
        documented below raises, which is why the extra pins ``arviz>=1.0``.
        Contains a ``posterior`` group with variables shaped
        ``(n_chains, n_draws, *param_shape)``.
        For WPPM: variable ``W`` of shape
        ``(n_chains, n_draws, d0, d1, d2, d3)``.

    Notes
    -----
    ArviZ handles arbitrary-dimensional variables natively — no
    flattening required. R-hat and ESS are computed per scalar element.

    If ``sampler_stats`` were provided at construction (e.g.
    ``acceptance_rate``), they are included in the ``sample_stats``
    group.

    Examples
    --------
    >>> idata = nuts_posterior.to_arviz()
    >>> import arviz as az
    >>> print(az.summary(idata))
    """
    try:
        import arviz as az
    except ImportError as e:
        raise ImportError(
            "ArviZ is required for MCMCPosterior.to_arviz(). "
            "Install it with: pip install 'psyphy[diagnostics]'\n"
            f"Note: to_arviz() needs arviz >= 1.0, and every arviz 1.x "
            f"requires Python >= 3.12 (this is {sys.version_info.major}."
            f"{sys.version_info.minor}). On an older interpreter the "
            "'diagnostics' extra resolves to nothing and this error appears "
            "even after installing it. The chains are still available via "
            "`.sample()` / `._samples` without ArviZ."
        ) from e

    # -- ArviZ 1.0 API boundary --------------------------------
    # az.from_dict(data) where data is a nested dict:
    #   {
    #     "posterior":    {var_name: np.ndarray of shape (n_chains, n_draws, *param_shape)},
    #     "sample_stats": {stat_name: np.ndarray of shape (n_chains, n_draws)},  # optional
    #   }
    #
    # For WPPM with 4 chains, 1000 draws, W shape (2,2,2,3):
    #   {"posterior": {"W": np.ndarray (4, 1000, 2, 2, 2, 3)}}
    #
    # Returns: xarray.DataTree  (ArviZ 1.0 — was InferenceData in ArviZ 0.x)
    #   idata["posterior"]["W"].values  -> np.ndarray (4, 1000, 2, 2, 2, 3)
    #   list(idata.children)            -> ["posterior", "sample_stats"]
    #   az.rhat(idata).ds.data_vars     -> {"W"}  (R-hat per scalar element)
    #
    # Note: must convert JAX arrays to numpy — ArviZ uses numpy/xarray internally.
    # ----------------------------------------------------------------
    data: dict = {"posterior": {k: np.asarray(v) for k, v in self._samples.items()}}

    if self.sampler_stats:
        # sample_stats values shape: (n_chains, n_draws)
        # e.g. {"acceptance_rate": (4, 1000), "diverging": (4, 1000), ...}
        data["sample_stats"] = {
            k: np.asarray(v) for k, v in self.sampler_stats.items()
        }

    return az.from_dict(data)

ParameterPosterior ¶

Bases: Protocol

Protocol for posterior distributions over model parameters p(θ | data).

Returned by InferenceEngine.fit(model, data). Used for diagnostics, parameter sampling, uncertainty.

Methods:

Name	Description
`sample`	Sample parameter vectors from p(θ \| data).

Attributes:

Name	Type	Description
`model`		Associated generative model.
`params`	`dict`	Point estimate or posterior mean parameters.

model ¶

model

Associated generative model.

Returns:

Type	Description
`Model`	The WPPM or (other model) instance used for predictions.

params ¶

params: dict

Point estimate or posterior mean parameters.

Returns:

Type	Description
`dict`	Parameter PyTree (e.g., {"log_diag": jnp.ndarray, ...})

Notes

MAP: θ_MAP
MCMC: posterior mean of samples

sample ¶

sample(n: int, *, key: Array) -> dict

Sample parameter vectors from p(θ | data).

Parameters:

Name	Type	Description	Default
`n`	`int`	Number of samples	required
`key`	`Array`	PRNG key for randomness	required

Returns:

Type	Description
`dict`	Parameter PyTree with leading dimension n.

Notes

MAP: returns repeated θ_MAP
Numpyro/Blackjax: returns stored samples (may subsample if n differs)

Source code in src/psyphy/posterior/parameter_posterior.py

def sample(self, n: int, *, key: jax.Array) -> dict:
    """
    Sample parameter vectors from p(θ | data).

    Parameters
    ----------
    n : int
        Number of samples
    key : jax.Array
        PRNG key for randomness

    Returns
    -------
    dict
        Parameter PyTree with leading dimension n.


    Notes
    -----
    - MAP: returns repeated θ_MAP
    - Numpyro/Blackjax: returns stored samples (may subsample if n differs)
    """
    ...

PredictivePosterior ¶

Bases: Protocol

Protocol for predictive distributions p(f(X*) | data) at test stimuli.

Returned by Model.posterior(X) for use in acquisition functions.

Methods:

Name	Description
`cov_field`	Posterior over perceptual covariance field Σ(X).
`rsample`	Reparameterized samples from p(f(X*) \| data).

Attributes:

Name	Type	Description
`mean`	`ndarray`	Posterior predictive mean E[f(X*) \| data].
`variance`	`ndarray`	Posterior predictive marginal variances Var[f(X*) \| data].

mean ¶

mean: ndarray

Posterior predictive mean E[f(X*) | data].

Returns:

Type	Description
`ndarray`	Shape (n_test,) for scalar outputs Shape (n_test, output_dim) for vector outputs (future)

Notes

Computed via Monte Carlo integration over parameter posterior.

:class:WPPMPredictivePosterior in threshold_pred=True mode is a documented exception to the shapes above: it returns matrix-valued moments, shape (n_test, input_dim, input_dim) -- a threshold covariance per test point, not a scalar or vector. See its own mean/variance docstrings.

variance ¶

variance: ndarray

Posterior predictive marginal variances Var[f(X*) | data].

Returns:

Type	Description
`ndarray`	Shape (n_test,) for scalar outputs Shape (n_test, output_dim) for vector outputs (future)

Notes

Captures both aleatoric (model) and epistemic (parameter) uncertainty.

Same WPPM threshold_pred=True exception as :attr:mean applies.

cov_field ¶

cov_field(X: ndarray) -> ndarray

Posterior over perceptual covariance field Σ(X).

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Test stimuli, shape (n_test, input_dim)	required

Returns:

Type	Description
`ndarray`	Posterior mean covariance E[Σ(X) \| data], shape (n_test, input_dim, input_dim)

Notes

WPPM-specific method for visualizing perceptual noise structure. This is NOT the predictive covariance - it's the model's internal representation of perceptual uncertainty.

Source code in src/psyphy/posterior/predictive_posterior.py

def cov_field(self, X: jnp.ndarray) -> jnp.ndarray:
    """
    Posterior over perceptual covariance field Σ(X).

    Parameters
    ----------
    X : jnp.ndarray
        Test stimuli, shape (n_test, input_dim)

    Returns
    -------
    jnp.ndarray
        Posterior mean covariance E[Σ(X) | data],
        shape (n_test, input_dim, input_dim)

    Notes
    -----
    WPPM-specific method for visualizing perceptual noise structure.
    This is NOT the predictive covariance - it's the model's
    internal representation of perceptual uncertainty.
    """
    ...

rsample ¶

rsample(sample_shape: tuple = (), *, key: Array) -> ndarray

Reparameterized samples from p(f(X*) | data).

Parameters:

Name	Type	Description	Default
`sample_shape`	`tuple`	Shape of sample batch	`()`
`key`	`Array`	PRNG key	required

Returns:

Type	Description
`ndarray`	Shape (sample_shape, n_test) for scalar outputs Shape (sample_shape, n_test, output_dim) for vector outputs

Notes

Enables gradient-based acquisition optimization via reparameterization trick.

Source code in src/psyphy/posterior/predictive_posterior.py

def rsample(self, sample_shape: tuple = (), *, key: jax.Array) -> jnp.ndarray:
    """
    Reparameterized samples from p(f(X*) | data).

    Parameters
    ----------
    sample_shape : tuple, default=()
        Shape of sample batch
    key : jax.Array
        PRNG key

    Returns
    -------
    jnp.ndarray
        Shape (*sample_shape, n_test) for scalar outputs
        Shape (*sample_shape, n_test, output_dim) for vector outputs

    Notes
    -----
    Enables gradient-based acquisition optimization via reparameterization trick.
    """
    ...

ThresholdConfig ¶

ThresholdConfig(n_theta: int = 16, n_length: int = 200, bounds: tuple[float, float] = (0.0005, 0.3), target_pc: float = 2.0 / 3.0, chunk: int = 2000)

Configuration for :class:WPPMPredictivePosterior's threshold_pred mode.

Controls the oddity-task inversion that recovers a threshold covariance Σ_thres(x) — the target_pc-correct discrimination contour — from a fitted noise field. There is no closed form for this inversion (see :class:~psyphy.model.likelihood.OddityTask), so it is done numerically: probe n_theta directions around each reference point, brute-force search n_length candidate distances per direction for the one closest to target_pc, then fit an ellipsoid to the resulting boundary points.

Attributes:

Name	Type	Description
`n_theta`	`int`	Number of directions probed around each reference point. Must be at least `d * (d + 1) // 2` (the number of free parameters of a symmetric `d x d` matrix) or the ellipsoid fit is underdetermined.
`n_length`	`int`	Number of candidate distances probed per direction. Brute-force search rather than bisection: P(correct) is only monotone in distance up to Monte Carlo noise, so a root-finder can walk off a noisy plateau where a dense-grid argmin cannot.
`bounds`	`tuple[float, float]`	`(min, max)` candidate distance, in the model's stimulus units.
`target_pc`	`float`	Target probability-correct defining the threshold (2/3 for the oddity task's standard criterion).
`chunk`	`int`	Batch size for the Monte Carlo sweep, bounding peak memory — `n_theta * n_length` candidate trials are evaluated per (posterior sample, reference point) pair.

Notes

Cost per (posterior sample, reference point) pair is n_theta * n_length calls to the task's MC-simulated predict, each itself an average over OddityTaskConfig.num_samples draws. Total cost for a call to :class:WPPMPredictivePosterior scales as n_samples * n_test * n_theta * n_length * num_samples — a warning is emitted when n_samples * n_test alone exceeds 200.

bounds ¶

bounds: tuple[float, float] = (0.0005, 0.3)

chunk ¶

chunk: int = 2000

n_length ¶

n_length: int = 200

n_theta ¶

n_theta: int = 16

target_pc ¶

target_pc: float = 2.0 / 3.0

WPPMPredictivePosterior ¶

WPPMPredictivePosterior(param_posterior: ParameterPosterior, X: ndarray, n_samples: int = 100, threshold_pred: bool = False, threshold_config: ThresholdConfig | None = None)

Predictive posterior for WPPM models.

Computes p(f(X*) | data) via Monte Carlo integration over parameter posterior p(θ | data).

Parameters:

Name	Type	Description	Default
`param_posterior`	`ParameterPosterior`	Posterior over model parameters	required
`X`	`ndarray`	Test stimuli. Shape depends on `threshold_pred`: `threshold_pred=False`: `(n_test, k_stimuli, input_dim)` — pre-paired (reference, comparison) stimuli, scored directly via `model.predict_prob`. `threshold_pred=True`: `(n_test, input_dim)` — bare reference points. Comparison stimuli are not supplied; the oddity-task inversion generates its own by sweeping directions and distances around each reference (see :class:`ThresholdConfig`). Passing the paired-stimulus shape here raises `ValueError`.	required
`n_samples`	`int`	Number of posterior samples for MC integration	`100`
`threshold_pred`	`bool`	Whether to compute threshold covariances (Σ_thres, the `target_pc`-correct discrimination contour) instead of predicted probabilities. Changes both `X`'s expected shape (above) and `mean`/`variance`'s returned shape (see their docstrings).	`False`
`threshold_config`	`ThresholdConfig`	Settings for the threshold inversion. Only used when `threshold_pred=True`; defaults to `ThresholdConfig()`.	`None`

Attributes:

Name	Type	Description
`param_posterior`	`ParameterPosterior`	Wrapped parameter posterior
`X`	`ndarray`	Test stimuli
`n_samples`	`int`	MC sample count
`threshold_pred`	`bool`	Whether this instance predicts thresholds rather than probabilities
`threshold_config`	`ThresholdConfig or None`	Threshold-inversion settings (`threshold_pred=True` only)

Notes

Uses lazy evaluation: moments computed on first access.

Methods:

Name	Description
`cov_field`	Posterior mean covariance field E[Σ(X) \| data].
`rsample`	Sample predictions from p(f(X*) \| data).

Source code in src/psyphy/posterior/predictive_posterior.py

def __init__(
    self,
    param_posterior: ParameterPosterior,
    X: jnp.ndarray,
    n_samples: int = 100,
    threshold_pred: bool = False,
    threshold_config: ThresholdConfig | None = None,
):
    self.param_posterior = param_posterior
    self.X = X
    self.n_samples = n_samples
    self.threshold_pred = threshold_pred
    self.threshold_config = threshold_config

    # Lazy evaluation cache. Annotated explicitly: without it the
    # initialiser alone makes these `None`, and the `mean`/`variance`
    # properties then look like they return None rather than an array.
    self._mean: jnp.ndarray | None = None
    self._variance: jnp.ndarray | None = None
    self._computed = False

X ¶

X = X

mean ¶

mean: ndarray

E[f(X*) | data].

Shape (n_test,) — predicted probabilities — unless threshold_pred=True, in which case shape is (n_test, input_dim, input_dim): the posterior-mean threshold covariance Σ_thres at each reference point, averaged arithmetically over posterior parameter samples (the PD cone is convex, so this stays PD; not a geometric/Bures mean).

n_samples ¶

n_samples = n_samples

param_posterior ¶

param_posterior = param_posterior

threshold_config ¶

threshold_config = threshold_config

threshold_pred ¶

threshold_pred = threshold_pred

variance ¶

variance: ndarray

Var[f(X*) | data].

Shape (n_test,) unless threshold_pred=True, in which case shape is (n_test, input_dim, input_dim): elementwise variance of the per-sample threshold covariances (not a variance of a matrix in any distributional sense — just jnp.var over the sample axis, entry by entry).

cov_field ¶

cov_field(X: ndarray) -> ndarray

Posterior mean covariance field E[Σ(X) | data].

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Test stimuli, shape (n_test, input_dim) OR (n_test, k_stim, input_dim) Note that standard use for OddityTask is 2D with X as test REFS, not all stimuli.	required

Returns:

Type	Description
`ndarray`	Covariance matrices, shape (n_test, input_dim, input_dim) OR (n_test, k_stim, input_dim, input_dim)

Notes

Averages local_covariance(x) over parameter posterior samples.

Source code in src/psyphy/posterior/predictive_posterior.py

def cov_field(self, X: jnp.ndarray) -> jnp.ndarray:
    """
    Posterior mean covariance field E[Σ(X) | data].

    Parameters
    ----------
    X : jnp.ndarray
        Test stimuli, shape (n_test, input_dim) OR (n_test, k_stim, input_dim)
        Note that standard use for OddityTask is 2D with X as test REFS, not all stimuli.

    Returns
    -------
    jnp.ndarray
        Covariance matrices, shape (n_test, input_dim, input_dim) OR (n_test, k_stim, input_dim, input_dim)

    Notes
    -----
    Averages local_covariance(x) over parameter posterior samples.
    """
    key = jr.PRNGKey(0)
    param_samples = self.param_posterior.sample(self.n_samples, key=key)

    model = self.param_posterior.model

    def cov_at_x(params, x):
        """Evaluate Σ(x) with given parameters."""
        return model.local_covariance(params, x)

    # Vectorized evaluation: (n_samples, n_test, k_stim, input_dim, input_dim)
    if jnp.ndim(X) == 3:  # more than 1 stimulus
        cov_samples = jax.vmap(
            lambda params: jax.vmap(jax.vmap(lambda s: cov_at_x(params, s)))(X)
        )(param_samples)
    elif jnp.ndim(X) == 2:  # only 1 stimulus
        cov_samples = jax.vmap(
            lambda params: jax.vmap(lambda s: cov_at_x(params, s))(X)
        )(param_samples)
    else:
        raise ValueError(
            "Incorrect input dimensionality"
            "Expected 2D (n_test, input_dim) or 3D (n_test, k_stim, input_dim)"
            f"Received {jnp.ndim(X)}D."
        )

    # Return posterior mean
    return jnp.mean(cov_samples, axis=0)

rsample ¶

rsample(sample_shape: tuple = (), *, key: Array) -> ndarray

Sample predictions from p(f(X*) | data).

Parameters:

Name	Type	Description	Default
`sample_shape`	`tuple`	Batch shape	`()`
`key`	`Array`	PRNG key	required

Returns:

Type	Description
`ndarray`	Shape (sample_shape, n_test) — or, when `threshold_pred=True`, (sample_shape, n_test, input_dim, input_dim): one recovered threshold covariance per draw, i.e. genuine samples of Σ_thres rather than a posterior-mean summary. Cost scales linearly in the number of draws requested — see :class:`ThresholdConfig`.

Source code in src/psyphy/posterior/predictive_posterior.py

def rsample(self, sample_shape: tuple = (), *, key: jax.Array) -> jnp.ndarray:
    """
    Sample predictions from p(f(X*) | data).

    Parameters
    ----------
    sample_shape : tuple
        Batch shape
    key : jax.Array
        PRNG key

    Returns
    -------
    jnp.ndarray
        Shape (*sample_shape, n_test) — or, when ``threshold_pred=True``,
        (*sample_shape, n_test, input_dim, input_dim): one recovered
        threshold covariance per draw, i.e. genuine samples of Σ_thres
        rather than a posterior-mean summary. Cost scales linearly in
        the number of draws requested — see :class:`ThresholdConfig`.
    """
    n = int(jnp.prod(jnp.array(sample_shape))) if sample_shape else 1
    param_samples = self.param_posterior.sample(n, key=key)

    model = self.param_posterior.model

    if self.threshold_pred:
        # cov_samples: (n, n_test, input_dim, input_dim)
        cov_samples = self._recover_threshold_covariances(param_samples, self.X)
        if sample_shape:
            return cov_samples.reshape(*sample_shape, *cov_samples.shape[1:])
        return cov_samples

    def predict_one(params):
        """Predict for all test points with given params."""
        return jax.vmap(lambda x: model.predict_prob(params, x))(self.X)

    samples = jax.vmap(predict_one)(param_samples)

    if sample_shape:
        return samples.reshape(*sample_shape, -1)
    return samples

Parameter Posterior Protocol¶

parameter_posterior ¶

parameter_posterior.py

Protocol and implementations for posterior distributions over model parameters.

This module defines the ParameterPosterior interface representing p(θ | data).

Design

Different inference engines produce different posterior representations: - MAP: delta distribution at θ_MAP - MCMC: collection of samples

All implement a common protocol for polymorphic use.

Classes:

Name	Description
`ParameterPosterior`	Protocol for posterior distributions over model parameters p(θ \| data).

ParameterPosterior ¶

Bases: Protocol

Protocol for posterior distributions over model parameters p(θ | data).

Returned by InferenceEngine.fit(model, data). Used for diagnostics, parameter sampling, uncertainty.

Methods:

Name	Description
`sample`	Sample parameter vectors from p(θ \| data).

Attributes:

Name	Type	Description
`model`		Associated generative model.
`params`	`dict`	Point estimate or posterior mean parameters.

model ¶

model

Associated generative model.

Returns:

Type	Description
`Model`	The WPPM or (other model) instance used for predictions.

params ¶

params: dict

Point estimate or posterior mean parameters.

Returns:

Type	Description
`dict`	Parameter PyTree (e.g., {"log_diag": jnp.ndarray, ...})

Notes

MAP: θ_MAP
MCMC: posterior mean of samples

sample ¶

sample(n: int, *, key: Array) -> dict

Sample parameter vectors from p(θ | data).

Parameters:

Name	Type	Description	Default
`n`	`int`	Number of samples	required
`key`	`Array`	PRNG key for randomness	required

Returns:

Type	Description
`dict`	Parameter PyTree with leading dimension n.

Notes

MAP: returns repeated θ_MAP
Numpyro/Blackjax: returns stored samples (may subsample if n differs)

Source code in src/psyphy/posterior/parameter_posterior.py

def sample(self, n: int, *, key: jax.Array) -> dict:
    """
    Sample parameter vectors from p(θ | data).

    Parameters
    ----------
    n : int
        Number of samples
    key : jax.Array
        PRNG key for randomness

    Returns
    -------
    dict
        Parameter PyTree with leading dimension n.


    Notes
    -----
    - MAP: returns repeated θ_MAP
    - Numpyro/Blackjax: returns stored samples (may subsample if n differs)
    """
    ...

Parameter Posterior Implementations¶

posterior ¶

posterior.py

Concrete ParameterPosterior implementations.

This module provides: - MAPPosterior: delta distribution at θ_MAP (point estimate)

Classes:

Name	Description
`MAPPosterior`	MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.

MAPPosterior ¶

MAPPosterior(params, model)

MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.

Represents a point estimate with no uncertainty.

Parameters:

Name	Type	Description	Default
`params`	`dict`	MAP parameter dictionary (θ_MAP)	required
`model`	`WPPM`	Model instance used for predictions	required

Notes

This implements the ParameterPosterior protocol.

Methods:

Name	Description
`sample`	Sample from delta distribution (returns repeated θ_MAP).

Attributes:

Name	Type	Description
`model`		Return the associated model.
`params`		Return the MAP parameters (θ_MAP).

Source code in src/psyphy/posterior/posterior.py

def __init__(self, params, model):
    self._params = params
    self._model = model

model ¶

model

Return the associated model.

params ¶

params

Return the MAP parameters (θ_MAP).

sample ¶

sample(n: int = 1, *, key=None)

Sample from delta distribution (returns repeated θ_MAP).

Parameters:

Name	Type	Description	Default
`n`	`int`	Number of samples	`1`
`key`	`Array`	PRNG key (unused for delta distribution)	`None`

Returns:

Type	Description
`dict`	Parameter PyTree with leading dimension n. Each array has shape (n, ...) with identical values.

Notes

Delta distribution has no randomness - returns repeated MAP estimate.

Source code in src/psyphy/posterior/posterior.py

def sample(self, n: int = 1, *, key=None):
    """
    Sample from delta distribution (returns repeated θ_MAP).

    Parameters
    ----------
    n : int, default=1
        Number of samples
    key : jax.Array, optional
        PRNG key (unused for delta distribution)

    Returns
    -------
    dict
        Parameter PyTree with leading dimension n.
        Each array has shape (n, ...) with identical values.

    Notes
    -----
    Delta distribution has no randomness - returns repeated MAP estimate.
    """
    return jax.tree.map(
        lambda x: jnp.tile(x[None, ...], (n,) + (1,) * x.ndim), self._params
    )

Predictive Posterior¶

predictive_posterior ¶

predictive_posterior.py

Predictive posterior distributions p(f(X*) | data) at test stimuli.

This module defines posteriors over predictions (not parameters), used by acquisition functions for Bayesian optimization.

Design

PredictivePosterior wraps a ParameterPosterior and computes predictions via: E[f(X*) | data] pprox (1/N) Σ_i f(X*; θ_i) where θ_i ~ p(θ | data)

This separates concerns: - ParameterPosterior: represents uncertainty over θ - PredictivePosterior: represents uncertainty over f(X*) (decision-making) - effectively decoupling how we FIT the model from how we USE the fitted model

Classes:

Name	Description
`PredictivePosterior`	Protocol for predictive distributions p(f(X*) \| data) at test stimuli.
`ThresholdConfig`	Configuration for :class:`WPPMPredictivePosterior`'s `threshold_pred` mode.
`WPPMPredictivePosterior`	Predictive posterior for WPPM models.

PredictivePosterior ¶

Bases: Protocol

Protocol for predictive distributions p(f(X*) | data) at test stimuli.

Returned by Model.posterior(X) for use in acquisition functions.

Methods:

Name	Description
`cov_field`	Posterior over perceptual covariance field Σ(X).
`rsample`	Reparameterized samples from p(f(X*) \| data).

Attributes:

Name	Type	Description
`mean`	`ndarray`	Posterior predictive mean E[f(X*) \| data].
`variance`	`ndarray`	Posterior predictive marginal variances Var[f(X*) \| data].

mean ¶

mean: ndarray

Posterior predictive mean E[f(X*) | data].

Returns:

Type	Description
`ndarray`	Shape (n_test,) for scalar outputs Shape (n_test, output_dim) for vector outputs (future)

Notes

Computed via Monte Carlo integration over parameter posterior.

:class:WPPMPredictivePosterior in threshold_pred=True mode is a documented exception to the shapes above: it returns matrix-valued moments, shape (n_test, input_dim, input_dim) -- a threshold covariance per test point, not a scalar or vector. See its own mean/variance docstrings.

variance ¶

variance: ndarray

Posterior predictive marginal variances Var[f(X*) | data].

Returns:

Type	Description
`ndarray`	Shape (n_test,) for scalar outputs Shape (n_test, output_dim) for vector outputs (future)

Notes

Captures both aleatoric (model) and epistemic (parameter) uncertainty.

Same WPPM threshold_pred=True exception as :attr:mean applies.

cov_field ¶

cov_field(X: ndarray) -> ndarray

Posterior over perceptual covariance field Σ(X).

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Test stimuli, shape (n_test, input_dim)	required

Returns:

Type	Description
`ndarray`	Posterior mean covariance E[Σ(X) \| data], shape (n_test, input_dim, input_dim)

Notes

WPPM-specific method for visualizing perceptual noise structure. This is NOT the predictive covariance - it's the model's internal representation of perceptual uncertainty.

Source code in src/psyphy/posterior/predictive_posterior.py

def cov_field(self, X: jnp.ndarray) -> jnp.ndarray:
    """
    Posterior over perceptual covariance field Σ(X).

    Parameters
    ----------
    X : jnp.ndarray
        Test stimuli, shape (n_test, input_dim)

    Returns
    -------
    jnp.ndarray
        Posterior mean covariance E[Σ(X) | data],
        shape (n_test, input_dim, input_dim)

    Notes
    -----
    WPPM-specific method for visualizing perceptual noise structure.
    This is NOT the predictive covariance - it's the model's
    internal representation of perceptual uncertainty.
    """
    ...

rsample ¶

rsample(sample_shape: tuple = (), *, key: Array) -> ndarray

Reparameterized samples from p(f(X*) | data).

Parameters:

Name	Type	Description	Default
`sample_shape`	`tuple`	Shape of sample batch	`()`
`key`	`Array`	PRNG key	required

Returns:

Type	Description
`ndarray`	Shape (sample_shape, n_test) for scalar outputs Shape (sample_shape, n_test, output_dim) for vector outputs

Notes

Enables gradient-based acquisition optimization via reparameterization trick.

Source code in src/psyphy/posterior/predictive_posterior.py

def rsample(self, sample_shape: tuple = (), *, key: jax.Array) -> jnp.ndarray:
    """
    Reparameterized samples from p(f(X*) | data).

    Parameters
    ----------
    sample_shape : tuple, default=()
        Shape of sample batch
    key : jax.Array
        PRNG key

    Returns
    -------
    jnp.ndarray
        Shape (*sample_shape, n_test) for scalar outputs
        Shape (*sample_shape, n_test, output_dim) for vector outputs

    Notes
    -----
    Enables gradient-based acquisition optimization via reparameterization trick.
    """
    ...

ThresholdConfig ¶

ThresholdConfig(n_theta: int = 16, n_length: int = 200, bounds: tuple[float, float] = (0.0005, 0.3), target_pc: float = 2.0 / 3.0, chunk: int = 2000)

Configuration for :class:WPPMPredictivePosterior's threshold_pred mode.

Controls the oddity-task inversion that recovers a threshold covariance Σ_thres(x) — the target_pc-correct discrimination contour — from a fitted noise field. There is no closed form for this inversion (see :class:~psyphy.model.likelihood.OddityTask), so it is done numerically: probe n_theta directions around each reference point, brute-force search n_length candidate distances per direction for the one closest to target_pc, then fit an ellipsoid to the resulting boundary points.

Attributes:

Name	Type	Description
`n_theta`	`int`	Number of directions probed around each reference point. Must be at least `d * (d + 1) // 2` (the number of free parameters of a symmetric `d x d` matrix) or the ellipsoid fit is underdetermined.
`n_length`	`int`	Number of candidate distances probed per direction. Brute-force search rather than bisection: P(correct) is only monotone in distance up to Monte Carlo noise, so a root-finder can walk off a noisy plateau where a dense-grid argmin cannot.
`bounds`	`tuple[float, float]`	`(min, max)` candidate distance, in the model's stimulus units.
`target_pc`	`float`	Target probability-correct defining the threshold (2/3 for the oddity task's standard criterion).
`chunk`	`int`	Batch size for the Monte Carlo sweep, bounding peak memory — `n_theta * n_length` candidate trials are evaluated per (posterior sample, reference point) pair.

Notes

Cost per (posterior sample, reference point) pair is n_theta * n_length calls to the task's MC-simulated predict, each itself an average over OddityTaskConfig.num_samples draws. Total cost for a call to :class:WPPMPredictivePosterior scales as n_samples * n_test * n_theta * n_length * num_samples — a warning is emitted when n_samples * n_test alone exceeds 200.

bounds ¶

bounds: tuple[float, float] = (0.0005, 0.3)

chunk ¶

chunk: int = 2000

n_length ¶

n_length: int = 200

n_theta ¶

n_theta: int = 16

target_pc ¶

target_pc: float = 2.0 / 3.0

WPPMPredictivePosterior ¶

WPPMPredictivePosterior(param_posterior: ParameterPosterior, X: ndarray, n_samples: int = 100, threshold_pred: bool = False, threshold_config: ThresholdConfig | None = None)

Predictive posterior for WPPM models.

Computes p(f(X*) | data) via Monte Carlo integration over parameter posterior p(θ | data).

Parameters:

Name	Type	Description	Default
`param_posterior`	`ParameterPosterior`	Posterior over model parameters	required
`X`	`ndarray`	Test stimuli. Shape depends on `threshold_pred`: `threshold_pred=False`: `(n_test, k_stimuli, input_dim)` — pre-paired (reference, comparison) stimuli, scored directly via `model.predict_prob`. `threshold_pred=True`: `(n_test, input_dim)` — bare reference points. Comparison stimuli are not supplied; the oddity-task inversion generates its own by sweeping directions and distances around each reference (see :class:`ThresholdConfig`). Passing the paired-stimulus shape here raises `ValueError`.	required
`n_samples`	`int`	Number of posterior samples for MC integration	`100`
`threshold_pred`	`bool`	Whether to compute threshold covariances (Σ_thres, the `target_pc`-correct discrimination contour) instead of predicted probabilities. Changes both `X`'s expected shape (above) and `mean`/`variance`'s returned shape (see their docstrings).	`False`
`threshold_config`	`ThresholdConfig`	Settings for the threshold inversion. Only used when `threshold_pred=True`; defaults to `ThresholdConfig()`.	`None`

Attributes:

Name	Type	Description
`param_posterior`	`ParameterPosterior`	Wrapped parameter posterior
`X`	`ndarray`	Test stimuli
`n_samples`	`int`	MC sample count
`threshold_pred`	`bool`	Whether this instance predicts thresholds rather than probabilities
`threshold_config`	`ThresholdConfig or None`	Threshold-inversion settings (`threshold_pred=True` only)

Notes

Uses lazy evaluation: moments computed on first access.

Methods:

Name	Description
`cov_field`	Posterior mean covariance field E[Σ(X) \| data].
`rsample`	Sample predictions from p(f(X*) \| data).

Source code in src/psyphy/posterior/predictive_posterior.py

def __init__(
    self,
    param_posterior: ParameterPosterior,
    X: jnp.ndarray,
    n_samples: int = 100,
    threshold_pred: bool = False,
    threshold_config: ThresholdConfig | None = None,
):
    self.param_posterior = param_posterior
    self.X = X
    self.n_samples = n_samples
    self.threshold_pred = threshold_pred
    self.threshold_config = threshold_config

    # Lazy evaluation cache. Annotated explicitly: without it the
    # initialiser alone makes these `None`, and the `mean`/`variance`
    # properties then look like they return None rather than an array.
    self._mean: jnp.ndarray | None = None
    self._variance: jnp.ndarray | None = None
    self._computed = False

X ¶

X = X

mean ¶

mean: ndarray

E[f(X*) | data].

Shape (n_test,) — predicted probabilities — unless threshold_pred=True, in which case shape is (n_test, input_dim, input_dim): the posterior-mean threshold covariance Σ_thres at each reference point, averaged arithmetically over posterior parameter samples (the PD cone is convex, so this stays PD; not a geometric/Bures mean).

n_samples ¶

n_samples = n_samples

param_posterior ¶

param_posterior = param_posterior

threshold_config ¶

threshold_config = threshold_config

threshold_pred ¶

threshold_pred = threshold_pred

variance ¶

variance: ndarray

Var[f(X*) | data].

Shape (n_test,) unless threshold_pred=True, in which case shape is (n_test, input_dim, input_dim): elementwise variance of the per-sample threshold covariances (not a variance of a matrix in any distributional sense — just jnp.var over the sample axis, entry by entry).

cov_field ¶

cov_field(X: ndarray) -> ndarray

Posterior mean covariance field E[Σ(X) | data].

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Test stimuli, shape (n_test, input_dim) OR (n_test, k_stim, input_dim) Note that standard use for OddityTask is 2D with X as test REFS, not all stimuli.	required

Returns:

Type	Description
`ndarray`	Covariance matrices, shape (n_test, input_dim, input_dim) OR (n_test, k_stim, input_dim, input_dim)

Notes

Averages local_covariance(x) over parameter posterior samples.

Source code in src/psyphy/posterior/predictive_posterior.py

def cov_field(self, X: jnp.ndarray) -> jnp.ndarray:
    """
    Posterior mean covariance field E[Σ(X) | data].

    Parameters
    ----------
    X : jnp.ndarray
        Test stimuli, shape (n_test, input_dim) OR (n_test, k_stim, input_dim)
        Note that standard use for OddityTask is 2D with X as test REFS, not all stimuli.

    Returns
    -------
    jnp.ndarray
        Covariance matrices, shape (n_test, input_dim, input_dim) OR (n_test, k_stim, input_dim, input_dim)

    Notes
    -----
    Averages local_covariance(x) over parameter posterior samples.
    """
    key = jr.PRNGKey(0)
    param_samples = self.param_posterior.sample(self.n_samples, key=key)

    model = self.param_posterior.model

    def cov_at_x(params, x):
        """Evaluate Σ(x) with given parameters."""
        return model.local_covariance(params, x)

    # Vectorized evaluation: (n_samples, n_test, k_stim, input_dim, input_dim)
    if jnp.ndim(X) == 3:  # more than 1 stimulus
        cov_samples = jax.vmap(
            lambda params: jax.vmap(jax.vmap(lambda s: cov_at_x(params, s)))(X)
        )(param_samples)
    elif jnp.ndim(X) == 2:  # only 1 stimulus
        cov_samples = jax.vmap(
            lambda params: jax.vmap(lambda s: cov_at_x(params, s))(X)
        )(param_samples)
    else:
        raise ValueError(
            "Incorrect input dimensionality"
            "Expected 2D (n_test, input_dim) or 3D (n_test, k_stim, input_dim)"
            f"Received {jnp.ndim(X)}D."
        )

    # Return posterior mean
    return jnp.mean(cov_samples, axis=0)

rsample ¶

rsample(sample_shape: tuple = (), *, key: Array) -> ndarray

Sample predictions from p(f(X*) | data).

Parameters:

Name	Type	Description	Default
`sample_shape`	`tuple`	Batch shape	`()`
`key`	`Array`	PRNG key	required

Returns:

Type	Description
`ndarray`	Shape (sample_shape, n_test) — or, when `threshold_pred=True`, (sample_shape, n_test, input_dim, input_dim): one recovered threshold covariance per draw, i.e. genuine samples of Σ_thres rather than a posterior-mean summary. Cost scales linearly in the number of draws requested — see :class:`ThresholdConfig`.

Source code in src/psyphy/posterior/predictive_posterior.py

def rsample(self, sample_shape: tuple = (), *, key: jax.Array) -> jnp.ndarray:
    """
    Sample predictions from p(f(X*) | data).

    Parameters
    ----------
    sample_shape : tuple
        Batch shape
    key : jax.Array
        PRNG key

    Returns
    -------
    jnp.ndarray
        Shape (*sample_shape, n_test) — or, when ``threshold_pred=True``,
        (*sample_shape, n_test, input_dim, input_dim): one recovered
        threshold covariance per draw, i.e. genuine samples of Σ_thres
        rather than a posterior-mean summary. Cost scales linearly in
        the number of draws requested — see :class:`ThresholdConfig`.
    """
    n = int(jnp.prod(jnp.array(sample_shape))) if sample_shape else 1
    param_samples = self.param_posterior.sample(n, key=key)

    model = self.param_posterior.model

    if self.threshold_pred:
        # cov_samples: (n, n_test, input_dim, input_dim)
        cov_samples = self._recover_threshold_covariances(param_samples, self.X)
        if sample_shape:
            return cov_samples.reshape(*sample_shape, *cov_samples.shape[1:])
        return cov_samples

    def predict_one(params):
        """Predict for all test points with given params."""
        return jax.vmap(lambda x: model.predict_prob(params, x))(self.X)

    samples = jax.vmap(predict_one)(param_samples)

    if sample_shape:
        return samples.reshape(*sample_shape, -1)
    return samples