cryo_em_minimizer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/huhlim/cg2all/blob/main/cryo_em_minimizer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FBOi_bUBW5dR"
      },
      "source": [
        "\n",
        "This notebook performs local optimization of a protein model structure against a cryo-EM density map using __cryo_em_minimizer.py__ This task is running very slow on the Google Colab __CPU__ runtime. Thus, it is highly recommended to use a __GPU__ runtime with enough VRAM. Because it only performs local optimization, the input structure should be roughly fitted to the input electron density map.\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pmJA_J_IlU7s",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Install cg2all package (takes 4-5 minutes)\n",
        "%%bash\n",
        "\n",
        "gpu_available=$(nvidia-smi | grep \"CUDA Version\" | wc -l)\n",
        "if [[ $gpu_available == 1 ]]; then\n",
        "    echo \"This notebook is running on a GPU runtime.\"\n",
        "    pip install  dgl -f https://data.dgl.ai/wheels/torch-2.3/cu121/repo.html &> /dev/null\n",
        "    #pip install dgl -f https://data.dgl.ai/wheels/cu116/repo.html &> /dev/null\n",
        "else\n",
        "    echo \"This notebook is running on a CPU runtime.\"\n",
        "fi\n",
        "\n",
        "pip install -q git+http://github.com/huhlim/cg2all@cuda-12 &> /dev/null\n",
        "pip install -q py3Dmol gdown mrcfile &> /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BaywXyY50z7O"
      },
      "outputs": [],
      "source": [
        "#@title Download a model checkpoint file (optional, takes a few seconds)\n",
        "#@markdown This step downloads a necessary PyTorch model checkpoint file. If you did not run this step and the file is missing, then the script will download it automatically later.\n",
        "\n",
        "import cg2all.lib.libmodel\n",
        "from cg2all.lib.libconfig import MODEL_HOME\n",
        "\n",
        "for model_type in [\"CalphaBasedModel\", \"ResidueBasedModel\"]:\n",
        "    ckpt_fn = MODEL_HOME / f\"{model_type}.ckpt\"\n",
        "    if not ckpt_fn.exists():\n",
        "        cg2all.lib.libmodel.download_ckpt_file(model_type, ckpt_fn)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2izhOiEl_LZz"
      },
      "outputs": [],
      "source": [
        "#@title Set up py3Dmol for structure display\n",
        "\n",
        "import mdtraj\n",
        "import py3Dmol\n",
        "\n",
        "def display(pdb_fn, representation=\"cartoon\", is_traj=False):\n",
        "    view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "    if is_traj:\n",
        "        view.addModelsAsFrames(open(pdb_fn,'r').read(),'pdb')\n",
        "    else:\n",
        "        view.addModel(open(pdb_fn,'r').read(),'pdb')\n",
        "\n",
        "    if representation == \"cartoon\":\n",
        "        view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "        #\n",
        "        view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom': [\"N\", \"C\", \"O\"],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':\"WhiteCarbon\",'radius':0.3}})\n",
        "        view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                        {'sphere':{'colorscheme':\"WhiteCarbon\",'radius':0.3}})\n",
        "        view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "    elif representation == \"ball+stick\":\n",
        "        view.setStyle({\"sphere\": {\"color\": \"spectrum\", \"radius\": 1.0}})\n",
        "\n",
        "    else:\n",
        "        raise NotImplementedError(representation)\n",
        "\n",
        "    view.zoomTo()\n",
        "    if is_traj:\n",
        "        view.animate({'loop': 'forward', \"interval\": 500})\n",
        "\n",
        "    return view"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8B3PrhVv2XjL"
      },
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CoX74ebi5u3b"
      },
      "outputs": [],
      "source": [
        "#@title Run local optimization\n",
        "\n",
        "# upload a PDB file\n",
        "import requests\n",
        "from google.colab import files\n",
        "\n",
        "coarse_grained_model_type = \"ResidueBasedModel\" #@param [\"CalphaBasedModel\", \"ResidueBasedModel\"]\n",
        "segmentation_method = \"chain\" #@param [\"None\", \"chain\", \"segment\"]\n",
        "if segmentation_method == \"None\":\n",
        "    segmentation_method = None\n",
        "use_example = True #@param {type:\"boolean\"}\n",
        "#@markdown - An example input PDB and electron density map files will be downloaded from [our repository](https://github.com/huhlim/cg2all/tree/main/tests).\n",
        "\n",
        "if use_example:\n",
        "    url = \"https://raw.githubusercontent.com/huhlim/cg2all/main/tests/3isr.af2.pdb\"\n",
        "    input_pdb = url.split(\"/\")[-1]\n",
        "    with open(input_pdb, \"wt\") as fout:\n",
        "        fout.write(requests.get(url).text)\n",
        "    #\n",
        "    url = \"https://raw.githubusercontent.com/huhlim/cg2all/main/tests/3isr_5.mrc\"\n",
        "    input_map = url.split(\"/\")[-1]\n",
        "    with open(input_map, \"wb\") as fout:\n",
        "        fout.write(requests.get(url).content)\n",
        "\n",
        "else:\n",
        "    input_pdb = files.upload()\n",
        "    input_pdb = list(input_pdb)[0]\n",
        "    input_map = files.upload()\n",
        "    input_map = list(input_map)[0]\n",
        "\n",
        "number_of_snapshots = 10 #@param {type: \"slider\", min:1, max:100}\n",
        "#@markdown - The number of minimization steps is 10 x number_of_snapshots.\n",
        "n_step = number_of_snapshots * 10\n",
        "\n",
        "import pathlib\n",
        "output_prefix = pathlib.Path(input_map).stem + \"+\" + pathlib.Path(input_pdb).stem\n",
        "\n",
        "!cryo_em_minimizer --pdb $input_pdb --map $input_map --output $output_prefix --cg $coarse_grained_model_type --all -n $n_step --output_freq 10 --segment $segmentation_method --uniform_restraint\n",
        "\n",
        "pdb_fn_s = list(pathlib.Path(output_prefix).glob(\"min.*.pdb\"))\n",
        "pdb_fn_s.sort(key=lambda fn: int(fn.name.split(\".\")[-2]))\n",
        "if len(pdb_fn_s) > 10:\n",
        "    pdb_fn_s = pdb_fn_s[::len(pdb_fn_s)//10]\n",
        "    print(f\"Displaying only {len(pdb_fn_s)}\")\n",
        "\n",
        "traj = mdtraj.load(pdb_fn_s)\n",
        "traj.save(\"display.pdb\")\n",
        "display(\"display.pdb\", representation=\"cartoon\", is_traj=True).show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "O_Bfse-lTcEA"
      },
      "outputs": [],
      "source": [
        "#@title Download the optimized files\n",
        "\n",
        "import os\n",
        "zip = f\"{output_prefix}.zip\"\n",
        "os.system(f\"zip -r {zip} {output_prefix}\")\n",
        "files.download(zip)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}