Move COVID-19 plugins into generic flepimop2-extras modules

[WestonVoglesonger]

Josh's custom flepimop2 plugins live in ACCIDDA/COVID19_USA (model_input/plugins/, ~28 modules registered via module="model_input.plugins.<name>" with PYTHONPATH=.). This tracks converting them into generic, namespaced modules in flepimop2-extras (packages/flepimop2-<name>/src/flepimop2/<namespace>/<module>/), one sub-issue per module.

Two foundations gate most of the work and should land first: extracting the shared prior/substitution registry that seven plugins import, and a fetch-if-not-cached data-sourcing helper.

Each module sub-issue is done when the corresponding COVID19_USA config points at the new extras module and reproduces its output. The wider COVID19_USA refactor (config alone drives the science, replacing committed data with data-product parameters, docs rewrite) is the end state we are building toward; it is tracked separately, not here.

Sub-issues

• Foundation: extract the shared prior/substitution registry #15 Foundation: extract the shared prior/substitution registry
• Foundation: fetch-if-not-cached data-sourcing helper #16 Foundation: fetch-if-not-cached data-sourcing helper
• Migrate the diffrax ODE engine to the engine namespace #17 Migrate the diffrax ODE engine to the engine namespace
• Migrate file and data-product-backed parameters #18 Migrate file and data-product-backed parameters
• Migrate prior/posterior-predictive simulation to the process namespace #19 Migrate prior/posterior-predictive simulation to the process namespace
• Migrate projection-calibration diagnostics to the process namespace #20 Migrate projection-calibration diagnostics to the process namespace
• Migrate the NHSN reporting-regime fit and coverage loader #21 Migrate the NHSN reporting-regime fit and coverage loader
• Resolve the config-driven runpy runners #22 Resolve the config-driven runpy runners

Open decisions for @pearsonca (shape the dependent sub-issues)

