Psyphy

---

psyphy

psyphy

Psychophysical modeling and adaptive trial placement.

This package implements the Wishart Process Psychophysical Model (WPPM) with modular components for priors, task likelihoods, and noise models, which can be fitted to incoming subject data and used to adaptively select new trials to present to the subject next. This is useful for efficiently estimating psychophysical parameters (e.g. threshold contours) with minimal trials.

---

Workflow

Core design

1. WPPM (model/wppm.py):
2. Structural definition of the psychophysical model.
3. Maintains parameterization of local covariance fields.
4. Computes discriminability between stimuli.

5. Delegates trial likelihoods and predictions to the task.

6. Prior (model/prior.py):

7. Defines the distribution over model parameters.

8. WPPM: structured prior over basis weights and decay_rate-controlled covariance fields.

9. TaskLikelihood (model/likelihood.py):

10. Encodes the psychophysical decision rule.

11. WPPM: loglik and predict implemented via Monte Carlo observer simulations, using the noise model explicitly.

12. NoiseModel (model/noise.py):

13. Defines the distribution of internal representation noise.
14. WPPM: GaussianNoise or StudentTNoise option.

Unified import style

Top-level (core models + session): from psyphy import WPPM, Prior, OddityTask, GaussianNoise, MAPOptimizer from psyphy import ExperimentSession, ResponseData, TrialBatch

Subpackages: from psyphy.model import WPPM, Prior, OddityTask, GaussianNoise, StudentTNoise from psyphy.inference import MAPOptimizer, LangevinSampler, LaplaceApproximation from psyphy.acquisition import expected_improvement, upper_confidence_bound, mutual_information from psyphy.acquisition import optimize_acqf, optimize_acqf_discrete, optimize_acqf_random from psyphy.trial_placement import GridPlacement, SobolPlacement from psyphy.utils import grid_candidates, sobol_candidates, custom_candidates, chebyshev_basis

Data flow

• A ResponseData object (psyphy.data) contains trial stimuli and responses.
• WPPM.init_params(prior) samples parameter initialization.
• Inference engines optimize the log posterior: log_posterior = task.loglik(params, data, model=WPPM, noise=NoiseModel) + prior.log_prob(params)
• Posterior predictions (p(correct), threshold ellipses) are always obtained through WPPM delegating to TaskLikelihood.

Extensibility

• To add a new task: subclass TaskLikelihood, implement predict/loglik.
• To add a new noise model: subclass NoiseModel, implement logpdf/sample.

---

Classes:

Name	Description
GaussianNoise
LangevinSampler	Langevin sampler (stub).
LaplaceApproximation	Laplace approximation around MAP estimate.
MAPOptimizer	MAP (Maximum A Posteriori) optimizer.
OddityTask	Three-alternative forced-choice oddity task (MC-based only).
OddityTaskConfig	Configuration for :class:OddityTask.
Prior	Prior distribution over WPPM parameters
ResponseData	Python-friendly incremental trial log.
StudentTNoise
TrialBatch	Container for a proposed batch of trials
WPPM	Wishart Process Psychophysical Model (WPPM).

GaussianNoise

GaussianNoise(sigma: float = 1.0)

Methods:

Name	Description
log_prob
sample_standard	Sample from standard Gaussian (mean=0, var=1).

Attributes:

Name	Type	Description
sigma	float

sigma

sigma: float = 1.0

log_prob

log_prob(residual: float) -> float

Source code in src/psyphy/model/noise.py

def log_prob(self, residual: float) -> float:
    _ = residual
    return -0.5

sample_standard

sample_standard(key: Array, shape: tuple[int, ...]) -> Array

Sample from standard Gaussian (mean=0, var=1).

Source code in src/psyphy/model/noise.py

def sample_standard(self, key: jax.Array, shape: tuple[int, ...]) -> jax.Array:
    """Sample from standard Gaussian (mean=0, var=1)."""
    return jr.normal(key, shape)

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, *, track_history: bool = True, log_every: int = 1, progress_every: int = 10, show_progress: bool = False, max_grad_norm: float | None = 1.0)

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.

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
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: if tqdm 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. This stabilizes optimization when gradients blow up.	1.0

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
log_every
loss_history	list[float]
loss_steps	list[int]
max_grad_norm
optimizer
progress_every
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,
    *,
    track_history: bool = True,
    log_every: int = 1,
    progress_every: int = 10,
    show_progress: bool = False,
    max_grad_norm: float | None = 1.0,
):
    """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).
    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: if tqdm 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. This stabilizes optimization when gradients blow up.
    """
    self.steps = steps
    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] = []

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

optimizer

optimizer = base_optimizer

progress_every

progress_every = max(1, int(progress_every))

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)

    # key is now an explicit argument so each JIT-compiled call receives a
    # distinct random key — a fresh MC noise realization per gradient step.
    @jax.jit
    def step(params, opt_state, key):
        loss, grads = jax.value_and_grad(
            lambda p: -model.log_posterior_from_data(p, data, key=key)
        )(params)
        updates, opt_state = self.optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, loss

    # clear any previous history
    if self.track_history:
        self.loss_steps.clear()
        self.loss_history.clear()

    # 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 = step(params, opt_state, subkey)

        # 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()

    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

OddityTask

OddityTask(config: OddityTaskConfig | None = None)

Bases: TaskLikelihood

Three-alternative forced-choice oddity task (MC-based only).

Implements the full 3-stimulus oddity task using Monte Carlo simulation: - Samples three internal representations per trial (z0, z1, z2) - Uses proper oddity decision rule with three pairwise distances - Suitable for complex covariance structures

