Skip to content

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. MVP: Gaussian prior over diagonal log-variances.

  9. Full WPPM mode: structured prior over basis weights and lengthscale-controlled covariance fields.

  10. TaskLikelihood (model/task.py):

  11. Encodes the psychophysical decision rule.
  12. MVP: OddityTask (3AFC) and TwoAFC with sigmoid mappings.

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

  14. NoiseModel (model/noise.py):

  15. Defines the distribution of internal representation noise.
  16. MVP: GaussianNoise (zero mean, isotropic).
  17. Full WPPM mode: add StudentTNoise option and beyond.
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, TwoAFC, GaussianNoise, StudentTNoise from psyphy.inference import MAPOptimizer, LangevinSampler, LaplaceApproximation from psyphy.posterior import Posterior, effective_sample_size, rhat 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, StaircasePlacement from psyphy.utils import grid_candidates, sobol_candidates, custom_candidates, chebyshev_basis from psyphy.utils import bootstrap_predictions, bootstrap_statistic, bootstrap_compare_models

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.
  • To upgrade from MVP -> Full WPPM mode: replace local_covariance and discriminability with basis-expansion Wishart process + MC simulation.
MVP vs Full WPPM mode
  • MVP is a diagonal-covariance, closed-form scaffold that runs out of the box.
  • Full WPPM mode matches the published research model:
  • Smooth covariance fields (Wishart process priors).
  • Monte Carlo likelihood evaluation.
  • Explicit noise model in predictions.

Classes:

Name Description
ExperimentSession

High-level experiment orchestrator.
GaussianNoise
LangevinSampler

Langevin sampler (stub).
LaplaceApproximation

Laplace approximation around MAP estimate.
MAPOptimizer

MAP (Maximum A Posteriori) optimizer.
OddityTask

Three-alternative forced-choice oddity task.
Prior

Prior distribution over WPPM parameters
ResponseData

Container for psychophysical trial data.
StudentTNoise
TrialBatch

Container for a proposed batch of trials
TwoAFC

2-alternative forced-choice task (MVP placeholder).
WPPM

Wishart Process Psychophysical Model (WPPM).

Attributes:

Name Type Description
Posterior

Posterior 

Posterior = MAPPosterior

ExperimentSession 

ExperimentSession(
    model, inference, placement, init_placement=None
)

High-level experiment orchestrator.

Parameters:

Name Type Description Default
model WPPM

(Psychophysical) model instance.

 required
inference InferenceEngine

Inference engine (MAP, Langevin, etc.).

 required
placement TrialPlacement

Adaptive trial placement strategy.

 required
init_placement TrialPlacement

Initial placement strategy (e.g., Sobol exploration).

 None

Attributes:

Name Type Description
data ResponseData

Stores all collected trials.
posterior Posterior or None

Current posterior estimate (None before initialization).

Methods:

Name Description
initialize

Fit an initial posterior before any adaptive placement.
next_batch

Propose the next batch of trials.
update

Refit posterior with accumulated data.
Source code in src/psyphy/session/experiment_session.py 
52
53
54
55
56
57
58
59
60
61
62
def __init__(self, model, inference, placement, init_placement=None):
    self.model = model
    self.inference = inference
    self.placement = placement
    self.init_placement = init_placement

    # Data store starts empty
    self.data = ResponseData()

    # Posterior will be set after initialize() or update()
    self.posterior = None

data 

data = ResponseData()

inference 

inference = inference

init_placement 

init_placement = init_placement

model 

model = model

placement 

placement = placement

posterior 

posterior = None

initialize 

initialize()

Fit an initial posterior before any adaptive placement.

Returns:

Type Description
Posterior

Posterior object wrapping fitted parameters.
Notes

MVP: Posterior is fitted to empty data (prior only). Full WPPM mode: Could use pilot data or pre-collected trials along grid etc.

Source code in src/psyphy/session/experiment_session.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def initialize(self):
    """
    Fit an initial posterior before any adaptive placement.

    Returns
    -------
    Posterior
        Posterior object wrapping fitted parameters.

    Notes
    -----
    MVP:
        Posterior is fitted to empty data (prior only).
    Full WPPM mode:
        Could use pilot data or pre-collected trials along grid etc.
    """
    self.posterior = self.inference.fit(self.model, self.data)
    return self.posterior

next_batch 

next_batch(batch_size: int)

Propose the next batch of trials.

Parameters:

Name Type Description Default
batch_size int

Number of trials to propose.

 required

Returns:

Type Description
TrialBatch

Batch of proposed (reference, probe) stimuli.
Notes

MVP: Always calls placement.propose() on current posterior. Full WPPM mode: Could support hybrid placement (init strategy -> adaptive strategy).

Source code in src/psyphy/session/experiment_session.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def next_batch(self, batch_size: int):
    """
    Propose the next batch of trials.

    Parameters
    ----------
    batch_size : int
        Number of trials to propose.

    Returns
    -------
    TrialBatch
        Batch of proposed (reference, probe) stimuli.

    Notes
    -----
    MVP:
        Always calls placement.propose() on current posterior.
    Full WPPM mode:
        Could support hybrid placement (init strategy -> adaptive strategy).
    """
    if self.posterior is None:
        raise RuntimeError("Posterior not initialized. Call initialize() first.")
    return self.placement.propose(self.posterior, batch_size)

update 

update()

Refit posterior with accumulated data.

Returns:

Type Description
Posterior

Updated posterior.
Notes

MVP: Re-optimizes from scratch using all data. Full WPPM mode: Could support warm-start or online parameter updates.

Source code in src/psyphy/session/experiment_session.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def update(self):
    """
    Refit posterior with accumulated data.

    Returns
    -------
    Posterior
        Updated posterior.

    Notes
    -----
    MVP:
        Re-optimizes from scratch using all data.
    Full WPPM mode:
        Could support warm-start or online parameter updates.
    """
    self.posterior = self.inference.fit(self.model, self.data)
    return self.posterior

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 
24
25
26
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 
28
29
30
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,
)

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 
36
37
38
39
40
41
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) -> Posterior

Fit model parameters with Langevin dynamics (stub).

Parameters:

Name Type Description Default
model WPPM

Model instance.

 required
data ResponseData

Observed trials.

 required

Returns:

Type Description
Posterior

Posterior wrapper (MVP: params from init).
Source code in src/psyphy/inference/langevin.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def fit(self, model, data) -> Posterior:
    """
    Fit model parameters with Langevin dynamics (stub).

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

    Returns
    -------
    Posterior
        Posterior wrapper (MVP: params from init).
    """
    return Posterior(params=model.init_params(None), model=model)

LaplaceApproximation

Laplace approximation around MAP estimate.

Methods:

Name Description
from_map

Construct a Gaussian approximation centered at MAP.

from_map 

from_map(map_posterior: Posterior) -> Posterior

Return posterior approximation from MAP.

Parameters:

Name Type Description Default
map_posterior Posterior

Posterior object from MAP optimization.

 required

Returns:

Type Description
Posterior

Same posterior object (MVP).
Source code in src/psyphy/inference/laplace.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def from_map(self, map_posterior: Posterior) -> Posterior:
    """
    Return posterior approximation from MAP.

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

    Returns
    -------
    Posterior
        Same posterior object (MVP).
    """
    return map_posterior

MAPOptimizer 

MAPOptimizer(
    steps: int = 500,
    learning_rate: float = 5e-05,
    momentum: float = 0.9,
    optimizer: GradientTransformation | None = None,
    *,
    track_history: bool = False,
    log_every: int = 10,
)

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.

 False
log_every int

