Support matrix population mixing in convolution scanner - #876
Support matrix population mixing in convolution scanner#876damonbayer wants to merge 16 commits into
Conversation
|
Thank you for your contribution @dylanhmorris 🚀! Your github-pages is ready for download 👉 here 👈! |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #876 +/- ##
=======================================
Coverage 98.74% 98.74%
=======================================
Files 59 59
Lines 2232 2236 +4
=======================================
+ Hits 2204 2208 +4
Misses 28 28
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
dylanhmorris
left a comment
There was a problem hiding this comment.
Various wording tweaks
Co-authored-by: Dylan H. Morris <dylanhmorris@users.noreply.github.com>
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
|
The matrix-convolution support in this PR is a necessary building block for contact mixing in PyRenew. Before presenting contact mixing as a supported feature, however, we should discuss how it fits into the high-level API. I may be missing context from while I was away. Could we clarify:
These choices affect both the epidemiological interpretation and the public API. Even when the modeling need is urgent, a short issue or design note would help us record the use case, agree on these decisions, and keep the implementation consistent with the rest of PyRenew. I suggest merging the low-level primitive once its semantics and validation are settled, deferring or relabeling the tutorial, and opening a short design issue for integration with SubpopulationInfections, PyrenewBuilder, and MultiSignalModel. We can add the end-user tutorial once that path exists. I’m happy to draft the design issue and work through the interface together. |
| - building_multisignal_models.md | ||
| - latent_infections.md | ||
| - latent_subpopulation_infections.md | ||
| - population_mixing.md |
There was a problem hiding this comment.
Defer this tutorial until we have a high-level API.
cdc-mitzimorris
left a comment
There was a problem hiding this comment.
These three comments about the multiplier dispatch, the contact-versus-reproduction-matrix interpretation, and coverage of the public function are separate from the high-level API questions in my earlier comment.
| convolved_history = jnp.einsum( | ||
| "i...,i...->...", array_to_convolve, history_subset | ||
| ) | ||
| if jnp.ndim(multiplier) < 2: |
There was a problem hiding this comment.
Could we make matrix application explicit rather than infer it solely from ndim? Previously, higher-dimensional multipliers were applied elementwise; this now interprets every rank-2-or-higher multiplier as a matrix. That can break callers whose state has multiple batch dimensions. A separate matrix scanner or an explicit mode would preserve the existing behavior and make the contract unambiguous.
| pmf vector. | ||
| Rt | ||
| Timeseries of $\\mathcal{R}(t)$ values | ||
| Timeseries of $\\mathcal{R}(t)$ values. Each timepoint can be a |
There was a problem hiding this comment.
Could we describe this as an effective reproduction or next-generation matrix and distinguish it from a raw contact matrix? As used here, each entry already represents transmission from a source to a target. Survey contact matrices generally require scaling for transmissibility, susceptibility/infectiousness, population sizes, and sometimes reciprocity before they can be used in this equation. Clarifying that distinction would help prevent users from passing raw contact data directly as Rt.
| assert_array_equal(new_val_mat, new_val_vec) | ||
|
|
||
|
|
||
| def test_convolve_scanner_with_population_mixing_matrix(): |
There was a problem hiding this comment.
Could we test matrix-valued input through compute_infections_from_rt, rather than only through the scanner? I would also like coverage for a multi-step time-varying matrix against an independent calculation, invalid matrix/state shapes, and JIT/gradient compatibility.
Summary
new_convolve_scannerM[i, j]generically as the contribution of componentjto componenti, while preserving scalar and vector behaviorM[target, source], constant and time-varying mixing, and two-population figuresTesting
uv run pytest -m 'not integration'(843 passed, 35 deselected)uv run pre-commit run --files docs/index.md docs/tutorials/.pages docs/tutorials/population_mixing.qmd pyrenew/convolve.py pyrenew/latent/infection_functions.py test/test_convolve.pyQUARTO_PYTHON=.venv/bin/python quarto render docs/tutorials/population_mixing.qmd