File size: 6,575 Bytes
d766458
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/af_single_sequential.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-AXy0s_4cKaK"
      },
      "outputs": [],
      "source": [
        "#@title af_single_sequential\n",
        "#@markdown implementing [@MartinPacesa](https://twitter.com/MartinPacesa/status/1640005299124707328)'s idea\n",
        "import os\n",
        "if not os.path.isdir(\"params\"):\n",
        "  print(\"installing AlphaFold...\")\n",
        "  # get code\n",
        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
        "  # for debugging\n",
        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
        "  print(\"downloading AlphaFold params...\")\n",
        "  # download params\n",
        "  os.system(\"mkdir params\")\n",
        "  os.system(\"apt-get install aria2 -qq\")\n",
        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-03-02.tar\")\n",
        "  os.system(\"tar -xf alphafold_params_2022-03-02.tar -C params\")\n",
        "\n",
        "import os\n",
        "from colabdesign import mk_afdesign_model, clear_mem\n",
        "from IPython.display import HTML\n",
        "from google.colab import files\n",
        "import numpy as np\n",
        "import tqdm.notebook\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "def get_pdb(pdb_code=\"\"):\n",
        "  if pdb_code is None or pdb_code == \"\":\n",
        "    upload_dict = files.upload()\n",
        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
        "    return \"tmp.pdb\"\n",
        "  elif os.path.isfile(pdb_code):\n",
        "    return pdb_code\n",
        "  elif len(pdb_code) == 4:\n",
        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
        "    return f\"{pdb_code}.pdb\"\n",
        "  else:\n",
        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
        "\n",
        "#@markdown Specify either pdb/pdb_chain or sequence:\n",
        "pdb = \"6D0T\" #@param {type:\"string\"}\n",
        "pdb_chain = \"A\" #@param {type:\"string\"}\n",
        "sequence = \"\" #@param {type:\"string\"}\n",
        "num_recycles = 0 #@param [\"0\", \"1\", \"2\", \"4\", \"8\", \"16\"] {type:\"raw\"}\n",
        "clear_mem()\n",
        "if sequence == \"\":\n",
        "  pdb_filename = get_pdb(pdb)\n",
        "  af_model = mk_afdesign_model(protocol=\"fixbb\")\n",
        "  af_model.prep_inputs(pdb_filename=pdb_filename, chain=pdb_chain)\n",
        "  af_model.set_seq(mode=\"wildtype\")\n",
        "  mode = \"pdb\"\n",
        "else:\n",
        "  af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
        "  af_model.prep_inputs(length=len(sequence))\n",
        "  af_model.set_seq(sequence)\n",
        "  mode = \"free\"\n",
        "\n",
        "import random, string\n",
        "path = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))\n",
        "while os.path.isdir(path):\n",
        "  path = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))\n",
        "\n",
        "seq = af_model._params[\"seq\"].copy()\n",
        "os.makedirs(path, exist_ok=True)\n",
        "with tqdm.notebook.tqdm(total=af_model._len, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "  for n in range(1,af_model._len+1):\n",
        "  \n",
        "    # replace sequence (this is just for visualization)\n",
        "    af_model._params[\"seq\"] = seq.copy()\n",
        "    af_model._params[\"seq\"][:,n:] = 0\n",
        "    \n",
        "    # mask part of the input features\n",
        "    af_model._inputs[\"seq_mask\"] = (np.arange(af_model._len) < n).astype(float)\n",
        "    af_model.predict(verbose=False, hard=False, num_recycles=num_recycles)\n",
        "    af_model._save_results(verbose=False)\n",
        "    af_model.save_current_pdb(f\"{path}/{n:04d}.pdb\")\n",
        "    pbar.update(1)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "HTML(af_model.animate(dpi=100))"
      ],
      "metadata": {
        "id": "lphWx8wkEpwP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.plot_pdb()"
      ],
      "metadata": {
        "id": "9h1uXK44TygM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Package and download results\n",
        "#@markdown If you are having issues downloading the result archive, \n",
        "#@markdown try disabling your adblocker and run this cell again. \n",
        "#@markdown  If that fails click on the little folder icon to the \n",
        "#@markdown  left, navigate to file: `name.result.zip`, \n",
        "#@markdown  right-click and select \\\"Download\\\" \n",
        "#@markdown (see [screenshot](https://pbs.twimg.com/media/E6wRW2lWUAEOuoe?format=jpg&name=small)).\n",
        "from google.colab import files\n",
        "!zip -r {path}.result.zip {path}/*\n",
        "files.download(f\"{path}.result.zip\")"
      ],
      "metadata": {
        "id": "zGE1CJxuTzyQ",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "q4qiU9I0QHSz"
      ],
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}