Spec: didactic marimo notebook recreating arXiv:2608.04664 (soft-noncrossing Bayesian panel QR) with GPJax v1.0 + NumPyro

[thomaspinder]

Problem Statement

I want to recreate Huber, Poon & Zhu (2026), "Soft-Noncrossing Bayesian Panel Quantile Regression for Measuring Climate Tail Risk" (arXiv:2608.04664), as a self-contained, didactic marimo notebook built on the GPJax v1.0 branch API together with NumPyro. The paper ships no replication package and leaves many estimation settings unstated (priors, Bernstein degree, MCMC configuration, covariate scaling, temperature-data variant), so recreating it requires a documented set of judgment calls. The notebook is a proof of concept that the v1.0 API composes into a serious hierarchical Bayesian workflow; it will migrate out of the GPJax repo once v1.0 merges to main.

Solution

A single marimo notebook, confined with its docs and data cache to one self-contained directory, that: downloads and caches both public datasets automatically; constructs the paper's four climate shocks from raw series; implements the full §2 model — Gaussian-process-smoothed coefficient paths over a 99-node quantile grid, a componentwise-monotone Bernstein population path, unit deviations delivering soft noncrossing, a quantile-invariant AR(1) common time effect, and the asymmetric-Laplace working likelihood — fits it with NUTS at horizons h=1 and h=4 for all 33 countries, and reproduces the paper's headline outputs: DIRF paths for the climate shock quartet, ΔGaR, and ΔES(0.05) with EM/AE comparisons. Every deviation from the paper is called out inline and recorded in an ADR.

User Stories

1. As a GPJax user, I want a runnable notebook recreating a published hierarchical Bayesian model end-to-end, so that I can see the v1.0 API composing with NumPyro in a realistic workflow.
2. As a reader, I want both datasets downloaded and cached automatically on first run, so that I can execute the notebook from a clean checkout with no manual data steps.
3. As a reader on a normal connection, I want the temperature series obtained by ranged extraction of the single required archive member, so that I never download the 2.7 GB climate bundle.
4. As an econometrician, I want the climate shock quartet constructed exactly by the paper's recipe (per-country AR(P=8, h=8) filter, four-quarter RMS volatility, PPP-GDP-weighted global component with the local component as residual, unit-variance standardisation), so that the regressors match eq. (9).
5. As a reader, I want a math-to-code map written in the project glossary's language (population path, unit deviation, common time effect, working likelihood), so that every code object traces to its equation.
6. As a reader, I want the population path monotone by construction through its parametrisation, so that soft noncrossing is visible in the prior rather than merely asserted.
7. As a skeptical reviewer, I want a synthetic-panel validation with known ground truth run through the same model code, so that I can see parameter recovery and a low realized crossing rate before trusting real-data results.
8. As a practitioner, I want estimation on the paper's full 99-node quantile grid, so that credible bands are comparable with the paper's figures without composite-likelihood reweighting.
9. As a macroeconomist, I want the common time effect estimated as a sum-to-zero latent AR(1) driven by the oil price and a PPP-GDP-weighted global short rate, so that aggregate shocks are separated from heterogeneous country responses as in the paper.
10. As a user with a time budget, I want a pilot timing cell that runs a truncated warmup in the exact production configuration and reports a lower-bound runtime projection, so that I am never trapped in a surprise multi-hour run.
11. As the notebook owner, I want the run to stop and ask when the projection exceeds two hours, so that budget decisions remain mine.
12. As a returning user, I want fitted posteriors cached to disk and the fits gated behind a run button, so that reopening the notebook never silently refits.
13. As a reader, I want convergence diagnostics — divergences, R-hat, ESS — plus a realized quantile-crossing rate, so that inference quality is assessable at a glance.
14. As a researcher, I want DIRF plots for all four climate shocks with posterior medians and 68/90% bands at h=1 and h=4, so that I can compare against the paper's Figures 4–8.
15. As a reader, I want a per-country results selector that only re-indexes stored posterior draws, so that exploring countries is instant and never triggers recomputation.
16. As a policy analyst, I want ΔGaR at τ=0.05 and 0.10 and ΔES(0.05) for all 33 countries with EM and AE averages, so that the paper's headline tail-risk numbers (−0.530 / −0.395 pp) can be checked.
17. As a careful reader, I want explicit callouts wherever the notebook deviates from the paper — prior measure on the monotone cone, sampler choice, the Lemma-1 domain caveat for signed climate regressors, composite-likelihood calibration — so that differences are explained rather than discovered.
18. As the repo owner, I want the notebook, its glossary, its ADRs, and its data cache confined to one directory, so that the whole context migrates out cleanly when v1.0 merges.
19. As a maintainer, I want the notebook to import the in-repo GPJax rather than a PyPI release, so that it exercises the branch API.
20. As a reproducibility-minded reader, I want fixed seeds and a footer recording package versions and wall time, so that results can be reproduced and cited.
21. As a downstream user, I want data-source citations and licenses stated (GVAR: CC BY 4.0 Mendeley; Gortan et al. 2024, Sci Data), so that reuse obligations are clear.
22. As a future implementer, I want ADRs recording the sampler and prior-measure decisions, so that nobody "fixes" deliberate deviations.
23. As a marimo user, I want the notebook to be a plain reactive Python module, so that I can open it with marimo edit, run it headlessly, or read it as source.
24. As an experimenter, I want every knob — grid, Bernstein degree, horizons, chains, draws, dataset variant strings — in one configuration cell, so that scale-up choices are one edit away.
25. As a researcher, I want the horizon set extensible to h=1…10 behind a constant, so that the paper's full local-projection sweep is reachable without code changes.
26. As a reader comparing to the paper, I want the expected-shortfall approximation to use exactly the grid nodes at or below α=0.05, so that eq. (17) is computed identically to the paper.

