testing_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pyro\n",
    "from pyro.optim import Adam, ClippedAdam\n",
    "import congas as cg\n",
    "from congas.models import LatentCategorical\n",
    "from pyro.infer import TraceMeanField_ELBO,TraceEnum_ELBO, TraceGraph_ELBO, Trace_ELBO\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "data_file = open(\"data.pkl\",'rb')\n",
    "\n",
    "data = pickle.loads(data_file.read())\n",
    "data_file.close()\n",
    "\n",
    "param_file =  open(\"parameters.pkl\",'rb')\n",
    "\n",
    "param = pickle.loads(param_file.read())\n",
    "param_file.close()\n",
    "\n",
    "\n",
    "interface = cg.Interface(CUDA = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'atac': {'exp': array([[  82.,  309.,  279., ...,  128.,   54.,   27.],\n",
       "         [  38.,  238.,  293., ...,  120.,  187.,  117.],\n",
       "         [ 330.,  178.,  122., ...,   79.,  294.,  201.],\n",
       "         ...,\n",
       "         [ 753., 1097.,  326., ...,  218.,  252., 1179.],\n",
       "         [ 121.,  205.,  189., ...,  302.,  391.,  506.],\n",
       "         [ 293.,  375.,  319., ...,  435.,  536.,  324.]]),\n",
       "  'MAF': array([[0.37053729, 0.32253585, 0.4158179 , ..., 0.28930478, 0.40497724,\n",
       "          0.2818496 ],\n",
       "         [0.24322206, 0.39135757, 0.34969442, ..., 0.37679794, 0.37214288,\n",
       "          0.36666653],\n",
       "         [0.41908695, 0.38388334, 0.35205534, ..., 0.38393695, 0.37789072,\n",
       "          0.22120664],\n",
       "         ...,\n",
       "         [0.39910356, 0.41747507, 0.47764789, ..., 0.3003187 , 0.23592494,\n",
       "          0.28347963],\n",
       "         [0.23674662, 0.33409717, 0.36924853, ..., 0.30981061, 0.33964248,\n",
       "          0.19481125],\n",
       "         [0.43433506, 0.31206319, 0.21098797, ..., 0.2784139 , 0.36360118,\n",
       "          0.37149347]]),\n",
       "  'labels': array([0, 0, 0, ..., 7, 7, 7]),\n",
       "  'norm_factors': array([1.58362581, 4.56342113, 4.84301328, ..., 3.0888184 , 1.29521766,\n",
       "         3.29757071]),\n",
       "  'theta_shape_atac': array([10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10.,\n",
       "         10., 10.]),\n",
       "  'theta_rate_atac': array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5,\n",
       "         0.5, 0.5])},\n",
       " 'rna': {'exp': array([[  92.,  655.,  189., ..., 1078.,  634.,  116.],\n",
       "         [ 367.,  322.,  184., ...,  342.,  781.,  513.],\n",
       "         [ 112.,  123.,  306., ...,  299.,  535.,  761.],\n",
       "         ...,\n",
       "         [ 298.,  654.,  925., ..., 3186., 3064.,  651.],\n",
       "         [ 864.,  345., 1692., ..., 2367.,  434., 1744.],\n",
       "         [  89.,  215.,  421., ..., 2174., 2756.,  355.]]),\n",
       "  'MAF': array([[0.33985936, 0.44677223, 0.30345976, ..., 0.34187006, 0.44803149,\n",
       "          0.41446912],\n",
       "         [0.35839273, 0.26479499, 0.24144486, ..., 0.26336874, 0.50528797,\n",
       "          0.33883761],\n",
       "         [0.3328056 , 0.31825513, 0.34169097, ..., 0.29650008, 0.32994033,\n",
       "          0.30046783],\n",
       "         ...,\n",
       "         [0.21471308, 0.40165591, 0.19828991, ..., 0.46651373, 0.37248749,\n",
       "          0.42640101],\n",
       "         [0.30309893, 0.37418362, 0.37127361, ..., 0.36217653, 0.27804225,\n",
       "          0.41785873],\n",
       "         [0.31980186, 0.34939168, 0.28904917, ..., 0.33034339, 0.37834006,\n",
       "          0.32590183]]),\n",
       "  'labels': array([0, 0, 0, ..., 7, 7, 7]),\n",
       "  'norm_factors': array([1.42659483, 4.82754723, 2.48500422, ..., 5.51542297, 1.73458245,\n",
       "         5.5646475 ]),\n",
       "  'theta_shape_rna': array([20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20.,\n",
       "         20., 20.]),\n",
       "  'theta_rate_rna': array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5,\n",
       "         0.5, 0.5])},\n",
       " 'CNA_subclones': {'Major':    0   1   2   3   4   5   6   7   8   9   10  11  12  13  14\n",
       "  0   2   3   2   2   2   4   1   2   4   1   2   1   1   3   1\n",
       "  1   3   3   3   1   1   2   4   3   3   1   3   2   2   1   2\n",
       "  2   1   1   3   4   1   4   1   3   2   2   1   2   1   1   4\n",
       "  3   1   3   2   4   3   4   4   4   2   4   3   3   1   1   1\n",
       "  4   3   1   2   3   1   2   1   3   2   1   4   4   2   2   3\n",
       "  5   4   4   3   2   1   2   4   1   3   1   2   4   3   2   2\n",
       "  6   1   3   3   4   1   1   4   1   3   1   2   4   1   4   4\n",
       "  7   4   4   3   4   3   4   1   4   2   2   3   2   1   2   4,\n",
       "  'minor':    0   1   2   3   4   5   6   7   8   9   10  11  12  13  14\n",
       "  0   1   2   1   1   0   0   0   1   2   0   0   0   0   1   0\n",
       "  1   1   1   2   0   0   0   2   0   2   0   1   0   1   0   0\n",
       "  2   0   0   2   3   0   0   0   1   1   1   0   1   0   0   2\n",
       "  3   0   2   1   1   0   3   1   1   1   0   2   1   0   0   0\n",
       "  4   0   0   1   0   0   0   0   0   1   0   2   3   0   1   1\n",
       "  5   1   0   2   0   0   0   3   0   2   0   0   3   1   0   1\n",
       "  6   0   2   2   2   0   0   3   0   2   0   1   0   0   2   2\n",
       "  7   2   2   1   0   1   0   0   0   1   0   2   1   0   1   2,\n",
       "  'total':    0   1   2   3   4   5   6   7   8   9   10  11  12  13  14\n",
       "  0   3   5   3   3   2   4   1   3   6   1   2   1   1   4   1\n",
       "  1   4   4   5   1   1   2   6   3   5   1   4   2   3   1   2\n",
       "  2   1   1   5   7   1   4   1   4   3   3   1   3   1   1   6\n",
       "  3   1   5   3   5   3   7   5   5   3   4   5   4   1   1   1\n",
       "  4   3   1   3   3   1   2   1   3   3   1   6   7   2   3   4\n",
       "  5   5   4   5   2   1   2   7   1   5   1   2   7   4   2   3\n",
       "  6   1   5   5   6   1   1   7   1   5   1   3   4   1   6   6\n",
       "  7   6   6   4   4   4   4   1   4   3   2   5   3   1   3   6},\n",
       " 'CNA_bulk':     Major  minor  total\n",
       " 0     2.0    1.0    3.0\n",
       " 1     3.0    1.0    4.0\n",
       " 2     3.0    2.0    5.0\n",
       " 3     3.0    1.0    4.0\n",
       " 4     2.0    0.0    2.0\n",
       " 5     3.0    0.0    3.0\n",
       " 6     2.0    1.0    3.0\n",
       " 7     3.0    0.0    3.0\n",
       " 8     3.0    2.0    5.0\n",
       " 9     2.0    0.0    2.0\n",
       " 10    2.0    1.0    3.0\n",
       " 11    3.0    1.0    4.0\n",
       " 12    2.0    0.0    2.0\n",
       " 13    2.0    1.0    3.0\n",
       " 14    3.0    1.0    4.0,\n",
       " 'hyperparams': {'ncells': 5000, 'is_multiome': False}}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "I  = 15\n",
    "sim_data = cg.simulate_data_congas(ncells = 5000, seed = 16, subclones=8, \n",
    "                                   mixture_proportion_rna = torch.ones(8) / 8,\n",
    "                                  mixture_proportion_atac = torch.ones(8) / 8, n_segments=I)\n",
    "sim_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5000, 15)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim_data[\"rna\"][\"exp\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {\n",
    "    \"data_atac\" : torch.tensor(sim_data[\"atac\"][\"exp\"]).T.float(),\n",
    "    \"data_rna\" : torch.tensor(sim_data[\"rna\"][\"exp\"]).T.float(),\n",
    "    \"norm_factor_atac\" : torch.tensor(sim_data[\"atac\"][\"norm_factors\"]).float(),\n",
    "    \"norm_factor_rna\" : torch.tensor(sim_data[\"rna\"][\"norm_factors\"]).float(),\n",
    "    \"pld\" : torch.tensor(sim_data['CNA_bulk']['total'].values).float(), \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "param[\"binom_prior_limits\"] = torch.tensor([10,1000])\n",
    "param[\"theta_shape_atac\"] = torch.tensor(sim_data[\"atac\"][\"theta_shape_atac\"]).float()\n",
    "param[\"theta_rate_atac\"] = torch.tensor(sim_data[\"atac\"][\"theta_rate_atac\"]).float()\n",
    "param[\"theta_shape_rna\"] = torch.tensor(sim_data[\"rna\"][\"theta_shape_rna\"]).float()\n",
    "param[\"theta_rate_rna\"] = torch.tensor(sim_data[\"rna\"][\"theta_rate_rna\"]).float()\n",
    "param['nb_size_init_atac'] = torch.tensor([150., 150., 150., 150.])\n",
    "param[\"K\"] = 3\n",
    "param['probs'] = torch.tensor([0.2000, 0.6000, 0.0500, 0.0250, 0.0250])\n",
    "param[\"likelihood_atac\"] = \"NB\"\n",
    "param[\"likelihood_rna\"] = \"NB\"\n",
    "param[\"temperature\"] = 10\n",
    "param[\"lambda\"] = 1\n",
    "param[\"purity\"] = None\n",
    "param[\"nb_size_init_atac\"] = torch.ones(I) * 150\n",
    "param[\"nb_size_init_rna\"] = torch.ones(I) * 150\n",
    "param[\"CUDA\"] = False\n",
    "param[\"multiome\"] = True\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "interface.set_model(LatentCategorical)\n",
    "interface.set_optimizer(ClippedAdam)\n",
    "interface.set_loss(Trace_ELBO)\n",
    "interface.initialize_model(data)\n",
    "interface.set_model_params(param)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ELBO: 6.972898562  : 100%|██████████| 20/20 [00:01<00:00, 14.15it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Done!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "ll = interface.run(steps = 20, param_optimizer = {\"lr\":0.1})\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Computing assignment probabilities\n",
      "Computing information criteria.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "lr = interface.learned_parameters()\n",
    "ICs = interface.calculate_ICs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1736"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(lr[\"assignment_atac\"] == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1736"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(lr[\"assignment_rna\"] == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "congas_reticulate",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}