feat(obs): zero-inflated negative binomial likelihood (gh#433)#434
Merged
Conversation
interpolated_value walked every knot to find the bracketing interval
(`for i in 0..times.len()-1`), while its sibling constant_value already
used binary_search_by. On many-knot interpolated forcings this scan is
the dominant per-step cost: a ~4.8k-knot/village table looked up ~480x
per RHS evaluation, ~96k RHS evaluations per solve.
Replace the scan with a binary search for the same bracket. Both the
Ok (exact knot) and Err (strictly-interior) cases resolve to the same
interval the linear scan found first, so the lerp operands and their
order are unchanged and the result is bit-for-bit identical -- verified
two ways:
- interp_tests (new): interpolated_value matches the former linear
scan bit-for-bit across exact-knot, interior, clamp, and
single/empty cases.
- full trajectory of camdl-garki joint_capacity_gam (ode, 1910-1976,
~2.4k forcing lookups/step) is byte-identical old vs new (identical
output SHA).
Measured on that model, standalone ode solve: ~37s -> ~14s. The
residual is the one-time rate_state_grad deserialize (unused on the
gradient-free MH path); per-solve compute drops ~6x, which is what a
fit sees since that load amortizes across solves. Benefits every
interpolated-forcing model, not just garki.
Regression: sim forcing/interpolation + golden byte-identity suite green
(interpolation, forcing_lag, forcing_coefficient_freeze, periodic_forcing,
cubic_spline, flat_eval_byte_identity, gate_trajectory_baseline,
smoke_all_golden, golden_simulate -- 37 tests, 0 failed).
Indexed parameter declarations now accept N dimensions, expanding to one scalar param per cell of the cartesian product (mu[village, season] -> mu_kwaru_wet, mu_kwaru_dry, ...). Frontend-only: the IR is flat scalars, so no ir/VERSION bump, no Rust change, and every existing model/golden is byte-identical (verified: regenerating all goldens changed nothing). The use side (resolution, arity, obs multi-column matching, scenario/init refs, --param-vec, collision detection, CAS, dimcheck) was already N-dim; only the declaration was 1-D. Four sites: the parser's param_decl variants (single IDENT -> separated_nonempty_list, no new menhir conflicts), the expander's expansion arm + lookup table (routed through the existing expand_indexed_decl_names cartesian producer), and a stale comment. Two guards on the merged arm: E331 (duplicate dimension) and E330 (unknown or empty dimension — also closes a latent 1-D dangling-Param bug where an unknown dim silently produced no diagnostic). Tests: 7 new compiler tests (expansion, 3-dim, use, E330/E331/E299 guards), green alongside the existing suite; the reactive E274 regression tests still pass (the fresh E330/E331 codes don't cross E274's meanings). Golden ocaml/golden/multi_index_beta.camdl (a region x age beta design matrix) pins the feature across gillespie/chain_binomial/ode + the ODE-state baseline. Proposal: docs/dev/proposals/2026-07-12-multi-index-parameters.md
Add a `zero_inflated(base = neg_binomial(mean = , r = ), pi = )` observation
likelihood for excess-zero count data. Focal vector surveillance (mosquito
catch by site × time) produces zeros from two processes — structural absence
and sampling zeros — that a single-component NegBinomial cannot represent;
forcing NB onto excess-zero data biases the dispersion and mean.
P(Y = 0) = pi + (1 - pi)·f_NB(0)
P(Y = k>0) = (1 - pi)·f_NB(k)
Design:
- The `zero_inflated(...)` wrapper (parse-time sugar, mirroring `diagnostic_test`)
desugars to a flat IR variant `ZeroInflatedNegBinomial { mean, dispersion, pi }`.
Base is restricted to `neg_binomial` for now (E324 otherwise; E325 on missing
base/pi; E251 on a stray sibling kwarg).
- Scoring-only. All three fields are bare exprs, not `Diffable`, so the family
carries no gradient. Gradient methods (PGAS/NUTS) refuse a model containing a
ZI stream at the capability gate — a MODEL-LEVEL refusal (any ZI stream ⇒
every estimated param refused), because a ZI stream's ∂loglik/∂θ is un-emitted
for every param, including a rate param that reaches it only through the shared
trajectory. MH / PMMH / IF2 / PF score it. The refusal names the family and
points to a gradient-free method.
- Run identity: mean/dispersion/pi are hashed explicitly in the canonical IR
hash (they have no `Diffable` to carry them via `diffables()`), so two ZI
models differing only in an argument do not collide.
No `ir/VERSION` bump. Unlike the prior schema bumps — which changed the wire
shape of existing fields and genuinely broke older IR reads — this is purely
additive: a new optional variant. Existing 0.29 IR reads unchanged under it, so
a bump would only invalidate every persisted 0.29 artifact (goldens, caches,
downstream stored IR) for no compatibility benefit. The IR-hash schema version
is likewise unchanged (the new variant/discriminant slots are append-only and
do not re-key any existing model).
Verified end to end with the installed binary: `simulate --obs` produces the
excess-zero signature (~69% zeros + overdispersed positives), `pfilter` scores
a finite loglik, `if2` estimates `pi_zero` + `beta` to convergence, and `pgas`
cleanly refuses (exit 1, no panic) with the gradient-free hint.
Tests: `zi_negbin_logpmf` (reduces to NB at pi=0; mixture math; sum-to-1
normalization; pi=1 boundary), a gate regression pinning the model-level
refusal (a rate param reaching a ZI stream only through the projection is
refused, not panicked), and the E251 stray-kwarg rejection. Golden:
`ocaml/golden/zinb_vector_catch` (auto-covered by smoke_all_golden).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a zero-inflated negative binomial observation likelihood for excess-zero count data (focal vector surveillance):
catch ~ zero_inflated(base = neg_binomial(mean = …, r = k), pi = pi_zero).P(Y=0) = π + (1−π)·f_NB(0),P(Y=k>0) = (1−π)·f_NB(k)diagnostic_test) → flat IR variantZeroInflatedNegBinomial { mean, dispersion, pi }. Base restricted toneg_binomial(E324/E325/E251 guard the surface).Diffable), so ZI models don't collide.No
ir/VERSIONbumpPurely additive (a new optional variant): existing 0.29 IR reads unchanged, so a bump would only invalidate every persisted 0.29 artifact for no compatibility benefit. IR-hash schema version unchanged (append-only slots).
Review + verification
Adversarial review cleared the silent-wrong classes (hash collision, silently-zero gradient, wildcard mishandling — all verified clean) and caught one bug: the gradient gate was per-param but ZI is wholly non-differentiable, so a rate param reaching ZI only through the projection hit a reachable
unreachable!panic. Fixed with a model-level refusal (+ regression test); also closed a loose-semantics hole (stray kwargs → E251).End-to-end (installed binary):
simulate --obsshows the excess-zero signature (~69% zeros + overdispersed positives),pfilterscores finite loglik,if2estimatespi_zero+betato convergence,pgascleanly refuses. Fullmake testgreen.Tests
zi_negbin_logpmf(reduces-to-NB, mixture math, sum-to-1 normalization, pi=1 boundary), gradient-gate refusal regression, E251 stray-kwarg rejection, andocaml/golden/zinb_vector_catch(smoke + trajectory baselines).Closes #433.