DREAMNets_BuildModel_Train_and_Predict.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "36168a15-1a5d-4783-9673-ee00667c8928",
   "metadata": {},
   "source": [
    "# Import libraries and functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c8cc7498-bf9d-4b12-96e9-20ce804d5ad0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/muntakimrafi/anaconda3/envs/dream/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import os\n",
    "from prixfixe.autosome import AutosomeDataProcessor, AutosomeFirstLayersBlock, AutosomeCoreBlock, AutosomeFinalLayersBlock, AutosomeTrainer, AutosomePredictor\n",
    "from prixfixe.bhi import BHIFirstLayersBlock,BHICoreBlock\n",
    "from prixfixe.unlockdna import UnlockDNACoreBlock\n",
    "from prixfixe.prixfixe import PrixFixeNet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c742eb34-bf37-4241-99e5-b24da9119744",
   "metadata": {},
   "source": [
    "# Initialize paths and variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c1914474-3bd0-434c-add7-283c313abcbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_DATA_PATH = \"data/demo_train.txt\" #change filename to actual training data\n",
    "VALID_DATA_PATH = \"data/demo_val.txt\" #change filename to actual validaiton data\n",
    "TRAIN_BATCH_SIZE = 512 # replace with 1024, if 1024 doesn't fit in gpu memory, decrease by order of 2 (512,256)\n",
    "BATCH_PER_EPOCH = 10 #replace with total amount of possible batches in the training data\n",
    "N_PROCS = 8\n",
    "VALID_BATCH_SIZE = 4096\n",
    "BATCH_PER_VALIDATION = 10 #replace with total amount of possible batches in the validaiton data\n",
    "PLASMID_PATH = \"data/plasmid.json\"\n",
    "SEQ_SIZE = 150\n",
    "NUM_EPOCHS = 5 #replace with 80\n",
    "CUDA_DEVICE_ID = 0\n",
    "lr = 0.005 # 0.001 for attention layers in coreBlock"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07429083-4ebb-42ee-a3b9-e9bdebd6d881",
   "metadata": {},
   "source": [
    "# DataProcessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ee7888fb-a823-411f-b091-dadc6569b080",
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = torch.Generator()\n",
    "generator.manual_seed(2147483647)\n",
    "\n",
    "dataprocessor = AutosomeDataProcessor(\n",
    "    path_to_training_data=TRAIN_DATA_PATH,\n",
    "    path_to_validation_data=VALID_DATA_PATH,\n",
    "    train_batch_size=TRAIN_BATCH_SIZE, \n",
    "    batch_per_epoch=BATCH_PER_EPOCH,\n",
    "    train_workers=N_PROCS,\n",
    "    valid_batch_size=VALID_BATCH_SIZE,\n",
    "    valid_workers=N_PROCS,\n",
    "    shuffle_train=True,\n",
    "    shuffle_val=False,\n",
    "    plasmid_path=PLASMID_PATH,\n",
    "    seqsize=SEQ_SIZE,\n",
    "    generator=generator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b1ac15af-6475-4bb3-adec-bd6786483361",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'x': tensor([[[1., 1., 0.,  ..., 1., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 1., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       " \n",
       "         [[1., 1., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 1., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 1., 1.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
       " \n",
       "         [[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 1., 0.,  ..., 1., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 1., 1.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       " \n",
       "         ...,\n",
       " \n",
       "         [[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 1., 0.,  ..., 1., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 1., 1.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
       " \n",
       "         [[1., 1., 0.,  ..., 1., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 1., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       " \n",
       "         [[1., 1., 0.,  ..., 1., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 1.,  ..., 0., 1., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.]]]),\n",
       " 'y_probs': tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 1.2799e-12, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         [0.0000e+00, 0.0000e+00, 6.1172e-39,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         ...,\n",
       "         [0.0000e+00, 0.0000e+00, 4.4842e-44,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 2.8665e-07, 1.2799e-12,\n",
       "          0.0000e+00],\n",
       "         [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 1.2799e-12, 0.0000e+00,\n",
       "          0.0000e+00]]),\n",
       " 'y': tensor([11.0000, 10.2122,  9.0000,  5.0000, 13.0000, 10.0000, 10.2233, 14.8991,\n",
       "         13.0000,  9.1253, 14.0000, 14.0000, 13.6456, 11.0000, 13.7995,  9.7855,\n",
       "          9.4371, 15.0000, 12.0000, 10.0000,  8.3589,  9.3933, 13.0000, 15.0000,\n",
       "         12.0832, 12.9071,  9.0000, 10.0000, 11.0000, 11.0000, 14.0000, 16.0000,\n",
       "          9.0000,  8.0000, 12.0000,  8.5612, 11.0000,  7.1086, 14.9155, 12.9907,\n",
       "         11.0000, 10.2565,  9.6457, 12.0000, 10.5164, 12.0000, 15.0000, 12.0000,\n",
       "          6.0374,  8.7169, 11.0000, 12.0000, 11.0000,  9.0000, 15.0000, 11.5860,\n",
       "          9.0000, 11.0000,  9.0000, 12.0000, 10.0000,  9.0000,  8.0000, 11.0000,\n",
       "          9.2093, 10.0000, 15.0000, 12.0000, 10.0000, 13.0000, 13.0000, 14.0000,\n",
       "         11.0000,  9.0000, 17.0000, 11.9322,  7.9480, 12.3483, 12.0000,  8.9042,\n",
       "          4.0000, 11.0000,  9.8783, 11.0000, 12.0000,  9.0000, 11.0000, 12.6791,\n",
       "         14.0000, 13.0000, 14.0000,  7.0000, 11.0000, 12.0000, 14.7725, 15.0000,\n",
       "         11.0000,  9.0000,  9.9331, 12.7995,  9.0000,  4.7989, 12.0000,  7.0000,\n",
       "          9.0000, 12.4249, 13.0000, 10.4019, 11.0000, 15.0000, 11.1048, 10.0061,\n",
       "         10.0000, 14.0000, 13.1594,  8.0000, 10.0000, 14.0000, 12.0832, 10.4421,\n",
       "         11.3648,  8.0000, 12.0000, 12.3479,  9.9713,  9.0000, 16.0000, 10.9119,\n",
       "         13.8590,  7.7495, 10.7512, 13.4173, 12.0000,  7.0000, 10.0973, 14.0000,\n",
       "         11.0000, 12.0000, 14.9850, 13.1677, 13.0851, 11.0000, 12.0000, 11.0000,\n",
       "          9.9572,  7.0000, 15.0000,  9.0000, 17.0000, 10.2827, 16.0000, 11.0000,\n",
       "         13.0000,  9.0588, 11.8454, 11.1615, 13.0000,  9.3508, 14.5088, 11.0000,\n",
       "         15.0000, 14.9388,  9.0000, 11.0000, 11.3286,  9.8668, 12.0000, 11.0000,\n",
       "         13.0000, 14.0000, 15.0000,  9.0000, 13.7391, 10.0033,  9.1304, 10.0000,\n",
       "          6.2071,  9.6986, 14.1973,  6.0000,  8.0000, 12.0000, 11.0000, 14.0000,\n",
       "         13.0000, 12.0000, 12.0000, 10.0000,  9.4496, 13.7995, 12.9626,  7.7428,\n",
       "          9.0000, 11.9719, 11.0000, 12.0000, 11.5718, 11.0000, 11.0000, 12.0907,\n",
       "          8.0000, 15.0000, 12.0000, 13.0000, 12.0000, 13.0000, 10.0000,  9.9124,\n",
       "         10.6767, 13.4724, 16.0000, 10.7419, 10.0799,  8.4934, 11.0000, 10.1803,\n",
       "          8.0000, 13.0000, 11.0696, 11.0000, 15.0000,  9.6185,  7.1645, 10.0000,\n",
       "         12.0000, 11.0000, 12.3997, 10.0000,  9.8742,  8.6688,  9.0000, 14.9531,\n",
       "         11.0000, 14.0000, 11.0000,  6.9809, 10.0000,  9.7173, 13.0133, 10.0000,\n",
       "         13.0000, 12.0000,  7.0832, 10.0000, 10.0000,  8.2024, 10.0000,  7.8199,\n",
       "         13.0000,  9.8067, 11.0842, 14.9550, 12.9907,  9.0000, 11.0000, 11.0000,\n",
       "         13.0000, 11.0000, 10.0000, 11.0000,  9.0000,  7.8655, 15.0000, 11.0000,\n",
       "         10.2258, 14.0000,  9.9941, 12.0000,  8.8170, 13.5223, 12.0000, 11.0000,\n",
       "          8.0000, 12.6316,  8.0000,  8.6824, 13.0000, 11.0000,  8.5461, 10.0000,\n",
       "         10.0000,  8.0000, 13.0000,  8.0000, 13.0000, 12.0000,  9.0000, 13.0000,\n",
       "         10.0000, 15.0000, 15.0000, 11.1583, 11.0000, 11.8440,  3.1496, 11.0000,\n",
       "         11.0000,  8.6556, 15.0000, 15.0000, 15.0000, 15.0000, 10.0000, 10.0000,\n",
       "         12.0000, 13.2898, 11.0397,  8.8170, 12.0000,  1.0000, 13.0000, 15.0000,\n",
       "          8.0207, 10.1471, 13.0000, 14.1537, 11.0000, 12.0182, 12.0000, 12.0000,\n",
       "         10.1774, 13.5064,  8.1080, 14.9421, 11.0000, 12.0569, 15.0000, 12.0000,\n",
       "         11.0000, 15.0000, 11.0000, 13.0000, 14.0000, 13.0000, 15.0000,  9.6748,\n",
       "         15.0000, 11.0000,  7.0895,  9.0000, 13.0033, 15.0000,  9.0000, 12.5089,\n",
       "         14.8559, 14.7970, 14.0000,  9.0929,  9.2448, 14.4877,  9.3193, 14.8254,\n",
       "         10.0000, 13.0000, 15.0000, 15.6326,  9.0000,  8.7147, 11.0000, 12.0000,\n",
       "          9.1460, 12.0000,  9.1772,  8.1990, 10.3765,  8.6044, 11.9188, 13.0000,\n",
       "         10.0000, 11.0000,  9.0000, 15.0000,  9.9627, 15.0000,  9.5899,  3.9440,\n",
       "         11.1868, 13.0000, 12.0000,  9.5933, 16.6340, 13.0000, 10.0000,  9.7565,\n",
       "         10.0000, 10.4631,  8.3652, 10.0000, 13.0000, 10.0000, 12.6562, 10.0000,\n",
       "          9.5508, 13.7793,  8.8396, 13.0000, 15.5655, 13.0000, 14.0000, 13.0000,\n",
       "         10.9557, 11.1257, 11.0000,  9.0000, 10.5368,  9.6164, 11.0000, 11.0000,\n",
       "          8.7269, 11.0000, 10.0000,  8.2351, 11.0000,  9.6336, 12.0000,  5.5282,\n",
       "         13.0000, 11.0000,  8.4520,  3.0089, 11.0000, 14.0000, 11.0000, 12.0000,\n",
       "          9.8664, 10.0000, 13.1078, 10.0000, 13.4036, 15.0000, 13.0000, 14.3622,\n",
       "          8.0278, 14.0000,  9.0000,  6.7825, 11.0000, 13.0000,  9.5646, 10.0000,\n",
       "         11.7428, 11.7428,  8.1222, 14.0000,  7.2215, 10.0000, 11.0000, 14.0000,\n",
       "         13.0000, 13.0000, 14.0000,  7.7189, 10.9729, 13.0000, 12.0000, 11.0000,\n",
       "          9.7217, 14.0000, 13.0000, 15.0000, 10.0000,  8.0000, 12.0000, 12.0000,\n",
       "         11.0000, 10.6970, 10.0000, 11.9467, 10.0000, 10.0000, 11.4457,  6.0000,\n",
       "         14.0000, 16.0000,  8.6609,  8.0000,  9.2311,  9.1526, 14.0000, 12.0000,\n",
       "          9.5432, 10.0000, 11.0000, 14.0000, 11.0000, 10.0000, 12.0000, 13.0000,\n",
       "         13.0000, 14.0000, 10.0000, 10.3212, 14.0000, 13.0000, 14.9529, 14.0000,\n",
       "          8.6841, 12.4968, 14.0000,  9.5520, 15.0000, 11.4367, 11.0000, 10.0000,\n",
       "         13.0000, 10.8579,  8.0000, 10.4931,  8.5128,  9.4376, 12.0000, 11.0000],\n",
       "        dtype=torch.float64)}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(dataprocessor.prepare_train_dataloader())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e50580fc-281f-410e-a67a-661bf85a9275",
   "metadata": {},
   "source": [
    "# Prix-Fixe Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fff673d6-f26f-4cdd-be57-f1a2be4c4916",
   "metadata": {},
   "source": [
    "### DREAM-CNN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "184d418d-70f9-4ee9-9a03-0dd437364677",
   "metadata": {},
   "outputs": [],
   "source": [
    "first = BHIFirstLayersBlock(\n",
    "    in_channels = dataprocessor.data_channels(),\n",
    "    out_channels = 320,\n",
    "    seqsize = dataprocessor.data_seqsize(),\n",
    "    kernel_sizes = [9, 15],\n",
    "    pool_size = 1,\n",
    "    dropout = 0.2\n",
    "    )\n",
    "\n",
    "core = AutosomeCoreBlock(in_channels=first.out_channels,\n",
    "                         out_channels =64,\n",
    "                         seqsize=first.infer_outseqsize())\n",
    "\n",
    "final = AutosomeFinalLayersBlock(in_channels=core.out_channels, \n",
    "                                 seqsize=core.infer_outseqsize())\n",
    "model = PrixFixeNet(\n",
    "    first=first,\n",
    "    core=core,\n",
    "    final=final,\n",
    "    generator=generator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "481b1551-219f-434d-8642-3c535d0d7ce7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "===============================================================================================\n",
       "Layer (type:depth-idx)                        Output Shape              Param #\n",
       "===============================================================================================\n",
       "PrixFixeNet                                   [1]                       --\n",
       "├─BHIFirstLayersBlock: 1-1                    --                        --\n",
       "│    └─ModuleList: 2-1                        --                        --\n",
       "│    │    └─ConvBlock: 3-1                    [1, 160, 150]             8,800\n",
       "│    │    └─ConvBlock: 3-2                    [1, 160, 150]             14,560\n",
       "├─AutosomeCoreBlock: 1-2                      --                        --\n",
       "│    └─ModuleDict: 2-2                        --                        --\n",
       "│    │    └─Sequential: 3-3                   [1, 320, 150]             420,048\n",
       "│    │    └─Sequential: 3-4                   [1, 128, 150]             573,696\n",
       "│    │    └─Sequential: 3-5                   [1, 128, 150]             173,856\n",
       "│    │    └─Sequential: 3-6                   [1, 128, 150]             229,632\n",
       "│    │    └─Sequential: 3-7                   [1, 128, 150]             87,072\n",
       "│    │    └─Sequential: 3-8                   [1, 64, 150]              114,816\n",
       "│    │    └─Sequential: 3-9                   [1, 64, 150]              45,968\n",
       "│    │    └─Sequential: 3-10                  [1, 64, 150]              57,472\n",
       "│    │    └─Sequential: 3-11                  [1, 64, 150]              45,968\n",
       "│    │    └─Sequential: 3-12                  [1, 64, 150]              57,472\n",
       "│    │    └─Sequential: 3-13                  [1, 64, 150]              45,968\n",
       "│    │    └─Sequential: 3-14                  [1, 64, 150]              57,472\n",
       "├─AutosomeFinalLayersBlock: 1-3               --                        --\n",
       "│    └─Sequential: 2-3                        [1, 18, 150]              --\n",
       "│    │    └─Conv1d: 3-15                      [1, 18, 150]              1,170\n",
       "===============================================================================================\n",
       "Total params: 1,933,970\n",
       "Trainable params: 1,933,970\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 264.97\n",
       "===============================================================================================\n",
       "Input size (MB): 0.00\n",
       "Forward/backward pass size (MB): 13.33\n",
       "Params size (MB): 7.74\n",
       "Estimated Total Size (MB): 21.07\n",
       "==============================================================================================="
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchinfo import summary\n",
    "summary(model, (1, 6, 150))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "013bcc9e-f130-4a04-bdaa-02b60a49a514",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MODEL_LOG_DIR = f\"prix_fixe_model_weights/0_1_0_0\"\n",
    "# model.load_state_dict(torch.load(os.path.join(MODEL_LOG_DIR, 'model_best.pth')))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f904bb7-984f-45b1-9bc2-51a3548e03fd",
   "metadata": {},
   "source": [
    "### DREAM-RNN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "42438091-8dd8-4728-b633-a15b3097e849",
   "metadata": {},
   "outputs": [],
   "source": [
    "first = BHIFirstLayersBlock(\n",
    "    in_channels = dataprocessor.data_channels(),\n",
    "    out_channels = 320,\n",
    "    seqsize = dataprocessor.data_seqsize(),\n",
    "    kernel_sizes = [9, 15],\n",
    "    pool_size = 1,\n",
    "    dropout = 0.2\n",
    "    )\n",
    "\n",
    "core = BHICoreBlock(\n",
    "    in_channels = first.out_channels,\n",
    "    out_channels = 320,\n",
    "    seqsize = first.infer_outseqsize(),\n",
    "    lstm_hidden_channels = 320,\n",
    "    kernel_sizes = [9, 15],\n",
    "    pool_size = 1,\n",
    "    dropout1 = 0.2,\n",
    "    dropout2 = 0.5\n",
    "    )\n",
    "\n",
    "final = AutosomeFinalLayersBlock(in_channels=core.out_channels, \n",
    "                                 seqsize=core.infer_outseqsize())\n",
    "model = PrixFixeNet(\n",
    "    first=first,\n",
    "    core=core,\n",
    "    final=final,\n",
    "    generator=generator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e76c2da7-eb1d-44d0-9384-3d5aa17908f4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "PrixFixeNet                              [1]                       --\n",
       "├─BHIFirstLayersBlock: 1-1               --                        --\n",
       "│    └─ModuleList: 2-1                   --                        --\n",
       "│    │    └─ConvBlock: 3-1               [1, 160, 150]             8,800\n",
       "│    │    └─ConvBlock: 3-2               [1, 160, 150]             14,560\n",
       "├─BHICoreBlock: 1-2                      --                        --\n",
       "│    └─LSTM: 2-2                         [1, 150, 640]             1,643,520\n",
       "│    └─ModuleList: 2-3                   --                        --\n",
       "│    │    └─ConvBlock: 3-3               [1, 160, 150]             921,760\n",
       "│    │    └─ConvBlock: 3-4               [1, 160, 150]             1,536,160\n",
       "│    └─Dropout: 2-4                      [1, 320, 150]             --\n",
       "├─AutosomeFinalLayersBlock: 1-3          --                        --\n",
       "│    └─Sequential: 2-5                   [1, 18, 150]              --\n",
       "│    │    └─Conv1d: 3-5                  [1, 18, 150]              5,778\n",
       "==========================================================================================\n",
       "Total params: 4,130,578\n",
       "Trainable params: 4,130,578\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 619.59\n",
       "==========================================================================================\n",
       "Input size (MB): 0.00\n",
       "Forward/backward pass size (MB): 1.56\n",
       "Params size (MB): 16.52\n",
       "Estimated Total Size (MB): 18.08\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchinfo import summary\n",
    "summary(model, (1, 6, 150))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "de38726f-c8f0-4fd6-95f8-f3973fb27aa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MODEL_LOG_DIR = f\"prix_fixe_model_weights/0_1_1_0\"\n",
    "# model.load_state_dict(torch.load(os.path.join(MODEL_LOG_DIR, 'model_best.pth')))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ed6d217-5811-4476-8aee-760c9669c03f",
   "metadata": {},
   "source": [
    "### DREAM-Attn Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c9eb2b96-6a36-4b98-9bfb-b75e30d46086",
   "metadata": {},
   "outputs": [],
   "source": [
    "first = AutosomeFirstLayersBlock(in_channels=dataprocessor.data_channels(),\n",
    "                                   out_channels=256, \n",
    "                                   seqsize=dataprocessor.data_seqsize())\n",
    "\n",
    "core = UnlockDNACoreBlock(\n",
    "    in_channels = first.out_channels, out_channels= first.out_channels, seqsize = dataprocessor.data_seqsize(), n_blocks = 4,\n",
    "                                     kernel_size = 15, rate = 0.1, num_heads = 8)\n",
    "\n",
    "final = AutosomeFinalLayersBlock(in_channels=core.out_channels, \n",
    "                                 seqsize=core.infer_outseqsize())\n",
    "model = PrixFixeNet(\n",
    "    first=first,\n",
    "    core=core,\n",
    "    final=final,\n",
    "    generator=generator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6529b6fb-f21b-4ccd-b8e3-e232e34d798e",
   "metadata": {},
   "outputs": [],
   "source": [
    "first = AutosomeFirstLayersBlock(in_channels=dataprocessor.data_channels(),\n",
    "                                   out_channels=256, \n",
    "                                   seqsize=dataprocessor.data_seqsize())\n",
    "\n",
    "core = UnlockDNACoreBlock(\n",
    "    in_channels = first.out_channels, out_channels= first.out_channels, seqsize = dataprocessor.data_seqsize(), n_blocks = 4,\n",
    "                                     kernel_size = 15, rate = 0.1, num_heads = 8)\n",
    "\n",
    "final = AutosomeFinalLayersBlock(in_channels=core.out_channels, \n",
    "                                 seqsize=core.infer_outseqsize())\n",
    "model = PrixFixeNet(\n",
    "    first=first,\n",
    "    core=core,\n",
    "    final=final,\n",
    "    generator=generator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c240c3bc-f496-4d8a-9347-0197af3deeba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "===============================================================================================\n",
       "Layer (type:depth-idx)                        Output Shape              Param #\n",
       "===============================================================================================\n",
       "PrixFixeNet                                   [1]                       --\n",
       "├─AutosomeFirstLayersBlock: 1-1               --                        --\n",
       "│    └─Sequential: 2-1                        [1, 256, 150]             --\n",
       "│    │    └─Conv1d: 3-1                       [1, 256, 150]             10,752\n",
       "│    │    └─BatchNorm1d: 3-2                  [1, 256, 150]             512\n",
       "│    │    └─SiLU: 3-3                         [1, 256, 150]             --\n",
       "├─UnlockDNACoreBlock: 1-2                     --                        --\n",
       "│    └─Embedding: 2-2                         [1, 150, 256]             38,400\n",
       "│    └─ModuleList: 2-3                        --                        --\n",
       "│    │    └─ConformerSASwiGLULayer: 3-4       [1, 256, 150]             1,121,280\n",
       "│    │    └─ConformerSASwiGLULayer: 3-5       [1, 256, 150]             1,121,280\n",
       "│    │    └─ConformerSASwiGLULayer: 3-6       [1, 256, 150]             1,121,280\n",
       "│    │    └─ConformerSASwiGLULayer: 3-7       [1, 256, 150]             1,121,280\n",
       "├─AutosomeFinalLayersBlock: 1-3               --                        --\n",
       "│    └─Sequential: 2-4                        [1, 18, 150]              --\n",
       "│    │    └─Conv1d: 3-8                       [1, 18, 150]              4,626\n",
       "===============================================================================================\n",
       "Total params: 4,539,410\n",
       "Trainable params: 4,539,410\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 47.28\n",
       "===============================================================================================\n",
       "Input size (MB): 0.00\n",
       "Forward/backward pass size (MB): 20.60\n",
       "Params size (MB): 13.95\n",
       "Estimated Total Size (MB): 34.55\n",
       "==============================================================================================="
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchinfo import summary\n",
    "summary(model, (1, 6, 150))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6d8b3d51-42b0-4d9f-a3db-a6685073a693",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # you neeed to download the model_weights\n",
    "# MODEL_LOG_DIR = f\"prix_fixe_model_weights/0_0_2_0\" \n",
    "# model.load_state_dict(torch.load(os.path.join(MODEL_LOG_DIR, 'model_best.pth')))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44c68ff-c05e-459a-9d4f-f100abaa9cc8",
   "metadata": {},
   "source": [
    "# Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "940f32b0-0735-493e-bdac-a3e21e24d456",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = AutosomeTrainer(\n",
    "    model,    \n",
    "    device=torch.device(f\"cuda:{CUDA_DEVICE_ID}\"), \n",
    "    model_dir=\"data/model_weights\",\n",
    "    dataprocessor=dataprocessor,\n",
    "    num_epochs=NUM_EPOCHS,\n",
    "    lr = lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "040dc14d-1f7f-4ae4-af38-ac8326cf6d80",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                     | 0/5 [00:00<?, ?it/s]\n",
      "Train epoch:   0%|                                       | 0/10 [00:00<?, ?it/s]\u001b[A\n",
      "Train epoch:  10%|███                            | 1/10 [00:00<00:06,  1.36it/s]\u001b[A\n",
      "Train epoch:  20%|██████▏                        | 2/10 [00:01<00:03,  2.16it/s]\u001b[A\n",
      "Train epoch:  30%|█████████▎                     | 3/10 [00:01<00:02,  2.63it/s]\u001b[A\n",
      "Train epoch:  40%|████████████▍                  | 4/10 [00:01<00:02,  2.96it/s]\u001b[A\n",
      "Train epoch:  50%|███████████████▌               | 5/10 [00:01<00:01,  3.14it/s]\u001b[A\n",
      "Train epoch:  60%|██████████████████▌            | 6/10 [00:02<00:01,  3.31it/s]\u001b[A\n",
      "Train epoch:  70%|█████████████████████▋         | 7/10 [00:02<00:00,  3.38it/s]\u001b[A\n",
      "Train epoch:  80%|████████████████████████▊      | 8/10 [00:02<00:00,  3.44it/s]\u001b[A\n",
      "Train epoch:  90%|███████████████████████████▉   | 9/10 [00:02<00:00,  3.44it/s]\u001b[A\n",
      "Train epoch: 100%|██████████████████████████████| 10/10 [00:03<00:00,  3.51it/s]\u001b[A\n",
      "                                                                                \u001b[A\n",
      "Validation:   0%|                                         | 0/4 [00:00<?, ?it/s]\u001b[A\n",
      "Validation:  25%|████████▎                        | 1/4 [00:02<00:08,  2.97s/it]\u001b[A\n",
      "Validation:  50%|████████████████▌                | 2/4 [00:03<00:03,  1.59s/it]\u001b[A\n",
      "Validation:  75%|████████████████████████▊        | 3/4 [00:04<00:01,  1.15s/it]\u001b[A\n",
      " 20%|█████████                                    | 1/5 [00:07<00:30,  7.66s/it]\u001b[A\n",
      "Train epoch:   0%|                                       | 0/10 [00:00<?, ?it/s]\u001b[A\n",
      "Train epoch:  10%|███                            | 1/10 [00:00<00:02,  3.67it/s]\u001b[A\n",
      "Train epoch:  20%|██████▏                        | 2/10 [00:00<00:02,  3.67it/s]\u001b[A\n",
      "Train epoch:  30%|█████████▎                     | 3/10 [00:00<00:01,  3.66it/s]\u001b[A\n",
      "Train epoch:  40%|████████████▍                  | 4/10 [00:01<00:01,  3.64it/s]\u001b[A\n",
      "Train epoch:  50%|███████████████▌               | 5/10 [00:01<00:01,  3.65it/s]\u001b[A\n",
      "Train epoch:  60%|██████████████████▌            | 6/10 [00:01<00:01,  3.63it/s]\u001b[A\n",
      "Train epoch:  70%|█████████████████████▋         | 7/10 [00:01<00:00,  3.64it/s]\u001b[A\n",
      "Train epoch:  80%|████████████████████████▊      | 8/10 [00:02<00:00,  3.57it/s]\u001b[A\n",
      "Train epoch:  90%|███████████████████████████▉   | 9/10 [00:02<00:00,  3.56it/s]\u001b[A\n",
      "Train epoch: 100%|██████████████████████████████| 10/10 [00:02<00:00,  3.49it/s]\u001b[A\n",
      "                                                                                \u001b[A\n",
      "Validation:   0%|                                         | 0/4 [00:00<?, ?it/s]\u001b[A\n",
      "Validation:  25%|████████▎                        | 1/4 [00:02<00:08,  2.80s/it]\u001b[A\n",
      "Validation:  50%|████████████████▌                | 2/4 [00:03<00:03,  1.54s/it]\u001b[A\n",
      "Validation:  75%|████████████████████████▊        | 3/4 [00:04<00:01,  1.13s/it]\u001b[A\n",
      " 40%|██████████████████                           | 2/5 [00:14<00:21,  7.33s/it]\u001b[A\n",
      "Train epoch:   0%|                                       | 0/10 [00:00<?, ?it/s]\u001b[A\n",
      "Train epoch:  10%|███                            | 1/10 [00:00<00:02,  3.54it/s]\u001b[A\n",
      "Train epoch:  20%|██████▏                        | 2/10 [00:00<00:02,  3.57it/s]\u001b[A\n",
      "Train epoch:  30%|█████████▎                     | 3/10 [00:00<00:01,  3.60it/s]\u001b[A\n",
      "Train epoch:  40%|████████████▍                  | 4/10 [00:01<00:01,  3.59it/s]\u001b[A\n",
      "Train epoch:  50%|███████████████▌               | 5/10 [00:01<00:01,  3.71it/s]\u001b[A\n",
      "Train epoch:  60%|██████████████████▌            | 6/10 [00:02<00:01,  2.22it/s]\u001b[A\n",
      "Train epoch:  70%|█████████████████████▋         | 7/10 [00:02<00:01,  2.49it/s]\u001b[A\n",
      "Train epoch:  80%|████████████████████████▊      | 8/10 [00:02<00:00,  2.72it/s]\u001b[A\n",
      "Train epoch:  90%|███████████████████████████▉   | 9/10 [00:03<00:00,  2.89it/s]\u001b[A\n",
      "                                                                                \u001b[A\n",
      "Validation:   0%|                                         | 0/4 [00:00<?, ?it/s]\u001b[A\n",
      "Validation:  25%|████████▎                        | 1/4 [00:02<00:08,  2.87s/it]\u001b[A\n",
      "Validation:  50%|████████████████▌                | 2/4 [00:03<00:03,  1.55s/it]\u001b[A\n",
      "Validation:  75%|████████████████████████▊        | 3/4 [00:04<00:01,  1.13s/it]\u001b[A\n",
      " 60%|███████████████████████████                  | 3/5 [00:22<00:14,  7.35s/it]\u001b[A\n",
      "Train epoch:   0%|                                       | 0/10 [00:00<?, ?it/s]\u001b[A\n",
      "Train epoch:  10%|███                            | 1/10 [00:00<00:02,  3.75it/s]\u001b[A\n",
      "Train epoch:  20%|██████▏                        | 2/10 [00:00<00:02,  3.67it/s]\u001b[A\n",
      "Train epoch:  30%|█████████▎                     | 3/10 [00:00<00:01,  3.70it/s]\u001b[A\n",
      "Train epoch:  40%|████████████▍                  | 4/10 [00:01<00:01,  3.67it/s]\u001b[A\n",
      "Train epoch:  50%|███████████████▌               | 5/10 [00:01<00:01,  3.58it/s]\u001b[A\n",
      "Train epoch:  60%|██████████████████▌            | 6/10 [00:01<00:01,  3.54it/s]\u001b[A\n",
      "Train epoch:  70%|█████████████████████▋         | 7/10 [00:01<00:00,  3.46it/s]\u001b[A\n",
      "Train epoch:  80%|████████████████████████▊      | 8/10 [00:02<00:00,  3.54it/s]\u001b[A\n",
      "Train epoch:  90%|███████████████████████████▉   | 9/10 [00:02<00:00,  3.53it/s]\u001b[A\n",
      "Train epoch: 100%|██████████████████████████████| 10/10 [00:02<00:00,  3.58it/s]\u001b[A\n",
      "                                                                                \u001b[A\n",
      "Validation:   0%|                                         | 0/4 [00:00<?, ?it/s]\u001b[A\n",
      "Validation:  25%|████████▎                        | 1/4 [00:02<00:08,  2.70s/it]\u001b[A\n",
      "Validation:  50%|████████████████▌                | 2/4 [00:03<00:02,  1.49s/it]\u001b[A\n",
      "Validation:  75%|████████████████████████▊        | 3/4 [00:03<00:01,  1.11s/it]\u001b[A\n",
      " 80%|████████████████████████████████████         | 4/5 [00:29<00:07,  7.21s/it]\u001b[A\n",
      "Train epoch:   0%|                                       | 0/10 [00:00<?, ?it/s]\u001b[A\n",
      "Train epoch:  10%|███                            | 1/10 [00:00<00:02,  3.52it/s]\u001b[A\n",
      "Train epoch:  20%|██████▏                        | 2/10 [00:00<00:02,  3.50it/s]\u001b[A\n",
      "Train epoch:  30%|█████████▎                     | 3/10 [00:00<00:01,  3.51it/s]\u001b[A\n",
      "Train epoch:  40%|████████████▍                  | 4/10 [00:01<00:01,  3.41it/s]\u001b[A\n",
      "Train epoch:  50%|███████████████▌               | 5/10 [00:01<00:01,  3.39it/s]\u001b[A\n",
      "Train epoch:  60%|██████████████████▌            | 6/10 [00:01<00:01,  3.40it/s]\u001b[A\n",
      "Train epoch:  70%|█████████████████████▋         | 7/10 [00:02<00:00,  3.41it/s]\u001b[A\n",
      "Train epoch:  80%|████████████████████████▊      | 8/10 [00:02<00:00,  3.43it/s]\u001b[A\n",
      "Train epoch:  90%|███████████████████████████▉   | 9/10 [00:02<00:00,  3.39it/s]\u001b[A\n",
      "Train epoch: 100%|██████████████████████████████| 10/10 [00:02<00:00,  3.40it/s]\u001b[A\n",
      "                                                                                \u001b[A\n",
      "Validation:   0%|                                         | 0/4 [00:00<?, ?it/s]\u001b[A\n",
      "Validation:  25%|████████▎                        | 1/4 [00:02<00:08,  2.76s/it]\u001b[A\n",
      "Validation:  50%|████████████████▌                | 2/4 [00:03<00:03,  1.52s/it]\u001b[A\n",
      "Validation:  75%|████████████████████████▊        | 3/4 [00:04<00:01,  1.11s/it]\u001b[A\n",
      "100%|█████████████████████████████████████████████| 5/5 [00:36<00:00,  7.26s/it]\u001b[A\n"
     ]
    }
   ],
   "source": [
    "trainer.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42462ab7-166f-438c-a5ae-022c3aeb0770",
   "metadata": {},
   "source": [
    "# Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "98fa2775-5b49-4096-bc52-dd14c64fd7db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10.054896354675293"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "predictor = AutosomePredictor(model=model, model_pth='data/model_weights/model_best.pth', device=torch.device(f\"cuda:0\"))\n",
    "dna = \"TGCATTTTTTTCACATC\"+ ''.join(random.choice('ACGT') for _ in range(80)) + \"GGTTACGGCTGTT\"\n",
    "predictor.predict(dna)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf378bf9-5d44-4924-bc2e-50662fb8542a",
   "metadata": {},
   "source": [
    "# Prediction on the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9b60c500-3751-4908-8a88-bdbe67640b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df = pd.read_csv('data/filtered_test_data_with_MAUDE_expression.txt', header=None, sep='\\t')\n",
    "\n",
    "from tqdm import tqdm\n",
    "pred_expr = []\n",
    "for seq in tqdm(test_df.iloc[:, 0]):\n",
    "    pred_expr.append(predictor.predict(seq))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "69d32e96-34b3-469b-804b-8ef298eabe13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import pearsonr, spearmanr\n",
    "print(pearsonr(pred_expr, list(test_df.iloc[:, 1])), spearmanr(pred_expr, list(test_df.iloc[:, 1])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d4705a0-d045-43f0-840e-8da9e93ee381",
   "metadata": {},
   "source": [
    "# Score your submission on DREAM Challenge test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "166d2022-69b3-40d1-9b07-227790187756",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "******************************************************\n",
      "Pearson Score: 0.7657255844881551\n",
      "\n",
      "Spearman Score: 0.8228750904214907\n",
      "\n",
      "******************************************************\n",
      "all r: 0.957144539749361\n",
      "\n",
      "all r²: 0.916125669972016\n",
      "\n",
      "all ρ: 0.961451653086994\n",
      "\n",
      "******************************************************\n",
      "high r: 0.6200899915391505\n",
      "\n",
      "low r: 0.6211738513565918\n",
      "\n",
      "yeast r: 0.8382821111688279\n",
      "\n",
      "random r: 0.9677444394489736\n",
      "\n",
      "challenging r: 0.9354983554787447\n",
      "\n",
      "SNVs r: 0.8227819183935022\n",
      "\n",
      "motif perturbation r: 0.9671482009080143\n",
      "\n",
      "motif tiling r: 0.9449999831802987\n",
      "\n",
      "******************************************************\n",
      "high ρ: 0.5754373259429003\n",
      "\n",
      "low ρ: 0.596033541641311\n",
      "\n",
      "yeast ρ: 0.839060331461191\n",
      "\n",
      "random ρ: 0.970287191964816\n",
      "\n",
      "challenging ρ: 0.9289802083256298\n",
      "\n",
      "SNVs ρ: 0.6775184531537061\n",
      "\n",
      "motif perturbation ρ: 0.9611406141596464\n",
      "\n",
      "motif tiling ρ: 0.9273541130425778\n",
      "\n",
      "******************************************************\n"
     ]
    }
   ],
   "source": [
    "pred_expr = pd.read_csv('data/sample_submission.txt', sep = '\\t', header = None).iloc[:,1]\n",
    "from prixfixe.evaluation import evaluate_predictions\n",
    "evaluate_predictions(pred_expr)"
   ]
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}