SpeCollate-Search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### This Notebook is for database search using peptide strings and experimental spectra as input to the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from pandas import ExcelWriter\n",
    "from pandas import ExcelFile\n",
    "import numpy as np\n",
    "import random as rand\n",
    "import math\n",
    "import queue\n",
    "import csv\n",
    "from collections import OrderedDict\n",
    "from IPython.display import clear_output\n",
    "import csv\n",
    "import re\n",
    "from heapq import merge\n",
    "from sklearn import preprocessing\n",
    "from operator import itemgetter\n",
    "import bisect\n",
    "import pickle\n",
    "import os\n",
    "import gc\n",
    "from configparser import ConfigParser\n",
    "import ast\n",
    "import sys\n",
    "from pathlib import Path\n",
    "import pickle\n",
    "import shutil\n",
    "import progressbar\n",
    "\n",
    "from sklearn.datasets import make_circles, make_moons\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.distributed as dist\n",
    "import torch.multiprocessing as mp\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "from IPython.core.debugger import set_trace\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torch.autograd import Variable\n",
    "from IPython.display import display\n",
    "from os import listdir\n",
    "from os.path import isfile, join\n",
    "\n",
    "from src.snaputils import simulatespectra as sim\n",
    "from src.snaptrain import process\n",
    "from src.snaputils import reader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Temporary config func. Original one in the project.\n",
    "class config:\n",
    "    \"\"\"Define constants\"\"\"\n",
    "    # AAMass = OrderedDict([('A', 71.037114), ('C', 103.009185), ('D', 115.026943), ('E', 129.042593),\n",
    "    #                       ('F', 147.068414), ('G', 57.021464), ('H', 137.058912), ('I', 113.084064),\n",
    "    #                       ('K', 128.094963), ('L', 113.084064), ('M', 131.040485), ('N', 114.042927),\n",
    "    #                       ('P', 97.052764), ('Q', 128.058578), ('R', 156.101111), ('S', 87.032028),\n",
    "    #                       ('T', 101.047679), ('V', 99.068414), ('W', 186.079313), ('Y', 163.0633),\n",
    "    #                       ('p', 79.97), ('c', 57.02), ('o', 15.99)\n",
    "    #                       #('r', -17.03), ('y', 43.01), ('d', -18.01), ('t', 26.02),\n",
    "    #                       #('h', 0.98), ('a', 42.01)\n",
    "    #                       ])\n",
    "\n",
    "    AAMass = OrderedDict([('A', 71.037114), ('C', 103.009185), ('D', 115.026943), ('E', 129.042593),\n",
    "                          ('F', 147.068414), ('G', 57.021464), ('H', 137.058912), ('I', 113.084064),\n",
    "                          ('K', 128.094963), ('L', 113.084064), ('M', 131.040485), ('N', 114.042927),\n",
    "                          ('P', 97.052764), ('Q', 128.058578), ('R', 156.101111), ('S', 87.032028),\n",
    "                          ('T', 101.047679), ('V', 99.068414), ('W', 186.079313), ('Y', 163.0633),\n",
    "                          ('p', 79.9663), ('o', 15.994915)\n",
    "                          ,('h', 0.9840), ('c', 57.02146), ('a', 42.0106),\n",
    "                          ('r', -17.026549), ('y', 43.00581), ('d', -18.010565), ('t', 26.02)])\n",
    "\n",
    "    ModMass = {\"Oxidation\": 15.994915, \"CAM\": 57.02146, \"Carbamidomethyl\": 57.02146, \"ICAT_light\": 227.12,\n",
    "               \"ICAT_heavy\": 236.12, \"AB_old_ICATd0\": 442.20, \"AB_old_ICATd8\": 450.20, \"Acetyl\": 42.0106,\n",
    "               \"Deamidation\": 0.9840, \"Pyro-cmC\": -17.026549, \"Pyro-glu\": -17.026549, \"Pyro_glu\": -18.010565,\n",
    "               \"Amide\": -0.984016, \"Phospho\": 79.9663, \"Methyl\": 14.0157, \"Carbamyl\": 43.00581}\n",
    "\n",
    "    ModCHAR = OrderedDict([(\"15.99\", \"o\"), (\"0.98\", \"h\"), (\"57.02\", \"c\"), (\"42.01\", \"a\"), (\"-17.03\", \"r\"),\n",
    "                           (\"79.97\", \"p\"), (\"43.01\", \"y\"), (\"-18.01\", \"d\"), (\"26.02\", \"t\")])\n",
    "    # ModCHAR = {\"15.99\": \"o\", \"0.98\": \"h\", \"57.02\": \"c\", \"42.01\": \"a\", \"-17.03\": \"r\", \"79.97\": \"p\"}\n",
    "    Ignore = [\"U\", \"X\"]\n",
    "    Mods = [{\"mod_char\": \"p\", \"aas\": [\"S\", \"T\", \"Y\"]}\n",
    "            # {\"mod_char\": \"o\", \"aas\": [\"nt\", \"M\"]}\n",
    "            ]\n",
    "    H2O = 18.010564683\n",
    "    NH3 = 17.031\n",
    "    PROTON = 1.00727647\n",
    "    \n",
    "    DEFAULT_PARAM_PATH = os.path.join(os.getcwd(), 'config.ini')\n",
    "    PARAM_PATH = None\n",
    "    l_config = None\n",
    "\n",
    "\n",
    "    def get_config(section='input', key=None):\n",
    "        \"\"\"Read the configuration parameters and return a dictionary.\"\"\"\n",
    "\n",
    "        # If file path is given use it otherwise use default.\n",
    "        file_path = config.PARAM_PATH if config.PARAM_PATH else config.DEFAULT_PARAM_PATH\n",
    "\n",
    "        # Read config and convert each value to appropriate type.\n",
    "        # Only for the first time.\n",
    "        if not config.l_config:\n",
    "            config.l_config = dict()\n",
    "            config_ = ConfigParser()\n",
    "            assert isinstance(file_path, str)\n",
    "            config_.read(file_path)\n",
    "            for section_ in config_.sections():\n",
    "                config.l_config[section_] = dict()\n",
    "                for key_ in config_[section_]:\n",
    "                    try:\n",
    "                        config.l_config[section_][key_] = ast.literal_eval(config_[section_][key_])\n",
    "                    except (ValueError, SyntaxError):\n",
    "                        config.l_config[section_][key_] = config_[section_][key_]\n",
    "\n",
    "        if section and section in config.l_config:\n",
    "            if key and key in config.l_config[section]:\n",
    "                return config.l_config[section][key]\n",
    "            return config.l_config[section]\n",
    "        return config.l_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "True\n",
      "1024\n",
      "80000\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "print(config.get_config(section='input', key='charge'))\n",
    "print(config.get_config(section='input', key='use_mods'))\n",
    "print(config.get_config(section='ml', key='batch_size'))\n",
    "print(config.get_config(section='input', key='spec_size'))\n",
    "print(config.get_config(section='search', key='num_mods'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, vocab_size, output_size=512, embedding_dim=512, hidden_lstm_dim=1024, lstm_layers=2):\n",
    "        super(Net, self).__init__()\n",
    "        self.spec_size = config.get_config(section='input', key='spec_size')\n",
    "        self.seq_len = config.get_config(section='ml', key='pep_seq_len')\n",
    "        self.output_size = output_size\n",
    "        self.lstm_layers = lstm_layers\n",
    "        self.hidden_lstm_dim = hidden_lstm_dim\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.vocab_size = vocab_size\n",
    "        \n",
    "        ################### Spectra branch ###################\n",
    "        # self.conv1_1    = nn.Conv1d(1, 32, 3, stride=1, padding=1)\n",
    "        # self.maxpool1_1 = nn.MaxPool1d(50, stride=50)\n",
    "        # self.linear1_1  = nn.Linear(int((self.spec_size*32)/50), 1024) # use with convo layers\n",
    "        self.linear1_1 = nn.Linear(self.spec_size, 512)\n",
    "        self.linear1_2 = nn.Linear(512, 256)\n",
    "        # Spectra charge sub-branch\n",
    "        # self.charge = config.get_config(section=\"input\", key=\"charge\")\n",
    "        # self.linear_charge1 = nn.Linear(1024, self.charge)\n",
    "\n",
    "        ################### Peptide branch ###################\n",
    "        self.embedding = nn.Embedding(vocab_size, embedding_dim)\n",
    "        self.lstm = nn.LSTM(embedding_dim, self.hidden_lstm_dim, self.lstm_layers,\n",
    "                            # dropout=0.5,\n",
    "                            batch_first=True, bidirectional=True)\n",
    "\n",
    "        self.linear2_1 = nn.Linear(self.hidden_lstm_dim * 2, 512) # 2048, 1024\n",
    "        self.linear2_2 = nn.Linear(512, 256)\n",
    "        # NOTE: should the two layers below be just one layer.\n",
    "        # self.linear_spec_mass1 = nn.Linear(1024, 256)\n",
    "        # self.linear_spec_mass2 = nn.Linear(256, 1)\n",
    "        # self.linear_spec_mass3 = nn.Linear(128, 1)\n",
    "        # self.linear_pep_mass1 = nn.Linear(1024, 256)\n",
    "        # self.linear_pep_mass2 = nn.Linear(256, 1)\n",
    "        # self.linear_pep_mass3 = nn.Linear(128, 1)\n",
    "        do = config.get_config(section=\"ml\", key=\"dropout\")\n",
    "        self.dropout_conv1_1 = nn.Dropout(do)\n",
    "        self.dropout1_1 = nn.Dropout(do)\n",
    "        self.dropout1_2 = nn.Dropout(do)\n",
    "        self.dropout1_3 = nn.Dropout(do)\n",
    "        self.dropout1_4 = nn.Dropout(do)\n",
    "        self.dropout1_5 = nn.Dropout(do)\n",
    "        self.dropout2_1 = nn.Dropout(do)\n",
    "        self.dropout2_2 = nn.Dropout(do)\n",
    "        self.dropout2_3 = nn.Dropout(do)\n",
    "        self.dropout2_4 = nn.Dropout(do)\n",
    "        self.dropout2_5 = nn.Dropout(do)\n",
    "        self.dropout2_6 = nn.Dropout(do)\n",
    "        self.dropout_charge1 = nn.Dropout(do)\n",
    "        self.dropout_spec_mass1 = nn.Dropout(do)\n",
    "        self.dropout_spec_mass2 = nn.Dropout(do)\n",
    "        self.dropout_spec_mass3 = nn.Dropout(do)\n",
    "        self.dropout_pep_mass1 = nn.Dropout(do)\n",
    "        self.dropout_pep_mass2 = nn.Dropout(do)\n",
    "        self.dropout_pep_mass3 = nn.Dropout(do)\n",
    "        print(\"dropout: {}\".format(do))\n",
    "        #self.dropout3 = nn.Dropout(0.3)\n",
    "        \n",
    "    def forward(self, data, data_type=None):\n",
    "        assert not data_type or data_type == \"specs\" or data_type == \"peps\"\n",
    "        res = []\n",
    "        if not data_type or data_type == \"specs\":\n",
    "            \n",
    "            specs = data[0]\n",
    "            # out = F.relu(self.conv1_1(specs.view(-1, 1, self.spec_size)))\n",
    "            # out = self.maxpool1_1(out)\n",
    "            # out = self.dropout_conv1_1(out)\n",
    "\n",
    "            # out = F.relu(self.linear1_1(out.view(out.size(0), -1))) # use this line with conv networks\n",
    "            out = F.relu(self.linear1_1(specs.view(-1, self.spec_size)))\n",
    "            out_spec = self.dropout1_1(out)\n",
    "            \n",
    "            out_spec = F.relu(self.linear1_2(out_spec))\n",
    "            out_spec = F.normalize(out_spec)\n",
    "            \n",
    "            # out_ch = self.dropout_charge1(out)\n",
    "            # out_ch = self.linear_charge1(out_ch)\n",
    "            # out_ch = F.relu(out_ch)\n",
    "            \n",
    "            # out_mass = self.dropout_spec_mass1(out)\n",
    "            # out_mass = self.linear_spec_mass1(out_mass)\n",
    "            # out_mass = F.relu(out_mass)\n",
    "            # out_mass = self.dropout_spec_mass2(out_mass)\n",
    "            # out_mass = self.linear_spec_mass2(out_mass)\n",
    "            # out_mass = F.relu(out_mass)\n",
    "            # out_mass = self.dropout_spec_mass3(out_mass)\n",
    "            # out_mass = self.linear_spec_mass3(out_mass)\n",
    "            res.append(out_spec)\n",
    "            # res.append(out_ch)\n",
    "            # res.append(out_mass)\n",
    "        if not data_type or data_type == \"peps\":\n",
    "            for peps in data[1:2]:\n",
    "                embeds = self.embedding(peps)\n",
    "                # embeds = self.one_hot_tensor(peps)\n",
    "                hidden = self.init_hidden(len(peps))\n",
    "                hidden = tuple([e.data for e in hidden])\n",
    "                lstm_out, _ = self.lstm(embeds, hidden)\n",
    "                lstm_out = lstm_out[:, -1, :]\n",
    "                out = lstm_out.contiguous().view(-1, self.hidden_lstm_dim * 2)\n",
    "                \n",
    "                out = self.dropout2_1(out)\n",
    "                out = F.relu((self.linear2_1(out)))\n",
    "                \n",
    "                out_pep = self.dropout2_2(out)\n",
    "                out_pep = F.relu(self.linear2_2(out_pep))\n",
    "                out_pep = F.normalize(out_pep)\n",
    "                \n",
    "                # out_mass = self.dropout_pep_mass1(out)\n",
    "                # out_mass = self.linear_pep_mass1(out_mass)\n",
    "                # out_mass = F.relu(out_mass)\n",
    "                # out_mass = self.dropout_pep_mass2(out_mass)\n",
    "                # out_mass = self.linear_pep_mass2(out_mass)\n",
    "                # out_mass = F.relu(out_mass)\n",
    "                # out_mass = self.dropout_pep_mass3(out_mass)\n",
    "                # out_mass = self.linear_pep_mass3(out_mass)\n",
    "                res.append(out_pep)\n",
    "                # res.append(out_mass)\n",
    "        return res\n",
    "    \n",
    "    def init_hidden(self, batch_size):\n",
    "        weight = next(self.parameters()).data\n",
    "        hidden = (weight.new(self.lstm_layers * 2, batch_size, self.hidden_lstm_dim).zero_(),\n",
    "                      weight.new(self.lstm_layers * 2, batch_size, self.hidden_lstm_dim).zero_())\n",
    "        return hidden\n",
    "\n",
    "    def one_hot_tensor(self, peps):\n",
    "        batch_size = len(peps)\n",
    "        src = torch.zeros((batch_size, self.seq_len), dtype=torch.float16, device='cuda')\n",
    "        src[peps > 0] = 1.0\n",
    "        one_hots = torch.zeros((batch_size, self.seq_len, self.vocab_size), dtype=torch.float16, \n",
    "                                device='cuda')\n",
    "        one_hots.scatter_(2, peps.view(batch_size, self.seq_len, 1), src.view(batch_size, self.seq_len, 1))\n",
    "        one_hots.requires_grad = True\n",
    "        return one_hots\n",
    "    \n",
    "    def name(self):\n",
    "        return \"Net\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_out_dir(dir_path, exist_ok=True):\n",
    "    out_path = Path(dir_path)\n",
    "    if out_path.exists() and out_path.is_dir():\n",
    "        if not exist_ok:\n",
    "            shutil.rmtree(out_path)\n",
    "            out_path.mkdir()\n",
    "    else:\n",
    "        out_path.mkdir()\n",
    "        \n",
    "    Path(join(out_path, 'spectra')).mkdir()\n",
    "    Path(join(out_path, 'peptides')).mkdir()\n",
    "    \n",
    "def verify_in_dir(dir_path, ext, ignore_list=[]):\n",
    "    in_path = Path(dir_path)\n",
    "    assert in_path.exists() and in_path.is_dir()\n",
    "    \n",
    "    files = [join(dir_path, f) for f in listdir(dir_path) if\n",
    "                 isfile(join(dir_path, f)) and not f.startswith('.') \n",
    "                 and f.split('.')[-1] == ext and f not in ignore_list]\n",
    "    assert len(files) > 0\n",
    "    return files\n",
    "\n",
    "def isfloat(str_float):\n",
    "    try:\n",
    "        float(str_float)\n",
    "        return True\n",
    "    except ValueError: \n",
    "        return False\n",
    "    \n",
    "def gray_code(num):\n",
    "    return num ^ (num >> 1)\n",
    "\n",
    "def decimal_to_binary_array(num, arr_len):\n",
    "    bin_arr = [float(i) for i in list('{0:0b}'.format(num))]\n",
    "    assert len(bin_arr) <= arr_len\n",
    "    res = [0.] * (arr_len - len(bin_arr)) + bin_arr\n",
    "    inds = [int(i) for i, _ in enumerate(res) if res[i] > 0.1] # greater than zero. 0.1 for the floating pointing errors.\n",
    "    vals = [1.0] * len(inds)\n",
    "    return inds, vals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mod_repl(match):\n",
    "    lookup = str(round(float(match.group(0)), 2))\n",
    "    return config.ModCHAR[lookup] if lookup in config.ModCHAR else \"\"\n",
    "\n",
    "def preprocess_mgfs(mgf_dir, out_dir):\n",
    "    \n",
    "    mgf_files = verify_in_dir(mgf_dir, \"mgf\")\n",
    "    create_out_dir(out_dir, exist_ok=False)\n",
    "        \n",
    "    print('reading {} files'.format(len(mgf_files)))\n",
    "    \n",
    "    spec_size = config.get_config(section='input', key='spec_size')\n",
    "    charge = config.get_config(section='input', key='charge')\n",
    "    use_mods = config.get_config(section='input', key='use_mods')\n",
    "    num_species = config.get_config(section='input', key='num_species')\n",
    "    seq_len = config.get_config(section='ml', key='pep_seq_len')\n",
    "    \n",
    "    ch = np.zeros(20)\n",
    "    modified = 0\n",
    "    unmodified = 0\n",
    "    unique_pep_set = set()\n",
    "    \n",
    "    pep_dict = {}\n",
    "    idx_spec_map = []\n",
    "    pep_spec = []\n",
    "    pep_idx = 0\n",
    "    \n",
    "    summ = np.zeros(spec_size)\n",
    "    sq_sum = np.zeros(spec_size)\n",
    "    N = 0\n",
    "    \n",
    "    tot_count = 0\n",
    "    max_peaks = max_moz = 0\n",
    "    for species_id, mgf_file in enumerate(mgf_files):\n",
    "        print('Reading: {}'.format(mgf_file))\n",
    "        \n",
    "        f = open(mgf_file, \"r\")\n",
    "        lines = f.readlines()\n",
    "        f.close()\n",
    "        \n",
    "        count = lcount = 0\n",
    "        \n",
    "        pep_list = []\n",
    "        dataset = []\n",
    "        label = []\n",
    "        \n",
    "        mass_ign = 0\n",
    "        pep_len_ign = 0\n",
    "        dup_ign = 0\n",
    "\n",
    "        print('len of file: ' + str(len(lines)))\n",
    "        limit = 200000\n",
    "        pep = []\n",
    "        spec = []\n",
    "        pep_set = set()\n",
    "        is_name = is_mw = is_charge = False\n",
    "        prev = 0\n",
    "        i = 0\n",
    "        while i < len(lines) and limit > 0:\n",
    "            line = lines[i]\n",
    "            i += 1\n",
    "\n",
    "            if line.startswith('PEPMASS'):\n",
    "                count += 1\n",
    "                mass = float(re.findall(r\"PEPMASS=([-+]?[0-9]*\\.?[0-9]*)\", line)[0])\n",
    "                is_mw = True\n",
    "            \n",
    "            if is_mw and line.startswith('CHARGE'):\n",
    "                l_charge = int(re.findall(r\"CHARGE=([-+]?[0-9]*\\.?[0-9]*)\", line)[0])\n",
    "                is_charge = True\n",
    "                mass = (mass - config.PROTON) * l_charge\n",
    "                \n",
    "            if is_mw and is_charge:\n",
    "\n",
    "                while not isfloat(re.split(' |\\t|=', lines[i])[0]):\n",
    "                    i += 1\n",
    "                    \n",
    "                spec_ind = []\n",
    "                spec_val = []\n",
    "                num_peaks = 0\n",
    "                while 'END IONS' not in lines[i].upper():\n",
    "                    if lines[i] == '\\n':\n",
    "                        i += 1\n",
    "                        continue\n",
    "                    mz_line = lines[i]\n",
    "                    i += 1\n",
    "                    num_peaks += 1\n",
    "                    mz_splits = re.split(' |\\t', mz_line)\n",
    "                    moz = round(float(mz_splits[0]) * 10) # + 32 # 32 because charge is len 8 and mass is len 24\n",
    "                    intensity = math.sqrt(float(mz_splits[1]) + 1.0) # adding 1 to avoid sqrt of zero\n",
    "                    if moz > max_moz:\n",
    "                        max_moz = moz\n",
    "                    if 0 < moz < spec_size:\n",
    "                        # spec[round(moz*10)] += round(intensity)\n",
    "                        if spec_ind and spec_ind[-1] == moz:\n",
    "                            spec_val[-1] = max(intensity, spec_val[-1])\n",
    "                        else:\n",
    "                            spec_ind.append(moz)\n",
    "                            spec_val.append(intensity) # adding one to avoid sqrt of zero\n",
    "                if num_peaks < 15:\n",
    "                    is_name = is_mw = is_charge = False\n",
    "                    continue\n",
    "                    \n",
    "                spec_ind = np.array(spec_ind)\n",
    "                spec_val = np.array(spec_val)\n",
    "                spec_val = (spec_val - np.amin(spec_val)) / (np.amax(spec_val) - np.amin(spec_val))\n",
    "                \n",
    "                gray_mass = gray_code(int(mass * 100))\n",
    "                mass_ind, mass_val = decimal_to_binary_array(gray_mass, 24) # mass is gray encoded into array of length 24\n",
    "                ch_ind,   ch_val   = [l_charge - 1], [1.0]\n",
    "#                 for ch_val in range(l_charge):\n",
    "#                     ind.append(ch_val+24)\n",
    "#                     val.append(1)\n",
    "                # Join mass, charge, and spec indicies and values.\n",
    "#                 ind = np.concatenate((np.array(mass_ind), np.array(ch_ind) + 24, spec_ind + 32))\n",
    "#                 val = np.concatenate((np.array(mass_val), np.array(ch_val), spec_val))\n",
    "\n",
    "                ind = spec_ind\n",
    "                val = spec_val\n",
    "            \n",
    "                assert len(ind) == len(val)\n",
    "                spec = np.array([ind, val])\n",
    "                \n",
    "                summ[ind] += val\n",
    "                sq_sum[ind] += val**2\n",
    "                N += 1\n",
    "\n",
    "                is_name = True\n",
    "\n",
    "            if is_name and is_mw and is_charge:\n",
    "                is_name = is_mw = is_charge = False\n",
    "\n",
    "                \"\"\"output the data to \"\"\"\n",
    "                spec_file_name = '{}-{}-{}.npy'.format(lcount, mass, l_charge)\n",
    "                np.save(join(out_dir, 'spectra', spec_file_name), spec)\n",
    "\n",
    "                lcount += 1\n",
    "                tot_count += 1\n",
    "                \n",
    "                pep = 0\n",
    "                spec = []\n",
    "                new = int((i / len(lines)) * 100)\n",
    "                if new >= prev + 10:\n",
    "                    #clear_output(wait=True)\n",
    "                    print('count: ' + str(lcount))\n",
    "                    print(str(new) + '%')\n",
    "                    prev = new\n",
    "\n",
    "        #print('max peaks: ' + str(max_peaks))\n",
    "        print('In current file, read {} out of {}'.format(lcount, count))\n",
    "        print(\"Ignored: large mass: {}, pep len: {}, dup: {}\".format(mass_ign, pep_len_ign, dup_ign))\n",
    "        print('overall running count: ' + str(tot_count))\n",
    "        print('max moz: ' + str(max_moz))\n",
    "#         return pep_list, dataset, label\n",
    "#         tmp_pep_list, tmp_dataset, tmp_labels = read_msp(msp_file, species_id, decoy)\n",
    "#         pep_list.extend(tmp_dataset)\n",
    "#         dataset.extend(tmp_dataset)\n",
    "#         label.extend(tmp_labels)\n",
    "\n",
    "    # save the map. this will be used to generate masks for hard positive/negative mining during training.\n",
    "    # np.save(join(out_dir, \"idx_spec_map.npy\"), idx_spec_map)\n",
    "    # with open(join(out_dir, 'pep_spec.pkl'), 'wb') as f:\n",
    "    #     pickle.dump(pep_spec, f)\n",
    "    \n",
    "    print(\"Statistics:\")\n",
    "    print(\"Charge distribution:\")\n",
    "    print(ch)\n",
    "    print(\"Modified:\\t{}\".format(modified))\n",
    "    print(\"Unmodified:\\t{}\".format(unmodified))\n",
    "    print(\"Unique Peptides:\\t{}\".format(len(unique_pep_set)))\n",
    "    print(\"Sum: {}\".format(summ))\n",
    "    print(\"Sum-Squared: {}\".format(sq_sum))\n",
    "    print(\"N: {}\".format(N))\n",
    "    means = summ / N\n",
    "    print(\"mean: {}\".format(means))\n",
    "    stds = np.sqrt((sq_sum / N) - means**2)\n",
    "    stds[stds < 0.0000001] = float(\"inf\")\n",
    "    print(\"std: {}\".format(stds))\n",
    "    np.save(join(out_dir, 'means.npy'), means)\n",
    "    np.save(join(out_dir, 'stds.npy'), stds)\n",
    "\n",
    "# return spectra, masses, charges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading 1 files\n",
      "Reading: /disk/raptor-2/mtari008/data/deepsnap/pxd001468/mgfs/b1906_293T_proteinID_01A_QE3_122212.mgf\n",
      "len of file: 3402889\n",
      "count: 5031\n",
      "10%\n",
      "count: 9900\n",
      "20%\n",
      "count: 14400\n",
      "30%\n",
      "count: 18501\n",
      "40%\n",
      "count: 22416\n",
      "50%\n",
      "count: 26270\n",
      "60%\n",
      "count: 29988\n",
      "70%\n",
      "count: 33771\n",
      "80%\n",
      "count: 37683\n",
      "90%\n",
      "In current file, read 41799 out of 41820\n",
      "Ignored: large mass: 0, pep len: 0, dup: 0\n",
      "overall running count: 41799\n",
      "max moz: 59520\n",
      "Statistics:\n",
      "Charge distribution:\n",
      "[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
      "Modified:\t0\n",
      "Unmodified:\t0\n",
      "Unique Peptides:\t0\n",
      "Sum: [0. 0. 0. ... 0. 0. 0.]\n",
      "Sum-Squared: [0. 0. 0. ... 0. 0. 0.]\n",
      "N: 41799\n",
      "mean: [0. 0. 0. ... 0. 0. 0.]\n",
      "std: [inf inf inf ... inf inf inf]\n"
     ]
    }
   ],
   "source": [
    "mgf_dir = \"/disk/raptor-2/mtari008/data/deepsnap/pxd001468/mgfs\"\n",
    "# mgf_dir = \"/disk/raptor-2/mtari008/data/deepsnap/pxd000612/human-hcd-phospho-mgf\"\n",
    "# mgf_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed-human-hcd-tryp-best/mgfs\"\n",
    "# out_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed-human-hcd-tryp-best/pts\"\n",
    "out_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed/pxd001468-122212-nochmass\"\n",
    "# out_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed/one\"\n",
    "preprocess_mgfs(mgf_dir, out_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_occurrences(s, ch):\n",
    "    return [i for i, letter in enumerate(s) if letter == ch]\n",
    "\n",
    "def apply_mod(peps, mod):\n",
    "    result_peps = set()\n",
    "    pep_seq_len = config.get_config(section=\"ml\", key=\"pep_seq_len\")\n",
    "    for pep in peps:\n",
    "        if len(pep) >= pep_seq_len:\n",
    "            continue\n",
    "        for mod_aa in mod[\"aas\"]:\n",
    "            if mod_aa == \"nt\" and not pep[1].islower():\n",
    "                result_peps.update([pep[0] + mod[\"mod_char\"] + pep[1:]])\n",
    "            elif mod_aa == \"ct\" and not pep[-1].islower():\n",
    "                result_peps.update([pep + mod[\"mod_char\"]])\n",
    "            else:\n",
    "                aa_indices = find_occurrences(pep, mod_aa)\n",
    "                for index in aa_indices:\n",
    "                    if index == len(pep) - 1 or not pep[index+1].islower():\n",
    "                        result_peps.update([pep[:index+1] + mod[\"mod_char\"] + (pep[index+1:] if index < len(pep) else \"\")])\n",
    "        \n",
    "    return result_peps\n",
    "\n",
    "def add_mods(pep, mods, num_mods):\n",
    "    mod_peps = set([pep])\n",
    "    result_peps = set([pep])\n",
    "    for i in range(num_mods):\n",
    "        temp_mod_peps = set()\n",
    "        for mod in mods:\n",
    "            temp_mod_peps.update(apply_mod(mod_peps, mod))\n",
    "        mod_peps.update(set(temp_mod_peps))\n",
    "        result_peps.update(mod_peps)\n",
    "    \n",
    "    return result_peps\n",
    "\n",
    "def load_peps(pep_dir):\n",
    "    fasta_files = verify_in_dir(pep_dir, \"fasta\")\n",
    "    \n",
    "    use_mods = config.get_config(key=\"use_mods\", section=\"input\")\n",
    "    mods_list = config.Mods\n",
    "    num_mods = config.get_config(key=\"num_mods\", section=\"search\")\n",
    "    pep_seq_len = config.get_config(key=\"pep_seq_len\", section=\"ml\")\n",
    "    \n",
    "    pep_set = set()\n",
    "    pep_list = []\n",
    "    masses = []\n",
    "    modifieds = []\n",
    "    prot_list = []\n",
    "    \n",
    "    tot_pep_count = 0\n",
    "    for fasta_file in fasta_files:\n",
    "        print('Reading: {}'.format(fasta_file))\n",
    "        \n",
    "        f = open(fasta_file, \"r\")\n",
    "        lines = f.readlines()\n",
    "        f.close()\n",
    "        \n",
    "        print(\"File length: {}\".format(len(lines)))\n",
    "        peps = []\n",
    "        temp_prot = \"\"\n",
    "        with progressbar.ProgressBar(max_value=len(lines)) as bar:\n",
    "            for i, line in enumerate(lines):\n",
    "                line = line.strip().replace(\"C\", \"Cc\")\n",
    "                if line.startswith(\">\"):\n",
    "                    temp_prot = line[1:].strip()\n",
    "                    continue\n",
    "                elif any(x in line for x in config.Ignore):\n",
    "                    continue\n",
    "                peps = add_mods(line, mods_list, num_mods) if use_mods else [line]\n",
    "                for pep in peps:\n",
    "                    pep = pep.strip()\n",
    "                    mass = sim.get_pep_mass(pep)\n",
    "                    modified = any(aa.islower() for aa in pep)\n",
    "                    if pep not in pep_set and len(pep) <= pep_seq_len:\n",
    "                        pep_set.add(pep)\n",
    "                        pep_list.append(pep)\n",
    "                        masses.append(mass)\n",
    "                        modifieds.append(modified)\n",
    "                        prot_list.append(temp_prot)\n",
    "                        tot_pep_count += 1\n",
    "                bar.update(i)\n",
    "        print(\"Peptides written: {}\".format(tot_pep_count))\n",
    "        \n",
    "        return pep_list, prot_list, masses, modifieds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_specs(spec_dir):\n",
    "    spec_size = config.get_config(key=\"spec_size\", section=\"input\")\n",
    "    spec_files = verify_in_dir(spec_dir, \"npy\")\n",
    "    spec_ids = []\n",
    "    spec_list = []\n",
    "    masses = []\n",
    "    charges = []\n",
    "    count = 0\n",
    "    with progressbar.ProgressBar(max_value=len(spec_files)) as bar:\n",
    "        for i, spec_file in enumerate(spec_files):\n",
    "            file_name = spec_file.split('/')[-1]\n",
    "            file_parts = re.search(r\"(\\d+)-(\\d+.\\d+)-(\\d+).[pt|npy]\", file_name)\n",
    "            spec_id = int(file_parts[1])\n",
    "            mass = round(float(file_parts[2]), 2)\n",
    "            charge = int(file_parts[3])\n",
    "            if charge > 5:\n",
    "                continue\n",
    "            spec_ids.append(spec_id)\n",
    "            np_spec = np.load(spec_file)\n",
    "            spec_list.append(np_spec)\n",
    "            masses.append(mass)\n",
    "            charges.append(charge)\n",
    "\n",
    "            count += 1\n",
    "            bar.update(i)\n",
    "    print(\"count: {}\".format(count))\n",
    "    return spec_ids, spec_list, masses, charges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils import data\n",
    "class SpectralDataset(data.Dataset):\n",
    "    'Characterizes a dataset for PyTorch'\n",
    "    def __init__(self, dir_path):\n",
    "        'Initialization'\n",
    "        \n",
    "        in_path = Path(dir_path)\n",
    "        assert in_path.exists()\n",
    "        assert in_path.is_dir()\n",
    "        \n",
    "        self.spec_path = join(dir_path, \"spectra\")\n",
    "        self.spec_size = config.get_config(section='input', key='spec_size')\n",
    "        \n",
    "        self.means = torch.from_numpy(np.load(join(dir_path, \"means.npy\"))).float()\n",
    "        self.stds  = torch.from_numpy(np.load(join(dir_path, \"stds.npy\"))).float()\n",
    "        \n",
    "        spec_ids, spec_lst, spec_mass_lst, spec_charge_lst = load_specs(self.spec_path)\n",
    "        all_sorts = list(zip(*sorted(zip(spec_ids, spec_lst, spec_mass_lst, spec_charge_lst), key=lambda x: x[2])))\n",
    "        self.spec_ids         = all_sorts[0]\n",
    "        self.spec_list        = all_sorts[1]\n",
    "        self.spec_mass_list   = all_sorts[2]\n",
    "        self.spec_charge_list = all_sorts[3]\n",
    "        print('Spectral Dataset Size: {}'.format(len(self.spec_list)))\n",
    "        \n",
    "\n",
    "    def __len__(self):\n",
    "        'Denotes the total number of samples'\n",
    "        return len(self.spec_list)\n",
    "\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        'Generates one sample of data'\n",
    "\n",
    "        # Load spectra\n",
    "        np_spec = self.spec_list[index]\n",
    "        ind = torch.LongTensor([[0]*np_spec.shape[1], np_spec[0]])\n",
    "        val = torch.FloatTensor(np_spec[1])\n",
    "        torch_spec = torch.sparse_coo_tensor(\n",
    "            ind, val, torch.Size([1, self.spec_size])).to_dense()\n",
    "        torch_spec = (torch_spec - self.means) / self.stds\n",
    "\n",
    "        return torch_spec\n",
    "\n",
    "\n",
    "class PeptideDataset(data.Dataset):\n",
    "    'Characterizes a dataset for PyTorch'\n",
    "    def __init__(self, dir_path, decoy=False):\n",
    "        'Initialization'\n",
    "        \n",
    "        in_path = Path(dir_path)\n",
    "        assert in_path.exists()\n",
    "        assert in_path.is_dir()\n",
    "\n",
    "        self.aas            = ['_PAD'] + list(config.AAMass.keys())# + list(config.ModCHAR.values())\n",
    "        self.aa2idx         = {a:i for i, a in enumerate(self.aas)}\n",
    "        self.idx2aa         = {i:a for i, a in enumerate(self.aas)}\n",
    "        \n",
    "        self.pep_path       = dir_path\n",
    "        self.vocab_size     = len(self.aa2idx) # + self.charge + self.num_species + 1\n",
    "        print(\"Vocabulary Size: {}\".format(self.vocab_size))\n",
    "        self.seq_len        = config.get_config(section='ml', key='pep_seq_len')\n",
    "        \n",
    "        pep_lst, prot_list, pep_mass_lst, pep_modified_lst = load_peps(self.pep_path)\n",
    "        \n",
    "        self.pep_lst_set = set(pep_lst)\n",
    "#         out_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed-human-hcd-tryp-best/pts/\"\n",
    "#         with open(join(out_dir, 'pep_pickle.pkl'), 'rb') as f:\n",
    "#             search_peps = pickle.load(f)\n",
    "#         added_counter = 0\n",
    "#         for s_pep in search_peps:\n",
    "#             if s_pep not in self.pep_lst_set:\n",
    "#                 self.pep_lst_set.add(s_pep)\n",
    "#                 added_counter += 1\n",
    "#                 pep_lst.append(s_pep)\n",
    "#                 prot_list.append(\"unknown\")\n",
    "#                 pep_mass_lst.append(sim.get_pep_mass(s_pep))\n",
    "#                 pep_modified_lst.append(any(aa.islower() for aa in s_pep))\n",
    "#         print(\"New peptides added: {}\".format(added_counter))\n",
    "            \n",
    "            \n",
    "        all_sorts = list(zip(*sorted(zip(pep_lst, prot_list, pep_mass_lst, pep_modified_lst), key=lambda x: x[2])))\n",
    "        self.pep_list          = all_sorts[0]\n",
    "        self.prot_list         = all_sorts[1]\n",
    "        self.pep_mass_list     = all_sorts[2]\n",
    "        self.pep_modified_list = all_sorts[3]\n",
    "        if decoy:\n",
    "            self.pep_list, self.prot_list, self.pep_mass_list, self.pep_modified_list = self.get_docoys()\n",
    "            \n",
    "        print('Peptide Dataset Size: {}'.format(len(self.pep_list)))\n",
    "        \n",
    "\n",
    "    def __len__(self):\n",
    "        'Denotes the total number of samples'\n",
    "        return len(self.pep_list)\n",
    "\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        'Generates one sample of data'\n",
    "        pep = self.pep_list[index].strip()\n",
    "        pepl = [self.aa2idx[aa] for aa in pep]\n",
    "        pepl = self.pad_left(pepl)\n",
    "        torch_pep = torch.tensor(pepl, dtype=torch.long)\n",
    "        return torch_pep\n",
    "        \n",
    "\n",
    "    def pad_left(self, arr):\n",
    "        out = np.zeros(self.seq_len)\n",
    "        out[-len(arr):] = arr\n",
    "        return out\n",
    "    \n",
    "    def get_docoys(self):\n",
    "        decoy_list = []\n",
    "        decoy_prot_list = []\n",
    "        decoy_mass_list = []\n",
    "        decoy_modified_list = []\n",
    "        for pep, prot, mass, modified in zip(self.pep_list, self.prot_list, self.pep_mass_list, self.pep_modified_list):\n",
    "            pep_parts = re.findall(r\"([A-Z][a-z]?)\", pep)\n",
    "            decoy_pep = pep_parts[0] + \"\".join(pep_parts[-2:0:-1]) + pep_parts[-1]\n",
    "            if decoy_pep not in self.pep_lst_set:\n",
    "                decoy_list.append(decoy_pep)\n",
    "                decoy_prot_list.append(prot)\n",
    "                decoy_mass_list.append(mass)\n",
    "                decoy_modified_list.append(modified)\n",
    "        return decoy_list, decoy_prot_list, decoy_mass_list, decoy_modified_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% (66858 of 66858) |##################| Elapsed Time: 0:00:14 Time:  0:00:14\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count: 66532\n",
      "Spectral Dataset Size: 66532\n",
      "Vocabulary Size: 30\n",
      "Reading: /disk/raptor-2/mtari008/data/deepsnap/peps/human-peptidome.fasta\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0% (13301 of 11384076) |               | Elapsed Time: 0:00:00 ETA:   0:02:51"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File length: 11384076\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% (11384076 of 11384076) |############| Elapsed Time: 0:03:04 Time:  0:03:04\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Peptides written: 11735409\n",
      "Peptide Dataset Size: 11735409\n",
      "Vocabulary Size: 30\n",
      "Reading: /disk/raptor-2/mtari008/data/deepsnap/peps/human-peptidome.fasta\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0% (13443 of 11384076) |               | Elapsed Time: 0:00:00 ETA:   0:02:49"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File length: 11384076\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% (11384076 of 11384076) |############| Elapsed Time: 0:02:50 Time:  0:02:50\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Peptides written: 11735409\n",
      "Peptide Dataset Size: 11731403\n"
     ]
    }
   ],
   "source": [
    "spec_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed/one\"\n",
    "spec_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed/human-hcd-phospho-prep\"\n",
    "# spec_dir = \"/disk/raptor-2/mtari008/data/deepsnap/preprocessed-human-hcd-tryp-best/pts\"\n",
    "pep_dir = \"/disk/raptor-2/mtari008/data/deepsnap/peps\"\n",
    "\n",
    "spec_dataset = SpectralDataset(spec_dir)\n",
    "pep_dataset = PeptideDataset(pep_dir)\n",
    "dec_dataset = PeptideDataset(pep_dir, decoy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_batch_size = 16384\n",
    "# pep_batch_size = 4096\n",
    "pep_batch_size = 2048\n",
    "\n",
    "def spec_collate(batch):\n",
    "    specs = torch.stack([item for item in batch], 0)\n",
    "    dummy_pep = np.zeros(config.get_config(section=\"ml\", key=\"pep_seq_len\"))\n",
    "    dummy_pep = torch.from_numpy(dummy_pep).long().unsqueeze(0)\n",
    "    return [specs, dummy_pep]\n",
    "\n",
    "def pep_collate(batch):\n",
    "    peps = torch.stack([item for item in batch], 0)\n",
    "    dummy_spec = np.zeros(config.get_config(section=\"input\", key=\"spec_size\"))\n",
    "    dummy_spec = torch.from_numpy(dummy_spec).float().unsqueeze(0)\n",
    "    dummy_pep = np.zeros(config.get_config(section=\"ml\", key=\"pep_seq_len\"))\n",
    "    dummy_pep = torch.from_numpy(dummy_pep).long().unsqueeze(0)\n",
    "    return [dummy_spec, peps, dummy_pep]\n",
    "\n",
    "spec_loader = torch.utils.data.DataLoader(\n",
    "        dataset=spec_dataset, batch_size=spec_batch_size,\n",
    "        collate_fn=spec_collate)\n",
    "pep_loader = torch.utils.data.DataLoader(\n",
    "        dataset=pep_dataset, batch_size=pep_batch_size,\n",
    "        collate_fn=pep_collate)\n",
    "dec_loader = torch.utils.data.DataLoader(\n",
    "        dataset=dec_dataset, batch_size=pep_batch_size,\n",
    "        collate_fn=pep_collate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def runModel(loader, s_model, in_type):\n",
    "    with torch.no_grad():\n",
    "        accurate_labels = 0\n",
    "        all_labels = 0\n",
    "        loss = 0\n",
    "        out_out = torch.Tensor().cpu()\n",
    "        with progressbar.ProgressBar(max_value=len(loader)) as bar:\n",
    "            for batch_idx, batch in enumerate(loader):\n",
    "                batch[0], batch[1] = batch[0].to(device), batch[1].to(device)\n",
    "                if in_type == \"specs\":\n",
    "                    out_ = s_model(batch, data_type=in_type)[0]\n",
    "                elif in_type == \"peps\":\n",
    "                    batch[2] = batch[2].to(device)\n",
    "                    out_ = s_model(batch, data_type=in_type)[0]\n",
    "                out_out = torch.cat((out_out, out_.to(\"cpu\")), dim=0)\n",
    "                bar.update(batch_idx)\n",
    "        return out_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n",
      "dropout: 0.3\n",
      "torch.Size([512, 80000])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "  (linear1_2): Linear(in_features=512, out_features=256, bias=True)\n",
       "  (embedding): Embedding(30, 512)\n",
       "  (lstm): LSTM(512, 512, num_layers=2, batch_first=True, bidirectional=True)\n",
       "  (linear2_1): Linear(in_features=1024, out_features=512, bias=True)\n",
       "  (linear2_2): Linear(in_features=512, out_features=256, bias=True)\n",
       "  (dropout_conv1_1): Dropout(p=0.3, inplace=False)\n",
       "  (dropout1_1): Dropout(p=0.3, inplace=False)\n",
       "  (dropout1_2): Dropout(p=0.3, inplace=False)\n",
       "  (dropout1_3): Dropout(p=0.3, inplace=False)\n",
       "  (dropout1_4): Dropout(p=0.3, inplace=False)\n",
       "  (dropout1_5): Dropout(p=0.3, inplace=False)\n",
       "  (dropout2_1): Dropout(p=0.3, inplace=False)\n",
       "  (dropout2_2): Dropout(p=0.3, inplace=False)\n",
       "  (dropout2_3): Dropout(p=0.3, inplace=False)\n",
       "  (dropout2_4): Dropout(p=0.3, inplace=False)\n",
       "  (dropout2_5): Dropout(p=0.3, inplace=False)\n",
       "  (dropout2_6): Dropout(p=0.3, inplace=False)\n",
       "  (dropout_charge1): Dropout(p=0.3, inplace=False)\n",
       "  (dropout_spec_mass1): Dropout(p=0.3, inplace=False)\n",
       "  (dropout_spec_mass2): Dropout(p=0.3, inplace=False)\n",
       "  (dropout_spec_mass3): Dropout(p=0.3, inplace=False)\n",
       "  (dropout_pep_mass1): Dropout(p=0.3, inplace=False)\n",
       "  (dropout_pep_mass2): Dropout(p=0.3, inplace=False)\n",
       "  (dropout_pep_mass3): Dropout(p=0.3, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(device)\n",
    "os.environ['MASTER_ADDR'] = 'localhost'\n",
    "os.environ['MASTER_PORT'] = '12350'\n",
    "#dist.init_process_group(backend='nccl', world_size=1, rank=0)\n",
    "# snap_model = torch.load('models/model-all-mass2L-172-0.3-0.0005-mse-0.00001.pt')\n",
    "snap_model = Net(vocab_size=30, embedding_dim=512, hidden_lstm_dim=512, lstm_layers=2).to(device)\n",
    "print(snap_model.linear1_1.weight.shape)\n",
    "snap_model = nn.parallel.DistributedDataParallel(snap_model, device_ids=[0])\n",
    "# snap_model.load_state_dict(torch.load('models/32-embed-2-lstm-SnapLoss2-noch-3k-1k-152.pt')['model_state_dict'])\n",
    "snap_model.load_state_dict(torch.load('models/512-embed-2-lstm-SnapLoss-noch-80k-nist-massive-52.pt')['model_state_dict'])\n",
    "snap_model = snap_model.module\n",
    "snap_model.eval()\n",
    "snap_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(snap_model.state_dict(), 'models/model-all-mass2L-71-0.3-0.0005-mse-0.00001_dict.pt')"
   ]
  },
  {
   "cell_type": "code",
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-9-590a5fc04306>:36: DeprecationWarning: an integer is required (got type numpy.float64).  Implicit conversion to integers using __int__ is deprecated, and may be removed in a future version of Python.\n",
      "  ind = torch.LongTensor([[0]*np_spec.shape[1], np_spec[0]])\n",
      "100% (5 of 5) |##########################| Elapsed Time: 0:00:40 Time:  0:00:40\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Spectra done!\n"
     ]
    }
   ],
   "source": [
    "torch.cuda.set_device(0)\n",
    "e_specs = runModel(spec_loader, snap_model, \"specs\")\n",
    "print(\"Spectra done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% (5731 of 5731) |####################| Elapsed Time: 1:22:07 Time:  1:22:07\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Peptides done!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% (5729 of 5729) |####################| Elapsed Time: 1:23:17 Time:  1:23:17\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decoys done!\n"
     ]
    }
   ],
   "source": [
    "e_peps = runModel(pep_loader, snap_model, \"peps\")\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()\n",
    "print(\"Peptides done!\")\n",
    "e_decs = runModel(dec_loader, snap_model, \"peps\")\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()\n",
    "print(\"Decoys done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(e_specs.shape, e_peps.shape, e_decs.shape)\n",
    "#torch.save(e_specs, \"/lclhome/mtari008/mtari008/data/deepsnap/embeds/epeps-sty/e-specs.pt\")\n",
    "torch.save(e_peps,  \"/lclhome/mtari008/mtari008/data/deepsnap/embeds/epeps-45-0.3-0.0005-mse-0.00001/e-peps.pt\")\n",
    "torch.save(e_decs,  \"/lclhome/mtari008/mtari008/data/deepsnap/embeds/epeps-45-0.3-0.0005-mse-0.00001/e-decs.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#e_specs = torch.load(\"/lclhome/mtari008/mtari008/data/deepsnap/embeds/epeps-sty/e-specs.pt\")\n",
    "e_peps  = torch.load(\"/lclhome/mtari008/mtari008/data/deepsnap/embeds/epeps-172-0.3-0.0005-mse-0.00001/e-peps.pt\")\n",
    "e_decs  = torch.load(\"/lclhome/mtari008/mtari008/data/deepsnap/embeds/epeps-172-0.3-0.0005-mse-0.00001/e-decs.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ppm(val, ppm_val):\n",
    "    return (ppm_val / 1000000.) * val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_search_mask(spec_masses, pep_masses, tol):\n",
    "    l_tol = tol\n",
    "    rows = []\n",
    "    cols = []\n",
    "    pep_min = pep_max = 0\n",
    "    for row_id, spec_mass in enumerate(spec_masses):\n",
    "        min_mass = max(spec_mass - l_tol, 0.0)\n",
    "        max_mass = spec_mass + l_tol\n",
    "        while (pep_min < len(pep_masses) and \n",
    "               min_mass > pep_masses[pep_min]):\n",
    "            pep_min += 1\n",
    "        while (pep_max < len(pep_masses) and \n",
    "               max_mass > pep_masses[pep_max]):\n",
    "            pep_max += 1\n",
    "            \n",
    "        rows.extend([row_id] * (pep_max - pep_min))\n",
    "        cols.extend(range(pep_min, pep_max))\n",
    "    \n",
    "    assert len(rows) == len(cols)\n",
    "    mask = torch.zeros(len(spec_masses), len(pep_masses))\n",
    "    mask[rows, cols] = 1\n",
    "    return mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(search_loader, tol):\n",
    "    pep_sort_inds = []\n",
    "    pep_sort_vals = []\n",
    "    dec_sort_inds = []\n",
    "    dec_sort_vals = []\n",
    "    with progressbar.ProgressBar(max_value=len(search_loader)) as bar:\n",
    "        for idx, spec_batch in enumerate(search_loader):\n",
    "            l_tol = tol\n",
    "            batch_size = search_loader.batch_size\n",
    "            st = idx * batch_size\n",
    "            en = st + batch_size\n",
    "            spec_masses = spec_dataset.spec_mass_list[st:en]\n",
    "            min_mass = max(spec_masses[0] - l_tol, 0)\n",
    "            max_mass = spec_masses[-1] + l_tol\n",
    "            pep_min = pep_max = 0\n",
    "\n",
    "            while (pep_min < len(pep_dataset.pep_mass_list) and \n",
    "                   min_mass - pep_dataset.pep_mass_list[pep_min] > 0.01):\n",
    "                pep_min += 1\n",
    "            while (pep_max < len(pep_dataset.pep_mass_list) and\n",
    "                   max_mass - pep_dataset.pep_mass_list[pep_max] >= 0.01):\n",
    "                pep_max += 1\n",
    "\n",
    "            dec_min = dec_max = 0\n",
    "            while (dec_min < len(dec_dataset.pep_mass_list) and\n",
    "                   min_mass - dec_dataset.pep_mass_list[dec_min] > 0.01):\n",
    "                dec_min += 1\n",
    "            while (dec_max < len(dec_dataset.pep_mass_list) and \n",
    "                   max_mass - dec_dataset.pep_mass_list[dec_max] >= 0.01):\n",
    "                dec_max += 1\n",
    "\n",
    "            pep_batch = e_peps[pep_min:pep_max]\n",
    "            pep_masses = pep_dataset.pep_mass_list[pep_min:pep_max]\n",
    "            dec_batch = e_decs[dec_min:dec_max]\n",
    "            dec_masses = dec_dataset.pep_mass_list[dec_min:dec_max]\n",
    "\n",
    "            spec_batch = spec_batch.to(device)\n",
    "            #print(\"pep batch len: {}\".format(len(pep_batch)))\n",
    "            l_pep_batch_size = 16384\n",
    "            # l_pep_batch_size = 32768\n",
    "            pep_loader = torch.utils.data.DataLoader(\n",
    "                dataset=pep_batch, batch_size=l_pep_batch_size)\n",
    "            l_pep_dist = []\n",
    "            for pep_idx, l_pep_batch in enumerate(pep_loader):\n",
    "                l_pep_batch = l_pep_batch.to(device)\n",
    "                l_st = pep_idx * l_pep_batch_size\n",
    "                l_en = l_st + l_pep_batch_size\n",
    "                l_pep_masses = pep_masses[l_st:l_en]\n",
    "                spec_pep_mask = get_search_mask(spec_masses, l_pep_masses, tol).to(device)\n",
    "                # spec_pep_mask[spec_pep_mask == 0] = float(\"inf\")\n",
    "                spec_pep_dist = 1 / process.pairwise_distances(spec_batch, l_pep_batch)\n",
    "                l_pep_dist.append((spec_pep_dist * spec_pep_mask).to(\"cpu\"))\n",
    "            \n",
    "            l_pep_dist.append(torch.zeros(len(spec_batch), keep_psms + 1))\n",
    "            pep_sort = torch.cat(l_pep_dist, 1)\n",
    "            pep_lcn = np.ma.masked_array(pep_sort, mask=pep_sort==0).min(1).data\n",
    "            pep_sort = pep_sort.sort(descending=True)\n",
    "            pep_sort_inds.append(pep_sort.indices[:, :keep_psms + 1] + pep_min) # offset for the global array\n",
    "            pep_sort_vals.append(torch.cat((pep_sort.values[:, :keep_psms + 1], \n",
    "                                            torch.from_numpy(pep_lcn).unsqueeze(1)), 1))\n",
    "            \n",
    "            dec_loader = torch.utils.data.DataLoader(\n",
    "                dataset=dec_batch, batch_size=l_pep_batch_size)\n",
    "            l_dec_dist = []\n",
    "            for dec_idx, l_dec_batch in enumerate(dec_loader):\n",
    "                l_dec_batch = l_dec_batch.to(device)\n",
    "                l_st = dec_idx * l_pep_batch_size\n",
    "                l_en = l_st + l_pep_batch_size\n",
    "                l_dec_masses = dec_masses[l_st:l_en]\n",
    "                \n",
    "                spec_dec_mask = get_search_mask(spec_masses, l_dec_masses, tol).to(device)\n",
    "                # spec_dec_mask[spec_dec_mask == 0] = float(\"inf\")\n",
    "                spec_dec_dist = 1 / process.pairwise_distances(spec_batch, l_dec_batch)\n",
    "                l_dec_dist.append((spec_dec_dist * spec_dec_mask).to(\"cpu\"))\n",
    "            \n",
    "            l_dec_dist.append(torch.zeros(len(spec_batch), keep_psms + 1))\n",
    "            dec_sort = torch.cat(l_dec_dist, 1)\n",
    "            dec_lcn = np.ma.masked_array(dec_sort, mask=dec_sort==0).min(1).data\n",
    "            dec_sort = dec_sort.sort(descending=True)\n",
    "            dec_sort_inds.append(dec_sort.indices[:, :keep_psms + 1] + dec_min) # offset for the global array\n",
    "            dec_sort_vals.append(torch.cat((dec_sort.values[:, :keep_psms + 1], \n",
    "                                            torch.from_numpy(dec_lcn).unsqueeze(1)), 1))\n",
    "            \n",
    "            bar.update(idx)\n",
    "    \n",
    "    pep_inds = torch.cat(pep_sort_inds, 0)\n",
    "    pep_vals = torch.cat(pep_sort_vals, 0)\n",
    "    dec_inds = torch.cat(dec_sort_inds, 0)\n",
    "    dec_vals = torch.cat(dec_sort_vals, 0)\n",
    "    return pep_inds, pep_vals, dec_inds, dec_vals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% (65 of 65) |########################| Elapsed Time: 0:24:56 Time:  0:24:56\n"
     ]
    }
   ],
   "source": [
    "search_spec_batch_size = 1024\n",
    "#search_spec_batch_size = 2048\n",
    "precursor_tolerance = 0.5\n",
    "keep_psms = 5\n",
    "search_loader = torch.utils.data.DataLoader(\n",
    "    dataset=e_specs, batch_size=search_spec_batch_size, shuffle=False)\n",
    "pep_inds, pep_vals, dec_inds, dec_vals = search(search_loader, precursor_tolerance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(pep_inds, \"/lclhome/mtari008/mtari008/data/deepsnap/results-human-hcd-tryp-best/pep-inds.pt\")\n",
    "torch.save(pep_vals, \"/lclhome/mtari008/mtari008/data/deepsnap/results-human-hcd-tryp-best/pep-vals.pt\")\n",
    "torch.save(dec_inds, \"/lclhome/mtari008/mtari008/data/deepsnap/results-human-hcd-tryp-best/dec-inds.pt\")\n",
    "torch.save(dec_vals, \"/lclhome/mtari008/mtari008/data/deepsnap/results-human-hcd-tryp-best/dec-vals.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "pep_inds = torch.load(\"/lclhome/mtari008/mtari008/data/deepsnap/results/pep-inds.pt\")\n",
    "pep_vals = torch.load(\"/lclhome/mtari008/mtari008/data/deepsnap/results/pep-vals.pt\")\n",
    "dec_inds = torch.load(\"/lclhome/mtari008/mtari008/data/deepsnap/results/dec-inds.pt\")\n",
    "dec_vals = torch.load(\"/lclhome/mtari008/mtari008/data/deepsnap/results/dec-vals.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_percolator_input(l_pep_inds, l_pep_vals, pd_dataset, res_type):\n",
    "    assert res_type == \"target\" or res_type == \"decoy\"\n",
    "    assert len(l_pep_inds) == len(l_pep_vals)\n",
    "    pin_charge = config.get_config(section=\"search\", key=\"charge\")\n",
    "    l_global_out = []\n",
    "    tot_count = 0\n",
    "    # max_snap = torch.max(l_pep_vals).item()\n",
    "    max_snap = 5.137566089630127\n",
    "    for idx, (pep_inds_row, pep_vals_row) in enumerate(zip(l_pep_inds, l_pep_vals)):\n",
    "        # Reminder: pep_inds_row length is one less than pep_vals_row\n",
    "        for iidx in range(len(pep_inds_row) - 1):\n",
    "            pep_ind = pep_inds_row[iidx]\n",
    "            pep_val = pep_vals_row[iidx]\n",
    "            pep_val_last = pep_vals_row[-1]\n",
    "            pep_val_next = pep_vals_row[iidx + 1]\n",
    "            if pep_val.item() > 0:\n",
    "                charge = [0] * pin_charge\n",
    "                charge[spec_dataset.spec_charge_list[idx] - 1] = 1\n",
    "                label = 1 if res_type == \"target\" else -1\n",
    "                out_row = [f\"{res_type}-{tot_count}\", label, idx]\n",
    "                spec_mass = spec_dataset.spec_mass_list[idx]\n",
    "                pep_mass = pd_dataset.pep_mass_list[pep_ind.item()]\n",
    "                out_row.append(spec_mass)\n",
    "                out_row.append(pep_mass)\n",
    "                out_row.append(((pep_val - pep_vals_row[-1]).item()) / max(pep_val.item(), 1.0))\n",
    "                out_row.append(((pep_val - pep_vals_row[iidx + 1]).item()) / max(pep_val.item(), 1.0))\n",
    "#                 out_row.append(((pep_val - pep_vals_row[-1]).item()))\n",
    "#                 out_row.append(((pep_val - pep_vals_row[iidx + 1]).item()))\n",
    "                out_row.append(pep_val.item())\n",
    "                out_row.extend(charge)\n",
    "                out_row.append(spec_mass - pep_mass)\n",
    "                out_row.append(abs(spec_mass - pep_mass))\n",
    "                \n",
    "                out_pep = pd_dataset.pep_list[pep_ind.item()]\n",
    "                out_pep_array = []\n",
    "                for aa in out_pep:\n",
    "                    if aa.islower():\n",
    "                        out_pep_array.append(\"[\"+str(config.AAMass[aa])+\"]\")\n",
    "                    else:\n",
    "                        out_pep_array.append(aa)\n",
    "                out_pep = \"\".join(out_pep_array)\n",
    "                out_row.append((out_pep.count(\"K\") + out_pep.count(\"R\")) - (out_pep.count(\"KP\") + out_pep.count(\"RP\")))\n",
    "                out_prot = pd_dataset.prot_list[pep_ind.item()]\n",
    "                pep_len = sum([a.isupper() for a in out_pep])\n",
    "                out_row.append(pep_len)\n",
    "                out_row.append(out_pep)\n",
    "                out_row.append(out_prot)\n",
    "                l_global_out.append(out_row)\n",
    "                tot_count += 1\n",
    "    return l_global_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_out = generate_percolator_input(pep_inds, pep_vals, pep_dataset, \"target\")\n",
    "pin_charge = config.get_config(section=\"search\", key=\"charge\")\n",
    "charge_cols = [f\"charge-{ch+1}\" for ch in range(pin_charge)]\n",
    "cols = [\"SpecId\", \"Label\", \"ScanNr\", \"ExpMass\", \"CalcMass\", \"deltLCn\", \"deltCn\", \"SNAP\"] + charge_cols + [\"dM\", \"absdM\", \"enzInt\", \"PepLen\", \"Peptide\", \"Proteins\"]\n",
    "df = pd.DataFrame(global_out, columns=cols)\n",
    "df.sort_values(by=\"SNAP\", inplace=True, ascending=False)\n",
    "df.to_csv(\"/lclhome/mtari008/mtari008/data/deepsnap/results-human-hcd-tryp-best/res-pin-target-29.pin\", sep=\"\\t\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_out = generate_percolator_input(dec_inds, dec_vals, dec_dataset, \"decoy\")\n",
    "pin_charge = config.get_config(section=\"search\", key=\"charge\")\n",
    "charge_cols = [f\"charge-{ch+1}\" for ch in range(pin_charge)]\n",
    "cols = [\"SpecId\", \"Label\", \"ScanNr\", \"ExpMass\", \"CalcMass\", \"deltLCn\", \"deltCn\", \"SNAP\"] + charge_cols + [\"dM\", \"absdM\", \"enzInt\", \"PepLen\", \"Peptide\", \"Proteins\"]\n",
    "df = pd.DataFrame(global_out, columns=cols)\n",
    "df.sort_values(by=\"SNAP\", inplace=True, ascending=False)\n",
    "df.to_csv(\"/lclhome/mtari008/mtari008/data/deepsnap/results-human-hcd-tryp-best/res-pin-decoy-29.pin\", sep=\"\\t\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this one is for combined file. for separate files use the above two.\n",
    "pin_charge = config.get_config(section=\"search\", key=\"charge\")\n",
    "charge_cols = [f\"charge-{ch+1}\" for ch in range(pin_charge)]\n",
    "cols = [\"SpecId\", \"Label\", \"ScanNr\", \"ExpMass\", \"CalcMass\", \"deltLCn\", \"deltCn\", \"SNAP\"] + charge_cols + [\"dM\", \"absdM\", \"enzInt\", \"PepLen\", \"Peptide\", \"Proteins\"]\n",
    "df = pd.DataFrame(global_out, columns=cols)\n",
    "df.sort_values(by=\"SNAP\", inplace=True, ascending=False)\n",
    "df.to_csv(\"/lclhome/mtari008/mtari008/data/deepsnap/results-human-hcd-tryp-best/res-pin20.tab\", sep=\"\\t\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.71178627,\n",
       " 0.82731843,\n",
       " 0.8814876,\n",
       " 0.9404478,\n",
       " 0.90577126,\n",
       " 0.7447554,\n",
       " 0.87054515,\n",
       " 0.81516623,\n",
       " 1.0211788,\n",
       " 0.84863305]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sort_col = [x[0] for x in np.array(pep_vals)]\n",
    "sort_col[-10:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'sort_col' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-27-86499b30da0a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m df = pd.DataFrame(list(zip(np.array(pep_inds), np.array(pep_vals), \n\u001b[0;32m----> 2\u001b[0;31m                            np.array(dec_inds), np.array(dec_vals), np.array(sort_col))), \n\u001b[0m\u001b[1;32m      3\u001b[0m                   columns=[\"pep_inds\", \"pep_vals\", \"dec_inds\", \"dec_vals\", \"sort_col\"])\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mdf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msort_values\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mby\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"sort_col\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minplace\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mascending\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'sort_col' is not defined"
     ]
    }
   ],
   "source": [
    "df = pd.DataFrame(list(zip(np.array(pep_inds), np.array(pep_vals), \n",
    "                           np.array(dec_inds), np.array(dec_vals), np.array(sort_col))), \n",
    "                  columns=[\"pep_inds\", \"pep_vals\", \"dec_inds\", \"dec_vals\", \"sort_col\"])\n",
    "\n",
    "df.sort_values(by=\"sort_col\", inplace=True, ascending=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv(\"/lclhome/mtari008/mtari008/data/deepsnap/results/output.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.0.5'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fastaToSpectraBatch(filePath, spectra_batch_size):    \n",
    "    f = open(filePath)\n",
    "    lines = f.readlines()\n",
    "    f.close()\n",
    "    masses = []\n",
    "    spectra_out = torch.Tensor().to(device)\n",
    "    peps = []\n",
    "    \n",
    "    dh = display('0%', display_id=True)\n",
    "    \n",
    "    start = 0\n",
    "    i = 0\n",
    "    while start < len(lines):\n",
    "        print('Batch: ' + str(i))\n",
    "        i += 1\n",
    "        \n",
    "        print('Generating spectra...')\n",
    "        spectra, l_masses, l_peps = fastatospectra(lines, start, spectra_batch_size, dh)\n",
    "        masses.extend(l_masses)\n",
    "        peps.extend(l_peps)\n",
    "        start = start + spectra_batch_size\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            print('Converting to tensor...')\n",
    "            '''dtype=torch.float'''\n",
    "            spectra = np.asarray(spectra)\n",
    "            spectraTensor = torch.as_tensor(spectra, dtype=torch.float)[:, None, :]\n",
    "            spectra_loader = torch.utils.data.DataLoader(dataset=spectraTensor, batch_size=batch_size, shuffle=False)\n",
    "            \n",
    "            print('Running the model...')\n",
    "            model_out = runModel(spectra_loader)\n",
    "            spectra_out = torch.cat((spectra_out, model_out), dim=0)\n",
    "            \n",
    "    return spectra_out, masses, peps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spectra_batch_size = 500000\n",
    "#targetPath = '/disk/raptor/lclhome/mtari008/Proteogenomics/DeepSNAP/data/rat/peptide-dbs/rat-peptides.target.txt'\n",
    "#decoyPath = '/disk/raptor/lclhome/mtari008/Proteogenomics/DeepSNAP/data/rat/peptide-dbs/rat-peptides.decoy.txt'\n",
    "targetPath = '/disk/raptor/lclhome/mtari008/Proteomics/crux/crux-output/generate-peptides.target.txt'\n",
    "decoyPath = '/disk/raptor/lclhome/mtari008/Proteomics/crux/crux-output/generate-peptides.decoy.txt'\n",
    "targetspectra_out, targetmasses, targetpeptides = fastaToSpectraBatch(targetPath, spectra_batch_size)\n",
    "decoyspectra_out, decoymasses, decoypeptides = fastaToSpectraBatch(decoyPath, spectra_batch_size)\n",
    "\n",
    "\n",
    "print(len(targetspectra_out), len(targetmasses))\n",
    "print(len(decoyspectra_out), len(decoymasses))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nearest_target_distances, nearest_target_indexes = process.pairwise_distances(\n",
    "    queryspectra_out, targetspectra_out).min(1)\n",
    "\n",
    "nearest_decoy_distances, nearest_decoy_indexes  = process.pairwise_distances(\n",
    "    queryspectra_out,  decoyspectra_out).min(1)\n",
    "\n",
    "nearest_distances, target_decoy = torch.cat((nearest_target_distances.view(1, -1), \n",
    "                                             nearest_decoy_distances.view(1, -1)), dim=0).min(0)\n",
    "\n",
    "print(target_decoy.to('cpu'))\n",
    "\n",
    "sorted_nearest_distances, snd_index = nearest_distances.sort()\n",
    "target_decoy = target_decoy[snd_index]\n",
    "\n",
    "print(target_decoy)\n",
    "\n",
    "## To lookup a peptide corresponding to a value in sorted_nearest_distances at index i:\n",
    "## td = target_deocy[i]            # a binary array. 0: target, 1: decoy\n",
    "## query_index = snd_index[i]        # lookup the index of query spectrum\n",
    "## pep_index = nearest_decoy_indexes[query_index] if td else nearest_target_indexes[query_index]\n",
    "## td_peptide = decoy_peptides[pep_index] if td else target_peptides[pep_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fdr(td, target_fdr):\n",
    "    target_count = 0\n",
    "    decoy_count = 0\n",
    "    q_values = [0] * len(td)\n",
    "    for idx, val in enumerate(td):\n",
    "        if val == 0:\n",
    "            target_count += 1\n",
    "        else:\n",
    "            decoy_count += 1\n",
    "        return_fdr = decoy_count / (idx + 1)\n",
    "        print('{} / {} = {}'.format(decoy_count, idx+1, decoy_count/(idx+1)))\n",
    "        #if return_fdr > target_fdr:\n",
    "        #    return target_count, return_fdr\n",
    "    return target_count, return_fdr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fdr(target_decoy, 0.01))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(targetmasses[0:10])\n",
    "print(len(decoypeptides))\n",
    "targetspectra_out, targetmasses, targetpeptides = list(zip(*sorted(zip(targetspectra_out, targetmasses, targetpeptides), \n",
    "                                                                   key=lambda pair: pair[1])))\n",
    "decoyspectra_out, decoymasses, decoypeptides = list(zip(*sorted(zip(decoyspectra_out, decoymasses, decoypeptides),\n",
    "                                                                key=lambda pair: pair[1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(targetmasses))\n",
    "print(targetmasses[1000:1010])\n",
    "print(decoymasses[1000:1010])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_range(arr, val, tolerance=3):\n",
    "    assert len(arr) > 0\n",
    "    assert tolerance > 0\n",
    "    assert val > 0\n",
    "    \n",
    "    left = val - tolerance if val - tolerance >= 0 else 0\n",
    "    right = val + tolerance\n",
    "    \n",
    "    ileft = bisect.bisect_left(arr, left)\n",
    "    iright = bisect.bisect_right(arr, right)\n",
    "    \n",
    "    return (ileft, iright-1) if iright - ileft > 0 else (-1, -1)\n",
    "    \n",
    "# Function to insert element \n",
    "def insert(lst, n): \n",
    "    assert not lst or len(lst[0]) == len(n)\n",
    "    \n",
    "    index = num_psm\n",
    "    # Searching for the position \n",
    "    for i in range(len(lst)): \n",
    "        if lst[i][0] > n[0]: \n",
    "            index = i \n",
    "            break\n",
    "      \n",
    "    # Inserting n in the list \n",
    "    lst = lst[:index] + [n] + lst[index:] \n",
    "    return lst[:num_psm]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_L2_Dist(queryspectra, querymasses, dbspectra, dbmasses, pre_tol=3, keep_psm=1):\n",
    "    prev = 0\n",
    "    psm_scores = []\n",
    "    psm_pepids = []\n",
    "    \n",
    "    for ind, (spec, spec_mass) in enumerate(zip(queryspectra, querymasses)):\n",
    "        left, right = find_range(dbmasses, spec_mass, pre_tol)\n",
    "        if left != -1 and right != -1:\n",
    "            candidate_specs = dbspectra[left:right+1]\n",
    "            l_scores = np.linalg.norm(np.subtract(candidate_specs, spec), axis=1)**2\n",
    "            psm_pepids.append(left + l_scores.argsort()[:keep_psm]) #add argsort later\n",
    "            psm_scores.append(np.sort(l_scores)[:keep_psm])\n",
    "            #print(ind)\n",
    "            new = int((ind/len(queryspectra)) * 100)\n",
    "            if new > prev:\n",
    "                clear_output(wait=True)\n",
    "                print(str(new) + '%')\n",
    "                prev = new\n",
    "        else:\n",
    "            print('***********error**************')\n",
    "            print('ind: ' + str(ind))\n",
    "            print('spec_mass: ' + str(spec_mass))\n",
    "            print('**********end error************')\n",
    "            \n",
    "    return psm_scores, psm_pepids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_xcorr(queryspectra, querymasses, dbspectra, dbmasses, pre_tol=3, keep_psm=1):\n",
    "    prev = 0\n",
    "    psm_scores = []\n",
    "    psm_pepids = []\n",
    "    \n",
    "    for ind, (spec, spec_mass) in enumerate(zip(queryspectra, querymasses)):\n",
    "        left, right = find_range(dbmasses, spec_mass, pre_tol)\n",
    "        if left != -1 and right != -1:\n",
    "            candidate_specs = dbspectra[left:right+1]\n",
    "            l_scores = np.dot(candidate_specs, spec)\n",
    "            psm_pepids.append(left + l_scores.argsort()[::-1][:keep_psm]) #add argsort later\n",
    "            psm_scores.append(np.sort(l_scores)[::-1][:keep_psm])\n",
    "            #print(ind)\n",
    "            new = int((ind/len(queryspectra)) * 100)\n",
    "            if new > prev:\n",
    "                clear_output(wait=True)\n",
    "                print(str(new) + '%')\n",
    "                prev = new\n",
    "        else:\n",
    "            print('***********error**************')\n",
    "            print('ind: ' + str(ind))\n",
    "            print('spec_mass: ' + str(spec_mass))\n",
    "            print('**********end error************')\n",
    "            \n",
    "    return psm_scores, psm_pepids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# targetpsm_scores, targetpsm_ids = calculate_L2_Dist(queryspectra_out, spectramasses, targetspectra_out,\n",
    "#                                      targetmasses, precursor_tolerance, num_psm)\n",
    "# decoypsm_scores, decoypsm_ids = calculate_L2_Dist(queryspectra_out, spectramasses, decoyspectra_out,\n",
    "#                                      decoymasses, precursor_tolerance, num_psm)\n",
    "targetpsm_scores, targetpsm_ids = calculate_xcorr(queryspectra_out, spectramasses, targetspectra_out,\n",
    "                                     targetmasses, precursor_tolerance, num_psm)\n",
    "decoypsm_scores, decoypsm_ids = calculate_xcorr(queryspectra_out, spectramasses, decoyspectra_out,\n",
    "                                     decoymasses, precursor_tolerance, num_psm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(targetmasses[2000000:2000010])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(targetspectra_out))\n",
    "print(len(decoyspectra_out))\n",
    "\n",
    "print(len(targetpsm_scores))\n",
    "print(len(decoypsm_scores))\n",
    "\n",
    "print(targetpsm_scores[0:10])\n",
    "print(decoypsm_scores[0:10])\n",
    "\n",
    "targetpsm_scores = np.asarray(targetpsm_scores).flatten()\n",
    "decoypsm_scores = np.asarray(decoypsm_scores).flatten()\n",
    "\n",
    "targetpsm_ids = np.asarray(targetpsm_ids).flatten()\n",
    "decoypsm_ids = np.asarray(decoypsm_ids).flatten()\n",
    "\n",
    "targetpsm_peps = np.asarray(targetpeptides)[targetpsm_ids]\n",
    "decoypsm_peps = np.asarray(decoypeptides)[decoypsm_ids]\n",
    "\n",
    "#targetpsm_ids = targetpsm_ids[np.argsort(np.asarray(targetpsm_scores).flatten())]\n",
    "#decoypsm_ids = decoypsm_ids[np.argsort(np.asarray(decoypsm_scores).flatten())]\n",
    "\n",
    "sorted_targets = np.sort(np.asarray(targetpsm_scores).flatten())[::-1]\n",
    "sorted_decoys = np.sort(np.asarray(decoypsm_scores).flatten())[::-1]\n",
    "\n",
    "print(targetpsm_ids[0:10])\n",
    "print(decoypsm_ids[0:10])\n",
    "\n",
    "print(sorted_targets[0:100])\n",
    "print(sorted_decoys[0:100])\n",
    "\n",
    "#targetpsm_scores = np.column_stack((targetpsm_scores, np.ones(len(targetpsm_scores))))\n",
    "#decoypsm_scores = np.column_stack((decoypsm_scores, np.zeros(len(decoypsm_scores))))\n",
    "#decoypsm_scores = np.asarray(decoypsm_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "targetpsm_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(\n",
    "    data={'spectrum_ids':np.arange(len(targetpsm_peps)), 'target_psms':targetpsm_peps, \n",
    "          'targetpsm_scores':targetpsm_scores, 'decoy_psms':decoypsm_peps, 'decoypsm_scores':decoypsm_scores})\n",
    "df.to_csv('deepsearch.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sortedpeps = np.array(sortedpeps)\n",
    "correct_matches = 0\n",
    "for ind, psms in enumerate(psm_scores):\n",
    "    #print('ind: ' + str(ind) + ', mass: ' + str(sortedspecmasses[ind]))\n",
    "    #print('psms: ' + str(psms) + ', psm mass: ' + str(sortedpepmasses[psms[0]]))\n",
    "    cand_peps = sortedpeps[psms]\n",
    "    #print(cand_peps)\n",
    "    for pep in cand_peps:\n",
    "        #print('pepidmass: ' + str(pepidmass[pep][0]))\n",
    "        #print('ids: ' + str(ids[ind]) + ', in: ' + str(pepidmass[pep][0]))\n",
    "        if ids[ind] in pepidmass[pep][0]:\n",
    "            correct_matches += 1\n",
    "            break\n",
    "print(len(psm_scores))\n",
    "print(correct_matches)\n",
    "print(str((correct_matches/len(psm_scores)) * 100) + '% correct matches in top ' + str(num_psm) + ' psms.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scratch Work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['charge-1', 'charge-2', 'charge-3', 'charge-4', 'charge-5']\n",
      "['a', 'b', 'c', 'charge-1', 'charge-2', 'charge-3', 'charge-4', 'charge-5']\n"
     ]
    }
   ],
   "source": [
    "pin_charge = config.get_config(section=\"search\", key=\"charge\")\n",
    "charge_cols = [f\"charge-{ch+1}\" for ch in range(pin_charge)]\n",
    "print(charge_cols)\n",
    "a = [\"a\", \"b\", \"c\"]\n",
    "b = a + charge_cols\n",
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = open('/disk/raptor/lclhome/mtari008/Proteogenomics/DeepSNAP/data/rat/peptide-dbs/rat-peptides.target.txt')\n",
    "tlines = file.readlines()\n",
    "print(len(tlines))\n",
    "tpeps = list((tline.split('\\t')[0], float(tline.split('\\t')[1])) for tline in tlines)\n",
    "tpeps = sorted(tpeps, key=lambda pair: pair[1])\n",
    "print(tpeps[10000:10010])\n",
    "file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = open('/disk/raptor/lclhome/mtari008/Proteogenomics/DeepSNAP/data/rat/peptide-dbs/rat-peptides.decoy.txt')\n",
    "tlines = file.readlines()\n",
    "print(len(tlines))\n",
    "tpeps = list((tline.split('\\t')[0], float(tline.split('\\t')[1])) for tline in tlines)\n",
    "tpeps = sorted(tpeps, key=lambda pair: pair[1])\n",
    "print(tpeps[10000:10010])\n",
    "file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.arange(150)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.ones(2, 3)\n",
    "z = (torch.ones(3) * 5).view(1, 3)\n",
    "print(x)\n",
    "print(z)\n",
    "# or this\n",
    "#z = torch.ones(2, 3, 1) * 5\n",
    "y = torch.cat((x, z), dim=0)\n",
    "print(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Keep this cell....\n",
    "## How to sort a 2D tensor by its first row.....\n",
    "x = torch.Tensor([5., 3., 4., 1., 2.]).view(1, -1)\n",
    "z = torch.Tensor([1., 3., 2., 5., 4.]).view(1, -1)\n",
    "print(x)\n",
    "print(z)\n",
    "# or this\n",
    "#z = torch.ones(2, 3, 1) * 5\n",
    "y = torch.cat((x, z), dim=0)\n",
    "_, sorted_index_y = y[0, :].sort()\n",
    "sorted_y = y[:, sorted_index_y]\n",
    "print(y)\n",
    "print(sorted_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regex = r\"Name\\:\\s(?P<pep>[a-zA-Z]+)\\/(?P<charge>\\d+)(?:\\_(?P<num_mods>\\d+)(?P<mods>.*))?\"\n",
    "test_str = \"Name: AAAAAEEGMEPR/2\"\n",
    "m = re.search(regex, test_str)\n",
    "print((m['pep']))\n",
    "print(type(int(m.group('charge'))))\n",
    "\n",
    "print(bool(m['pep']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = float(re.findall(r\"MW\\:\\s([-+]?[0-9]*\\.?[0-9]*)\", \"MW: 1261.5610\")[0])\n",
    "print(type(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['ABC', 'CD', 'GHTE']\n"
     ]
    }
   ],
   "source": [
    "dic = {'ABC':(123.5, True), 'CD':(544.2, False), 'GHTE':(7654.445, True)}\n",
    "df = pd.DataFrame.from_dict(dic, orient=\"index\", columns=[\"mass\", \"modified\"])\n",
    "df.index.name = \"pep\"\n",
    "print(list(df.index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ADaCiBE\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "pep = \"ABCiDaE\"\n",
    "pep_parts = re.findall(r\"([A-Z][a-z]?)\", pep)\n",
    "pep = pep_parts[0] + \"\".join(pep_parts[-2:0:-1]) + pep_parts[-1]\n",
    "print(pep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pairwise_distances(x, y=None):\n",
    "    \"\"\"\n",
    "    Input: x is a Nxd matrix\n",
    "           y is an optional Mxd matirx\n",
    "    Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]\n",
    "            if y is not given then use 'y=x'.\n",
    "    i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2\n",
    "    \"\"\"\n",
    "    x_norm = (x ** 2).sum(1).view(-1, 1)\n",
    "    print(\"x norm:\")\n",
    "    print(x_norm)\n",
    "    if y is not None:\n",
    "        y_t = torch.transpose(y, 0, 1)\n",
    "        print(\"y_t:\")\n",
    "        print(y_t)\n",
    "        y_norm = (y ** 2).sum(1).view(1, -1)\n",
    "        print(\"y_norm:\")\n",
    "        print(y_norm)\n",
    "    else:\n",
    "        y_t = torch.transpose(x, 0, 1)\n",
    "        y_norm = x_norm.view(1, -1)\n",
    "\n",
    "    dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t)\n",
    "    # Ensure diagonal is zero if x=y\n",
    "    if y is None:\n",
    "        dist = dist - torch.diag(dist.diag())\n",
    "    dist[dist != dist] = 0  # set all nan values to zero\n",
    "    return torch.clamp(dist, 0.0, np.inf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[3., 2., 0., 1., 6.],\n",
      "        [8., 4., 5., 1., 6.],\n",
      "        [0., 7., 2., 1., 6.]])\n",
      "tensor([[0., 2., 0., 1., 1.],\n",
      "        [1., 4., 5., 5., 6.],\n",
      "        [0., 7., 2., 1., 6.]])\n",
      "tensor([5.8310, 8.0623, 0.0000])\n"
     ]
    }
   ],
   "source": [
    "x = torch.tensor([[3., 2., 0., 1., 6.], [8., 4., 5., 1., 6.], [0., 7., 2., 1., 6.]])\n",
    "y = torch.tensor([[0., 2., 0., 1., 1.], [1., 4., 5., 5., 6.], [0., 7., 2., 1., 6.]])\n",
    "print(x)\n",
    "print(y)\n",
    "z = torch.norm(x - y, p=2, dim=1)\n",
    "print(z)\n",
    "#x = set(map(tuple, x))\n",
    "#y = set(map(tuple, y))\n",
    "#z = x - x.intersection(y)\n",
    "#print(torch.tensor(list(z), dtype=float))\n",
    "#print(set.intersection(set(*map(frozenset, x))))\n",
    "#z = torch.cat((x, y), axis=1)\n",
    "#print(z)\n",
    "#print(x)\n",
    "#print(torch.from_numpy(np.ma.masked_array(x, mask=x==0).min(1).data).unsqueeze(1))\n",
    "#x[range(3), torch.from_numpy(np.ma.masked_array(x, mask=x==0).min(1))]\n",
    "# for idx in np.nonzero(x):\n",
    "#     print(idx)\n",
    "#     print(x[tuple(idx)])\n",
    "# y = torch.tensor([[9, 9], [9, 9], [9, 9]])\n",
    "# z = torch.cat((x, y), 1)\n",
    "# print(z.shape[1])\n",
    "# x = torch.tensor([3., 2., 0., 1., 6., 0., 1, 9.])\n",
    "# x[np.nonzero(x)]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}