Notes

MC simulation in loglik() (full 3-stimulus oddity): 1. Sample three internal representations: z_ref, z_refprime ~ N(ref, Σ_ref), z_comparison ~ N(comparison, Σ_comparison) 2. Compute average covariance: Σ_avg = (2/3) Σ_ref + (1/3) Σ_comparison 3. Compute three pairwise Mahalanobis distances: - d^2(z_ref, z_refprime) = distance between two reference samples - d^2(z_ref, z_comparison) = distance from ref to comparison - d^2(z_refprime, z_comparison) = distance from reference_prime to comparison 4. Apply oddity decision rule: delta = min(d^2(z_ref,z_comparison), d^2(z_refprime,z_comparison)) - d^2(z_ref,z_refprime) 5. Logistic smoothing: P(correct) pprox logistic.cdf(delta / bandwidth) 6. Average over samples

Examples:

>>> from psyphy.model.likelihood import OddityTask
>>> from psyphy.model.likelihood import OddityTaskConfig
>>> from psyphy.model import WPPM, Prior
>>> from psyphy.model.noise import GaussianNoise
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>>
>>> # Create task and model (task-owned MC controls)
>>> likelihood = OddityTask(
...     config=OddityTaskConfig(num_samples=1000, bandwidth=1e-2)
... )
>>> model = WPPM(
...     input_dim=2,
...     prior=Prior(input_dim=2),
...     likelihood=task,
...     noise=GaussianNoise(),
... )
>>> params = model.init_params(jr.PRNGKey(0))

>>> # MC simulation
>>> from psyphy.data.dataset import ResponseData
>>> data = ResponseData()
>>> data.add_trial(ref, comparison, resp=1)
>>> ll_mc = likelihood.loglik(params, data, model, key=jr.PRNGKey(42))
>>> print(f"Log-likelihood (MC): {ll_mc:.4f}")

Methods:

Name	Description
loglik	Compute Bernoulli log-likelihood over a batch of trials.
predict	Return p(correct) for a single (ref, comparison) trial via MC simulation.
simulate	Simulate observed binary responses for a batch of trials.

Attributes:

Name	Type	Description
config

Source code in src/psyphy/model/likelihood.py

def __init__(self, config: OddityTaskConfig | None = None) -> None:
    # No analytical parameters in MC-only mode.
    self.config = config or OddityTaskConfig()

config

config = config or OddityTaskConfig()

loglik

loglik(params: Any, data: Any, model: Any, *, key: Any = None) -> ndarray

Compute Bernoulli log-likelihood over a batch of trials.

This is a concrete base-class method: it vmaps predict over trials then applies the Bernoulli log-likelihood formula. Subclasses only need to implement predict.

Parameters:

Name	Type	Description	Default
params	Any	Model parameters.	required
data	Any	Object with .refs, .comparisons, .responses array attributes.	required
model	Any	Model instance.	required
key	KeyArray	PRNG key. Passed as independent per-trial subkeys to predict. When None, falls back to key=jr.PRNGKey(0) (deterministic).	None

Returns:

Type	Description
ndarray	Scalar sum of Bernoulli log-likelihoods over all trials.

Source code in src/psyphy/model/likelihood.py

def loglik(
    self,
    params: Any,
    data: Any,
    model: Any,
    *,
    key: Any = None,
) -> jnp.ndarray:
    """Compute Bernoulli log-likelihood over a batch of trials.

    This is a concrete base-class method: it vmaps ``predict`` over trials
    then applies the Bernoulli log-likelihood formula. Subclasses only need
    to implement ``predict``.

    Parameters
    ----------
    params : Any
        Model parameters.
    data : Any
        Object with ``.refs``, ``.comparisons``, ``.responses`` array attributes.
    model : Any
        Model instance.
    key : jax.random.KeyArray, optional
        PRNG key. Passed as independent per-trial subkeys to ``predict``.
        When None, falls back to ``key=jr.PRNGKey(0)`` (deterministic).

    Returns
    -------
    jnp.ndarray
        Scalar sum of Bernoulli log-likelihoods over all trials.
    """
    refs = jnp.asarray(data.refs)
    comparisons = jnp.asarray(data.comparisons)
    responses = jnp.asarray(data.responses)
    n_trials = int(refs.shape[0])

    base_key = key if key is not None else jr.PRNGKey(0)
    trial_keys = jr.split(base_key, n_trials)

    probs = jax.vmap(
        lambda ref, comparison, k: self.predict(
            params, ref, comparison, model, key=k
        )
    )(refs, comparisons, trial_keys)

    log_likelihoods = jnp.where(
        responses == 1,
        jnp.log(probs),
        jnp.log(1.0 - probs),
    )
    return jnp.sum(log_likelihoods)

predict

predict(params: Any, ref: ndarray, comparison: ndarray, model: Any, *, key: Any = None) -> ndarray

Return p(correct) for a single (ref, comparison) trial via MC simulation.

MC controls (num_samples, bandwidth) are read from :class:OddityTaskConfig. Pass key to control randomness; when None, config.default_key_seed is used.

Source code in src/psyphy/model/likelihood.py

