add redistribution to ODE models

UGentBiomath · May 14, 2024 · faa5fd1 · faa5fd1
1 parent 1f4219e
commit faa5fd1
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/src/covid19_DTM/models/ODE_models.py b/src/covid19_DTM/models/ODE_models.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 from pySODM.models.base import ODE
-from covid19_DTM.models.jit_utils import jit_matmul_2D_1D, jit_matmul_3D_2D, jit_matmul_klm_m, jit_matmul_klmn_n, matmul_q_2D
+from covid19_DTM.models.jit_utils import jit_matmul_2D_1D, jit_matmul_3D_2D, jit_matmul_klm_m, jit_matmul_klmn_n, matmul_q_2D, redistribute_infections
 
 class simple_multivariant_SIR(ODE):
     """
@@ -396,15 +396,23 @@ def integrate(t, S, E, I, A, M_R, M_H, C_R, C_D, C_icurec, ICU_R, ICU_D, R, D, M
         multip_work = beta*np.expand_dims(jit_matmul_3D_2D(Nc_home, infpop_work), axis=2)
         multip_rest = beta*np.expand_dims(jit_matmul_3D_2D(Nc-Nc_home, infpop_home), axis=2)
 
-        # Compute rates of change
-        dS_inf = (S_work * multip_work + S_post_vacc * multip_rest)*e_s
+        # Compute total number of 
+        dS_work = S_work * multip_work * e_s
+        dS_rest = S_post_vacc * multip_rest * e_s
+
+        # We have the number of new infections happening on a visited spatial patch
+        # --> These need to be transformed back into the place of residency
+        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+        # Convert the number of infections to float and then copy to guarantee the array is contiguous for numba performance
+        dS_work = np.rint(np.nan_to_num(redistribute_infections(S, dS_work, place_eff), nan=0.0))
 
         ############################
         ## Compute system of ODEs ##
         ############################
 
-        dS  = dS - dS_inf
-        dE  = dS_inf - (1/sigma)*E
+        dS  = dS - dS_work - dS_rest
+        dE  = dS_work + dS_rest - (1/sigma)*E
         dI = (1/sigma)*E - (1/omega)*I
         dA = (a/omega)*I - (1/da)*A
         dM_R = (1-h_acc)*((1-a)/omega)*I - (1/dm)*M_R