Seceptible-Infectious-Recovered-Sceptible (SIRS) model coupled to climate.
The old problematic MERRA2 weekly clim data
/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/merra_weekly_clim.nc
is replaced by the updated one
/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/INDIVIDUAL_YEARS/MERRA_DATA/merra_weekly_climo.nc.
!ncdump -h /tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/merra_weekly_clim.nc
!ncdump -h /tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/INDIVIDUAL_YEARS/MERRA_DATA/merra_weekly_climo.nc
# init
%matplotlib inline
import sys
wython = '/tiger/scratch/gpfs/GEOCLIM/wenchang/wython'
if wython not in sys.path:
sys.path.append(wython)
%run -im wystart
# https://www.dataquest.io/blog/jupyter-notebook-tips-tricks-shortcuts/
# from IPython.core.interactiveshell import InteractiveShell
# InteractiveShell.ast_node_interactivity = 'last' # options: 'all', 'none'
# https://stackoverflow.com/questions/41125690/matplotlib-notebook-showing-a-blank-histogram
# in case to switch back to the notebook backend, run the 2 lines below,
# reload(plt)
# %matplotlib notebook
# import xaddon, xfilter, xtc
import scipy as sp
import matplotlib as mpl
print()
print(f'{sp.__name__}({sp.__version__})')
from scipy.integrate import odeint
from scipy.integrate import solve_ivp
#lat and lon for multiple cities
regions = dict(
London=(51.5074, 0.1278),
NYC=(40.7128, -74.0060+360),
LA=(34.0522, -118.2437+360),
Wuhan=(30.5928, 114.3055),
Delhi=(28.7041, 77.1025),
Maracaibo=(10.6427, -71.6125+360),
Kinshasa=(-4.4419, 15.2663),
Jakarta=(-6.2088, 106.8456),
Johannesburg=(-26.2041, 28.0473),
Buenos_Aires=(-34.6037, -58.3816+360),
Melbourne=(-37.8136, 144.9631)
)
from numpy import pi as Pi
def get_q(case='cos', qmin=0.002, qmax=0.014, qrange=None, phi0=None, qyear=None):
'''obtain specific humidity annual cycle'''
# case = 'merra2'
# qmin, qmax = 0.002, 0.018# units: kg/kg
if case == 'cos':
if phi0 is None:
phi0 = -Pi
if qrange is None:
qrange = qmax - qmin
else: #overide qmax according to qmin and qrange
qmax = qmin + qrange
t = np.arange(0, 365)
t = xr.DataArray(t, dims='time', coords=[t])
t.attrs['units'] = 'day'
qmean = (qmin + qmax)/2
cycle = np.cos(2*Pi*t/365 + phi0) # normalized to have range of [0, 1]
q = qmean + (qrange/2)*cycle
q.name = 'q'
q.attrs['units'] = 'kg/kg'
s = f'cos-curve q: '
if not hasattr(qmin, 'ndim'):
s += f'qmin = {qmin}; '
if not hasattr(qmax, 'ndim'):
s += f'qmax = {qmax}; '
if not hasattr(phi0, 'ndim'):
s += f'phi0 = {phi0/Pi:.1f}xPi; '
q.attrs['note'] = s
elif case == 'era5':
ifile = '/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/clim/q2m.run7clim.2010-2019.nc'
q = (
xr.open_dataset(ifile)
['q2m']
.rename(dayofyear='time')
.load()
)
q.attrs['note'] = case
elif case == 'era5_single_year':
ifiles = f'/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/era5.2m_specific_humidity.daily.{qyear}-??.nc'
q = (
xr.open_mfdataset(ifiles, combine='by_coords')
['q2m']
# .load()
)
q.attrs['note'] = f'{case} {qyear}'
elif case == 'merra2':
# ifile = '/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/merra_weekly_clim.nc'
ifile = '/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/INDIVIDUAL_YEARS/MERRA_DATA/merra_weekly_climo.nc'
q = (
xr.open_dataset(ifile)
.rename(WEEK_QV2M='q2m', TWEEK='time', LAT='latitude', LON='longitude')
['q2m']
.pipe(lambda x: x.assign_coords(time=x.time.dt.dayofyear))
.reindex(time=np.arange(1, 366), method='nearest')
.load()
)
# shift from 0-center to 180E-center
if q.longitude.min() < 0:
q = q.roll(longitude=q.longitude.size//2)
q = q.assign_coords(
longitude=q.longitude.where(q.longitude>=0, other=q.longitude+360).values
)
q.attrs['note'] = case
return q
q = get_q()
q.plot()
plt.title(q.note)
# q.isel(time=0).plot()
# plt.title(case)
# plt.figure()
# for city,(latx, lonx) in regions.items():
# q.sel(latitude=latx, longitude=lonx, method='nearest').plot(label=city)
# print()
# plt.legend()
# plt.title(case)
climdepens ={
'HKU1': dict(a=-227.5, R0_min=1.5, R0_max=2.5),
'OC43': dict(a=-32.5, R0_min=1.5, R0_max=2.5)
}
def q2R0(q, a, R0_min, R0_max):
'''R0 dependence on climate (specific humidity)'''
R0 = (R0_max - R0_min) * np.exp(a*q) + R0_min
if isinstance(R0, xr.DataArray):
R0.name = 'R0'
return R0
def R0_at_t(R0, t, cycle=True):
'''R0 as function of t (time) and specific humidity (q).
q has dims of (n_time, n_grid).'''
ndays_per_year = 364 # 52x7, to be consistent with Rachael's
#R0_min, R0_max = 1.4, 2.5
#a = -227.5 for HKU1 and -32.5 for OC43
if cycle:# cycle through the first 365 daily specific humidity
R0t = R0.isel(time=int( np.mod(t, ndays_per_year) ), drop=True)
else:
R0t = R0.isel(time=int(t), drop=True)
return R0t
# func_R0(0, R0t)
Ldis = {
'HKU1': 66.25*7,
'OC43': 62.5*7
}
def get_ii_from_gabe(dis, city):
#ifile = f'/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/FIGURE_CONTROL_AllRachel_RK4/{dis}_D5_nosd_null_null_null.nc'
ifile = f'/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/FIGURE_CONTROL_AllRachel_RK4/{dis}_D5_nosd_null_null_null.nc'
ds_gabe = xr.open_dataset(ifile)
new_names = {'LON': 'longitude', 'LAT': 'latitude', 'TTAFA': 't'}
latx, lonx = regions[city]
if lonx >= 180:
lonx = lonx - 360
da = ds_gabe.IT.sel(LAT=latx, LON=lonx, method='nearest').rename('ii').rename(new_names).load()
ds = da.to_dataset()
ds['year'] = da.t/364
return ds
# ds = get_ii_from_gabe('HKU1', 'NYC')
# ds
odeint¶def model_climSIRS(y, t, L, D, R0):
'''SIRS model right hand side from Baker et al., 2020.
Input:
-------
y: y[0:y.size//2] is S/N; y[y.size//2:] is I/N
t: time (units is 'day')
L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
D: infection period, 5 days
R0: beta*D, scalar, 1-D array or time-dependent N-D DataArray
Return:
--------
dydt: 1-D array with half d(S/N)/dt and half d(I/N)/dt
'''
N = y.size
ss = y[:N//2]
ii = y[N//2:]
if hasattr(R0, 'time'): # R0 is time-dependent
_R0 = R0_at_t(R0, t)
if _R0.ndim < 1:
_R0 = _R0.item()
else:
_R0 = R0
dydt = np.array([
(1-ss-ii)/L - _R0*ii*ss/D,
_R0*ii*ss/D - ii/D
]).ravel()
return dydt
from scipy.integrate import odeint
def run_climSIRS(ii0=1/8e6, ss0=None, T=5*365, dt=1, tvec=None, L=None, D=5, R0=2, dis='HKU1'):
'''Integrate SIRS model from Baker et al., 2020.
Input:
-------
I0: initial I/N
T: time stop (start is 1; units is 'day')
dt: time step
tvec: array, will override T and dt if not None;
L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
D: infection period, 5 days
R0 = beta*D, ignored if q is not None
dis: disease name, e.g. "HKU1" or "OC43"
Return:
--------
ds: solution for the SIRS model.
'''
ndays_of_year = 364
if L is None and dis is not None:
L = Ldis[dis]
if hasattr(R0, 'time'): # R0 time-dependent
if R0.ndim > 1:#
grid = R0.stack(grid=R0.isel(time=0).dims).grid # grid info is saved
zeros = np.zeros(grid.size) # broadcast array by R0.dims[1:]
else: # only time dimension
zeros = ii0*L*D*0 # broadcast array by ii0*L*D*a
else: # R0 time-independent
zeros = ii0*L*D*R0*0
# default initial S/N value determined by initial I/N value
if ss0 is None:
ss0 = 1 - ii0
# broadcast initial values by adding zeros
ii0 = ii0 + zeros
ss0 = ss0 + zeros
y0 = np.hstack([ss0, ii0])
if tvec is None:
t = np.arange(0, T, dt)
else:
t = tvec
sol = odeint(model_climSIRS, y0, t, args=(L, D, R0))
ss = sol[:, 0:y0.size//2]
ii = sol[:, y0.size//2:]
da_ss = xr.DataArray(ss, dims=['t', 'grid'],coords=[t, np.arange(1, y0.size//2+1)])
da_ss.attrs['long_name'] = 'S/N'
da_ii = xr.DataArray(ii, dims=['t', 'grid'],coords=[t, np.arange(1, y0.size//2+1)])
da_ii.attrs['long_name'] = 'I/N'
if hasattr(R0, 'time') and R0.ndim > 1: # R0 has more than "time" dimensions
da_ss = da_ss.assign_coords(grid=grid).unstack()
da_ii = da_ii.assign_coords(grid=grid).unstack()
ds = xr.Dataset(dict(ss=da_ss, ii=da_ii, year=da_ss.t/ndays_of_year))
ds.t.attrs['units'] = 'day'
return ds
# latx, lonx = regions['NYC']
# q = get_q('era5').sel(latitude=latx, longitude=lonx, method='nearest')
# q = get_q()
ds = run_climSIRS()
ds.ii.plot()
ds
solve_ivp¶# t instead of y is the first argument required by scipy.integrate.solve_ivp
def model_climSIRSpy(t, y, L, D, R0):
return model_climSIRS(y=y, t=t, L=L, D=D, R0=R0)
# scipy.integrate.solve_ivp backend
from scipy.integrate import solve_ivp
def run_climSIRSpy(ii0=1/8e6, ss0=None, T=5*365, dt=1, tvec=None, L=None, D=5, R0=2, dis='HKU1'):
'''Integrate SIRS model from Baker et al., 2020.
Input:
-------
I0: initial I/N
T: time stop (start is 1; units is 'day')
dt: time step
tvec: will override T and dt is not None;
L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
D: infection period, 5 days
R0 = beta*D, scalar, 1-D array or time-dependent N-D DataArray
dis: disease name, e.g. "HKU1" or "OC43"
Return:
--------
ds: solution for the SIRS model.
'''
ndays_of_year = 364
if L is None and dis is not None:
L = Ldis[dis]
if hasattr(R0, 'time'): # R0 time-dependent
if R0.ndim > 1:#
R0 = R0.stack(grid=R0.isel(time=0).dims)
grid = R0.grid # grid info is saved
zeros = np.zeros(grid.size) # broadcast array by R0.dims[1:]
else: # only time dimension
zeros = ii0*L*D*0 # broadcast array by ii0*L*D*a
else: # R0 time-independent
zeros = ii0*L*D*R0*0
# default initial S/N value determined by initial I/N value
if ss0 is None:
ss0 = 1 - ii0
# broadcast initial values by adding zeros
ii0 = ii0 + zeros
ss0 = ss0 + zeros
y0 = np.hstack([ss0, ii0])
if tvec is None:
t = np.arange(0, T, dt)
else:
t = tvec
# # scipy.integrate.odeint backend; call FORTRAN code ultimately
# sol = odeint(model_climSIRS, y0, t, args=(L, D, R0))
# ss = sol[:, 0:y0.size//2]
# ii = sol[:, y0.size//2:]
# da_ss = xr.DataArray(ss, dims=['t', 'grid'],coords=[t, np.arange(1, y0.size//2+1)])
# da_ss.attrs['long_name'] = 'S/N'
# da_ii = xr.DataArray(ii, dims=['t', 'grid'],coords=[t, np.arange(1, y0.size//2+1)])
# da_ii.attrs['long_name'] = 'I/N'
# scipy.integrate.solve_ivp backend, pure Python code
# print(t)
sol = solve_ivp(model_climSIRSpy,[0, T], y0, t_eval=t, args=(L, D, R0))
# print(sol)
ss = sol.y[0:y0.size//2, :]
ii = sol.y[y0.size//2:, :]
# print(ss.shape)
da_ss = xr.DataArray(ss, dims=['grid', 't'],coords=[np.arange(1, y0.size//2+1), t])
da_ss.attrs['long_name'] = 'S/N'
da_ii = xr.DataArray(ii, dims=['grid', 't'],coords=[np.arange(1, y0.size//2+1), t])
da_ii.attrs['long_name'] = 'I/N'
if hasattr(R0, 'time') and R0.ndim > 1: # R0 has more than "time" dimensions
da_ss = da_ss.assign_coords(grid=grid).unstack()
da_ii = da_ii.assign_coords(grid=grid).unstack()
ds = xr.Dataset(dict(ss=da_ss, ii=da_ii, year=da_ss.t/ndays_of_year))
ds.t.attrs['units'] = 'day'
return ds
# latx, lonx = regions['NYC']
# q = get_q('era5').sel(latitude=latx, longitude=lonx, method='nearest')
# q = get_q()
ds = run_climSIRSpy()
ds.ii.plot()
ds
RK4 explicitly¶def model_climSIRSwy(ss, ii, t, L, D, R0):
'''SIRS model right hand side from Baker et al., 2020.
Input:
-------
ss: S/N
ii: I/N
t: time (units is 'day')
L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
D: infection period, 5 days
R0 = beta*D, scalar, 1-D array or time-dependent N-D DataArray
Return:
--------
tuple (dsdt, didt)
dsdt: S/N tendency
didt: I/N tendency
'''
if hasattr(R0, 'time'): # R0 is time-dependent
_R0 = R0_at_t(R0, t)
if _R0.ndim < 1:
_R0 = _R0.item()
else:
_R0 = R0
dsdt = (1-ss-ii)/L - _R0*ii*ss/D
didt = _R0*ii*ss/D - ii/D
return dsdt, didt
def rk4_climSIRSwy(ss, ii, t, dt, L, D, R0):
'''Step change from classic Runge–Kutta (RK4) method for SIRS model of Baker et al., 2020.
Input:
-------
ss: S/N
ii: I/N
t: time (units is 'day')
dt: time step
L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
D: infection period, 5 days
R0 = beta*D, scalar, 1-D array or N-D time-dependent DataArray
Return:
--------
tuple (dss, dii)
dss: S/N step increase
dii: I/N step increase
'''
model_climSIRS = model_climSIRSwy
k1s, k1i = model_climSIRS(ss=ss, ii=ii, t=t, L=L, D=D, R0=R0)
k2s, k2i = model_climSIRS(ss=ss+k1s*dt/2, ii=ii+k1i*dt/2, t=t+dt/2, L=L, D=D, R0=R0)
k3s, k3i = model_climSIRS(ss=ss+k2s*dt/2, ii=ii+k2i*dt/2, t=t+dt/2, L=L, D=D, R0=R0)
k4s, k4i = model_climSIRS(ss=ss+k3s*dt, ii=ii+k3i*dt, t=t+dt, L=L, D=D, R0=R0)
dss = dt*(k1s + 2*k2s + 2*k3s + k4s)/6
dii = dt*(k1i + 2*k2i + 2*k3i + k4i)/6
return dss, dii
def run_climSIRSwy(ii0=1/8e6, ss0=None, T=5*365, dt=1, tvec=None, L=None, D=5, R0=2, dis='HKU1'):
'''Integrate SIRS model from Baker et al., 2020.
Input:
-------
I0: initial I/N
T: time stop (start is 1; units is 'day')
dt: time step
tvec: will override T and dt if not None;
L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
D: infection period, 5 days
R0: beta*D, scalar, 1-D array or time-dependent n-D DataArray
dis: disease name, e.g. "HKU1" or "OC43"
Return:
--------
ds: solution for the climSIRSwy model.
'''
ndays_of_year = 364
if tvec is None:
t_vec = np.arange(0, T, dt)
else:
t_vec = tvec
if L is None and dis is not None:
L = Ldis[dis]
if hasattr(R0, 'time'): # R0 time-dependent
zeros = ii0*L*D*R0.isel(time=0, drop=True)*0 # broadcast array by R0.dims[1:]
else: # R0 time-independent
zeros = ii0*L*D*R0*0
# default initial S/N value determined by initial I/N value
if ss0 is None:
ss0 = 1 - ii0
# broadcast initial values by adding zeros
ii = ii0 + zeros
ss = ss0 + zeros
ss_vec = [ss,]
ii_vec = [ii,]
print('time step (units: day): ', end='')
for t in t_vec[1:]:
if np.mod(t, ndays_of_year) == 0:
print(f'{ndays_of_year}x{t//ndays_of_year}', end='; ')
dss, dii = rk4_climSIRSwy(ss=ss, ii=ii, t=t, dt=dt, L=L, D=D, R0=R0)
ss = ss + dss
ii = ii + dii
ss_vec.append(ss)
ii_vec.append(ii)
print()
if isinstance(ss, xr.DataArray):
try:
da_ss = xr.concat(ss_vec, dim=pd.Index(t_vec, name='t'))
da_ii = xr.concat(ii_vec, dim=pd.Index(t_vec, name='t'))
except MemoryError:
ss_vec = [_ss.expand_dims('dum').chunk({'dum':1}).squeeze() for _ss in ss_vec]
ii_vec = [_ii.expand_dims('dum').chunk({'dum':1}).squeeze() for _ii in ii_vec]
da_ss = xr.concat(ss_vec, dim=pd.Index(t_vec, name='t'))
da_ii = xr.concat(ii_vec, dim=pd.Index(t_vec, name='t'))
else:
da_ss = xr.DataArray(ss_vec, dims=['t'], coords=[t_vec])
da_ii = xr.DataArray(ii_vec, dims=['t'], coords=[t_vec])
da_ss.attrs['long_name'] = 'S/N'
da_ii.attrs['long_name'] = 'I/N'
ds = xr.Dataset(dict(ss=da_ss, ii=da_ii, year=da_ss.t/ndays_of_year))
ds.t.attrs['units'] = 'day'
return ds
ds = run_climSIRSwy()
# ds.ii.sel(latitude=latx, longitude=lonx, method='nearest').plot()
ds.ii.assign_coords(t=ds.year).rename(t='year').plot()
# da.sel(si='I/N').plot()
dis = 'HKU1'
T = 15*365
ii0 = 1/8e6
q = get_q()
q.plot()
R0 = q2R0(q, **climdepens[dis])
plt.figure()
R0.plot()
plt.figure()
ds = run_climSIRS(dis=dis, R0=R0, T=T, ii0=ii0)
ds.ii.assign_coords(t=ds.year).rename(t='year').plot()
ds = run_climSIRSpy(dis=dis, R0=R0, T=T, ii0=ii0)
ds.ii.assign_coords(t=ds.year).rename(t='year').plot()
ds = run_climSIRSwy(dis=dis, R0=R0, T=T, ii0=ii0)
ds.ii.assign_coords(t=ds.year).rename(t='year').plot()
plt.legend(['odeint', 'solve_ivp', 'RK4'])
dt¶dis = 'HKU1'
ii0 = 1/8e6
city = 'NYC'
latx, lonx = regions[city]
q = get_q('era5').sel(latitude=latx, longitude=lonx, method='nearest')
q.plot()
R0 = q2R0(q, **climdepens[dis])
plt.figure()
R0.plot()
plt.figure()
dss = []
ds = run_climSIRSwy(dis=dis, R0=R0, T=10*365, ii0=ii0, dt=0.5)
dss.append(ds)
ds.ii.assign_coords(t=ds.year).rename(t='year').plot()
ds = run_climSIRSwy(dis=dis, R0=R0, T=10*365, ii0=ii0, dt=1)
dss.append(ds)
ds.ii.assign_coords(t=ds.year).rename(t='year').plot()
plt.title(f'{city} {dis}')
plt.legend(['RK4, dt = 0.5', 'RK4, dt = 1.0'])
# sensitivity to time step
for i,dt in enumerate([0.5, 1.0]):
# ds = run_climSIRSwy(L=66.25*7, D=5, I0=1e-5, T=5*365, dt=dt,
# q=q2m.sel(latitude=latx, longitude=lonx, method='nearest'), R0=2)
ds = dss[i]
ds.ii.sel(t=slice(0,365)).plot(color=f'C{i}', label=f'dt = {dt}')
dss.append(ds)
plt.legend()
plt.title(f'{city} {dis} (RK4)')
# sensitivity to time step
for i,dt in enumerate([0.5, 1.0]):
ds = dss[i]
ds.ii.sel(t=slice(72,83)).plot(color=f'C{i}', label=f'dt = {dt}')
dss.append(ds)
plt.legend()
plt.title(f'{city} {dis} (RK4)')
q2m = get_q('merra2')
q2m
dis = 'HKU1'
city = 'NYC'
latx, lonx = regions[city]
q = q2m.sel(latitude=latx, longitude=lonx, method='nearest')
q.rename(time='day_of_year').plot()
R0 = q2R0(q, **climdepens[dis])
R0.rename('day_of_year').plot()
dss = {}
ii0 = 1/8e6
T = 10*364
dt = 1
dss['odeint'] = run_climSIRS(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
dt = 1
dss['rk4'] = run_climSIRSwy(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
dt = 7
dss['rk4_dt7'] = run_climSIRSwy(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
dss['Gabe'] = get_ii_from_gabe(dis=dis, city=city)
label = 'odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, color='k')
label = 'rk4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
label = 'rk4_dt7'
ds = dss[label]
# ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls=':')
label = 'Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label='Gabe', ls=':')
plt.legend()
label = 'odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, color='k')
label = 'rk4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
label = 'rk4_dt7'
ds = dss[label]
# ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
label = 'Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label='Gabe', ls=':')
plt.legend()
plt.xlim(0.18, .3)
label = 'odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, color='k')
label = 'rk4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
label = 'rk4_dt7'
ds = dss[label]
# ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
label = 'Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label='Gabe', ls=':')
plt.legend()
plt.xlim(1.9, 2.1)
dis = 'HKU1'
city = 'Wuhan'
latx, lonx = regions[city]
q = q2m.sel(latitude=latx, longitude=lonx, method='nearest')
q.rename(time='day_of_year').plot()
R0 = q2R0(q, **climdepens[dis])
R0.rename('day_of_year').plot()
if 'dss' not in globals():
dss = {}
ii0 = 1/8e6
T = 11*364
dt = 1
dss['odeint'] = run_climSIRS(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
dt = 1
dss['rk4'] = run_climSIRSwy(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
dt = 7
dss['rk4_dt7'] = run_climSIRSwy(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
dss['Gabe'] = get_ii_from_gabe(dis=dis, city=city)
label = 'odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, color='k')
label = 'rk4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
label = 'rk4_dt7'
ds = dss[label]
# ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
label = 'Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label='Gabe', ls=':')
plt.legend()
plt.title(f'{city} {dis}', loc='left', y=1.05)
label = 'odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, color='k')
label = 'rk4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
label = 'rk4_dt7'
ds = dss[label]
# ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
label = 'Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label='Gabe', ls=':')
plt.legend()
plt.xlim(0.2, .3)
if 'dss' not in globals():
dss = {}
q2m = get_q('merra2')
ii0 = 1/8e6
T = 10*364
diseases = ['HKU1', 'OC43']
for dis in diseases:
print(dis)
for city,(latx, lonx) in regions.items():
print(city, f'lat = {latx}', f'lon = {lonx}')
q = q2m.sel(latitude=latx, longitude=lonx, method='nearest')
R0 = q2R0(q, **climdepens[dis])
# odeint
dt = 1
label = f'{dis}_{city}_odeint'
print(label)
dss[label] = run_climSIRS(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
# RK4
dt = 1
label = f'{dis}_{city}_RK4'
print(label)
dss[label] = run_climSIRSwy(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
# RK4 dt = 7
dt = 7
label = f'{dis}_{city}_RK4dt7'
print(label)
dss[label] = run_climSIRSwy(ii0=ii0, T=T, dt=dt, R0=R0, dis=dis)
#Gabe
label = f'{dis}_{city}_Gabe'
print(label)
dss[label] = get_ii_from_gabe(dis=dis, city=city)
# plot
diseases = ['HKU1', 'OC43']
for dis in diseases:
for city,(latx, lonx) in regions.items():
plt.figure()
# odeint
label = f'{dis}_{city}_odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
# RK4
label = f'{dis}_{city}_RK4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
# RK4dt7
label = f'{dis}_{city}_RK4dt7'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls=':')
# Gabe
label = f'{dis}_{city}_Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
plt.legend()
if dis == 'HKU1':
plt.ylim(None, 0.15)
elif dis == 'OC43':
plt.ylim(None, 0.25)
plt.tight_layout()
figname = f'climSIRS/figs/R0min1p5_newMERRA2clim_comparison_10years_{dis}_{city}.png'
plt.savefig(figname)
print('[saved]:', figname)
# plot
diseases = ['HKU1', 'OC43']
for dis in diseases:
for city,(latx, lonx) in regions.items():
plt.figure()
# odeint
label = f'{dis}_{city}_odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
# RK4
label = f'{dis}_{city}_RK4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
# RK4dt7
label = f'{dis}_{city}_RK4dt7'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls=':')
# Gabe
label = f'{dis}_{city}_Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
plt.legend()
plt.xlim(0, 1)
if dis == 'HKU1':
plt.ylim(None, 0.15)
elif dis == 'OC43':
plt.ylim(None, 0.25)
plt.tight_layout()
figname = f'climSIRS/figs/R0min1p5_newMERRA2clim_comparison_firstyear_{dis}_{city}.png'
plt.savefig(figname)
print('[saved]:', figname)
# plot
diseases = ['HKU1', 'OC43']
for dis in diseases:
for city,(latx, lonx) in regions.items():
plt.figure()
# odeint
label = f'{dis}_{city}_odeint'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
# RK4
label = f'{dis}_{city}_RK4'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label)
# RK4dt7
label = f'{dis}_{city}_RK4dt7'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls=':')
# Gabe
label = f'{dis}_{city}_Gabe'
ds = dss[label]
ds.ii.assign_coords(t=ds.year).rename(t='year').plot(label=label, ls='--')
plt.legend()
plt.xlim(1, 3)
if dis == 'HKU1':
plt.ylim(None, 0.15)
elif dis == 'OC43':
plt.ylim(None, 0.25)
plt.tight_layout()
figname = f'climSIRS/figs/R0min1p5_newMERRA2clim_comparison_year1to3_{dis}_{city}.png'
plt.savefig(figname)
print('[saved]:', figname)
%%html
<script src="https://cdn.rawgit.com/parente/4c3e6936d0d7a46fd071/raw/65b816fb9bdd3c28b4ddf3af602bfd6015486383/code_toggle.js">
</script>