Merge pull request #131 from hmd101/likelihood-api

### Likelihood refactor:

- subclasses of likelihood now only need to implement predict, because log_likelihood vmaps over predict
- added a simulate method that creates synthetic data based on a task likelihood

Author: hmd101

docs/examples/wppm/full_wppm_fit_example.md

@@ -69,8 +69,7 @@ For OddityTask, we store trials as (ref, comparison) even though the task involv
 
 Note on data used in this script: here, we simulate data (and hence have a ground truth to compare agains). To how to conveniently simulate data yourself, checkout  the [script](https://github.com/flatironinstitute/psyphy/blob/main/docs/examples/wppm/full_wppm_fit_example.py)
 
-<details>
-<summary>### 2 ways of representing data in `psyphy` (important)</summary>
+### 2 ways of representing data in `psyphy` (important)
 `psyphy` provides two lightweight containers for trial data (defined in [`src/psyphy/data/dataset.py`](https://github.com/flatironinstitute/psyphy/blob/main/src/psyphy/data/dataset.py)):
 
 **`TrialData` (compute-first; used for fitting):**
@@ -96,7 +95,7 @@ Avoid repeatedly converting Python lists -> JAX arrays inside tight loops.
 
 
 Note that here, we simlulate data, for details check out  [`full_wppm_fit_example.py`](full_wppm_fit_example.py) directly resulting in a `TrialData` object.
-</details>
+
 
 ---
 
@@ -109,8 +108,7 @@ The WPPM parameters are basis weights stored as a dict:
 where `W` is a tensor of Chebyshev-basis coefficients.
 
 
-<details>
-  <summary>### Prior distribution over weights:</summary>
+### Prior distribution over weights:
 
 `Prior.sample_params(key)` samples weights `W` from a **zero-mean Gaussian** with a *degree-dependent variance*.
 
@@ -146,7 +144,7 @@ W_{ijde} \sim \mathcal{N}(0, \sigma^2_{ij}).
 \]
 
 This is the  state of the  WPPM: **before any data**, WPPM draws smooth random fields because high-frequency coefficients are shrunk by the decay.
-</details>
+
 
 ---
 
@@ -240,43 +238,14 @@ To see how we generate the covariance field figures, checkout the plotting code
 ---
 
 ## To recap: Minimal recipe (copy/paste mental model)
+We are Bayesian, so we need to define the Prior and Likelihood and choose an inference method (here MAP) that will hand us the posterior distribution over the parameters.
 
-To use WPPM on your own data, these are the essential calls:
-
-**1. Create** task + noise + prior:
-
-   - `task = OddityTask()`
-
-   - `noise = GaussianNoise(sigma=...)`
-
-   - `prior = Prior(input_dim=..., basis_degree=..., extra_embedding_dims=..., decay_rate=..., variance_scale=...)`
-
-**2. Create** WPPM:
-
-   - `model = WPPM(input_dim=..., prior=prior, task=task, noise=noise, diag_term=...)`
-
-**3. Initialize** parameters:
-
-   - `params0 = model.init_params(jax.random.PRNGKey(...))`  (draws from `Prior.sample_params`)
-
-**4. Load/build** a dataset:
-
-   - `data = TrialData(refs=..., comparisons=..., responses=...)`
-
-**5. Fit**:
-
-   - `map = MAPOptimizer(...).fit(model, data, init_params=params0, ...)`
-
-**6. Inspect** $\Sigma(x)$:
-
-   - `field = WPPMCovarianceField(model, map.params)`
-   - `Sigmas = field(xs)`
-
+For an even more minimal code setup that doesn't require a GPU but will run on your CPU in < 1 min, you may find [`quickstart`](https://flatironinstitute.github.io/psyphy/examples/wppm/quick_start/) helpful.
 ---
 
 ## Notes and pitfalls
 
-- **CPU vs GPU:** this example can be heavy because the oddity likelihood uses Monte Carlo. A GPU can help a lot.
+- **CPU vs GPU:** this example can be heavy because the oddity likelihood uses Monte Carlo. A GPU can help a lot, see [`quickstart`](https://flatironinstitute.github.io/psyphy/examples/wppm/quick_start/) for a CPU friendly version.
 - **Positive definiteness:** `diag_term` is important. If you ever see a non-PD covariance, increase `diag_term` slightly.
 - **MC variance:** optimization stability depends on `MC_SAMPLES`. Too small means noisy gradients.
 
@@ -308,7 +277,7 @@ instead of using relative filesystem paths.
 - MAP fitting: [`src/psyphy/inference/map_optimizer.py`](https://github.com/flatironinstitute/psyphy/blob/main/src/psyphy/inference/map_optimizer.py) (see `MAPOptimizer`)
 - Data container: [`src/psyphy/data/dataset.py`](https://github.com/flatironinstitute/psyphy/blob/main/src/psyphy/data/dataset.py) (see `ResponseData`)
 
-If you want to “follow the call graph”:
+If you want to "follow the call graph":
 
 1. `WPPM.init_params(...)` (defined in [`src/psyphy/model/wppm.py`](https://github.com/flatironinstitute/psyphy/blob/main/src/psyphy/model/wppm.py)) → delegates to the prior’s `Prior.sample_params(...)` (defined in [`src/psyphy/model/prior.py`](https://github.com/flatironinstitute/psyphy/blob/main/src/psyphy/model/prior.py)).
 2. `OddityTask.predict_with_kwargs(...)` / `OddityTask.loglik(...)` (defined in [`src/psyphy/model/likelihood.py`](https://github.com/flatironinstitute/psyphy/blob/main/src/psyphy/model/likelihood.py)) → calls into the model to get $\Sigma(x)$ and then runs the task’s decision rule (Monte Carlo in the full model).

docs/examples/wppm/full_wppm_fit_example.py

@@ -209,7 +209,7 @@ def _cov_to_points(cov: jnp.ndarray, center: jnp.ndarray) -> jnp.ndarray:
 #
 
 
-# --8<-- [start:simulate_data]
+##### Simulate data
 num_trials_per_ref = NUM_TRIALS_Per_Ref  # (trials per reference point)
 n_ref_grid = 5  # NUM_GRID_PTS
 ref_grid = jnp.linspace(-1, 1, n_ref_grid)  # [-1,1] space
@@ -234,7 +234,7 @@ def _cov_to_points(cov: jnp.ndarray, center: jnp.ndarray) -> jnp.ndarray:
 Sigmas_ref = truth_field(refs)  # (N, 2, 2)
 
 # Sample unit directions on the circle.
-k_dir, k_pred, k_y = jr.split(key, 3)
+k_dir, k_sim = jr.split(key)
 angles = jr.uniform(k_dir, shape=(num_trials_total,), minval=0.0, maxval=2.0 * jnp.pi)
 unit_dirs = jnp.stack([jnp.cos(angles), jnp.sin(angles)], axis=1)  # (N, 2)
 
@@ -250,30 +250,9 @@ def _cov_to_points(cov: jnp.ndarray, center: jnp.ndarray) -> jnp.ndarray:
 deltas = MAHAL_RADIUS * jnp.einsum("nij,nj->ni", L, unit_dirs)  # (N, 2)
 comparisons = jnp.clip(refs + deltas, -1.0, 1.0)
 
-# Compute p(correct) in batch. We vmap the single-trial predictor.
-trial_pred_keys = jr.split(k_pred, num_trials_total)
-
-
-# we use task as the generative model to create observations (user responses)
-def _p_correct_one(ref: jnp.ndarray, comp: jnp.ndarray, kk: jnp.ndarray) -> jnp.ndarray:
-    # Task MC settings (num_samples/bandwidth) come from OddityTaskConfig.
-    # Only the randomness is threaded dynamically.
-    return task._simulate_trial_mc(
-        params=truth_params,
-        ref=ref,
-        comparison=comp,
-        model=truth_model,
-        noise=truth_model.noise,
-        num_samples=int(task.config.num_samples),
-        bandwidth=float(task.config.bandwidth),
-        key=kk,
-    )
-
-
-p_correct = jax.vmap(_p_correct_one)(refs, comparisons, trial_pred_keys)
-
-# Sample observed y ~ Bernoulli(p_correct) in batch.
-ys = jr.bernoulli(k_y, p_correct, shape=(num_trials_total,)).astype(jnp.int32)
+# --8<-- [start:simulate_data]
+# Simulate observed responses using the likelihood implied by the task.
+ys, p_correct = task.simulate(truth_params, refs, comparisons, truth_model, key=k_sim)
 
 # Build the canonical batched dataset for compute.
 #

docs/examples/wppm/plots/quick_start_ellipses.png

@@ -134,7 +134,7 @@ def _cov_to_points(cov: jnp.ndarray, center: jnp.ndarray) -> jnp.ndarray:
 # Step 2 — Simulate data at a *single* reference point
 # ---------------------------------------------------------------------------
 
-# --8<-- [start:simulate_data]
+
 # Single reference point at the centre of the stimulus space.
 ref_point = jnp.array([[0.0, 0.0]])  # shape (1, 2) — kept as a batch for generality
 
@@ -149,7 +149,7 @@ def _cov_to_points(cov: jnp.ndarray, center: jnp.ndarray) -> jnp.ndarray:
 Sigmas_ref = truth_field(refs)  # (NUM_TRIALS, 2, 2)
 
 # Sample unit directions and build covariance-scaled probe displacements.
-k_dir, k_pred, k_y = jr.split(key, 3)
+k_dir, k_sim = jr.split(key)
 angles = jr.uniform(k_dir, shape=(NUM_TRIALS,), minval=0.0, maxval=2.0 * jnp.pi)
 unit_dirs = jnp.stack([jnp.cos(angles), jnp.sin(angles)], axis=1)  # (N, 2)
 
@@ -160,27 +160,9 @@ def _cov_to_points(cov: jnp.ndarray, center: jnp.ndarray) -> jnp.ndarray:
 deltas = MAHAL_RADIUS * jnp.einsum("nij,nj->ni", L, unit_dirs)  # (N, 2)
 comparisons = jnp.clip(refs + deltas, -1.0, 1.0)
 
-# Compute p(correct) via MC simulation of the oddity task.
-trial_pred_keys = jr.split(k_pred, NUM_TRIALS)
-
-
-def _p_correct_one(ref: jnp.ndarray, comp: jnp.ndarray, kk: jnp.ndarray) -> jnp.ndarray:
-    return task._simulate_trial_mc(
-        params=truth_params,
-        ref=ref,
-        comparison=comp,
-        model=truth_model,
-        noise=truth_model.noise,
-        num_samples=int(task.config.num_samples),
-        bandwidth=float(task.config.bandwidth),
-        key=kk,
-    )
-
-
-p_correct = jax.vmap(_p_correct_one)(refs, comparisons, trial_pred_keys)
-
-# Sample observed responses y ~ Bernoulli(p_correct).
-ys = jr.bernoulli(k_y, p_correct, shape=(NUM_TRIALS,)).astype(jnp.int32)
+# --8<-- [start:simulate_data]
+# Simulate observed responses using the likelihood implied by the task
+ys, p_correct = task.simulate(truth_params, refs, comparisons, truth_model, key=k_sim)
 # --8<-- [end:simulate_data]
 
 # --8<-- [start:data]