03_prepare_tasic_heterogeneous.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Tasic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle\n",
    "import scanpy as sc\n",
    "import anndata\n",
    "import openTSNE\n",
    "from sklearn.decomposition import PCA\n",
    "from openTSNE.initialization import rescale as rescale_pca\n",
    "sc.settings.verbosity = 2\n",
    "from tqdm import tqdm\n",
    "from readcount_tools import get_tag, compute_residuals, select_hvgs, scanpy_preproc_baseline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Quick load Tasic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/anndata/_core/anndata.py:121: ImplicitModificationWarning: Transforming to str index.\n",
      "  warnings.warn(\"Transforming to str index.\", ImplicitModificationWarning)\n",
      "/usr/local/lib/python3.8/dist-packages/anndata/_core/anndata.py:1830: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"var\")\n"
     ]
    }
   ],
   "source": [
    "basepath = 'data/tasic/'\n",
    "adata = anndata.read_h5ad(f'{basepath}adata.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(23822, 42776)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "tcmalloc: large alloc 1729986560 bytes == 0xb0518000 @ \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "filtered out 4266 genes that are detected in less than 5 cells\n",
      "(23822, 38510)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/anndata/_core/anndata.py:1830: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"var\")\n"
     ]
    }
   ],
   "source": [
    "print(adata.shape)\n",
    "sc.pp.filter_genes(adata,min_cells=5)\n",
    "print(adata.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compound Pearson residuals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5 [00:00<?, ?it/s]tcmalloc: large alloc 3669549056 bytes == 0x13e360000 @ \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha1.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "tcmalloc: large alloc 3669549056 bytes == 0x22a7ca000 @ \n",
      "tcmalloc: large alloc 3669549056 bytes == 0x305358000 @ \n",
      "tcmalloc: large alloc 3669549056 bytes == 0x3e06e6000 @ \n",
      "tcmalloc: large alloc 3669549056 bytes == 0x4bb274000 @ \n",
      "tcmalloc: large alloc 3669549056 bytes == 0x4bb274000 @ \n",
      "tcmalloc: large alloc 3669549056 bytes == 0x13e360000 @ \n",
      " 20%|██        | 1/5 [00:18<01:15, 18.85s/it]tcmalloc: large alloc 3669549056 bytes == 0x13e360000 @ \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha10.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "tcmalloc: large alloc 3669549056 bytes == 0x22a7ca000 @ \n",
      "tcmalloc: large alloc 3669549056 bytes == 0x3e06e6000 @ \n",
      " 40%|████      | 2/5 [00:36<00:55, 18.36s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha100.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 3/5 [00:54<00:36, 18.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha1000.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 4/5 [01:11<00:17, 17.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha50.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [01:27<00:00, 17.49s/it]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[None, None, None, None, None]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alphas = np.append(np.logspace(0,3,num=4),50)\n",
    "[compute_residuals(adata,alpha=alpha,theta=100) for alpha in tqdm(alphas)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#recompute means because we did some subsetting\n",
    "adata.var['gene_mean'] = np.array(np.mean(adata.X,axis=0)).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [00:00<00:00, 135.12it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[None, None, None, None, None]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[select_hvgs(adata,alpha=alpha,theta=100,n_hvgs=3000,clipping=True) for alpha in tqdm(alphas)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run PCA with 1000 PCs for all alphas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Layers with keys: pr_theta100_alpha1.0, pr_theta100_alpha10.0, pr_theta100_alpha100.0, pr_theta100_alpha1000.0, pr_theta100_alpha50.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5 [00:00<?, ?it/s]/usr/local/lib/python3.8/dist-packages/anndata/_core/anndata.py:1830: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"var\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha1.0_afterHVG\n",
      "pr_theta100_alpha1.0_afterHVG\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 1/5 [00:16<01:07, 16.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha10.0_afterHVG\n",
      "pr_theta100_alpha10.0_afterHVG\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 2/5 [00:33<00:50, 16.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha100.0_afterHVG\n",
      "pr_theta100_alpha100.0_afterHVG\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 3/5 [00:49<00:32, 16.25s/it]/usr/local/lib/python3.8/dist-packages/anndata/_core/anndata.py:1830: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"var\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha1000.0_afterHVG\n",
      "pr_theta100_alpha1000.0_afterHVG\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 4/5 [01:05<00:16, 16.10s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pr_theta100_alpha50.0_afterHVG\n",
      "pr_theta100_alpha50.0_afterHVG\n",
      "running control exp on pr_theta100_alpha50.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [01:32<00:00, 18.59s/it]\n"
     ]
    }
   ],
   "source": [
    "ads_hvg = []\n",
    "theta=100\n",
    "ncomp=1000\n",
    "\n",
    "for alpha in tqdm(alphas):\n",
    "    \n",
    "    tag = get_tag(alpha=alpha, theta=100,clipping=True)\n",
    "    \n",
    "    #subset to HVGs\n",
    "    ad = adata[:,adata.var[f'top3000_{tag}']].copy()\n",
    "    ad.uns['alpha'] = alpha\n",
    "    \n",
    "    #recompute residuals\n",
    "    tag_suffix='_afterHVG'\n",
    "    full_tag = tag + tag_suffix\n",
    "    print(full_tag)\n",
    "    compute_residuals(ad,theta=theta,alpha=alpha,tag_suffix=tag_suffix)\n",
    "    \n",
    "    #compute PCA\n",
    "    pca = PCA(random_state=42)\n",
    "    ad.obsm['pca'] = rescale_pca(pca.fit_transform(ad.layers[full_tag]))\n",
    "    ad.obsm[f'pca{ncomp}'] = ad.obsm['pca'][:,:ncomp]\n",
    "    \n",
    "    #if alpha = 50, also compute PCA without recomputing residuals (control experiment)\n",
    "    if alpha==50:\n",
    "        print('running control exp on', tag)\n",
    "        pca = PCA(random_state=42)\n",
    "        ad.obsm['pca_control'] = rescale_pca(pca.fit_transform(ad.layers[tag]))\n",
    "        ad.obsm[f'pca{ncomp}_control'] = ad.obsm['pca_control'][:,:ncomp]\n",
    "    \n",
    "    ads_hvg.append(ad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### set first two PCs as shared init for suppl. figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#define joint PCA init for suppl figure\n",
    "init_alpha = 10\n",
    "for ad in ads_hvg:\n",
    "    if ad.uns['alpha']==init_alpha:\n",
    "        #use first 2 PCs to initialize the init\n",
    "        pca_init = ad.obsm['pca'][:,:2]\n",
    "        adata.obsm['pca_shared_init'] = pca_init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### compute tSNEs\n",
    "Suppl figure tSNEs with shared PCA init, main figure tSNE with default init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "TSNE(n_jobs=38, random_state=42, verbose=True)\n",
      "--------------------------------------------------------------------------------\n",
      "===> Finding 90 nearest neighbors using Annoy approximate search using euclidean distance...\n",
      "   --> Time elapsed: 7.27 seconds\n",
      "===> Calculating affinity matrix...\n",
      "   --> Time elapsed: 0.33 seconds\n",
      "===> Running optimization with exaggeration=12.00, lr=1985.17 for 250 iterations...\n",
      "Iteration   50, KL divergence 5.9971, 50 iterations in 1.7796 sec\n",
      "Iteration  100, KL divergence 4.6969, 50 iterations in 1.6073 sec\n",
      "Iteration  150, KL divergence 4.4263, 50 iterations in 1.5620 sec\n",
      "Iteration  200, KL divergence 4.2650, 50 iterations in 1.6657 sec\n",
      "Iteration  250, KL divergence 4.1859, 50 iterations in 1.6947 sec\n",
      "   --> Time elapsed: 8.31 seconds\n",
      "===> Running optimization with exaggeration=1.00, lr=1985.17 for 500 iterations...\n",
      "Iteration   50, KL divergence 3.0453, 50 iterations in 1.6123 sec\n",
      "Iteration  100, KL divergence 2.6213, 50 iterations in 1.7060 sec\n",
      "Iteration  150, KL divergence 2.4315, 50 iterations in 1.6303 sec\n",
      "Iteration  200, KL divergence 2.3261, 50 iterations in 2.6629 sec\n",
      "Iteration  250, KL divergence 2.2603, 50 iterations in 2.8507 sec\n",
      "Iteration  300, KL divergence 2.2159, 50 iterations in 3.7701 sec\n",
      "Iteration  350, KL divergence 2.1853, 50 iterations in 4.4227 sec\n",
      "Iteration  400, KL divergence 2.1642, 50 iterations in 6.1565 sec\n",
      "Iteration  450, KL divergence 2.1485, 50 iterations in 6.1260 sec\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 1/5 [01:04<04:18, 64.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration  500, KL divergence 2.1353, 50 iterations in 6.4367 sec\n",
      "   --> Time elapsed: 37.38 seconds\n",
      "--------------------------------------------------------------------------------\n",
      "TSNE(n_jobs=38, random_state=42, verbose=True)\n",
      "--------------------------------------------------------------------------------\n",
      "===> Finding 90 nearest neighbors using Annoy approximate search using euclidean distance...\n",
      "   --> Time elapsed: 8.09 seconds\n",
      "===> Calculating affinity matrix...\n",
      "   --> Time elapsed: 0.25 seconds\n",
      "===> Running optimization with exaggeration=12.00, lr=1985.17 for 250 iterations...\n",
      "Iteration   50, KL divergence 5.5512, 50 iterations in 1.5163 sec\n",
      "Iteration  100, KL divergence 4.4510, 50 iterations in 1.3975 sec\n",
      "Iteration  150, KL divergence 4.1906, 50 iterations in 1.5069 sec\n",
      "Iteration  200, KL divergence 4.0650, 50 iterations in 1.3944 sec\n",
      "Iteration  250, KL divergence 3.9927, 50 iterations in 1.7286 sec\n",
      "   --> Time elapsed: 7.54 seconds\n",
      "===> Running optimization with exaggeration=1.00, lr=1985.17 for 500 iterations...\n",
      "Iteration   50, KL divergence 2.9647, 50 iterations in 1.6612 sec\n",
      "Iteration  100, KL divergence 2.5013, 50 iterations in 1.4085 sec\n",
      "Iteration  150, KL divergence 2.2674, 50 iterations in 2.3829 sec\n",
      "Iteration  200, KL divergence 2.1277, 50 iterations in 3.2839 sec\n",
      "Iteration  250, KL divergence 2.0371, 50 iterations in 4.1520 sec\n",
      "Iteration  300, KL divergence 1.9755, 50 iterations in 4.7782 sec\n",
      "Iteration  350, KL divergence 1.9326, 50 iterations in 5.0926 sec\n",
      "Iteration  400, KL divergence 1.8989, 50 iterations in 6.6313 sec\n",
      "Iteration  450, KL divergence 1.8758, 50 iterations in 8.9298 sec\n",
      "Iteration  500, KL divergence 1.8577, 50 iterations in 9.3564 sec\n",
      "   --> Time elapsed: 47.68 seconds\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 2/5 [02:08<03:12, 64.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "TSNE(n_jobs=38, random_state=42, verbose=True)\n",
      "--------------------------------------------------------------------------------\n",
      "===> Finding 90 nearest neighbors using Annoy approximate search using euclidean distance...\n",
      "   --> Time elapsed: 7.21 seconds\n",
      "===> Calculating affinity matrix...\n",
      "   --> Time elapsed: 0.25 seconds\n",
      "===> Running optimization with exaggeration=12.00, lr=1985.17 for 250 iterations...\n",
      "Iteration   50, KL divergence 5.5577, 50 iterations in 1.6917 sec\n",
      "Iteration  100, KL divergence 4.4898, 50 iterations in 1.7616 sec\n",
      "Iteration  150, KL divergence 4.2380, 50 iterations in 1.7966 sec\n",
      "Iteration  200, KL divergence 4.1201, 50 iterations in 1.7783 sec\n",
      "Iteration  250, KL divergence 4.0503, 50 iterations in 1.8125 sec\n",
      "   --> Time elapsed: 8.84 seconds\n",
      "===> Running optimization with exaggeration=1.00, lr=1985.17 for 500 iterations...\n",
      "Iteration   50, KL divergence 3.0064, 50 iterations in 1.4606 sec\n",
      "Iteration  100, KL divergence 2.5189, 50 iterations in 1.3315 sec\n",
      "Iteration  150, KL divergence 2.2698, 50 iterations in 1.7080 sec\n",
      "Iteration  200, KL divergence 2.1207, 50 iterations in 3.4346 sec\n",
      "Iteration  250, KL divergence 2.0227, 50 iterations in 4.2213 sec\n",
      "Iteration  300, KL divergence 1.9557, 50 iterations in 4.8652 sec\n",
      "Iteration  350, KL divergence 1.9073, 50 iterations in 6.0965 sec\n",
      "Iteration  400, KL divergence 1.8710, 50 iterations in 6.4767 sec\n",
      "Iteration  450, KL divergence 1.8438, 50 iterations in 8.4656 sec\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 3/5 [03:12<02:08, 64.07s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration  500, KL divergence 1.8211, 50 iterations in 9.3200 sec\n",
      "   --> Time elapsed: 47.38 seconds\n",
      "--------------------------------------------------------------------------------\n",
      "TSNE(n_jobs=38, random_state=42, verbose=True)\n",
      "--------------------------------------------------------------------------------\n",
      "===> Finding 90 nearest neighbors using Annoy approximate search using euclidean distance...\n",
      "   --> Time elapsed: 7.27 seconds\n",
      "===> Calculating affinity matrix...\n",
      "   --> Time elapsed: 0.27 seconds\n",
      "===> Running optimization with exaggeration=12.00, lr=1985.17 for 250 iterations...\n",
      "Iteration   50, KL divergence 5.7264, 50 iterations in 1.5823 sec\n",
      "Iteration  100, KL divergence 4.5457, 50 iterations in 1.6690 sec\n",
      "Iteration  150, KL divergence 4.2805, 50 iterations in 1.4719 sec\n",
      "Iteration  200, KL divergence 4.1568, 50 iterations in 1.5615 sec\n",
      "Iteration  250, KL divergence 4.0849, 50 iterations in 1.4904 sec\n",
      "   --> Time elapsed: 7.78 seconds\n",
      "===> Running optimization with exaggeration=1.00, lr=1985.17 for 500 iterations...\n",
      "Iteration   50, KL divergence 3.0247, 50 iterations in 1.4652 sec\n",
      "Iteration  100, KL divergence 2.5359, 50 iterations in 1.4521 sec\n",
      "Iteration  150, KL divergence 2.2865, 50 iterations in 1.9517 sec\n",
      "Iteration  200, KL divergence 2.1372, 50 iterations in 3.0983 sec\n",
      "Iteration  250, KL divergence 2.0385, 50 iterations in 3.9785 sec\n",
      "Iteration  300, KL divergence 1.9723, 50 iterations in 4.6674 sec\n",
      "Iteration  350, KL divergence 1.9243, 50 iterations in 6.5171 sec\n",
      "Iteration  400, KL divergence 1.8897, 50 iterations in 7.2973 sec\n",
      "Iteration  450, KL divergence 1.8631, 50 iterations in 8.2776 sec\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 4/5 [04:15<01:03, 63.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration  500, KL divergence 1.8411, 50 iterations in 8.3345 sec\n",
      "   --> Time elapsed: 47.05 seconds\n",
      "--------------------------------------------------------------------------------\n",
      "TSNE(n_jobs=38, random_state=42, verbose=True)\n",
      "--------------------------------------------------------------------------------\n",
      "===> Finding 90 nearest neighbors using Annoy approximate search using euclidean distance...\n",
      "   --> Time elapsed: 7.37 seconds\n",
      "===> Calculating affinity matrix...\n",
      "   --> Time elapsed: 0.26 seconds\n",
      "===> Running optimization with exaggeration=12.00, lr=1985.17 for 250 iterations...\n",
      "Iteration   50, KL divergence 5.9875, 50 iterations in 1.6773 sec\n",
      "Iteration  100, KL divergence 4.5360, 50 iterations in 1.7754 sec\n",
      "Iteration  150, KL divergence 4.2422, 50 iterations in 1.6580 sec\n",
      "Iteration  200, KL divergence 4.1145, 50 iterations in 1.6969 sec\n",
      "Iteration  250, KL divergence 4.0366, 50 iterations in 1.6444 sec\n",
      "   --> Time elapsed: 8.45 seconds\n",
      "===> Running optimization with exaggeration=1.00, lr=1985.17 for 500 iterations...\n",
      "Iteration   50, KL divergence 2.9852, 50 iterations in 1.4676 sec\n",
      "Iteration  100, KL divergence 2.5050, 50 iterations in 1.5665 sec\n",
      "Iteration  150, KL divergence 2.2589, 50 iterations in 1.9188 sec\n",
      "Iteration  200, KL divergence 2.1119, 50 iterations in 3.0230 sec\n",
      "Iteration  250, KL divergence 2.0148, 50 iterations in 3.2992 sec\n",
      "Iteration  300, KL divergence 1.9485, 50 iterations in 4.6277 sec\n",
      "Iteration  350, KL divergence 1.9005, 50 iterations in 5.6443 sec\n",
      "Iteration  400, KL divergence 1.8656, 50 iterations in 6.8104 sec\n",
      "Iteration  450, KL divergence 1.8403, 50 iterations in 8.0626 sec\n",
      "Iteration  500, KL divergence 1.8194, 50 iterations in 10.0655 sec\n",
      "   --> Time elapsed: 46.49 seconds\n",
      "--------------------------------------------------------------------------------\n",
      "TSNE(n_jobs=38, random_state=42, verbose=True)\n",
      "--------------------------------------------------------------------------------\n",
      "===> Finding 90 nearest neighbors using Annoy approximate search using euclidean distance...\n",
      "   --> Time elapsed: 7.86 seconds\n",
      "===> Calculating affinity matrix...\n",
      "   --> Time elapsed: 0.32 seconds\n",
      "===> Running optimization with exaggeration=12.00, lr=1985.17 for 250 iterations...\n",
      "Iteration   50, KL divergence 5.7141, 50 iterations in 1.5945 sec\n",
      "Iteration  100, KL divergence 4.5132, 50 iterations in 1.4951 sec\n",
      "Iteration  150, KL divergence 4.2369, 50 iterations in 1.5187 sec\n",
      "Iteration  200, KL divergence 4.1140, 50 iterations in 1.4216 sec\n",
      "Iteration  250, KL divergence 4.0415, 50 iterations in 1.5916 sec\n",
      "   --> Time elapsed: 7.62 seconds\n",
      "===> Running optimization with exaggeration=1.00, lr=1985.17 for 500 iterations...\n",
      "Iteration   50, KL divergence 2.9925, 50 iterations in 1.7071 sec\n",
      "Iteration  100, KL divergence 2.5035, 50 iterations in 1.8334 sec\n",
      "Iteration  150, KL divergence 2.2524, 50 iterations in 2.6039 sec\n",
      "Iteration  200, KL divergence 2.1031, 50 iterations in 3.6772 sec\n",
      "Iteration  250, KL divergence 2.0057, 50 iterations in 4.4823 sec\n",
      "Iteration  300, KL divergence 1.9402, 50 iterations in 5.8268 sec\n",
      "Iteration  350, KL divergence 1.8925, 50 iterations in 6.7649 sec\n",
      "Iteration  400, KL divergence 1.8577, 50 iterations in 8.0154 sec\n",
      "Iteration  450, KL divergence 1.8307, 50 iterations in 9.7310 sec\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [06:29<00:00, 77.83s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration  500, KL divergence 1.8087, 50 iterations in 10.5529 sec\n",
      "   --> Time elapsed: 55.20 seconds\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "ncomp = 1000\n",
    "for ad in tqdm(ads_hvg):\n",
    "    \n",
    "    if ad.uns['alpha'] in [1,10,100,1000]:\n",
    "        #use shared PCA init\n",
    "        init=adata.obsm['pca_shared_init']\n",
    "    elif ad.uns['alpha']==50:\n",
    "        #use default PCA init (first 2 PCs)\n",
    "        init=ad.obsm['pca'][:,:2]\n",
    "            \n",
    "    pca_data_after_HVG = ad.obsm[f'pca{ncomp}']\n",
    "    tsne = openTSNE.TSNE(random_state=42,verbose=True,n_jobs=38)\n",
    "    ad.obsm['tsne'] = np.array(tsne.fit(X=pca_data_after_HVG,initialization=init))\n",
    "    \n",
    "    if ad.uns['alpha'] == 50:\n",
    "        #compute control tSNE where residuals where not re-computed after HVG selection / before PCA\n",
    "        pca_data_after_HVG = ad.obsm[f'pca{ncomp}_control']\n",
    "        tsne = openTSNE.TSNE(random_state=42,verbose=True,n_jobs=38)\n",
    "        ad.obsm['tsne_hvg_control'] = np.array(tsne.fit(X=pca_data_after_HVG,initialization=init))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run scanpy baseline preprocessing for comparion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/anndata/_core/anndata.py:1830: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"var\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "normalizing counts per cell\n",
      "    finished (0:00:01)\n",
      "If you pass `n_top_genes`, all cutoffs are ignored.\n",
      "extracting highly variable genes\n",
      "    finished (0:00:05)\n",
      "normalizing counts per cell\n",
      "    finished (0:00:00)\n",
      "--------------------------------------------------------------------------------\n",
      "TSNE(n_jobs=38, random_state=42, verbose=True)\n",
      "--------------------------------------------------------------------------------\n",
      "===> Finding 90 nearest neighbors using Annoy approximate search using euclidean distance...\n",
      "   --> Time elapsed: 7.04 seconds\n",
      "===> Calculating affinity matrix...\n",
      "   --> Time elapsed: 0.28 seconds\n",
      "===> Running optimization with exaggeration=12.00, lr=1985.17 for 250 iterations...\n",
      "Iteration   50, KL divergence 5.4656, 50 iterations in 1.4626 sec\n",
      "Iteration  100, KL divergence 4.4254, 50 iterations in 1.5978 sec\n",
      "Iteration  150, KL divergence 4.1669, 50 iterations in 1.5271 sec\n",
      "Iteration  200, KL divergence 4.0382, 50 iterations in 1.6023 sec\n",
      "Iteration  250, KL divergence 3.9628, 50 iterations in 1.6826 sec\n",
      "   --> Time elapsed: 7.87 seconds\n",
      "===> Running optimization with exaggeration=1.00, lr=1985.17 for 500 iterations...\n",
      "Iteration   50, KL divergence 3.0145, 50 iterations in 1.6652 sec\n",
      "Iteration  100, KL divergence 2.6037, 50 iterations in 1.5923 sec\n",
      "Iteration  150, KL divergence 2.4109, 50 iterations in 1.6699 sec\n",
      "Iteration  200, KL divergence 2.3022, 50 iterations in 2.5603 sec\n",
      "Iteration  250, KL divergence 2.2343, 50 iterations in 3.6713 sec\n",
      "Iteration  300, KL divergence 2.1874, 50 iterations in 3.7426 sec\n",
      "Iteration  350, KL divergence 2.1560, 50 iterations in 4.7974 sec\n",
      "Iteration  400, KL divergence 2.1335, 50 iterations in 5.8488 sec\n",
      "Iteration  450, KL divergence 2.1152, 50 iterations in 6.7618 sec\n",
      "Iteration  500, KL divergence 2.1000, 50 iterations in 6.8048 sec\n",
      "   --> Time elapsed: 39.12 seconds\n"
     ]
    }
   ],
   "source": [
    "ad_hvg_seurat = scanpy_preproc_baseline(adata,n_hvgs=3000,n_comps=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with open(f'{basepath}ads_hvg.pickle','wb') as f:\n",
    "    pickle.dump(ads_hvg,f)\n",
    "adata.write_h5ad(f'{basepath}adata_residuals.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "ad_hvg_seurat.uns['hvg_criterion'] =ad_hvg_seurat.uns['hvg_criterion'].values\n",
    "ad_hvg_seurat.write(f'{basepath}ad_hvg_seurat.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}