Add MRP model

BENCHMARKING.md

@@ -23,7 +23,7 @@ Each model is wrapped in an asv `ModelBench` class that captures four metrics:
 - **`time_eval`** — steady-state per-call time, measured by asv's native timing
   machinery.
 
-The four metrics are tracked on a curated subset of 25 models listed in
+The four metrics are tracked on a curated subset of 26 models listed in
 [`BENCHMARK_CORE.md`](./BENCHMARK_CORE.md), chosen to give broad coverage
 (hierarchical, GP, scan, linear algebra, survival, ODE, mixtures, discrete).
 

BENCHMARK_CORE.md

@@ -1,10 +1,10 @@
-# Benchmark Core (25 models)
+# Benchmark Core (26 models)
 
 A curated subset of the catalogue chosen to give diverse coverage for
 benchmarking and optimization research: small ↔ large, centered ↔ noncentered
 hierarchical, dense linear algebra, `scan`-based time series and ODEs, GPs,
-mixtures, spatial CAR, survival, and one discrete-vars model so the logp-only
-codepath is exercised.
+mixtures, spatial CAR, survival, wide gather-heavy survey models, and one
+discrete-vars model so the logp-only codepath is exercised.
 
 ## Tiny / canonical baselines
 1. `models/eight_schools_noncentered.py` — smallest hierarchical, noncentered reparameterization
@@ -44,5 +44,9 @@ codepath is exercised.
 23. `models/nyc_bym_traffic.py` — CAR/BYM spatial, large
 24. `models/excess_deaths.py` — larger real dataset, structural model
 
+## Survey / categorical
+25. `models/mrp_survey_zerosum_categorical.py` — many zero-sum effect tables gathered
+    per respondent, Categorical likelihood over a wide (N, G, K) logit array
+
 ## Discrete free variables
-25. `models_discrete/occupancy_crossbill.py` — exercises the logp-only path
+26. `models_discrete/occupancy_crossbill.py` — exercises the logp-only path

README.md

@@ -106,7 +106,7 @@ Authoring conventions and the extraction template live in
 ## Benchmarking & the timeline
 
 The catalogue backs an ASV suite that tracks four metrics — `rewrite_time`,
