Appease pre-commit + sum incidence across strains

CDCgov · Nov 3, 2023 · 23568c0 · 23568c0
1 parent a2982a0
commit 23568c0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/mechanistic_compartments.py b/mechanistic_compartments.py
@@ -7,7 +7,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import numpyro
-import numpyro.distributions as dist
 import pandas as pd
 from diffrax import ODETerm, SaveAt, Tsit5, diffeqsolve
 from jax.random import PRNGKey

diff --git a/model_odes/seir_model_v5.py b/model_odes/seir_model_v5.py
@@ -122,7 +122,8 @@ def _seirw_ode(state, _, parameters):
     di = jnp.add(jnp.zeros(i.shape), jnp.add(de_to_i, -di_to_r))
     dr = jnp.add(jnp.zeros(r.shape), jnp.add(di_to_r, -dr_to_w))
 
-    dw_to_e = jnp.sum(dw_to_e, axis=2)
+    ds_to_e = jnp.sum(ds_to_e, axis=1) # across strains
+    dw_to_e = jnp.sum(dw_to_e, axis=[1, 2]) # across strains & w compartments
     incidence = ds_to_e + dw_to_e
     return [(ds, de, di, dr, dw), incidence]