Quickstart

Active-learning-driven adaptive experimentation in psychophysics

Installation | Documentation | Examples | Contributing

---

Quick-start walkthrough — fit your first covariance ellipse

The snippet below shows the minimal end-to-end workflow: simulate a handful of oddity-task trials at a single reference point, fit the WPPM with MAP optimization, and visualize the result. No GPU needed — runs in under 2 min on CPU.

The complete runnable script is quick_start.py. A step-by-step explanation lives in the Quick-start example.

Imports

from psyphy.data import TrialData  # batched trial container
from psyphy.inference import MAPOptimizer  # fitter
from psyphy.model import (
    WPPM,
    GaussianNoise,
    OddityTask,
    OddityTaskConfig,
    Prior,
    WPPMCovarianceField,  # fast Σ(x) evaluation
)
from psyphy.viz import plot_ellipses  # covariance-ellipse plotting

Compute settings

MC_SAMPLES = 50  # MC samples per trial in the likelihood (full example: 500)
NUM_TRIALS = 1000  # total simulated trials (full example: 4000 × 25)
NUM_STEPS = 600  # optimizer steps (full example: 2000)

learning_rate = 1e-4  # full example: 5e-5. The smaller the lr, the more steps
# are required.
#
# These four are not free choices -- they were swept (across both the jax
# version floor CI pins, jax==0.4.28, and a current jax, and across several
# seeds each) for the cheapest setting that recovers the ground truth reliably:
# fitted/true ellipse area ratio within roughly [0.6, 1.6] and no sign of
# collapse or divergence, ~4-13 s CPU depending on jax version.
#
#   * NUM_TRIALS is the binding constraint. All trials sit at a *single*
#     reference point, so the likelihood is weak, and the outcome at low
#     trial counts is not a clean "collapses to zero" failure -- it is
#     chaotic: depending on the exact random draw (which differs across jax
#     versions for the same PRNGKey, not just across seeds) the same trial
#     count can collapse toward zero, overshoot by several times, or blow up
#     numerically. Below ~500 trials this instability shows up often enough
#     to make the fit unreliable; 1000 clears it with margin on every
#     jax/seed combination tested.
#   * NUM_STEPS below ~600 overshoots rather than under-fits.
#   * learning_rate above ~2e-3 diverges to a non-finite loss within 25 steps.
#
# MAPOptimizer defaults to reduction="mean" (a per-trial objective), so this
# learning rate does not need rescaling when you change NUM_TRIALS.

Ground-truth model + simulate data

task = OddityTask(config=OddityTaskConfig(num_samples=int(MC_SAMPLES)))
noise = GaussianNoise(sigma=0.1)

# Set all Wishart process hyperparameters in Prior
truth_prior = Prior()
truth_model = WPPM(
    prior=truth_prior,
    likelihood=task,
    noise=noise,
)

# Sample ground-truth Wishart process weights
truth_params = truth_model.init_params(jax.random.PRNGKey(123))

# Simulate observed responses using the likelihood implied by the task
ys, prob_params = task.simulate(truth_params, stimuli, truth_model, key=k_sim)
p_correct = prob_params[0]  # <- for Bernoulli tasks, p_correct is the only prob_param

Build model and fit

prior = Prior()

model = WPPM(
    prior=prior,
    likelihood=task,
    noise=noise,  # we use the same Gaussian noise as for the ground truth
)

inference = MAPOptimizer(
    steps=NUM_STEPS,
    learning_rate=learning_rate,
    track_history=True,
    log_every=1,
)

map_estimate = inference.fit(model, data, init_params=init_params, seed=4)
# Protocol: ParameterPosterior, here point estimate

# optional: for visualization:
map_cov_field = WPPMCovarianceField(model, map_estimate.params)
# OUTPUT: Covariance Matrices (N, 2, 2) for plotting

Results

Ground truth (black), prior sample (blue), and MAP-fitted (red) covariance ellipses at the single reference point.

Negative log-likelihood over optimizer steps.