train_and_get_grads.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import pearsonr,zscore\n",
    "import hickle as hkl\n",
    "from sklearn.metrics import r2_score\n",
    "from numpy.linalg import solve, svd, norm\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_embedding = hkl.load(\"embeddings/final_X_tcga_processed.hkl\")\n",
    "cell_embedding /= norm(cell_embedding, axis=1).reshape(-1, 1)\n",
    "gene_effects_df = hkl.load(\"datasets/2023/CRISPRGeneEffect_processed.hkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_individual_rfm_cell():\n",
    "    device = \"cuda:3\"\n",
    "    bandwidth = 1\n",
    "    reg = 1e-5\n",
    "\n",
    "    X = torch.tensor(cell_embedding.values).to(device).float()\n",
    "\n",
    "    num_cells = X.shape[0]\n",
    "    knockouts = gene_effects_df.columns\n",
    "    num_knockouts = len(knockouts)\n",
    "\n",
    "    cell_distances = kernel.euclidean_distances(X,X).to(device)\n",
    "    dist_ko = cell_distances.fill_diagonal_(0)\n",
    "    y = torch.tensor(gene_effects_df.values).to(device).float()\n",
    "    sol = torch.linalg.solve(torch.exp(-bandwidth*(dist_ko)**0.5).to(device) + reg*torch.eye(dist_ko.shape[0],device=device),y)\n",
    "\n",
    "    return sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sol = train_individual_rfm_cell()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get Grads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def euclidean_distances(samples, centers, M=None,squared=True,diag_only=False):\n",
    "    '''Calculate the pointwise distance.\n",
    "    Args:\n",
    "        samples: of shape (n_sample, n_feature).\n",
    "        centers: of shape (n_center, n_feature).\n",
    "        squared: boolean.\n",
    "    Returns:\n",
    "        pointwise distances (n_sample, n_center).\n",
    "    '''\n",
    "    if M is None:\n",
    "        samples_norm = torch.sum(samples**2, dim=1, keepdim=True)\n",
    "    else:\n",
    "        if diag_only:\n",
    "            samples_norm = (samples * M) * samples\n",
    "            # samples_cpu = samples.detach().cpu()\n",
    "            # M_cpu = M.detach().cpu()\n",
    "            # samples_norm = samples_cpu**2 * M_cpu\n",
    "            # samples_norm = samples_norm.to(\"cuda:1\")\n",
    "        else:\n",
    "            samples_norm = (samples @ M) * samples\n",
    "        samples_norm = torch.sum(samples_norm, dim=1, keepdims=True)\n",
    "\n",
    "    if samples is centers:\n",
    "        centers_norm = samples_norm\n",
    "    else:\n",
    "        if M is None:\n",
    "            centers_norm = torch.sum(centers**2, dim=1, keepdims=True)\n",
    "        else:\n",
    "            # centers_norm = (centers.diag() * M).diag() * centers\n",
    "            if diag_only:\n",
    "                centers_norm = (centers * M) * centers\n",
    "            else:\n",
    "                centers_norm = (centers @ M) * centers\n",
    "            centers_norm = torch.sum(centers_norm, dim=1, keepdims=True)\n",
    "    centers_norm = torch.reshape(centers_norm, (1, -1))\n",
    "\n",
    "    distances = samples.mm(torch.t(centers))\n",
    "    distances.mul_(-2)\n",
    "    distances.add_(samples_norm)\n",
    "    distances.add_(centers_norm)\n",
    "    if not squared:\n",
    "        distances.clamp_(min=0)\n",
    "        distances.sqrt_()\n",
    "\n",
    "    return distances\n",
    "\n",
    "def laplace_kernel(samples, centers, bandwidth,M=None, diag_only=False):\n",
    "    '''Laplacian kernel.\n",
    "    Args:\n",
    "        samples: of shape (n_sample, n_feature).\n",
    "        centers: of shape (n_center, n_feature).\n",
    "        bandwidth: kernel bandwidth.\n",
    "    Returns:\n",
    "        kernel matrix of shape (n_sample, n_center).\n",
    "    '''\n",
    "    assert bandwidth > 0\n",
    "    kernel_mat = euclidean_distances(samples, centers, M=M, squared=False, diag_only=diag_only)\n",
    "    kernel_mat.clamp_(min=0)\n",
    "    gamma = 1. / bandwidth\n",
    "    kernel_mat.mul_(-gamma)\n",
    "    kernel_mat.exp_()\n",
    "    return kernel_mat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_grads(X, sol, P, L=1,centering=False,diag_only=True):\n",
    "    \n",
    "    K = laplace_kernel(X, X, bandwidth=1, M=P, diag_only=diag_only)\n",
    "\n",
    "    dist = euclidean_distances(X, X, M=P, squared=False, diag_only=diag_only)\n",
    "    dist.clamp_(min=0)\n",
    "    dist[dist < 1e-10] = 0\n",
    "\n",
    "    with np.errstate(divide='ignore'):\n",
    "        K = K/dist\n",
    "\n",
    "    K[K == float(\"Inf\")] = 0.\n",
    "    n,d = X.shape\n",
    "    num_kos,n = sol.shape\n",
    "\n",
    "    grads = torch.zeros((d,num_kos)).to(X.device)\n",
    "    for i in tqdm(range(num_kos)):\n",
    "        weight = sol[i,:].reshape((-1,1))\n",
    "\n",
    "        step2 = K @ (weight * X)\n",
    "        step3 = (weight.T @ K).T * X\n",
    "        G = (step2 - step3) * -1/L\n",
    "        G = torch.sum(G**2,axis=0)\n",
    "        grads[:,i] = G/n\n",
    "\n",
    "    return grads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda:3\"\n",
    "X = torch.tensor(cell_embedding.values).to(device).float()\n",
    "n,d = X.shape\n",
    "P = torch.ones(d).double().to(device)\n",
    "grads = get_grads(X,sol.T,L=1,P=P,centering=False,diag_only=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Add PCC Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_pcc():\n",
    "    cell_embedding = hkl.load(\"embeddings/final_X_tcga_processed.hkl\")\n",
    "    cell_embedding /= norm(cell_embedding, axis=1).reshape(-1, 1)   \n",
    "\n",
    "    exp_cols = [e for e in cell_embedding.columns if e.split(\"_\")[-1] == \"exp\"]\n",
    "\n",
    "    std_val = cell_embedding[exp_cols].std(axis=0).replace(0,1)\n",
    "    zscore = (cell_embedding[exp_cols] - cell_embedding[exp_cols].mean(axis=0))/std_val\n",
    "    cell_embedding[exp_cols] *= (np.abs(zscore) < 3).fillna(0).astype(int)\n",
    "\n",
    "    normalized_cell_embedding = cell_embedding - cell_embedding.mean(axis=0)\n",
    "    normalized_gene_effects_df = gene_effects_df - gene_effects_df.mean(axis=0)\n",
    "\n",
    "    cell_norms = (normalized_cell_embedding**2).sum(axis=0).values\n",
    "    gene_norms = (normalized_gene_effects_df**2).sum(axis=0).values\n",
    "\n",
    "    pcc = (normalized_cell_embedding.T @ normalized_gene_effects_df)/(cell_norms.reshape((-1,1)) @ gene_norms.reshape((1,-1)))**0.5\n",
    "    \n",
    "    pcc = pd.DataFrame(pcc,columns=knockouts,index=features)\n",
    "    hkl.dump(pcc,\"datasets/pcc.hkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pcc = get_pcc().fillna(0)\n",
    "pcc = pcc.loc[grads.index]\n",
    "mut = [x for x in pcc.index if x.split(\"_\")[-1] != \"exp\"]\n",
    "pcc.loc[mut] = -(pcc.loc[mut].clip(upper=0))\n",
    "\n",
    "exp = [x for x in pcc.index if x.split(\"_\")[-1] == \"exp\"]\n",
    "pcc.loc[exp] = abs(pcc.loc[exp])\n",
    "\n",
    "feature_importance_df = grads * pcc\n",
    "hkl.dump(feature_importance_df,\"datasets/feature_importances.hkl\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}