Record every N steps (also records the last step).

 10

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]
optimizer
steps
track_history
Source code in src/psyphy/inference/map_optimizer.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(
    self,
    steps: int = 500,
    learning_rate: float = 5e-5,
    momentum: float = 0.9,
    optimizer: optax.GradientTransformation | None = None,
    *,
    track_history: bool = False,
    log_every: int = 10,
):
    """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).
    """
    self.steps = steps
    self.optimizer = optimizer or optax.sgd(
        learning_rate=learning_rate, momentum=momentum
    )
    self.track_history = track_history
    self.log_every = max(1, int(log_every))
    # 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] = []

optimizer 

optimizer = optimizer or sgd(
    learning_rate=learning_rate, momentum=momentum
)

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. 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 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
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. If None, defaults to 0.

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

    def loss_fn(params):
        return -model.log_posterior_from_data(params, data)

    # Initialize parameters
    if init_params is not None:
        params = init_params
    else:
        rng_seed = 0 if seed is None else int(seed)
        params = model.init_params(jax.random.PRNGKey(rng_seed))
    opt_state = self.optimizer.init(params)

    @jax.jit
    def step(params, opt_state):
        # Ensure params and opt_state are JAX PyTrees for JIT compatibility
        loss, grads = jax.value_and_grad(loss_fn)(params)  # auto-diff
        updates, opt_state = self.optimizer.update(
            grads, opt_state, params
        )  # optimizer update
        params = optax.apply_updates(params, updates)  # apply updates
        # Only return JAX-compatible types (PyTrees of arrays, scalars)
        return params, opt_state, loss

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

    for i in range(self.steps):
        params, opt_state, loss = step(params, opt_state)
        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:
                # Best-effort; do not break fitting if logging fails
                pass

    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 
148
149
150
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(slope: float = 1.5)

Bases: TaskLikelihood

Three-alternative forced-choice oddity task.

In an oddity task, the observer is presented with three stimuli: two identical references and one comparison. The task is to identify which stimulus is the "odd one out" (the comparison). Performance depends on the discriminability between reference and comparison.

This class provides two likelihood computation methods:

  1. Analytical approximation (MVP mode):
  2. predict(): maps discriminability to P(correct) via tanh
  3. loglik(): Bernoulli likelihood using analytical predictions

  4. Fast, differentiable, suitable for gradient-based optimization

  5. Monte Carlo simulation (Full WPPM mode):

  6. loglik_mc(): simulates the full 3-stimulus oddity task
  7. Samples three internal representations per trial (z0, z1, z2)
  8. Uses proper oddity decision rule with three pairwise distances
  9. More accurate for complex covariance structures
  10. Suitable for validation and benchmarking

Parameters:

Name Type Description Default
slope float

Slope parameter for analytical tanh mapping in predict(). Controls steepness of discriminability -> performance relationship.

 1.5

Attributes:

Name Type Description
chance_level float

Chance performance for oddity task (1/3)
performance_range float

Range from chance to perfect performance (2/3)
Notes

The analytical approximation in predict() uses: P(correct) = 1/3 + 2/3 * (1 + tanh(slope * d)) / 2

