DeconvolutionPrediagnostics.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d685bf5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
    "plotwidth = 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2198e6b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lollipop as ll\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from matplotlib.colors import ListedColormap\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bf5e25c",
   "metadata": {},
   "source": [
    "# Globals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c61315a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source of inspiration from covariatns, see:\n",
    "# https://github.com/hodcroftlab/covariants/blob/master/web/data/clusters.json\n",
    "#\n",
    "# Keep in sync with covspectrum, see:\n",
    "# https://github.com/GenSpectrum/cov-spectrum-website/blob/develop/src/models/wasteWater/constants.tsv\n",
    "color_map = {\n",
    "    \"B.1.1.7\": \"#D16666\",\n",
    "    \"B.1.351\": \"#FF6665\",\n",
    "    \"P.1\": \"#FFB3B3\",\n",
    "    \"B.1.617.1\": \"#66C265\",\n",
    "    \"B.1.617.2\": \"#66A366\",\n",
    "    \"BA.1\": \"#A366A3\",\n",
    "    \"BA.2\": \"#cfafcf\",\n",
    "    \"BA.4\": \"#8a66ff\",\n",
    "    \"BA.5\": \"#585eff\",\n",
    "    \"BA.2.12.1\": \"#0400e0\",\n",
    "    \"BA.2.75\": \"#008fe0\",\n",
    "    \"BA.2.75.2\": \"#208fe0\",  # improv\n",
    "    \"BQ.1.1\": \"#8fe000\",  # improv\n",
    "    \"XBB\": \"#dd6bff\",\n",
    "    \"undetermined\": \"#969696\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94081008",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Overwrite globals set by WwDec.main:\n",
    "# temporary, globals\n",
    "\n",
    "# zst needs python's Zstandard\n",
    "tally_data = \"./work-vp-test/variants/tallymut.tsv.zst\"  # zst needs python's Zstandard # \"./tallymut_line.tsv\"\n",
    "out_dir = \"./out\"\n",
    "\n",
    "import yaml\n",
    "\n",
    "# load variants configuations\n",
    "with open(\"work-vp-test/variant_config.yaml\", \"r\") as file:\n",
    "    conf_yaml = yaml.load(file, Loader=yaml.SafeLoader)\n",
    "variants_list = conf_yaml.get(\"variants_list\", None)\n",
    "variants_pangolin = conf_yaml[\"variants_pangolin\"]\n",
    "variants_not_reported = conf_yaml.get(\"variants_not_reported\", [])\n",
    "start_date = conf_yaml.get(\"start_date\", None)\n",
    "end_date = conf_yaml.get(\n",
    "    \"end_date\"\n",
    ")  # optionnal, usually absent in ongoing surveillance, and present in articles with subset of historical data\n",
    "\n",
    "to_drop = conf_yaml.get(\"to_drop\", [])\n",
    "locations_list = conf_yaml.get(\"locations_list\", None)\n",
    "\n",
    "# var dates\n",
    "with open(\"work-vp-test/var_dates.yaml\", \"r\") as file:\n",
    "    conf_yaml.update(yaml.load(file, Loader=yaml.SafeLoader))\n",
    "\n",
    "if not variants_list:\n",
    "    # build list of all variants from var_dates (if we did lack one)\n",
    "    conf_yaml[\"variants_list\"] = variants_list = sorted(\n",
    "        list(set([var for lst in conf_yaml[\"var_dates\"].values() for var in lst]))\n",
    "    )\n",
    "\n",
    "\n",
    "# display the current config\n",
    "conf_yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb8868cd",
   "metadata": {},
   "source": [
    "# Load and preprocess data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b412ffce",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_tally = pd.read_csv(\n",
    "    tally_data, sep=\"\\t\", parse_dates=[\"date\"], dtype={\"location_code\": \"str\"}\n",
    ")  # .drop(columns=['proto'])\n",
    "df_tally.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8a6b4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not locations_list:\n",
    "    # remember to remove empty cells: nan or empty cells\n",
    "    conf_yaml[\"locations_list\"] = locations_list = list(\n",
    "        set(df_tally[\"location\"].unique()) - {\"\", np.nan}\n",
    "    )\n",
    "    display(locations_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7bef731",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    set(df_tally.columns)\n",
    "    - set(variants_pangolin.keys())\n",
    "    - {\n",
    "        \"base\",\n",
    "        \"batch\",\n",
    "        \"cov\",\n",
    "        \"date\",\n",
    "        \"frac\",\n",
    "        \"gene\",\n",
    "        \"location_code\",\n",
    "        \"location\",\n",
    "        \"pos\",\n",
    "        \"proto\",\n",
    "        \"sample\",\n",
    "        \"var\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ab53021",
   "metadata": {},
   "outputs": [],
   "source": [
    "preproc = ll.DataPreprocesser(df_tally)\n",
    "preproc = preproc.general_preprocess(\n",
    "    variants_list=variants_list,\n",
    "    variants_pangolin=variants_pangolin,\n",
    "    variants_not_reported=variants_not_reported,\n",
    "    to_drop=[\"subset\"],\n",
    "    start_date=start_date,\n",
    "    remove_deletions=True,\n",
    ")\n",
    "t_df_tally = preproc.df_tally\n",
    "# split into v41 and not v41, filter mutations and join\n",
    "df_tally_vcurrent = preproc.df_tally[preproc.df_tally.proto == \"v532\"]\n",
    "print(df_tally_vcurrent.shape)\n",
    "preproc.df_tally = preproc.df_tally[preproc.df_tally.proto != \"v532\"]\n",
    "preproc = preproc.filter_mutations()\n",
    "print(preproc.df_tally.shape)\n",
    "\n",
    "preproc.df_tally = pd.concat([preproc.df_tally, df_tally_vcurrent])\n",
    "print(preproc.df_tally.shape)\n",
    "# preproc.df_tally['']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3690b89",
   "metadata": {},
   "source": [
    "# Look at design of mutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a90edf29",
   "metadata": {},
   "outputs": [],
   "source": [
    "des_matrix = (\n",
    "    preproc.df_tally[variants_list + [\"undetermined\", \"mutations\"]]\n",
    "    .drop_duplicates(\"mutations\")\n",
    "    .set_index(\"mutations\")\n",
    ")\n",
    "des_matrix_mut = des_matrix[~des_matrix.index.str.startswith(\"-\")]\n",
    "des_matrix_wt = des_matrix[des_matrix.index.str.startswith(\"-\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f5c637d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(ncols=1, nrows=2, figsize=(plotwidth * 0.5, plotwidth / 8))\n",
    "cmap_binary = ListedColormap([\"white\", \"red\"])\n",
    "# sns.heatmap(des_matrix.T, square=True, cmap=cmap_binary, cbar=False)\n",
    "\n",
    "sns.heatmap(des_matrix_mut.T, square=True, cmap=cmap_binary, cbar=False, ax=axes[0])\n",
    "sns.heatmap(des_matrix_wt.T, square=True, cmap=cmap_binary, cbar=False, ax=axes[1])\n",
    "\n",
    "# axes[0].tick_params(labelsize=9)\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f536544",
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.linalg.cond(des_matrix_mut.drop('undetermined', axis=1))\n",
    "print(\"condition number = {:.2f}\".format(np.linalg.cond(des_matrix)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f2b01d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(22, 7))\n",
    "\n",
    "common_mut = des_matrix_mut.T.dot(des_matrix_mut)\n",
    "sns.heatmap(common_mut, square=True, cmap=\"viridis\", annot=common_mut, ax=axes[0])\n",
    "axes[0].set_title(\"common mutations\")\n",
    "\n",
    "corr_mut = (des_matrix_mut).corr()\n",
    "sns.heatmap(\n",
    "    corr_mut, square=True, cmap=\"viridis\", annot=corr_mut, ax=axes[1], fmt=\".1g\"\n",
    ")\n",
    "axes[1].set_title(\"correlation\")\n",
    "\n",
    "from sklearn.metrics.pairwise import pairwise_distances\n",
    "\n",
    "jac_sim = 1 - pairwise_distances(des_matrix_mut.T, metric=\"hamming\")\n",
    "jac_sim = pd.DataFrame(\n",
    "    jac_sim, index=des_matrix_mut.columns, columns=des_matrix_mut.columns\n",
    ")\n",
    "sns.heatmap(jac_sim, square=True, cmap=\"viridis\", annot=jac_sim, ax=axes[2])\n",
    "axes[2].set_title(\"jaccard similarity ((A∩B)/(A∪B))\")\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c7c75c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "locations_1 = locations_list\n",
    "# locations_1 = ['Lugano (TI)',\n",
    "#               'Zürich (ZH)',\n",
    "#               'Genève (GE)',\n",
    "#               'Chur (GR)',\n",
    "#               'Altenrhein (SG)',\n",
    "#               'Laupen (BE)',\n",
    "#               'Lausanne (Vidy)',\n",
    "#               'Sion (VS)',\n",
    "#               'Porrentruy (JU)',\n",
    "#               'Basel (catchment area ARA Basel)']\n",
    "# print(set(locations_1)-set(cities_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c84b7348",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_conds_df = []\n",
    "for proto in preproc.df_tally.proto.unique():\n",
    "    for location in locations_1:\n",
    "        t_df_tally_zh = preproc.df_tally[preproc.df_tally.location == location]\n",
    "        t_df_tally_zh = t_df_tally_zh[\n",
    "            (t_df_tally_zh.proto == proto) & (t_df_tally_zh[\"cov\"] >= 5)\n",
    "        ]\n",
    "\n",
    "        conds = []\n",
    "        for date in t_df_tally_zh.date.unique():\n",
    "            des_matrix = (\n",
    "                t_df_tally_zh[(t_df_tally_zh.date == date)][\n",
    "                    variants_list + [\"undetermined\", \"mutations\"]\n",
    "                ]\n",
    "                .drop_duplicates(\"mutations\")\n",
    "                .set_index(\"mutations\")\n",
    "            )\n",
    "            des_matrix_mut = des_matrix[~des_matrix.index.str.startswith(\"-\")]\n",
    "            des_matrix_wt = des_matrix[des_matrix.index.str.startswith(\"-\")]\n",
    "\n",
    "            #             print((location, date))\n",
    "\n",
    "            jac_sim = 1 - pairwise_distances(des_matrix_mut.T, metric=\"hamming\")\n",
    "            jac_sim = pd.DataFrame(jac_sim)\n",
    "            jac_arr = jac_sim.values\n",
    "            np.fill_diagonal(jac_arr, np.nan)\n",
    "            maxjac = np.nanmax(jac_arr)\n",
    "\n",
    "            corr_mut = (des_matrix_mut).corr()\n",
    "            corr_arr = corr_mut.values\n",
    "            np.fill_diagonal(corr_arr, np.nan)\n",
    "            maxcorr = np.nanmax(corr_arr)\n",
    "\n",
    "            conds.append(\n",
    "                {\n",
    "                    \"n_mut\": des_matrix_mut.shape[0],\n",
    "                    \"cond_number\": np.linalg.cond(des_matrix),\n",
    "                    \"cond_number_omicron\": np.linalg.cond(\n",
    "                        des_matrix[[\"BA.1\", \"BA.2\", \"BA.4\", \"BA.5\"]]\n",
    "                    ),\n",
    "                    \"max_jac\": maxjac,\n",
    "                    \"max_corr\": maxcorr,\n",
    "                    \"location\": location,\n",
    "                }\n",
    "            )\n",
    "\n",
    "        conds_df = pd.DataFrame(conds, index=t_df_tally_zh.date.unique())\n",
    "        conds_df[\"proto\"] = proto\n",
    "        all_conds_df.append(conds_df)\n",
    "        # print(np.linalg.cond(des_matrix_mut.drop('undetermined', axis=1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e6d4ec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_conds_df_conc = pd.concat(all_conds_df)\n",
    "all_conds_df_conc = all_conds_df_conc.reset_index()\n",
    "all_conds_df_conc.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "735e3791",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(5, 2, figsize=(14, 12))\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i, location in enumerate(all_conds_df_conc.location.unique()):\n",
    "    tmp_df = all_conds_df_conc[all_conds_df_conc.location == location]\n",
    "    h = sns.lineplot(\n",
    "        x=tmp_df[\"index\"], y=tmp_df[\"max_jac\"], hue=tmp_df[\"proto\"], ax=axes[i]\n",
    "    )\n",
    "    # h.set_ylim(top=20)\n",
    "    h.set_xlim(left=np.datetime64(\"2021-12-01\"))\n",
    "    axes[i].set_title(location)\n",
    "    axes[i].set_ylabel(\"max jaccard sim\")\n",
    "    axes[i].set_xlabel(\"\")\n",
    "    #     axes[i].set_xticks(rotation = 45) # Rotates X-Axis Ticks by 45-degrees\n",
    "\n",
    "    for tick in axes[i].get_xticklabels():\n",
    "        tick.set_rotation(45)\n",
    "\n",
    "fig.tight_layout()  # Or equivalently,  \"plt.tight_layout()\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed52348a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "prediag",
   "language": "python",
   "name": "prediag"
  },
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}