{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# climSIRS\n", "Seceptible-Infectious-Recovered-Sceptible (SIRS) model coupled to climate.\n", "* Wenchang Yang (wenchang@prnceton.edu)\n", "* Department of Geosciences, Princeton University" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Go to comparison directly" ] }, { "cell_type": "code", "execution_count": 59, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T19:49:02.112741Z", "start_time": "2020-04-26T19:49:02.088214Z" }, "code_folding": [], "scrolled": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "[2020-04-26_15:49:02]: wystart\n", "\n", "[imported]: os.path, sys, os, datetime, glob\n", "[imported]: xarray(0.15.1) as xr, numpy(1.18.1) as np, pandas(1.0.3) as pd, matplotlib(3.1.1) as mpl\n", "[imported]: import matplotlib.pyplot as plt\n", "[imported]: from matplotlib.pyplot import plot, figure, close\n", "[executed]: plt.ion()\n", "[config]: plt.rcParams['figure.dpi'] = 128\n", "[config]: plt.rcParams['figure.figsize'] = [6.4, 6.4*9/16]:\n", "[imported]: import misc.colormaps\n", "[config]: xr.set_options(cmap_sequential=\"parula\")\n", "[config]: xr.set_options(cmap_divergent=\"turbo\")\n", "[iPython config]: InlineBackend.figure_format ='retina'\n", "\n", "[2020-04-26_15:49:02]: wystart done\n", "[time used]: **0** seconds\n", "[imported]: import scipy(1.4.1) as sp\n" ] } ], "source": [ "# init\n", "%matplotlib inline\n", "import sys\n", "wython = '/tiger/scratch/gpfs/GEOCLIM/wenchang/wython'\n", "if wython not in sys.path:\n", " sys.path.append(wython)\n", "%run -im wystart\n", "# https://www.dataquest.io/blog/jupyter-notebook-tips-tricks-shortcuts/\n", "# from IPython.core.interactiveshell import InteractiveShell\n", "# InteractiveShell.ast_node_interactivity = 'last' # options: 'all', 'none'\n", "# https://stackoverflow.com/questions/41125690/matplotlib-notebook-showing-a-blank-histogram\n", "# in case to switch back to the notebook backend, run the 2 lines below, \n", "# reload(plt)\n", "# %matplotlib notebook\n", "# import xaddon, xfilter, xtc\n", "import scipy as sp\n", "print(f'[imported]: import scipy({sp.__version__}) as sp')" ] }, { "cell_type": "code", "execution_count": 60, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T19:49:13.991420Z", "start_time": "2020-04-26T19:49:13.987298Z" } }, "outputs": [], "source": [ "from scipy.integrate import odeint\n", "from scipy.integrate import solve_ivp\n", "\n", "#lat and lon for multiple cities\n", "regions = dict(\n", " London=(51.5074, 0.1278),\n", " NYC=(40.7128, -74.0060+360),\n", " LA=(34.0522, -118.2437+360),\n", " Wuhan=(30.5928, 114.3055),\n", " Delhi=(28.7041, 77.1025),\n", " Maracaibo=(10.6427, -71.6125+360),\n", " Kinshasa=(-4.4419, 15.2663),\n", " Jakarta=(-6.2088, 106.8456),\n", " Johannesburg=(-26.2041, 28.0473),\n", " Buenos_Aires=(-34.6037, -58.3816+360),\n", " Melbourne=(-37.8136, 144.9631)\n", " \n", ")" ] }, { "cell_type": "code", "execution_count": 61, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T19:49:17.572003Z", "start_time": "2020-04-26T19:49:17.553483Z" } }, "outputs": [], "source": [ "from numpy import pi as Pi\n", "def get_q(case='cos', qmin=0.002, qmax=0.014, qrange=None, phi0=None, qyear=None):\n", " '''obtain specific humidity annual cycle'''\n", "# case = 'merra2'\n", "# qmin, qmax = 0.002, 0.018# units: kg/kg \n", " if case == 'cos':\n", " if phi0 is None:\n", " phi0 = -Pi\n", " if qrange is None:\n", " qrange = qmax - qmin\n", " else: #overide qmax according to qmin and qrange\n", " qmax = qmin + qrange\n", " t = np.arange(0, 365)\n", " t = xr.DataArray(t, dims='time', coords=[t])\n", " t.attrs['units'] = 'day'\n", " qmean = (qmin + qmax)/2\n", " cycle = np.cos(2*Pi*t/365 + phi0) # normalized to have range of [0, 1]\n", " q = qmean + (qrange/2)*cycle\n", " q.name = 'q'\n", " q.attrs['units'] = 'kg/kg'\n", " s = f'cos-curve q: '\n", " if not hasattr(qmin, 'ndim'):\n", " s += f'qmin = {qmin}; '\n", " if not hasattr(qmax, 'ndim'):\n", " s += f'qmax = {qmax}; '\n", " if not hasattr(phi0, 'ndim'):\n", " s += f'phi0 = {phi0/Pi:.1f}xPi; '\n", " q.attrs['note'] = s\n", " elif case == 'era5':\n", " ifile = '/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/clim/q2m.run7clim.2010-2019.nc'\n", " q = ( \n", " xr.open_dataset(ifile)\n", " ['q2m']\n", " .rename(dayofyear='time')\n", " .load()\n", " )\n", " q.attrs['note'] = case\n", " elif case == 'era5_single_year':\n", " ifiles = f'/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/era5.2m_specific_humidity.daily.{qyear}-??.nc'\n", " q = ( \n", " xr.open_mfdataset(ifiles, combine='by_coords')\n", " ['q2m']\n", "# .load()\n", " )\n", " q.attrs['note'] = f'{case} {qyear}'\n", " elif case == 'merra2':\n", " ifile = '/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/merra_weekly_clim.nc'\n", " q = ( \n", " xr.open_dataset(ifile)\n", " .rename(WEEK_QV2M='q2m', TT='time', LAT='latitude', LON='longitude')\n", " ['q2m']\n", " .pipe(lambda x: x.assign_coords(time=x.time.dt.dayofyear))\n", " .reindex(time=np.arange(1, 366), method='nearest')\n", " .load()\n", " )\n", " # shift from 0-center to 180E-center\n", " if q.longitude.min() < 0:\n", " q = q.roll(longitude=q.longitude.size//2)\n", " q = q.assign_coords(\n", " longitude=q.longitude.where(q.longitude>=0, other=q.longitude+360).values\n", " )\n", " q.attrs['note'] = case\n", " \n", " return q\n", "# q = get_q()\n", "# q.plot()\n", "# plt.title(q.note)\n", "\n", "# q.isel(time=0).plot()\n", "# plt.title(case)\n", "# plt.figure()\n", "# for city,(latx, lonx) in regions.items():\n", "# q.sel(latitude=latx, longitude=lonx, method='nearest').plot(label=city)\n", "# print()\n", "# plt.legend()\n", "# plt.title(case)" ] }, { "cell_type": "code", "execution_count": 62, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T19:49:19.920695Z", "start_time": "2020-04-26T19:49:19.912032Z" } }, "outputs": [], "source": [ "climdepens ={\n", " 'HKU1': dict(a=-227.5, R0_min=1.4, R0_max=2.5),\n", " 'OC43': dict(a=-32.5, R0_min=1.4, R0_max=2.5)\n", "}\n", "def q2R0(q, a, R0_min, R0_max):\n", " '''R0 dependence on climate (specific humidity)'''\n", " R0 = (R0_max - R0_min) * np.exp(a*q) + R0_min\n", " if isinstance(R0, xr.DataArray):\n", " R0.name = 'R0'\n", " \n", " return R0" ] }, { "cell_type": "code", "execution_count": 63, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T19:49:22.422651Z", "start_time": "2020-04-26T19:49:22.418844Z" } }, "outputs": [], "source": [ "def R0_at_t(R0, t, cycle=True):\n", " '''R0 as function of t (time) and specific humidity (q).\n", " q has dims of (n_time, n_case).'''\n", "# ndays_per_year = 364 # 52x7, to be consistent with Rachael's\n", " ndays_per_year = 365\n", " #R0_min, R0_max = 1.4, 2.5\n", " #a = -227.5 for HKU1 and -32.5 for OC43\n", " if cycle:# cycle through the first 365 daily specific humidity\n", " R0t = R0.isel(time=int( np.mod(t, ndays_per_year) ), drop=True)\n", " else:\n", " R0t = R0.isel(time=int(t), drop=True)\n", " return R0t\n", "# func_R0(0, R0t)" ] }, { "cell_type": "code", "execution_count": 64, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T19:49:24.163717Z", "start_time": "2020-04-26T19:49:24.161003Z" } }, "outputs": [], "source": [ "Ldis = {\n", " 'HKU1': 66.25*7,\n", " 'OC43': 62.5*7\n", "}" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## climSIRS using `odeint`" ] }, { "cell_type": "code", "execution_count": 112, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T20:28:09.624083Z", "start_time": "2020-04-26T20:28:09.610715Z" } }, "outputs": [], "source": [ "def model_climSIRS(y, t, L, D, R0, sdR0=None, sdIN=None):\n", " '''SIRS model right hand side from Baker et al., 2020.\n", " Input:\n", " -------\n", " y: y[0:y.size//2] is S/N; y[y.size//2:] is I/N\n", " t: time (units is 'day')\n", " L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)\n", " D: infection period, 5 days\n", " R0: beta*D, scalar, 1-D array or time-dependent N-D DataArray\n", "\n", " Return:\n", " --------\n", " dydt: 1-D array with half d(S/N)/dt and half d(I/N)/dt\n", " '''\n", " N = y.size\n", " ss = y[:N//2]\n", " ii = y[N//2:]\n", " if hasattr(R0, 'time'): # R0 is time-dependent\n", " _R0 = R0_at_t(R0, t).values\n", " if _R0.ndim < 1:\n", " _R0 = _R0.item()\n", " else:\n", " _R0 = R0\n", " \n", " # adaptive social distance control when ii > I/N threshold sdIN\n", " if sdR0 is not None:\n", " _R0 = _R0 + ii*0\n", " _R0[ii>sdIN] = sdR0\n", " if ii[0] > sdIN:\n", " print(_R0)\n", " print(ii)\n", "# sys.exit()\n", " \n", " dydt = np.array([\n", " (1-ss-ii)/L - _R0*ii*ss/D, \n", " _R0*ii*ss/D - ii/D\n", " ]).ravel()\n", " \n", " return dydt" ] }, { "cell_type": "code", "execution_count": 113, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T20:28:11.195252Z", "start_time": "2020-04-26T20:28:11.185721Z" }, "code_folding": [] }, "outputs": [], "source": [ "from scipy.integrate import odeint\n", "def run_climSIRS(ii0=1e-8, ss0=None, T=5*365, dt=1, tvec=None, L=None, D=5, R0=2, dis='HKU1', sdR0=None, sdIN=None):\n", " '''Integrate SIRS model from Baker et al., 2020.\n", "\n", " Input:\n", " -------\n", " I0: initial I/N\n", " T: time stop (start is 1; units is 'day')\n", " dt: time step\n", " tvec: array, will override T and dt if not None;\n", " L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)\n", " D: infection period, 5 days\n", " R0 = beta*D, ignored if q is not None\n", " dis: disease name, e.g. \"HKU1\" or \"OC43\"\n", "\n", " Return:\n", " --------\n", " ds: solution for the SIRS model.\n", " '''\n", " ndays_of_year = 365\n", " \n", " if L is None and dis is not None:\n", " L = Ldis[dis]\n", " if hasattr(R0, 'time'): # R0 time-dependent\n", " if R0.ndim > 1:# \n", " #case = R0.stack(case=R0.isel(time=0).dims).case # case info is saved\n", " R0 = R0.stack(case=R0.isel(time=0).dims)\n", " case = R0.case\n", " zeros = np.zeros(case.size) # broadcast array by R0.dims[1:]\n", " else: # only time dimension\n", " zeros = ii0*L*D*0 # broadcast array by ii0*L*D*a\n", " else: # R0 time-independent\n", " zeros = ii0*L*D*R0*0\n", " # default initial S/N value determined by initial I/N value\n", " if ss0 is None:\n", " ss0 = 1 - ii0\n", " # broadcast initial values by adding zeros\n", " ii0 = ii0 + zeros\n", " ss0 = ss0 + zeros\n", " \n", " y0 = np.hstack([ss0, ii0])\n", " if tvec is None:\n", " t = np.arange(0, T, dt)\n", " else:\n", " t = tvec\n", " sol = odeint(model_climSIRS, y0, t, args=(L, D, R0, sdR0, sdIN))\n", " ss = sol[:, 0:y0.size//2]\n", " ii = sol[:, y0.size//2:]\n", " da_ss = xr.DataArray(ss, dims=['t', 'case'],coords=[t, np.arange(1, y0.size//2+1)])\n", " da_ss.attrs['long_name'] = 'S/N'\n", " da_ii = xr.DataArray(ii, dims=['t', 'case'],coords=[t, np.arange(1, y0.size//2+1)])\n", " da_ii.attrs['long_name'] = 'I/N'\n", " \n", " if hasattr(R0, 'time') and R0.ndim > 1: # R0 has more than \"time\" dimensions\n", " da_ss = da_ss.assign_coords(case=case).unstack()\n", " da_ii = da_ii.assign_coords(case=case).unstack()\n", " \n", " ds = xr.Dataset(dict(ss=da_ss, ii=da_ii, year=da_ss.t/ndays_of_year))\n", " ds.t.attrs['units'] = 'day'\n", " \n", " return ds\n", "\n", "# ds = run_climSIRS()\n", "# ds.ii.plot()\n", "# ds" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### social distance" ] }, { "cell_type": "code", "execution_count": 73, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T20:08:07.778572Z", "start_time": "2020-04-26T20:07:59.403219Z" } }, "outputs": [], "source": [ "q2m = get_q('era5')" ] }, { "cell_type": "code", "execution_count": 74, "metadata": { "ExecuteTime": { "end_time": "2020-04-26T20:09:01.584382Z", "start_time": "2020-04-26T20:09:01.411935Z" } }, "outputs": [ { "data": { "text/html": [ "
array([1.9579365, 1.9751573, 1.9822502, 2.0056412, 2.0240772, 2.0302188,\n",
" 2.026044 , 2.0184145, 1.9842691, 1.9696474, 1.9608523, 1.9542973,\n",
" 1.9473044, 1.9433839, 1.9422187, 1.9640894, 1.972503 , 1.9748982,\n",
" 1.9789406, 1.9832816, 1.9797525, 1.9864752, 1.992496 , 1.9911659,\n",
" 1.988812 , 1.9876965, 1.9892855, 2.0063925, 2.0103629, 2.0070922,\n",
" 2.0153694, 2.01311 , 2.005015 , 2.0058172, 1.9960301, 1.9861073,\n",
" 1.9952822, 1.9993043, 1.9994841, 2.0060463, 2.0126576, 2.018021 ,\n",
" 2.0181656, 1.9981111, 1.9961119, 1.9953272, 1.9922931, 1.9786885,\n",
" 1.9600965, 1.9566598, 1.9525721, 1.9257393, 1.9031277, 1.9040514,\n",
" 1.9090778, 1.9147029, 1.9098585, 1.9112151, 1.9326644, 1.9568337,\n",
" 1.9577582, 1.9590927, 1.9643064, 1.9713023, 1.9721831, 1.9545491,\n",
" 1.9320221, 1.9115003, 1.9005009, 1.8836106, 1.8722453, 1.8694613,\n",
" 1.8744315, 1.8851918, 1.891075 , 1.8881253, 1.8858285, 1.8815114,\n",
" 1.8816046, 1.8824384, 1.8780087, 1.8709122, 1.8734574, 1.8808477,\n",
" 1.8762941, 1.8603528, 1.8498205, 1.8399961, 1.831815 , 1.8135428,\n",
" 1.802345 , 1.8068511, 1.8128409, 1.805934 , 1.7943326, 1.792876 ,\n",
" 1.7962652, 1.7894248, 1.7758343, 1.7584101, 1.7445357, 1.7447374,\n",
" 1.7377334, 1.7232758, 1.7226999, 1.7190459, 1.7152574, 1.720166 ,\n",
" 1.7142522, 1.707699 , 1.7046545, 1.698156 , 1.6894509, 1.6828017,\n",
" 1.6746387, 1.675951 , 1.6821706, 1.6819735, 1.6685638, 1.6615276,\n",
" 1.6539338, 1.6484323, 1.6358763, 1.6229521, 1.6129779, 1.6138082,\n",
" 1.6105924, 1.6098368, 1.6084002, 1.6140037, 1.6128385, 1.6057254,\n",
" 1.5950505, 1.5893732, 1.5860945, 1.578205 , 1.5655041, 1.5608941,\n",
" 1.5594257, 1.5555644, 1.5503889, 1.5445226, 1.5359672, 1.527229 ,\n",
" 1.5159074, 1.5066247, 1.4990131, 1.4920794, 1.484925 , 1.484003 ,\n",
" 1.4861549, 1.4918623, 1.4947493, 1.4984734, 1.5071251, 1.5140661,\n",
" 1.5154192, 1.5127066, 1.5052032, 1.5041525, 1.4994279, 1.4946598,\n",
" 1.4958271, 1.4946532, 1.4903505, 1.4863793, 1.4828602, 1.4837495,\n",
" 1.4807388, 1.4749533, 1.4693863, 1.4676753, 1.4680111, 1.4677149,\n",
" 1.4671227, 1.46634 , 1.4644711, 1.4659333, 1.466193 , 1.4634342,\n",
" 1.4603645, 1.4576874, 1.456451 , 1.4534458, 1.4502288, 1.447356 ,\n",
" 1.4467868, 1.4470153, 1.4464221, 1.4452622, 1.4460379, 1.4461184,\n",
" 1.4457166, 1.4451663, 1.4437969, 1.4426558, 1.4409823, 1.439363 ,\n",
" 1.4391438, 1.4398232, 1.4393443, 1.4394294, 1.4411511, 1.4441558,\n",
" 1.4464748, 1.4463936, 1.4462537, 1.4486835, 1.4508765, 1.4492143,\n",
" 1.4473603, 1.4469191, 1.4471405, 1.4482466, 1.4477535, 1.4466182,\n",
" 1.4478955, 1.4485525, 1.4494361, 1.4506372, 1.4497443, 1.4483439,\n",
" 1.4491113, 1.4492189, 1.4491173, 1.447744 , 1.4465768, 1.4470313,\n",
" 1.4483842, 1.447731 , 1.4461104, 1.448202 , 1.4534237, 1.4574604,\n",
" 1.4611331, 1.4624591, 1.4637419, 1.4680192, 1.4687314, 1.4652607,\n",
" 1.4634396, 1.4601393, 1.4582206, 1.4583747, 1.4581735, 1.4575866,\n",
" 1.4593892, 1.4610425, 1.4650955, 1.4685552, 1.4697784, 1.4697286,\n",
" 1.4723785, 1.4747244, 1.4784322, 1.4820826, 1.4884608, 1.4956896,\n",
" 1.5012226, 1.5027221, 1.5028858, 1.5013895, 1.5012628, 1.4989105,\n",
" 1.4978976, 1.499554 , 1.5000969, 1.4968541, 1.4961768, 1.4959567,\n",
" 1.4986564, 1.5020243, 1.5069189, 1.5104561, 1.5199901, 1.5293927,\n",
" 1.5307627, 1.5332108, 1.5334818, 1.5342232, 1.5378382, 1.5442075,\n",
" 1.5515901, 1.5601497, 1.5665505, 1.5741739, 1.5854038, 1.6007547,\n",
" 1.6089385, 1.608313 , 1.6180427, 1.6236452, 1.6325978, 1.6346076,\n",
" 1.6313438, 1.634131 , 1.6394868, 1.64094 , 1.640937 , 1.6465394,\n",
" 1.6461192, 1.6514664, 1.652313 , 1.6599029, 1.6741153, 1.6869276,\n",
" 1.6832776, 1.692019 , 1.710773 , 1.7284026, 1.7371184, 1.744926 ,\n",
" 1.7561178, 1.7814823, 1.8010862, 1.8017663, 1.7952026, 1.7997406,\n",
" 1.8014463, 1.8049881, 1.8021691, 1.8050618, 1.8034714, 1.8185769,\n",
" 1.8303093, 1.8337866, 1.8349179, 1.8308182, 1.826284 , 1.8311089,\n",
" 1.8279217, 1.8138473, 1.8066587, 1.806497 , 1.8161862, 1.8128556,\n",
" 1.8000591, 1.8041039, 1.829734 , 1.8408358, 1.8449667, 1.8507686,\n",
" 1.8611925, 1.8800784, 1.8763185, 1.8616614, 1.8706262, 1.8786378,\n",
" 1.8828707, 1.888712 , 1.9007581, 1.903017 , 1.8892637, 1.8663638,\n",
" 1.8570467, 1.8498833, 1.8496565, 1.8405936, 1.8376889, 1.8604977,\n",
" 1.8812679, 1.8970602, 1.9182673, 1.9240311, 1.9394858, 1.9613364],\n",
" dtype=float32)array(40.75, dtype=float32)
array(286., dtype=float32)
array([ 1, 2, 3, ..., 364, 365, 366])