1. Registry placement: flepimop2 core (a ParameterABC prior-collection hook) vs a shared private extras layer? Gates the parameter/predictive/reporting modules.
2. Collapse prior-predictive and posterior-predictive into one module with a mode switch, or keep two?
3. Cache/artifact-store contract for the data-sourcing helper (cache dir, content-hash keys, re-fetch policy).
4. How much COVID/SMH-specific logic (age labels, regime dates, per-100k denominators, NC schemas) is genericized vs kept in a COVID-specific package?
5. Does the reporting-regime fit become a declared config process step, and one package or two (regime + intrinsic)?
6. Runpy runners: one generic module, fold into core as subprocess, or drop in favor of real fetch/transform modules? Is in-process execution a hard requirement?
7. Do the ~11 model_input/preprocess/*.py scripts each become their own modules, or fold into fewer fetch+transform modules?

Note

Carl's meeting list had a "readme quick start" item; that is prose plus bash (the Quick Start docs), not a migratable module, so it is part of the end-state docs rewrite rather than a flepimop2-extras sub-issue.

[pearsonca]

@jc-macdonald please chime in when you get the chance to note if these are the "right" items to modularize.

[WestonVoglesonger]

Revised plan of attack (after Carl's feedback)

Carl has commented across all eight subs. The through-line is that most of these are not 1:1 plugin ports; several dissolve into a few generic flepimop2 capabilities, some of which live in core and do not exist yet. Below is where each sub actually landed, the shared substrate they now sit on, and what can move today.

Cross-cutting themes and where they are tracked

• Parameter substrate -> flepimop2 PR #314 ("Parameters consuming parameters"). Carl's steer on Foundation: extract the shared prior/substitution registry #15 is that some pieces get pulled into core flepimop2 alongside PR #314, while the backend-specific prior implementations become modules (a NumPyroParameter, plus a hypothetical ScipyParameter he floated as the test of what belongs in core vs a module). On Migrate file and data-product-backed parameters #18 his steer is that transforms stop being fixed shape/scale fields and become functions that can reference other parameters, which lines up with #314's mechanism. This gates Foundation: extract the shared prior/substitution registry #15, Migrate file and data-product-backed parameters #18, Migrate prior/posterior-predictive simulation to the process namespace #19, Migrate the NHSN reporting-regime fit and coverage loader #21.
• Optimization targets -> flepimop2 #110 (TargetABC). The calibration/scoring and separate-estimation work all resolve into "define a target, score against a truth product." Carl explicitly paused Migrate projection-calibration diagnostics to the process namespace #20 on this, and his Migrate the NHSN reporting-regime fit and coverage loader #21 reframe (functional form + data + prior spec + likelihood -> posterior file) points at the same abstraction. #110 is currently a stub, so this is a core spec that needs to be written before Migrate projection-calibration diagnostics to the process namespace #20 and Migrate the NHSN reporting-regime fit and coverage loader #21 can be shaped. Suggest we treat #110 as a milestone that Migrate projection-calibration diagnostics to the process namespace #20/Migrate the NHSN reporting-regime fit and coverage loader #21 (and the posterior/target half of Migrate prior/posterior-predictive simulation to the process namespace #19) block on.
• Process depends-on-process / a build DAG -> genuine core gap. Both Foundation: fetch-if-not-cached data-sourcing helper #16 (fetch + transform, three-state caching) and Resolve the config-driven runpy runners #22 (Carl's process: {stepB: {depends: stepA}} sketch with acyclic checking and cache no-ops) point at a make/just-style process scheduler that flepimop2 does not have yet. This is the one theme with no existing tracker. Proposal: file a dedicated flepimop2 core issue for a process-step dependency/DAG runner, and hang Foundation: fetch-if-not-cached data-sourcing helper #16 and Resolve the config-driven runpy runners #22 off it rather than solving caching twice.
• diffrax -> op_engine, owned by Josh. Per Carl and Josh on Migrate the diffrax ODE engine to the engine namespace #17, the diffrax engine does not become a flepimop2-extras package; Carl reads it as an op_engine modification and Josh has taken it, noting there is likely already a related op_engine issue that he will confirm. This makes Migrate the diffrax ODE engine to the engine namespace #17 an outlier that leaves the extras workflow entirely.

Per-sub decided direction

• Foundation: extract the shared prior/substitution registry #15 (registry foundation) - Split: some pieces pulled into core alongside PR #314; NumPyro-specific priors become a NumPyroParameter module. The core-vs-extras boundary is still open, and is the question Carl's ScipyParameter thought experiment is meant to resolve.
• Foundation: fetch-if-not-cached data-sourcing helper #16 (data-sourcing helper) - Reshaped from one resolver into two config-declared process modules, fetch and transform, sharing three-state (absent / present / present-with-force) caching, with target names allowed to be files. Content-hash key scheme and cache dir still open.
• Migrate the diffrax ODE engine to the engine namespace #17 (diffrax engine) - Redirected to op_engine, owned by Josh. Outlier vs its siblings. Awaiting Josh's op_engine issue link and plan.
• Migrate file and data-product-backed parameters #18 (file / data-product parameters) - Posterior loader becomes an arviz/xarray connector rather than a hand-rolled reader; transforms become functions-of-parameters (#314); caching is removed from the parameter and left to the consumer. FileLookupSpec scalar-from-CSV/JSON intent unchanged.
• Migrate prior/posterior-predictive simulation to the process namespace #19 (prior/posterior-predictive) - Superseded by Carl's four-part decomposition: fittable parameters carrying both prior and posterior, simulator = system + engine, a generic "sample prior -> run simulator -> emit target" process, and a separate generic time-series plotting process. The original collapse-vs-two-modules question is now moot.
• Migrate projection-calibration diagnostics to the process namespace #20 (calibration diagnostics) - Explicitly paused by Carl until optimization targets (#110) are specified. Bias/PIT/coverage is the scoring side of a target, so it cannot be designed generically first.
• Migrate the NHSN reporting-regime fit and coverage loader #21 (NHSN reporting-regime fit) - Reframed upward into a reusable separate-estimation capability (functional form + data source + prior spec + likelihood spec -> posterior file), aimed at Josh for confirmation. Lands on TargetABC/#110 plus the general distributional spec from Foundation: extract the shared prior/substitution registry #15. The narrow "config step? one package or two?" questions are subsumed.
• Resolve the config-driven runpy runners #22 (runpy runners) - Reframed from "pick one of three runner options" into building a process step-dependency DAG with caching no-ops and acyclic checking. The original 3-way decision and in-process-vs-subprocess question are unanswered and now secondary to the DAG design.

Sequencing: actionable now vs blocked vs awaiting Josh

Actionable now

• Foundation: fetch-if-not-cached data-sourcing helper #16 fetch + transform modules. No upstream code dependency, and Carl's three-state caching behavior is concrete enough to start. Open item to pin down in parallel: the cache/artifact-store contract (cache dir, content-hash keys, re-fetch policy) which the presence-based behavior did not fully specify. This is the first thing we should build, and its caching model should be designed to feed the Resolve the config-driven runpy runners #22 DAG later so we do not implement caching twice.

Blocked on core work

• Foundation: extract the shared prior/substitution registry #15 / Migrate file and data-product-backed parameters #18 on PR #314 landing and Carl's core-vs-extras boundary call.
• Migrate prior/posterior-predictive simulation to the process namespace #19 on #314 and #110 (target output), plus the Foundation: extract the shared prior/substitution registry #15 registry decision.
• Migrate projection-calibration diagnostics to the process namespace #20 hard-paused on #110 (TargetABC) being specified.
• Migrate the NHSN reporting-regime fit and coverage loader #21 on #110 plus Foundation: extract the shared prior/substitution registry #15/Foundation: fetch-if-not-cached data-sourcing helper #16, if the generic-estimation direction is adopted.
• Resolve the config-driven runpy runners #22 on a not-yet-filed process-DAG core capability (see the new tracker proposal above).

Awaiting Josh

• Migrate the diffrax ODE engine to the engine namespace #17 entirely (op_engine issue link and plan).
• Parent-level scope sign-off (below), which several subs inherit.

For @jc-macdonald

Two asks:

1. Confirm this is the right set of items to modularize, and whether any COVID19_USA plugin is missing or wrongly split (the preprocess/*.py scripts in particular: individual modules vs folded into the Foundation: fetch-if-not-cached data-sourcing helper #16 fetch/transform pair).
2. On Migrate the diffrax ODE engine to the engine namespace #17, please link the related op_engine issue and confirm the intended approach (a new engine vs a modification of existing op_engine code), so we can move Migrate the diffrax ODE engine to the engine namespace #17 out of this tracker cleanly.

If the framing above looks right, I will start on #16 and file the process-DAG core issue that #16 and #22 should share.

[WestonVoglesonger]

Handoff: I’m leaving ACCIDDA next week and will not be able to continue driving this tracker. @pearsonca and @jc-macdonald, please take ownership of the remaining scope/architecture decisions. #16 now has draft PR #25; #17 belongs in op_engine; #20 remains paused on the target abstraction; and the other children should follow the revised decomposition in the discussion rather than being ported 1:1.