def predict(
    self,
    params: Any,
    ref: jnp.ndarray,
    comparison: jnp.ndarray,
    model: Any,
    *,
    key: Any = None,
) -> jnp.ndarray:
    """Return p(correct) for a single (ref, comparison) trial via MC simulation.

    MC controls (``num_samples``, ``bandwidth``) are read from
    :class:`OddityTaskConfig`. Pass ``key`` to control randomness; when
    None, ``config.default_key_seed`` is used.
    """
    num_samples = int(self.config.num_samples)
    bandwidth = float(self.config.bandwidth)
    if key is None:
        key = jr.PRNGKey(int(self.config.default_key_seed))

    return self._simulate_trial_mc(
        params=params,
        ref=ref,
        comparison=comparison,
        model=model,
        num_samples=num_samples,
        bandwidth=bandwidth,
        key=key,
    )

simulate

simulate(params: Any, refs: ndarray, comparisons: ndarray, model: Any, *, key: Any) -> tuple[ndarray, ndarray]

Simulate observed binary responses for a batch of trials.

Parameters:

Name	Type	Description	Default
params	Any	Model parameters.	required
refs	(ndarray, shape(n_trials, input_dim))	Reference stimuli.	required
comparisons	(ndarray, shape(n_trials, input_dim))	Comparison stimuli.	required
model	Any	Model instance.	required
key	KeyArray	PRNG key (required; split internally for prediction and sampling).	required

Returns:

Name	Type	Description
responses	jnp.ndarray, shape (n_trials,), dtype int32	Simulated binary responses (1 = correct, 0 = incorrect).
p_correct	(ndarray, shape(n_trials))	Estimated P(correct) per trial used to draw the responses.

Source code in src/psyphy/model/likelihood.py

