adding logic to split chains

dauparas · Jun 17, 2022 · 4ec7228 · 4ec7228
1 parent 06478cf
commit 4ec7228
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/colab_notebooks/quickdemo_wAF2.ipynb b/colab_notebooks/quickdemo_wAF2.ipynb
@@ -21,7 +21,7 @@
         "\n",
         "Examples: \n",
         "1.   pdb: `6MRR`, homomer: `False`, designed_chain: `A`\n",
-        "2.   pdb: `3DJI`, homomer: `True`, designed_chain: `A,B,C` \n",
+        "2.   pdb: `1X2I`, homomer: `True`, designed_chain: `A,B` \n",
         "     (for correct symmetric tying lenghts of homomer chains should be the same)"
       ]
     },
@@ -462,7 +462,7 @@
         "\n",
         "  return jax.jit(runner), model_params, {\"inputs\":inputs, \"length\":max_length}\n",
         "\n",
-        "def save_pdb(outs, filename):\n",
+        "def save_pdb(outs, filename, Ls=None):\n",
         "  '''save pdb coordinates'''\n",
         "  p = {\"residue_index\":outs[\"residue_idx\"] + 1,\n",
         "       \"aatype\":outs[\"seq\"],\n",
@@ -473,7 +473,12 @@
         "  b_factors = 100 * p.pop(\"plddt\")[:,None] * p[\"atom_mask\"]\n",
         "  p = alphafold_protein.Protein(**p,b_factors=b_factors)\n",
         "  pdb_lines = alphafold_protein.to_pdb(p)\n",
-        "  with open(filename, 'w') as f: f.write(pdb_lines)"
+        "  with open(filename, 'w') as f:\n",
+        "    f.write(pdb_lines)\n",
+        "  if Ls is not None:\n",
+        "    pdb_lines = cf.read_pdb_renum(filename, Ls)\n",
+        "    with open(filename, 'w') as f:\n",
+        "      f.write(pdb_lines)"
       ],
       "metadata": {
         "cellView": "form",
@@ -536,7 +541,7 @@
         "      plddts[-1].append(O[\"plddt\"][:length])\n",
         "      paes[-1].append(O[\"pae\"][:length,:length])\n",
         "      outs[-1].append(O)\n",
-        "      save_pdb(outs[-1][-1], f\"out_seq_{s}_model_{n}.pdb\")\n",
+        "      save_pdb(outs[-1][-1], f\"out_seq_{s}_model_{n}.pdb\", Ls=LS[-1])\n",
         "      tmscores = mmalign(pdb_path, f\"out_seq_{s}_model_{n}.pdb\")\n",
         "      print(f\"{s} {n}\\t{plddts[-1][-1].mean():.3}\\t{paes[-1][-1].mean():.3}\\t{tmscores[-1]:.3}\")"
       ],