-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup_fix01.py
173 lines (136 loc) · 6.39 KB
/
setup_fix01.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import numpy as np
import pandas as pd
import geopandas as gpd
import datetime, os
from shapely.geometry import Point, Polygon
from COVIDScenarioPipeline.SEIR import seir
ncomp = 7
S, E, I1, I2, I3, R, cumI = np.arange(ncomp)
class Setup():
"""
This class hold a setup model setup.
"""
def __init__(self, setup_name, spatial_setup, nsim, ti, tf,
interactive = True, write_csv = False,
dt = 1/6, nbetas = None):
self.setup_name = setup_name
self.nsim = nsim
self.dt = dt
self.ti = ti
self.tf = tf
self.interactive = interactive
self.write_csv = write_csv
if nbetas is None:
nbetas = nsim
self.nbetas = nbetas
self.spatset = spatial_setup
self.build_setup()
self.dynfilter = - np.ones((self.t_span,self.nnodes))
if self.write_csv:
self.timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
self.datadir = f'model_output/{self.setup_name}/'
if not os.path.exists(self.datadir):
os.makedirs(self.datadir)
def build_setup(self):
self.t_span = (self.tf - self.ti).days
self.t_inter = np.arange(0, self.t_span+0.0001, self.dt)
self.nnodes = self.spatset.nnodes
self.popnodes = self.spatset.popnodes
self.mobility = self.spatset.mobility
def buildIC(self, seeding_places, seeding_amount):
self.y0 = np.zeros((ncomp, self.nnodes))#, dtype = 'int64')
self.y0[S,:] = self.popnodes
for i, pl in enumerate(seeding_places):
self.y0[S, pl] = self.popnodes[pl] - seeding_amount[i]
self.y0[I1, pl] = seeding_amount[i]
return self.y0
def buildICfromfilter(self):
y0 = np.zeros((ncomp, self.nnodes))#, dtype = 'int64')
draw = np.random.poisson(5*self.dynfilter[31]+0.1)
y0[S,:] = self.popnodes - draw
y0[E, :] = (draw/4).astype(np.int)
y0[I1, :] = (draw/4).astype(np.int)
y0[I2, :] = (draw/4).astype(np.int)
y0[I3, :] = (draw/4).astype(np.int)
y0[cumI, :] = (3*draw/4).astype(np.int)
return y0
def set_filter(self, dynfilter):
self.dynfilter = dynfilter
class COVID19Parameters():
""" Class to hold parameters for COVID19 transmission.
When temporal rates, unit is [d^-1]
"""
def __init__(self, s):
self.s = s
# https://github.com/midas-network/COVID-19/tree/master/parameter_estimates/2019_novel_coronavirus
# incubation period 5.2 days based on an estimate from Lauer et al. 2020
self.sigma = 1/5.2
# Number of infected compartiments
n_Icomp = 3
# time from symptom onset to recovery per compartiment
self.gamma = np.random.uniform(1/6, 1/2.6, s.nbetas) * n_Icomp # range of serial from 8.2 to 6.5
if 'Red' in s.setup_name:
self.gamma = np.random.uniform(1/4.1, 1/2.6, s.nbetas) * n_Icomp
if 'low' in s.setup_name: self.R0s = np.random.uniform(.6, 1.3, s.nbetas) # np.random.uniform(1.5, 2, nbetas)
if 'mid' in s.setup_name: self.R0s = np.random.uniform(2, 3, s.nbetas)
if 'South' in s.setup_name: self.R0s = np.random.uniform(.6, 1.3, s.nbetas)
self.betas = np.multiply(self.R0s, self.gamma) / n_Icomp
self.betas = np.vstack([self.betas]*len(s.t_inter))
self.gamma = np.vstack([self.gamma]*len(s.t_inter))
self.sigma = np.hstack([self.sigma]*len(s.t_inter))
self.betas = np.dstack([self.betas]*s.nnodes)
self.gamma = np.dstack([self.gamma]*s.nnodes)
self.sigma = np.vstack([self.sigma]*s.nnodes)
def draw(self, beta_id):
""" for speed, to use with numba JIT compilation"""
return(np.array([self.betas[:,beta_id%self.s.nbetas], self.sigma.T, self.gamma[:,beta_id%self.s.nbetas]]))
def addNPIfromcsv(self, filename):
npi = pd.read_csv(filename).T
npi.columns = npi.iloc[0]
npi = npi.drop('Unnamed: 0')
npi.index = pd.to_datetime(npi.index)
npi = npi.resample(str(self.s.dt*24) + 'H').ffill()
for i in range(self.s.nbetas):
self.betas[:,i,:] = np.multiply(self.betas[:,i,:], np.ones_like(self.betas[:,i,:]) - npi.to_numpy())
print (f'>>> Added NPI as specicied in file {filename}')
def addNPIfromR(self, npi):
npi.index = pd.to_datetime(npi.index.astype(str))
npi = npi.resample(str(self.s.dt*24) + 'H').ffill()
for i in range(self.s.nbetas):
self.betas[:,i,:] = np.multiply(self.betas[:,i,:], np.ones_like(self.betas[:,i,:]) - npi.to_numpy())
def parameters_quick_draw(s, npi):
sigma = 1/5.2
n_Icomp = 3
gamma = np.random.uniform(1/6, 1/2.6) * n_Icomp # range of serial from 8.2 to 6.5
if 'Red' in s.setup_name:
gamma = np.random.uniform(1/4.1, 1/2.6) * n_Icomp
if 'low' in s.setup_name: R0s = np.random.uniform(1.5, 2) # np.random.uniform(1.5, 2, nbetas)
if 'mid' in s.setup_name:
R0s = np.random.uniform(2, 3)
if 'South' in s.setup_name: R0s = np.random.uniform(.6, 1.3)
print('gamma=%.2f, R0=%.2f'%(gamma, R0s))
beta = np.multiply(R0s, gamma) / n_Icomp
#print(beta, gamma, sigma)
beta = np.hstack([beta]*len(s.t_inter))
gamma = np.hstack([gamma]*len(s.t_inter))
sigma = np.hstack([sigma]*len(s.t_inter))
beta = np.vstack([beta]*s.nnodes)
gamma = np.vstack([gamma]*s.nnodes)
sigma = np.vstack([sigma]*s.nnodes)
#print(beta.shape, gamma.shape, sigma.shape)
npi.index = pd.to_datetime(npi.index.astype(str))
npi = npi.resample(str(s.dt*24) + 'H').ffill()
beta = np.multiply(beta, np.ones_like(beta) - npi.to_numpy().T)
return(np.array([beta.T, sigma.T, gamma.T]))
class CaliforniaSpatialSetup():
"""
Setup for california at the county scale.
"""
def __init__(self):
folder = 'california/'
self.data = pd.read_csv(f'data/{folder}geodata.csv')
self.mobility = np.loadtxt(f'data/{folder}mobility.txt')
self.popnodes = self.data['new_pop'].to_numpy()
self.nnodes = len(self.data)
self.counties_shp = gpd.read_file(f'data/{folder}california-counties-shp/california-counties.shp')
self.counties_shp.sort_values('GEOID', inplace=True)