notebooks/mutrans_backtesting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Backtesting Plots for mutation growth rate paper\n",
    "\n",
    "This notebook generates plots for the [paper/backtesting](paper/backtesting) directory. This assumes you've alread run\n",
    "```sh\n",
    "make update                       # Downloads data (~1hour).\n",
    "make preprocess-usher             # Preprocesses usher tree\n",
    "make backtesting                  # Fits backtesting models\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%load_ext autoreload\n",
    "#%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import math\n",
    "import os\n",
    "import pickle\n",
    "import re\n",
    "import logging\n",
    "from collections import Counter, OrderedDict, defaultdict\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import torch\n",
    "import pyro.distributions as dist\n",
    "from pyrocov import mutrans, pangolin, stats\n",
    "from pyrocov.stats import normal_log10bf\n",
    "from pyrocov.util import pretty_print, pearson_correlation\n",
    "import seaborn as sns\n",
    "import matplotlib.colors as mcolors\n",
    "import matplotlib.cm as cm\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import colorcet as cc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matplotlib.rcParams[\"figure.dpi\"] = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# configure logging\n",
    "logging.basicConfig(format=\"%(relativeCreated) 9d %(message)s\", level=logging.INFO)\n",
    "# This line can be used to modify logging as required later in the notebook\n",
    "#logging.getLogger().setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set matplotlib params\n",
    "#matplotlib.rcParams[\"figure.dpi\"] = 200\n",
    "#matplotlib.rcParams['figure.figsize'] = [8, 8]\n",
    "matplotlib.rcParams[\"axes.edgecolor\"] = \"gray\"\n",
    "matplotlib.rcParams[\"savefig.bbox\"] = \"tight\"\n",
    "matplotlib.rcParams['font.family'] = 'sans-serif'\n",
    "matplotlib.rcParams['font.sans-serif'] = ['Arial', 'Avenir', 'DejaVu Sans']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load input data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the entire constant dataset\n",
    "max_num_clades = 3000\n",
    "min_num_mutations = 1\n",
    "min_region_size = 50\n",
    "ambiguous = False\n",
    "columns_filename=f\"results/columns.{max_num_clades}.pkl\"\n",
    "features_filename=f\"results/features.{max_num_clades}.{min_num_mutations}.pt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dataset = mutrans.load_gisaid_data(\n",
    "        device=\"cpu\",\n",
    "        columns_filename=columns_filename,\n",
    "        features_filename=features_filename,\n",
    "        min_region_size=min_region_size\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load trained models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fits = torch.load(\"results/mutrans.backtesting.pt\", map_location=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'We have loaded {len(fits)} models')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print info on available models and what the keys are\n",
    "if True:\n",
    "    for key in fits:\n",
    "        print(f'{key} -- {fits[key][\"weekly_clades_shape\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scale `coef` by 1/100 in all results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ALREADY_SCALED = set()\n",
    "\n",
    "def scale_tensors(x, names={\"coef\"}, scale=0.01, prefix=\"\", verbose=True):\n",
    "    if id(x) in ALREADY_SCALED:\n",
    "        return\n",
    "    if isinstance(x, dict):\n",
    "        for k, v in list(x.items()):\n",
    "            if k in names:\n",
    "                if verbose:\n",
    "                    print(f\"{prefix}.{k}\")\n",
    "                x[k] = v * scale\n",
    "            elif k == \"diagnostics\":\n",
    "                continue\n",
    "            else:\n",
    "                scale_tensors(v, names, scale, f\"{prefix}.{k}\", verbose=verbose)\n",
    "    ALREADY_SCALED.add(id(x))\n",
    "                \n",
    "scale_tensors(fits, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "forecast_dir_prefix = \"paper/backtesting/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Forecasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weekly_clades_to_lineages(weekly_clades, clade_id_to_lineage_id, n_model_lineages):\n",
    "    weekly_lineages = weekly_clades.new_zeros(weekly_clades.shape[:-1] + (n_model_lineages,)).scatter_add_(\n",
    "        -1, clade_id_to_lineage_id.expand_as(weekly_clades), weekly_clades)\n",
    "    return weekly_lineages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plusminus(mean, std):\n",
    "    p95 = 1.96 * std\n",
    "    return torch.stack([mean - p95, mean, mean + p95])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyrocov.util import (\n",
    "    pretty_print, pearson_correlation, quotient_central_moments, generate_colors\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split(a, n):\n",
    "    k, m = divmod(len(a), n)\n",
    "    return (a[i*k+min(i, m):(i+1)*k+min(i+1, m)] for i in range(n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def select_lineages_for_plot(\n",
    "        weekly_lineages,\n",
    "        num_lineages,\n",
    "        lineage_id_inv,\n",
    "        location_ids, # location ids\n",
    "        nbins = 10,\n",
    "        additional_lineages = [],\n",
    "    ):\n",
    "    \"\"\"Return names of lineages for plot\"\"\"\n",
    "    \n",
    "    keep_per_bin = math.ceil(num_lineages / nbins)\n",
    "    T = weekly_lineages.shape[0]\n",
    "    time_intervals = list(split(np.arange(T), nbins))\n",
    "    lineage_ids = []\n",
    "    for interval in time_intervals:\n",
    "        kept_lineage_ids = weekly_lineages[interval][:, location_ids].sum([0, 1]).sort(-1, descending=True).indices[:keep_per_bin]\n",
    "        lineage_ids.append(kept_lineage_ids)\n",
    "    lineage_ids = torch.cat(lineage_ids)\n",
    "    additional_indexes = list(lineage_id_inv.index(x) for x in additional_lineages)\n",
    "    lineage_ids = torch.cat((lineage_ids, torch.tensor(additional_indexes))).tolist()\n",
    "\n",
    "    return sorted(set(lineage_id_inv[x] for x in lineage_ids)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_colors_from_lineage_names(lineage_names):\n",
    "    \n",
    "    standard_colors_dict = {\n",
    "        'BA.1': cc.glasbey[0],\n",
    "        'BA.2': cc.glasbey[1],\n",
    "        'BA.1.1': cc.glasbey[2],\n",
    "        'AY.4': cc.glasbey[3],\n",
    "        'B.1.1.7': cc.glasbey[4],\n",
    "        'B.1.1': cc.glasbey[5],\n",
    "        'B.1.177': cc.glasbey[6],\n",
    "    }\n",
    "    \n",
    "    glasbey_offset = len(standard_colors_dict)\n",
    "    \n",
    "    colors = []\n",
    "    for lineage_name in lineage_names:\n",
    "        try:\n",
    "            color = standard_colors_dict[lineage_name]\n",
    "        except KeyError:\n",
    "            color = cc.glasbey[glasbey_offset]\n",
    "            glasbey_offset += 1\n",
    "        colors.append(color)\n",
    "        \n",
    "    return colors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_forecast2(fit, input_dataset, queries, num_lineages=10, filenames=[], verbose=False, additional_lineages = ['BA.2'], nbins=5, legend_out=False, figsize_x = None):\n",
    "    # Convert queries to array if only only string\n",
    "    if isinstance(queries, str):\n",
    "        queries = [queries]\n",
    "    \n",
    "    # Get dimensions of the model fit (T,P,L) these are probabilities\n",
    "    n_model_periods, n_model_places, n_model_lineages = fit['mean']['probs'].shape\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'n_model_periods: {n_model_periods}')\n",
    "        print(f'n_model_places: {n_model_places}')\n",
    "        print(f'n_model_lineages: {n_model_lineages}')\n",
    "    \n",
    "    # Get dimensions of weekly_cases (T,P) these are JHU counts\n",
    "    weekly_cases_fit = fit['weekly_cases']\n",
    "    n_cases_periods, n_cases_places = weekly_cases_fit.shape\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'n_cases_periods: {n_cases_periods}')\n",
    "        print(f'n_cases_places: {n_cases_places}')\n",
    "    \n",
    "    # Some checks\n",
    "    assert n_cases_places == n_model_places\n",
    "    assert n_model_periods > n_cases_periods\n",
    "    \n",
    "    # Calculate how many periods are forecasted (i.e. are beyond the input to the model)\n",
    "    n_forecast_steps = n_model_periods - n_cases_periods\n",
    "    if (verbose):\n",
    "        print(f'n_forecast_steps: {n_forecast_steps}')\n",
    "        \n",
    "    # Weekly case counts by time place and clade obtained from the fit\n",
    "    weekly_clades_fit = fit['weekly_clades'] # T, P, C\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'weekly_clades_fit shape: {weekly_clades_fit.shape}')\n",
    "    \n",
    "    # Weekly case counts by time place and clade obtain from the input data\n",
    "    # This has more time point and more regions than the one from the fit\n",
    "    weekly_clades_data = input_dataset['weekly_clades']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'weekly_clades_data shape: {weekly_clades_data.shape}')\n",
    "    \n",
    "    # Mapping from clades to lineages, a tensor of indexes\n",
    "    # This is valid for both the fit and the input_data\n",
    "    clade_id_to_lineage_id = input_dataset['clade_id_to_lineage_id']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'clade_id_to_lineage_id length: {len(clade_id_to_lineage_id)}')\n",
    "        \n",
    "    # We don't have clade_id_to_lineage_id in the fit -- it should in principle be the same\n",
    "    \n",
    "    # Summarize the counts of the weekly_clades (from data or fit) to the number of lineages in the model\n",
    "    weekly_lineages_data = weekly_clades_to_lineages(weekly_clades_data, clade_id_to_lineage_id, n_model_lineages)\n",
    "    weekly_lineages_fit = weekly_clades_to_lineages(weekly_clades_fit, clade_id_to_lineage_id, n_model_lineages)\n",
    "    \n",
    "    # Add CI to the probs\n",
    "    probs = plusminus(fit['mean']['probs'], fit['std']['probs']) # [3,T,P,L]\n",
    "    \n",
    "    # Expand weekly_cases_fit (JHU counts) from the model to cover the steps we are forecasting\n",
    "    padding = 1 + weekly_cases_fit.mean(0, keepdim=True).expand(n_forecast_steps, -1)\n",
    "    weekly_cases_fit_ = torch.cat([weekly_cases_fit, padding], 0)\n",
    "    weekly_cases_fit_.add_(10)\n",
    "    # Generate predictions\n",
    "    # Note: For the evaluation maybe we are better off comparing probabilities not counts\n",
    "    predicted = probs * weekly_cases_fit_[..., None]\n",
    "    \n",
    "    # This is an array of strings listing the locations for the data\n",
    "    location_id_inv_data = input_dataset['location_id_inv']\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'location_id_inv_data length: {len(location_id_inv_data)}')\n",
    "    \n",
    "    # This is an array of strings listing the locations for the fit\n",
    "    location_id_inv_fit = fit['location_id_inv']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'location_id_inv_fit length: {len(location_id_inv_fit)}')\n",
    "    \n",
    "    # Get the location indexes that we want to keep based on query for the data\n",
    "    ids_fit = torch.tensor([i for i, name in enumerate(location_id_inv_fit) if any(q in name for q in queries)])\n",
    "    \n",
    "    # These are the lineage labels, we can get them from either the fit or the dataset. \n",
    "    # We assume that these are identical and we assert this below\n",
    "    lineage_id_inv_fit = fit['lineage_id_inv']\n",
    "    lineage_id_inv_data = input_dataset['lineage_id_inv']\n",
    "    assert lineage_id_inv_fit == lineage_id_inv_data\n",
    "    \n",
    "    # Subset weekly_lineages_fit to those location sum over time and place and get the indices in descending order\n",
    "    plot_lineages_ids_inv_fit = select_lineages_for_plot(\n",
    "        weekly_lineages = weekly_lineages_fit,\n",
    "        num_lineages = num_lineages,\n",
    "        lineage_id_inv = lineage_id_inv_fit,\n",
    "        location_ids = ids_fit, \n",
    "        nbins = nbins,\n",
    "        additional_lineages = additional_lineages,\n",
    "    )\n",
    "    \n",
    "    # tbw\n",
    "    plot_lineages_ids_inv_pred = select_lineages_for_plot(\n",
    "        weekly_lineages = fit['mean']['probs'],\n",
    "        num_lineages = num_lineages,\n",
    "        lineage_id_inv = lineage_id_inv_fit,\n",
    "        location_ids = ids_fit, \n",
    "        nbins = nbins,\n",
    "        additional_lineages = additional_lineages,\n",
    "    )\n",
    "\n",
    "    # Same thing for the data\n",
    "    ids_data = torch.tensor([ i for i, name in enumerate(location_id_inv_data) if any(q in name for q in queries)])\n",
    "    \n",
    "    # Subset weekly_lineages_fit to those location sum over time and place and get the indices in descending order\n",
    "    plot_lineages_ids_inv_data = select_lineages_for_plot(\n",
    "        weekly_lineages = weekly_lineages_data,\n",
    "        num_lineages = num_lineages,\n",
    "        lineage_id_inv = lineage_id_inv_data,\n",
    "        location_ids = ids_data, \n",
    "        nbins = nbins,\n",
    "        additional_lineages = additional_lineages,\n",
    "    )\n",
    "    \n",
    "    # merge the lineage name from datset and fit to get a single list\n",
    "    lineage_name_to_index_map_data = { l:i for i, l in enumerate(lineage_id_inv_data)}\n",
    "    lineage_name_to_index_map_fit = { l:i for i, l in enumerate(lineage_id_inv_fit)}\n",
    "    \n",
    "    plot_lineages_ids_inv_joint = sorted(\n",
    "        set(plot_lineages_ids_inv_fit)\n",
    "            .union(plot_lineages_ids_inv_data)\n",
    "            .union(plot_lineages_ids_inv_pred))\n",
    "    \n",
    "    lineage_ids_fit = list(map(lineage_name_to_index_map_fit.get, plot_lineages_ids_inv_joint))\n",
    "    lineage_ids_data = list(map(lineage_name_to_index_map_data.get, plot_lineages_ids_inv_joint))\n",
    "    assert lineage_ids_fit == lineage_ids_data\n",
    "    \n",
    "    # we may have a few plotted lineages now...\n",
    "    num_lineages = len(plot_lineages_ids_inv_joint)\n",
    "    \n",
    "    # Get some colors to plot with\n",
    "    colors = generate_colors_from_lineage_names(plot_lineages_ids_inv_joint)\n",
    "    assert len(colors) >= num_lineages\n",
    "    light = '#bbbbbb'\n",
    "    dark = '#444444'\n",
    "    \n",
    "    # Generate Figure\n",
    "    if figsize_x is None:\n",
    "        figsize_x = 8\n",
    "    \n",
    "    fig, axes = plt.subplots(len(queries), figsize=(figsize_x, 0.5 + 2.5 * len(queries)), sharex=True)\n",
    "    if not isinstance(axes, (list, np.ndarray)):\n",
    "        axes = [axes]\n",
    "    \n",
    "    # Get x axis dates for plotting\n",
    "    dates = matplotlib.dates.date2num(mutrans.date_range(len(fit[\"mean\"][\"probs\"])))\n",
    "\n",
    "    # Query (region) plotting loop\n",
    "    for row, (query, ax) in enumerate(zip(queries, axes)):\n",
    "        # location ids for this query (some queries are made of multiple regions)\n",
    "        ids_fit = torch.tensor([i for i, name in enumerate(location_id_inv_fit) if query in name])\n",
    "        if verbose:\n",
    "            print('---')\n",
    "            print(f\"{query} matched {len(ids_fit)} regions in the fit\")\n",
    "        \n",
    "        # location ids for this query in the data\n",
    "        ids_data = torch.tensor([i for i, name in enumerate(location_id_inv_data) if query in name])\n",
    "        \n",
    "        if len(axes) > 1:\n",
    "            # Plot weekly cases total\n",
    "            counts = weekly_cases_fit[:, ids_fit].sum(1)\n",
    "            if verbose:\n",
    "                print(f\"{query}: max {counts.max():g}, total {counts.sum():g}\")\n",
    "            counts /= counts.max()\n",
    "            ax.plot(dates[:len(counts)], counts, \"k-\", color=light, lw=0.8, zorder=-20)\n",
    "            \n",
    "            # Plot weekly lineages total we are getting the data from the fit not the dataset\n",
    "            counts = weekly_lineages_fit[:, ids_fit].sum([1, 2])\n",
    "            counts /= counts.max()\n",
    "            ax.plot(dates[:len(counts)], counts, \"k--\", color=light, lw=1, zorder=-20)\n",
    "            \n",
    "        # Get the predictions for the relevant regions, normalize\n",
    "        pred = predicted.index_select(-2, ids_fit).sum(-2)\n",
    "        pred /= pred[1].sum(-1, True).clamp_(min=1e-20)\n",
    "        \n",
    "        # Get the observations for the relevant regions\n",
    "        obs = weekly_lineages_fit[:, ids_fit].sum(1)\n",
    "        obs /= obs.sum(-1, True).clamp_(min=1e-9)\n",
    "        \n",
    "        # Observations from the data -- this extends further in the time dimension\n",
    "        obs_data = weekly_lineages_data[:, ids_data].sum(1)\n",
    "        obs_data /= obs_data.sum(-1, True).clamp(min=1e-9)\n",
    "        \n",
    "        # lineage plotting loop\n",
    "        for s, color in zip(lineage_ids_fit, colors):\n",
    "            lb, mean, ub = pred[..., s]\n",
    "            ax.fill_between(dates, lb, ub, color=color, alpha=0.2, zorder=-10)\n",
    "            ax.plot(dates, mean, color=color, lw=1, zorder=-9)\n",
    "            # Get the lineage label\n",
    "            lineage = lineage_id_inv_fit[s]\n",
    "            ax.plot(dates[:len(obs)], obs[:, s], color=color, lw=0, marker='o', markersize=3,\n",
    "                    label=lineage if row == 0 else None)\n",
    "        \n",
    "        # Plot observations from the dataset for all the forecast points\n",
    "        # TODO: Fix colors to match (we probably want to fix \"sort(-1, descending=True)\" to be a matching permutation instead)\n",
    "        for s, color in zip(lineage_ids_data, colors):\n",
    "            lineage = lineage_id_inv_data[s]\n",
    "            max_time_step = min((len(obs)+n_forecast_steps), obs_data.shape[0]-1)\n",
    "            \n",
    "            ax.plot(dates[len(obs):max_time_step], obs_data[len(obs):max_time_step, s], label='_nolegend_',\n",
    "                    color=color, lw=0, marker='x', markersize=2)\n",
    "            \n",
    "        # Add shading for the forecast region\n",
    "        ax.axvline(dates[len(obs)], linestyle='--', lw=1, color=(0.5, 0.5, 0.5))\n",
    "        ax.axvspan(dates[len(obs)],dates[len(obs)+n_forecast_steps-1], color=(0.5, 0.5, 0.5), alpha=0.2)\n",
    "        \n",
    "        # Set axis ticks\n",
    "        ax.set_ylim(0, 1)\n",
    "        ax.set_yticks(())\n",
    "        ax.set_ylabel(query.replace(\" / \", \"\\n\"))\n",
    "        ax.set_xlim(dates.min(), dates.max())\n",
    "        \n",
    "        # Print legend\n",
    "        if legend_out:\n",
    "            \n",
    "            if row == 0:\n",
    "                ax.legend(loc=\"upper left\", bbox_to_anchor=(1.01, 1.04), fontsize=10)\n",
    "            elif row == 1:\n",
    "                ax.plot([], \"k--\", color=light, lw=1, label=\"relative #samples\")\n",
    "                ax.plot([], \"k-\", color=light, lw=0.8, label=\"relative #cases\")\n",
    "                ax.plot([], lw=0, marker='o', markersize=3, color='gray',\n",
    "                        label=\"observed portion\")\n",
    "                ax.fill_between([], [], [], color='gray', label=\"predicted portion\")\n",
    "                ax.legend(loc=\"upper left\")\n",
    "        else:\n",
    "            \n",
    "            if row == 0:\n",
    "                ax.legend(loc=\"upper left\", fontsize=8 * (10 / num_lineages) ** 0.8)\n",
    "            elif row == 1:\n",
    "                ax.plot([], \"k--\", color=light, lw=1, label=\"relative #samples\")\n",
    "                ax.plot([], \"k-\", color=light, lw=0.8, label=\"relative #cases\")\n",
    "                ax.plot([], lw=0, marker='o', markersize=3, color='gray',\n",
    "                        label=\"observed portion\")\n",
    "                ax.fill_between([], [], [], color='gray', label=\"predicted portion\")\n",
    "                ax.legend(loc=\"upper left\",)\n",
    "          \n",
    "    # Setup the date axis correctly\n",
    "    ax.xaxis.set_major_locator(matplotlib.dates.MonthLocator())\n",
    "    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter(\"%b %Y\"))\n",
    "    plt.xticks(rotation=90)\n",
    "    plt.subplots_adjust(hspace=0)\n",
    "    \n",
    "    for filename in filenames:\n",
    "        plt.savefig(filename, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate all Forecasting Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if False:\n",
    "    for model_key in list(fits.keys()):\n",
    "        fit_n = fits[model_key]\n",
    "        plot_forecast2(\n",
    "            fit_n, \n",
    "            input_dataset, \n",
    "            queries=[\"England\", \"USA / Ma\", \"Brazil\"],\n",
    "            num_lineages=10,\n",
    "            verbose=False,\n",
    "            filenames =  [f'{forecast_dir_prefix}/backtesting_day_{model_key[9]}.png']\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Selected Forecast Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = list(fits.keys())[7]\n",
    "print(k[9])\n",
    "fit_n = fits[k]\n",
    "plot_forecast2(\n",
    "    fit_n, \n",
    "    input_dataset, \n",
    "    queries=[\"England\"],\n",
    "    num_lineages=14,\n",
    "    verbose=False,\n",
    "    additional_lineages = ['AY.4'],\n",
    "    filenames =  [f'{forecast_dir_prefix}/backtesting_day_{k[9]}_early_prediction_england.png'],\n",
    "    figsize_x = 528 / 752 * 8,\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = list(fits.keys())[23]\n",
    "print(k[9])\n",
    "fit_n = fits[k]\n",
    "plot_forecast2(\n",
    "    fit_n, \n",
    "    input_dataset, \n",
    "    queries=[\"England\"],\n",
    "    num_lineages=13,\n",
    "    verbose=False,\n",
    "    additional_lineages = ['BA.1'],\n",
    "    filenames =  [f'{forecast_dir_prefix}/backtesting_day_{k[9]}_early_prediction_england.png'],\n",
    "    figsize_x = 8\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Country Specific Ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = list(fits.keys())[len( list(fits.keys()))-1]\n",
    "print(k[9])\n",
    "fit_n = fits[k]\n",
    "plot_forecast2(\n",
    "    fit_n, \n",
    "    input_dataset, \n",
    "    queries=[\"USA\",\"France\",\"England\",\"Brazil\",\"Australia\",\"Russia\"],\n",
    "    num_lineages=20,\n",
    "    verbose=False,\n",
    "    additional_lineages = ['BA.1'],\n",
    "    filenames =  [f'{forecast_dir_prefix}/backtesting_day_{k[9]}_early_prediction_USA_France_England_Brazil_Australia_Russia.png'],\n",
    "    legend_out = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = list(fits.keys())[len( list(fits.keys()))-1]\n",
    "print(k[9])\n",
    "fit_n = fits[k]\n",
    "plot_forecast2(\n",
    "    fit_n, \n",
    "    input_dataset, \n",
    "    queries=[\"Asia\",\"Europe\",\"Africa\"],\n",
    "    num_lineages=13,\n",
    "    verbose=False,\n",
    "    additional_lineages = ['BA.1'],\n",
    "    filenames =  [f'{forecast_dir_prefix}/backtesting_day_{k[9]}_early_prediction_Asia_Europe_Africa.png']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the forecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_forecast2(fit, input_dataset, queries, num_lineages=10, filenames=[], verbose=False, data_region = None):\n",
    "    # Convert queries to array if only only string\n",
    "    if isinstance(queries, str):\n",
    "        queries = [queries]\n",
    "    \n",
    "    # Get dimensions of the model fit (T,P,L) these are probabilities\n",
    "    n_model_periods, n_model_places, n_model_lineages = fit['mean']['probs'].shape\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'n_model_periods: {n_model_periods}')\n",
    "        print(f'n_model_places: {n_model_places}')\n",
    "        print(f'n_model_lineages: {n_model_lineages}')\n",
    "    \n",
    "    # Get dimensions of weekly_cases (T,P) these are JHU counts\n",
    "    weekly_cases_fit = fit['weekly_cases']\n",
    "    n_cases_periods, n_cases_places = weekly_cases_fit.shape\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'n_cases_periods: {n_cases_periods}')\n",
    "        print(f'n_cases_places: {n_cases_places}')\n",
    "    \n",
    "    # Some checks\n",
    "    assert n_cases_places == n_model_places\n",
    "    assert n_model_periods > n_cases_periods\n",
    "    \n",
    "    # Calculate how many periods are forecasted (i.e. are beyond the input to the model)\n",
    "    n_forecast_steps = n_model_periods - n_cases_periods\n",
    "    if (verbose):\n",
    "        print(f'n_forecast_steps: {n_forecast_steps}')\n",
    "        \n",
    "    # Weekly case counts by time place and clade obtained from the fit\n",
    "    weekly_clades_fit = fit['weekly_clades'] # T, P, C\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'weekly_clades_fit shape: {weekly_clades_fit.shape}')\n",
    "    \n",
    "    # Weekly case counts by time place and clade obtain from the input data\n",
    "    # This has more time point and more regions than the one from the fit\n",
    "    weekly_clades_data = input_dataset['weekly_clades']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'weekly_clades_data shape: {weekly_clades_data.shape}')\n",
    "    \n",
    "    # Mapping from clades to lineages, a tensor of indexes\n",
    "    # This is valid for both the fit and the input_data\n",
    "    clade_id_to_lineage_id = input_dataset['clade_id_to_lineage_id']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'clade_id_to_lineage_id length: {len(clade_id_to_lineage_id)}')\n",
    "        \n",
    "    # We don't have clade_id_to_lineage_id in the fit -- it should in principle be the same\n",
    "    \n",
    "    # Summarize the counts of the weekly_clades (from data or fit) to the number of lineages in the model\n",
    "    weekly_lineages_data = weekly_clades_to_lineages(weekly_clades_data, clade_id_to_lineage_id, n_model_lineages)\n",
    "    weekly_lineages_fit = weekly_clades_to_lineages(weekly_clades_fit, clade_id_to_lineage_id, n_model_lineages)\n",
    "    \n",
    "    # Get the probs\n",
    "    probs = fit['mean']['probs']\n",
    "    #probs = plusminus(fit['mean']['probs'], fit['std']['probs']) # [3,T,P,L]\n",
    "    \n",
    "    # Expand weekly_cases_fit (JHU counts) from the model to cover the steps we are forecasting\n",
    "    #padding = 1 + weekly_cases_fit.mean(0, keepdim=True).expand(n_forecast_steps, -1)\n",
    "    #weekly_cases_fit_ = torch.cat([weekly_cases_fit, padding], 0)\n",
    "    # Generate predictions\n",
    "    # Note: For the evaluation maybe we are better off comparing probabilities not counts\n",
    "    #predicted = probs * weekly_cases_fit_[..., None]\n",
    "    \n",
    "    # This is an array of strings listing the locations for the data\n",
    "    location_id_inv_data = input_dataset['location_id_inv']\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'location_id_inv_data length: {len(location_id_inv_data)}')\n",
    "    \n",
    "    # This is an array of strings listing the locations for the fit\n",
    "    location_id_inv_fit = fit['location_id_inv']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'location_id_inv_fit length: {len(location_id_inv_fit)}')\n",
    "    \n",
    "    # Get the location indexes that we want to keep based on query for the data\n",
    "    ids_fit = torch.tensor([i for i, name in enumerate(location_id_inv_fit) if any(q in name for q in queries)])\n",
    "    \n",
    "    # Subset weekly_lineages_fit to those location sum over time and place and get the indices in descending order\n",
    "    lineage_ids_fit = weekly_lineages_fit[:, ids_fit].sum([0, 1]).sort(-1, descending=True).indices\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'lineage_ids_fit shape: {lineage_ids_fit.shape}')\n",
    "    # Keep only the top n number of lineages we want to plot\n",
    "    lineage_ids_fit = lineage_ids_fit[:num_lineages]\n",
    "\n",
    "    # This is problematic without fixing the above permutation\n",
    "    # TODO: Add assert that they are the same set / eliminate code\n",
    "    # Check if order of \n",
    "    lineage_ids_data = lineage_ids_fit[:num_lineages]\n",
    "    \n",
    "    # These are the lineage labels, we can get them from either the fit or the dataset. \n",
    "    # We assume that these are identical and we assert this below\n",
    "    lineage_id_inv_fit = fit['lineage_id_inv']\n",
    "    lineage_id_inv_data = input_dataset['lineage_id_inv']\n",
    "    assert lineage_id_inv_fit == lineage_id_inv_data\n",
    "    \n",
    "    # Get shared locations between full dataset and fit dataset\n",
    "    common_regions = list(set(location_id_inv_fit).intersection(set(location_id_inv_data)))\n",
    "    \n",
    "    if data_region is not None:\n",
    "        common_regions = list(set(common_regions).intersection(set(data_region)))\n",
    "    \n",
    "    # Get indexes of these common regions for each set\n",
    "    common_regions_fit_inv_map = []\n",
    "    common_regions_data_inv_map = []\n",
    "    for r in common_regions:\n",
    "        common_regions_fit_inv_map.append(location_id_inv_fit.index(r))\n",
    "        common_regions_data_inv_map.append(location_id_inv_data.index(r))\n",
    "        \n",
    "    # We want to compare empirical and predicted probabilities for the forecast interval\n",
    "    probs = probs[n_cases_periods:,common_regions_fit_inv_map,:]\n",
    "    \n",
    "    # Subset observed to relevant periods and regions\n",
    "    obs_data = weekly_lineages_data[n_cases_periods:n_cases_periods+n_forecast_steps,common_regions_data_inv_map,:]\n",
    "    empirical_probs = obs_data / obs_data.sum(-1,True).clamp_(min=1e-9)\n",
    "    \n",
    "    # Truncate to availanle data\n",
    "    probs = probs[:empirical_probs.shape[0],]\n",
    "    \n",
    "    # Calculate errors\n",
    "    l1_error = (probs - empirical_probs).abs().sum([-1,-2]) / probs.shape[-2]\n",
    "    l2_error = (probs - empirical_probs).pow(2).sum([-1,-2]).sqrt() / probs.shape[-2]\n",
    "\n",
    "    # consider spearman error\n",
    "    # correlations on the probabilities (average over time)\n",
    "    # precision at k\n",
    "    return {\n",
    "        'L1_error': l1_error,\n",
    "        'L2_error': l2_error,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_forecast_eval(fits, input_dataset, data_region = None, queries = None):\n",
    "    model_keys = list(fits.keys())\n",
    "    \n",
    "    if not queries:\n",
    "        queries = input_dataset['location_id_inv']\n",
    "\n",
    "    forecast_start_days = []\n",
    "    period_forecast_ahead = []\n",
    "    l1_error = []\n",
    "    l2_error = []\n",
    "    \n",
    "    period_length = 14\n",
    "\n",
    "    for key in model_keys:\n",
    "        forecast_start_day = key[9]\n",
    "        fit_n = fits[key]\n",
    "        # Get forecast error for all independent location ids \n",
    "        forecast_error = evaluate_forecast2(\n",
    "            fit_n, \n",
    "            input_dataset, \n",
    "            queries = queries,\n",
    "            num_lineages=100,\n",
    "            data_region = data_region,\n",
    "        verbose=False)\n",
    "        n_periods_forecast = len(forecast_error['L1_error'].tolist())\n",
    "        forecast_start_days.extend([forecast_start_day] * n_periods_forecast)\n",
    "        period_forecast_ahead.extend(list(range(1,n_periods_forecast+1)))\n",
    "        l1_error.extend(forecast_error['L1_error'].tolist())\n",
    "        l2_error.extend(forecast_error['L2_error'].tolist())\n",
    "        \n",
    "    df1 = pd.DataFrame({\n",
    "    'forecast_start_days': forecast_start_days,\n",
    "    'period_forecast_ahead': period_forecast_ahead, \n",
    "    'l1_error': l1_error, \n",
    "    'l2_error': l2_error})\n",
    "    \n",
    "    df1['day_of_forecast'] = df1['forecast_start_days'] + df1['period_forecast_ahead'] * 14\n",
    "    \n",
    "    return df1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_region_forecast = generate_forecast_eval(fits, input_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matplotlib.rcParams['figure.figsize'] = [6,4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = sns.boxplot(x=\"period_forecast_ahead\", y=\"l1_error\", data=all_region_forecast, palette='rainbow')\n",
    "ax.set(xlabel = '2-week period forecast ahead', ylabel=\"L1 Error\")\n",
    "ax.set_ylim([0.0,2.0])\n",
    "plt.savefig('paper/backtesting/L1_error_barplot_all.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Top 100 region forecasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get top covered regions\n",
    "top_region_idx = input_dataset['weekly_clades'].sum([0,2]).sort(-1, descending=True).indices[:100].tolist()\n",
    "regions = list(input_dataset['location_id_inv'][x] for x in top_region_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_region_forecast = generate_forecast_eval(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = sns.boxplot(x=\"period_forecast_ahead\", y=\"l1_error\", data=top_region_forecast, palette='rainbow')\n",
    "ax.set(xlabel = '2-week period forecast ahead', ylabel=\"L1 Error\")\n",
    "ax.set_ylim([0.0,2.0])\n",
    "plt.savefig('paper/backtesting/L1_error_barplot_top100.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Top 100-200 region forecasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get top covered regions\n",
    "top_region_idx = input_dataset['weekly_clades'].sum([0,2]).sort(-1, descending=True).indices[100:1000].tolist()\n",
    "regions = list(input_dataset['location_id_inv'][x] for x in top_region_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_region_forecast = generate_forecast_eval(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = sns.boxplot(x=\"period_forecast_ahead\", y=\"l1_error\", data=top_region_forecast, palette='rainbow')\n",
    "ax.set(xlabel = '2-week period forecast ahead', ylabel=\"L1 Error\")\n",
    "ax.set_ylim([0.0,2.0])\n",
    "plt.savefig('paper/backtesting/L1_error_barplot_top100-1000.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Other region forecasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get top covered regions\n",
    "bottom_region_idx = input_dataset['weekly_clades'].sum([0,2]).sort(-1, descending=True).indices[100:].tolist()\n",
    "regions = list(input_dataset['location_id_inv'][x] for x in bottom_region_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bottom_region_forecast = generate_forecast_eval(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = sns.boxplot(x=\"period_forecast_ahead\", y=\"l1_error\", data=bottom_region_forecast, palette='rainbow')\n",
    "ax.set(xlabel = '2-week period forecast ahead', ylabel=\"L1 Error\")\n",
    "ax.set_ylim([0.0,2.0])\n",
    "plt.savefig('paper/backtesting/L1_error_barplot_other.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation of forecasting accuracy\n",
    "\n",
    " - What are we trying to do? For a given region and for all models get a % of how often we predict the correct strain n intervals ahead\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_forecast3(fit, input_dataset, queries, num_lineages=10, verbose=False, data_region = None):    \n",
    "    # Convert queries to array if only only string\n",
    "    if isinstance(queries, str):\n",
    "        queries = [queries]\n",
    "    \n",
    "    # Get dimensions of the model fit (T,P,L) these are probabilities\n",
    "    n_model_periods, n_model_places, n_model_lineages = fit['mean']['probs'].shape\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'n_model_periods: {n_model_periods}')\n",
    "        print(f'n_model_places: {n_model_places}')\n",
    "        print(f'n_model_lineages: {n_model_lineages}')\n",
    "    \n",
    "    # Get dimensions of weekly_cases (T,P) these are JHU counts\n",
    "    weekly_cases_fit = fit['weekly_cases']\n",
    "    n_cases_periods, n_cases_places = weekly_cases_fit.shape\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'n_cases_periods: {n_cases_periods}')\n",
    "        print(f'n_cases_places: {n_cases_places}')\n",
    "    \n",
    "    # Some checks\n",
    "    assert n_cases_places == n_model_places\n",
    "    assert n_model_periods > n_cases_periods\n",
    "    \n",
    "    # Calculate how many periods are forecasted (i.e. are beyond the input to the model)\n",
    "    n_forecast_steps = n_model_periods - n_cases_periods\n",
    "    if (verbose):\n",
    "        print(f'n_forecast_steps: {n_forecast_steps}')\n",
    "        \n",
    "    # Weekly case counts by time place and clade obtained from the fit\n",
    "    weekly_clades_fit = fit['weekly_clades'] # T, P, C\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'weekly_clades_fit shape: {weekly_clades_fit.shape}')\n",
    "    \n",
    "    # Weekly case counts by time place and clade obtain from the input data\n",
    "    # This has more time point and more regions than the one from the fit\n",
    "    weekly_clades_data = input_dataset['weekly_clades']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'weekly_clades_data shape: {weekly_clades_data.shape}')\n",
    "    \n",
    "    # Mapping from clades to lineages, a tensor of indexes\n",
    "    # This is valid for both the fit and the input_data\n",
    "    clade_id_to_lineage_id = input_dataset['clade_id_to_lineage_id']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'clade_id_to_lineage_id length: {len(clade_id_to_lineage_id)}')\n",
    "        \n",
    "    # We don't have clade_id_to_lineage_id in the fit -- it should in principle be the same\n",
    "    \n",
    "    # Summarize the counts of the weekly_clades (from data or fit) to the number of lineages in the model\n",
    "    weekly_lineages_data = weekly_clades_to_lineages(weekly_clades_data, clade_id_to_lineage_id, n_model_lineages)\n",
    "    weekly_lineages_fit = weekly_clades_to_lineages(weekly_clades_fit, clade_id_to_lineage_id, n_model_lineages)\n",
    "    \n",
    "    # Get the probs\n",
    "    probs = fit['mean']['probs']\n",
    "    #probs = plusminus(fit['mean']['probs'], fit['std']['probs']) # [3,T,P,L]\n",
    "    \n",
    "    # Expand weekly_cases_fit (JHU counts) from the model to cover the steps we are forecasting\n",
    "    #padding = 1 + weekly_cases_fit.mean(0, keepdim=True).expand(n_forecast_steps, -1)\n",
    "    #weekly_cases_fit_ = torch.cat([weekly_cases_fit, padding], 0)\n",
    "    # Generate predictions\n",
    "    # Note: For the evaluation maybe we are better off comparing probabilities not counts\n",
    "    #predicted = probs * weekly_cases_fit_[..., None]\n",
    "    \n",
    "    # This is an array of strings listing the locations for the data\n",
    "    location_id_inv_data = input_dataset['location_id_inv']\n",
    "    if (verbose):\n",
    "        print('---')\n",
    "        print(f'location_id_inv_data length: {len(location_id_inv_data)}')\n",
    "    \n",
    "    # This is an array of strings listing the locations for the fit\n",
    "    location_id_inv_fit = fit['location_id_inv']\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'location_id_inv_fit length: {len(location_id_inv_fit)}')\n",
    "    \n",
    "    # Get the location indexes that we want to keep based on query for the data\n",
    "    ids_fit = torch.tensor([i for i, name in enumerate(location_id_inv_fit) if any(q in name for q in queries)])\n",
    "    \n",
    "    # Subset weekly_lineages_fit to those location sum over time and place and get the indices in descending order\n",
    "    lineage_ids_fit = weekly_lineages_fit[:, ids_fit].sum([0, 1]).sort(-1, descending=True).indices\n",
    "    if verbose:\n",
    "        print('---')\n",
    "        print(f'lineage_ids_fit shape: {lineage_ids_fit.shape}')\n",
    "    # Keep only the top n number of lineages we want to plot\n",
    "    lineage_ids_fit = lineage_ids_fit[:num_lineages]\n",
    "\n",
    "    # This is problematic without fixing the above permutation\n",
    "    # TODO: Add assert that they are the same set / eliminate code\n",
    "    # Check if order of \n",
    "    lineage_ids_data = lineage_ids_fit[:num_lineages]\n",
    "    \n",
    "    # These are the lineage labels, we can get them from either the fit or the dataset. \n",
    "    # We assume that these are identical and we assert this below\n",
    "    lineage_id_inv_fit = fit['lineage_id_inv']\n",
    "    lineage_id_inv_data = input_dataset['lineage_id_inv']\n",
    "    assert lineage_id_inv_fit == lineage_id_inv_data\n",
    "    \n",
    "    # Get shared locations between full dataset and fit dataset\n",
    "    common_regions = list(set(location_id_inv_fit).intersection(set(location_id_inv_data)))\n",
    "    \n",
    "    if data_region is not None:\n",
    "        common_regions = list(set(common_regions).intersection(set(data_region)))\n",
    "    \n",
    "    # Get indexes of these common regions for each set\n",
    "    common_regions_fit_inv_map = []\n",
    "    common_regions_data_inv_map = []\n",
    "    for r in common_regions:\n",
    "        common_regions_fit_inv_map.append(location_id_inv_fit.index(r))\n",
    "        common_regions_data_inv_map.append(location_id_inv_data.index(r))\n",
    "        \n",
    "    # We want to compare empirical and predicted probabilities for the forecast interval\n",
    "    probs = probs[n_cases_periods:,common_regions_fit_inv_map,:]\n",
    "    \n",
    "    # Subset observed to relevant periods and regions\n",
    "    obs_data = weekly_lineages_data[n_cases_periods:n_cases_periods+n_forecast_steps,common_regions_data_inv_map,:]\n",
    "    empirical_probs = obs_data / obs_data.sum(-1,True).clamp_(min=1e-9)\n",
    "    \n",
    "    # Truncate to availanle data\n",
    "    probs = probs[:empirical_probs.shape[0]-1,]\n",
    "    \n",
    "    return {\n",
    "        'probs': probs,\n",
    "        'empirical_probs': empirical_probs,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_forecast_eval_percent(fits, input_dataset, data_region = None, queries = None):\n",
    "    model_keys = list(fits.keys())\n",
    "    \n",
    "\n",
    "    match_4wk = []\n",
    "    match_8wk = []\n",
    "    \n",
    "    if queries is None:\n",
    "        queries = input_dataset['location_id_inv']\n",
    "\n",
    "    for key in model_keys:\n",
    "        forecast_start_day = key[9]\n",
    "        fit_n = fits[key]\n",
    "        # Get forecast error for all independent location ids \n",
    "        probs_dict = evaluate_forecast3(\n",
    "            fit_n, \n",
    "            input_dataset, \n",
    "            queries = queries,\n",
    "            num_lineages=101,\n",
    "            data_region = data_region,\n",
    "        verbose=False)\n",
    "        \n",
    "        try:\n",
    "            period_index_4wk = 1\n",
    "            predicted_4wk = probs_dict['probs'][period_index_4wk,:].sum(-2).argmax(0).item()\n",
    "            actual_4wk = probs_dict['empirical_probs'][period_index_4wk,:].sum(-2).argmax(0).item()      \n",
    "\n",
    "            period_index_8wk = 3\n",
    "            predicted_8wk = probs_dict['probs'][period_index_8wk,:].sum(-2).argmax(0).item()\n",
    "            actual_8wk = probs_dict['empirical_probs'][period_index_8wk,].sum(-2).argmax(0).item()\n",
    "\n",
    "            match_4wk.append(predicted_4wk == actual_4wk)\n",
    "            match_8wk.append(predicted_8wk == actual_8wk)\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "    return {\n",
    "        'match_4wk': match_4wk,\n",
    "        'match_8wk': match_8wk,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### USA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'USA'\n",
    "regions = list(x for x in input_dataset['location_id'].keys() if query in x)\n",
    "selected_region_forecast = generate_forecast_eval_percent(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_4wk']).sum() / len(selected_region_forecast['match_4wk']) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_8wk']).sum() / len(selected_region_forecast['match_8wk']) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### France"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'France'\n",
    "regions = list(x for x in input_dataset['location_id'].keys() if query in x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_region_forecast = generate_forecast_eval_percent(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_4wk']).sum() / len(selected_region_forecast['match_4wk']) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_8wk']).sum() / len(selected_region_forecast['match_8wk']) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### England"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'England'\n",
    "regions = list(x for x in input_dataset['location_id'].keys() if query in x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_region_forecast = generate_forecast_eval_percent(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_4wk']).sum() / len(selected_region_forecast['match_4wk']) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_8wk']).sum() / len(selected_region_forecast['match_8wk']) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Brazil"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'Brazil'\n",
    "regions = list(x for x in input_dataset['location_id'].keys() if query in x)\n",
    "selected_region_forecast = generate_forecast_eval_percent(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_4wk']).sum() / len(selected_region_forecast['match_4wk']) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_8wk']).sum() / len(selected_region_forecast['match_8wk']) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Australia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'Australia'\n",
    "regions = list(x for x in input_dataset['location_id'].keys() if query in x)\n",
    "selected_region_forecast = generate_forecast_eval_percent(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_4wk']).sum() / len(selected_region_forecast['match_4wk']) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_8wk']).sum() / len(selected_region_forecast['match_8wk']) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Russia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'Russia'\n",
    "regions = list(x for x in input_dataset['location_id'].keys() if query in x)\n",
    "selected_region_forecast = generate_forecast_eval_percent(fits, input_dataset, data_region = regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_4wk']).sum() / len(selected_region_forecast['match_4wk']) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(selected_region_forecast['match_8wk']).sum() / len(selected_region_forecast['match_8wk']) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}