Vectorize r_approx_from_R for array-valued R

[developer-rpai]

Summary

r_approx_from_R in pyrenew.math previously accepted only a scalar reproduction number. Passing an array of R values raised TypeError: mul got incompatible shapes for broadcasting, because neg_MGF and neg_MGF_del_r summed over all axes and assumed a scalar rate. This PR generalizes r_approx_from_R, neg_MGF, and neg_MGF_del_r so R (and r) may be arrays of any shape, with the MGF sums taken over the weights axis only. Each entry gets an independent Newton solve, and the output shape matches the input shape. Scalar inputs behave exactly as before, confirmed by the existing tests.

Closes #486.

Root cause

In neg_MGF/neg_MGF_del_r, the expression w * jnp.exp(-r * t_vec) followed by a full jnp.sum assumed r was scalar. With vector r, shapes (k,) and (n,) do not broadcast, so the Newton iteration inside r_approx_from_R could never run on vector R. The in-tree workaround was jax.vmap around the scalar function in SubpopulationInfections.sample.

Changes

• pyrenew/math.py: neg_MGF, neg_MGF_del_r, and r_approx_from_R now accept array inputs. The rate is aligned with [..., None] and the weighted sum is taken over the last axis, so each entry is solved independently. Docstrings and type hints updated, scalar behavior unchanged.
• pyrenew/latent/subpopulation_infections.py: replaced the jax.vmap(partial(r_approx_from_R, ...)) workaround with a direct vectorized call, and removed the now unused jax and functools.partial imports.
• test/test_math.py: five new regression tests covering 1D and 2D vector R, the defining equation residual per entry, batched neg_MGF/neg_MGF_del_r, and scalar backward compatibility.

Tests

• Pristine reproduction: on unmodified code, r_approx_from_R with vector R raises TypeError; scalar calls satisfy R * M_-(r) - 1 == 0 with residuals near zero.
• The five new tests fail on pristine code (verified with the fix stashed) and pass with the fix.
• test/test_math.py: 16 passed.
• test/test_subpopulation_infections.py and test/test_population_infections.py: 57 passed.
• ruff check and the ruff format check: clean on all changed files.

Limitations

• g, the generation interval PMF, remains a single 1D vector shared across all entries of R. Per-entry PMFs are out of scope for this PR.
• numpydoc validation is not installed in this environment, so docstring validation still needs to run in CI, along with the full test suite.

[dylanhmorris]

Thanks, @developer-rpai! Going to merge with one small tweak for test readability.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.82%. Comparing base (e6e7645) to head (08680d3).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #895   +/-   ##
=======================================
  Coverage   98.82%   98.82%           
=======================================
  Files          58       58           
  Lines        2124     2128    +4     
=======================================
+ Hits         2099     2103    +4     
  Misses         25       25           

Flag	Coverage Δ
unittests	98.82% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.