Remove jurisdiction as a threaded dimension

[smwindecker]

Summary

Part 1: remove jurisdiction as a threaded dimension. jurisdiction was previously threaded through nearly every function as an explicit second dimension: matrices became date x jurisdiction, convolution operators became lapply-built lists of matrices, inits were assembled via abind into 3D arrays, and greta arrays needed sweep/apply(MARGIN=2) calls to broadcast per jurisdiction-column. This restructures the pipeline into three layers so jurisdiction is a dimension only in the last one, while preserving (not losing) hierarchical/partially-pooled multi-jurisdiction fitting:

1. Per-jurisdiction, per-stream prep (prepare_observation_data()) — now purely 1-D vectors in, 1-D vectors out.
2. Per-jurisdiction stream-grouping (define_observation_model()) — bundles cases/hospitalisations/sero for one jurisdiction, still 1-D. User-facing define_observation_data()/define_sero_data() signatures are unchanged.
3. Multi-jurisdiction stacking + fit — new stack_jurisdictions(), called internally by fit_waves(), takes a named list of per-jurisdiction bundles and stacks them into the matrices create_infection_timeseries()/create_observation_model()/create_dow_priors() already expect. Those three functions keep their existing shared-hyperparameter, partial-pooling behaviour essentially unchanged — a length-1 list collapses to today's single-jurisdiction numbers.

A master-date-axis invariant is enforced explicitly: every per-jurisdiction vector coming out of layer 1 is guaranteed length(target_infection_dates) and positionally aligned to it at construction time (as_matrix() now requires target_infection_dates and reindexes onto it, rather than inferring a date range from the data, which only worked before because all jurisdictions were pivoted together in one call). This is what makes stacking safe across jurisdictions with different/staggered data coverage.

Part 2: adopt epiwave.params's native discrete series objects. epiwave.params now has native discrete_pmf_series/discrete_weights_series objects (new_discrete_series(), with their own date-based subsetting/validation). epiwave's own epiwave_massfun_timeseries wrapper predated these and was purely redundant — prepare_observation_data() was coercing delays into that wrapper and immediately reconstructing a discrete_pmf_series from it. That wrapper is removed; prepare_observation_data() now consumes discrete_pmf/discrete_weights (or their already-time-varying _series forms) directly. new_convolution_matrix()/evaluate() are widened to accept discrete_weights alongside discrete_pmf, since seroconversion (unlike a one-time notification event) is typically persistent and better modelled as an unnormalised weights curve than a normalised PMF.

Bugs fixed along the way (all in files this refactor already rewrites):

• inits_by_jurisdiction()'s obs_prop[n_juris_ID] used linear indexing into a matrix, silently returning one scalar instead of a jurisdiction's proportion series.
• A delay pushing an inferred infection date off the front edge of the date window could hit an NA-subscript error or silently write to the wrong row.
• The seroprevalence pathway was previously unreachable: create_small_sero_model()'s call in fit_waves() was commented out, it called new_convolution_matrix() with a stale 3-argument signature, and it read a size_mat field that was never actually set (size_vec was stored instead). All three are fixed as a natural consequence of every stream (cases, hospitalisations, sero) flowing through the same unified pipeline.
• greta::initials(gp_lengthscale = rep(0.5, n_jurisdictions)) would have produced a dimension mismatch on the first real n>1 fit, since gp_lengthscale is a scalar shared kernel hyperparameter regardless of n_jurisdictions.
• inits_by_jurisdiction()'s expected_delay_vals was computed as sum(x$delay * x$mass), but discrete_pmf objects have columns step/prob, not delay/mass — this silently evaluated to 0 always (mean delay treated as zero). Now uses epiwave.params::mean.discrete_pmf() (generalised to discrete_weights via normalise() first).
• tests/test_workflow/{testing.R,test2.R} called epiwave.params::add_distributions(), which no longer exists in the current epiwave.params (renamed to add_discrete()/+).

Also declares cli (Imports) and distributional (Suggests) in DESCRIPTION — both were used directly in new_convolution_matrix.R without being formally declared.

Full design writeup and rationale for part 1: see the plan this branch implements (.claude/plans/logical-waddling-reef.md in the local Claude Code plan directory, not part of the repo).

Test plan

No automated test suite exists yet (tests/testthat/ isn't populated). Verified via devtools::load_all() + synthetic reprexes at each stage:

• devtools::check() — 0 errors, 0 warnings (remaining 3 NOTEs are pre-existing, in files this branch doesn't touch)
• Single-jurisdiction synthetic fit (flat_prior, small MCMC settings): correct output shapes, finite/non-negative posterior draws
• Two-jurisdiction fit with staggered, non-identical date coverage + a seroprevalence stream + dow_model = TRUE: correct shapes, hierarchical GP and DOW pooling both execute, and the critical alignment check confirms row i means the same calendar date in every jurisdiction's column (verified at a date only jurisdiction A has data, and a date only jurisdiction B has data)
• Discrete-series adoption: a fit exercising all three delay-input shapes together (raw discrete_pmf, a pre-built discrete_pmf_series, and a discrete_weights persistence curve for sero) runs end-to-end; confirmed the sero convolution matrix is genuinely unnormalised (row sums ≠ 1) and mean.discrete_pmf()-based inits now compute the correct non-zero delay shift (previously silently 0)
• tests/test_workflow/testing.R/test2.R updated to the new call pattern and parse cleanly (can't run to completion — depend on local, not-synced data)

🤖 Generated with Claude Code