guide x cell matrix.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b9c5d4b9-86b2-4ea7-856b-6a2744a7868f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pertpy as pt\n",
    "import scanpy as sc\n",
    "import pandas as pd\n",
    "import pandas as pd\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scvi\n",
    "from _quasiSCVI import QuasiSCVI\n",
    "from _quasivae import QuasiVAE\n",
    "from anndata import AnnData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e3b1c333-f625-402d-af88-20fe73512ca1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/anndata/_core/anndata.py:430: FutureWarning: The dtype argument is deprecated and will be removed in late 2024.\n",
      "  warnings.warn(\n",
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/anndata/_core/anndata.py:430: FutureWarning: The dtype argument is deprecated and will be removed in late 2024.\n",
      "  warnings.warn(\n",
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/anndata/_core/anndata.py:430: FutureWarning: The dtype argument is deprecated and will be removed in late 2024.\n",
      "  warnings.warn(\n",
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/anndata/_core/anndata.py:430: FutureWarning: The dtype argument is deprecated and will be removed in late 2024.\n",
      "  warnings.warn(\n",
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/mudata/_core/mudata.py:491: UserWarning: Cannot join columns with the same name because var_names are intersecting.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# Load a sample dataset\n",
    "mdata = pt.dt.papalexi_2021()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "30f45b78-1e86-4b22-84f8-5c66620f9ec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "gdo = mdata.mod[\"gdo\"]\n",
    "# Extract the guide count matrix\n",
    "guide_count_matrix = gdo.X\n",
    "\n",
    "guide_count_df = pd.DataFrame(guide_count_matrix.toarray(), index=gdo.obs.index, columns=gdo.var.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0eef5b36-e9df-4f8d-922f-78d6686c6fd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AnnData object with n_obs × n_vars = 20729 × 111\n",
      "    obs: 'replicate', 'Phase', 'perturbation', 'guide'\n"
     ]
    }
   ],
   "source": [
    "guide_adata = sc.AnnData(X=guide_count_df.values)\n",
    "guide_adata.obs.index = guide_count_df.index\n",
    "guide_adata.var.index = guide_count_df.columns\n",
    "guide_adata.obs['replicate'] = mdata[\"rna\"].obs['replicate']\n",
    "guide_adata.obs['Phase'] = mdata[\"rna\"].obs['Phase']\n",
    "guide_adata.obs['perturbation'] = mdata[\"rna\"].obs['perturbation']\n",
    "guide_adata.obs['guide']= mdata[\"rna\"].obs['NT']\n",
    "# Inspect the AnnData object\n",
    "print(guide_adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a20f9965-4117-4867-bbde-8362eb457e95",
   "metadata": {},
   "outputs": [],
   "source": [
    "scvi.model.SCVI.setup_anndata(guide_adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ecaad243-00c1-4e51-b6cc-27aa92a63f28",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/lightning/pytorch/trainer/setup.py:187: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.\n",
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=9` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 386/386: 100%|███████████████████████████████████████████| 386/386 [04:07<00:00,  1.57it/s, v_num=1, train_loss_step=144, train_loss_epoch=146]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=386` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 386/386: 100%|███████████████████████████████████████████| 386/386 [04:07<00:00,  1.56it/s, v_num=1, train_loss_step=144, train_loss_epoch=146]\n"
     ]
    }
   ],
   "source": [
    "scvi_ref = scvi.model.SCVI(\n",
    "    guide_adata,\n",
    "    use_layer_norm=\"both\",\n",
    "    use_batch_norm=\"none\",\n",
    "    encode_covariates=True,\n",
    "    dropout_rate=0.2,\n",
    "    n_layers=2,\n",
    ")\n",
    "scvi_ref.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "80976f8f-9823-44bc-ac50-b155af0083ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of adata_ref.X: (20729, 18649)\n",
      "Shape of guide_embeddings_df.values: (20729, 10)\n",
      "(20729, 18659)\n"
     ]
    }
   ],
   "source": [
    "guide_embeddings = scvi_ref.get_latent_representation()\n",
    "guide_embeddings_df = pd.DataFrame(guide_embeddings, index=guide_adata.obs.index)\n",
    "adata_ref = mdata[\"rna\"].copy()\n",
    "\n",
    "print(\"Shape of adata_ref.X:\", adata_ref.X.shape)\n",
    "print(\"Shape of guide_embeddings_df.values:\", guide_embeddings_df.values.shape)\n",
    "if not isinstance(adata_ref.X, np.ndarray):\n",
    "    adata_ref.X = adata_ref.X.toarray()\n",
    "\n",
    "# Ensure guide_embeddings_df.values is a numpy array\n",
    "guide_embeddings_array = guide_embeddings_df.values\n",
    "\n",
    "# Concatenate gene expression data (X) with guide embeddings along columns\n",
    "combined_X = np.concatenate([adata_ref.X, guide_embeddings_array], axis=1)\n",
    "print(combined_X.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa407825-38e5-43cf-86c3-bbbda95ed73c",
   "metadata": {},
   "outputs": [],
   "source": [
    "QuasiSCVI.setup_anndata(adata=combined_X, batch_key='replicate', labels_key='perturbation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f5211f-eb22-477f-b9a9-8db7b069261d",
   "metadata": {},
   "outputs": [],
   "source": [
    "guide_model = QuasiSCVI(\n",
    "    combined_X,\n",
    "    use_layer_norm=\"both\",\n",
    "    use_batch_norm=\"none\",\n",
    "    encode_covariates=True,\n",
    "    dropout_rate=0.2,\n",
    "    n_layers=2,\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "guide_model.train(max_epochs=100, plan_kwargs={\"lr\": 1e-4}, accelerator='gpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d053ad2-4981-4547-951c-59210977cfe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "SCVI_LATENT_KEY = \"X_scVI\"\n",
    "\n",
    "combined_X.obsm[SCVI_LATENT_KEY] = guide_model.get_latent_representation()\n",
    "sc.pp.neighbors(combined_X, use_rep=SCVI_LATENT_KEY)\n",
    "sc.tl.leiden(combined_X)\n",
    "sc.tl.umap(combined_X)\n",
    "sc.pl.umap(\n",
    "    combined_X,\n",
    "    color=[\"replicate\", \"Phase\", \"perturbation\" ],\n",
    "    frameon=False,\n",
    "    ncols=1,\n",
    ")"
   ]
  }
 ],
 "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}