MC simulation in loglik_mc() (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:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
>>> from psyphy.model.task import OddityTask
>>> 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 = OddityTask(slope=1.5)
>>> model = WPPM(
...     input_dim=2, prior=Prior(input_dim=2), task=task, noise=GaussianNoise()
... )
>>> params = model.init_params(jr.PRNGKey(0))
>>>
>>> # Analytical prediction
>>> ref = jnp.array([0.0, 0.0])
>>> comparison = jnp.array([0.5, 0.5])
>>> p_correct = task.predict(params, (ref, comparison), model, model.noise)
>>> print(f"P(correct) pprox {p_correct:.3f}")
>>>
>>> # MC simulation (more accurate)
>>> from psyphy.data.dataset import ResponseData
>>> data = ResponseData()
>>> data.add_trial(ref, comparison, resp=1)
>>> ll_mc = task.loglik_mc(
...     params, data, model, model.noise, num_samples=1000, key=jr.PRNGKey(42)
... )
>>> print(f"Log-likelihood (MC): {ll_mc:.4f}")

Methods:

Name Description
loglik

Compute log-likelihood using analytical predictions.
loglik_mc

Compute log-likelihood via Monte Carlo observer simulation.
predict

Predict probability of correct response using analytical approximation.
Source code in src/psyphy/model/task.py 
153
154
155
156
def __init__(self, slope: float = 1.5) -> None:
    self.slope = float(slope)
    self.chance_level: float = 1.0 / 3.0
    self.performance_range: float = 1.0 - self.chance_level

chance_level 

chance_level: float = 1.0 / 3.0

performance_range 

performance_range: float = 1.0 - chance_level

slope 

slope = float(slope)

loglik 

loglik(
    params: Any, data: Any, model: Any, noise: Any
) -> ndarray

Compute log-likelihood using analytical predictions.

Parameters:

Name Type Description Default
params dict

Model parameters

 required
data ResponseData

Trial data containing refs, comparisons, responses

 required
model WPPM

Model instance

 required
noise NoiseModel

Observer noise model

 required

Returns:

Type Description
ndarray

Scalar sum of log-likelihoods over all trials
Notes

Uses Bernoulli log-likelihood: LL = Σ [y * log(p) + (1-y) * log(1-p)] where p comes from predict() (analytical approximation)

Source code in src/psyphy/model/task.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def loglik(self, params: Any, data: Any, model: Any, noise: Any) -> jnp.ndarray:
    """
    Compute log-likelihood using analytical predictions.

    Parameters
    ----------
    params : dict
        Model parameters
    data : ResponseData
        Trial data containing refs, comparisons, responses
    model : WPPM
        Model instance
    noise : NoiseModel
        Observer noise model

    Returns
    -------
    jnp.ndarray
        Scalar sum of log-likelihoods over all trials

    Notes
    -----
    Uses Bernoulli log-likelihood:
        LL = Σ [y * log(p) + (1-y) * log(1-p)]
    where p comes from predict() (analytical approximation)
    """
    refs, comparisons, responses = data.to_numpy()
    ps = jnp.array(
        [
            self.predict(params, (r, p), model, noise)
            for r, p in zip(refs, comparisons)
        ]
    )
    eps = 1e-9
    return jnp.sum(
        jnp.where(responses == 1, jnp.log(ps + eps), jnp.log(1.0 - ps + eps))
    )

loglik_mc 

loglik_mc(
    params: Any,
    data: Any,
    model: Any,
    noise: Any,
    num_samples: int = 1000,
    bandwidth: float = 0.01,
    key: Any = None,
) -> ndarray

Compute log-likelihood via Monte Carlo observer simulation.

This method implements the FULL 3-stimulus oddity task. Instead of using an analytical approximation, we: 1. Sample three internal noisy representations per trial: - z_ref, z_refprime ~ N(ref, Σ_ref) [two samples from reference] - z_comparison ~ N(comparison, Σ_comparison) [one sample from comparison] 2. Compute three pairwise Mahalanobis distances 3. Apply oddity decision rule: comparison is odd if it's farther from BOTH ref and reference_prime 4. Apply logistic smoothing to approximate P(correct) 5. Average over MC samples

Parameters:

Name Type Description Default
params dict

Model parameters (must contain 'W' for WPPM basis coefficients)

 required
data ResponseData

Trial data with refs, comparisons, and responses

 required
model WPPM

Model instance providing compute_U() for covariance computation

 required
noise NoiseModel

Observer noise model (provides sigma for diagonal noise term)

 required
num_samples int

Number of Monte Carlo samples per trial. - Use 1000-5000 for accurate likelihood estimation - Larger values reduce MC variance but increase compute time

 1000
bandwidth float

Smoothing parameter for logistic CDF approximation. - Smaller values -> sharper transition (closer to step function) - Larger values -> smoother approximation - Typical range: [1e-3, 5e-2]

 1e-2
key PRNGKey

Random key for reproducible sampling. If None, uses PRNGKey(0) (deterministic but not recommended for production)

 None

Returns:

Type Description
ndarray

Scalar sum of log-likelihoods over all trials. Same shape and interpretation as loglik().

Raises:

Type Description
ValueError

If num_samples <= 0
Notes

Full 3-stimulus oddity task algorithm:

For each trial (ref, comparison, response): 1. Compute covariances: - Σ_ref = U_ref @ U_ref.T + σ^2 I - Σ_comparison = U_comparison @ U_comparison.T + σ^2 I - Σ_avg = (2/3) Σ_ref + (1/3) Σ_comparison [weighted by stimulus frequency]

  1. Sample three internal representations:
  2. z_ref, z_refprime ~ N(ref, Σ_ref) [2 samples from reference, num_samples times each]

  3. z_comparison ~ N(comparison, Σ_comparison) [1 sample from comparison, num_samples times]

  4. Compute three pairwise Mahalanobis distances:

  5. d^2(z_ref, z_refprime) = (z_ref - z_refprime).T @ Σ_avg^{-1} @ (z_ref - z_refprime) [ref vs reference_prime]
  6. d^2(z_ref, z_comparison) = (z_ref - z_comparison).T @ Σ_avg^{-1} @ (z_ref - z_comparison) [ref vs comparison]

  7. d^2(z_refprime, z_comparison) = (z_refprime - z_comparison).T @ Σ_avg^{-1} @ (z_refprime - z_comparison) [reference_prime vs comparison]

  8. Apply oddity decision rule:

  9. delta = min(d^2(z_ref,z_comparison), d^2(z_refprime,z_comparison)) - d^2(z_ref,z_refprime)

  10. delta > 0 means comparison is farther from BOTH ref and reference_prime -> correct identification

  11. Apply logistic smoothing:

  12. P(correct) pprox mean(logistic.cdf(delta / bandwidth))

  13. Bernoulli log-likelihood:

  14. LL = Σ [y * log(p) + (1-y) * log(1-p)]

Performance: - Time complexity: O(n_trials * num_samples * input_dim³) - Memory: O(num_samples * input_dim) per trial - Vectorized across trials using jax.vmap for GPU acceleration - Can be JIT-compiled for additional speed (future optimization)

Comparison to analytical: - MC implements full 3-stimulus oddity (more realistic) - MC is more accurate for complex Σ(x) structures - Analytical is faster and differentiable - Use MC for validation and benchmarking, analytical for optimization

Examples:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> from psyphy.model import WPPM, Prior
>>> from psyphy.model.task import OddityTask
>>> from psyphy.model.noise import GaussianNoise
>>> from psyphy.data.dataset import ResponseData
>>>
>>> # Setup
>>> model = WPPM(
...     input_dim=2,
...     prior=Prior(input_dim=2, basis_degree=3),
...     task=OddityTask(),
...     noise=GaussianNoise(sigma=0.03),
... )
>>> params = model.init_params(jr.PRNGKey(0))
>>>
>>> # Create trial data
>>> data = ResponseData()
>>> data.add_trial(
...     ref=jnp.array([0.0, 0.0]), comparison=jnp.array([0.3, 0.2]), resp=1
... )
>>>
>>> # Compare analytical vs MC
>>> ll_analytical = model.task.loglik(params, data, model, model.noise)
>>> ll_mc = model.task.loglik_mc(
...     params,
...     data,
...     model,
...     model.noise,
...     num_samples=5000,
...     bandwidth=1e-3,
...     key=jr.PRNGKey(42),
... )
>>> print(f"Analytical: {ll_analytical:.4f}")
>>> print(f"MC (N=5000): {ll_mc:.4f}")
>>> print(f"Difference: {abs(ll_mc - ll_analytical):.4f}")
See Also

loglik : Analytical log-likelihood (faster, differentiable) predict : Analytical prediction for single trial

Source code in src/psyphy/model/task.py 
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def loglik_mc(
    self,
    params: Any,
    data: Any,
    model: Any,
    noise: Any,
    num_samples: int = 1000,
    bandwidth: float = 1e-2,
    key: Any = None,
) -> jnp.ndarray:
    """
    Compute log-likelihood via Monte Carlo observer simulation.

    This method implements the FULL 3-stimulus oddity task. Instead of using
    an analytical approximation, we:
    1. Sample three internal noisy representations per trial:
       - z_ref, z_refprime ~ N(ref, Σ_ref)  [two samples from reference]
       - z_comparison ~ N(comparison, Σ_comparison)           [one sample from comparison]
    2. Compute three pairwise Mahalanobis distances
    3. Apply oddity decision rule: comparison is odd if it's farther from BOTH ref and reference_prime
    4. Apply logistic smoothing to approximate P(correct)
    5. Average over MC samples

    Parameters
    ----------
    params : dict
        Model parameters (must contain 'W' for WPPM basis coefficients)
    data : ResponseData
        Trial data with refs, comparisons, and responses
    model : WPPM
        Model instance providing compute_U() for covariance computation
    noise : NoiseModel
        Observer noise model (provides sigma for diagonal noise term)
    num_samples : int, default=1000
        Number of Monte Carlo samples per trial.
        - Use 1000-5000 for accurate likelihood estimation
        - Larger values reduce MC variance but increase compute time
    bandwidth : float, default=1e-2
        Smoothing parameter for logistic CDF approximation.
        - Smaller values -> sharper transition (closer to step function)
        - Larger values -> smoother approximation
        - Typical range: [1e-3, 5e-2]
    key : jax.random.PRNGKey, optional
        Random key for reproducible sampling.
        If None, uses PRNGKey(0) (deterministic but not recommended for production)

    Returns
    -------
    jnp.ndarray
        Scalar sum of log-likelihoods over all trials.
        Same shape and interpretation as loglik().

    Raises
    ------
    ValueError
        If num_samples <= 0

    Notes
    -----
    **Full 3-stimulus oddity task algorithm:**

    For each trial (ref, comparison, response):
    1. Compute covariances:
       - Σ_ref = U_ref @ U_ref.T + σ^2 I
       - Σ_comparison = U_comparison @ U_comparison.T + σ^2 I
       - Σ_avg = (2/3) Σ_ref + (1/3) Σ_comparison  [weighted by stimulus frequency]

    2. Sample three internal representations:
       - z_ref, z_refprime ~ N(ref, Σ_ref)  [2 samples from reference, num_samples times each]
       - z_comparison ~ N(comparison, Σ_comparison)           [1 sample from comparison, num_samples times]

    3. Compute three pairwise Mahalanobis distances:
       - d^2(z_ref, z_refprime) = (z_ref - z_refprime).T @ Σ_avg^{-1} @ (z_ref - z_refprime)  [ref vs reference_prime]
       - d^2(z_ref, z_comparison) = (z_ref - z_comparison).T @ Σ_avg^{-1} @ (z_ref - z_comparison)  [ref vs comparison]
       - d^2(z_refprime, z_comparison) = (z_refprime - z_comparison).T @ Σ_avg^{-1} @ (z_refprime - z_comparison)  [reference_prime vs 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)
       - delta > 0 means comparison is farther from BOTH ref and reference_prime -> correct identification

    5. Apply logistic smoothing:
       - P(correct) \approx mean(logistic.cdf(delta / bandwidth))

    6. Bernoulli log-likelihood:
       - LL = Σ [y * log(p) + (1-y) * log(1-p)]

    Performance:
    - Time complexity: O(n_trials * num_samples * input_dim³)
    - Memory: O(num_samples * input_dim) per trial
    - Vectorized across trials using jax.vmap for GPU acceleration
    - Can be JIT-compiled for additional speed (future optimization)

    Comparison to analytical:
    - MC implements full 3-stimulus oddity (more realistic)
    - MC is more accurate for complex Σ(x) structures
    - Analytical is faster and differentiable
    - Use MC for validation and benchmarking, analytical for optimization

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> import jax.random as jr
    >>> from psyphy.model import WPPM, Prior
    >>> from psyphy.model.task import OddityTask
    >>> from psyphy.model.noise import GaussianNoise
    >>> from psyphy.data.dataset import ResponseData
    >>>
    >>> # Setup
    >>> model = WPPM(
    ...     input_dim=2,
    ...     prior=Prior(input_dim=2, basis_degree=3),
    ...     task=OddityTask(),
    ...     noise=GaussianNoise(sigma=0.03),
    ... )
    >>> params = model.init_params(jr.PRNGKey(0))
    >>>
    >>> # Create trial data
    >>> data = ResponseData()
    >>> data.add_trial(
    ...     ref=jnp.array([0.0, 0.0]), comparison=jnp.array([0.3, 0.2]), resp=1
    ... )
    >>>
    >>> # Compare analytical vs MC
    >>> ll_analytical = model.task.loglik(params, data, model, model.noise)
    >>> ll_mc = model.task.loglik_mc(
    ...     params,
    ...     data,
    ...     model,
    ...     model.noise,
    ...     num_samples=5000,
    ...     bandwidth=1e-3,
    ...     key=jr.PRNGKey(42),
    ... )
    >>> print(f"Analytical: {ll_analytical:.4f}")
    >>> print(f"MC (N=5000): {ll_mc:.4f}")
    >>> print(f"Difference: {abs(ll_mc - ll_analytical):.4f}")

    See Also
    --------
    loglik : Analytical log-likelihood (faster, differentiable)
    predict : Analytical prediction for single trial



    """
    # Validate inputs
    if num_samples <= 0:
        raise ValueError(f"num_samples must be > 0, got {num_samples}")

    # Default key for reproducibility (warn: not secure)
    if key is None:
        key = jr.PRNGKey(0)

    # Unpack trial data
    refs, comparisons, responses = data.to_numpy()
    n_trials = len(refs)

    # Split keys for each trial (ensures independent sampling)
    trial_keys = jr.split(key, n_trials)

    # Vectorized computation of P(correct) for all trials
    # This processes all trials in parallel using jax.vmap
    # Note: probabilities are already clipped in _simulate_trial_mc()
    probs = self._simulate_trials_mc_vectorized(
        params=params,
        refs=refs,
        comparisons=comparisons,
        model=model,
        noise=noise,
        num_samples=num_samples,
        bandwidth=bandwidth,
        trial_keys=trial_keys,
    )

    # Bernoulli log-likelihood: LL = Σ [y log(p) + (1-y) log(1-p)]
    # Probabilities are already clipped to [eps, 1-eps] so log is safe
    log_likelihoods = jnp.where(
        responses == 1,
        jnp.log(probs),  # Correct response
        jnp.log(1.0 - probs),  # Incorrect response
    )

    return jnp.sum(log_likelihoods)

predict 

predict(
    params: Any, stimuli: Stimulus, model: Any, noise: Any
) -> ndarray

Predict probability of correct response using analytical approximation.

Parameters:

Name Type Description Default
params dict

Model parameters (e.g., W for WPPM)

 required
stimuli tuple[ndarray, ndarray]

(reference, comparison) stimulus pair

 required
model WPPM

Model instance providing discriminability()

 required
noise NoiseModel

Observer noise model (currently unused in analytical version)

 required

Returns:

Type Description
ndarray

Scalar probability of correct response, in range [1/3, 1]
Notes

Uses tanh mapping: P(correct) = 1/3 + 2/3 * sigmoid(slope * d) where d is discriminability from model.discriminability()

Source code in src/psyphy/model/task.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def predict(
    self, params: Any, stimuli: Stimulus, model: Any, noise: Any
) -> jnp.ndarray:
    """
    Predict probability of correct response using analytical approximation.

    Parameters
    ----------
    params : dict
        Model parameters (e.g., W for WPPM)
    stimuli : tuple[jnp.ndarray, jnp.ndarray]
        (reference, comparison) stimulus pair
    model : WPPM
        Model instance providing discriminability()
    noise : NoiseModel
        Observer noise model (currently unused in analytical version)

    Returns
    -------
    jnp.ndarray
        Scalar probability of correct response, in range [1/3, 1]

    Notes
    -----
    Uses tanh mapping: P(correct) = 1/3 + 2/3 * sigmoid(slope * d)
    where d is discriminability from model.discriminability()
    """
    d = model.discriminability(params, stimuli)
    g = 0.5 * (jnp.tanh(self.slope * d) + 1.0)
    return self.chance_level + self.performance_range * g

Prior 

Prior(
    input_dim: int,
    scale: float = 0.5,
    basis_degree: int | None = None,
    variance_scale: float = 1.0,
    decay_rate: float = 0.5,
    lengthscale: float = 1.0,
    extra_embedding_dims: int = 0,
)

Prior distribution over WPPM parameters

Parameters:

Name Type Description Default
input_dim int

Dimensionality of the model space (same as WPPM.input_dim)

 required
scale float

Stddev of Gaussian prior for log_diag entries (MVP only).

 0.5
basis_degree int | None

Degree of Chebyshev basis for Wishart process. If None, uses MVP mode with log_diag parameters. 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
lengthscale float

Alias for decay_rate (kept for backward compatibility). If both specified, decay_rate takes precedence.

 1.0
extra_embedding_dims int

Additional latent dimensions in U matrices beyond input dimensions. Allows richer ellipsoid shapes in Wishart mode.

 0

Methods:

Name Description
default

Convenience constructor with MVP defaults.
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 | None
decay_rate float
extra_embedding_dims int
input_dim int
lengthscale float
scale float
variance_scale float

basis_degree 

basis_degree: int | None = None

decay_rate 

decay_rate: float = 0.5

extra_embedding_dims 

extra_embedding_dims: int = 0

input_dim 

input_dim: int

lengthscale 

lengthscale: float = 1.0

scale 

scale: float = 0.5

variance_scale 

variance_scale: float = 1.0

default 

default(input_dim: int, scale: float = 0.5) -> Prior

Convenience constructor with MVP defaults.

Source code in src/psyphy/model/prior.py 
84
85
86
87
@classmethod
def default(cls, input_dim: int, scale: float = 0.5) -> Prior:
    """Convenience constructor with MVP defaults."""
    return cls(input_dim=input_dim, scale=scale)

log_prob 

log_prob(params: Params) -> ndarray

Compute log prior density (up to a constant)

MVP mode: Isotropic Gaussian on log_diag

Wishart mode: 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 
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def log_prob(self, params: Params) -> jnp.ndarray:
    """
    Compute log prior density (up to a constant)

    MVP mode:
        Isotropic Gaussian on log_diag

    Wishart mode:
        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 "log_diag" in params:
        # MVP mode
        log_diag = params["log_diag"]
        var = self.scale**2
        return -0.5 * jnp.sum((log_diag**2) / var)

    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 either 'log_diag' (MVP) or 'W' (Wishart)")

sample_params 

sample_params(key: Any) -> Params

Sample initial parameters from the prior.

MVP mode (basis_degree=None): Returns {"log_diag": shape (input_dim,)}

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

1
2
3
Note: The 3rd dimension is input_dim (output space dimension).
This matches the einsum in _compute_U:
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 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def sample_params(self, key: Any) -> Params:
    """
    Sample initial parameters from the prior.

    MVP mode (basis_degree=None):
        Returns {"log_diag": shape (input_dim,)}

    Wishart mode (basis_degree set):
        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_U:
        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:
        # MVP mode: simple diagonal covariance
        log_diag = jr.normal(key, shape=(self.input_dim,)) * self.scale
        return {"log_diag": log_diag}

    # Wishart mode: 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()

Container for psychophysical trial data.

Attributes:

Name Type Description
refs List[Any]

List of reference stimuli.
comparisons List[Any]

List of comparison stimuli.
responses List[int]

List of subject responses (e.g., 0/1 or categorical).

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.
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.
Source code in src/psyphy/data/dataset.py 
41
42
43
44
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 
63
64
65
66
67
68
69
70
71
72
73
74
75
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 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
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 
203
204
205
206
207
208
209
210
211
212
213
214
215
216
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:

1
2
3
4
>>> # 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)

1
2
3
4
>>> # 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 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@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

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 
170
171
172
173
174
175
176
177
178
179
180
181
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 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
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.

Returns:

Name Type Description
refs ndarray
comparisons ndarray
responses ndarray
Source code in src/psyphy/data/dataset.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Return refs, comparisons, responses as numpy arrays.

    Returns
    -------
    refs : np.ndarray
    comparisons : np.ndarray
    responses : np.ndarray
    """
    return (
        np.array(self.refs),
        np.array(self.comparisons),
        np.array(self.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 
38
39
40
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 
42
43
44
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 
229
230
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 
232
233
234
235
236
237
@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)

TwoAFC 

TwoAFC(slope: float = 2.0)

Bases: TaskLikelihood

2-alternative forced-choice task (MVP placeholder).

Methods:

Name Description
loglik
predict

Attributes:

Name Type Description
chance_level float
performance_range float
slope
Source code in src/psyphy/model/task.py 
798
799
800
801
def __init__(self, slope: float = 2.0) -> None:
    self.slope = float(slope)
    self.chance_level: float = 0.5
    self.performance_range: float = 1.0 - self.chance_level

chance_level 

chance_level: float = 0.5

performance_range 

performance_range: float = 1.0 - chance_level

slope 

slope = float(slope)

loglik 

loglik(
    params: Any, data: Any, model: Any, noise: Any
) -> ndarray
Source code in src/psyphy/model/task.py 
809
810
811
812
813
814
815
816
817
818
819
820
def loglik(self, params: Any, data: Any, model: Any, noise: Any) -> jnp.ndarray:
    refs, comparisons, responses = data.to_numpy()
    ps = jnp.array(
        [
            self.predict(params, (r, p), model, noise)
            for r, p in zip(refs, comparisons)
        ]
    )
    eps = 1e-9
    return jnp.sum(
        jnp.where(responses == 1, jnp.log(ps + eps), jnp.log(1.0 - ps + eps))
    )

predict 

predict(
    params: Any, stimuli: Stimulus, model: Any, noise: Any
) -> ndarray
Source code in src/psyphy/model/task.py 
803
804
805
806
807
def predict(
    self, params: Any, stimuli: Stimulus, model: Any, noise: Any
) -> jnp.ndarray:
    d = model.discriminability(params, stimuli)
    return self.chance_level + self.performance_range * jnp.tanh(self.slope * d)

WPPM 

WPPM(
    input_dim: int,
    prior: Prior,
    task: TaskLikelihood,
    noise: Any | None = None,
    *,
    extra_dims: int = 0,
    variance_scale: float = 1.0,
    lengthscale: float = 1.0,
    diag_term: float = 1e-06,
    **kwargs,
)

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 probe live in R^{input_dim}.

 required
prior Prior

Prior distribution over model parameters. Controls basis_degree for Wishart mode (basis expansion) vs MVP mode (diagonal covariance). The WPPM delegates to prior.basis_degree to ensure consistency between parameter sampling and basis evaluation.

 required
task TaskLikelihood

Psychophysical task mapping that defines how discriminability translates to p(correct) and how log-likelihood of responses is computed. (e.g., OddityTask, TwoAFC)

 required
noise Any

Noise model describing internal representation noise (e.g., GaussianNoise). Not used in MVP mapping but passed to the task interface for future MC sims.

 None
Forward-compatible hyperparameters

extra_dims : int, default=0 Additional embedding dimensions for basis expansions (beyond input_dim). In Wishart mode, embedding_dim = input_dim + extra_dims. variance_scale : float, default=1.0 Global scaling factor for covariance magnitude (unused in MVP). lengthscale : float, default=1.0 Smoothness/length-scale for spatial covariance variation (unused in MVP). (formerly "decay_rate") diag_term : float, default=1e-6 Small positive value added to the covariance diagonal for numerical stability. MVP uses this in matrix solves; the research model will also use it.

Methods:

Name Description
condition_on_observations

Update model with new observations (online learning).
discriminability

Compute scalar discriminability d >= 0 for a (reference, probe) pair
fit

Fit model to data.
init_params

Sample initial parameters from the prior.
local_covariance

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

Compute the log-likelihood for arrays of trials.
log_likelihood_from_data

Compute log-likelihood directly from a ResponseData object.
log_posterior_from_data

Convenience helper if you want log posterior in one call (MVP).
posterior

Return posterior distribution.
predict_prob

Predict probability of a correct response for a single stimulus.
predict_with_params

Evaluate model at specific parameter values (no marginalization).

Attributes:

Name Type Description
basis_degree int | None

Chebyshev polynomial degree for Wishart process basis expansion.
diag_term
embedding_dim int

Dimension of the embedding space (perceptual space).
extra_dims
input_dim
lengthscale
noise
online_config
prior
task
variance_scale
Source code in src/psyphy/model/wppm.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def __init__(
    self,
    input_dim: int,
    prior: Prior,
    task: TaskLikelihood,
    noise: Any | None = None,
    *,  # everything after here is keyword-only
    extra_dims: int = 0,
    variance_scale: float = 1.0,
    lengthscale: float = 1.0,
    diag_term: float = 1e-6,
    **kwargs,  # Accept online_config from model base
) -> None:
    # Initialize Model base class
    super().__init__(**kwargs)

    # --- core components ---
    self.input_dim = int(input_dim)  # stimulus-space dimensionality
    self.prior = prior  # prior over parameter PyTree
    self.task = task  # task mapping and likelihood
    self.noise = noise  # noise model

    # --- forward-compatible hyperparameters (stubs in MVP) ---
    self.extra_dims = int(extra_dims)
    self.variance_scale = float(variance_scale)
    self.lengthscale = float(lengthscale)
    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.) None indicates MVP mode (no basis expansion)
Notes

WPPM gets its basis_degree parameter from Prior.basis_degree.

diag_term 

diag_term = float(diag_term)

embedding_dim 

embedding_dim: int

Dimension of the embedding space (perceptual 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 (in Wishart mode) input_dim (in MVP mode, extra_dims ignored)
Notes

This is a computed property, not a constructor parameter.

extra_dims 

extra_dims = int(extra_dims)

input_dim 

input_dim = int(input_dim)

lengthscale 

lengthscale = float(lengthscale)

noise 

noise = noise

online_config 

online_config = online_config or OnlineConfig()

prior 

prior = prior

task 

task = task

variance_scale 

variance_scale = float(variance_scale)

condition_on_observations 

condition_on_observations(X: ndarray, y: ndarray) -> Model

Update model with new observations (online learning).

Behavior depends on self.online_config.strategy: - "full": Accumulate all data, refit periodically - "sliding_window": Keep only recent window_size trials - "reservoir": Random sampling of window_size trials - "none": Refit from scratch (no caching)

Returns a NEW model instance (immutable update).

Parameters:

Name Type Description Default
X ndarray

New stimuli

 required
y ndarray

New responses

 required

Returns:

Type Description
Model

Updated model (new instance)

Examples:

1
2
3
4
5
>>> # Online learning loop
>>> model = WPPM(...).fit(X_init, y_init)
>>> for X_new, y_new in stream:
...     model = model.condition_on_observations(X_new, y_new)
...     # Model automatically manages memory and refitting
Source code in src/psyphy/model/base.py 
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
def condition_on_observations(self, X: jnp.ndarray, y: jnp.ndarray) -> Model:
    """
    Update model with new observations (online learning).

    Behavior depends on self.online_config.strategy:
    - "full": Accumulate all data, refit periodically
    - "sliding_window": Keep only recent window_size trials
    - "reservoir": Random sampling of window_size trials
    - "none": Refit from scratch (no caching)

    Returns a NEW model instance (immutable update).

    Parameters
    ----------
    X : jnp.ndarray
        New stimuli
    y : jnp.ndarray
        New responses

    Returns
    -------
    Model
        Updated model (new instance)

    Examples
    --------
    >>> # Online learning loop
    >>> model = WPPM(...).fit(X_init, y_init)
    >>> for X_new, y_new in stream:
    ...     model = model.condition_on_observations(X_new, y_new)
    ...     # Model automatically manages memory and refitting
    """
    from psyphy.data import ResponseData

    # Convert new data
    new_data = ResponseData.from_arrays(X, y)

    # Update data buffer according to strategy
    if self.online_config.strategy == "none":
        data_to_fit = new_data

    elif self.online_config.strategy == "full":
        if self._data_buffer is None:
            self._data_buffer = ResponseData()
        self._data_buffer.merge(new_data)
        data_to_fit = self._data_buffer

    elif self.online_config.strategy == "sliding_window":
        if self._data_buffer is None:
            self._data_buffer = ResponseData()
        self._data_buffer.merge(new_data)

        # Keep only last N trials
        window_size = self.online_config.window_size
        assert window_size is not None, "window_size must be set for sliding_window"
        if len(self._data_buffer) > window_size:
            self._data_buffer = self._data_buffer.tail(window_size)
        data_to_fit = self._data_buffer

    elif self.online_config.strategy == "reservoir":
        if self._data_buffer is None:
            self._data_buffer = ResponseData()

        # Reservoir sampling
        window_size = self.online_config.window_size
        assert window_size is not None, "window_size must be set for reservoir"
        self._data_buffer = self._reservoir_update(
            self._data_buffer,
            new_data,
            window_size,
        )
        data_to_fit = self._data_buffer
    else:
        raise ValueError(f"Unknown strategy: {self.online_config.strategy}")

    # Decide whether to refit
    self._n_updates += 1
    should_refit = self._n_updates % self.online_config.refit_interval == 0

    if not should_refit:
        # Return clone with updated buffer but old posterior
        new_model = self._clone()
        new_model._data_buffer = data_to_fit
        return new_model

    # Refit with optional warm start
    inference = self._inference_engine
    assert inference is not None, (
        "Model must be fit before condition_on_observations"
    )

    if self.online_config.warm_start and self._posterior is not None:
        # TODO: Add warm-start support to inference engines
        # For now, just refit from cached params
        pass

    new_model = self._clone()
    new_model._data_buffer = data_to_fit
    new_model._posterior = inference.fit(new_model, data_to_fit)
    new_model._inference_engine = inference

    return new_model

discriminability 

discriminability(
    params: Params, stimulus: Stimulus
) -> ndarray

Compute scalar discriminability d >= 0 for a (reference, probe) pair

MVP mode: d = sqrt( (probe - ref)^T Σ(ref)^{-1} (probe - ref) ) with Σ(ref) the local covariance at the reference in stimulus space.

Wishart mode (rectangular U design) if extra_dims > 0: d = sqrt( (probe - ref)^T Σ(ref)^{-1} (probe - ref) ) where Σ(ref) is directly computed in stimulus space (input_dim, input_dim) via U(x) @ U(x)^T with U rectangular.

The discrimination task only depends on observable stimulus dimensions. The rectangular U design means local_covariance() already returns the stimulus covariance - no block extraction needed.

Future (full WPPM mode): d is implicit via Monte Carlo simulation of internal noisy responses under the task's decision rule (no closed form). In that case, tasks will directly implement predict/loglik with MC, and this method may be used only for diagnostics.

Parameters:

Name Type Description Default
params dict

Model parameters.

 required
stimulus tuple

(reference, probe) arrays of shape (input_dim,).

 required

Returns:

Name Type Description
d ndarray

Nonnegative scalar discriminability.
Source code in src/psyphy/model/wppm.py 
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
def discriminability(self, params: Params, stimulus: Stimulus) -> jnp.ndarray:
    """
    Compute scalar discriminability d >= 0 for a (reference, probe) pair

    MVP mode:
        d = sqrt( (probe - ref)^T Σ(ref)^{-1} (probe - ref) )
        with Σ(ref) the local covariance at the reference in stimulus space.

    Wishart mode (rectangular U design) if extra_dims > 0:
        d = sqrt( (probe - ref)^T Σ(ref)^{-1} (probe - ref) )
        where Σ(ref) is directly computed in stimulus space (input_dim, input_dim)
        via U(x) @ U(x)^T with U rectangular.

    The discrimination task only depends on observable stimulus dimensions.
    The rectangular U design means local_covariance() already returns
    the stimulus covariance - no block extraction needed.

    Future (full WPPM mode):
        d is implicit via Monte Carlo simulation of internal noisy responses
        under the task's decision rule (no closed form). In that case, tasks
        will directly implement predict/loglik with MC, and this method may be
        used only for diagnostics.

    Parameters
    ----------
    params : dict
        Model parameters.
    stimulus : tuple
        (reference, probe) arrays of shape (input_dim,).

    Returns
    -------
    d : jnp.ndarray
        Nonnegative scalar discriminability.
    """
    ref, probe = stimulus

    # Delta is in stimulus space (input_dim)
    delta = probe - ref

    # Get stimulus covariance at reference
    # (rectangular U design: already returns (input_dim, input_dim))
    Sigma = self.local_covariance(params, ref)

    # Add jitter for stable solve; diag_term is configurable
    jitter = self.diag_term * jnp.eye(self.input_dim)

    # Solve (Σ + jitter)^{-1} delta using a PD-aware solver
    x = jax.scipy.linalg.solve(Sigma + jitter, delta, assume_a="pos")
    d2 = jnp.dot(delta, x)  # quadratic form

    # Guard against tiny negative values from numerical error
    return jnp.sqrt(jnp.maximum(d2, 0.0))

fit 

fit(
    X: ndarray,
    y: ndarray,
    *,
    inference: InferenceEngine | str = "laplace",
    inference_config: dict | None = None,
) -> Model

Fit model to data.

Parameters:

Name Type Description Default
X ndarray

Stimuli, shape (n_trials, 2, input_dim) for (ref, probe) pairs or (n_trials, input_dim) for references only

 required
y ndarray

Responses, shape (n_trials,)

 required
inference InferenceEngine | str

Inference engine or string key ("map", "laplace", "langevin")

 "laplace"
inference_config dict | None

Hyperparameters for string-based inference. Examples: {"steps": 500, "lr": 1e-3} for MAP

 None

Returns:

Type Description
Model

Self for method chaining

Examples:

1
2
>>> # Simple: use defaults
>>> model.fit(X, y)

1
2
3
>>> # Explicit optimizer
>>> from psyphy.inference import MAPOptimizer
>>> model.fit(X, y, inference=MAPOptimizer(steps=500))

1
2
>>> # String + config (for experiment tracking)
>>> model.fit(X, y, inference="map", inference_config={"steps": 500})
Source code in src/psyphy/model/base.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
def fit(
    self,
    X: jnp.ndarray,
    y: jnp.ndarray,
    *,
    inference: InferenceEngine | str = "laplace",
    inference_config: dict | None = None,
) -> Model:
    """
    Fit model to data.

    Parameters
    ----------
    X : jnp.ndarray
        Stimuli, shape (n_trials, 2, input_dim) for (ref, probe) pairs
        or (n_trials, input_dim) for references only
    y : jnp.ndarray
        Responses, shape (n_trials,)
    inference : InferenceEngine | str, default="laplace"
        Inference engine or string key ("map", "laplace", "langevin")
    inference_config : dict | None
        Hyperparameters for string-based inference.
        Examples: {"steps": 500, "lr": 1e-3} for MAP

    Returns
    -------
    Model
        Self for method chaining

    Examples
    --------
    >>> # Simple: use defaults
    >>> model.fit(X, y)

    >>> # Explicit optimizer
    >>> from psyphy.inference import MAPOptimizer
    >>> model.fit(X, y, inference=MAPOptimizer(steps=500))

    >>> # String + config (for experiment tracking)
    >>> model.fit(X, y, inference="map", inference_config={"steps": 500})
    """
    from psyphy.data import ResponseData
    from psyphy.inference import INFERENCE_ENGINES, InferenceEngine

    # Resolve inference engine
    is_string_inference = isinstance(inference, str)

    if is_string_inference:
        config = inference_config or {}
        inference_key: str = inference  # type: ignore[assignment]
        if inference_key not in INFERENCE_ENGINES:
            available = ", ".join(INFERENCE_ENGINES.keys())
            raise ValueError(
                f"Unknown inference: '{inference}'. Available: {available}"
            )
        inference_engine: InferenceEngine = INFERENCE_ENGINES[inference_key](
            **config
        )
    elif isinstance(inference, InferenceEngine):
        inference_engine = inference
    else:
        raise TypeError(
            f"inference must be InferenceEngine or str, got {type(inference)}"
        )

    if inference_config is not None and not is_string_inference:
        raise ValueError(
            "Cannot pass inference_config with InferenceEngine instance"
        )

    # Convert data
    data = ResponseData.from_arrays(X, y)

    # Fit
    self._posterior = inference_engine.fit(self, data)
    self._inference_engine = inference_engine
    self._data_buffer = data  # Initialize buffer
    return self

init_params 

init_params(key: KeyArray) -> Params

Sample initial parameters from the prior.

MVP parameters: {"log_diag": shape (input_dim,)} which defines a constant diagonal covariance across the space.

Returns:

Name Type Description
params dict[str, ndarray]
Source code in src/psyphy/model/wppm.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
def init_params(self, key: jr.KeyArray) -> Params:
    """
    Sample initial parameters from the prior.

    MVP parameters:
        {"log_diag": shape (input_dim,)}
    which defines a constant diagonal covariance across the space.

    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.

MVP mode (basis_degree=None): Σ(x) = diag(exp(log_diag)), constant across x. - Positive-definite because exp(log_diag) > 0.

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: - MVP: {"log_diag": (input_dim,)} - Wishart: {"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 
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
def local_covariance(self, params: Params, x: jnp.ndarray) -> jnp.ndarray:
    """
    Return local covariance Σ(x) at stimulus location x.

    MVP mode (basis_degree=None):
        Σ(x) = diag(exp(log_diag)), constant across x.
        - Positive-definite because exp(log_diag) > 0.

    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:
        - MVP: {"log_diag": (input_dim,)}
        - Wishart: {"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.
    """
    # MVP mode: constant diagonal covariance
    if "log_diag" in params:
        log_diag = params["log_diag"]
        diag = jnp.exp(log_diag)
        return jnp.diag(diag)

    # Wishart mode: spatially-varying covariance
    if "W" in params:
        U = self._compute_U(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 either 'log_diag' (MVP) or 'W' (Wishart)")

log_likelihood 

log_likelihood(
    params: Params,
    refs: ndarray,
    probes: ndarray,
    responses: ndarray,
) -> ndarray

Compute the log-likelihood for arrays of trials.

IMPORTANT: We delegate to the TaskLikelihood to avoid duplicating Bernoulli (MPV) or MC likelihood logic in multiple places.

Parameters:

Name Type Description Default
params dict

Model parameters.

 required
refs (ndarray, shape(N, input_dim))
required
probes (ndarray, shape(N, input_dim))
required
responses (ndarray, shape(N))

Typically 0/1; task may support richer encodings.

 required

Returns:

Name Type Description
loglik ndarray

Scalar log-likelihood (task-only; add prior outside if needed)
Source code in src/psyphy/model/wppm.py 
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
def log_likelihood(
    self,
    params: Params,
    refs: jnp.ndarray,
    probes: jnp.ndarray,
    responses: jnp.ndarray,
) -> jnp.ndarray:
    """
    Compute the log-likelihood for arrays of trials.

    IMPORTANT:
        We delegate to the TaskLikelihood to avoid duplicating Bernoulli (MPV)
        or MC likelihood logic in multiple places.

    Parameters
    ----------
    params : dict
        Model parameters.
    refs : jnp.ndarray, shape (N, input_dim)
    probes : jnp.ndarray, shape (N, input_dim)
    responses : jnp.ndarray, shape (N,)
        Typically 0/1; task may support richer encodings.

    Returns
    -------
    loglik : jnp.ndarray
        Scalar log-likelihood (task-only; add prior outside if needed)
    """
    # We need a ResponseData-like object. To keep this method usable from
    # array inputs, we construct one on the fly. If you already have a
    # ResponseData instance, prefer `log_likelihood_from_data`.
    from psyphy.data.dataset import ResponseData  # local import to avoid cycles

    data = ResponseData()
    # ResponseData.add_trial(ref, probe, resp)
    for r, p, y in zip(refs, probes, responses):
        data.add_trial(r, p, int(y))
    return self.task.loglik(params, data, self, self.noise)

log_likelihood_from_data 

log_likelihood_from_data(
    params: Params, data: Any
) -> ndarray

Compute log-likelihood directly from a ResponseData object.

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

Parameters:

Name Type Description Default
params dict

Model parameters.

 required
data ResponseData

Collected trial data.

 required

Returns:

Name Type Description
loglik ndarray

scalar log-likelihood (task-only; add prior outside if needed)
Source code in src/psyphy/model/wppm.py 
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
def log_likelihood_from_data(self, params: Params, data: Any) -> jnp.ndarray:
    """
    Compute log-likelihood directly from a ResponseData object.

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

    Parameters
    ----------
    params : dict
        Model parameters.
    data : ResponseData
        Collected trial data.

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

log_posterior_from_data 

log_posterior_from_data(
    params: Params, data: Any
) -> ndarray

Convenience helper if you want log posterior in one call (MVP).

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

Returns:

Type Description
jnp.ndarray : scalar log posterior = loglik(params | data) + log_prior(params)
Source code in src/psyphy/model/wppm.py 
558
559
560
561
562
563
564
565
566
567
568
569
def log_posterior_from_data(self, params: Params, data: Any) -> jnp.ndarray:
    """
    Convenience helper if you want log posterior in one call (MVP).

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

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

posterior 

posterior(
    X: ndarray | None = None,
    *,
    probes: ndarray | None = None,
    kind: str = "predictive",
) -> PredictivePosterior | ParameterPosterior

Return posterior distribution.

Parameters:

Name Type Description Default
X ndarray | None

Test stimuli (references), shape (n_test, input_dim). Required for predictive posteriors, optional for parameter posteriors.

 None
probes ndarray | None

Test probes, shape (n_test, input_dim). Required for predictive posteriors.

 None
kind ('predictive', 'parameter')

Type of posterior to return: - "predictive": PredictivePosterior over f(X*) [for acquisitions] - "parameter": ParameterPosterior over θ [for diagnostics]

 "predictive"

Returns:

Type Description
PredictivePosterior | ParameterPosterior

Posterior distribution

Raises:

Type Description
RuntimeError

If model has not been fit yet

Examples:

1
2
3
4
>>> # For acquisition functions
>>> pred_post = model.posterior(X_candidates, probes=X_probes)
>>> mean = pred_post.mean
>>> var = pred_post.variance

1
2
3
>>> # For diagnostics
>>> param_post = model.posterior(kind="parameter")
>>> samples = param_post.sample(100, key=jr.PRNGKey(42))
Source code in src/psyphy/model/base.py 
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
def posterior(
    self,
    X: jnp.ndarray | None = None,
    *,
    probes: jnp.ndarray | None = None,
    kind: str = "predictive",
) -> PredictivePosterior | ParameterPosterior:
    """
    Return posterior distribution.

    Parameters
    ----------
    X : jnp.ndarray | None
        Test stimuli (references), shape (n_test, input_dim).
        Required for predictive posteriors, optional for parameter posteriors.
    probes : jnp.ndarray | None
        Test probes, shape (n_test, input_dim).
        Required for predictive posteriors.
    kind : {"predictive", "parameter"}
        Type of posterior to return:
        - "predictive": PredictivePosterior over f(X*) [for acquisitions]
        - "parameter": ParameterPosterior over θ [for diagnostics]

    Returns
    -------
    PredictivePosterior | ParameterPosterior
        Posterior distribution

    Raises
    ------
    RuntimeError
        If model has not been fit yet

    Examples
    --------
    >>> # For acquisition functions
    >>> pred_post = model.posterior(X_candidates, probes=X_probes)
    >>> mean = pred_post.mean
    >>> var = pred_post.variance

    >>> # For diagnostics
    >>> param_post = model.posterior(kind="parameter")
    >>> samples = param_post.sample(100, key=jr.PRNGKey(42))
    """
    if self._posterior is None:
        raise RuntimeError("Must call fit() before posterior()")

    if kind == "parameter":
        return self._posterior
    elif kind == "predictive":
        if X is None:
            raise ValueError("X is required for predictive posteriors")
        from psyphy.posterior import WPPMPredictivePosterior

        return WPPMPredictivePosterior(self._posterior, X, probes=probes)
    else:
        raise ValueError(
            f"Unknown kind: '{kind}'. Use 'predictive' or 'parameter'."
        )

predict_prob 

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

Predict probability of a correct response for a single stimulus.

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

Parameters:

Name Type Description Default
params dict
required
stimulus (reference, probe)
required

Returns:

Name Type Description
p_correct ndarray
Source code in src/psyphy/model/wppm.py 
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
def predict_prob(self, params: Params, stimulus: Stimulus) -> jnp.ndarray:
    """
    Predict probability of a correct response for a single stimulus.

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

    Parameters
    ----------
    params : dict
    stimulus : (reference, probe)

    Returns
    -------
    p_correct : jnp.ndarray
    """
    return self.task.predict(params, stimulus, self, self.noise)

predict_with_params 

predict_with_params(
    X: ndarray,
    probes: ndarray | None,
    params: dict[str, ndarray],
) -> ndarray

Evaluate model at specific parameter values (no marginalization).

This is useful for: - Threshold uncertainty estimation (evaluate at sampled parameters) - Parameter sensitivity analysis - Debugging and diagnostics

NOT for making predictions (use .posterior() instead, which marginalizes over parameter uncertainty).

Parameters:

Name Type Description Default
X (ndarray, shape(n_test, input_dim))

Test stimuli (references)

 required
probes (ndarray, shape(n_test, input_dim))

Probe stimuli (for discrimination tasks)

 required
params dict[str, ndarray]

Specific parameter values to evaluate at. Keys and shapes depend on the model (e.g., WPPM has "W", "noise_scale", etc.)

 required

Returns:

Name Type Description
predictions (ndarray, shape(n_test))

Predicted probabilities at each test point, given these parameters

Examples:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
>>> # Sample parameters and evaluate
>>> param_post = model.posterior(kind="parameter")
>>> samples = param_post.sample(100, key=jr.PRNGKey(0))
>>>
>>> # Evaluate at first parameter sample
>>> params_0 = {k: v[0] for k, v in samples.items()}
>>> predictions = model.predict_with_params(X_test, probes, params_0)
>>>
>>> # Use for threshold uncertainty estimation
>>> threshold_locs = []
>>> for i in range(100):
...     params_i = {k: v[i] for k, v in samples.items()}
...     preds_i = model.predict_with_params(X_grid, probes_grid, params_i)
...     threshold_idx = jnp.argmin(jnp.abs(preds_i - 0.75))
...     threshold_locs.append(threshold_idx)
Notes

This bypasses the posterior marginalization. For acquisition functions, always use .posterior() which properly accounts for parameter uncertainty.

Source code in src/psyphy/model/base.py 
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
def predict_with_params(
    self,
    X: jnp.ndarray,
    probes: jnp.ndarray | None,
    params: dict[str, jnp.ndarray],
) -> jnp.ndarray:
    """
    Evaluate model at specific parameter values (no marginalization).

    This is useful for:
    - Threshold uncertainty estimation (evaluate at sampled parameters)
    - Parameter sensitivity analysis
    - Debugging and diagnostics

    NOT for making predictions (use .posterior() instead, which
    marginalizes over parameter uncertainty).

    Parameters
    ----------
    X : jnp.ndarray, shape (n_test, input_dim)
        Test stimuli (references)
    probes : jnp.ndarray, shape (n_test, input_dim), optional
        Probe stimuli (for discrimination tasks)
    params : dict[str, jnp.ndarray]
        Specific parameter values to evaluate at.
        Keys and shapes depend on the model (e.g., WPPM has "W", "noise_scale", etc.)

    Returns
    -------
    predictions : jnp.ndarray, shape (n_test,)
        Predicted probabilities at each test point, given these parameters

    Examples
    --------
    >>> # Sample parameters and evaluate
    >>> param_post = model.posterior(kind="parameter")
    >>> samples = param_post.sample(100, key=jr.PRNGKey(0))
    >>>
    >>> # Evaluate at first parameter sample
    >>> params_0 = {k: v[0] for k, v in samples.items()}
    >>> predictions = model.predict_with_params(X_test, probes, params_0)
    >>>
    >>> # Use for threshold uncertainty estimation
    >>> threshold_locs = []
    >>> for i in range(100):
    ...     params_i = {k: v[i] for k, v in samples.items()}
    ...     preds_i = model.predict_with_params(X_grid, probes_grid, params_i)
    ...     threshold_idx = jnp.argmin(jnp.abs(preds_i - 0.75))
    ...     threshold_locs.append(threshold_idx)

    Notes
    -----
    This bypasses the posterior marginalization. For acquisition functions,
    always use .posterior() which properly accounts for parameter uncertainty.
    """
    return self._forward(X, probes, params)