Implementation Decisions

• Single marimo notebook in a dedicated, self-contained directory alongside its own glossary (CONTEXT.md), ADRs, and gitignored data cache; no changes to the GPJax library, its lockfile, root glossary, root ADRs, CI, or docs build. marimo is run via an ephemeral uv run --with marimo invocation against the repo environment.

• GPJax v1.0 branch API (in-repo checkout, version metadata 0.18.0); GPJax supplies the Matern12 kernel gram over the quantile grid (the paper's exponential kernel, K(τ,τ)=1 at unit variance). NumPyro supplies distributions, transforms, and NUTS.

• Inference by NUTS on the full joint posterior (per ADR-0001 recorded in the notebook directory): the asymmetric-Laplace working likelihood is used directly via AsymmetricLaplaceQuantile (no Kozumi–Kobayashi augmentation), and the monotone cone is reparametrized away with an ordered transform rather than sampled as a truncated MVN (Botev). Algorithm-level fidelity to the paper's Gibbs sampler is a non-goal.

• Known deviation, to be recorded as ADR-0002: the ordered-transformed Gaussian prior on Bernstein rows has the same support as the paper's truncated MVN on the monotone cone but a different measure.

• Model structure per the paper's §2 and eq. (9): p=8 coefficient paths per country (intercept, 4 climate responses, 3 macro slopes); separable prior covariance (exponential kernel over the quantile grid ⊗ within-quantile covariance Σ); componentwise-monotone Bernstein population path over all 8 rows (paper-faithful; the noncrossing guarantee holds only on nonnegative regressor support, called out inline); quantile-invariant sum-to-zero AR(1) common time effect with w = [oil price, PPP-GDP-weighted global short rate]; a fully separate model per horizon.

• Prior wiring validated in a session prototype (0 divergences, factor recovery corr 0.97, population path monotone in every draw); the decision-rich core:

  K_lam = Matern12(lengthscale=lam, variance=1.0).gram(taus[:, None]).as_matrix()
  Sigma = (sd[:, None] * L_corr) @ (sd[:, None] * L_corr).T      # LKJ(2) + HalfNormal(0.25) scales
  Gamma = sample(TransformedDistribution(Normal(0, 2).expand([p, M + 1]).to_event(2), OrderedTransform()))
  Theta = sample(MultivariateNormal(mu.reshape(-1), kron(K_lam, Sigma)).expand([N]))
  y ~ AsymmetricLaplaceQuantile(loc=q, scale=sigma_tau, quantile=taus)   # one term per (i, t, τ)

• Settings resolved during the design interview (the paper states none of these): quantile grid = the paper's full 99 nodes (0.01…0.99); Bernstein degree M=10 fixed; priors — ordered Normal(0,2) Bernstein rows, Σ = LKJ(2) with HalfNormal(0.25) scales (small scales deliver soft noncrossing), lengthscale LogNormal(log 0.3, 0.5) sampled, quantile-specific ALD scales HalfNormal(2), AR coefficient Uniform(−0.99, 0.99), predictor loadings N(0,1), factor innovation HalfNormal(0.2) kept tight per the paper's §3.2 remark.

• Data: GVAR 2023 vintage (Mendeley, DOI 10.17632/kfp5fhgkvf.1) for the macro panel, PPP-GDP weights, and oil price; Gortan et al. temperature data as ERA5, unweighted, country-level (gadm0), monthly, obtained by HTTP-Range extraction of the single CSV member from the Figshare archive (verified 206 support); the GCS bucket behind the project dashboard is private and is not used.

• Preprocessing: deseasonalisation by monthly-dummy regression; macro covariates (lagged growth, inflation, Δ real exchange rate) mapped to [0,1] by per-country empirical rank transform; climate quartet standardised, not rank-transformed; outcome is cumulative annualised growth per eq. (3).

• Estimation defaults: all 33 countries; horizons h ∈ {1, 4}; 4 parallel chains × (2500 warmup + 1500 draws), target_accept 0.9, float64. Pilot timing cell runs a truncated (~200-draw) warmup in production configuration, extrapolates from its last quarter, reports the result as a lower bound, and stops for user confirmation if the projection exceeds 2 hours.

• Interactivity policy: no widget may trigger sampling or heavy linear algebra; the only widgets are post-fit selectors that re-index stored draws. MCMC is gated behind a run button with disk-cached posteriors.

• Outputs: DIRF figures (four shocks; panel average and per-country), eq.-13 hypothesis check, ΔGaR(0.05, 0.10) and ΔES(0.05) tables with EM/AE averages, diagnostics, and a reproducibility footer.

Testing Decisions

• The notebook is a self-validating artifact with a single seam: the model-and-pipeline functions are exercised end-to-end by the notebook's own synthetic-validation cell, which simulates a small panel with known coefficient paths and runs the identical model code used for the real fits. No new seams are introduced; no separate test files are added to the repo's tests/ suite (proof of concept, excluded from CI by decision).
• Good tests here assert external behavior only: true coefficient paths fall inside posterior credible bands; the realized crossing rate is small and consistent with the Theorem-1 bound; the common factor is recovered (correlation with truth); NUTS reports zero or near-zero divergences; the population path is monotone in every draw (a construction invariant).
• Prior art: the session prototype that validated this exact composition (Matern12 gram ⊗ Σ, ordered Bernstein rows, ALD likelihood, scan-based AR(1)) with 0 divergences and 0.97 factor recovery.
• Real-data cells validate by reconciliation rather than assertion: qualitative agreement with the paper's Figures 4–8 and Table 1 signs/magnitudes, with explained deviations.

Out of Scope

• The pseudo-out-of-sample forecasting exercise (§3.6) and its CSQR benchmark.
• Horizons beyond h ∈ {1, 4} in the default run (the h=1…10 sweep stays one constant away).
• Algorithm-level replication of the paper's Gibbs sampler (Kozumi–Kobayashi augmentation, Botev truncated-MVN sampling, precision sampler).
• The paper's §4 robustness checks.
• Sandwich/calibration corrections for the composite working likelihood (called out, not implemented).
• CI integration, docs-site rendering, additions to the repo test suite, or any change outside the notebook's directory.
• Bernstein-degree sensitivity studies (M=10 fixed by decision).

Further Notes

• No replication package exists for the paper (verified: no code/data availability statement; posted 2026-08-05), so this spec's judgment calls are the documented reconstruction; the temperature-variant choice (ERA5, unweighted) is the most plausible reading of "monthly average surface temperatures" and is a switchable constant.
• Working-likelihood caveat: each observation contributes one ALD term per grid node, so nominal posterior precision is inflated at any grid size; the paper's own bands inherit this. The notebook states it; matching the paper is the goal, calibration is not.
• Licenses/citations: GVAR (Mohaddes & Raissi 2024, CC BY 4.0); Gortan et al. (2024, Scientific Data 11:533). Both cited in the notebook header.
• Glossary and ADR-0001 already exist in the notebook directory; ADR-0002 (prior measure) is written at build time.
• Runtime expectation on the target machine (M3 Max, CPU JAX): the full-grid fit is large (~26k latent dimensions); the pilot cell exists precisely because a naive projection can undershoot by 5–10× once tree depth deepens.