-`compile_time`, `n_rewrites`, `time_eval` — on a curated 25-model core
+`compile_time`, `n_rewrites`, `time_eval` — on a curated 26-model core
 ([`BENCHMARK_CORE.md`](./BENCHMARK_CORE.md)) across every pymc release, published
 to a [dashboard](https://pymc-devs.github.io/pymc-model-catalogue/dashboard.html).
 A separate [experiments](https://pymc-devs.github.io/pymc-model-catalogue/experiments.html)

models/mrp_survey_zerosum_categorical.py

@@ -0,0 +1,68 @@
+"""
+Model: Survey MRP step - zero-sum demographic effects with a categorical rating likelihood
+Source: synthetic stress test written for this catalogue, shaped after an MRP-style survey
+    step: a "feeling thermometer" battery where every respondent rates every party on a
+    fixed answer scale.
+Authors: pymc-model-catalogue
+Description: N = 2500 respondents rate G = 8 parties on a K = 10 point scale. The latent
+    rating per (respondent, question) sums 10 demographic main effects and 9 pairwise
+    interactions, each a ZeroSumNormal (levels x G) table scaled by its own HalfNormal and
+    gathered at the respondent's category. A monotone loading built from the cumulative sums
+    of a Dirichlet maps that latent onto K logits, which are added to zero-sum per-question
+    base rates and fed to a Categorical likelihood over (N, G, K). 1429 parameters.
+
+Benchmark results:
+- Original:  logp = -47408.0753, grad norm = 146.9804, 7460.3 us/call (2069 evals)
+- Frozen:    logp = -47408.0753, grad norm = 146.9804, 7655.0 us/call (2036 evals)
+"""
+
+import numpy as np
+import pymc as pm
+import pytensor.tensor as pt
+
+
+def build_model():
+    N, G, K = 2500, 8, 10  # respondents, rating questions each answers, scale points
+    MAIN = {"age": 7, "gender": 2, "education": 4, "nationality": 3, "urban": 2,
+            "income": 6, "region": 24, "methods": 3, "vote_intent": 8, "party_pref": 10}
+    INTER = [("age", "education"), ("gender", "nationality"), ("age", "gender"),
+             ("age", "nationality"), ("education", "nationality"), ("nationality", "urban"),
+             ("gender", "urban"), ("age", "urban"), ("education", "urban")]
+
+    rng = np.random.default_rng(0)
+    codes = {d: rng.integers(0, L, size=N) for d, L in MAIN.items()}
+    y = rng.integers(0, K, size=(N, G))  # respondent i's scale point for question j
+
+    with pm.Model(check_bounds=False) as model:
+        # 1. hierarchical block: one zero-sum effect table per demographic, gathered per respondent
+        mu = pt.zeros((N, G))
+        tables = [(d, MAIN[d], codes[d]) for d in MAIN]
+        tables += [(f"{a}_x_{b}", MAIN[a] * MAIN[b], codes[a] * MAIN[b] + codes[b])
+                   for a, b in INTER]
+        for name, levels, idx in tables:
+            sd = pm.HalfNormal(f"sd_{name}", 0.3)  # how much this demographic matters at all
+            off = pm.ZeroSumNormal(f"a_{name}", sigma=1.0, shape=(levels, G), n_zerosum_axes=1)
+            mu = mu + (sd * off)[idx]  # (N, G) latent rating
+        # No per-question intercept: adding one shifts the logits by c_g * phi_g, whose
+        # mean over the scale cancels under the softmax and whose remainder is already
+        # in alpha's zero-sum space, so it would be a redundant degree of freedom.
+
+        # 2. likelihood: a monotone loading phi turns the single latent into K logits.
+        # phi_g = (0, ..., 1) from a Dirichlet's cumulative sums, so mu orders the scale points
+        # rather than shifting all of them equally (a constant shift cancels under the softmax).
+        phi_diffs = pm.Dirichlet("phi_diffs", np.ones(K - 1), shape=(G, K - 1))
+        phi = pt.concatenate([pt.zeros((G, 1)), pt.cumsum(phi_diffs, axis=-1)], axis=-1)  # (G, K)
+        # Per-question base rates, zero-sum over the scale so they don't fight the softmax.
+        alpha = pm.ZeroSumNormal("alpha", sigma=2.0, shape=(G, K), n_zerosum_axes=1)
+        pm.Categorical("y", logit_p=alpha[None, :, :] + phi[None, :, :] * mu[:, :, None],
+                       observed=y)
+
+    ip = model.initial_point()
+    model.rvs_to_initial_values = {rv: None for rv in model.free_RVs}
+    return model, ip
+
+
+if __name__ == "__main__":
+    from _benchmark import run_benchmark
+
+    run_benchmark(build_model)

scripts/core_models.py

@@ -8,11 +8,11 @@
 from __future__ import annotations
 
 CORE_MODELS: tuple[str, ...] = (
-    # Trimmed to 10 diverse models, picked to keep broad coverage
-    # across the 25-model catalogue described in BENCHMARK_CORE.md
-    # while keeping a single backfill dispatch cheap enough to finish
-    # in reasonable time on a 2-vCPU runner. One per category:
-    # tiny/hierarchical/linalg/time-series/GP/mixture/ODE/discrete.
+    # Trimmed to 11 diverse models, picked to keep broad coverage
+    # across the catalogue described in BENCHMARK_CORE.md while keeping
+    # a single backfill dispatch cheap enough to finish in reasonable
+    # time on a 2-vCPU runner. One per category:
+    # tiny/hierarchical/linalg/time-series/GP/mixture/survey/discrete.
     "models.eight_schools_noncentered",                     # tiny hierarchical
     "models.BEST",                                          # trivial two-group
     "models.GLM_hierarchical_binomial_rat_tumor",           # hierarchical binomial
@@ -22,6 +22,7 @@
     "models.bayesian_var_ireland",                          # VAR (scan + linalg)
     "models.gp_marginal_matern52",                          # small marginal GP
     "models.marginalized_gaussian_mixture_model",           # mixture
+    "models.mrp_survey_zerosum_categorical",                # zero-sum tables + categorical
     "models_discrete.occupancy_crossbill",                  # logp-only discrete path
 )