AutoEncoder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adeb6dce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "plt.switch_backend('agg')\n",
    "\n",
    "import os\n",
    "import re\n",
    "import random\n",
    "\n",
    "import sys\n",
    "sys.path.insert(1, '/gscratch/ml4ml/sidlak/superfold/superfold');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c0aea8cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pyrosetta\n",
    "# from pyrosetta.rosetta.core.pose import Pose\n",
    "\n",
    "# pyrosetta.init(\"-ignore_unrecognized_res true\")\n",
    "\n",
    "# def get_rmsd(design: Pose, prediction: Pose) -> float:\n",
    "#     \"\"\"Calculate Ca-RMSD of prediction to design\"\"\"\n",
    "#     import pyrosetta\n",
    "#     rmsd_calc = pyrosetta.rosetta.core.simple_metrics.metrics.RMSDMetric()\n",
    "#     # https://graylab.jhu.edu/PyRosetta.documentation/pyrosetta.rosetta.core.scoring.html?highlight=rmsd_atoms#pyrosetta.rosetta.core.scoring.rmsd_atoms\n",
    "#     rmsd_calc.set_rmsd_type(pyrosetta.rosetta.core.scoring.rmsd_atoms(3)) # change to the rmsd atom type desired, 3 = Ca only\n",
    "#     rmsd_calc.set_run_superimpose(True)\n",
    "#     rmsd_calc.set_comparison_pose(design)\n",
    "#     rmsd = float(rmsd_calc.calculate(prediction))\n",
    "#     return rmsd\n",
    "\n",
    "# def get_tm_score(design: Pose, prediction: Pose) -> float:\n",
    "#     \"\"\"Calculate Ca-RMSD of prediction to design\"\"\"\n",
    "#     import pyrosetta\n",
    "#     tm_align = pyrosetta.rosetta.protocols.hybridization.TMalign()\n",
    "#     tm_align.apply(prediction, design)\n",
    "#     tm_score = tm_align.TMscore(len(design))\n",
    "#     return tm_score\n",
    "\n",
    "\n",
    "# def fold(list_seq, list_name, chain_id):\n",
    "#     ofile = open(\"./output/fastas/e0.fasta\", \"w+\")\n",
    "#     for i in range(1, len(list_seq), 2):\n",
    "#         ofile.write(\">\" + list_name[i] + \"\\n\" +list_seq[i] + \"\\n\")\n",
    "#     ofile.close()\n",
    "\n",
    "#     os.system(\"python3 ../run_superfold.py ./output/fastas/e0.fasta --overwrite --models 5\")\n",
    "#     sum_tm = 0\n",
    "#     count_tm = 0\n",
    "#     for i in range(0, len(list_name), 2):\n",
    "#         pdb1 = 0\n",
    "#         success = True\n",
    "#         while success:\n",
    "#             try:\n",
    "#                 pdb1 = pyrosetta.toolbox.rcsb.pose_from_rcsb(chain_id)\n",
    "#                 success = False\n",
    "#             except:\n",
    "#                 print('error')\n",
    "#                 success = True\n",
    "#         pdb2 = pyrosetta.pose_from_file('./output/' + list_name[i + 1] + '_model_5_ptm_seed_0_unrelaxed.pdb')\n",
    "#         tm = 0\n",
    "#         if (len(list_seq[i]) > len(list_seq[i + 1])):\n",
    "#             tm = get_tm_score(pdb1, pdb2)\n",
    "#         else:\n",
    "#             tm = get_tm_score(pdb2, pdb1)\n",
    "#         sum_tm += tm\n",
    "#         count_tm += 1\n",
    "\n",
    "#     return sum_tm / count_tm\n",
    "\n",
    "# def levenshtein_distance(seq_a, seq_b) -> int:\n",
    "#     \"\"\"\n",
    "#     :param seq_a: first sequence to compare.\n",
    "#     :param seq_b: second sequence to compare.\n",
    "#     :return: levenshtein distance between the two sequences.\n",
    "#     Calculate the levenshtein distance between two sequences.\n",
    "#     \"\"\"\n",
    "#     # https://en.wikipedia.org/wiki/Levenshtein_distance\n",
    "#     # initialize distance matrix\n",
    "#     distance_matrix = numpy.zeros(seq_a.shape[0] + 1, seq_b.shape[0] + 1)\n",
    "#     for id1 in range(len(seq_a) + 1):\n",
    "#         distance_matrix[id1][0] = id1\n",
    "#     for id2 in range(len(seq_b) + 1):\n",
    "#         distance_matrix[0][id2] = id2\n",
    "#     a = 0\n",
    "#     b = 0\n",
    "#     c = 0\n",
    "#     for id1 in range(1, len(seq_a) + 1):\n",
    "#         for id2 in range(1, len(seq_b) + 1):\n",
    "#             if torch.eq(seq_a[id1 - 1], seq_b[id2 - 1]):\n",
    "#                 distance_matrix[id1][id2] = distance_matrix[id1 - 1][id2 - 1]\n",
    "#             else:\n",
    "#                 a = distance_matrix[id1][id2 - 1]\n",
    "#                 b = distance_matrix[id1 - 1][id2]\n",
    "#                 c = distance_matrix[id1 - 1][id2 - 1]\n",
    "#                 if a <= b and a <= c:\n",
    "#                     distance_matrix[id1][id2] = a + 1\n",
    "#                 elif b <= a and b <= c:\n",
    "#                     distance_matrix[id1][id2] = b + 1\n",
    "#                 else:\n",
    "#                     distance_matrix[id1][id2] = c + 1\n",
    "#     levenshtein_distance = int(distance_matrix[id1][id2])\n",
    "#     return levenshtein_distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8afff11a",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "MAX_LENGTH = 600\n",
    "DATA_DIR = './data/PDB-2021AUG02.csv'\n",
    "\n",
    "import time\n",
    "import math\n",
    "\n",
    "\n",
    "def asMinutes(s):\n",
    "    m = math.floor(s / 60)\n",
    "    s -= m * 60\n",
    "    return '%dm %ds' % (m, s)\n",
    "\n",
    "\n",
    "def timeSince(since, percent):\n",
    "    now = time.time()\n",
    "    s = now - since\n",
    "    es = s / (percent)\n",
    "    rs = es - s\n",
    "    return '%s (- %s)' % (asMinutes(s), asMinutes(rs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "65a99e1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_dataset(max_length, data_dir=''):\n",
    "    print (\"loading dataset...\")\n",
    "    data = pd.read_csv(data_dir)\n",
    "    dirname = data_dir.split('.')\n",
    "    dirfilename = (dirname[1].split('/'))[-1]\n",
    "    print(dirfilename)\n",
    "    print(data.keys())\n",
    "    lines = list(set((x[0:4], y) for x, y in zip(data['CHAINID'].tolist(), data['SEQUENCE'].tolist())))\n",
    "    lines = [l for l in lines if ('X' not in l[1])]\n",
    "    \n",
    "    lines = [l for l in lines if (len(l[1]) <= max_length)]\n",
    "    lines = [(l[0], tuple(l[1] + '0'*(MAX_LENGTH - len(l[1])))) for l in lines] # pad with 0\n",
    "    print(\"loaded {} lines in dataset\".format(len(lines)))\n",
    "    np.random.shuffle(lines) \n",
    "    return lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3e5e67c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading dataset...\n",
      "(505296, 8)\n",
      "(97616, 8)\n",
      "('MGSSHHHHHHSSGLEVLFQGPEENGAHTIANNHTDMMEVDGDVEIPSNKAVVLRGHESEVFICAWNPVSDLLASGSGDSTARIWNLSENSTSGPTQLVLRHCIREGGQDVPSNKDVTSLDWNSEGTLLATGSYDGFARIWTKDGNLASTLGQHKGPIFALKWNKKGNFILSAGVDKTTIIWDAHTGEAKQQFPFHSAPALDVDWQSNNTFASCSTDMCIHVCKLGQDRPIKTFQGHTNEVNAIKWDPTGNLLASCSDDMTLKIWSMKQDNCVHDLQAHNKEIYTIKWSPTGPGTNNPNANLMLASASFDSTVRLWDVDRGICIHTLTKHQEPVYSVAFSPDGRYLASGSFDKCVHIWNTQTGALVHSYRGTGGIFEVCWNAAGDKVGASASDGSVCVLDLRK------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', 'CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCECCHHHEEEECCCCCCEEEEEECCCCCEEEEEECCCEEEEEECCCCCCCCCEEEEEECCCCCCCCCCCCCCCEEEEEECCCCCEEEEEECCCEEEEEECCCCEEEEEEECCCCEEEEEECCCCCEEEEEECCCCEEEEECCCCCEEEEECCCCCCEEEEEECCCCEEEEEECCCEEEEEECCCCCCCEEEECCCCCEEEEEECCCCCEEEEEECCCCEEEEECCCCCCCEEECCCCCCEEEEEECCCCCCCCCCCCCCEEEEEECCCCEEEEECCCCEEEEEECCCCCCEEEEEECCCCCEEEEEECCCEEEEEECCCCCEEEEEECCCCEEEEEECCCCCEEEEEECCCCEEEEECCC------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1]], device='cuda:0'), tensor([[1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1]], device='cuda:0'))\n",
      "('MTATDNARQVTIIGAGLAGTLVARLLARNGWQVNLFERRPDPRIETGARGRSINLALAERGAHALRLAGLEREVLAEAVMMRGRMVHVPGTPPNLQPYGRDDSEVIWSINRDRLNRILLDGAEAAGASIHFNLGLDSVDFARQRLTLSNVSGERLEKRFHLLIGADGCNSAVRQAMASVVDLGEHLETQPHGYKELQITPEASAQFNLEPNALHIWPHGDYMCIALPNLDRSFTVTLFLHHQSPAAQPASPSFAQLVDGHAARRFFQRQFPDLSPMLDSLEQDFEHHPTGKLATLRLTTWHVGGQAVLLGDAAHPMVPFHGQGMNCALEDAVALAEHLQSAADNASALAAFTAQRQPDALAIQAMALENYVEMSSKVASPTYLLERELGQIMAQRQPTRFIPRYSMVTFSRLPYAQAMARGQIQEQLLKFAVANHSDLTSINLDAVEHEVTRCLPPLSHLS-------------------------------------------------------------------------------------------------------------------------------------------', 'CCCCCCCCEEEEECCCHHHHHHHHHHHHCCCEEEEECCCCCCCCCCCCCCCCCEEEECHHHHHHHHHHCCHHHHHHHEEEECEEEECCCCCCCEEEECCCCCCCCEEEEEHHHHHHHHHHHHHHHCCEEECCCEEEEEECCCCEEEEECCCCCEEEEECCCEEECCCCCCHHHHHHHHCCCCCCEEEECCEEEEEEEECHHHHHHCCCCCCCEEEEEECCEEEEEEECCCCCEEEEEEEECCCCCCCCCCCCHHHCCCHHHHHHHHHHHCCCCHHHCCCHHHHHHHCCCECCEEEECCCCEECCCEEECHHHCCCCCCCCCCHHHHHHHHHHHHHHHHHHCCCHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHCCCCCCHHHHHHHHHHHHHHHCCCCCCCHHHHHHHCCCCHHHHHHHHHHHHHHHHHHHCCCCCHHHCCHHHHHHHHHHHCCCCCCCC-------------------------------------------------------------------------------------------------------------------------------------------', tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [1, 0, 0,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1]], device='cuda:0'), tensor([[1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1]], device='cuda:0'))\n",
      "('GSMSEQSICQARAAVMVYDDANKKWVPAGGSTGFSRVHIYHHTGNNTFRVVGRKIQDHQVVINCAIPKGLKYNQATQTFHQWRDARQVYGLNFGSKEDANVFASAMMHALEVL-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', 'CCCCEEECCCEEEEEEEEECCCCEEEEHHHCCCEEEEEEEEECCCCEEEEEEEECCCCCEEEEEECCCCCCCEEEECCEEEEECCCCEEEEEECCHHHHHHHHHHHHHHHHHC-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1]], device='cuda:0'), tensor([[1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1]], device='cuda:0'))\n",
      "('MSAKDERAREILRGFKLNWMNLRDAETGKILWQGTEDLSVPGVEHEARVPKKILKCKAVSRELNFSSTEQMEKFRLEQKVYFKGQCLEEWFFEFGFVIPNSTNTWQSLIEAAPESQMMPASVLTGNVIIETKFFDDDLLVSTSRVRLFYV------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', 'CCHHHHHHHHHHHCEEEEEEEEEECCCCCEEEEECCCCCCCCCEEEEEECHHHHHCCEEEEEEEEEECCCEEEEEEEEEEEECCEEEEEEEEEEEEECCCEEEEEEEEEECCCCCCCCCHHHHCCCEEEEEEEEECCEEEEEEEEEEEEC------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [1, 0, 0,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1]], device='cuda:0'), tensor([[1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        [0, 0, 1, 0],\n",
      "        ...,\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1]], device='cuda:0'))\n",
      "('GITIPRNPGCPNSEDKNFPRTVMVNLNIHNRNTNTNPKRSSDYYNRSTSPWNLHRNEDPERYPSVIWEAKCRHLGCINADGNVDYHMNSVPIQQEILVLRREPPHCPNSFRLEKILVSVGCTCVTPIVHHVA------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', 'CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCHHHHCCCCEEEEEEECCCEECCEEEEEEECCCCEECCCCCEECCCEEEEEEEEEEEEEECCCCCCCCEEEEEEEEEEEEEEECCCCCCCC------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------', tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1],\n",
      "        [0, 0, 0,  ..., 0, 0, 1]], device='cuda:0'), tensor([[1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        [1, 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1],\n",
      "        [0, 0, 0, 1]], device='cuda:0'))\n"
     ]
    }
   ],
   "source": [
    "class ProteinDataset(Dataset):\n",
    "    def __init__(self, data_file):\n",
    "        print (\"loading dataset...\")\n",
    "        self.data = pd.read_csv(data_file)\n",
    "        print(self.data.shape)\n",
    "        self.data = self.data.drop_duplicates(subset=['SEQUENCE'])\n",
    "        self.data = self.data[~(self.data['s4pred_truth'].str.contains('X'))]\n",
    "        print(self.data.shape)\n",
    "        sequence_map = {}\n",
    "        for i, c in enumerate('ARNDCQEGHILKMFPSTWYV-'):\n",
    "            sequence_map[c] = i;\n",
    "        \n",
    "        structure_map = {}\n",
    "        for i, c in enumerate('CEH-'):\n",
    "            structure_map[c] = i;\n",
    "        \n",
    "        self.data['SEQUENCE'] = self.data['SEQUENCE'].str.pad(width=600, side='right', fillchar='-')\n",
    "        self.data['SEQUENCE_encode'] = self.data['SEQUENCE'].map(lambda a: self.transform(a, sequence_map))\n",
    "        self.data['s4pred_truth'] = self.data['s4pred_truth'].str.pad(width=600, side='right', fillchar='-')\n",
    "        self.data['s4pred_truth_encode'] = self.data['s4pred_truth'].map(lambda a: self.transform(a, structure_map))\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.data.shape[0]\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "        sample = self.data.iloc[idx]['SEQUENCE']\n",
    "        structure = self.data.iloc[idx]['s4pred_truth']\n",
    "        sample_e = self.data.iloc[idx]['SEQUENCE_encode']\n",
    "        structure_e = self.data.iloc[idx]['s4pred_truth_encode']\n",
    "        return (sample, structure, sample_e, structure_e)\n",
    "\n",
    "    def transform(self, a, source):\n",
    "        indexes = [source[c] for c in list(a)]\n",
    "        return F.one_hot(torch.tensor(indexes, dtype=torch.long, device=device), num_classes=len(source))\n",
    "        \n",
    "        \n",
    "\n",
    "pdb = ProteinDataset('./data/PDB-2021AUG02_noX_dssp_ss.csv')\n",
    "for i in range(5):\n",
    "    sample = pdb[i]\n",
    "    print(sample)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b4818ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# class Lang:\n",
    "#     def __init__(self, name):\n",
    "#         self.name = name\n",
    "#         self.char2index = {}\n",
    "#         self.char2count = {}\n",
    "#         self.char2word = {}\n",
    "#         self.n_chars = 0\n",
    "\n",
    "#     def addSequence(self, seq):\n",
    "#         for c in list(seq):\n",
    "#             self.addChar(c)\n",
    "\n",
    "#     def addChar(self, c):\n",
    "#         if c not in self.char2index:\n",
    "#             self.char2index[c] = self.n_chars\n",
    "#             self.char2count[c] = 1\n",
    "#             self.char2word[self.n_chars] = c\n",
    "#             self.n_chars += 1\n",
    "#         else:\n",
    "#             self.char2count[c] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a3ce36a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading dataset...\n",
      "PDB-2021AUG02\n",
      "Index(['CHAINID', 'DEPOSITION', 'RESOLUTION', 'HASH', 'CLUSTER', 'SEQUENCE'], dtype='object')\n",
      "loaded 242971 lines in dataset\n",
      "('6jm9', ('A', 'K', 'T', 'R', 'S', 'S', 'R', 'A', 'G', 'L', 'Q', 'F', 'P', 'V', 'G', 'R', 'V', 'H', 'R', 'L', 'L', 'R', 'K', 'G', 'N', 'Y', 'A', 'E', 'R', 'V', 'G', 'A', 'G', 'A', 'P', 'V', 'Y', 'L', 'A', 'A', 'V', 'L', 'E', 'Y', 'L', 'T', 'A', 'E', 'I', 'L', 'E', 'L', 'A', 'G', 'N', 'A', 'A', 'R', 'D', 'N', 'K', 'K', 'T', 'R', 'I', 'I', 'P', 'R', 'H', 'L', 'Q', 'L', 'A', 'V', 'R', 'N', 'D', 'E', 'E', 'L', 'N', 'K', 'L', 'L', 'G', 'R', 'V', 'T', 'I', 'A', 'Q', 'G', 'G', 'V', 'L', 'P', 'N', 'I', 'Q', 'S', 'V', 'L', 'L', 'P', 'K', 'K', 'T', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0'))\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 23.64 GiB total capacity; 22.41 GiB already allocated; 1.38 MiB free; 22.62 GiB reserved in total by PyTorch)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_82322/3985219100.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     19\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindexes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlong\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 21\u001b[0;31m \u001b[0minput_lang\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_lang\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpairs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprepare_data\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     22\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_lang\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchar2word\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mtensorsFromPair\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpair\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[0;32m/tmp/ipykernel_82322/3985219100.py\u001b[0m in \u001b[0;36mprepare_data\u001b[0;34m(max_len, data_dir)\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlines\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\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      9\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0ms\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mlines\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m         \u001b[0mretlines\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mone_hot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensorFromSequence\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlang\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_classes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlang\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mn_chars\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\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     11\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mretlines\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\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     12\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mlang\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlang\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0ms\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mretlines\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[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 23.64 GiB total capacity; 22.41 GiB already allocated; 1.38 MiB free; 22.62 GiB reserved in total by PyTorch)"
     ]
    }
   ],
   "source": [
    "# def prepare_data(max_len=MAX_LENGTH, data_dir=DATA_DIR):\n",
    "#     lines = load_dataset(max_len, data_dir)\n",
    "#     lang = Lang(\"PDB\")\n",
    "#     lang.addSequence(\"ARNDCQEGHILKMFPSTWYV0\")\n",
    "#     for line in lines:\n",
    "#         lang.addSequence(line[1])\n",
    "#     retlines = []\n",
    "#     print(lines[0])\n",
    "#     for s in lines:\n",
    "#         retlines.append((s[0], F.one_hot(tensorFromSequence(lang, s), num_classes=lang.n_chars).float()))\n",
    "#     print((retlines[0])[1].size())\n",
    "#     return (lang, lang, [[s, s] for s in retlines])\n",
    "\n",
    "# def indexesFromSequence(lang, sequence):\n",
    "#     return [lang.char2index[c] for c in list(sequence[1])]\n",
    "\n",
    "# def tensorFromSequence(lang, sequence):\n",
    "#     indexes = indexesFromSequence(lang, sequence)\n",
    "#     return torch.tensor(indexes, dtype=torch.long, device=device)\n",
    "\n",
    "# input_lang, output_lang, pairs = prepare_data()\n",
    "# print(input_lang.char2word)\n",
    "# def tensorsFromPair(pair):\n",
    "#     input_tensor = tensorFromSequence(input_lang, pair[0])\n",
    "#     target_tensor = tensorFromSequence(output_lang, pair[1])\n",
    "#     return (input_tensor, target_tensor)\n",
    "\n",
    "# def decodeTensor(lang, t):\n",
    "#     oneHot = torch.argmax((t[1]).view(MAX_LENGTH, lang.n_chars), dim=1)\n",
    "#     return (''.join([lang.char2word[i.item()] for i in oneHot])).replace('0', '')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34543932",
   "metadata": {},
   "outputs": [],
   "source": [
    "# class EncoderRNN(nn.Module):\n",
    "#     def __init__(self, input_size, hidden_size, num_layers=1):\n",
    "#         super(EncoderRNN, self).__init__()\n",
    "#         self.hidden_size = hidden_size\n",
    "#         self.num_layers = num_layers\n",
    "#         self.embedding = nn.Embedding(input_size, hidden_size)\n",
    "#         self.gru = nn.GRU(hidden_size, hidden_size, num_layers=num_layers)\n",
    "\n",
    "#     def forward(self, input, hidden):\n",
    "#         embedded = self.embedding(input).view(1, 1, -1)\n",
    "#         output = embedded\n",
    "#         if hidden.shape[0] != self.num_layers:\n",
    "#             hidden = hidden.repeat(self.num_layers, 1, 1)\n",
    "#         output, hidden = self.gru(output, hidden)\n",
    "#         return output, hidden\n",
    "\n",
    "#     def initHidden(self):\n",
    "#         return torch.zeros(1, 1, self.hidden_size, device=device)\n",
    "\n",
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, input_size=21*MAX_LENGTH, hidden_size=100):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.layers = nn.Sequential(nn.Linear(input_size, 10000), \n",
    "                                    nn.LeakyReLU(), \n",
    "                                    nn.Linear(10000, 5000), \n",
    "                                    nn.LeakyReLU(),\n",
    "                                    nn.Linear(5000, 1000), \n",
    "                                    nn.LeakyReLU(),\n",
    "                                    nn.Linear(1000, 500), \n",
    "                                    nn.LeakyReLU(),\n",
    "                                    nn.Linear(500, 300), \n",
    "                                    nn.LeakyReLU(),\n",
    "                                    nn.Linear(300, hidden_size), \n",
    "                                    nn.LeakyReLU())\n",
    "\n",
    "    def forward(self, input):\n",
    "        output = self.layers(input)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1e61ebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size=100, output_size=21*MAX_LENGTH):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.layers = nn.Sequential(nn.Linear(hidden_size, 300), \n",
    "                                    nn.LeakyReLU(), \n",
    "                                    nn.Linear(500, 1000), \n",
    "                                    nn.LeakyReLU(),\n",
    "                                    nn.Linear(1000, 5000), \n",
    "                                    nn.LeakyReLU(),\n",
    "                                    nn.Linear(5000, 10000),\n",
    "                                    nn.LeakyReLU())\n",
    "        self.out = nn.Linear(10000, output_size)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, input):\n",
    "        output = self.layers(input)\n",
    "        output = self.sigmoid(self.out(output))\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "464ecbfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Created on Thu Jun  4 13:44:17 2020\n",
    "@author:\n",
    "    Lewis Moffat\n",
    "    Bioinformatics Group - Comp. Sci. Dep., University College London (UCL)\n",
    "    Github: CraftyColossus\n",
    "Inference Only Version of S4PRED - Single Sequence Secondary Structure Pred\n",
    "This is culled down to exclude the various DropConnect/Dropout etc. from the \n",
    "training methods so that it is more clear.\n",
    "If you'd like a training version of the model please raise an issue or submit a PR.\n",
    "The AWD-GRU training script model was a tweak on the offical Salesforce \n",
    "AWD-LSTM (https://github.com/salesforce/awd-lstm-lm/). It needed to be adapted to \n",
    "take multiple layers of RNNs. \n",
    "\"\"\"\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class ResidueEmbedding(nn.Embedding):\n",
    "    def __init__(self, vocab_size=21, embed_size=128, padding_idx=None):\n",
    "        super().__init__(vocab_size, embed_size, padding_idx=padding_idx)\n",
    "\n",
    "        \n",
    "        \n",
    "class GRUnet(nn.Module):\n",
    "    def __init__(self,lstm_hdim=1024, embed_size=128, num_layers=3,bidirectional=True,lstm=False,outsize=3):\n",
    "        super().__init__()\n",
    "        \"\"\"\n",
    "            This version of the model has all the bells & whistles (e.g. \n",
    "            dropconnect) ripped out so its slimmed down for inference\n",
    "            \n",
    "        \"\"\"\n",
    "        \n",
    "        self.lstm_hdim = lstm_hdim\n",
    "        self.embed=ResidueEmbedding(vocab_size=22, embed_size=embed_size, padding_idx=21)\n",
    "        self.lstm = nn.GRU(128, 1024, num_layers=3, bidirectional=True, batch_first=True,dropout=0.0)\n",
    "        self.outlayer = nn.Linear(lstm_hdim*2, outsize)\n",
    "        self.finalact=F.log_softmax\n",
    "\n",
    "    \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "            Assumes a batch size of one currently but can be changed\n",
    "        \"\"\"\n",
    "        x=self.embed(x)\n",
    "        x, _ = self.lstm(x)\n",
    "        x=self.outlayer(x)\n",
    "        x=self.finalact(x,dim=-1)\n",
    "        return x.squeeze()        \n",
    "        \n",
    "        \n",
    "class S4PRED(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \"\"\"\n",
    "            This loads the ensemble of models in a lazy way but its clear and \n",
    "            leaves the weight loading out of the run_model script. \n",
    "        \"\"\"\n",
    "                                            \n",
    "        # Manually listing for clarity and hot swapping in future\n",
    "        self.model_1=GRUnet()\n",
    "        self.model_2=GRUnet()\n",
    "        self.model_3=GRUnet()\n",
    "        self.model_4=GRUnet()\n",
    "        self.model_5=GRUnet()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        y_1=self.model_1(x)\n",
    "        y_2=self.model_2(x)\n",
    "        y_3=self.model_3(x)\n",
    "        y_4=self.model_4(x)\n",
    "        y_5=self.model_5(x)\n",
    "        y_out=y_1*0.2+y_2*0.2+y_3*0.2+y_4*0.2+y_5*0.2\n",
    "        return y_out\n",
    "    \n",
    "s4pred = S4PRED().to(device)\n",
    "s4pred.requires_grad=True\n",
    "scriptdir = '../s4pred'\n",
    "weight_files=['/weights/weights_1.pt',\n",
    "              '/weights/weights_2.pt',\n",
    "              '/weights/weights_3.pt',\n",
    "              '/weights/weights_4.pt',\n",
    "              '/weights/weights_5.pt']\n",
    "\n",
    "# Manually listing for clarity and hot swapping in future\n",
    "# Inelegant, ugly ugly, to be cleaned up in the future\n",
    "s4pred.model_1.load_state_dict(torch.load(scriptdir + weight_files[0], map_location=lambda storage, loc: storage))\n",
    "s4pred.model_2.load_state_dict(torch.load(scriptdir + weight_files[1], map_location=lambda storage, loc: storage))\n",
    "s4pred.model_3.load_state_dict(torch.load(scriptdir + weight_files[2], map_location=lambda storage, loc: storage))\n",
    "s4pred.model_4.load_state_dict(torch.load(scriptdir + weight_files[3], map_location=lambda storage, loc: storage))\n",
    "s4pred.model_5.load_state_dict(torch.load(scriptdir + weight_files[4], map_location=lambda storage, loc: storage))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7faef832",
   "metadata": {},
   "outputs": [],
   "source": [
    "# class Model(nn.Module):\n",
    "#     def __init__(self, input_size, encoding_size, hidden=[], h_act=nn.ReLU(), out_act=nn.Tanh()):\n",
    "#         super(Model, self).__init__()\n",
    "#         self.encoder = Encoder(input_size, encoding_size, hidden, h_act, out_act)\n",
    "#         self.decoder = Decoder(encoding_size, input_size, hidden, h_act)\n",
    "    \n",
    "#     def forward(self, x):\n",
    "#         seq_len = x.shape[0]\n",
    "#         x = self.encoder(x)\n",
    "#         x = self.decoder(x, seq_len)\n",
    "#         return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b07819ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def decodeTensorValue(lang, t):\n",
    "    # torch.argmax((t[1]).view(MAX_LENGTH, lang.n_chars), dim=1)\n",
    "    input = (t[1]).view(MAX_LENGTH, lang.n_chars)\n",
    "    *_, n = input.shape\n",
    "    input = nn.functional.softmax(100 * input, dim=-1)\n",
    "    indices = torch.linspace(0, 1, n).to(device)\n",
    "    result = torch.sum((n - 1) * input * indices, dim=-1)\n",
    "    return result\n",
    "\n",
    "def train(input_tensor, target_tensor, chain_id, encoder, decoder, encoder_optimizer, decoder_optimizer, ss_optimizer, criterion, iter, e):\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "    ss_optimizer.zero_grad()\n",
    "    \n",
    "    encoder_hidden = encoder(input_tensor)\n",
    "    decoder_output = decoder(encoder_hidden)\n",
    "    \n",
    "    inp = decodeTensorValue(input_lang, (chain_id, input_tensor)).int()\n",
    "    out = decodeTensorValue(input_lang, (chain_id, decoder_output)).int()\n",
    "#     list_name = ['r', 'f']\n",
    "#     ofile = open(\"./output/fastas/i0.fas\", \"w+\")\n",
    "#     ofile.write(\">\" + list_name[1] + \"\\n\" + list_seqs[1] + \"\\n\")\n",
    "#     ofile.close()\n",
    "#     os.system(\"python3 ../s4pred/run_model.py --outfmt fas ./output/fastas/i0.fas > ./output/i0ss.fas\")\n",
    "#     f1 = open('./output/i0ss.fas')\n",
    "#     s1 = f1.readlines()[2]\n",
    "#     pose = 0\n",
    "#     success = True\n",
    "#     while success:\n",
    "#         try:\n",
    "#             pose = pyrosetta.toolbox.rcsb.pose_from_rcsb(chain_id)\n",
    "#             success = False\n",
    "#         except:\n",
    "#             print(chain_id)\n",
    "#             success = True\n",
    "#     ss = pyrosetta.rosetta.core.scoring.dssp.Dssp(pose) \n",
    "#     s2 = \"\".join([ss.get_dssp_secstruct(pos) for pos in range(1, len(pose.sequence()))])\n",
    "#     s2.replace('L', 'C')\n",
    "    # loss = torch.mul(loss, levenshtein_distance(s1, s2))\n",
    "\n",
    "    ss_inp = s4pred(inp[None, :])\n",
    "    ss_out = s4pred(out[None, :])\n",
    "        \n",
    "    loss = criterion(input_tensor, decoder_output)\n",
    "    \n",
    "    \n",
    "#     if (iter % 1000 == 1):\n",
    "#         list_seqs = [decodeTensor(input_lang, (chain_id, input_tensor)), decodeTensor(input_lang, (chain_id, decoder_output))]\n",
    "#         list_name = ['r' + str(e), 'f' + str(e) + '_' + chain_id + '_']\n",
    "#         alpha = fold(list_seqs, list_name, chain_id)\n",
    "#         print(alpha)\n",
    "#         alpha = 1.0 / alpha\n",
    "    \n",
    "#     loss = torch.mul(loss, alpha)\n",
    "    loss.backward()\n",
    "\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "    #ss_optimizer.step()\n",
    "\n",
    "    return decoder_output, loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f6f3e21",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainIters(encoder, decoder, n_iters, epochs, print_every=1, plot_every=1, learning_rate=1):\n",
    "    start = time.time()\n",
    "    with open('./output/seqs.txt', 'w+') as f:\n",
    "        print(\"Sequences: out, in\\n\", file=f)\n",
    "    with open('./output/losses.txt', 'w+') as f:\n",
    "        print(\"Losses\\n\", file=f)\n",
    "    plot_losses = []\n",
    "    losses = []\n",
    "    print_loss_total = 0  # Reset every print_every\n",
    "    plot_loss_total = 0  # Reset every plot_every\n",
    "\n",
    "    encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "    decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "    ss_optimizer = optim.SGD(s4pred.parameters(), lr=learning_rate)\n",
    "    criterion = nn.BCELoss()\n",
    "    \n",
    "    for e in range(1, epochs + 1):\n",
    "        if e % print_every == 0:\n",
    "            print_loss_avg = print_loss_total / n_iters\n",
    "            print_loss_total = 0\n",
    "            print('%s (%d %d%%) Loss: %.4f' % (timeSince(start, e / epochs),\n",
    "                                         e, e / epochs * 100, print_loss_avg))\n",
    "#             f = open('./output/' + str(e) + '_e_weights.pt', 'w+')\n",
    "#             g = open('./output/' + str(e) + '_d_weights.pt', 'w+')\n",
    "            with open('./output/encoder.pt', 'w+') as f:\n",
    "                print({\n",
    "                    'epoch': e,\n",
    "                    'model_state_dict': encoder.state_dict(),\n",
    "                    'optimizer_state_dict': encoder_optimizer.state_dict(),\n",
    "                    'loss': print_loss_avg,\n",
    "                }, file=f)\n",
    "            with open('./output/decoder.pt', 'w+') as f:\n",
    "                print({\n",
    "                    'epoch': e,\n",
    "                    'model_state_dict': decoder.state_dict(),\n",
    "                    'optimizer_state_dict': decoder_optimizer.state_dict(),\n",
    "                    'loss': print_loss_avg,\n",
    "                }, file=f)\n",
    "            \n",
    "        training_pairs = [random.choice(pairs) for i in range(n_iters)]\n",
    "        for iter in range(1, n_iters + 1):\n",
    "            training_pair = training_pairs[iter - 1]\n",
    "            input_tensor = training_pair[0][1].view(-1, MAX_LENGTH * input_lang.n_chars)\n",
    "            target_tensor = training_pair[1][1].view(-1, MAX_LENGTH * input_lang.n_chars)\n",
    "            chain_id = training_pair[0][0]\n",
    "\n",
    "            output_tensor, loss = train(input_tensor, target_tensor, chain_id, encoder,\n",
    "                         decoder, encoder_optimizer, decoder_optimizer, ss_optimizer, criterion, iter, e)\n",
    "            print_loss_total += loss\n",
    "            plot_loss_total += loss\n",
    "            with torch.no_grad():\n",
    "                if iter % 1000 == 0:\n",
    "                    with open('./output/seqs.txt', 'a') as f:\n",
    "                        print('\\tIteration %d Loss: %.4f' % (iter, loss))\n",
    "                        print(decodeTensor(input_lang, (chain_id, output_tensor)), file=f)\n",
    "                        print(decodeTensor(input_lang, training_pair[1]), file=f)\n",
    "    \n",
    "                if iter % 1000 == 0:\n",
    "                    losses.append(loss)\n",
    "                    fig, ax = plt.subplots(1, figsize=(15,10))\n",
    "                    ax.plot(range(1, len(losses) + 1), losses)\n",
    "                    ax.grid()\n",
    "                    ax.set_title('Losses of Autoencoder')\n",
    "                    ax.set_xlabel('Iteration (every 1000)')\n",
    "                    ax.set_ylabel('Loss')\n",
    "                    fig.savefig(\"./output/losses_iter.png\")\n",
    "                    with open('./output/losses.txt', 'a') as f:\n",
    "                        print('\\tIteration %d Loss: %.4f' % (iter, loss), file=f)\n",
    "\n",
    "        if e % plot_every == 0:\n",
    "            plot_loss_avg = plot_loss_total / n_iters\n",
    "            plot_losses.append(plot_loss_avg)\n",
    "            plot_loss_total = 0\n",
    "            fig, ax = plt.subplots(1, figsize=(15,10))\n",
    "            ax.plot(range(1, len(plot_losses) + 1), plot_losses)\n",
    "            ax.grid()\n",
    "            ax.set_title('Losses of Autoencoder')\n",
    "            ax.set_xlabel('Epoch')\n",
    "            ax.set_ylabel('Loss')\n",
    "            fig.savefig(\"./output/losses_epoch.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343f8a60",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder1 = nn.DataParallel(EncoderRNN()).to(device)\n",
    "decoder1 = nn.DataParallel(DecoderRNN()).to(device)\n",
    "\n",
    "trainIters(encoder1, decoder1, 100000, 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d1b9e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_losses = []\n",
    "\n",
    "with open('./output/seqs.txt', 'r') as f:\n",
    "    lines = f.readlines()\n",
    "    for i in range(2, len(lines), 2):\n",
    "        ss_inp = s4pred(lines[i])\n",
    "        ss_out = s4pred(lines[i+1])\n",
    "        plot_losses.append(nn.MSELoss(ss_inp, ss_out))\n",
    "        print(nn.MSELoss(ss_inp, ss_out))\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(15,10))\n",
    "ax.plot(range(1, len(plot_losses) + 1), plot_losses)\n",
    "ax.grid()\n",
    "ax.set_title('Losses of Autoencoder')\n",
    "ax.set_xlabel('Epoch')\n",
    "ax.set_ylabel('Loss')\n",
    "fig.savefig(\"./output/losses_test.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c340c5e1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}