stan_sirv.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "#from owid_downloader import GenerateTrainingData\n",
    "#from utils import date_today, gravity_law_commute_dist\n",
    "\n",
    "os.environ['NUMEXPR_MAX_THREADS'] = '16'\n",
    "os.environ['NUMEXPR_NUM_THREADS'] = '8'\n",
    "\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import dgl\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from model_stan import STAN\n",
    "\n",
    "import sklearn\n",
    "from sklearn.metrics import mean_absolute_error\n",
    "from sklearn.metrics import mean_squared_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#GenerateTrainingData().download_jhu_data('2020-08-01', '2020-12-01')\n",
    "start_date = '2021-01-01'\n",
    "end_date = '2021-05-31'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data processing\n",
    "raw_data = pd.read_csv('https://raw.githubusercontent.com/owid/covid-19-data/master/public/data/owid-covid-data.csv', usecols=[\"location\", \"date\", \"total_cases\", \"new_cases_smoothed\", \"total_deaths\",\n",
    "                    \"new_deaths\", \"total_vaccinations\", \"people_fully_vaccinated\", \"new_vaccinations\", \"population\"])\n",
    "raw_data['date'] = pd.to_datetime(raw_data['date'])\n",
    "mask = (raw_data['date'] >= start_date) & (raw_data['date'] <= end_date) # & (raw_data['location'].isin(countries))\n",
    "raw_data = raw_data.loc[mask]\n",
    "#print(raw_data[raw_data['location'] == 'United States']['total_cases'].values[0])\n",
    "countries = []\n",
    "loc_list = list(raw_data['location'].unique())\n",
    "# only include countries that have more than 1000 total cases on start date and at least 1 death\n",
    "for loc in loc_list:\n",
    "    if raw_data[raw_data['location'] == loc][\"total_cases\"].values[0] > 1000 and \\\n",
    "        raw_data[raw_data['location'] == loc][\"total_deaths\"].values[0] > 0:\n",
    "        countries.append(loc)\n",
    "# hard-coded; these are problematic locations (non-countries) that need to be removed\n",
    "countries.remove(\"European Union\")\n",
    "countries.remove(\"Europe\")\n",
    "countries.remove(\"Africa\")\n",
    "countries.remove(\"Asia\")\n",
    "countries.remove(\"North America\")\n",
    "countries.remove(\"Oceania\")\n",
    "countries.remove(\"South America\")\n",
    "countries.remove(\"World\")\n",
    "countries.remove(\"Tajikistan\")\n",
    "mask = raw_data['location'].isin(countries)\n",
    "raw_data = raw_data.loc[mask]\n",
    "print(len(raw_data['location'].unique()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Generate Graph\n",
    "# add flight neighbors\n",
    "# for now, add a connection if there is any flight between the two countries between start and end date\n",
    "loc_list = list(raw_data['location'].unique())\n",
    "flight_counts = pd.read_csv('processed_flights/flight_counts_2021_all.csv')\n",
    "adj_map = {}\n",
    "for each_loc in loc_list:\n",
    "    df = flight_counts.loc[flight_counts[\"origin_country\"] == each_loc]\n",
    "    adj_map[each_loc] = set(df[\"destination_country\"].unique())\n",
    "flight_counts['day'] = pd.to_datetime(flight_counts['day'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# add land neighbors\n",
    "import csv\n",
    "neighbor_reader = csv.reader(open('neighbors.csv', 'r'))\n",
    "neighbors = {}\n",
    "for row in neighbor_reader:\n",
    "   neighbors[row[0]] = row[1].split(',')\n",
    "for each_loc,connected in adj_map.items():\n",
    "    for neighbor in neighbors[each_loc]:\n",
    "        if neighbor in loc_list:\n",
    "            connected.add(neighbor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create graph\n",
    "rows = []\n",
    "cols = []\n",
    "for each_loc in adj_map:\n",
    "    for each_loc2 in adj_map[each_loc]:\n",
    "        if each_loc in loc_list and each_loc2 in loc_list:\n",
    "            rows.append(loc_list.index(each_loc))\n",
    "            cols.append(loc_list.index(each_loc2))\n",
    "#print(rows)\n",
    "#print(cols)\n",
    "g = dgl.graph((rows, cols))\n",
    "print(g.number_of_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import networkx as nx\n",
    "# nx_g = g.to_networkx()\n",
    "# pos = nx.kamada_kawai_layout(nx_g)\n",
    "# plt.figure(1,figsize=(8,8)) \n",
    "# nx.draw(nx_g, pos, with_labels=True, node_color=[[.7, .7, .7]])\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Preprocess features\n",
    "\n",
    "#active_cases = []\n",
    "confirmed_cases = []\n",
    "new_cases = []\n",
    "new_vaccinations = []\n",
    "fully_vaccinated = []\n",
    "death_cases = []\n",
    "static_feat = []\n",
    "\n",
    "for i, each_loc in enumerate(loc_list):\n",
    "    confirmed_cases.append(raw_data[raw_data['location'] == each_loc]['total_cases'])\n",
    "    new_cases.append(raw_data[raw_data['location'] == each_loc]['new_cases_smoothed'])\n",
    "    new_vaccinations.append(raw_data[raw_data['location'] == each_loc]['new_vaccinations'])\n",
    "    fully_vaccinated.append(raw_data[raw_data['location'] == each_loc]['people_fully_vaccinated'])\n",
    "    death_cases.append(raw_data[raw_data['location'] == each_loc]['total_deaths'])\n",
    "    static_feat.append(np.array(raw_data[raw_data['location'] == each_loc][['population']]))\n",
    "confirmed_cases = np.nan_to_num(np.array(confirmed_cases))\n",
    "death_cases = np.nan_to_num(np.array(death_cases)[:, 14:])\n",
    "new_cases = np.nan_to_num(np.array(new_cases)[:, 14:])\n",
    "new_vaccinations = np.nan_to_num(np.array(new_vaccinations)[:, 14:])\n",
    "fully_vaccinated = np.nan_to_num(np.array(fully_vaccinated))\n",
    "static_feat = np.nan_to_num(np.array(static_feat)[:, 0, :])\n",
    "\n",
    "import copy\n",
    "# active = confirmed(today) - confirmed(14 days ago)\n",
    "cases_copy = copy.deepcopy(confirmed_cases)\n",
    "active = []\n",
    "for loc in confirmed_cases:\n",
    "    active_loc = []\n",
    "    for i in range(14, len(loc)):\n",
    "        active_loc.append(loc[i] - loc[i-14])\n",
    "    active.append(active_loc)\n",
    "active_cases = np.array(active)\n",
    "\n",
    "confirmed_cases = confirmed_cases[:, 14:]\n",
    "fully_vaccinated = fully_vaccinated[:, 14:]\n",
    "\n",
    "recovered_cases = confirmed_cases - active_cases - death_cases + 0.94*fully_vaccinated\n",
    "susceptible_cases = np.expand_dims(static_feat[:, 0], -1) - active_cases - recovered_cases\n",
    "\n",
    "# Batch_feat: new_cases(dI), dR, dS\n",
    "#dI = np.array(new_cases)\n",
    "dI = np.concatenate((np.zeros((active_cases.shape[0],1), dtype=np.float32), np.diff(active_cases)), axis=-1)\n",
    "dR = np.concatenate((np.zeros((recovered_cases.shape[0],1), dtype=np.float32), np.diff(recovered_cases)), axis=-1)\n",
    "dS = np.concatenate((np.zeros((susceptible_cases.shape[0],1), dtype=np.float32), np.diff(susceptible_cases)), axis=-1)\n",
    "# number of new fully vaccinated each day\n",
    "Vt = np.concatenate((np.zeros((fully_vaccinated.shape[0],1), dtype=np.float32), np.diff(fully_vaccinated)), axis=-1)\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Build normalizer\n",
    "normalizer = {'S':{}, 'I':{}, 'R':{}, 'dS':{}, 'dI':{}, 'dR':{}, 'Vt':{}}\n",
    "\n",
    "for i, each_loc in enumerate(loc_list):\n",
    "    normalizer['S'][each_loc] = (np.mean(susceptible_cases[i]), np.std(susceptible_cases[i]))\n",
    "    normalizer['I'][each_loc] = (np.mean(active_cases[i]), np.std(active_cases[i]))\n",
    "    normalizer['R'][each_loc] = (np.mean(recovered_cases[i]), np.std(recovered_cases[i]))\n",
    "    normalizer['dI'][each_loc] = (np.mean(dI[i]), np.std(dI[i]))\n",
    "    normalizer['dR'][each_loc] = (np.mean(dR[i]), np.std(dR[i]))\n",
    "    normalizer['dS'][each_loc] = (np.mean(dS[i]), np.std(dS[i]))\n",
    "    normalizer['Vt'][each_loc] = (np.mean(Vt[i]), np.std(Vt[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_data(data, sum_I, sum_R, Vt, edges_df, start, history_window=8, pred_window=14, slide_step=7):\n",
    "    # Data shape n_loc, timestep, n_feat\n",
    "    # Reshape to n_loc, t, history_window*n_feat\n",
    "    n_loc = data.shape[0]\n",
    "    timestep = data.shape[1]\n",
    "    n_feat = data.shape[2]\n",
    "    \n",
    "    x = []\n",
    "    y_I = []\n",
    "    y_R = []\n",
    "    last_I = []\n",
    "    last_R = []\n",
    "    concat_I = []\n",
    "    concat_R = []\n",
    "    concat_Vt = []\n",
    "    edges = []\n",
    "    for i in range(0, timestep, slide_step):\n",
    "        if i+history_window+pred_window-1 >= timestep or i+history_window >= timestep:\n",
    "            break\n",
    "        x.append(data[:, i:i+history_window, :].reshape((n_loc, history_window*n_feat)))\n",
    "        \n",
    "        concat_I.append(data[:, i+history_window-1, 0])\n",
    "        concat_R.append(data[:, i+history_window-1, 1])\n",
    "        last_I.append(sum_I[:, i+history_window-1])\n",
    "        last_R.append(sum_R[:, i+history_window-1])\n",
    "\n",
    "        y_I.append(data[:, i+history_window:i+history_window+pred_window, 0])\n",
    "        y_R.append(data[:, i+history_window:i+history_window+pred_window, 1])\n",
    "\n",
    "        concat_Vt.append(Vt[:, i+history_window:i+history_window+pred_window])\n",
    "        \n",
    "        # e_matrix = np.zeros((n_loc, n_loc)) # edge weight matrix for every time period \n",
    "        # df = edges_df.groupby([\"origin_country\"])\n",
    "        # for loc in range(n_loc):\n",
    "        #     try:\n",
    "        #         src_df = df.get_group(loc_list[loc])\n",
    "        #         src_df = src_df.groupby([\"destination_country\"])\n",
    "        #         for loc2 in range(n_loc):\n",
    "        #             try:\n",
    "        #                 dst_df = src_df.get_group(loc_list[loc2])\n",
    "        #                 dst_df = dst_df.loc[(dst_df['day'] >= (pd.to_datetime(start) + pd.DateOffset(days=i)))]\n",
    "        #                 dst_df = dst_df.loc[(dst_df['day'] < (pd.to_datetime(start) + pd.DateOffset(days=i+history_window-1)))]\n",
    "        #                 e_matrix[loc, loc2] = dst_df['flight_count'].sum()\n",
    "        #             except:\n",
    "        #                 continue\n",
    "        #     except:\n",
    "        #         continue\n",
    "        # # normalize matrix (doubly stochastic, see https://arxiv.org/pdf/1809.02709.pdf)\n",
    "        # # step 1: row normalize\n",
    "        # norm = np.sum(e_matrix, axis=1, keepdims=True)\n",
    "        # norm[norm==0] = 1e-10\n",
    "        # norm = 1.0 / norm\n",
    "        # P = e_matrix * norm\n",
    "\n",
    "        # step 2: P @ P^T / column_norm\n",
    "        # norm = np.sum(P, axis=0, keepdims=True)\n",
    "        # norm[norm==0] = 1e-10\n",
    "        # norm = 1.0 / norm\n",
    "\n",
    "        # PT = np.transpose(P, (1, 0))\n",
    "        # P = np.multiply(P, norm)\n",
    "        # T = np.matmul(P, PT)\n",
    "        # edges.append(T) # n_rows = # countries, n_cols = # countries\n",
    "        \n",
    "    \n",
    "    x = np.array(x, dtype=np.float32).transpose((1, 0, 2))\n",
    "    last_I = np.array(last_I, dtype=np.float32).transpose((1, 0))\n",
    "    last_R = np.array(last_R, dtype=np.float32).transpose((1, 0))\n",
    "    concat_I = np.array(concat_I, dtype=np.float32).transpose((1, 0))\n",
    "    concat_R = np.array(concat_R, dtype=np.float32).transpose((1, 0))\n",
    "    y_I = np.array(y_I, dtype=np.float32).transpose((1, 0, 2))\n",
    "    y_R = np.array(y_R, dtype=np.float32).transpose((1, 0, 2))\n",
    "    concat_Vt = np.array(concat_Vt, dtype=np.float32).transpose((1, 0, 2))\n",
    "    return x, last_I, last_R, concat_I, concat_R, y_I, y_R, concat_Vt, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def squish_edges(E):\n",
    "    # use src and dst (rows and cols) to make E matrix\n",
    "    edges = []\n",
    "    for M in E:\n",
    "        edges_flat = []\n",
    "        for i in range(len(rows)):\n",
    "            edges_flat.append(M[rows[i]][cols[i]])\n",
    "        edges.append(edges_flat)\n",
    "    return np.array(edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "valid_window = 29\n",
    "test_window = 29\n",
    "\n",
    "history_window=14 # days of information\n",
    "pred_window=14 # predicts 14 days\n",
    "slide_step=3 # increment by one week\n",
    "\n",
    "dynamic_feat = np.concatenate((np.expand_dims(dI, axis=-1), np.expand_dims(dR, axis=-1), np.expand_dims(dS, axis=-1)), axis=-1)\n",
    "    \n",
    "#Normalize\n",
    "for i, each_loc in enumerate(loc_list):\n",
    "    dynamic_feat[i, :, 0] = (dynamic_feat[i, :, 0] - normalizer['dI'][each_loc][0]) / normalizer['dI'][each_loc][1]\n",
    "    dynamic_feat[i, :, 1] = (dynamic_feat[i, :, 1] - normalizer['dR'][each_loc][0]) / normalizer['dR'][each_loc][1]\n",
    "    dynamic_feat[i, :, 2] = (dynamic_feat[i, :, 2] - normalizer['dS'][each_loc][0]) / normalizer['dS'][each_loc][1]\n",
    "    # vaccinations don't need to be normalized\n",
    "    #mean_vax = normalizer['Vt'][each_loc][0]\n",
    "    #if mean_vax != 0:\n",
    "    #   Vt[i] = (Vt[i] - mean_vax) / normalizer['Vt'][each_loc][1]\n",
    "dI_mean = []\n",
    "dI_std = []\n",
    "dR_mean = []\n",
    "dR_std = []\n",
    "\n",
    "for i, each_loc in enumerate(loc_list):\n",
    "    dI_mean.append(normalizer['dI'][each_loc][0])\n",
    "    dR_mean.append(normalizer['dR'][each_loc][0])\n",
    "    dI_std.append(normalizer['dI'][each_loc][1])\n",
    "    dR_std.append(normalizer['dR'][each_loc][1])\n",
    "\n",
    "dI_mean = np.array(dI_mean)\n",
    "dI_std = np.array(dI_std)\n",
    "dR_mean = np.array(dR_mean)\n",
    "dR_std = np.array(dR_std)\n",
    "\n",
    "#Split train-test\n",
    "train_feat = dynamic_feat[:, :-valid_window-test_window, :]\n",
    "val_feat = dynamic_feat[:, -valid_window-test_window:-test_window, :]\n",
    "test_feat = dynamic_feat[:, -test_window:, :]\n",
    "\n",
    "valid_start_date = pd.to_datetime(end_date) + pd.DateOffset(days=-valid_window) + pd.DateOffset(days=-test_window)\n",
    "test_start_date = pd.to_datetime(end_date) + pd.DateOffset(days=-test_window)\n",
    "\n",
    "train_edges = flight_counts[(flight_counts[\"day\"] >= start_date) & (flight_counts[\"day\"] < valid_start_date)]\n",
    "val_edges = flight_counts[(flight_counts[\"day\"] >= valid_start_date) & (flight_counts[\"day\"] < test_start_date)]\n",
    "test_edges = flight_counts[(flight_counts[\"day\"] >= test_start_date) & (flight_counts[\"day\"] < end_date)]\n",
    "\n",
    "train_x, train_I, train_R, train_cI, train_cR, train_yI, train_yR, train_Vt, train_edges = prepare_data(train_feat, active_cases[:, :-valid_window-test_window], recovered_cases[:, :-valid_window-test_window], Vt[:, :-valid_window-test_window], train_edges, start_date, history_window, pred_window, slide_step)\n",
    "\n",
    "val_x, val_I, val_R, val_cI, val_cR, val_yI, val_yR, val_Vt, val_edges = prepare_data(val_feat, active_cases[:, -valid_window-test_window:-test_window], recovered_cases[:, -valid_window-test_window:-test_window], Vt[:, -valid_window-test_window:-test_window], val_edges, valid_start_date, history_window, pred_window, slide_step)\n",
    "\n",
    "test_x, test_I, test_R, test_cI, test_cR, test_yI, test_yR, test_Vt, test_edges = prepare_data(test_feat, active_cases[:, -test_window:], recovered_cases[:, -test_window:], Vt[:, -test_window:], test_edges, test_start_date, history_window, pred_window, slide_step)\n",
    "\n",
    "# train_edges = squish_edges(train_edges)\n",
    "# val_edges = squish_edges(val_edges)\n",
    "# test_edges = squish_edges(test_edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Build STAN model\n",
    "\n",
    "in_dim = 3*history_window\n",
    "hidden_dim1 = 32\n",
    "hidden_dim2 = 32\n",
    "gru_dim = 32\n",
    "num_heads = 1\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "g = g.to(device)\n",
    "model = STAN(g, in_dim, hidden_dim1, hidden_dim2, gru_dim, num_heads, pred_window, device).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)\n",
    "criterion = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x = torch.tensor(train_x).to(device)\n",
    "train_I = torch.tensor(train_I).to(device)\n",
    "train_R = torch.tensor(train_R).to(device)\n",
    "train_cI = torch.tensor(train_cI).to(device)\n",
    "train_cR = torch.tensor(train_cR).to(device)\n",
    "train_yI = torch.tensor(train_yI).to(device)\n",
    "train_yR = torch.tensor(train_yR).to(device)\n",
    "train_Vt = torch.tensor(train_Vt).to(device)\n",
    "#train_edges = torch.tensor(train_edges).to(device)\n",
    "\n",
    "val_x = torch.tensor(val_x).to(device)\n",
    "val_I = torch.tensor(val_I).to(device)\n",
    "val_R = torch.tensor(val_R).to(device)\n",
    "val_cI = torch.tensor(val_cI).to(device)\n",
    "val_cR = torch.tensor(val_cR).to(device)\n",
    "val_yI = torch.tensor(val_yI).to(device)\n",
    "val_yR = torch.tensor(val_yR).to(device)\n",
    "val_Vt = torch.tensor(val_Vt).to(device)\n",
    "#val_edges = torch.tensor(val_edges).to(device)\n",
    "\n",
    "test_x = torch.tensor(test_x).to(device)\n",
    "test_I = torch.tensor(test_I).to(device)\n",
    "test_R = torch.tensor(test_R).to(device)\n",
    "test_cI = torch.tensor(test_cI).to(device)\n",
    "test_cR = torch.tensor(test_cR).to(device)\n",
    "test_yI = torch.tensor(test_yI).to(device)\n",
    "test_yR = torch.tensor(test_yR).to(device)\n",
    "test_Vt = torch.tensor(test_Vt).to(device)\n",
    "#test_edges = torch.tensor(test_edges).to(device)\n",
    "\n",
    "dI_mean = torch.tensor(dI_mean, dtype=torch.float32).to(device).reshape((dI_mean.shape[0], 1, 1))\n",
    "dI_std = torch.tensor(dI_std, dtype=torch.float32).to(device).reshape((dI_mean.shape[0], 1, 1))\n",
    "dR_mean = torch.tensor(dR_mean, dtype=torch.float32).to(device).reshape((dI_mean.shape[0], 1, 1))\n",
    "dR_std = torch.tensor(dR_std, dtype=torch.float32).to(device).reshape((dI_mean.shape[0], 1, 1))\n",
    "\n",
    "N = torch.tensor(static_feat[:, 0], dtype=torch.float32).to(device).unsqueeze(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_adaptive_loss(pred_I, pred_I_sir, pred_R, pred_R_sir, true_I, true_R, pred_window=14):\n",
    "    total_loss = 0\n",
    "    for timestep in range(len(pred_I)):\n",
    "        for day in range(1, pred_window+1):\n",
    "            sir_weight = day/(pred_window+1)\n",
    "            pred_weight = 1-sir_weight\n",
    "            sir_loss = sir_weight*((pred_I_sir[timestep][day-1] - true_I[timestep][day-1])**2) + sir_weight*((pred_R_sir[timestep][day-1] - true_R[timestep][day-1])**2)\n",
    "            pred_loss = pred_weight*criterion(pred_I[timestep][day-1], true_I[timestep][day-1]) + pred_weight*criterion(pred_R[timestep][day-1], true_R[timestep][day-1])\n",
    "            total_loss += sir_loss + pred_loss\n",
    "    return total_loss / pred_window\n",
    "\n",
    "def val_adaptive_loss(pred_I, pred_I_sir, true_I, pred_window=14):\n",
    "    total_loss = 0\n",
    "    for day in range(1, pred_window+1):\n",
    "        sir_weight = day/(pred_window+1)\n",
    "        pred_weight = 1-sir_weight\n",
    "        sir_loss = sir_weight*((pred_I_sir[day-1] - true_I[day-1])**2)\n",
    "        pred_loss = pred_weight*((pred_I[day-1] - true_I[day-1])**2)\n",
    "        total_loss += sir_loss + pred_loss\n",
    "    return total_loss / pred_window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Train STAN\n",
    "all_loss = []\n",
    "file_name = './save/stan'\n",
    "min_loss = 1e10\n",
    "\n",
    "loc_name = 'Mexico' # United States, Brazil, Germany, Japan\n",
    "cur_loc = loc_list.index(loc_name)\n",
    "\n",
    "for epoch in range(60):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    active_pred, recovered_pred, phy_active, phy_recover, _ = model(train_x, train_cI[cur_loc], train_cR[cur_loc], N[cur_loc], train_I[cur_loc], train_R[cur_loc], V=train_Vt[cur_loc])\n",
    "    phy_active = (phy_active - dI_mean[cur_loc]) / dI_std[cur_loc]\n",
    "    phy_recover = (phy_recover - dR_mean[cur_loc]) / dR_std[cur_loc]\n",
    "    # SIR loss = (day) / (pred_window + 1); dynamics loss = 1 - SIR loss \n",
    "    loss = criterion(active_pred.squeeze(), train_yI[cur_loc].squeeze())+criterion(recovered_pred.squeeze(), train_yR[cur_loc].squeeze()) \\\n",
    "        + 0.1*criterion(phy_active.squeeze(), train_yI[cur_loc].squeeze())+0.1*criterion(phy_recover.squeeze(), train_yR[cur_loc].squeeze())\n",
    "    #loss = train_adaptive_loss(active_pred.squeeze(), phy_active.squeeze(), recovered_pred.squeeze(), phy_recover.squeeze(), train_yI[cur_loc].squeeze(), train_yR[cur_loc].squeeze(), pred_window)\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    all_loss.append(loss.item())\n",
    "    \n",
    "    model.eval()\n",
    "    _, _, _, _, prev_h = model(train_x, train_cI[cur_loc], train_cR[cur_loc], N[cur_loc], train_I[cur_loc], train_R[cur_loc], V=train_Vt[cur_loc])\n",
    "    val_active_pred, val_recovered_pred, val_phy_active, val_phy_recover, _ = model(val_x, val_cI[cur_loc], val_cR[cur_loc], N[cur_loc], val_I[cur_loc], val_R[cur_loc], prev_h, V=val_Vt[cur_loc])\n",
    "    \n",
    "    val_phy_active = (val_phy_active - dI_mean[cur_loc]) / dI_std[cur_loc]\n",
    "    # SIR loss = (day) / (pred_window + 1); dynamics loss = 1 - SIR loss \n",
    "    # change loss here \n",
    "    val_loss = criterion(val_active_pred.squeeze(), val_yI[cur_loc].squeeze()) + 0.1*criterion(val_phy_active.squeeze(), val_yI[cur_loc].squeeze())\n",
    "    #val_loss = val_adaptive_loss(val_active_pred.squeeze(), val_phy_active.squeeze(), val_yI[cur_loc].squeeze(), pred_window)\n",
    "    if val_loss < min_loss:    \n",
    "        state = {\n",
    "            'state': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "        }\n",
    "        torch.save(state, file_name)\n",
    "        min_loss = val_loss\n",
    "        print('-----Save best model-----')\n",
    "    \n",
    "    print('Epoch %d, Loss %.2f, Val loss %.2f'%(epoch, all_loss[-1], val_loss.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Pred with STAN\n",
    "file_name = './save/stan'\n",
    "checkpoint = torch.load(file_name)\n",
    "model.load_state_dict(checkpoint['state'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "model.eval()\n",
    "\n",
    "\n",
    "prev_x = torch.cat((train_x, val_x), dim=1)\n",
    "prev_I = torch.cat((train_I, val_I), dim=1)\n",
    "prev_R = torch.cat((train_R, val_R), dim=1)\n",
    "prev_cI = torch.cat((train_cI, val_cI), dim=1)\n",
    "prev_cR = torch.cat((train_cR, val_cR), dim=1)\n",
    "prev_Vt = torch.cat((train_Vt, val_Vt), dim=1)\n",
    "#prev_edges = torch.cat((train_edges, val_edges), dim=0)\n",
    "prev_active_pred, _, prev_phy_active_pred, _, h = model(prev_x, prev_cI[cur_loc], prev_cR[cur_loc], N[cur_loc], prev_I[cur_loc], prev_R[cur_loc], V=prev_Vt[cur_loc])\n",
    "\n",
    "test_pred_active, test_pred_recovered, test_pred_phy_active, test_pred_phy_recover, _ = model(test_x, test_cI[cur_loc], test_cR[cur_loc], N[cur_loc], test_I[cur_loc], test_R[cur_loc], h, V=test_Vt[cur_loc])\n",
    "\n",
    "#_, _, _, _, h = model(train_x, train_cI[cur_loc], train_cR[cur_loc], N[cur_loc], train_I[cur_loc], train_R[cur_loc], V=train_Vt[cur_loc], e_weights=train_edges)\n",
    "\n",
    "#test_pred_active, test_pred_recovered, test_pred_phy_active, test_pred_phy_recover, _ = model(val_x, val_cI[cur_loc], val_cR[cur_loc], N[cur_loc], val_I[cur_loc], val_R[cur_loc], h, V=val_Vt[cur_loc], e_weights=val_edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Estimated beta in SIR model is %.2f'%model.alpha_scaled)\n",
    "print('Estimated gamma in SIR model is %.2f'%model.beta_scaled)\n",
    "#print('Estimated theta in SIR model is %.2f'%model.theta_scaled)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Cumulate predicted dI\n",
    "pred_I = []\n",
    "sir_I = []\n",
    "\n",
    "for i in range(test_pred_active.size(1)):\n",
    "    # below is regular prediction\n",
    "    cur_pred = (test_pred_active[0, i, :].detach().cpu().numpy() * dI_std[cur_loc].reshape(1, 1).detach().cpu().numpy()) + dI_mean[cur_loc].reshape(1, 1).detach().cpu().numpy()\n",
    "    # below is SIR model prediction\n",
    "    sir_pred = test_pred_phy_active[0, i, :].detach().cpu().numpy()\n",
    "    # below is average of the two predictions\n",
    "    #cur_pred = (cur_pred + test_pred_phy_active[0, i, :].detach().cpu().numpy()) / 2\n",
    "    cur_pred = np.cumsum(cur_pred)\n",
    "    cur_pred = cur_pred + test_I[cur_loc, i].detach().cpu().item()\n",
    "    pred_I.append(cur_pred)\n",
    "\n",
    "    sir_pred = np.cumsum(sir_pred)\n",
    "    sir_pred = sir_pred + test_I[cur_loc, i].detach().cpu().item()\n",
    "    sir_I.append(sir_pred)\n",
    "pred_I = np.array(pred_I)\n",
    "sir_I = np.array(sir_I)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#I_true = get_real_y(active_cases[:], history_window, pred_window, slide_step)\n",
    "I_true = active_cases[cur_loc]\n",
    "I_true = I_true[-test_window:] # test\n",
    "I_true = I_true[history_window:history_window+pred_window]\n",
    "print(I_true)\n",
    "\n",
    "# test_pred_phy_active = (test_pred_phy_active - dI_mean[cur_loc]) / dI_std[cur_loc]\n",
    "# change loss here \n",
    "# loss = criterion(test_pred_active.squeeze(), val_yI[cur_loc].squeeze()) + 0.1*criterion(test_pred_phy_active.squeeze(), val_yI[cur_loc].squeeze())\n",
    "# #print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save results\n",
    "#pickle.dump(pred_I, open(f\"stan_results/{loc_name}_stan_results_new\", \"wb\"), 2)\n",
    "#pickle.dump(pred_I, open(f\"stan_results/{loc_name}_stan_results_sirv\", \"wb\"), 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.dates as mdates\n",
    "dates = [test_start_date + pd.DateOffset(days=i+history_window) for i in range(pred_window)]\n",
    "plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%m/%d/%Y'))\n",
    "#plt.plot(I_true[cur_loc, -1, :],c='r', label='Ground truth')\n",
    "#stan_res = pickle.load(open(f\"stan_results/{loc_name}_stan_results_new\", \"rb\"))\n",
    "#stan_sirv_res = pickle.load(open(f\"stan_results/{loc_name}_stan_results_sirv\", \"rb\"))\n",
    "plt.plot(dates, I_true, c='r', label='Ground truth')\n",
    "#plt.plot(dates, stan_res[-1, :], c='g', label=\"STAN\")\n",
    "plt.plot(dates, pred_I[-1, :],c='b', label='Proposed')\n",
    "plt.gcf().autofmt_xdate()\n",
    "#plt.plot(dates, sir_I[-1, :], c='purple', label='SIRV Prediction')\n",
    "plt.legend()\n",
    "plt.title(loc_name)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# final graph visualization\n",
    "import networkx as nx\n",
    "import seaborn as sns\n",
    "\n",
    "def plot(g, attention, ax, nodes_to_plot=None, nodes_labels=None,\n",
    "         edges_to_plot=None, nodes_pos=None, nodes_colors=None,\n",
    "         edge_colormap=plt.cm.viridis):\n",
    "    \"\"\"\n",
    "    Visualize edge attentions by coloring edges on the graph.\n",
    "    g: nx.DiGraph\n",
    "        Directed networkx graph\n",
    "    attention: list\n",
    "        Attention values corresponding to the order of sorted(g.edges())\n",
    "    ax: matplotlib.axes._subplots.AxesSubplot\n",
    "        ax to be used for plot\n",
    "    nodes_to_plot: list\n",
    "        List of node ids specifying which nodes to plot. Default to\n",
    "        be None. If None, all nodes will be plot.\n",
    "    nodes_labels: list, numpy.array\n",
    "        nodes_labels[i] specifies the label of the ith node, which will\n",
    "        decide the node color on the plot. Default to be None. If None,\n",
    "        all nodes will have the same canonical label. The nodes_labels\n",
    "        should contain labels for all nodes to be plot.\n",
    "    edges_to_plot: list of 2-tuples (i, j)\n",
    "        List of edges represented as (source, destination). Default to\n",
    "        be None. If None, all edges will be plot.\n",
    "    nodes_pos: dictionary mapping int to numpy.array of size 2\n",
    "        Default to be None. Specifies the layout of nodes on the plot.\n",
    "    nodes_colors: list\n",
    "        Specifies node color for each node class. Its length should be\n",
    "        bigger than number of node classes in nodes_labels.\n",
    "    edge_colormap: plt.cm\n",
    "        Specifies the colormap to be used for coloring edges.\n",
    "    \"\"\"\n",
    "    if nodes_to_plot is None:\n",
    "        nodes_to_plot = sorted(g.nodes())\n",
    "    if edges_to_plot is None:\n",
    "        edges_to_plot = sorted(g.edges())\n",
    "    nx.draw_networkx_edges(g, nodes_pos, edgelist=edges_to_plot,\n",
    "                           edge_color=attention, edge_cmap=edge_colormap,\n",
    "                           width=2, alpha=0.5, ax=ax, edge_vmin=0,\n",
    "                           edge_vmax=1)\n",
    "\n",
    "    if nodes_colors is None:\n",
    "        nodes_colors = sns.color_palette(\"deep\", max(nodes_labels) + 1)\n",
    "\n",
    "    nx.draw_networkx_nodes(g, nodes_pos, nodelist=nodes_to_plot, ax=ax, node_size=20,\n",
    "                           node_color=[nodes_colors[nodes_labels[v - 1]] for v in nodes_to_plot],\n",
    "                            alpha=0.9)\n",
    "\n",
    "# nx_g = model.g.cpu().to_networkx(edge_attrs=['e'])\n",
    "# #print(nx_g.nodes())\n",
    "# attention = []\n",
    "# for e in sorted(nx_g.edges()):\n",
    "#     attention.append(nx_g.get_edge_data(e[0], e[1])[0]['e'].item())\n",
    "# labels = [0 for i in range(len(nx_g.nodes()))]\n",
    "# pos = nx.spring_layout(nx_g, k=0.25, iterations=20)  # positions for all nodes\n",
    "\n",
    "# fig, ax = plt.subplots(figsize=(16, 12))\n",
    "# plot(nx_g, attention, ax, nodes_pos=pos, nodes_labels=labels, nodes_colors=[\"red\" for i in range(len(nx_g.nodes()))])\n",
    "# ax.set_axis_off()\n",
    "# sm = plt.cm.ScalarMappable(cmap=plt.cm.viridis, norm=plt.Normalize(vmin=0, vmax=1))\n",
    "# sm.set_array([])\n",
    "# plt.colorbar(sm, fraction=0.046, pad=0.01)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "cdfa9ace6613655d28fa2370fe498fbea8375f41a5fc4e643d7f8a581612fdbd"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 64-bit ('tf-gpu': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "metadata": {
   "interpreter": {
    "hash": "cdfa9ace6613655d28fa2370fe498fbea8375f41a5fc4e643d7f8a581612fdbd"
   }
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}