skyline_ne.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 353,
   "id": "2adc8e88",
   "metadata": {},
   "outputs": [],
   "source": [
    "import baltic as bt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "from matplotlib import pyplot as plt\n",
    "from datetime import datetime as dt\n",
    "from datetime import timedelta\n",
    "import time\n",
    "#import pymc3\n",
    "import math\n",
    "import arviz as az\n",
    "#from hpd import hpd\n",
    "import scipy.stats as stats\n",
    "from io import StringIO\n",
    "import altair as alt\n",
    "from altair import datum\n",
    "alt.data_transformers.disable_max_rows()\n",
    "import seaborn as sns\n",
    "\n",
    "from zipfile import ZipFile\n",
    "import scipy as sp\n",
    "\n",
    "\n",
    "import sys, subprocess, glob, os, shutil, re, importlib\n",
    "from subprocess import call\n",
    "import imp\n",
    "\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib as mpl\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.patheffects as path_effects\n",
    "import matplotlib.lines as mlines\n",
    "from matplotlib.font_manager import FontProperties\n",
    "import matplotlib.colors as clr\n",
    "from matplotlib import rc\n",
    "import textwrap as textwrap\n",
    "from textwrap import wrap\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from scipy.special import binom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 354,
   "id": "cd166b80",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_file_path = \"../results/mascot_skyline_updated_100_nes_1.log\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 355,
   "id": "3639eac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_in_Ne_changes_mascot(log_file_path):\n",
    "    \n",
    "    Ne_skyline_dict = {\"sample\":[]}\n",
    "    \n",
    "    with open(log_file_path, \"r\") as infile:\n",
    "        line_number = 0\n",
    "        for line in infile:\n",
    "            line_number += 1\n",
    "            if not line.startswith(\"#\"):  # log combiner will sometimes put the entire xml at the start of the log file\n",
    "                # use the first line to find the migration rate columns\n",
    "                #print(line)\n",
    "            # use the first line to find the migration rate columns\n",
    "                if \"posterior\" in line:\n",
    "                    all_cols = line.split(\"\\t\")\n",
    "                    Ne_column_indices = []   # list to store column indices\n",
    "                    Nes_key = {}   # dictionary to store the column index to map to column name\n",
    "\n",
    "                    for i in range(len(all_cols)):\n",
    "                        col = all_cols[i]\n",
    "                        if \"Ne.\" in col:\n",
    "                            Ne_column_indices.append(i)\n",
    "\n",
    "                    # make an empty dictionary to store Nes and generate dictionary to convert index to name\n",
    "                    for n in Ne_column_indices:\n",
    "                        name = line.split(\"\\t\")[n]\n",
    "                        deme = name.split(\".\")[1]# the syntax here is \"Ne_region.1\" where region is deme and 1 is interval 1\n",
    "                        interval = name.split(\".\")[2]\n",
    "                       \n",
    "                        Nes_key[n] = name\n",
    "                        Ne_skyline_dict[name] = []\n",
    "\n",
    "\n",
    "                # read in actual parameter estimates and store in dictionary\n",
    "                else:\n",
    "                    sample = line.split(\"\\t\")[0]\n",
    "                    Ne_skyline_dict[\"sample\"].append(sample)\n",
    "\n",
    "                    for index in Ne_column_indices:\n",
    "                        name = Nes_key[index]\n",
    "                        Ne_skyline_dict[name].append(line.split(\"\\t\")[index])\n",
    "                    \n",
    "                \n",
    "    return(Ne_skyline_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 356,
   "id": "28d616f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ne_skyline = read_in_Ne_changes_mascot(log_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 357,
   "id": "54e3d511",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'dict' object has no attribute 'head'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[357], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mNe_skyline\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mhead\u001b[49m()\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'dict' object has no attribute 'head'"
     ]
    }
   ],
   "source": [
    "Ne_skyline.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "840295d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a new dataframe that summarizes the 95% HPD estimate with mean for each deme and interval \n",
    "def generate_summary_df(input_df):\n",
    "    \n",
    "    \n",
    "    new_df = pd.DataFrame()\n",
    "\n",
    "    for i in input_df.columns.tolist():\n",
    "        if \"Ne.\" in i:\n",
    "            deme = i.split(\".\")[1]\n",
    "           # print(deme)\n",
    "            #interval = \n",
    "            #print(interval)\n",
    "            if \"\\n\" in i.split(\".\")[2]:\n",
    "                interval = i.split(\".\")[2][0:2]\n",
    "            else:\n",
    "                interval = i.split(\".\")[2]\n",
    "           # print(interval)\n",
    "            #print(interval)\n",
    "            #print(i)\n",
    "            #next_interval = int(interval)+1\n",
    "            local_series = input_df[i].astype('float').to_numpy()\n",
    "            #print(local_series)\n",
    "            mean_log = local_series.mean()\n",
    "            median_log = np.median(local_series)\n",
    "            mean_linear = 10**mean_log\n",
    "            hpd_95 = az.hdi(local_series, 0.95)\n",
    "            lower_hpd_log_95 = hpd_95[0]\n",
    "            lower_hpd_linear_95 = math.exp(lower_hpd_log_95)\n",
    "            upper_hpd_log_95 = hpd_95[1]\n",
    "            upper_hpd_linear_95 = math.exp(upper_hpd_log_95)\n",
    "            hpd_50 = az.hdi(local_series, 0.50)\n",
    "            lower_hpd_log_50 = hpd_50[0]\n",
    "            lower_hpd_linear_50 = math.exp(lower_hpd_log_50)\n",
    "            upper_hpd_log_50 = hpd_50[1]\n",
    "            upper_hpd_linear_50 = math.exp(upper_hpd_log_50)\n",
    "            \n",
    "#             try:\n",
    "#                 next_local_series = input_df[\"SkylineNe\"+\".\"+ str(deme) +\".\" + str(next_interval)].astype('float').to_numpy()\n",
    "#                 diff_series = np.subtract(local_series, next_local_series)\n",
    "#                 #print(local_series)\n",
    "#                 #print(next_local_series)\n",
    "#                 #print(diff_series)\n",
    "#                 diff_mean_log = diff_series.mean()\n",
    "#                 diff_median_log = np.median(diff_series)\n",
    "#                 diff_hpd_95 = az.hdi(diff_series, 0.95)\n",
    "#                 diff_lower_hpd_log_95 = diff_hpd_95[0]\n",
    "#                 diff_lower_hpd_linear_95 = math.exp(diff_lower_hpd_log_95)\n",
    "#                 diff_upper_hpd_log_95 = diff_hpd_95[1]\n",
    "#                 diff_upper_hpd_linear_95 = math.exp(diff_upper_hpd_log_95)\n",
    "#                 diff_hpd_50 = az.hdi(diff_series, 0.50)\n",
    "#                 diff_lower_hpd_log_50 = diff_hpd_50[0]\n",
    "#                 diff_lower_hpd_linear_50 = math.exp(diff_lower_hpd_log_50)\n",
    "#                 diff_upper_hpd_log_50 = diff_hpd_50[1]\n",
    "#                 diff_upper_hpd_linear_50 = math.exp(diff_upper_hpd_log_50)\n",
    "#             except KeyError:\n",
    "#                 pass   \n",
    "            \n",
    "            try:\n",
    "                local_df = pd.DataFrame.from_dict({\"deme\":deme, \"interval\":interval, \"mean_Ne_log\":mean_log,\"mean_Ne_linear\":mean_linear, \n",
    "                                                   \"median_Ne_log\" : median_log, \n",
    "                                                   \"upper_hpd_log_95\":upper_hpd_log_95,\"lower_hpd_log_95\":[lower_hpd_log_95], \n",
    "                                                   \"upper_hpd_log_50\":upper_hpd_log_50,\"lower_hpd_log_50\":lower_hpd_log_50,\n",
    "                                                   \"upper_hpd_linear\":upper_hpd_linear_95,\"lower_hpd_linear\":lower_hpd_linear_95,\n",
    "                                                   \"upper_hpd_linear_50\":upper_hpd_linear_50, \"lower_hpd_linear_50\":lower_hpd_linear_50,\n",
    "                                                  })\n",
    "                new_df = new_df.append(local_df)\n",
    "                #print(new_df)\n",
    "            except:\n",
    "                pass\n",
    "            \n",
    "    return(new_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "795fc880",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ne_df = pd.DataFrame.from_dict(Ne_skyline)\n",
    "print(len(Ne_df))\n",
    "Ne_df\n",
    "\n",
    "burnin_percent = 0.5\n",
    "print(len(Ne_df))\n",
    "rows_to_remove = int(len(Ne_df)* burnin_percent)\n",
    "Ne_df = Ne_df.iloc[rows_to_remove:]\n",
    "\n",
    "print(len(Ne_df))\n",
    "Ne_df = Ne_df.reset_index()\n",
    "Ne_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74b28c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_ne_df = Ne_df.copy()\n",
    "# Select only numeric columns excluding the first two columns\n",
    "numeric_columns = exp_ne_df.columns[1:]  # Exclude the first two columns\n",
    "\n",
    "# Apply the exponential function only to numeric columns\n",
    "exp_ne_df[numeric_columns] = exp_ne_df[numeric_columns].apply(pd.to_numeric, errors='coerce').apply(np.exp)\n",
    "\n",
    "exp_ne_df\n",
    "\n",
    "# exp_ne_df= np.exp(Ne_df.iloc[:, 1:])\n",
    "# exp_ne_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d68582d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#calculating transmission rate\n",
    "def generate_summary_diff_df(input_df):\n",
    "    \n",
    "    \n",
    "    new_df = pd.DataFrame()\n",
    "   \n",
    "    for i in input_df.columns.tolist():\n",
    "        if \"Ne.\" in i:\n",
    "            deme = i.split(\".\")[1]\n",
    "           # print(deme)\n",
    "            #interval = \n",
    "            #print(interval)\n",
    "            if \"\\n\" in i.split(\".\")[2]:\n",
    "                interval = i.split(\".\")[2][0:2]\n",
    "            else:\n",
    "                interval = i.split(\".\")[2]\n",
    "            next_interval = int(interval)+2 #averaging over two weeks to reduce noise\n",
    "            local_series = input_df[i].astype('float').to_numpy()\n",
    "            #print(local_series)\n",
    "           \n",
    "            try:\n",
    "                new_df[\"Ne\"+\".\"+ str(deme) +\".diff.\" + str(interval)] = (29/3)*(np.log(input_df[i].astype(\"float\")) - np.log(input_df[\"SkylineNe\"+\".\"+ str(deme) +\".\" + str(next_interval)].astype('float')))\n",
    "            \n",
    "            \n",
    "            except KeyError:\n",
    "                pass \n",
    "            \n",
    "            \n",
    "    return(new_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "832defcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "ne_diff_summary = generate_summary_diff_df(exp_ne_df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "894d271f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ne_diff_summary.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e8cf4a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "uninfectious_rate = 365/4.5\n",
    "\n",
    "#taken from https://www.medrxiv.org/content/10.1101/2022.08.17.22278897v1.full.pdf\n",
    "incubation_period = 365/8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "386bdb64",
   "metadata": {},
   "outputs": [],
   "source": [
    "seir_growth_rate = ((ne_diff_summary*2 + uninfectious_rate + incubation_period)**2 - (incubation_period- uninfectious_rate)**2)/(4*incubation_period)\n",
    "\n",
    "\n",
    "\n",
    "seir_growth_rate.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9806a42b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ne_summary = generate_summary_df(Ne_df)\n",
    "ne_summary.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b270a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = ne_summary\n",
    "test['days'] = (test.interval.astype(int) -1 )*10.76\n",
    "test['date'] = dt.strptime(\"2024-05-01\",  \"%Y-%m-%d\") - test.days.map(timedelta)\n",
    "#test.date = test.date.astype(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2f91aaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = test.reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd151c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "rolling_ne =test.reset_index().groupby([\"deme\"])[\"mean_Ne_log\", \"upper_hpd_log_95\",\"lower_hpd_log_95\", \"upper_hpd_log_50\", \"lower_hpd_log_50\", \"lower_hpd_linear_50\", \"upper_hpd_linear_50\", \"lower_hpd_linear\", \"upper_hpd_linear\"  ].rolling(4, min_periods =1).mean().reset_index()\n",
    "smoothed_ne = pd.merge(rolling_ne, test, left_on = \"level_1\",right_index = True, how = \"right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fca27cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#domain = ['CentralEurope', 'NorthAmerica', 'SouthAmerica' ,\"SouthernEurope\", \"WesternEurope\"]\n",
    "range_ = ['#EEC060', '#2664A5', '#A76BB1', \"#EEA160\", \"#356D4C\"]\n",
    "line = alt.Chart(smoothed_ne).mark_area(interpolate='monotone', opacity = 0.8).encode(\n",
    "    alt.X('date:T', axis=alt.Axis(title=\"\", grid=False,tickCount = \"month\",  format=\"%B %Y\")),\n",
    "    alt.Y('lower_hpd_log_50_x',axis=alt.Axis(title=\"Log(Effective Population Size)\", grid=False)),#,scale=alt.Scale(domain=(0, 13))),\n",
    "    alt.Y2('upper_hpd_log_50_x' ),\n",
    "    color=alt.Color('deme_x:N', title = \"Region\") \n",
    ").properties(\n",
    "    width=850,\n",
    "    height=400\n",
    ")\n",
    "band = alt.Chart(smoothed_ne).mark_area(interpolate='monotone', opacity = 0.5).encode(\n",
    "    alt.X('date:T', axis=alt.Axis(title=\"\", grid=False, format=\"%B %Y\")),\n",
    "    alt.Y('lower_hpd_log_95_x',axis=alt.Axis(title=\"Effective Population Size\", grid=False)),#,scale=alt.Scale(domain=(0, 13))),\n",
    "    alt.Y2('upper_hpd_log_95_x' ),\n",
    "    color=alt.Color('deme_x:N') \n",
    ").properties(\n",
    "    width=850,\n",
    "    height=400\n",
    ")\n",
    "\n",
    "line  \n",
    "\n",
    "# band = alt.Chart(test).mark_area(\n",
    "#     opacity=0.3, interpolate='monotone'\n",
    "# ).encode(\n",
    "#     alt.X('date:T', axis=alt.Axis(title=\"\", grid=False)),\n",
    "#     alt.Y('lower_hpd_linear'),#axis=None),#, scale=alt.Scale(domain=(0, 13))),\n",
    "#     alt.Y2('upper_hpd_linear'),\n",
    "#     color=alt.Color('deme:N', title = \"Region\", legend=alt.Legend(orient = \"right\",offset = -135, labelFontSize = 16, titleFontSize = 16, symbolSize = 110))\n",
    "# ).properties(\n",
    "#     width=850,\n",
    "#     height=400\n",
    "# ).transform_filter(\n",
    "#     (datum.upper_hpd_linear < 20) &(datum.upper_hpd_linear_50 < 30)\n",
    "# )\n",
    "# band\n",
    "# ne_plot = line + band\n",
    "# (ne_plot).configure_axis(\n",
    "#     labelFontSize=20,\n",
    "#     titleFontSize=20\n",
    "# )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00cf0f41",
   "metadata": {},
   "outputs": [],
   "source": [
    "#domain = ['CentralEurope', 'NorthAmerica', 'SouthAmerica' ,\"SouthernEurope\", \"WesternEurope\"]\n",
    "range_ = ['#EEC060', '#2664A5', '#A76BB1', \"#EEA160\", \"#356D4C\"]\n",
    "line = alt.Chart(smoothed_ne).mark_area(interpolate='monotone', opacity = 0.9, clip = True).encode(\n",
    "    alt.X('date:T', axis=alt.Axis(title=\"\", grid=False,tickCount = \"month\",  format=\"%b %Y\"), scale=alt.Scale(domain=(\"2022-02-01\", \"2024-05-01\"))),\n",
    "    alt.Y('lower_hpd_linear_50_x',axis=alt.Axis(title=\"Effective Population Size\", grid=False)),#,scale=alt.Scale(domain=(0, 13))),\n",
    "    alt.Y2('upper_hpd_linear_50_x' ),\n",
    "    color=alt.Color('deme_x:N', title = \"Region\") \n",
    ").properties(\n",
    "    width=850,\n",
    "    height=300\n",
    ")\n",
    "band = alt.Chart(smoothed_ne).mark_area(interpolate='monotone', opacity = 0.2).encode(\n",
    "    alt.X('date:T', axis=alt.Axis(title=\"\", grid=False, format=\"%b %Y\")),\n",
    "    alt.Y('lower_hpd_linear_x',axis=alt.Axis(title=\"Effective Population Size\", grid=False)),#,scale=alt.Scale(domain=(0, 13))),\n",
    "    alt.Y2('upper_hpd_linear_x' ),\n",
    "    color=alt.Color('deme_x:N') \n",
    ").properties(\n",
    "    width=850,\n",
    "    height=300\n",
    ").transform_filter(\n",
    "     (datum.upper_hpd_linear_x < 150) \n",
    " )\n",
    "\n",
    "line  \n",
    "\n",
    "# band = alt.Chart(smoothed_ne).mark_area(\n",
    "#     opacity=0.3, interpolate='monotone'\n",
    "# ).encode(\n",
    "#     alt.X('date:T', axis=alt.Axis(title=\"\", grid=False)),\n",
    "#     alt.Y('lower_hpd_linear_x'),#axis=None),#, scale=alt.Scale(domain=(0, 13))),\n",
    "#     alt.Y2('upper_hpd_linear_x'),\n",
    "#     color=alt.Color('deme_x:N', title = \"Region\", legend=alt.Legend(orient = \"right\",offset = -135, labelFontSize = 16, titleFontSize = 16, symbolSize = 110))\n",
    "# ).properties(\n",
    "#     width=850,\n",
    "#     height=400\n",
    "# ).transform_filter(\n",
    "#     (datum.upper_hpd_linear < 20) &(datum.upper_hpd_linear_50 < 30)\n",
    "# )\n",
    "# band\n",
    "ne_plot = line \n",
    "(ne_plot).configure_axis(\n",
    "    labelFontSize= 15,\n",
    "    titleFontSize=20\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5457abf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "smoothed_ne.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943972cb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41177c28",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_in_forward_migration_rates_mascot(log_file_path):\n",
    "    \n",
    "    mig_rates_dict = {\"sample\":[]}\n",
    "    \n",
    "    with open(log_file_path, \"r\") as infile:\n",
    "        line_number = 0\n",
    "        for line in infile:\n",
    "            #print(line_number)\n",
    "            line_number += 1\n",
    "            if not line.startswith(\"#\"):  # log combiner will sometimes put the entire xml at the start of the log file\n",
    "                # use the first line to find the migration rate columns\n",
    "                \n",
    "            # use the first line to find the migration rate columns\n",
    "                if \"posterior\" in line:\n",
    "                    all_cols = line.split(\"\\t\")\n",
    "                    mig_column_indices = []   # list to store column indices\n",
    "                    mig_key = {}   # dictionary to store the column index to map to column name\n",
    "\n",
    "                    for i in range(len(all_cols)):\n",
    "                        col = all_cols[i]\n",
    "                        if \"f_migration\" in col:\n",
    "                            mig_column_indices.append(i)\n",
    "\n",
    "                    # make an empty dictionary to store Nes and generate dictionary to convert index to name\n",
    "                    for n in mig_column_indices:\n",
    "                        name = line.split(\"\\t\")[n]\n",
    "                        deme = name.split(\".\")[1]# the syntax here is \"NeLog.state01\" where 0 is deme and 1 is interval 1\n",
    "                        #interval = name.split(\".\")[2]\n",
    "                       \n",
    "                        mig_key[n] = name\n",
    "                        mig_rates_dict[name] = []\n",
    "\n",
    "\n",
    "                # read in actual parameter estimates and store in dictionary\n",
    "                else:\n",
    "                    sample = line.split(\"\\t\")[0]\n",
    "                    mig_rates_dict[\"sample\"].append(sample)\n",
    "\n",
    "                    for index in mig_column_indices:\n",
    "                        name = mig_key[index]\n",
    "                        mig_rates_dict[name].append(line.split(\"\\t\")[index])\n",
    "                    \n",
    "                \n",
    "                \n",
    "                \n",
    "    return(mig_rates_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01e35af6",
   "metadata": {},
   "outputs": [],
   "source": [
    "migration_rates_f = read_in_forward_migration_rates_mascot(log_file_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f67caf8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mig_df_f = pd.DataFrame.from_dict(migration_rates_f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a093876",
   "metadata": {},
   "outputs": [],
   "source": [
    "burnin_percent = 0.6\n",
    "print(len(mig_df_f))\n",
    "rows_to_remove = int(len(mig_df_f)* burnin_percent)\n",
    "mig_df_f = mig_df_f.iloc[rows_to_remove:]\n",
    "\n",
    "print(len(mig_df_f))\n",
    "mig_df_f = mig_df_f.reset_index()\n",
    "mig_df_f.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40696fab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a new dataframe that summarizes the 95% HPD estimate with mean for each deme and interval \n",
    "def generate_summary_df(input_df):\n",
    "    \n",
    "    \n",
    "    new_df = pd.DataFrame()\n",
    "\n",
    "    for i in input_df.columns.tolist():\n",
    "        if \"f_migrationRatesSkyline.\" in i:\n",
    "            deme = i.split(\".\")[1]\n",
    "            #print(deme)\n",
    "            #interval = \n",
    "            #print(interval\n",
    "            #print(interval)\n",
    "            #print(i)\n",
    "            #next_interval = int(interval)+1\n",
    "            local_series = input_df[i].astype('float').to_numpy()\n",
    "            #print(local_series)\n",
    "            mean_log = local_series.mean()\n",
    "            median_log = np.median(local_series)\n",
    "            hpd_95 = az.hdi(local_series, 0.95)\n",
    "            lower_hpd_log_95 = hpd_95[0]\n",
    "            upper_hpd_log_95 = hpd_95[1]\n",
    "            hpd_50 = az.hdi(local_series, 0.50)\n",
    "            lower_hpd_log_50 = hpd_50[0]\n",
    "            upper_hpd_log_50 = hpd_50[1]\n",
    "            \n",
    "#             try:\n",
    "#                 next_local_series = input_df[\"SkylineNe\"+\".\"+ str(deme) +\".\" + str(next_interval)].astype('float').to_numpy()\n",
    "#                 diff_series = np.subtract(local_series, next_local_series)\n",
    "#                 #print(local_series)\n",
    "#                 #print(next_local_series)\n",
    "#                 #print(diff_series)\n",
    "#                 diff_mean_log = diff_series.mean()\n",
    "#                 diff_median_log = np.median(diff_series)\n",
    "#                 diff_hpd_95 = az.hdi(diff_series, 0.95)\n",
    "#                 diff_lower_hpd_log_95 = diff_hpd_95[0]\n",
    "#                 diff_lower_hpd_linear_95 = math.exp(diff_lower_hpd_log_95)\n",
    "#                 diff_upper_hpd_log_95 = diff_hpd_95[1]\n",
    "#                 diff_upper_hpd_linear_95 = math.exp(diff_upper_hpd_log_95)\n",
    "#                 diff_hpd_50 = az.hdi(diff_series, 0.50)\n",
    "#                 diff_lower_hpd_log_50 = diff_hpd_50[0]\n",
    "#                 diff_lower_hpd_linear_50 = math.exp(diff_lower_hpd_log_50)\n",
    "#                 diff_upper_hpd_log_50 = diff_hpd_50[1]\n",
    "#                 diff_upper_hpd_linear_50 = math.exp(diff_upper_hpd_log_50)\n",
    "#             except KeyError:\n",
    "#                 pass   \n",
    "            \n",
    "            try:\n",
    "                local_df = pd.DataFrame.from_dict({\"deme\":deme, \"mean_mig\":mean_log, \n",
    "                                                   \"median_mig\" : median_log, \n",
    "                                                   \"upper_hpd_log_95\":upper_hpd_log_95,\"lower_hpd_log_95\":[lower_hpd_log_95], \n",
    "                                                   \"upper_hpd_log_50\":upper_hpd_log_50,\"lower_hpd_log_50\":lower_hpd_log_50,\n",
    "                                                   \n",
    "                                                  })\n",
    "                new_df = new_df.append(local_df)\n",
    "                #print(new_df)\n",
    "            except:\n",
    "                pass\n",
    "            \n",
    "    return(new_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ff0f61c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mig_summary = generate_summary_df(mig_df_f).reset_index()\n",
    "mig_summary.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34d222a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "mig_summary[\"origin\"] = mig_summary.deme.str.split(\"_\").str[0].tolist()\n",
    "mig_summary[\"destination\"] = mig_summary.deme.str.split(\"_\").str[2].tolist()\n",
    "mig_summary.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62c99dc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix = mig_summary.pivot(index='origin', columns='destination', values='mean_mig')\n",
    "matrix = matrix.fillna(0)\n",
    "matrix.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e32826",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_backwards_mig_rates(mig_df_f, targetName):    \n",
    "    mig_rates = {}\n",
    "    target = targetName\n",
    "    target_columns = []\n",
    "    ne_interval= 0\n",
    "    mig_list = []\n",
    "\n",
    "    for name in mig_df_f.columns:\n",
    "        try:\n",
    "            if target in name.split(\"_\")[3]:\n",
    "                target_columns.append(name)\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "    for name in exp_ne_df.columns:\n",
    "        if target in name:\n",
    "            ne_interval+=1\n",
    "\n",
    "    for interval in range(1, ne_interval,1):\n",
    "        mig_rates[str(interval)+ \".\" + \"bmr\" + \".\" +  str(target)] = []\n",
    "        for index_1, row_1 in mig_df_f.iterrows():\n",
    "            mig_list = []\n",
    "            for direction in target_columns:\n",
    "                #print(direction)\n",
    "                other_loc = direction.split(\".\")[1].split(\"_\")[0]\n",
    "                #print(other_loc)\n",
    "                mig_rate_b = pd.to_numeric(row_1[direction]) * ((pd.to_numeric(exp_ne_df.loc[index_1,\"SkylineNe.\"+ str(other_loc) + \".\" + str(interval)]))/(pd.to_numeric(exp_ne_df.loc[index_1, \"SkylineNe.\"+ str(target) + \".\" + str(interval)])))\n",
    "                mig_list.append(mig_rate_b)\n",
    "            combined_mig = np.sum(np.array(mig_list))\n",
    "            mig_rates[str(interval)+ \".\" + \"bmr\" + \".\" + str(target)].append(combined_mig)\n",
    "\n",
    "    return(mig_rates) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40f430d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### if the forward migration rate is \n",
    "#### from a -> b then the backwards migration rate is from b - > a and interpreted as the rate\n",
    "#### at which a lineage in b has originated from a\n",
    "#### so our target is not necessarily location a but rather location b.\n",
    "\n",
    "unique_column= []\n",
    "unique_target = []\n",
    "for name in mig_df_f.columns:\n",
    "    try:\n",
    "        if name.split(\"_\")[3] not in unique_target:\n",
    "            #print(name.split(\"_\")[3])\n",
    "            unique_column.append(name)\n",
    "            unique_target.append(name.split(\"_\")[3])\n",
    "    except:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "555de8d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "mig_rates = []\n",
    "for target in unique_target:\n",
    "        mig_rates.append(calc_backwards_mig_rates(mig_df_f, target))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9887ba60",
   "metadata": {},
   "outputs": [],
   "source": [
    "mr_b_df = pd.DataFrame()\n",
    "for x in mig_rates:\n",
    "    x_df = pd.DataFrame(x)\n",
    "    mr_b_df = pd.concat([mr_b_df, x_df], axis=1)\n",
    "mr_b_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9bf105e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## percentage of cases due to introductions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1208155",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_percent_intro_df(input_df):\n",
    "    \n",
    "    temp_df = pd.DataFrame()\n",
    "    new_df = pd.DataFrame()\n",
    "   \n",
    "    for i in input_df.columns.tolist():\n",
    "        \n",
    "        interval = i.split(\".\")[0]\n",
    "        deme = i.split(\".\")[2]\n",
    "\n",
    "        try:\n",
    "            temp_df[\"total.\"+ str(interval)] = seir_growth_rate[\"Ne.\" + str(deme) + \".diff.\" + str(interval)].astype(\"float\") +  input_df[i].astype(\"float\")\n",
    "            new_df[\"intro.percent\"+\".\" + str(deme) + \".\" + str(interval)] = input_df[i].astype(\"float\").div(temp_df[\"total.\"+ str(interval)], axis = 0) \n",
    "\n",
    "        except KeyError: #this was added because not all regions have equal time periods for their epidemics and it was throwing an error everytime it had to switch deme\n",
    "            pass \n",
    "                  \n",
    "    return(new_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cc11e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "percent_df = generate_percent_intro_df(mr_b_df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea5abd00",
   "metadata": {},
   "outputs": [],
   "source": [
    "percent_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fb0a088",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a new dataframe that summarizes the 95% HPD estimate with mean for each deme and interval \n",
    "def generate_summary_df(input_df):\n",
    "    \n",
    "    \n",
    "    new_df = pd.DataFrame()\n",
    "\n",
    "    for i in input_df.columns.tolist():\n",
    "        if \"percent\" in i:\n",
    "            deme = i.split(\".\")[2]\n",
    "            interval = i.split(\".\")[3]\n",
    "            local_series = input_df[i].astype('float').to_numpy()\n",
    "            mean_percent = local_series.mean()\n",
    "            hpd_95 = az.hdi(local_series, 0.95)\n",
    "            lower_hpd_log_95 = hpd_95[0]\n",
    "            upper_hpd_log_95 = hpd_95[1]\n",
    "            hpd_50 = az.hdi(local_series, 0.50)\n",
    "            lower_hpd_log_50 = hpd_50[0]\n",
    "            upper_hpd_log_50 = hpd_50[1]\n",
    "            \n",
    "\n",
    "            \n",
    "            \n",
    "            try:\n",
    "                local_df = pd.DataFrame.from_dict({\"deme\":deme, \"interval\":interval, \"mean_percent\":mean_percent, \n",
    "                                                   \"upper_hpd_log_95\":upper_hpd_log_95,\"lower_hpd_log_95\":[lower_hpd_log_95], \n",
    "                                                   \"upper_hpd_log_50\":upper_hpd_log_50,\"lower_hpd_log_50\":lower_hpd_log_50})\n",
    "                new_df = new_df.append(local_df)\n",
    "            except:\n",
    "                pass\n",
    "            \n",
    "    return(new_df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c380f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_percent_df = generate_summary_df(percent_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "224e62c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_percent_df = summary_percent_df.reset_index()\n",
    "summary_percent_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0cf86f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#making sure that any numbers >1 are excluded as outliers from stochastics noise\n",
    "percent_df =pd.DataFrame(np.where(percent_df <1, percent_df, 1), columns=percent_df.columns )\n",
    "percent_df =pd.DataFrame(np.where(percent_df >0, percent_df, 0), columns=percent_df.columns )\n",
    "percent_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1f8c20e",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_percent_df['days'] = (summary_percent_df.interval.astype(int)-1)*10.76 #the 1.5 adjustment is made due to the fact that we take the difference of Nes over three weeks \n",
    "summary_percent_df['date'] = dt.strptime(\"2024-05-01\",  \"%Y-%m-%d\") - summary_percent_df.days.map(timedelta)\n",
    "\n",
    "summary_percent_df.date = summary_percent_df.date.astype(str)\n",
    "summary_percent_df.deme = summary_percent_df.deme.replace({\"CentralEurope\": \"Central Europe\", \"NorthAmerica\": \"North America\", \"SouthAmerica\":\"South America\", \"SouthernEurope\":\"Southern Europe\", \"WesternEurope\": \"Western Europe\" })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c611a913",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#smoothing via 2 week rolling average \n",
    "rolling_final_percent =summary_percent_df.groupby([\"deme\"])[\"mean_percent\", \"upper_hpd_log_95\",\"lower_hpd_log_95\", \"upper_hpd_log_50\", \"lower_hpd_log_50\" ].rolling(2, min_periods =1).mean().reset_index()\n",
    "#rolling_final_percent =final_north_df.groupby([\"deme\"])[\"mean_percent\", \"upper_hpd_log_95\",\"lower_hpd_log_95\", \"upper_hpd_log_50\", \"lower_hpd_log_50\" ].rolling(2, min_periods =1).mean().reset_index()\n",
    "smoothed_percent = pd.merge(rolling_final_percent, summary_percent_df, left_on = \"level_1\",right_index = True, how = \"right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e79c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "line1 = alt.Chart(smoothed_percent).mark_area(interpolate='monotone', opacity = 0.9, clip = True).encode(\n",
    "    alt.X('date:T', axis=alt.Axis(title=\"\",tickCount = \"month\",grid=False,  format=\"%b %Y\"), scale=alt.Scale(domain=(\"2022-02-01\", \"2024-05-01\"))),\n",
    "    alt.Y('lower_hpd_log_50_x',title = \"Percent of new cases due to introductions\", axis=alt.Axis( grid=False, format='%', values = list(np.arange (0, 1, 0.1))), scale=alt.Scale(domain=(0.0, 1))),\n",
    "    alt.Y2('upper_hpd_log_50_x' ), \n",
    "    alt.Color(\"deme_x\",title = \"Region\")\n",
    ").properties(\n",
    "    width=850,\n",
    "    height=300\n",
    ").transform_filter(\n",
    "    (datum.lower_hpd_log_50_x >0) & (datum.upper_hpd_log_50_x < 1) & (datum.upper_hpd_log_95_x < 1)\n",
    ")\n",
    "line1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ca92d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "band1 = alt.Chart(smoothed_percent).mark_area(\n",
    "    opacity=0.3, interpolate='monotone', clip = True\n",
    ").encode(\n",
    "    alt.X('date:T', axis=alt.Axis(title=\"\", grid=False)),\n",
    "    alt.Y('lower_hpd_log_95_x',title = \"Percent of new cases due to introductions\", axis=alt.Axis( grid=False)),\n",
    "    alt.Y2('upper_hpd_log_95_x'),\n",
    "    alt.Color(\"deme_x\")\n",
    ").properties(\n",
    "    width=850,\n",
    "    height=300\n",
    ").transform_filter(\n",
    "    (datum.lower_hpd_log_95_x >0) & (datum.upper_hpd_log_95_x < 0.9)\n",
    ")\n",
    "percent_plot= band1 + line1\n",
    "(percent_plot).configure_axis(\n",
    "    labelFontSize=20,\n",
    "    titleFontSize=20\n",
    ")\n",
    "\n",
    "band1 +line1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bc85be5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}