prettyify the model declaration somewhat

tutorials/IDD/spatial_SIR/models.py

@@ -8,30 +8,6 @@
 import numpy as np
 from pySODM.models.base import ODE, JumpProcess
 
-# helper function
-def matmul_2D_3D_matrix(X, W):
-    """
-    Computes the product of a 2D matrix (size n x m) and a 3D matrix (size m x m x n) as an n-dimensional stack of (1xm) and (m,m) products.
-
-    input
-    =====
-    X: np.ndarray
-        Matrix of size (n,m).
-    W : np.ndarray
-        2D or 3D matrix:
-        - If 2D: Shape (m, m). Expanded to size (m, m, n).
-        - If 3D: Shape (m, m, n).
-          Represents n stacked (m x m) matrices.
-
-    output
-    ======
-    X_out : np.ndarray
-        Matrix product of size (n, m). 
-        Element-wise equivalent operation: O_{ij} = \sum_{l} [ s_{il} * w_{lji} ]
-    """
-    W = np.atleast_3d(W)
-    return np.einsum('ik,kji->ij', X, np.broadcast_to(W, (W.shape[0], W.shape[0], X.shape[0])))
-
 ###################
 ## Deterministic ##
 ###################
@@ -93,10 +69,8 @@ def compute_rates(t, S, S_work, I, R, beta, gamma, f_v, N, M):
         T_v = matmul_2D_3D_matrix(T, M) # M can  be of size (n_loc, n_loc) or (n_loc, n_loc, n_age), representing a different OD matrix in every age group
         I_v = matmul_2D_3D_matrix(I, M)
 
-        # create a size dummy 
-        G = S.shape[0] # age stratification
-        H = S.shape[1] # spatial stratification
-        size_dummy = np.ones([G,H], np.float64)
+        # create a size "dummy"
+        size_dummy = np.ones(S.shape, np.float64)
 
         rates = {
 
@@ -106,8 +80,6 @@ def compute_rates(t, S, S_work, I, R, beta, gamma, f_v, N, M):
 
             }
         return rates
-
-
 
     @ staticmethod
     def apply_transitionings(t, tau, transitionings, S, S_work, I, R, 
@@ -119,14 +91,37 @@ def apply_transitionings(t, tau, transitionings, S, S_work, I, R,
         # the resulting matrix is an N x M matrix with each element of row n, column m representing the total number of people infected from work that need to be returned to age-group n, location m. 
 
 
-        ###################################################
-        ##### CREATE THE NEW VALUES FOR S, Swork, I AND R
-        ###################################################
+        #################################################
+        ## CREATE THE NEW VALUES FOR S, Swork, I AND R ##
+        #################################################
+
         S_new = S - transitionings['S'][0] - S_work_to_home[0]
-        # print("- transitionings S",-transitionings['S'][0],  "- S_work_to_home[0]", -S_work_to_home[0])
         S_work_new =  matmul_2D_3D_matrix(S_new, M)
         I_new = I + transitionings['S'][0] + S_work_to_home[0] - transitionings['I'][0]
-        # print("- transitionings I",-transitionings['I'][0])
         R_new = R + transitionings['I'][0]
 
-        return(S_new,S_work_new, I_new, R_new)
+        return(S_new,S_work_new, I_new, R_new)
+
+# helper function
+def matmul_2D_3D_matrix(X, W):
+    """
+    Computes the product of a 2D matrix (size n x m) and a 3D matrix (size m x m x n) as an n-dimensional stack of (1xm) and (m,m) products.
+
+    input
+    =====
+    X: np.ndarray
+        Matrix of size (n,m).
+    W : np.ndarray
+        2D or 3D matrix:
+        - If 2D: Shape (m, m). Expanded to size (m, m, n).
+        - If 3D: Shape (m, m, n).
+          Represents n stacked (m x m) matrices.
+
+    output
+    ======
+    X_out : np.ndarray
+        Matrix product of size (n, m). 
+        Element-wise equivalent operation: O_{ij} = \sum_{l} [ s_{il} * w_{lji} ]
+    """
+    W = np.atleast_3d(W)
+    return np.einsum('ik,kji->ij', X, np.broadcast_to(W, (W.shape[0], W.shape[0], X.shape[0])))