Dual Model - A Divorce Story.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "01227c9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run this block *first* every time your kernel starts/restarts\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"  # either 3 or 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ffba0753",
   "metadata": {},
   "outputs": [],
   "source": [
    "GENOMES = { \"mouse\" : \"/users/kcochran/genomes/mm10_no_alt_analysis_set_ENCODE.fasta\",\n",
    "            \"human\" : \"/users/kcochran/genomes/GRCh38_no_alt_analysis_set_GCA_000001405.15.fasta\" }\n",
    "\n",
    "ROOT = \"/users/kcochran/projects/cs197_cross_species_domain_adaptation/\"\n",
    "DATA_DIR = ROOT + \"data/\"\n",
    "\n",
    "EPSILON = 0.00001\n",
    "\n",
    "SPECIES = [\"mouse\", \"human\"]\n",
    "\n",
    "ACCESSIBILITY_FILE = \"/users/angezhao/cs197_covariates/data/human/accessibility/reads.bigWig\"\n",
    "POS_BIND = \"/users/angezhao/cs197_covariates/data/human/CTCF/chr3toY_pos_shuf.bed.gz\"\n",
    "NEG_BIND = \"/users/angezhao/cs197_covariates/data/human/ctcf/chr3toY_neg_shuf.bed.gz\"\n",
    "#WINDOWS = np.array(mouse_false_negatives_windows)\n",
    "train_species = \"mouse\"\n",
    "tf = \"CTCF\"\n",
    "\n",
    "TFS = [\"CTCF\", \"CEBPA\", \"HNF4A\", \"RXRA\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cdbd60c",
   "metadata": {},
   "source": [
    "# Data Loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ae190e90",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.8.1'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gzip\n",
    "import random\n",
    "import numpy as np\n",
    "from pyfaidx import Fasta\n",
    "from torch.utils.data import Dataset\n",
    "import pyBigWig\n",
    "import torch\n",
    "torch.device('cuda')\n",
    "# from torch.utils.data import Dataset\n",
    "#import pytorch\n",
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "123218aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expand_window(start, end, target_len):\n",
    "    midpoint = (start + end) / 2\n",
    "    if not midpoint.is_integer() and target_len % 2 == 0:\n",
    "        midpoint = midpoint - 0.5\n",
    "    if midpoint.is_integer() and target_len % 2 != 0:\n",
    "        midpoint = midpoint - 0.5\n",
    "    new_start = midpoint - target_len / 2\n",
    "    new_end = midpoint + target_len / 2\n",
    "    \n",
    "    assert new_start.is_integer(), new_start\n",
    "    assert new_end.is_integer(), new_end\n",
    "    assert new_start >= 0\n",
    "    assert new_end - new_start == target_len, (new_end, new_start, target_len)\n",
    "    \n",
    "    return int(new_start), int(new_end)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9fc653cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(Dataset):\n",
    "    letter_dict = {\n",
    "        'a':[1,0,0,0],'c':[0,1,0,0],'g':[0,0,1,0],'t':[0,0,0,1],\n",
    "        'n':[0,0,0,0],'A':[1,0,0,0],'C':[0,1,0,0],'G':[0,0,1,0],\n",
    "        'T':[0,0,0,1],'N':[0,0,0,0]}\n",
    "\n",
    "    def __init__(self, species, tf, train_val_test,\n",
    "                 seq_len = 2114, profile_len = 1000, return_labels = True):\n",
    "        \n",
    "        assert train_val_test in [\"train\", \"val\", \"test\"]\n",
    "        print(\"Beginning Init\")\n",
    "        self.pos_b = POS_BIND\n",
    "        self.neg_b = NEG_BIND\n",
    "        self.accessibility = ACCESSIBILITY_FILE #BIGWIGS_DIR + species + \"/\" + tf + \"/final.pos.bigWig\"\n",
    "        self.prof_len = profile_len\n",
    "        self.max_jitter = 0\n",
    "        self.return_labels = return_labels\n",
    "        \n",
    "        #self.genome_file = GENOMES[species]\n",
    "        self.genome_file = GENOMES[species]\n",
    "        self.seq_len = seq_len\n",
    "\n",
    "        self.set_len()\n",
    "        self.coords = self.get_coords()\n",
    "        #self.seqs_onehot = self.convert(self.coords)\n",
    "        self.profiles, self.logcounts_sum, self.logcounts_mean, self.logcounts_max = self.get_profiles_and_logcounts(self.coords)\n",
    "\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.len\n",
    "    \n",
    "    \n",
    "    def set_len(self):\n",
    "        with open(self.pos_b) as f:\n",
    "            gzip_fd = gzip.open(self.pos_b)\n",
    "            self.len = sum([1 for _ in gzip_fd])\n",
    "\n",
    "\n",
    "    def get_coords(self):\n",
    "        coords = []\n",
    "        #coords_tmp = WINDOWS\n",
    "        with open(self.pos_b) as pos:\n",
    "            posf = gzip.open(self.pos_b)\n",
    "            #coords_tmp = [line.split()[:3] for line in posf]\n",
    "            coords_tmp = [line.decode().split()[:3] for line in posf]  # expecting bed file format\n",
    "        #with open(self.neg_b) as neg:\n",
    "        #    negf = gzip.open(self.neg_b)\n",
    "            #coords_tmp = [line.split()[:3] for line in posf]\n",
    "        #    coords_tmp += [line.decode().split()[:3] for line in negf]  # expecting bed file format\n",
    "        \n",
    "        for coord in coords_tmp:\n",
    "            #print(coord)\n",
    "            chrom, start, end = coord[0], int(coord[1]), int(coord[2])\n",
    "            window_start, window_end = expand_window(start, end,\n",
    "                                                     self.seq_len + 2 * self.max_jitter)\n",
    "            coords.append((coord[2], window_start, window_end))  # no strand consideration\n",
    "\n",
    "            \n",
    "        return coords\n",
    "            \n",
    "\n",
    "    def get_profiles_and_logcounts(self, coords):\n",
    "        print(\"inside profiles\")\n",
    "        profiles = []\n",
    "        logcounts_sum = []\n",
    "        logcounts_mean = []\n",
    "        logcounts_max = []\n",
    "\n",
    "        with pyBigWig.open(self.accessibility) as accessibility_reader:\n",
    "            for chrom, start, end in coords:\n",
    "                print (chrom, \" is chrom. \", start, \" is start. \", end, \" is end.\")\n",
    "                prof_start, prof_end = expand_window(start, end,\n",
    "                                                 self.prof_len + 2 * self.max_jitter)\n",
    "                pos_profile = np.array(\n",
    "                    accessibility_reader.values(chrom, prof_start, prof_end))\n",
    "                pos_profile[np.isnan(pos_profile)] = 0 # Should this be epsilon? - JC\n",
    "                profile = np.array(pos_profile)\n",
    "\n",
    "                profiles.append(profile)\n",
    "#                 logcounts_sum.append(np.array(np.sum(np.log((pos_profile) + EPSILON))))\n",
    "#                 logcounts_mean.append(np.array(np.mean(np.log((pos_profile) + EPSILON))))\n",
    "#                 logcounts_max.append(np.array(np.max(np.log((pos_profile) + EPSILON))))\n",
    "\n",
    "                logcounts_sum.append(np.array(np.sum(pos_profile)))\n",
    "                logcounts_mean.append(np.array(np.mean(pos_profile)))\n",
    "                logcounts_max.append(np.array(np.max(pos_profile)))\n",
    "                \n",
    "        print(\"Finished going thorugh it all. \")        \n",
    "        profiles = np.array(profiles)\n",
    "        logcounts_sum = np.array(logcounts_sum)\n",
    "        logcounts_mean = np.array(logcounts_mean)\n",
    "        logcounts_max = np.array(logcounts_max)\n",
    "\n",
    "        \n",
    "        print(profiles.shape, \" is 0th in profiles\")\n",
    "        \n",
    "        return profiles, logcounts_sum, logcounts_mean, logcounts_max\n",
    "                \n",
    "\n",
    "    def convert(self, coords):\n",
    "        seqs_onehot = []\n",
    "        with Fasta(self.genome_file) as converter:\n",
    "            for coord in coords:\n",
    "                chrom, start, stop = coord\n",
    "                #print (chrom, \" is chrom\")\n",
    "                assert chrom in converter\n",
    "                seq = converter[chrom][start:stop].seq\n",
    "                seq_onehot = np.array([self.letter_dict.get(x,[0,0,0,0]) for x in seq])\n",
    "                seqs_onehot.append(seq_onehot)\n",
    "\n",
    "        seqs_onehot = np.array(seqs_onehot)\n",
    "        return seqs_onehot\n",
    "\n",
    "\n",
    "    def __getitem__(self, batch_index):\t\n",
    "        # get coordinates\n",
    "        onehot = self.seqs_onehot[batch_index]\n",
    "        assert onehot.shape[0] > 0, onehot.shape\n",
    "\n",
    "        onehot = torch.tensor(onehot, dtype=torch.float).permute(1, 0)\n",
    "        \n",
    "        if not self.return_labels:\n",
    "            return onehot\n",
    "        else:\n",
    "            # get profiles and logcounts for the two strands\n",
    "            profiles = self.profiles[batch_index]\n",
    "            logcounts = self.logcounts[batch_index]\n",
    "\n",
    "            profiles = torch.tensor(profiles, dtype=torch.float)\n",
    "            logcounts = torch.tensor(logcounts, dtype=torch.float)\n",
    "            return onehot, profiles, logcounts\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7ddf7800",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainGenerator(Dataset):\n",
    "    # must also load in accessibility here\n",
    "    # see profile model starter code for train_generator\n",
    "    # remove the control track, replace the bit w/ not control track with accessibility\n",
    "    # print out shapes\n",
    "    \n",
    "    letter_dict = {\n",
    "        'a':[1,0,0,0],'c':[0,1,0,0],'g':[0,0,1,0],'t':[0,0,0,1],\n",
    "        'n':[0,0,0,0],'A':[1,0,0,0],'C':[0,1,0,0],'G':[0,0,1,0],\n",
    "        'T':[0,0,0,1],'N':[0,0,0,0]}\n",
    "\n",
    "    def __init__(self, species, tf):\n",
    "        self.posfile = DATA_DIR + species + \"/\" + tf + \"/chr3toY_pos_shuf.bed.gz\"\n",
    "        self.negfile = DATA_DIR + species + \"/\" + tf + \"/chr3toY_neg_shuf_run1_1E.bed.gz\"\n",
    "        self.converter = Fasta(GENOMES[species])\n",
    "        self.batchsize = 400\n",
    "        self.halfbatchsize = self.batchsize // 2\n",
    "        self.current_epoch = 1\n",
    "\n",
    "        self.get_coords()\n",
    "        self.on_epoch_end()\n",
    "\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.steps_per_epoch\n",
    "\n",
    "\n",
    "    def get_coords(self):\n",
    "        with gzip.open(self.posfile) as posf:\n",
    "            pos_coords_tmp = [line.decode().split()[:3] for line in posf]  # expecting bed file format\n",
    "            self.pos_coords = [(coord[0], int(coord[1]), int(coord[2])) for coord in pos_coords_tmp]  # no strand consideration\n",
    "        with gzip.open(self.negfile) as negf:\n",
    "            neg_coords_tmp = [line.decode().split()[:3] for line in negf]\n",
    "            self.neg_coords = [(coord[0], int(coord[1]), int(coord[2])) for coord in neg_coords_tmp]\n",
    "            \n",
    "        self.steps_per_epoch = int(len(self.pos_coords) / self.halfbatchsize)\n",
    "        print(self.steps_per_epoch)\n",
    "                \n",
    "\n",
    "    def convert(self, coords):\n",
    "        seqs_onehot = []\n",
    "        for coord in coords:\n",
    "            chrom, start, stop = coord\n",
    "            seq = self.converter[chrom][start:stop].seq\n",
    "            seq_onehot = np.array([self.letter_dict.get(x,[0,0,0,0]) for x in seq])\n",
    "            seqs_onehot.append(seq_onehot)\n",
    "\n",
    "        seqs_onehot = np.array(seqs_onehot)\n",
    "        return seqs_onehot\n",
    "\n",
    "\n",
    "    def __getitem__(self, batch_index):\t\n",
    "        # First, get chunk of coordinates\n",
    "        pos_coords_batch = self.pos_coords[batch_index * self.halfbatchsize : (batch_index + 1) * self.halfbatchsize]\n",
    "        neg_coords_batch = self.neg_coords[batch_index * self.halfbatchsize : (batch_index + 1) * self.halfbatchsize]\n",
    "\n",
    "        # if train_steps calculation is off, lists of coords may be empty\n",
    "        assert len(pos_coords_batch) > 0, len(pos_coords_batch)\n",
    "        assert len(neg_coords_batch) > 0, len(neg_coords_batch)\n",
    "\n",
    "        # Second, convert the coordinates into one-hot encoded sequences\n",
    "        pos_onehot = self.convert(pos_coords_batch)\n",
    "        neg_onehot = self.convert(neg_coords_batch)\n",
    "\n",
    "        # seqdataloader returns empty array if coords are empty list or not in genome\n",
    "        assert pos_onehot.shape[0] > 0, pos_onehot.shape[0]\n",
    "        assert neg_onehot.shape[0] > 0, neg_onehot.shape[0]\n",
    "\n",
    "        # Third, combine bound and unbound sites into one large array, and create label vector\n",
    "        # We don't need to shuffle here because all these examples will correspond\n",
    "        # to a simultaneous gradient update for the whole batch\n",
    "        all_seqs = np.concatenate((pos_onehot, neg_onehot))\n",
    "        labels = np.concatenate((np.ones(pos_onehot.shape[0],), np.zeros(neg_onehot.shape[0],)))\n",
    "\n",
    "        all_seqs = torch.tensor(all_seqs, dtype=torch.float).permute(0, 2, 1)\n",
    "        labels = torch.tensor(labels, dtype=torch.float)\n",
    "        assert all_seqs.shape[0] == self.batchsize, all_seqs.shape[0]\n",
    "        return all_seqs, labels\n",
    "\n",
    "\n",
    "    def on_epoch_end(self):\n",
    "        # switch to next set of negative examples\n",
    "        prev_epoch = self.current_epoch\n",
    "        next_epoch = prev_epoch + 1\n",
    "\n",
    "        # update file where we will retrieve unbound site coordinates from\n",
    "        prev_negfile = self.negfile\n",
    "        next_negfile = prev_negfile.replace(str(prev_epoch) + \"E\", str(next_epoch) + \"E\")\n",
    "        self.negfile = next_negfile\n",
    "\n",
    "        # load in new unbound site coordinates\n",
    "        with gzip.open(self.negfile) as negf:\n",
    "            neg_coords_tmp = [line.decode().split()[:3] for line in negf]\n",
    "            self.neg_coords = [(coord[0], int(coord[1]), int(coord[2])) for coord in neg_coords_tmp]\n",
    "\n",
    "        # then shuffle positive examples\n",
    "        random.shuffle(self.pos_coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b75b7d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ValGenerator(Dataset):\n",
    "    letter_dict = {\n",
    "        'a':[1,0,0,0],'c':[0,1,0,0],'g':[0,0,1,0],'t':[0,0,0,1],\n",
    "        'n':[0,0,0,0],'A':[1,0,0,0],'C':[0,1,0,0],'G':[0,0,1,0],\n",
    "        'T':[0,0,0,1],'N':[0,0,0,0]}\n",
    "\n",
    "    def __init__(self, species, tf, return_labels = True):\n",
    "        self.valfile = DATA_DIR + species + \"/\" + tf + \"/chr1_random_1m.bed.gz\"\n",
    "        self.converter = Fasta(GENOMES[species])\n",
    "        self.batchsize = 1000  # arbitrarily large number that will fit into memory\n",
    "        self.return_labels = return_labels\n",
    "        self.get_coords_and_labels()\n",
    "\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.steps_per_epoch\n",
    "\n",
    "\n",
    "    def get_coords_and_labels(self):\n",
    "        with gzip.open(self.valfile) as f:\n",
    "            coords_tmp = [line.decode().split()[:4] for line in f]  # expecting bed file format\n",
    "        \n",
    "        self.labels = [int(coord[3]) for coord in coords_tmp]\n",
    "        self.coords = [(coord[0], int(coord[1]), int(coord[2])) for coord in coords_tmp]  # no strand consideration\n",
    "        \n",
    "        self.steps_per_epoch = int(len(self.coords) / self.batchsize)\n",
    "\n",
    "    def convert(self, coords):\n",
    "        seqs_onehot = []\n",
    "        for coord in coords:\n",
    "            chrom, start, stop = coord\n",
    "            seq = self.converter[chrom][start:stop].seq\n",
    "            seq_onehot = np.array([self.letter_dict.get(x,[0,0,0,0]) for x in seq])\n",
    "            seqs_onehot.append(seq_onehot)\n",
    "\n",
    "        seqs_onehot = np.array(seqs_onehot)\n",
    "        return seqs_onehot\n",
    "\n",
    "\n",
    "    def __getitem__(self, batch_index):\t\n",
    "        # First, get chunk of coordinates\n",
    "        batch_start = batch_index * self.batchsize\n",
    "        batch_end = (batch_index + 1) * self.batchsize\n",
    "        coords_batch = self.coords[batch_start : batch_end]\n",
    "\n",
    "        # if train_steps calculation is off, lists of coords may be empty\n",
    "        assert len(coords_batch) > 0, len(coords_batch)\n",
    "\n",
    "        # Second, convert the coordinates into one-hot encoded sequences\n",
    "        onehot = self.convert(coords_batch)\n",
    "\n",
    "        # array will be empty if coords are not found in the genome\n",
    "        assert onehot.shape[0] > 0, onehot.shape[0]\n",
    "\n",
    "        onehot = torch.tensor(onehot, dtype=torch.float).permute(0, 2, 1)\n",
    "        \n",
    "        if self.return_labels:\n",
    "            labels = self.labels[batch_start : batch_end]\n",
    "            labels = torch.tensor(labels, dtype=torch.float)\n",
    "            return onehot, labels\n",
    "        else:\n",
    "            return onehot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2779d41c",
   "metadata": {},
   "source": [
    "# Model Training And Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0d093ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Performance metric functions\n",
    "\n",
    "from sklearn.metrics import average_precision_score, roc_auc_score, confusion_matrix, log_loss\n",
    "\n",
    "\n",
    "def print_metrics(preds, labels):\n",
    "    preds = np.array(preds)\n",
    "    labels = np.array(labels)\n",
    "    preds = preds.squeeze()\n",
    "\n",
    "    # this is the binary cross-entropy loss, same as in training\n",
    "    print(\"Loss:\\t\", log_loss(labels, preds))\n",
    "    print(\"auROC:\\t\", roc_auc_score(labels, preds))\n",
    "    auPRC = average_precision_score(labels, preds)\n",
    "    print(\"auPRC:\\t\", auPRC)\n",
    "    print_confusion_matrix(preds, labels)\n",
    "    return auPRC\n",
    "\n",
    "def print_confusion_matrix(preds, labels):\n",
    "    npthresh = np.vectorize(lambda t: 1 if t >= 0.5 else 0)\n",
    "    preds_binarized = npthresh(preds)\n",
    "    conf_matrix = confusion_matrix(labels, preds_binarized)\n",
    "    print(\"Confusion Matrix (at t = 0.5):\\n\", conf_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b5f416b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "\n",
    "\n",
    "class BasicModel(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BasicModel, self).__init__()\n",
    "        self.input_seq_len = 500\n",
    "        num_conv_filters = 240\n",
    "        lstm_hidden_units = 32\n",
    "        fc_layer1_units = 1024\n",
    "        fc_layer2_units = 512\n",
    "        \n",
    "        #How do we know which species of data we should use???\n",
    "        #Our added on data.\n",
    "        self.access_data = Generator(train_species, tf, \"val\")\n",
    "        \n",
    "        # Defining the layers to go into our model\n",
    "        # (see the forward function for how they fit together)\n",
    "        self.conv = torch.nn.Conv1d(4, num_conv_filters, kernel_size=20, padding=0)\n",
    "        # self.conv = torch.nn.Conv1d(1, num_conv_filters, kernel_size=5, padding=0) -> accessibility\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        self.maxpool = torch.nn.MaxPool1d(15, stride=15, padding=0)\n",
    "        self.lstm = torch.nn.LSTM(input_size=num_conv_filters,\n",
    "                                  hidden_size=lstm_hidden_units,\n",
    "                                  batch_first=True)\n",
    "        self.fc1 = torch.nn.Linear(in_features=lstm_hidden_units,\n",
    "                                   out_features=fc_layer1_units)\n",
    "        self.dropout = torch.nn.Dropout(p=0.5)\n",
    "        self.fc2 = torch.nn.Linear(in_features=fc_layer1_units,\n",
    "                                   out_features=fc_layer2_units)\n",
    "        self.fc_final = torch.nn.Linear(in_features=fc_layer2_units,\n",
    "                                        out_features=1)\n",
    "        self.sigmoid = torch.nn.Sigmoid()\n",
    "\n",
    "        # The loss function we'll use -- binary cross-entropy\n",
    "        # (this is the standard loss to use for binary classification)\n",
    "        self.loss = torch.nn.BCELoss()\n",
    "\n",
    "        # We'll store performance metrics during training in these lists\n",
    "        self.train_loss_by_epoch = []\n",
    "        self.source_val_loss_by_epoch = []\n",
    "        self.source_val_auprc_by_epoch = []\n",
    "        self.target_val_loss_by_epoch = []\n",
    "        self.target_val_auprc_by_epoch = []\n",
    "\n",
    "        # We'll record the best model we've seen yet each epoch\n",
    "        self.best_state_so_far = self.state_dict()\n",
    "        self.best_auprc_so_far = 1\n",
    "\n",
    "\n",
    "    def forward(self, X):\n",
    "        Y = np.array(self.accessData.logcounts_mean)\n",
    "        Y_1 = self.relu(self.conv(Y))\n",
    "        Y_2 = self.maxpool(Y_1).permute(0, 2, 1)\n",
    "        \n",
    "        Y_3, _ = self.lstm(Y_2)\n",
    "        Y_4 = Y_3[:, -1]\n",
    "        X_1 = self.relu(self.conv(X))\n",
    "        # LSTM is expecting input of shape (batches, seq_len, conv_filters)\n",
    "        X_2 = self.maxpool(X_1).permute(0, 2, 1)\n",
    "        X_3, _ = self.lstm(X_2)\n",
    "        X_4 = X_3[:, -1]  # only need final output of LSTM\n",
    "        XY_4 = torch.cat((X_4, Y_4))\n",
    "        # pass in 2 xs to this function\n",
    "        # concat the two xs at X_4, and then return 1 y\n",
    "        \n",
    "        \n",
    "        X_5 = self.relu(self.fc1(XY_4))\n",
    "        X_6 = self.dropout(X_5)\n",
    "        X_7 = self.sigmoid(self.fc2(X_6))\n",
    "        y = self.sigmoid(self.fc_final(X_7)).squeeze()\n",
    "        return y\n",
    "    \n",
    "    def validation(self, data_loader):\n",
    "        # only run this within torch.no_grad() context!\n",
    "        losses = []\n",
    "        preds = []\n",
    "        labels = []\n",
    "        for seqs_onehot_batch, labels_batch in data_loader:\n",
    "            # push batch through model, get predictions, calculate loss\n",
    "            preds_batch = self(seqs_onehot_batch.squeeze().cuda())\n",
    "            labels_batch = labels_batch.squeeze()\n",
    "            loss_batch = self.loss(preds_batch, labels_batch.cuda())\n",
    "            losses.append(loss_batch.item())\n",
    "\n",
    "            # storing labels + preds for auPRC calculation later\n",
    "            labels.extend(labels_batch.detach().numpy())  \n",
    "            preds.extend(preds_batch.cpu().detach().numpy())\n",
    "            \n",
    "        return np.array(losses), np.array(preds), np.array(labels)\n",
    "\n",
    "\n",
    "    def fit(self, train_gen, source_val_data_loader, target_val_data_loader,\n",
    "            optimizer, epochs=15):\n",
    "        \n",
    "        for epoch in range(epochs):\n",
    "            torch.cuda.empty_cache()  # clear memory to keep stuff from blocking up\n",
    "            \n",
    "            print(\"=== Epoch \" + str(epoch + 1) + \" ===\")\n",
    "            print(\"Training...\")\n",
    "            self.train()\n",
    "            \n",
    "            # using a batch size of 1 here because the generator returns\n",
    "            # many examples in each batch\n",
    "            train_data_loader = DataLoader(train_gen,\n",
    "                               batch_size = 1, shuffle = True)\n",
    "\n",
    "            train_losses = []\n",
    "            train_preds = []\n",
    "            train_labels = []\n",
    "            for seqs_onehot_batch, labels_batch in train_data_loader:\n",
    "                # reset the optimizer; need to do each batch after weight update\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # push batch through model, get predictions, and calculate loss\n",
    "                preds = self(seqs_onehot_batch.squeeze().cuda())\n",
    "                labels_batch = labels_batch.squeeze()\n",
    "                loss_batch = self.loss(preds, labels_batch.cuda())\n",
    "\n",
    "                # brackpropogate the loss and update model weights accordingly\n",
    "                loss_batch.backward()\n",
    "                optimizer.step()\n",
    "                \n",
    "                train_losses.append(loss_batch.item())\n",
    "                train_labels.extend(labels_batch)\n",
    "                train_preds.extend(preds.cpu().detach().numpy())\n",
    "\n",
    "            self.train_loss_by_epoch.append(np.mean(train_losses))\n",
    "            print_metrics(train_preds, train_labels)\n",
    "            \n",
    "            # load new set of negative examples for next epoch\n",
    "            train_gen.on_epoch_end()\n",
    "\n",
    "            \n",
    "            # Assess model performance on same-species validation set\n",
    "            print(\"Evaluating on source validation data...\")\n",
    "            \n",
    "            # Since we don't use gradients during model evaluation,\n",
    "            # the following two lines let the model predict for many examples\n",
    "            # more efficiently (without having to keep track of gradients)\n",
    "            self.eval()\n",
    "            with torch.no_grad():\n",
    "                source_val_losses, source_val_preds, source_val_labels = self.validation(source_val_data_loader)\n",
    "\n",
    "                print(\"Validation loss:\", np.mean(source_val_losses))\n",
    "                self.source_val_loss_by_epoch.append(np.mean(source_val_losses))\n",
    "\n",
    "                # calc auPRC over source validation set\n",
    "                source_val_auprc = print_metrics(source_val_preds, source_val_labels)\n",
    "                self.source_val_auprc_by_epoch.append(source_val_auprc)\n",
    "\n",
    "                # check if this is the best performance we've seen so far\n",
    "                # if yes, save the model weights -- we'll use the best model overall\n",
    "                # for later analyses\n",
    "                if source_val_auprc < self.best_auprc_so_far:\n",
    "                    self.best_auprc_so_far = source_val_auprc\n",
    "                    self.best_state_so_far = self.state_dict()\n",
    "                \n",
    "                \n",
    "                # now repeat for target species data \n",
    "                print(\"Evaluating on target validation data...\")\n",
    "                \n",
    "                target_val_losses, target_val_preds, target_val_labels = self.validation(target_val_data_loader)\n",
    "\n",
    "                print(\"Validation loss:\", np.mean(target_val_losses))\n",
    "                self.target_val_loss_by_epoch.append(np.mean(target_val_losses))\n",
    "\n",
    "                # calc auPRC over source validation set\n",
    "                target_val_auprc = print_metrics(target_val_preds, target_val_labels)\n",
    "                self.target_val_auprc_by_epoch.append(target_val_auprc)\n",
    "                "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfcb58d5",
   "metadata": {},
   "source": [
    "# Setup + Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3c2e68d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1365\n"
     ]
    }
   ],
   "source": [
    "# setup generators / data loaders for training and validation\n",
    "\n",
    "# we'll make the training data loader in the training loop,\n",
    "# since we need to update some of the examples used each epoch\n",
    "train_gen = TrainGenerator(\"mouse\", \"CTCF\")\n",
    "\n",
    "source_val_gen = ValGenerator(\"mouse\", \"CTCF\")\n",
    "# using a batch size of 1 here because the generator returns\n",
    "# many examples in each batch\n",
    "source_val_data_loader = DataLoader(source_val_gen, batch_size = 1, shuffle = False)\n",
    "\n",
    "target_val_gen = ValGenerator(\"human\", \"CTCF\")\n",
    "target_val_data_loader = DataLoader(target_val_gen, batch_size = 1, shuffle = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "64b180f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beginning Init\n",
      "inside profiles\n",
      "152461100  is chrom.  152459793  is start.  152461907  is end.\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "Invalid interval bounds!",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_21873/3749853005.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# initialize the model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mBasicModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0moptimizer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAdam\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.001\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;31m# train!\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_21873/3956815000.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     15\u001b[0m         \u001b[0;31m#How do we know which species of data we should use???\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m         \u001b[0;31m#Our added on data.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maccess_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mGenerator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_species\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"val\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m         \u001b[0;31m# Defining the layers to go into our model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_21873/4022516070.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, species, tf, train_val_test, seq_len, profile_len, return_labels)\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcoords\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_coords\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;31m#self.seqs_onehot = self.convert(self.coords)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiles\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlogcounts_sum\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlogcounts_mean\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlogcounts_max\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_profiles_and_logcounts\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcoords\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_21873/4022516070.py\u001b[0m in \u001b[0;36mget_profiles_and_logcounts\u001b[0;34m(self, coords)\u001b[0m\n\u001b[1;32m     73\u001b[0m                                                  self.prof_len + 2 * self.max_jitter)\n\u001b[1;32m     74\u001b[0m                 pos_profile = np.array(\n\u001b[0;32m---> 75\u001b[0;31m                     accessibility_reader.values(chrom, prof_start, prof_end))\n\u001b[0m\u001b[1;32m     76\u001b[0m                 \u001b[0mpos_profile\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misnan\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpos_profile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;31m# Should this be epsilon? - JC\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     77\u001b[0m                 \u001b[0mprofile\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpos_profile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Invalid interval bounds!"
     ]
    }
   ],
   "source": [
    "# initialize the model\n",
    "model = BasicModel()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# train!\n",
    "model.cuda()\n",
    "model.fit(train_gen, source_val_data_loader, target_val_data_loader, optimizer, epochs = 1)\n",
    "model.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c647a90",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb9df27",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "575732fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e992bc40",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:cs197-env] *",
   "language": "python",
   "name": "conda-env-cs197-env-py"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}