notebooks/10_generate_protein_sequences.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 220
    },
    "colab_type": "code",
    "id": "HCWw1IgNHvhm",
    "outputId": "fcb1a352-1ca8-4708-c3af-11bb0f9a3c57"
   },
   "source": [
    "## Summary\n",
    "\n",
    "**Parameters**\n",
    "\n",
    "- `SEQUENCE_GENERATION_METHOD`\n",
    "- `STRUCTURE_ID`\n",
    "- `SLURM_ARRAY_TASK_ID`\n",
    "\n",
    "\n",
    "**Notes:**\n",
    "\n",
    "- `astar` method should be given >= 64G memory in order to generate 200k sequences.\n",
    "- `astar` cannot be ran in parallel.\n",
    "\n",
    "**SLURM scripts**\n",
    "\n",
    "```bash\n",
    "export STRUCTURE_ID=\"4beuA02\"\n",
    "\n",
    "SEQUENCE_GENERATION_METHOD=\"astar\" sbatch --mem 64G --time 72:00:00 ./scripts/run_notebook_gpu.sh $(realpath notebooks/10_generate_protein_sequences.ipynb)\n",
    "\n",
    "SEQUENCE_GENERATION_METHOD=\"expectimax\" sbatch --mem 32G --time 24:00:00 --array=1-3 ./scripts/run_notebook_gpu.sh $(realpath notebooks/10_generate_protein_sequences.ipynb)\n",
    "\n",
    "SEQUENCE_GENERATION_METHOD=\"randexpectimax\" sbatch --mem 32G --time 24:00:00 --array=1-3 ./scripts/run_notebook_gpu.sh $(realpath notebooks/10_generate_protein_sequences.ipynb)\n",
    "```\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eoBMUoW2Hvhp",
    "toc-hr-collapsed": true,
    "toc-nb-collapsed": true
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import heapq\n",
    "import io\n",
    "import json\n",
    "import os\n",
    "import shutil\n",
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "from IPython.display import HTML\n",
    "from IPython.display import display\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import kmtools.sci_tools\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import proteinsolver\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n",
    "import torch\n",
    "import torch_geometric\n",
    "from kmbio import PDB\n",
    "from torch_geometric.data import Batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "NOTEBOOK_NAME = \"generate_protein_sequences\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab1/strokach/workspace/proteinsolver/notebooks/generate_protein_sequences')"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NOTEBOOK_PATH = Path(NOTEBOOK_NAME).resolve()\n",
    "NOTEBOOK_PATH.mkdir(exist_ok=True)\n",
    "NOTEBOOK_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "UNIQUE_ID = \"191f05de\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "BEST_STATE_FILES = {\n",
    "    #\n",
    "    \"191f05de\": \"protein_train/191f05de/e53-s1952148-d93703104.state\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'4beuA02'"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# STRUCTURE_ID = os.getenv(\"STRUCTURE_ID\", \"1n5uA03\")\n",
    "# STRUCTURE_ID = os.getenv(\"STRUCTURE_ID\", \"4z8jA00\")\n",
    "# STRUCTURE_ID = os.getenv(\"STRUCTURE_ID\", \"4unuA00\")\n",
    "STRUCTURE_ID = os.getenv(\"STRUCTURE_ID\", \"4beuA02\")\n",
    "\n",
    "STRUCTURE_ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab1/strokach/workspace/proteinsolver/data/inputs/4beuA02.pdb')"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "STRUCTURE_FILE = Path(\n",
    "    os.getenv(\n",
    "        \"STRUCTURE_FILE\",\n",
    "        NOTEBOOK_PATH.parent.parent / \"proteinsolver\" / \"data\" / \"inputs\" / f\"{STRUCTURE_ID}.pdb\",\n",
    "    )\n",
    ").resolve()\n",
    "\n",
    "STRUCTURE_FILE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.25"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "min_expected_proba_preset = {\n",
    "    #\n",
    "    \"1n5uA03\": 0.20,\n",
    "    \"4z8jA00\": 0.29,\n",
    "    \"4unuA00\": 0.25,\n",
    "    \"4beuA02\": 0.25,\n",
    "}\n",
    "MIN_EXPECTED_PROBA = min_expected_proba_preset.get(STRUCTURE_ID, 0.15)\n",
    "\n",
    "MIN_EXPECTED_PROBA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'astar'"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SEQUENCE_GENERATION_METHOD = os.getenv(\"SEQUENCE_GENERATION_METHOD\", \"astar\")\n",
    "\n",
    "assert SEQUENCE_GENERATION_METHOD in (\"astar\", \"expectimax\", \"randexpectimax\", \"root2expectimax\", \"root10expectimax\")\n",
    "SEQUENCE_GENERATION_METHOD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "START_FILE_INDEX = int(os.getenv(\"SLURM_ARRAY_TASK_ID\", 0)) * 1000\n",
    "\n",
    "START_FILE_INDEX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0'"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def design_sequence(net, data, random_position=False, value_selection_strategy=\"map\", num_categories=None):\n",
    "    assert value_selection_strategy in (\"map\", \"multinomial\", \"ref\")\n",
    "\n",
    "    if num_categories is None:\n",
    "        num_categories = data.x.max().item()\n",
    "\n",
    "    if hasattr(data, \"batch\"):\n",
    "        batch_size = data.batch.max().item() + 1\n",
    "    else:\n",
    "        print(\"Defaulting to batch size of one.\")\n",
    "        batch_size = 1\n",
    "\n",
    "    if value_selection_strategy == \"ref\":\n",
    "        x_ref = data.y if hasattr(data, \"y\") and data.y is not None else data.x\n",
    "\n",
    "    x = torch.ones_like(data.x) * num_categories\n",
    "    x_proba = torch.zeros_like(x).to(torch.float)\n",
    "    index_array_ref = torch.arange(x.size(0))\n",
    "    mask_ref = x == num_categories\n",
    "    while mask_ref.any():\n",
    "        output = net(x, data.edge_index, data.edge_attr)\n",
    "        output_proba_ref = torch.softmax(output, dim=1)\n",
    "        output_proba_max_ref, _ = output_proba_ref.max(dim=1)\n",
    "\n",
    "        for i in range(batch_size):\n",
    "            mask = mask_ref\n",
    "            if batch_size > 1:\n",
    "                mask = mask & (data.batch == i)\n",
    "\n",
    "            index_array = index_array_ref[mask]\n",
    "            max_probas = output_proba_max_ref[mask]\n",
    "\n",
    "            if random_position:\n",
    "                selected_residue_subindex = torch.randint(0, max_probas.size(0), (1,)).item()\n",
    "                max_proba_index = index_array[selected_residue_subindex]\n",
    "            else:\n",
    "                selected_residue_subindex = max_probas.argmax().item()\n",
    "                max_proba_index = index_array[selected_residue_subindex]\n",
    "\n",
    "            assert x[max_proba_index] == num_categories\n",
    "            assert x_proba[max_proba_index] == 0\n",
    "            category_probas = output_proba_ref[max_proba_index]\n",
    "\n",
    "            if value_selection_strategy == \"map\":\n",
    "                chosen_category_proba, chosen_category = category_probas.max(dim=0)\n",
    "            elif value_selection_strategy == \"multinomial\":\n",
    "                chosen_category = torch.multinomial(category_probas, 1).item()\n",
    "                chosen_category_proba = category_probas[chosen_category]\n",
    "            else:\n",
    "                assert value_selection_strategy == \"ref\"\n",
    "                chosen_category = x_ref[max_proba_index]\n",
    "                chosen_category_proba = category_probas[chosen_category]\n",
    "\n",
    "            assert chosen_category != num_categories\n",
    "            x[max_proba_index] = chosen_category\n",
    "            x_proba[max_proba_index] = chosen_category_proba\n",
    "        mask_ref = x == num_categories\n",
    "        del output, output_proba_ref, output_proba_max_ref\n",
    "    return x.cpu(), x_proba.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from dataclasses import field\n",
    "from typing import Any\n",
    "\n",
    "\n",
    "def load_heap_dump(heap_file):\n",
    "    try:\n",
    "        pfile = pq.ParquetFile(heap_file)\n",
    "        heap = []\n",
    "        for row_group in pfile.num_row_groups:\n",
    "            df = pfile.read_row_group(row_group).to_parquet()\n",
    "            heap = heap + [\n",
    "                proteinsolver.utils.PrioritizedItem(\n",
    "                    tup.p, torch.tensor(tup.x, dtype=torch.int8), tup.total_proba, tup.total_logproba\n",
    "                )\n",
    "                for tup in df.itertuples()\n",
    "            ]\n",
    "    except Exception as e:\n",
    "        print(f\"Encountered error loading heap file '{heap_file}': '{e}'.\")\n",
    "\n",
    "    heap_file_bak = heap_file.with_suffix(\".parquet.bak\")\n",
    "    try:\n",
    "        pfile = pq.ParquetFile(heap_file)\n",
    "        heap = []\n",
    "        for row_group in pfile.num_row_groups:\n",
    "            df = pfile.read_row_group(row_group).to_parquet()\n",
    "            heap = heap + [\n",
    "                proteinsolver.utils.PrioritizedItem(\n",
    "                    tup.p, torch.tensor(tup.x, dtype=torch.int8), tup.total_proba, tup.total_logproba\n",
    "                )\n",
    "                for tup in df.itertuples()\n",
    "            ]\n",
    "    except Exception as e:\n",
    "        print(f\"Encountered error loading heap file '{heap_file_bak}': '{e}'.\")\n",
    "\n",
    "\n",
    "def update_heap_dump(heap_file, heap):\n",
    "    try:\n",
    "        shutil.copy2(heap_file, heap_file.with_suffix(\".parquet.bak\"))\n",
    "    except FileNotFoundError:\n",
    "        pass\n",
    "\n",
    "    df = pd.DataFrame(\n",
    "        [\n",
    "            {\"p\": pi.p, \"x\": pi.x.data.tolist(), \"total_proba\": pi.total_proba, \"total_logproba\": pi.total_logproba}\n",
    "            for pi in heap\n",
    "        ]\n",
    "    )\n",
    "    chunk_size = 100_000\n",
    "    writer = None\n",
    "    for start in range(0, len(df), chunk_size):\n",
    "        df_chunk = df[start : start + chunk_size]\n",
    "        table = pa.Table.from_pandas(df_chunk, preserve_index=False)\n",
    "        if writer is None:\n",
    "            writer = pq.ParquetWriter(heap_file, table.schema)\n",
    "        writer.write_table(table)\n",
    "    writer.close()\n",
    "\n",
    "\n",
    "def get_descendents(net, x, total_proba, total_logproba, edge_index, edge_attr, cutoff):\n",
    "    index_array = torch.arange(x.size(0))\n",
    "    mask = x == 20\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = net(x, edge_index, edge_attr).cpu()\n",
    "        output = torch.softmax(output, dim=1)\n",
    "\n",
    "    output = output[mask]\n",
    "    index_array = index_array[mask]\n",
    "\n",
    "    max_proba, max_index = output.max(dim=1)[0].max(dim=0)\n",
    "    row_with_max_proba = output[max_index]\n",
    "\n",
    "    assert total_logproba <= 0, total_logproba\n",
    "\n",
    "    children = []\n",
    "    for i, p in enumerate(row_with_max_proba):\n",
    "        x_clone = x.clone()\n",
    "        assert x_clone[index_array[max_index]] == 20\n",
    "        x_clone[index_array[max_index]] = i\n",
    "        total_proba_clone = total_proba - cutoff + p.item()\n",
    "        total_logproba_clone = total_logproba - np.log(cutoff) + np.log(p.item())\n",
    "        children.append((x_clone, total_proba_clone, total_logproba_clone))\n",
    "    return children\n",
    "\n",
    "\n",
    "def design_sequence_astar(\n",
    "    net, data, cutoff, num_categories=20, max_results=5_000, max_heap_size=100_000_000, heap=None\n",
    "):\n",
    "    # TODO: keep only total probabilities and log-probabilities instead of the entire array\n",
    "\n",
    "    assert num_categories < 128  # So that we can store x as int8\n",
    "    total_proba = cutoff * data.x.size(0)\n",
    "    total_logproba = np.log(cutoff) * data.x.size(0)\n",
    "    if heap is None:\n",
    "        heap = [\n",
    "            proteinsolver.utils.PrioritizedItem(\n",
    "                -total_logproba, data.x.cpu().to(torch.int8), total_proba, total_logproba\n",
    "            )\n",
    "        ]\n",
    "    results = []\n",
    "\n",
    "    pbar = tqdm(total=max_results)\n",
    "    while len(results) < max_results:\n",
    "        try:\n",
    "            item = heapq.heappop(heap)\n",
    "        except IndexError:\n",
    "            break\n",
    "        if not (item.x == num_categories).any():\n",
    "            assert item.x.dtype == torch.int8\n",
    "            results.append((item.x.data, item.total_proba, item.total_logproba))\n",
    "            pbar.update(1)\n",
    "        else:\n",
    "            children = get_descendents(\n",
    "                net,\n",
    "                item.x.to(torch.long).to(device),\n",
    "                item.total_proba,\n",
    "                item.total_logproba,\n",
    "                data.edge_index,\n",
    "                data.edge_attr,\n",
    "                cutoff,\n",
    "            )\n",
    "            for x, total_proba, total_logproba in children:\n",
    "                heapq.heappush(\n",
    "                    heap,\n",
    "                    proteinsolver.utils.PrioritizedItem(\n",
    "                        -total_logproba, x.cpu().to(torch.int8), total_proba, total_logproba\n",
    "                    ),\n",
    "                )\n",
    "        if len(heap) > max_heap_size:\n",
    "            heap = heap[: len(heap) // 2]\n",
    "            heapq.heapify(heap)\n",
    "    return results, heap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "structure_all = PDB.load(STRUCTURE_FILE)\n",
    "structure = PDB.Structure(STRUCTURE_FILE.name + \"A\", structure_all[0].extract('A'))\n",
    "assert len(list(structure.chains)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e44b5f4c000b47fe9b3801e19e3618d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view = PDB.view_structure(structure)\n",
    "\n",
    "view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e3kWuRaxr89h"
   },
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run protein_train/{UNIQUE_ID}/model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 1\n",
    "num_features = 20\n",
    "adj_input_size = 2\n",
    "hidden_size = 128\n",
    "frac_present = 0.5\n",
    "frac_present_valid = frac_present\n",
    "info_size= 1024"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'protein_train/191f05de/e53-s1952148-d93703104.state'"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state_file = BEST_STATE_FILES[UNIQUE_ID]\n",
    "state_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net(\n",
    "    x_input_size=num_features + 1, adj_input_size=adj_input_size, hidden_size=hidden_size, output_size=num_features\n",
    ")\n",
    "net.load_state_dict(torch.load(state_file, map_location=device))\n",
    "net.eval()\n",
    "net = net.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Design pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load protein sequence and geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ProteinData(sequence='LGQFQSNIEQFKSHMNANTKICAIMKADAYGNGIRGLMPTIIAQGIPCVGVASNAEARAVRESGFKGELIRVRSASLSEMSSALDLNIEELIGTHQQALDLAELAKQSGKTLKVHIALNDGGMGRNGIDMTTEAGKKEAVSIATQPSLSVVGIMTHFPNYNADEVRAKLAQFKESSTWLMQQANLKREEITLHVANSYTALNVPEAQLDMVRPGGVL', row_index=array([  0,   0,   0, ..., 214, 214, 215]), col_index=array([  1,   2,   3, ..., 215, 216, 216]), distances=array([1.32022909, 3.15466319, 2.63357794, ..., 1.32993585, 3.03500283,\n",
      "       1.33465257]))\n"
     ]
    }
   ],
   "source": [
    "pdata = proteinsolver.utils.extract_seq_and_adj(structure, 'A')\n",
    "print(pdata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "217 LGQFQSNIEQFKSHMNANTKICAIMKADAYGNGIRGLMPTIIAQGIPCVGVASNAEARAVRESGFKGELIRVRSASLSEMSSALDLNIEELIGTHQQALDLAELAKQSGKTLKVHIALNDGGMGRNGIDMTTEAGKKEAVSIATQPSLSVVGIMTHFPNYNADEVRAKLAQFKESSTWLMQQANLKREEITLHVANSYTALNVPEAQLDMVRPGGVL\n"
     ]
    }
   ],
   "source": [
    "sequence_ref = pdata.sequence\n",
    "print(len(sequence_ref), sequence_ref)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert data to suitable format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = proteinsolver.datasets.protein.row_to_data(pdata)\n",
    "data = proteinsolver.datasets.protein.transform_edge_attr(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic model statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_stats = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to batch size of one.\n"
     ]
    }
   ],
   "source": [
    "residues, residue_probas = design_sequence(\n",
    "    net, data.to(device), random_position=False, value_selection_strategy=\"map\", num_categories=20\n",
    ")\n",
    "\n",
    "model_stats.update(\n",
    "    {\n",
    "        \"map_sequence_identity\": sum(\n",
    "            proteinsolver.utils.AMINO_ACIDS[r] == sequence_ref[i] for (i, r) in enumerate(residues)\n",
    "        )\n",
    "        / len(sequence_ref),\n",
    "        \"map_proba\": residue_probas.mean().item(),\n",
    "        \"map_logproba\": residue_probas.log().mean().item(),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to batch size of one.\n"
     ]
    }
   ],
   "source": [
    "residues, residue_probas = design_sequence(\n",
    "    net, data.to(device), random_position=False, value_selection_strategy=\"ref\", num_categories=20\n",
    ")\n",
    "\n",
    "model_stats.update(\n",
    "    {\n",
    "        \"ref_sequence_identity\": sum(\n",
    "            proteinsolver.utils.AMINO_ACIDS[r] == sequence_ref[i] for (i, r) in enumerate(residues)\n",
    "        )\n",
    "        / len(sequence_ref),\n",
    "        \"ref_proba\": residue_probas.mean().item(),\n",
    "        \"ref_logproba\": residue_probas.log().mean().item(),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'map_sequence_identity': 0.41013824884792627,\n",
       " 'map_proba': 0.5304911136627197,\n",
       " 'map_logproba': -0.6995142102241516,\n",
       " 'ref_sequence_identity': 1.0,\n",
       " 'ref_proba': 0.3164460062980652,\n",
       " 'ref_logproba': -1.7113533020019531}"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab1/strokach/workspace/proteinsolver/notebooks/generate_protein_sequences/stats-191f05de-4beuA02.json')"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_stats_file = NOTEBOOK_PATH.joinpath(f\"stats-{UNIQUE_ID}-{STRUCTURE_FILE.stem}.json\")\n",
    "with model_stats_file.open(\"wt\") as fout:\n",
    "    json.dump(model_stats, fout)\n",
    "    \n",
    "model_stats_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run protein design using expectimax search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "amino_acids = proteinsolver.utils.AMINO_ACIDS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_output_file(file_index):\n",
    "    return NOTEBOOK_PATH.joinpath(f\"designs-{UNIQUE_ID}-{SEQUENCE_GENERATION_METHOD}-{STRUCTURE_FILE.stem}-{file_index}.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file_index = START_FILE_INDEX\n",
    "\n",
    "while get_output_file(file_index).is_file():\n",
    "    file_index += 1\n",
    "    \n",
    "file_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if SEQUENCE_GENERATION_METHOD == \"astar\":\n",
    "\n",
    "    data.y = data.x\n",
    "    data.x = torch.ones_like(data.x) * 20\n",
    "\n",
    "    heap_file = NOTEBOOK_PATH.joinpath(f\"heap-{UNIQUE_ID}-{SEQUENCE_GENERATION_METHOD}-{STRUCTURE_FILE.stem}.parquet\")\n",
    "    #     heap = load_heap_dump(heap_file)\n",
    "    heap = None\n",
    "\n",
    "    while True:\n",
    "        results, heap = design_sequence_astar(\n",
    "            net, data, cutoff=MIN_EXPECTED_PROBA, num_categories=20, max_results=20_000, heap=heap\n",
    "        )\n",
    "        print(len(heap))\n",
    "        update_heap_dump(heap_file, heap)\n",
    "        results = [(\"\".join(amino_acids[i] for i in r[0]),) + r[1:] for r in results]\n",
    "        df = pd.DataFrame(results, columns=[\"sequence\", \"probas_sum\", \"probas_log_sum\"])\n",
    "        table = pa.Table.from_pandas(df, preserve_index=False)\n",
    "        pq.write_table(table, get_output_file(file_index))\n",
    "        file_index += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if SEQUENCE_GENERATION_METHOD.endswith(\"expectimax\"):\n",
    "\n",
    "    random_position = SEQUENCE_GENERATION_METHOD.startswith(\"rand\")\n",
    "    print(f\"random_position: {random_position}\")\n",
    "\n",
    "    batch_size = int(\n",
    "        512\n",
    "        * (3586 / (data.x.size(0) + data.edge_attr.size(0)))\n",
    "        * (torch.cuda.get_device_properties(device).total_memory / 12_650_217_472)\n",
    "    )\n",
    "    print(f\"batch_size: {batch_size}\")\n",
    "\n",
    "    data_batch = Batch.from_data_list([data.clone() for _ in range(batch_size)]).to(device)\n",
    "    data_batch.x = torch.ones_like(data_batch.x) * 20\n",
    "\n",
    "    while True:\n",
    "        results = []\n",
    "        for _ in tqdm(range(int(np.ceil(5_000 / batch_size))), desc=str(file_index)):\n",
    "            batch_values, batch_probas = design_sequence(\n",
    "                net, data_batch, random_position=random_position, value_selection_strategy=\"multinomial\"\n",
    "            )\n",
    "            for i in range(batch_size):\n",
    "                values = batch_values[data_batch.batch == i]\n",
    "                probas = batch_probas[data_batch.batch == i]\n",
    "                sequence = \"\".join(amino_acids[i] for i in values)\n",
    "                probas_sum = probas.sum().item()\n",
    "                probas_log_sum = probas.log().sum().item()\n",
    "                results.append((sequence, probas_sum, probas_log_sum))\n",
    "        df = pd.DataFrame(results, columns=[\"sequence\", \"probas_sum\", \"probas_log_sum\"])\n",
    "        table = pa.Table.from_pandas(df, preserve_index=False)\n",
    "        pq.write_table(table, get_output_file(file_index))\n",
    "        file_index += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "1fXJz1l9HvhW",
    "u9q4iTifHvhl"
   ],
   "name": "Copy of 2019-03-30-sudoku-4xEdgeConv-09862+.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "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.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": false,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}