From 191175a70956299022edbe9f11e90ba1311cdc1f Mon Sep 17 00:00:00 2001 From: Tijs Alleman Date: Fri, 14 Jun 2024 16:11:06 -0400 Subject: [PATCH] remove bugs from model --- tutorials/IDD/spatial_SIR/models.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/tutorials/IDD/spatial_SIR/models.py b/tutorials/IDD/spatial_SIR/models.py index 3550d2b..652cef0 100644 --- a/tutorials/IDD/spatial_SIR/models.py +++ b/tutorials/IDD/spatial_SIR/models.py @@ -6,18 +6,12 @@ __copyright__ = "Copyright (c) 2024 by T.W. Alleman, IDD Group, Johns Hopkins Bloomberg School of Public Health. All Rights Reserved." import numpy as np - -################### -## Deterministic ## -################### - -# import the ODE class -from pySODM.models.base import ODE +from pySODM.models.base import ODE, JumpProcess # helper function def matmul_2D_3D_matrix(X, W): """ - Computes the matrix product of a 2D matrix (size n x m) and a 3D matrix (size m x m x n). + 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 ===== @@ -37,14 +31,17 @@ def matmul_2D_3D_matrix(X, W): """ return np.einsum('ik,kji->ij', X, np.broadcast_to(np.atleast_3d(W), (W.shape[0], W.shape[0], X.shape[0]))) -# Define the model equations +################### +## Deterministic ## +################### + class spatial_ODE_SIR(ODE): """ SIR model with a spatial stratification """ states = ['S','I','R'] - parameters = ['beta','gamma'] + parameters = ['beta','gamma', 'f_v', 'N', 'M'] dimensions = ['age', 'location'] @staticmethod @@ -59,11 +56,11 @@ def integrate(t, S, I, R, beta, gamma, f_v, N, M): I_v = matmul_2D_3D_matrix(I, M) # compute number of new infections on home patch and visited patch - dI_h = beta * S * np.transpose(matmul_2D_3D_matrix(np.transpose(I/T), N)) # N can be of size (n_age, n_age) or (n_age, n_age, n_loc), representing a different contact matrix in every spatial patch - dI_v = beta * S_v * np.transpose(matmul_2D_3D_matrix(np.transpose(I_v/T), N)) + dI_h = beta * S * np.transpose(matmul_2D_3D_matrix(np.transpose(I/T), (1-f_v)*N)) # N can be of size (n_age, n_age) or (n_age, n_age, n_loc), representing a different contact matrix in every spatial patch + dI_v = beta * S_v * np.transpose(matmul_2D_3D_matrix(np.transpose(I_v/T_v), f_v*N)) # distribute the number of new infections on visited patch to the home patch - dI_h = S * np.transpose(M @ np.transpose(dI_h/S_v)) + dI_v = S * np.transpose(M @ np.transpose(dI_v/S_v)) # Calculate differentials dS = - (dI_h + dI_v)