slc22a5_train_potts.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/songlab-cal/slc22a5/blob/main/slc22a5_train_potts.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ih_SIgoA7kpc"
      },
      "source": [
        "# Install mogwai"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EktPZrIN6Euz",
        "outputId": "fe5cf218-6de3-40ff-d9df-69d9a3fe1d07"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[K     |████████████████████████████████| 2.3 MB 6.2 MB/s \n",
            "\u001b[K     |████████████████████████████████| 584 kB 35.0 MB/s \n",
            "\u001b[K     |████████████████████████████████| 31.8 MB 28 kB/s \n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[K     |████████████████████████████████| 408 kB 42.2 MB/s \n",
            "\u001b[K     |████████████████████████████████| 136 kB 42.5 MB/s \n",
            "\u001b[K     |████████████████████████████████| 596 kB 45.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.1 MB 41.5 MB/s \n",
            "\u001b[K     |████████████████████████████████| 94 kB 3.3 MB/s \n",
            "\u001b[K     |████████████████████████████████| 271 kB 43.4 MB/s \n",
            "\u001b[K     |████████████████████████████████| 144 kB 47.1 MB/s \n",
            "\u001b[?25h  Building wheel for mogwai-protein (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for biotite (PEP 517) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "# takes ~3 mins\n",
        "\n",
        "!pip install -q git+https://github.com/nickbhat/mogwai.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bUK5ZHiE7ZNY"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import pytorch_lightning as pl\n",
        "import torch\n",
        "from tqdm import tqdm\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "asyA_LMUBy8S"
      },
      "source": [
        "# Download alignments and wildtype"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R5UIRrrn8Pv1"
      },
      "outputs": [],
      "source": [
        "# wildtype\n",
        "!wget -q http://s3.amazonaws.com/songlabdata/slc22a5/O76082.fasta\n",
        "\n",
        "# alignments\n",
        "!wget -q http://s3.amazonaws.com/songlabdata/slc22a5/alignments/mammals_30.a3m\n",
        "!wget -q http://s3.amazonaws.com/songlabdata/slc22a5/alignments/hhblits.a3m\n",
        "!wget -q http://s3.amazonaws.com/songlabdata/slc22a5/alignments/vertebrates_100.a3m\n",
        "!wget -q http://s3.amazonaws.com/songlabdata/slc22a5/alignments/eve.a3m\n",
        "!wget -q http://s3.amazonaws.com/songlabdata/slc22a5/alignments/deepsequence.a3m"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RDRxhz3WFkIg"
      },
      "source": [
        "# Train model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qEC0gePJBxaj"
      },
      "outputs": [],
      "source": [
        "import mogwai.models as models\n",
        "import mogwai.data_loading as data_loading\n",
        "import os\n",
        "\n",
        "def torch_to_numpy(state_dict, keys):\n",
        "    \"\"\"Convert `keys` in the `state_dict` to numpy arrays.\"\"\"\n",
        "    numpy_dict = dict()\n",
        "    for key in keys:\n",
        "        numpy_dict[key] = state_dict[key].numpy()\n",
        "    return numpy_dict\n",
        "\n",
        "\n",
        "def get_outfile_path(alignment_path):\n",
        "    prefix, suffix = alignment_path.rsplit('.', 1)\n",
        "    return os.path.join(prefix + '_potts_state_dict.npz')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BVpOeyOYBcva"
      },
      "outputs": [],
      "source": [
        "def train_potts(alignment_path, max_steps, gpus):\n",
        "    print('Reading {}'.format(alignment_path))\n",
        "    # Load msa\n",
        "    msa_dm = data_loading.MSADataModule(alignment_path, batch_size=4096)\n",
        "    msa_dm.setup()\n",
        "\n",
        "    # Initialize model\n",
        "    num_seqs, msa_length, msa_counts = msa_dm.get_stats()\n",
        "    model = models.Gremlin(num_seqs, msa_length, msa_counts)\n",
        "\n",
        "    # Initialize Trainer\n",
        "    trainer = pl.Trainer(min_steps=50, max_steps=max_steps, gpus=gpus)\n",
        "\n",
        "    # Train model\n",
        "    trainer.fit(model, msa_dm)\n",
        "\n",
        "    # save model\n",
        "\n",
        "    model_dict = torch_to_numpy(model.state_dict(), ['weight', 'bias'])\n",
        "\n",
        "    outfile = get_outfile_path(alignment_path)\n",
        "\n",
        "    np.savez(outfile, **model_dict)\n",
        "    print('Saved model state dict to {}'.format(outfile))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 399,
          "referenced_widgets": [
            "e63b9f38ee944f299dbb2366fc5683ac",
            "2b591156b964454586f345c076ca9eb5",
            "1c6a0e8645d344b398bd89ed75221956",
            "615693ff34ce4670ba31c08bae6dae9a",
            "e1b93a94271b4508afdb00ec4005f523",
            "0927196e970a4a81a3d072e23a60082a",
            "04f01be937554c2e8fd8100b93edba3a",
            "ea69a33130734680bca701d457f5b18c",
            "5df2826436e24b1dbca98c16ec040be8",
            "ca5bf6f457e247718fc5f81967dc4d3a",
            "82d29e08585e4736913e8d1b2e3cac0e"
          ]
        },
        "id": "-Vq9hw51Ghar",
        "outputId": "910b663f-3af3-4164-c78d-34f8979c2661"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reading deepsequence.a3m\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/datamodule.py:149: LightningDeprecationWarning: DataModule property `dims` was deprecated in v1.5 and will be removed in v1.7.\n",
            "  rank_zero_deprecation(\"DataModule property `dims` was deprecated in v1.5 and will be removed in v1.7.\")\n",
            "GPU available: True, used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n",
            "HPU available: False, using: 0 HPUs\n",
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
            "\n",
            "  | Name | Type | Params\n",
            "------------------------------\n",
            "------------------------------\n",
            "124 M     Trainable params\n",
            "0         Non-trainable params\n",
            "124 M     Total params\n",
            "496.443   Total estimated model params size (MB)\n",
            "/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py:490: UserWarning: This DataLoader will create 4 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
            "  cpuset_checked))\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Training: 0it [00:00, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "e63b9f38ee944f299dbb2366fc5683ac"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saved model state dict to deepsequence_potts_state_dict.npz\n"
          ]
        }
      ],
      "source": [
        "# this takes about 10 mins.\n",
        "alignment_path = 'deepsequence.a3m'\n",
        "max_steps = 500  # 500\n",
        "gpus = [0]  # CPU training is extremely slow, not recommended.\n",
        "train_potts(alignment_path=alignment_path, max_steps=max_steps, gpus=gpus)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "slc22a5_train_potts.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyPENR+t1WwUpXHCm6HcKJyY",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "e63b9f38ee944f299dbb2366fc5683ac": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_2b591156b964454586f345c076ca9eb5",
              "IPY_MODEL_1c6a0e8645d344b398bd89ed75221956",
              "IPY_MODEL_615693ff34ce4670ba31c08bae6dae9a"
            ],
            "layout": "IPY_MODEL_e1b93a94271b4508afdb00ec4005f523"
          }
        },
        "2b591156b964454586f345c076ca9eb5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0927196e970a4a81a3d072e23a60082a",
            "placeholder": "​",
            "style": "IPY_MODEL_04f01be937554c2e8fd8100b93edba3a",
            "value": "Epoch 0:  55%"
          }
        },
        "1c6a0e8645d344b398bd89ed75221956": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "danger",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ea69a33130734680bca701d457f5b18c",
            "max": 910,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_5df2826436e24b1dbca98c16ec040be8",
            "value": 500
          }
        },
        "615693ff34ce4670ba31c08bae6dae9a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ca5bf6f457e247718fc5f81967dc4d3a",
            "placeholder": "​",
            "style": "IPY_MODEL_82d29e08585e4736913e8d1b2e3cac0e",
            "value": " 500/910 [09:49&lt;08:03,  1.18s/it, loss=8.42e+06, v_num=1]"
          }
        },
        "e1b93a94271b4508afdb00ec4005f523": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": "inline-flex",
            "flex": null,
            "flex_flow": "row wrap",
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": "100%"
          }
        },
        "0927196e970a4a81a3d072e23a60082a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "04f01be937554c2e8fd8100b93edba3a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ea69a33130734680bca701d457f5b18c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": "2",
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5df2826436e24b1dbca98c16ec040be8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "ca5bf6f457e247718fc5f81967dc4d3a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "82d29e08585e4736913e8d1b2e3cac0e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}