Quick start — fit your first covariance ellipse

Goal: run the full psyphy workflow — simulate data, fit a model, inspect the result — in one short script with no GPU required.

The complete runnable script is quick_start.py. For a spatially-varying field over a 2-D stimulus grid, see the full example.

---

Runtime

Hardware	Approximate time
GPU (any modern CUDA device)	< 5 s
CPU (laptop / M-series Mac)	< 1 min

The three knobs that control runtime:

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.

---

Step 0 — 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

---

Step 1 — Define a ground-truth model and sample parameters

We create a WPPM with known parameters to act as the synthetic observer. Data will be generated from it so we have a ground truth to compare against.

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

---

Step 2 — Simulate trials at a single reference point

We generate NUM_TRIALS oddity-task responses at a single reference stimulus ref = [0, 0]. Probe displacements are scaled by the local covariance (constant Mahalanobis radius), so trial difficulty stays roughly uniform.

# 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

The TrialData container is the canonical input for fitting:

data = TrialData(stimuli=stimuli, responses=ys)  # contains 2 JAX arrays

---

Step 3 — Build the model to fit

We build a fresh WPPM with the same hyperparameters but independent random weights, then take one draw from the prior as the starting point for optimization.

prior = Prior()

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

# Initialize parameters at a sample from the prior
init_params = model.init_params(
    jax.random.PRNGKey(42)
)  # intitialize with a draw from the prior
prior_field = WPPMCovarianceField(model, init_params)
# Evaluate prior covariance at the reference point
covs_prior = prior_field(ref_point)  # (1, 2, 2)

---

Step 4 — Fit with MAP optimization

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

MAPOptimizer runs SGD + momentum and returns a MAPPosterior — a point estimate at \(W_\text{MAP}\).

---

Step 5 — Inspect the fitted covariance ellipse

WPPMCovarianceField binds a (model, params) pair into a single callable that returns \(\Sigma(x)\) for any stimulus x:

# Evaluate any covariance-field object at a single point or a batch of points.
covs_truth = truth_field(ref_point)  # (N, 2, 2)
covs_prior = prior_field(ref_point)  # (N, 2, 2)
covs_map = map_cov_field(ref_point)  # (N, 2, 2)
# here: N=1 for fast computation

The ellipse plot below overlays the ground truth (black), the prior initialization (blue), and the MAP fit (red) at the single reference point:

# All three fields in one call, at true size. `scale="auto"` is for *grids* of
# reference points -- it needs at least two centers to measure their spacing --
# and there is only one here. At this model's scale the ellipses are readable
# unmagnified anyway, which also means the figure can be read for absolute size.
plot_ellipses(
    ref_point,
    [covs_truth, covs_prior, covs_map],
    ax=ax,
    scale=1.0,
    colors=["k", "b", "r"],
    linestyles=["-", "--", "-"],
    linewidths=[2.0, 1.5, 2.0],
    labels=["Ground Truth", "Prior Sample (init)", "Fitted (MAP)"],
    alpha=0.8,
    show_centers=True,
)

Ground truth (black), prior sample (blue), and MAP-fitted (red) covariance ellipses at the single reference point. Note how the red ellipse has moved from the blue (init) closer to the ground truth (black).

---

Step 6 — Learning curve

steps_hist, loss_hist = inference.get_history()

Negative log-likelihood over optimizer steps.

---

Next steps

• Spatially-varying field: scale up to a full 2-D grid → full example.
• Your own data: replace the simulated TrialData with your own inputs and responses arrays. (currently calculations are only supported when the inputs are "refs" and "comparisons.")
• API reference: see MAPOptimizer, WPPM, and WPPMCovarianceField.