joint-annotation-embedding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bcab4487-1c91-4ad0-a4fb-c41e632411d6",
   "metadata": {},
   "source": [
    "# Multiple Annotation Embedding\n",
    "\n",
    "In this notebook we are showing how the annotation embedding avoids batch effects in pairwise comparisons of health and tumor samples from [Mair et al., 2022, Nature](https://www.nature.com/articles/s41586-022-04718-w)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a6e09287-c0ed-48b8-8c87-8d0ecc999988",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jscatter\n",
    "import math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import re\n",
    "import transformation\n",
    "import colors\n",
    "\n",
    "from glob import glob\n",
    "from openTSNE.sklearn import TSNE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3f66a2cb-86aa-4c30-8470-dd7c18fdafc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name_tissue = 'TISSUE_138'\n",
    "dataset_name_tumor = 'TUMOR_007'\n",
    "\n",
    "dataset_tissue = glob(f'data/mair-2022/{dataset_name_tissue}*')[0]\n",
    "dataset_tumor = glob(f'data/mair-2022/{dataset_name_tumor}*')[0]\n",
    "\n",
    "sample_tissue = dataset_tissue[16:-25]\n",
    "sample_tumor = dataset_tumor[16:-25]\n",
    "\n",
    "df_tissue = pd.read_parquet(dataset_tissue)\n",
    "df_tumor = pd.read_parquet(dataset_tumor)\n",
    "\n",
    "if len(df_tissue) < len(df_tumor):\n",
    "    df_tumor = df_tumor.sample(n=len(df_tissue), random_state=42)\n",
    "else:\n",
    "    df_tissue = df_tissue.sample(n=len(df_tumor), random_state=42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d40aec9-5a59-424d-8896-3dbfab9cd73d",
   "metadata": {},
   "source": [
    "**Joint Embedding:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62fdaa16-2feb-456b-92bb-17d7b7fb27b9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "reload(transformation)\n",
    "\n",
    "df_ann_embed_tsne = transformation.transform_embed(\n",
    "    [df_tissue, df_tumor],\n",
    "    embeddor=TSNE,\n",
    "    embeddor_random_state = 42,\n",
    "    save_as=f'{dataset_name_tissue}_{dataset_name_tumor}_tsne_ann',\n",
    ")\n",
    "\n",
    "df_raw_embed_tsne = transformation.transform_embed(\n",
    "    [df_tissue, df_tumor],\n",
    "    embeddor=TSNE,\n",
    "    embeddor_random_state=42,\n",
    "    embed_raw=True,\n",
    "    save_as=f'{dataset_name_tissue}_{dataset_name_tumor}_tsne_raw',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54fa8de8-896c-4710-bfec-45375ba8caf8",
   "metadata": {},
   "source": [
    "**Visualize Embeddings:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ece876eb-3354-4299-8ec0-a49cc3c1e9db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0fcef4152a984080a67b1369d6e16d6b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "GridBox(children=(HBox(children=(VBox(children=(Button(button_style='primary', icon='arrows', layout=Layout(wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Uncomment the line below to load previously embedded data\n",
    "df_ann_embed_tsne = pd.read_parquet(f'data/{dataset_name_tissue}_{dataset_name_tumor}_tsne_ann.pq')\n",
    "df_raw_embed_tsne = pd.read_parquet(f'data/{dataset_name_tissue}_{dataset_name_tumor}_tsne_raw.pq')\n",
    "\n",
    "cell_type_color_map = colors.get_cmap(len(df_tissue.faustLabels.unique()), mode='dark')\n",
    "sample_color_map = dict(\n",
    "    sample_tissue='#0072B2', # blue\n",
    "    sample_tumor='#E69F00', # orange\n",
    ")\n",
    "\n",
    "base_view_config = dict(x='x', y='y', opacity_unselected=0.1, background_color='black', axes=False)\n",
    "cell_type_view_config = dict(color_by='cellType', color_map=cell_type_color_map, **base_view_config)\n",
    "sample_view_config = dict(color_by='sampleOfOrigin', color_map=sample_color_map, **base_view_config)\n",
    "marker_view_config = dict(color_by='ICOS_Windsorized', color_map='viridis', **base_view_config)\n",
    "\n",
    "compose_config = dict(sync_selection=True, sync_hover=True, row_height=400, rows=2)\n",
    "\n",
    "plot_ann_embed_tsne_ct = jscatter.Scatter(data=df_ann_embed_tsne, **cell_type_view_config)\n",
    "plot_ann_embed_tsne_sp = jscatter.Scatter(data=df_ann_embed_tsne, **sample_view_config)\n",
    "plot_ann_embed_tsne_mx = jscatter.Scatter(data=df_ann_embed_tsne, **marker_view_config)\n",
    "\n",
    "plot_raw_embed_tsne_ct = jscatter.Scatter(data=df_raw_embed_tsne, **cell_type_view_config)\n",
    "plot_raw_embed_tsne_sp = jscatter.Scatter(data=df_raw_embed_tsne, **sample_view_config)\n",
    "plot_raw_embed_tsne_mx = jscatter.Scatter(data=df_raw_embed_tsne, **marker_view_config)\n",
    "\n",
    "jscatter.compose(\n",
    "    [\n",
    "        plot_ann_embed_tsne_ct,\n",
    "        plot_ann_embed_tsne_sp,\n",
    "        plot_ann_embed_tsne_mx,\n",
    "        plot_raw_embed_tsne_ct,\n",
    "        plot_raw_embed_tsne_sp,\n",
    "        plot_raw_embed_tsne_mx,\n",
    "    ],\n",
    "    **compose_config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "073642f7-b61c-474e-9e5a-66f584382486",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<jscatter.jscatter.Scatter at 0x1bc3aed00>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plot_ann_embed_tsne_ct.selection(df_ann_embed_tsne.query(\n",
    "    'cellType == \"CD4+CD8-CD3+CD45RA-CD27+CD19-CD103-CD28+CD69-PD1+HLADR+GranzymeB-CD25+ICOS+TCRgd-CD38+CD127-Tim3+\"'\n",
    ").index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ab4cd9-5320-4b26-8b73-9efc0e22f600",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}