def simulate(
    self,
    params: Any,
    refs: jnp.ndarray,
    comparisons: jnp.ndarray,
    model: Any,
    *,
    key: Any,
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Simulate observed binary responses for a batch of trials.

    Parameters
    ----------
    params : Any
        Model parameters.
    refs : jnp.ndarray, shape (n_trials, input_dim)
        Reference stimuli.
    comparisons : jnp.ndarray, shape (n_trials, input_dim)
        Comparison stimuli.
    model : Any
        Model instance.
    key : jax.random.KeyArray
        PRNG key (required; split internally for prediction and sampling).

    Returns
    -------
    responses : jnp.ndarray, shape (n_trials,), dtype int32
        Simulated binary responses (1 = correct, 0 = incorrect).
    p_correct : jnp.ndarray, shape (n_trials,)
        Estimated P(correct) per trial used to draw the responses.
    """
    refs = jnp.asarray(refs)
    comparisons = jnp.asarray(comparisons)
    n_trials = int(refs.shape[0])

    k_pred, k_bernoulli = jr.split(key)
    trial_keys = jr.split(k_pred, n_trials)

    p_correct = jax.vmap(
        lambda ref, comparison, k: self.predict(
            params, ref, comparison, model, key=k
        )
    )(refs, comparisons, trial_keys)

    responses = jr.bernoulli(k_bernoulli, p_correct).astype(jnp.int32)
    return responses, p_correct

OddityTaskConfig

OddityTaskConfig(num_samples: int = 1000, bandwidth: float = 0.01, default_key_seed: int = 0)

Configuration for :class:OddityTask.

This is the single source of truth for MC likelihood controls.

Attributes:

Name	Type	Description
num_samples	int	Number of Monte Carlo samples per trial.
bandwidth	float	Logistic CDF smoothing bandwidth.
default_key_seed	int	Seed used when no key is provided (keeps behavior deterministic by default while allowing reproducibility control upstream).

bandwidth

bandwidth: float = 0.01

default_key_seed

default_key_seed: int = 0

num_samples

num_samples: int = 1000

Prior

Prior(input_dim: int = 2, basis_degree: int = 4, variance_scale: float = 0.004, decay_rate: float = 0.4, extra_embedding_dims: int = 1)

Prior distribution over WPPM parameters

Parameters:

Name	Type	Description	Default
input_dim	int	Dimensionality of the model space (same as WPPM.input_dim)	2
basis_degree	int | None	Degree of Chebyshev basis for Wishart process. If set, uses Wishart mode with W coefficients.	None
variance_scale	float	Prior variance for degree-0 (constant) coefficient in Wishart mode. Controls overall scale of covariances.	1.0
decay_rate	float	Geometric decay rate for prior variance over higher-degree coefficients. Prior variance for degree-d coefficient = variance_scale * (decay_rate^d). Smaller decay_rate -> stronger smoothness prior.	0.5
extra_embedding_dims	int	Additional latent dimensions in U matrices beyond input dimensions. Allows richer ellipsoid shapes in Wishart mode.	0

Methods:

Name	Description
log_prob	Compute log prior density (up to a constant)
sample_params	Sample initial parameters from the prior.

Attributes:

Name	Type	Description
basis_degree	int
decay_rate	float
extra_embedding_dims	int
input_dim	int
variance_scale	float

basis_degree

basis_degree: int = 4

decay_rate

decay_rate: float = 0.4

extra_embedding_dims

extra_embedding_dims: int = 1

input_dim

input_dim: int = 2

variance_scale

variance_scale: float = 0.004

log_prob

log_prob(params: Params) -> ndarray

Compute log prior density (up to a constant)

Gaussian prior on W with smoothness via decay_rate log p(W) = Σ_ij log N(W_ij | 0, σ_ij^2) where σ_ij^2 = prior variance

Parameters:

Name	Type	Description	Default
params	dict	Parameter dictionary	required

Returns:

Name	Type	Description
log_prob	float	Log prior probability (up to normalizing constant)

Source code in src/psyphy/model/prior.py

def log_prob(self, params: Params) -> jnp.ndarray:
    """
    Compute log prior density (up to a constant)

    Gaussian prior on W with smoothness via decay_rate
        log p(W) = Σ_ij log N(W_ij | 0, σ_ij^2) where σ_ij^2 = prior variance

    Parameters
    ----------
    params : dict
        Parameter dictionary

    Returns
    -------
    log_prob : float
        Log prior probability (up to normalizing constant)
    """

    if "W" in params:
        # Wishart mode
        W = params["W"]
        variances = self._compute_W_prior_variances()

        # Gaussian log probability for each entry
        # log N(x | 0, σ^2) = -0.5 * (x^2/σ^2 + log(2πσ^2))
        # Up to constant: -0.5 * x^2/σ^2

        if self.input_dim == 2:
            # Each W[i,j,:,:] ~ Normal(0, variance[i,j] * I)
            return -0.5 * jnp.sum((W**2) / (variances[:, :, None, None] + 1e-10))
        elif self.input_dim == 3:
            return -0.5 * jnp.sum((W**2) / (variances[:, :, :, None, None] + 1e-10))

    raise ValueError("params must contain weights 'W'")

sample_params

sample_params(key: Any) -> Params

Sample initial parameters from the prior.

Returns {"W": shape (degree+1, degree+1, input_dim, embedding_dim)} for 2D, where embedding_dim = input_dim + extra_embedding_dims

Note: The 3rd dimension is input_dim (output space dimension). This matches the einsum in _compute_sqrt: U = einsum("ijde,ij->de", W, phi) where d indexes input_dim.

Parameters:

Name	Type	Description	Default
key	JAX random key		required

Returns:

Name	Type	Description
params	dict	Parameter dictionary

Source code in src/psyphy/model/prior.py

def sample_params(self, key: Any) -> Params:
    """
    Sample initial parameters from the prior.


    Returns {"W": shape (degree+1, degree+1, input_dim, embedding_dim)}
    for 2D, where embedding_dim = input_dim + extra_embedding_dims

    Note: The 3rd dimension is input_dim (output space dimension).
    This matches the einsum in _compute_sqrt:
    U = einsum("ijde,ij->de", W, phi) where d indexes input_dim.

    Parameters
    ----------
    key : JAX random key

    Returns
    -------
    params : dict
        Parameter dictionary
    """
    if self.basis_degree is None:
        raise ValueError(
            "'basis_degree' is None; please set "
            "`Prior.basis_degree` to an integer >0."
        )

    # Basis function coefficients W
    variances = self._compute_W_prior_variances()
    embedding_dim = self.input_dim + self.extra_embedding_dims

    if self.input_dim == 2:
        # Sample W ~ Normal(0, variances) for each matrix entry
        # Shape: (degree+1, degree+1, input_dim, embedding_dim)
        # Note: degree+1 to match number of basis functions [T_0, ..., T_degree]
        W = jnp.sqrt(variances)[:, :, None, None] * jr.normal(
            key,
            shape=(
                self.basis_degree + 1,
                self.basis_degree + 1,
                self.input_dim,
                embedding_dim,
            ),
        )
    elif self.input_dim == 3:
        # Shape: (degree+1, degree+1, degree+1, input_dim, embedding_dim)
        W = jnp.sqrt(variances)[:, :, :, None, None] * jr.normal(
            key,
            shape=(
                self.basis_degree + 1,
                self.basis_degree + 1,
                self.basis_degree + 1,
                self.input_dim,
                embedding_dim,
            ),
        )
    else:
        raise NotImplementedError(
            f"Wishart process only supports 2D and 3D. Got input_dim={self.input_dim}"
        )

    return {"W": W}

ResponseData

ResponseData()

Python-friendly incremental trial log.

This container is convenient for adaptive trial placement and I/O (e.g., CSV), but it is not a compute-efficient representation for JAX.

Use :class:TrialData for model fitting and likelihood evaluation.

Methods:

Name	Description
add_batch	Append responses for a batch of trials.
add_trial	append a single trial.
copy	Create a deep copy of this dataset.
from_arrays	Construct ResponseData from arrays.
from_trial_data	Build a ResponseData log from a :class:TrialData batch.
merge	Merge another dataset into this one (in-place).
tail	Return last n trials as a new ResponseData.
to_numpy	Return refs, comparisons, responses as NumPy arrays.
to_trial_data	Convert this log into the canonical JAX batch (:class:TrialData).

Attributes:

Name	Type	Description
comparisons	list[Any]
refs	list[Any]
responses	list[int]
trials	list[tuple[Any, Any, int]]	Return list of (ref, comparison, response) tuples.

Source code in src/psyphy/data/dataset.py

def __init__(self) -> None:
    self.refs: list[Any] = []
    self.comparisons: list[Any] = []
    self.responses: list[int] = []

comparisons

comparisons: list[Any] = []

refs

refs: list[Any] = []

responses

responses: list[int] = []

trials

trials: list[tuple[Any, Any, int]]

Return list of (ref, comparison, response) tuples.

Returns:

Type	Description
list[tuple]	Each element is (ref, comparison, resp)

add_batch

add_batch(responses: list[int], trial_batch: TrialBatch) -> None

Append responses for a batch of trials.

Parameters:

Name	Type	Description	Default
responses	List[int]	Responses corresponding to each (ref, comparison) in the trial batch.	required
trial_batch	TrialBatch	The batch of proposed trials.	required

Source code in src/psyphy/data/dataset.py

def add_batch(self, responses: list[int], trial_batch: TrialBatch) -> None:
    """
    Append responses for a batch of trials.

    Parameters
    ----------
    responses : List[int]
        Responses corresponding to each (ref, comparison) in the trial batch.
    trial_batch : TrialBatch
        The batch of proposed trials.
    """
    for (ref, comparison), resp in zip(trial_batch.stimuli, responses):
        self.add_trial(ref, comparison, resp)

add_trial

add_trial(ref: Any, comparison: Any, resp: int) -> None

append a single trial.

Parameters:

Name	Type	Description	Default
ref	Any	Reference stimulus (numpy array, list, etc.)	required
comparison	Any	Probe stimulus	required
resp	int	Subject response (binary or categorical)	required

Source code in src/psyphy/data/dataset.py

def add_trial(self, ref: Any, comparison: Any, resp: int) -> None:
    """
    append a single trial.

    Parameters
    ----------
    ref : Any
        Reference stimulus (numpy array, list, etc.)
    comparison : Any
        Probe stimulus
    resp : int
        Subject response (binary or categorical)
    """
    self.refs.append(ref)
    self.comparisons.append(comparison)
    self.responses.append(resp)

copy

copy() -> ResponseData

Create a deep copy of this dataset.

Returns:

Type	Description
ResponseData	New dataset with copied data

Source code in src/psyphy/data/dataset.py

def copy(self) -> ResponseData:
    """
    Create a deep copy of this dataset.

    Returns
    -------
    ResponseData
        New dataset with copied data
    """
    new_data = ResponseData()
    new_data.refs = list(self.refs)
    new_data.comparisons = list(self.comparisons)
    new_data.responses = list(self.responses)
    return new_data

from_arrays

from_arrays(X: ndarray | ndarray, y: ndarray | ndarray, *, comparisons: ndarray | ndarray | None = None) -> ResponseData

Construct ResponseData from arrays.

Parameters:

Name	Type	Description	Default
X	(array, shape(n_trials, 2, input_dim) or (n_trials, input_dim))	Stimuli. If 3D, second axis is [reference, comparison]. If 2D, comparisons must be provided separately.	required
y	(array, shape(n_trials))	Responses	required
comparisons	(array, shape(n_trials, input_dim))	Probe stimuli. Only needed if X is 2D.	None

Returns:

Type	Description
ResponseData	Data container

Examples:

>>> # From paired stimuli
>>> X = jnp.array([[[0, 0], [1, 0]], [[1, 1], [2, 1]]])
>>> y = jnp.array([1, 0])
>>> data = ResponseData.from_arrays(X, y)

>>> # From separate refs and comparisons
>>> refs = jnp.array([[0, 0], [1, 1]])
>>> comparisons = jnp.array([[1, 0], [2, 1]])
>>> data = ResponseData.from_arrays(refs, y, comparisons=comparisons)

Source code in src/psyphy/data/dataset.py

@classmethod
def from_arrays(
    cls,
    X: jnp.ndarray | np.ndarray,
    y: jnp.ndarray | np.ndarray,
    *,
    comparisons: jnp.ndarray | np.ndarray | None = None,
) -> ResponseData:
    """
    Construct ResponseData from arrays.

    Parameters
    ----------
    X : array, shape (n_trials, 2, input_dim) or (n_trials, input_dim)
        Stimuli. If 3D, second axis is [reference, comparison].
        If 2D, comparisons must be provided separately.
    y : array, shape (n_trials,)
        Responses
    comparisons : array, shape (n_trials, input_dim), optional
        Probe stimuli. Only needed if X is 2D.

    Returns
    -------
    ResponseData
        Data container

    Examples
    --------
    >>> # From paired stimuli
    >>> X = jnp.array([[[0, 0], [1, 0]], [[1, 1], [2, 1]]])
    >>> y = jnp.array([1, 0])
    >>> data = ResponseData.from_arrays(X, y)

    >>> # From separate refs and comparisons
    >>> refs = jnp.array([[0, 0], [1, 1]])
    >>> comparisons = jnp.array([[1, 0], [2, 1]])
    >>> data = ResponseData.from_arrays(refs, y, comparisons=comparisons)
    """
    data = cls()

    X = np.asarray(X)
    y = np.asarray(y)

    if X.ndim == 3:
        # X is (n_trials, 2, input_dim)
        refs = X[:, 0, :]
        comparisons_arr = X[:, 1, :]
    elif X.ndim == 2 and comparisons is not None:
        refs = X
        comparisons_arr = np.asarray(comparisons)
    else:
        raise ValueError(
            "X must be shape (n_trials, 2, input_dim) or "
            "(n_trials, input_dim) with comparisons argument"
        )

    for ref, comparison, response in zip(refs, comparisons_arr, y):
        data.add_trial(ref, comparison, int(response))

    return data

from_trial_data

from_trial_data(data: TrialData) -> ResponseData

Build a ResponseData log from a :class:TrialData batch.

Source code in src/psyphy/data/dataset.py

@classmethod
def from_trial_data(cls, data: TrialData) -> ResponseData:
    """Build a ResponseData log from a :class:`TrialData` batch."""
    refs = np.asarray(data.refs)
    comps = np.asarray(data.comparisons)
    ys = np.asarray(data.responses)
    out = cls()
    for r, c, y in zip(refs, comps, ys):
        out.add_trial(r, c, int(y))
    return out

merge

merge(other: ResponseData) -> None

Merge another dataset into this one (in-place).

Parameters:

Name	Type	Description	Default
other	ResponseData	Dataset to merge	required

Source code in src/psyphy/data/dataset.py

def merge(self, other: ResponseData) -> None:
    """
    Merge another dataset into this one (in-place).

    Parameters
    ----------
    other : ResponseData
        Dataset to merge
    """
    self.refs.extend(other.refs)
    self.comparisons.extend(other.comparisons)
    self.responses.extend(other.responses)

tail

tail(n: int) -> ResponseData

Return last n trials as a new ResponseData.

Parameters:

Name	Type	Description	Default
n	int	Number of trials to keep	required

Returns:

Type	Description
ResponseData	New dataset with last n trials

Source code in src/psyphy/data/dataset.py

def tail(self, n: int) -> ResponseData:
    """
    Return last n trials as a new ResponseData.

    Parameters
    ----------
    n : int
        Number of trials to keep

    Returns
    -------
    ResponseData
        New dataset with last n trials
    """
    new_data = ResponseData()
    new_data.refs = self.refs[-n:]
    new_data.comparisons = self.comparisons[-n:]
    new_data.responses = self.responses[-n:]
    return new_data

to_numpy

to_numpy() -> tuple[ndarray, ndarray, ndarray]

Return refs, comparisons, responses as NumPy arrays.

Source code in src/psyphy/data/dataset.py

def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return refs, comparisons, responses as NumPy arrays."""
    return (
        np.asarray(self.refs),
        np.asarray(self.comparisons),
        np.asarray(self.responses),
    )

to_trial_data

to_trial_data() -> TrialData

Convert this log into the canonical JAX batch (:class:TrialData).

Source code in src/psyphy/data/dataset.py

def to_trial_data(self) -> TrialData:
    """Convert this log into the canonical JAX batch (:class:`TrialData`)."""
    refs, comparisons, responses = self.to_numpy()
    return TrialData(
        refs=jnp.asarray(refs),
        comparisons=jnp.asarray(comparisons),
        responses=jnp.asarray(responses),
    )

StudentTNoise

StudentTNoise(df: float = 3.0, scale: float = 1.0)

Methods:

Name	Description
log_prob
sample_standard	Sample from standard Student-t (df=self.df).

Attributes:

Name	Type	Description
df	float
scale	float

df

df: float = 3.0

scale

scale: float = 1.0

log_prob

log_prob(residual: float) -> float

Source code in src/psyphy/model/noise.py

def log_prob(self, residual: float) -> float:
    _ = residual
    return -0.5

sample_standard

sample_standard(key: Array, shape: tuple[int, ...]) -> Array

Sample from standard Student-t (df=self.df).

Source code in src/psyphy/model/noise.py

def sample_standard(self, key: jax.Array, shape: tuple[int, ...]) -> jax.Array:
    """Sample from standard Student-t (df=self.df)."""
    return jr.t(key, self.df, shape)

TrialBatch

TrialBatch(stimuli: list[tuple[Any, Any]])

Container for a proposed batch of trials

Attributes:

Name	Type	Description
stimuli	List[Tuple[Any, Any]]	Each trial is a (reference, comparison) tuple.

Methods:

Name	Description
from_stimuli	Construct a TrialBatch from a list of stimuli (ref, comparison) pairs.

Source code in src/psyphy/data/dataset.py

def __init__(self, stimuli: list[tuple[Any, Any]]) -> None:
    self.stimuli = list(stimuli)

stimuli

stimuli = list(stimuli)

from_stimuli

from_stimuli(pairs: list[tuple[Any, Any]]) -> TrialBatch

Construct a TrialBatch from a list of stimuli (ref, comparison) pairs.

Source code in src/psyphy/data/dataset.py

@classmethod
def from_stimuli(cls, pairs: list[tuple[Any, Any]]) -> TrialBatch:
    """
    Construct a TrialBatch from a list of stimuli (ref, comparison) pairs.
    """
    return cls(pairs)

WPPM

WPPM(prior: Prior, likelihood: TaskLikelihood, noise: Any | None = None, *, input_dim: int = 2, extra_dims: int = 1, variance_scale: float = 0.004, decay_rate: float = 0.4, diag_term: float = 1e-06, **model_kwargs: Any)

Bases: Model

Wishart Process Psychophysical Model (WPPM).

Parameters:

Name	Type	Description	Default
input_dim	int	Dimensionality of the input stimulus space (e.g., 2 for isoluminant plane, 3 for RGB). Both reference and comparison live in R^{input_dim}.	2
prior	Prior	Prior distribution over model parameters. Controls basis_degree in WPPM (basis expansion). The WPPM delegates to prior.basis_degree to ensure consistency between parameter sampling and basis evaluation.	required
likelihood	TaskLikelihood	Psychophysical task mapping that defines how discriminability translates to p(correct) and how log-likelihood of responses is computed. (e.g., OddityTask)	required
noise	Any	Noise model describing internal representation noise (e.g., GaussianNoise).	None

hyperparameters

extra_dims : int, default=0 Additional embedding dimensions for basis expansions (beyond input_dim). embedding_dim = input_dim + extra_dims. variance_scale : float, default=1.0 Global scaling factor for covariance magnitude decay_rate : float, default=1.0 Smoothness/length-scale for spatial covariance variation diag_term : float, default=1e-6 Small positive value added to the covariance diagonal for numerical stability.

model_kwargs : Any Reserved for future keyword arguments accepted by the base Model.__init__. Do not pass WPPM math knobs or task/likelihood knobs here.

Methods:

Name	Description
init_params	Sample initial parameters from the prior.
local_covariance	Return local covariance Σ(x) at stimulus location x.
log_likelihood_from_data	Compute log-likelihood directly from a batched data object.
log_posterior_from_data	Compute log posterior from data.
predict_prob	Predict probability of a correct response for a single stimulus.

Attributes:

Name	Type	Description
basis_degree	int | None	Chebyshev polynomial degree for Wishart process basis expansion.
decay_rate
diag_term
embedding_dim	int	Dimension of the embedding space.
extra_dims
input_dim
likelihood
noise
prior
variance_scale

Source code in src/psyphy/model/wppm.py

def __init__(
    self,
    prior: Prior,
    likelihood: TaskLikelihood,
    noise: Any | None = None,
    *,  # everything after here is keyword-only
    input_dim: int = 2,
    extra_dims: int = 1,
    variance_scale: float = 4e-3,
    decay_rate: float = 0.4,
    diag_term: float = 1e-6,
    **model_kwargs: Any,
) -> None:
    # Base-model configuration.
    #
    # `model_kwargs` is reserved for *future* base `Model.__init__` kwargs.
    # It should NOT be used for WPPM-specific math (e.g. alternative covariance
    # parameterizations) or for task-specific likelihood knobs.
    if model_kwargs:
        known_misuses = {"num_samples", "bandwidth", "online_config"}
        bad = sorted(known_misuses.intersection(model_kwargs.keys()))
        if bad:
            raise TypeError(
                "Do not pass task-specific kwargs via WPPM(..., **model_kwargs). "
                f"Move {bad} into the task config (e.g. OddityTaskConfig)."
            )

    super().__init__(**model_kwargs)

    # --- core components ---
    self.input_dim = int(input_dim)  # stimulus-space dimensionality
    self.prior = prior  # prior over parameter PyTree

    if self.prior.input_dim != self.input_dim:
        raise ValueError(
            f"Dimension mismatch: Model initialized with input_dim={self.input_dim}, "
            f"but Prior expects input_dim={self.prior.input_dim}."
        )

    self.likelihood = likelihood  # task mapping and likelihood
    self.noise = noise  # noise model

    self.extra_dims = int(extra_dims)
    self.variance_scale = float(variance_scale)
    self.decay_rate = float(decay_rate)
    self.diag_term = float(diag_term)

basis_degree

basis_degree: int | None

Chebyshev polynomial degree for Wishart process basis expansion.

This property delegates to self.prior.basis_degree to ensure consistency between parameter sampling and basis evaluation.

Returns:

Type	Description
int | None	Degree of Chebyshev polynomial basis (0 = constant, 1 = linear, etc.)

Notes

WPPM gets its basis_degree parameter from Prior.basis_degree.

decay_rate

decay_rate = float(decay_rate)

diag_term

diag_term = float(diag_term)

embedding_dim

embedding_dim: int

Dimension of the embedding space.

embedding_dim = input_dim + extra_dims. this represents the full perceptual space where: - First input_dim dimensions correspond to observable stimulus features - Remaining extra_dims are latent dimensions

Returns:

Type	Description
int	input_dim + extra_dims

Notes

This is a computed property, not a constructor parameter.

extra_dims

extra_dims = int(extra_dims)

input_dim

input_dim = int(input_dim)

likelihood

likelihood = likelihood

noise

noise = noise

prior

prior = prior

variance_scale

variance_scale = float(variance_scale)

init_params

init_params(key: Array) -> Params

Sample initial parameters from the prior.

Returns:

Name	Type	Description
params	dict[str, ndarray]

Source code in src/psyphy/model/wppm.py

def init_params(self, key: jax.Array) -> Params:
    """Sample initial parameters from the prior.

    Returns
    -------
    params : dict[str, jnp.ndarray]
    """
    return self.prior.sample_params(key)

local_covariance

local_covariance(params: Params, x: ndarray) -> ndarray

Return local covariance Σ(x) at stimulus location x.

Wishart mode (basis_degree set): Σ(x) = U(x) @ U(x)^T + diag_term * I where U(x) is rectangular (input_dim, embedding_dim) if extra_dims > 0. - Varies smoothly with x - Guaranteed positive-definite - Returns stimulus covariance directly (input_dim, input_dim)

Parameters:

Name	Type	Description	Default
params	dict	Model parameters: - WPPM: {"W": (degree+1, ..., input_dim, embedding_dim)}	required
x	(ndarray, shape(input_dim))	Stimulus location	required

Returns:

Type	Description
Σ : jnp.ndarray, shape (input_dim, input_dim)	Covariance matrix in stimulus space.

Source code in src/psyphy/model/wppm.py

def local_covariance(self, params: Params, x: jnp.ndarray) -> jnp.ndarray:
    """
    Return local covariance Σ(x) at stimulus location x.


    Wishart mode (basis_degree set):
        Σ(x) = U(x) @ U(x)^T + diag_term * I
        where U(x) is rectangular (input_dim, embedding_dim) if extra_dims > 0.
        - Varies smoothly with x
        - Guaranteed positive-definite
        - Returns stimulus covariance directly (input_dim, input_dim)

    Parameters
    ----------
    params : dict
        Model parameters:
        - WPPM: {"W": (degree+1, ..., input_dim, embedding_dim)}
    x : jnp.ndarray, shape (input_dim,)
        Stimulus location

    Returns
    -------
    Σ : jnp.ndarray, shape (input_dim, input_dim)
        Covariance matrix in stimulus space.
    """
    # WPPM: spatially-varying covariance
    if "W" in params:
        U = self._compute_sqrt(params, x)  # (input_dim, embedding_dim)
        # Σ(x) = U(x) @ U(x)^T + diag_term * I
        # Result is (input_dim, input_dim)
        Sigma = U @ U.T + self.diag_term * jnp.eye(self.input_dim)
        return Sigma

    raise ValueError("params must contain 'W' (weights of WPPM)")

log_likelihood_from_data

log_likelihood_from_data(params: Params, data: Any, *, key: Array | None = None) -> ndarray

Compute log-likelihood directly from a batched data object.

Why delegate to the likelihood? - The likelihood knows the decision rule (oddity, 2AFC, ...). - The likelihood can use the model (this WPPM) to fetch discriminabilities. - The likelihood can use the noise model if it needs MC simulation.

Parameters:

Name	Type	Description	Default
params	dict	Model parameters.	required
data	TrialData (or any object with refs/comparisons/responses arrays)	Collected trial data.	required
key	Array | None	JAX random key for MC likelihood evaluation. When provided, a fresh noise realization is drawn every call — required for correct stochastic gradient estimates during optimization. When None, the task falls back to OddityTaskConfig.default_key_seed (useful for fixed evaluation and testing, but should not be used during gradient-based optimization).	None

Returns:

Name	Type	Description
loglik	ndarray	Scalar log-likelihood (task-only; add prior outside if needed).

Source code in src/psyphy/model/wppm.py

def log_likelihood_from_data(
    self, params: Params, data: Any, *, key: jax.Array | None = None
) -> jnp.ndarray:
    """Compute log-likelihood directly from a batched data object.

    Why delegate to the likelihood?
        - The likelihood knows the decision rule (oddity, 2AFC, ...).
        - The likelihood can use the model (this WPPM) to fetch discriminabilities.
        - The likelihood can use the noise model if it needs MC simulation.

    Parameters
    ----------
    params : dict
        Model parameters.
    data : TrialData (or any object with refs/comparisons/responses arrays)
        Collected trial data.
    key : jax.Array | None, optional
        JAX random key for MC likelihood evaluation. When provided, a fresh
        noise realization is drawn every call — required for correct stochastic
        gradient estimates during optimization. When None, the task falls back
        to ``OddityTaskConfig.default_key_seed`` (useful for fixed evaluation
        and testing, but should not be used during gradient-based optimization).

    Returns
    -------
    loglik : jnp.ndarray
        Scalar log-likelihood (task-only; add prior outside if needed).
    """
    return self.likelihood.loglik(params, data, self, key=key)

log_posterior_from_data

log_posterior_from_data(params: Params, data: Any, *, key: Array | None = None) -> ndarray

Compute log posterior from data.

This simply adds the prior log-probability to the task log-likelihood. Inference engines (e.g., MAP optimizer) typically optimize this quantity.

Parameters:

Name	Type	Description	Default
params	dict	Model parameters.	required
data	TrialData	Collected trial data.	required
key	Array | None	JAX random key for the MC likelihood. Must be provided during optimization so each gradient step uses a fresh noise realization. When None, falls back to OddityTaskConfig.default_key_seed.	None

Returns:

Type	Description
ndarray	Scalar log posterior = loglik(params | data) + log_prior(params).

Source code in src/psyphy/model/wppm.py

def log_posterior_from_data(
    self, params: Params, data: Any, *, key: jax.Array | None = None
) -> jnp.ndarray:
    """Compute log posterior from data.

    This simply adds the prior log-probability to the task log-likelihood.
    Inference engines (e.g., MAP optimizer) typically optimize this quantity.

    Parameters
    ----------
    params : dict
        Model parameters.
    data : TrialData
        Collected trial data.
    key : jax.Array | None, optional
        JAX random key for the MC likelihood. Must be provided during
        optimization so each gradient step uses a fresh noise realization.
        When None, falls back to ``OddityTaskConfig.default_key_seed``.

    Returns
    -------
    jnp.ndarray
        Scalar log posterior = loglik(params | data) + log_prior(params).
    """
    return self.log_likelihood_from_data(
        params, data, key=key
    ) + self.prior.log_prob(params)

predict_prob

predict_prob(params: Params, stimulus: Stimulus, **likelihood_kwargs: Any) -> ndarray

Predict probability of a correct response for a single stimulus.

Design choice: WPPM computes discriminability & covariance; the LIKELIHOOD defines how that translates to performance. We therefore delegate to: likelihood.predict(params, stimulus, model=self, noise=self.noise)

Parameters:

Name	Type	Description	Default
params	dict		required
stimulus	tuple[ndarray, ndarray]	(reference, comparison) pair in model space.	required

Returns:

Name	Type	Description
p_correct	ndarray

Source code in src/psyphy/model/wppm.py

def predict_prob(
    self, params: Params, stimulus: Stimulus, **likelihood_kwargs: Any
) -> jnp.ndarray:
    """
    Predict probability of a correct response for a single stimulus.

    Design choice:
        WPPM computes discriminability & covariance; the LIKELIHOOD defines how
        that translates to performance. We therefore delegate to:
            likelihood.predict(params, stimulus, model=self, noise=self.noise)

    Parameters
    ----------
    params : dict
    stimulus : tuple[jnp.ndarray, jnp.ndarray]
         (reference, comparison) pair in model space.

    Returns
    -------
    p_correct : jnp.ndarray
    """
    # Strict task-owned configuration:
    # - MC control knobs (e.g. num_samples/bandwidth) live in the task config.
    # - WPPM.predict_prob therefore does not accept task-specific kwargs.
    if likelihood_kwargs:
        raise TypeError(
            f"WPPM.predict_prob got unexpected kwargs: {likelihood_kwargs}. "
            "Configure likelihood behavior via the TaskLikelihood object itself."
        )

    ref, comparison = stimulus
    return self.likelihood.predict(params, ref, comparison, self)

---