Fix infected not going back home (#393)

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_post_vacc, dS_work.astype(float).copy(), 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

src/covid19_DTM/models/SDE_models.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2023 by T.W. Alleman BIOMATH, Ghent University. All Rights Reserved.
 
 import numpy as np
-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
 from pySODM.models.base import JumpProcess
 
 class COVID19_SEIQRD_hybrid_vacc_sto(JumpProcess):
@@ -284,7 +284,6 @@ def compute_rates(t, S, S_work, E, I, A, M_R, M_H, C_R, C_D, C_icurec, ICU_R, IC
         multip_work = beta*np.expand_dims(jit_matmul_3D_2D(Nc - Nc_home, infpop_work), axis=2) # All contacts minus home contacts on visited patch
         multip_rest = beta*np.expand_dims(jit_matmul_3D_2D(Nc_home, infpop_home), axis=2) # Home contacts always on home patch
 
-
         ################################
         ## Compute the transitionings ##
         ################################
@@ -388,6 +387,12 @@ def apply_transitionings(t, tau, transitionings, S, S_work, E, I, A, M_R, M_H, C
         dS[:,:,3] = dS[:,:,3] + N_vacc[:,:,3]*f_S
         dR[:,:,3] = dR[:,:,3] + N_vacc[:,:,3]*f_R
 
+        # 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
+        transitionings['S_work'][0] = np.rint(np.nan_to_num(redistribute_infections(S, transitionings['S_work'][0].astype(float).copy(), place_eff), nan=0.0))
 
         # Update the system
         # ~~~~~~~~~~~~~~~~~

src/covid19_DTM/models/TDPF.py

@@ -1363,7 +1363,7 @@ def get_mentality_summer_2020_lockdown(self, summer_rescaling_F, summer_rescalin
                                                     1.12,               # Bxl
                                                     1, 1.13,            # F: W-Fla: 1.45
                                                     2.15, 1.93,         # W
-                                                    0.9,                  # F: Lim: 1.43
+                                                    0.9,                # F: Lim: 1.43
                                                     1.16, 1.30])        # W
             # Rescale Flanders and Wallonia/Bxl seperately based on two parameters
             mentality_summer_2020_lockdown[idx_F] *= summer_rescaling_F
@@ -1406,11 +1406,11 @@ def get_mentality_summer_2020_lockdown(self, summer_rescaling_F, summer_rescalin
         return mentality_summer_2020_lockdown, idx_Hainaut
 
     def policies_no_lockdown(self, t, states, param, nc):
-        mat = self.__call__(datetime(2020, 1, 1), eff_home=1, eff_schools=1, eff_work=1, eff_rest=1, school=1)
+        mat = self.__call__(datetime(2020, 1, 1), eff_home=1, eff_schools=1, eff_work=1, eff_rest=1, school_primary_secundary=1, school_tertiary=1)
         return nc[:, np.newaxis, np.newaxis]*mat
 
     def policies_no_lockdown_home(self, t, states, param, nc):
-        mat = self.__call__(datetime(2020, 1, 1), eff_home=1, eff_schools=0, eff_work=0, eff_rest=0, school=0) 
+        mat = self.__call__(datetime(2020, 1, 1), eff_home=1, eff_schools=0, eff_work=0, eff_rest=0, school_primary_secundary=0, school_tertiary=0) 
         return nc[:, np.newaxis, np.newaxis]*mat
 
     def policies_all_spatial(self, t, states, param, l1, l2, eff_schools, eff_work, eff_rest, eff_home, mentality, k, summer_rescaling_F, summer_rescaling_W, summer_rescaling_B):

src/covid19_DTM/models/jit_utils.py

@@ -75,4 +75,44 @@ def matmul_q_2D(A,B):
             for j in range(m):
                 for k in range(f):
                     out[i, j, q] += A[i, k] * b[k, j]
-    return out
+    return out
+
+@jit(nopython=True)
+def redistribute_infections(S, I, G):
+    """ 
+    A jitted function to redistribute the number of drawn infections happening on the visited patch across the residency patches
+
+    Inputs
+    ------
+
+    S : np.ndarray (11,10,4)
+        Susceptibles
+
+    I : np.ndarray (11,10,4)
+        Number of infections happening on the visited patch
+
+    G : np.ndarray (11,11)
+        Origin-destination matrix
+    
+    Outputs
+    -------
+
+    I_star : np.ndarray (11,10,4)
+        Number of infections happening on visited patch, redistributed across residencies
+    """
+
+    # Compute ratio susceptibles on home over visited patch
+    div = S / matmul_q_2D(np.transpose(G), S)
+    # Compute fraction of susceptibles in spatial patch i coming from spatial patch j
+    out = np.zeros(S.shape)
+    # age
+    for j in range(S.shape[1]):
+        # vaccine
+        for k in range(S.shape[2]):
+            # Construct matrix
+            mat = G * div[:,j,k][:, np.newaxis]
+            # Matrix multiply
+            res = mat @ I[:,j,k]
+            # Assing to output
+            out[:,j,k] = res
+    return out