rna_velocity_steady_state_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b4022f45-13b7-4644-8730-327f48d4a222",
   "metadata": {},
   "source": [
    "Starting from a population of cells at different stages of development process, trajectory inference algorithms aim to reconstruct the developmental sequence of transcriptional changes leading to potential cell fates."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb335f94-0906-47eb-9080-bd151dd45edf",
   "metadata": {},
   "source": [
    "Unspliced pre-mRNAs and mature spliced mRNAs, the former can detectable by presence of introns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "63ed1630-663f-4c3e-9e46-5a48eb5f2009",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import tensorflow_probability as tfp\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "from utils import *\n",
    "import anndata\n",
    "from scipy.sparse import csr_matrix, issparse, coo_matrix\n",
    "import seaborn as sns\n",
    "from scanpy import Neighbors\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bea5c9fc-dba3-4407-a1f0-aed3a5ed268a",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = anndata.read(\"../processed_pancreas.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "958d421f-3b84-474c-a46c-3971c02bc4ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_weight(x, y=None, perc=95):\n",
    "    \"\"\"\n",
    "    gets the masking weights for gene count in upper and lower percentile \n",
    "    does it by normalizing the each gene vector by the max of each gene vector uses clipping \n",
    "    and then creates a mask selecting only the top percentile for computing the constant valued \n",
    "    degradation rate \n",
    "    \"\"\"\n",
    "    xy_norm = np.array(x.A if issparse(x) else x)\n",
    "    if y is not None:\n",
    "        if issparse(y):\n",
    "            y = y.A\n",
    "        xy_norm = xy_norm / np.clip(np.max(xy_norm, axis=0), 1e-3, None)\n",
    "        xy_norm += y / np.clip(np.max(y, axis=0), 1e-3, None)\n",
    "    if isinstance(perc, int):\n",
    "        weights = xy_norm >= np.percentile(xy_norm, perc, axis=0)\n",
    "    else:\n",
    "        lb, ub = np.percentile(xy_norm, perc, axis=0)\n",
    "        weights = (xy_norm <= lb) | (xy_norm >= ub)\n",
    "    return weights\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "559d53af-650a-4831-841a-f53600406f4f",
   "metadata": {},
   "source": [
    "shape is cell x gene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a9725c97-41c1-40d1-9172-e74ccfd13813",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 3696 × 1945\n",
       "    obs: 'clusters_coarse', 'clusters', 'S_score', 'G2M_score', 'initial_size_spliced', 'initial_size_unspliced', 'initial_size', 'n_counts', 'velocity_self_transition'\n",
       "    var: 'highly_variable_genes', 'gene_count_corr', 'means', 'dispersions', 'dispersions_norm', 'highly_variable', 'velocity_gamma', 'velocity_qreg_ratio', 'velocity_r2', 'velocity_genes'\n",
       "    uns: 'clusters_coarse_colors', 'clusters_colors', 'day_colors', 'neighbors', 'pca', 'velocity_graph', 'velocity_graph_neg', 'velocity_params'\n",
       "    obsm: 'X_pca', 'X_umap', 'velocity_umap'\n",
       "    layers: 'Ms', 'Mu', 'spliced', 'unspliced', 'velocity'\n",
       "    obsp: 'connectivities', 'distances'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d794df9e-5808-4047-93be-cf78b286fd0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "spliced = adata.layers[\"Ms\"]\n",
    "unspliced = adata.layers[\"Mu\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4608f322-98ed-4a89-a2f9-5b20a723bbd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "w =  csr_matrix(get_weight(spliced, unspliced))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "80a58e8b-a820-4988-a201-ce3e22bf7fcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = w.multiply(spliced).tocsr(), w.multiply(unspliced).tocsr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bd8daa43-491f-4223-beba-2ac009882618",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = x.A\n",
    "y = y.A"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a41ad902-148e-465b-a562-eced11bb6c97",
   "metadata": {},
   "source": [
    "## Kinetics are - \n",
    "$\\frac{du(t)}{dt} = \\alpha^k(t) - \\beta u(t)$  \n",
    "$\\frac{ds(t)}{dt} = \\beta u(t) - \\gamma s(t)$\n",
    "\n",
    "\n",
    "Now at s.s. the term $\\frac{ds(t)}{dt} = 0$ this implies that $\\beta u(t) - \\gamma s(t) = 0$. $\\gamma$ represents the degradation rate of mRNA and hence we can infer $\\gamma$ from this condition. For this we take the gene values in the upper percentile as they are the ones most likely to be at a steady state and use linear regression to infer the $\\gamma$ values from those list of values  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23ca55fb-0cdd-46dd-afcf-192db7034ae6",
   "metadata": {},
   "source": [
    "## Note on linear regression \n",
    "\n",
    "$Y = \\beta*X + \\epsilon$  In our case $Y \\sim u $ and $X \\sim s$ and we actually infer $\\frac{\\gamma}{\\beta}$\n",
    "\n",
    "A closed form solution is \n",
    "\n",
    "$\\beta = (X^T X)^{-1} X^T Y$\n",
    "\n",
    "\n",
    "To derive the formula used just write the $err = \\sum_{i}^{nobs} (y_i - w_i x_i)^2$ and take derivatives w.r.t $w_i$\n",
    "\n",
    "We can include a positive offset for the basal transcription rate in each cell "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6a398ebb-18a6-437d-a493-1a66d34f07cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "xx = np.multiply(x, x).sum(axis = 0)\n",
    "xy = np.multiply(x, y).sum(axis = 0) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b46def5c-02db-418d-9d63-70b0bcf821ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma = xy/xx\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "294318be-98f2-412e-9559-d2d7b6d5ba73",
   "metadata": {},
   "source": [
    "### compare with scvelo implementation \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "97a91295-c94d-4be1-9b26-8749c62e1e21",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gamma_scanpy = np.asarray(adata.var[\"velocity_gamma\"].tolist())\n",
    "np.allclose(gamma, gamma_scanpy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1026ebe-d7ee-4486-8b0f-dbf8a309d686",
   "metadata": {},
   "source": [
    "### velocity computations \n",
    "\n",
    "the velocity is given by the deviation of the ratio of unspliced and spliced from the ss ratio\n",
    "\n",
    "$\\nu_i = u_i - \\gamma_i' s_i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4a865f89-ae2a-4dab-9033-bcb36cb89bcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "velocity = unspliced - gamma*spliced ## this is ds/dt\n",
    "alpha = unspliced ## this is du/dt\n",
    "r2 = R_squared(velocity, total=unspliced - unspliced.mean(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c7556d3-2f42-4b6d-bd45-a8dc62d91ff1",
   "metadata": {},
   "source": [
    "### compare with scvelo implementation "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "186523bc-4844-4d40-a5f3-d7e7af3fb15c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(velocity, adata.layers[\"velocity\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "92ea77a5-243a-4886-870d-effb75f40ddc",
   "metadata": {},
   "outputs": [],
   "source": [
    "### filter genes by min likelihood comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9ac72d79-fe61-46ce-8949-b2810f46a34e",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_r2 = 1e-2\n",
    "min_ratio = 1e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e279aa32-343e-44f4-83e2-148d8a5454b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_genes(r2, gamma, min_r2, min_ratio, spliced, unspliced):\n",
    "    velocity_mask = (\n",
    "                (r2 > min_r2)\n",
    "                & (gamma > min_ratio)\n",
    "                & (np.max(spliced > 0, 0) > 0)\n",
    "                & (np.max(unspliced > 0, 0) > 0)\n",
    "            )\n",
    "    return velocity_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bc46479f-5de9-4aea-9225-2169866dbf30",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = mask_genes(r2, gamma, min_r2, min_ratio, spliced, unspliced)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 387,
   "id": "7f2d90af-bd13-469f-ab64-c9c7ec5013d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "## filtering step using min likelihood threshold "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "1fc1a9b3-e605-4f0f-b84b-452f47264597",
   "metadata": {},
   "outputs": [],
   "source": [
    "vgenes = adata.var[\"velocity_genes\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "224c3cd4-eceb-41a4-acf7-f610c9a8fd78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vgenes) == sum(vgenes==mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf916b4d-4771-4e67-9026-55cf567e36c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "kwargs = dict(frameon=False, ylabel='cell cycle genes')\n",
    "scv.pl.scatter(adata, \"\", **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0b991adc-b8dc-4b1b-ad3e-46e71a41411b",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'adata' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-bad2acd93aa9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0madata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvar\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'adata' is not defined"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62d99fbb-ff15-4f98-9fbf-63f140a9a188",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "819a145f-2b3d-4f74-895f-161cebd329c5",
   "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.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}