diff --git a/notebooks/tmol_how_to_guide.ipynb b/notebooks/tmol_how_to_guide.ipynb new file mode 100644 index 000000000..a27fb62b6 --- /dev/null +++ b/notebooks/tmol_how_to_guide.ipynb @@ -0,0 +1,2914 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "0NSgz0L8CeFA" + }, + "source": [ + "# Contents:\n", + "\n", + "Setup:\n", + "- Download tmol wheel (colab)\n", + "- Download a few input files\n", + "- Hello World of working with tmol: create a PoseStack and score it\n", + "\n", + "Fundamentals:\n", + "- Initialize the default ParameterDatabase\n", + "- Load alternate block types\n", + "- Initialize a PackedBlockTypes object from the default database\n", + "- Initialize a PackedBlockTypes object from a custom database\n", + "- Create a CanonicalOrdering from a ParameterDatabase\n", + "- Create a PackedBlockTypes object from a subset of the block types in a ParameterDatabase\n", + "\n", + "Input\n", + "- Initialize a single pose PoseStack from a PDB\n", + "- Initialize a pose stack from a subset of residues in a PDB\n", + "- Initialize a single pose PoseStack from a PDB file using biotite\n", + "- Initialize a single pose PoseStack from an .mmcif file (.cif)\n", + "- Initialize a single pose PoseStack from an OpenMM set of tensors\n", + "- Initialize a single pose PoseStack with missing residues\n", + "- Load a bunch of PoseStacks from different sources and then concatenate them to a single PoseStack\n", + "- Make many copies of the same single-pose PoseStack\n", + "- Add custom residue type\n", + "- Set the scoring parameters for a new residue type (elec, etc)\n", + "- Build an extended pose from sequence\n", + "- (Ligand features??)\n", + "- Create ligand block type from .params\n", + "- Create ligand block type from .mol2\n", + "- Create ligand block type from .cif that contains ligands\n", + "- Create BiotitePoseBuildContext with RefinedResidueTypes for repeat loading of ligand PDBs\n", + "- Load in multiple ligands and then process PDBs containing those ligands\n", + "\n", + "Output:\n", + "- Write a single pose PoseStack to a .pdb file\n", + "- Write a multi-pose stack out as a multi-model PDB file\n", + "- Write a multi-pose stack out to separate PDB files\n", + "- Write a single pose PoseStack to an .mmcif file\n", + "- Write a rotamer set out as a multi-model PDB\n", + "- Write ligand-containing PoseStack out to .cif file using BiotitePoseBuildContext\n", + "- Write out ligand .params file\n", + "\n", + "Kinematics:\n", + "- Create an N->C fold tree for a PoseStack\n", + "- Create a simple fold tree for a multi-chain PoseStack\n", + "- Create a simple fold tree for a PoseStack with missing residues\n", + "- Create a dandelion fold tree for a PoseStack\n", + "- Create a MoveMap that enables minimization for named torsions\n", + "- Create a MoveMap that enables backbone minimization for some residues but not all\n", + "- Apply a perturbation to the rigid-body DOFs between two chains\n", + "- Assign dihedral values to all the residues in a PoseStack and calculate the coordinates\n", + "\n", + "Scoring:\n", + "- Create the default score function\n", + "- Create the default score function from a custom Database\n", + "- Create the soft-rep version of the score function\n", + "- Create an empty score function\n", + "- Turn on a few terms in a score function\n", + "- Turn off a term in a score function\n", + "- Score a PoseStack\n", + "- Score a PoseStack and back-propagate through the coordinates\n", + "- Score a PoseStack and return per-residue weighted energies\n", + "- Score a PoseStack and return per-residue un-weighted energies\n", + "- Score a PoseStack and return per-residue weighted energies, weight them according to some principle, and then back-propagate the total energy.\n", + "- Add constraints to a PoseStack\n", + "- Add the same constraints to all the Poses in a PoseStack\n", + "- Add coordinate constraints to the current coordinates\n", + "- Alter the parameters for the cart-bonded energy function & rescore\n", + "\n", + "\n", + "Optimization\n", + "- Create a DunbrackSampler\n", + "- Add hydrogens\n", + "- Fill in side chains for a model that lacks them\n", + "- Perform fixed-sequence side-chain optimization: repack\n", + "- Repack with extra rotamers\n", + "- Add hydrogens and back-propagate\n", + "- Create a new PackerPalette subclass to handle logic of new block types\n", + "- Run minimization in double precision\n", + "- Perform cartesian minimization\n", + "- Perform kinematic minimization\n", + "- Relax a PoseStack w/ kinematic minimization\n", + "- Relax a PoseStack w/ cartesian minimization\n", + "- Relax structures generated one-at-a-time in batch format\n", + "- Idealize a structure from the PDB\n", + "- Idealize a structure from a dandelion\n", + "- Idealize just the backbone of a dandelion\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "nU1-4FoBbApM", + "outputId": "68e50ec9-25eb-450e-f1b2-882c310165c2", + "scrolled": true + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "tmol is already importable; skipping the wheel install.\n" + ] + } + ], + "source": [ + "# For Colab: install tmol from the pinned github wheel.\n", + "# On a local checkout (or a second run of this notebook) tmol is already\n", + "# importable, so the install is skipped and the notebook stays runnable\n", + "# in both environments.\n", + "try:\n", + " import tmol # noqa: F401\n", + "\n", + " print(\"tmol is already importable; skipping the wheel install.\")\n", + "except ImportError:\n", + " !pip install https://github.com/uw-ipd/tmol/releases/download/v0.1.36/tmol-0.1.36+cu128torch2.10-cp312-cp312-linux_x86_64.whl" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "y3f4dcA0mV8S", + "outputId": "29bc90e1-4fad-4e2a-a724-58d57032e1a2" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "1ubq.pdb: already present\n", + "1s78.pdb: already present\n", + "1qys.pdb: already present\n", + "3plc.pdb: already present\n", + "10VB.pdb: already present\n", + "1BL8.cif: already present\n", + "openfold_ubq_and_sumo.pt: already present\n", + "ligand.tmol: already present\n" + ] + } + ], + "source": [ + "# Download some structure files so we have something to work with.\n", + "# Files that are already present are left alone, so re-running this cell\n", + "# (or running from a tmol checkout) is cheap.\n", + "import os\n", + "from urllib.request import urlretrieve\n", + "\n", + "BASE = \"https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data\"\n", + "\n", + "INPUT_FILES = {\n", + " \"1ubq.pdb\": f\"{BASE}/pdb/1ubq.pdb\",\n", + " \"1s78.pdb\": f\"{BASE}/pdb/1s78.pdb\",\n", + " \"1qys.pdb\": f\"{BASE}/pdb/1qys.pdb\",\n", + " \"3plc.pdb\": f\"{BASE}/pdb/3plc.pdb\",\n", + " \"10VB.pdb\": f\"{BASE}/pdb/10VB.pdb\",\n", + " \"1BL8.cif\": f\"{BASE}/cif/1BL8.cif\",\n", + " \"openfold_ubq_and_sumo.pt\": f\"{BASE}/openfold/openfold_ubq_and_sumo.pt\",\n", + " # a ligand parameter file, used by the \"custom residue type\" cells below\n", + " \"ligand.tmol\": f\"{BASE}/protein_ligand_test/gr.xtal-lig.mmff94.tmol\",\n", + "}\n", + "\n", + "for name, url in INPUT_FILES.items():\n", + " if os.path.isfile(name):\n", + " print(f\"{name}: already present\")\n", + " continue\n", + " print(f\"{name}: downloading ...\")\n", + " urlretrieve(url, name)\n", + " print(f\"{name}: {os.path.getsize(name)} bytes\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "A_hmNzYjZLH0", + "outputId": "3477568c-0a01-4d71-dad9-6385092be812" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "ATOM 1 N MET A 1 27.340 24.430 2.614 1.00 9.67 N \r\n", + "ATOM 2 CA MET A 1 26.266 25.413 2.842 1.00 10.38 C \r\n", + "ATOM 3 C MET A 1 26.913 26.639 3.531 1.00 9.62 C \r\n", + "ATOM 4 O MET A 1 27.886 26.463 4.263 1.00 9.62 O \r\n", + "ATOM 5 CB MET A 1 25.112 24.880 3.649 1.00 13.77 C \r\n", + "ATOM 6 CG MET A 1 25.353 24.860 5.134 1.00 16.29 C \r\n", + "ATOM 7 SD MET A 1 23.930 23.959 5.904 1.00 17.17 S \r\n", + "ATOM 8 CE MET A 1 24.447 23.984 7.620 1.00 16.11 C \r\n", + "ATOM 9 1H MET A 1 26.961 23.619 2.168 1.00 0.00 H \r\n", + "ATOM 10 2H MET A 1 28.043 24.834 2.029 1.00 0.00 H \r\n" + ] + } + ], + "source": [ + "# Make sure the files downloaded correctly; this should print the first 10 atoms\n", + "# from methionine in ubiquitin\n", + "!head 1ubq.pdb" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "env: TMOL_USE_JIT=1\n" + ] + } + ], + "source": [ + "%env TMOL_USE_JIT=1" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "import tmol" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "['/home/jflat06/rosetta/tmol/tmol']\n" + ] + } + ], + "source": [ + "print(tmol.__path__)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "Mi9gshJmbVVC", + "outputId": "6a3b6001-415b-4570-a8b3-2b56561a709d" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "tensor([235.4992], device='cuda:0')\n" + ] + } + ], + "source": [ + "# The hello world of working with tmol:\n", + "# Load a PDB in from disk and score it\n", + "\n", + "import tmol\n", + "import torch\n", + "import os\n", + "\n", + "device = torch.device(\"cuda\", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device(\"cpu\")\n", + "\n", + "# Create a pose stack from a PDB.\n", + "pose_stack = tmol.pose_stack_from_pdb('1ubq.pdb', device=device)\n", + "# A PoseStack represents a batch of molecular systems (in this case, just a single structure - ubiquitin)\n", + "# PoseStacks are optimized for compactness for efficient processing on the GPU.\n", + "# Behind the scenes, pose_stack_from_pdb uses the default ParameterDatabase;\n", + "# which currently contains the parameters necessary to treat standard\n", + "# proteins, but little else. We will see more about the ParameterDatabase later.\n", + "\n", + "# Create our score function.\n", + "sfxn = tmol.beta2016_score_function(device=device)\n", + "# This tmol score function is based on the Rosetta energy function.\n", + "# The score function is composed of many terms and weights for those terms.\n", + "# In this particular case, the score function terms and weights are set to\n", + "# match the beta2016_cart score function from Rosetta3. Again, we\n", + "# are relying on the default ParameterDatabase in the background.\n", + "\n", + "# Create our scoring module.\n", + "scorer = sfxn.render_whole_pose_scoring_module(pose_stack)\n", + "# The scoring module is what does the actual score evaluation.\n", + "# This is separate from the ScoreFunction because it also needs details\n", + "# about the PoseStack being scored - mainly the Residue Types being used.\n", + "# The scoring module needs those Types because each score term must assemble\n", + "# compact tensors with the data necessary to score the Residue Types that are\n", + "# in use.\n", + "\n", + "# Score the PoseStack and print the output.\n", + "print(scorer(pose_stack.coords))\n", + "# Return a tensor with the score of each pose in the stack (in this case, just 1 value)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Jp_iT-4DjP_x" + }, + "source": [ + "Fundamentals" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "id": "lBCg3XA8jMq8" + }, + "outputs": [], + "source": [ + "# - Initialize the default ParameterDatabase\n", + "default_param_db = tmol.ParameterDatabase.get_default()" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "id": "9nyxfNgojaSB" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n block types in the default database: 142\n", + "HIS block types: ['HIS', 'HIS:cterm', 'HIS:cterm:nterm', 'HIS:nterm', 'HIS_D', 'HIS_D:cterm', 'HIS_D:cterm:nterm', 'HIS_D:nterm', 'HIS_POS', 'HIS_POS:cterm', 'HIS_POS:cterm:nterm', 'HIS_POS:nterm']\n", + "nucleic acid block types: ['DA', 'DC', 'DG', 'DT']\n" + ] + } + ], + "source": [ + "# - Load alternate block types\n", + "#\n", + "# The default ParameterDatabase ships with many more block types (\"residue\n", + "# types\") than the 20 canonical amino acids: protonation and tautomer\n", + "# variants, termini patches, D-amino acids and nucleic acids. ResidueTypeSet\n", + "# turns the raw chemical database into the RefinedResidueType objects the rest\n", + "# of tmol consumes.\n", + "\n", + "# TEMP: ResidueTypeSet is not yet part of the top-level tmol API\n", + "from tmol.chemical import ResidueTypeSet\n", + "\n", + "restype_set = ResidueTypeSet.from_database(default_param_db.chemical)\n", + "all_bt_names = [rt.name for rt in restype_set.residue_types]\n", + "print(\"n block types in the default database:\", len(all_bt_names))\n", + "\n", + "# the alternate protonation / tautomer states of histidine, for example\n", + "print(\"HIS block types:\", sorted(n for n in all_bt_names if n.startswith(\"HIS\")))\n", + "\n", + "# ... and the nucleic acid block types\n", + "print(\n", + " \"nucleic acid block types:\",\n", + " sorted(n for n in all_bt_names if n in (\"DA\", \"DC\", \"DG\", \"DT\", \"A\", \"C\", \"G\", \"U\")),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "id": "xClcO442jcaO" + }, + "outputs": [], + "source": [ + "# - Initialize a PackedBlockTypes object from the default database\n", + "pbt = tmol.default_packed_block_types(device)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": { + "id": "PyhxBHq9pLMx" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "block types added by ligand.tmol: ['LG1']\n", + "packed block types: ['ALA', 'GLY', 'SER', 'LG1']\n" + ] + } + ], + "source": [ + "# - Initialize a PackedBlockTypes object from a custom database\n", + "#\n", + "# inject_params_file returns a *new* ParameterDatabase that also contains the\n", + "# block type(s) described by a .tmol ligand-parameter file. A PackedBlockTypes\n", + "# built from that database can then represent the ligand.\n", + "\n", + "# TEMP: the ligand params machinery is not yet part of the top-level tmol API\n", + "from tmol.ligand import inject_params_file\n", + "\n", + "custom_param_db = inject_params_file(default_param_db, \"ligand.tmol\")\n", + "custom_restype_set = ResidueTypeSet.from_database(custom_param_db.chemical)\n", + "\n", + "added = sorted(\n", + " {rt.name for rt in custom_restype_set.residue_types}\n", + " - {rt.name for rt in restype_set.residue_types}\n", + ")\n", + "print(\"block types added by ligand.tmol:\", added)\n", + "\n", + "# pack a small set: a few amino acids plus the new ligand block type(s)\n", + "custom_subset = [\n", + " rt\n", + " for rt in custom_restype_set.residue_types\n", + " if rt.name in (\"ALA\", \"GLY\", \"SER\") or rt.name in added\n", + "]\n", + "custom_pbt = tmol.PackedBlockTypes.from_restype_list(\n", + " custom_param_db.chemical, custom_restype_set, custom_subset, device\n", + ")\n", + "print(\"packed block types:\", [bt.name for bt in custom_pbt.active_block_types])" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": { + "id": "OO32waIlp4-5" + }, + "outputs": [], + "source": [ + "# - Create a default CanonicalOrdering object\n", + "canonical_ordering = tmol.default_canonical_ordering()" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": { + "id": "L-uVsXPLjiDC" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n residue-type equivalence classes: 30\n" + ] + } + ], + "source": [ + "# - Create a CanonicalOrdering from a ParameterDatabase\n", + "#\n", + "# The canonical ordering assigns each residue-type equivalence class and each\n", + "# of its atoms a fixed index; it is the interchange format tmol uses when\n", + "# talking to the outside world. default_canonical_ordering() is exactly this\n", + "# call applied to the default database.\n", + "canonical_ordering_from_db = tmol.CanonicalOrdering.from_chemdb(\n", + " default_param_db.chemical\n", + ")\n", + "print(\n", + " \"n residue-type equivalence classes:\",\n", + " canonical_ordering_from_db.n_restype_io_equiv_classes,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": { + "id": "oXpe5H0Rjmfr" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n packed block types: 6\n", + "names: ['ALA', 'GLY', 'LEU', 'PHE', 'SER', 'VAL']\n" + ] + } + ], + "source": [ + "# - Create a PackedBlockTypes object from a subset of the block types\n", + "# in a ParameterDatabase\n", + "#\n", + "# Restricting the packed set keeps the score-term tensors small; it is useful\n", + "# when you know a priori which block types can appear.\n", + "subset_names = (\"ALA\", \"GLY\", \"SER\", \"LEU\", \"VAL\", \"PHE\")\n", + "subset_restypes = [rt for rt in restype_set.residue_types if rt.name in subset_names]\n", + "\n", + "pbt_subset = tmol.PackedBlockTypes.from_restype_list(\n", + " default_param_db.chemical, restype_set, subset_restypes, device\n", + ")\n", + "print(\"n packed block types:\", len(pbt_subset.active_block_types))\n", + "print(\"names:\", [bt.name for bt in pbt_subset.active_block_types])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "42PMj1JkkTaD" + }, + "source": [ + "Input" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "k_2XAIzBjqKn", + "outputId": "45d58612-7efc-4968-d77c-783e7404a70d" + }, + "outputs": [], + "source": [ + "# - Initialize a single pose PoseStack from a PDB\n", + "pose_1ubq = tmol.pose_stack_from_pdb(\"1ubq.pdb\", device=device)\n", + "assert pose_1ubq.max_n_blocks == 76" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": { + "id": "496P7xj4OCcr" + }, + "outputs": [], + "source": [ + "# - Initialize a pose stack from a subset of residues in a PDB\n", + "\n", + "# We have to tell the PDB reader that the first residue is not an\n", + "# N-terminus but merely is not connected to the residue that preceeds it\n", + "# and that the last residue is not a C-terminus.\n", + "res_not_connected = torch.zeros([1, 31, 2], dtype=torch.bool, device=device)\n", + "res_not_connected[0, 0, 0] = True\n", + "res_not_connected[0, 30, 1] = True\n", + "\n", + "pose_1ubq_20to50 = tmol.pose_stack_from_pdb(\"1ubq.pdb\", device=device, residue_start=20, residue_end=51, res_not_connected=res_not_connected)\n", + "assert pose_1ubq_20to50.inter_residue_connections[0, 0, 0, 0] == -1\n", + "assert pose_1ubq_20to50.inter_residue_connections[0, 30, 1, 0] == -1\n" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": { + "id": "N8RfPFCWYkfy" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n blocks: 76\n" + ] + } + ], + "source": [ + "# - Initialize a single pose PoseStack from a PDB file using biotite\n", + "import biotite.structure\n", + "import biotite.structure.io.pdb\n", + "\n", + "# TEMP: pose_stack_from_biotite is not yet part of the top-level tmol API\n", + "from tmol.io import pose_stack_from_biotite\n", + "\n", + "bt_pdb_file = biotite.structure.io.pdb.PDBFile.read(\"1ubq.pdb\")\n", + "bt_struct = bt_pdb_file.get_structure()\n", + "if isinstance(bt_struct, biotite.structure.AtomArrayStack):\n", + " bt_struct = bt_struct[0]\n", + "\n", + "pose_1ubq_bt = pose_stack_from_biotite(bt_struct, device)\n", + "print(\"n blocks:\", pose_1ubq_bt.max_n_blocks)" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "ysGCJuh8kXRc", + "outputId": "ea818011-19e4-4b37-c638-34ef10bef4d7" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n blocks: 388\n" + ] + }, + { + "output_type": "stream", + "name": "stderr", + "text": [ + "Unrecognized 3lc K\n" + ] + } + ], + "source": [ + "# - Initialize a single pose PoseStack from an mmcif file (.cif)\n", + "import biotite.structure\n", + "from biotite.structure.io.pdbx import CIFFile, set_structure\n", + "\n", + "# TEMP: pose_stack_from_biotite is not yet part of the top-level tmol API\n", + "from tmol.io import pose_stack_from_biotite\n", + "\n", + "bt_struct = biotite.structure.io.load_structure(\n", + " \"1BL8.cif\", extra_fields=[\"occupancy\", \"b_factor\"]\n", + ")\n", + "if isinstance(bt_struct, biotite.structure.AtomArrayStack):\n", + " bt_struct = bt_struct[0]\n", + "\n", + "pose_1bl8 = pose_stack_from_biotite(bt_struct, device)\n", + "print(\"n blocks:\", pose_1bl8.max_n_blocks)" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": { + "id": "bNOuYq4QkX7A" + }, + "outputs": [], + "source": [ + "# - Initialize a single pose PoseStack from an OpenFold set of tensors\n", + "\n", + "# here is a dictionary with the set of tensors that OpenFold produces when\n", + "# asked to predict the structures of ubiquitin and sumo; in particular, tmol\n", + "# reads from the \"aatype\", \"positions\" and \"chain_index\" tensors.\n", + "\n", + "openfold_ubq_and_sumo_pred = torch.load(\"openfold_ubq_and_sumo.pt\", map_location=device)\n", + "\n", + "ps_ubq_sumo = tmol.pose_stack_from_openfold(openfold_ubq_and_sumo_pred)\n", + "assert len(ps_ubq_sumo) == 2\n", + "assert ps_ubq_sumo.max_n_blocks == openfold_ubq_and_sumo_pred[\"positions\"].shape[2]\n", + "assert ps_ubq_sumo.coords.device == device\n" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": { + "id": "CgVHFNkVkYUZ" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n residues: 991 n chains: 3\n", + "chain-final blocks (no upper conn): [554, 768, 990]\n", + "chain-initial blocks (no lower conn): [0, 555, 769]\n", + "stretched peptide bonds (block_i, block_j, C-N dist): [(100, 101, 14.05)]\n" + ] + } + ], + "source": [ + "# - Initialize a single pose PoseStack with missing residues\n", + "pose_1s78 = tmol.pose_stack_from_pdb(\"1s78.pdb\", device=device)\n", + "\n", + "# Which blocks start / end a polymer run? The \"up\" and \"down\" connection\n", + "# indices are properties of each *block type*, so look them up per block\n", + "# rather than assuming a fixed connection number.\n", + "n_res_1s78 = int(pose_1s78.n_res_per_pose[0])\n", + "pbt_1s78 = pose_1s78.packed_block_types\n", + "bt64_1s78 = pose_1s78.block_type_ind64[0, :n_res_1s78]\n", + "irc_1s78 = pose_1s78.inter_residue_connections[0, :n_res_1s78]\n", + "ar_1s78 = torch.arange(n_res_1s78, device=device)\n", + "\n", + "\n", + "def _conn_absent(conn_inds):\n", + " conn_inds = conn_inds.to(torch.int64)\n", + " absent = torch.ones(n_res_1s78, dtype=torch.bool, device=device)\n", + " present = conn_inds >= 0\n", + " absent[present] = irc_1s78[ar_1s78[present], conn_inds[present], 0] == -1\n", + " return absent\n", + "\n", + "\n", + "no_up = _conn_absent(pbt_1s78.up_conn_inds[bt64_1s78])\n", + "no_down = _conn_absent(pbt_1s78.down_conn_inds[bt64_1s78])\n", + "n_chains_1s78 = len(torch.unique(pose_1s78.chain_id[0, :n_res_1s78]))\n", + "\n", + "print(\"n residues:\", n_res_1s78, \" n chains:\", n_chains_1s78)\n", + "print(\"chain-final blocks (no upper conn):\", torch.nonzero(no_up).flatten().tolist())\n", + "print(\"chain-initial blocks (no lower conn):\", torch.nonzero(no_down).flatten().tolist())\n", + "assert int(no_up.sum()) == n_chains_1s78\n", + "assert int(no_down.sum()) == n_chains_1s78\n", + "\n", + "# Here is the catch worth knowing about: chain A of 1s78 is missing residues\n", + "# 102-110, but pose_stack_from_pdb does *not* infer that from the residue\n", + "# numbering. It bonds the flanking residues to each other, which leaves a\n", + "# badly stretched peptide bond rather than a chain break. Such gaps are easy\n", + "# to find geometrically:\n", + "def _atom_xyz(pose_stack, block, atom_name):\n", + " bt = pose_stack.block_type(0, block)\n", + " if atom_name not in bt.atom_to_idx:\n", + " return None\n", + " offset = int(pose_stack.block_coord_offset64[0, block])\n", + " return pose_stack.coords[0, offset + bt.atom_to_idx[atom_name]]\n", + "\n", + "\n", + "stretched = []\n", + "for b in range(n_res_1s78 - 1):\n", + " if bool(no_up[b]):\n", + " continue # a chain end: genuinely not bonded to b + 1\n", + " c_xyz = _atom_xyz(pose_1s78, b, \"C\")\n", + " n_xyz = _atom_xyz(pose_1s78, b + 1, \"N\")\n", + " if c_xyz is None or n_xyz is None:\n", + " continue\n", + " dist = float(torch.linalg.norm(c_xyz - n_xyz))\n", + " if dist > 2.0: # an ideal peptide bond is ~1.33 A\n", + " stretched.append((b, b + 1, round(dist, 2)))\n", + "\n", + "print(\"stretched peptide bonds (block_i, block_j, C-N dist):\", stretched)\n", + "assert stretched, \"expected to detect the chain-A gap\"\n", + "\n", + "# To model the gap as a real chain break, pass res_not_connected when reading\n", + "# the file, exactly as in the residue-subset cell above." + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": { + "id": "HX9jTnCrkYry" + }, + "outputs": [], + "source": [ + "# - Load a bunch of PoseStacks from different sources and then concatenate\n", + "# them to a single PoseStack\n", + "ps_3plc = tmol.pose_stack_from_pdb(\"3plc.pdb\", device=device)\n", + "ps_1qys = tmol.pose_stack_from_pdb(\"1qys.pdb\", device=device)\n", + "\n", + "# TEMP: PoseStackBuilder is not yet part of the top-level tmol API\n", + "from tmol.pose import PoseStackBuilder\n", + "\n", + "pose_stack_3 = PoseStackBuilder.from_poses([pose_1ubq, ps_3plc, ps_1qys], device=device)\n", + "assert pose_stack_3.n_poses == 3" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": { + "id": "IpP5VgmbkZBz" + }, + "outputs": [], + "source": [ + "# - Make many copies of the same single-pose PoseStack\n", + "# You can create a list of shallow copies of a single-pose PoseStack\n", + "# and the PoseStackBuilder will expand them into complete and fully\n", + "# independent poses.\n", + "ten_1ubqs = PoseStackBuilder.from_poses([pose_1ubq] * 10, device=device)\n", + "assert ten_1ubqs.n_poses == 10\n", + "assert ten_1ubqs.coords.shape[0] == 10" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": { + "id": "zIacYunXkZX4" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n residue types described by ligand.tmol: 1\n", + "name3: LG1\n", + "n atoms: 57\n", + "n bonds: 60\n", + "first few atoms: ['C1', 'O1', 'C2', 'O2', 'C3', 'O3', 'C4', 'O4']\n", + "added block types: ['LG1']\n" + ] + } + ], + "source": [ + "# - Add custom residue type\n", + "#\n", + "# A .tmol parameter file describes one or more non-canonical block types:\n", + "# the chemical graph (atoms/bonds/torsions), the partial charges used by the\n", + "# elec term, and its cartbonded parameters. load_params_file parses it into\n", + "# LigandPreparation records, which are what actually get injected into a\n", + "# ParameterDatabase.\n", + "\n", + "# TEMP: the ligand params machinery is not yet part of the top-level tmol API\n", + "from tmol.ligand import load_params_file\n", + "from tmol.ligand._registry import inject_ligand_preparations\n", + "\n", + "preparations = load_params_file(\"ligand.tmol\")\n", + "print(\"n residue types described by ligand.tmol:\", len(preparations))\n", + "\n", + "prep = preparations[0]\n", + "print(\"name3:\", prep.residue_type.name3)\n", + "print(\"n atoms:\", len(prep.residue_type.atoms))\n", + "print(\"n bonds:\", len(prep.residue_type.bonds))\n", + "print(\"first few atoms:\", [a.name for a in prep.residue_type.atoms[:8]])\n", + "\n", + "# inject_ligand_preparations is the single chokepoint that extends a database;\n", + "# inject_params_file (used above) is just a convenience wrapper around it.\n", + "param_db_w_custom_bt = inject_ligand_preparations(default_param_db, preparations)\n", + "custom_names = {rt.name for rt in param_db_w_custom_bt.chemical.residues} - {\n", + " rt.name for rt in default_param_db.chemical.residues\n", + "}\n", + "print(\"added block types:\", sorted(custom_names))" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": { + "id": "B9BJmHWCkZxQ" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n partial charges: 57\n", + "a few original partial charges: [('C1', -0.288), ('O1', -0.57), ('C2', -0.136), ('O2', -0.68), ('C3', 0.541)]\n", + "partial charge on C1: -0.288 -> 0.0\n", + "cartbonded params present: True\n" + ] + } + ], + "source": [ + "# - Set the scoring parameters for a new residue type (elec, etc)\n", + "#\n", + "# The LigandPreparation carries the per-atom partial charges that the fa_elec\n", + "# term will use. It is a frozen dataclass, so dataclasses.replace() gives us a\n", + "# modified copy to inject, which is how you override the scoring parameters of\n", + "# a new block type.\n", + "import dataclasses\n", + "\n", + "print(\"n partial charges:\", len(prep.partial_charges))\n", + "some = list(prep.partial_charges.items())[:5]\n", + "print(\"a few original partial charges:\", some)\n", + "\n", + "# neutralize one atom's charge as a demonstration\n", + "altered_charges = dict(prep.partial_charges)\n", + "first_atom = next(iter(altered_charges))\n", + "altered_charges[first_atom] = 0.0\n", + "\n", + "altered_prep = dataclasses.replace(prep, partial_charges=altered_charges)\n", + "param_db_w_altered_elec = inject_ligand_preparations(\n", + " default_param_db, [altered_prep]\n", + ")\n", + "\n", + "print(\n", + " f\"partial charge on {first_atom}: \"\n", + " f\"{prep.partial_charges[first_atom]} -> {altered_prep.partial_charges[first_atom]}\"\n", + ")\n", + "\n", + "# the cartbonded parameters for the new block type travel with it too\n", + "print(\"cartbonded params present:\", prep.cartbonded_params is not None)" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": { + "id": "MTd5eVotkaJ9" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n poses: 2\n", + "residues per pose: [20, 10]\n", + "scores: tensor([8233.6143, 53.6222], device='cuda:0')\n" + ] + } + ], + "source": [ + "# - Build an extended pose from sequence\n", + "#\n", + "# extended_pose_stack_from_sequences builds ideal-geometry poses with extended\n", + "# backbone torsions. Pass a list of sequences to build a batch in one call;\n", + "# the poses are padded to the length of the longest one.\n", + "extended_ps = tmol.extended_pose_stack_from_sequences(\n", + " [\"ACDEFGHIKLMNPQRSTVWY\", \"GGGGSGGGGS\"], device=device\n", + ")\n", + "print(\"n poses:\", extended_ps.n_poses)\n", + "print(\"residues per pose:\", extended_ps.n_res_per_pose.tolist())\n", + "\n", + "sfxn_ext = tmol.beta2016_score_function(device)\n", + "wpsm_ext = sfxn_ext.render_whole_pose_scoring_module(extended_ps)\n", + "print(\"scores:\", wpsm_ext(extended_ps.coords))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "R1LJAlYJktGE" + }, + "source": [ + "Output" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": { + "id": "WAYlCkZ8kNh_" + }, + "outputs": [], + "source": [ + "# - Write a single pose PoseStack to a .pdb file\n", + "tmol.write_pose_stack_pdb(pose_1ubq, \"1ubq_out.pdb\")\n", + "assert os.path.isfile(\"1ubq_out.pdb\")" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "metadata": { + "id": "a58GBcGrk0WW" + }, + "outputs": [], + "source": [ + "# - Write a multi-pose stack out as a multi-model PDB file\n", + "tmol.write_pose_stack_pdb(ten_1ubqs, \"ten_1ubqs.pdb\")\n", + "assert os.path.isfile(\"ten_1ubqs.pdb\")\n", + "def nlines_from_file(fname):\n", + " with open(fname) as fid:\n", + " lines = fid.readlines()\n", + " return len(lines)\n", + "nlines_1ubq = nlines_from_file(\"1ubq.pdb\")\n", + "nlines_ten_1ubqs = nlines_from_file(\"ten_1ubqs.pdb\")\n", + "assert 10 * nlines_1ubq <= nlines_ten_1ubqs" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": { + "id": "sPGfDWNrk0zL" + }, + "outputs": [], + "source": [ + "# - Write a multi-pose stack out to separate PDB files\n", + "\n", + "# take advantage of PoseStack's split() method to create a single-pose PoseStack\n", + "for i in range(10):\n", + " pose_i = ten_1ubqs.split(i)\n", + " tmol.write_pose_stack_pdb(pose_i, f\"1ubq_{i:04}.pdb\")\n", + "\n", + "for i in range(10):\n", + " assert os.path.isfile(f\"1ubq_{i:04}.pdb\")" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": { + "id": "wYa6C8EGk1GT" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "wrote 1ubq_out.cif; 96653 bytes\n" + ] + } + ], + "source": [ + "# - Write a single pose PoseStack to an .mmcif file\n", + "#\n", + "# biotite_from_pose_stack converts back to a biotite AtomArray, which biotite\n", + "# can then write in any format it supports, including mmCIF.\n", + "\n", + "# TEMP: biotite_from_pose_stack is not yet part of the top-level tmol API\n", + "from tmol.io import biotite_from_pose_stack\n", + "from biotite.structure.io.pdbx import CIFFile, set_structure\n", + "\n", + "bt_out = biotite_from_pose_stack(pose_1ubq)\n", + "cif_out = CIFFile()\n", + "set_structure(cif_out, bt_out)\n", + "cif_out.write(\"1ubq_out.cif\")\n", + "\n", + "assert os.path.isfile(\"1ubq_out.cif\")\n", + "print(\"wrote 1ubq_out.cif;\", os.path.getsize(\"1ubq_out.cif\"), \"bytes\")" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": { + "id": "gcNdsybLk1dR" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "built 6 rotamers for block 22\n", + "wrote rotamers.pdb with 6 models\n" + ] + } + ], + "source": [ + "# - Write a rotamer set out as a multi-model PDB\n", + "#\n", + "# build_rotamers produces every conformer the packer would consider for a\n", + "# task. Here we restrict the task to a single residue, build its rotamers,\n", + "# and write each one as its own MODEL record.\n", + "\n", + "# TEMP: the rotamer machinery is not yet part of the top-level tmol API\n", + "from tmol.pack import PackerPalette, PackerTask, SetPackerTask\n", + "from tmol.pack.rotamer import FixedAAChiSampler, build_rotamers\n", + "from tmol.pack.rotamer.dunbrack import create_dunbrack_sampler_from_database\n", + "\n", + "rotset_sampler = create_dunbrack_sampler_from_database(default_param_db, device)\n", + "\n", + "ROT_BLOCK = 22 # an arbitrary residue to enumerate rotamers for\n", + "rot_task = PackerTask(pose_1ubq, PackerPalette())\n", + "rot_task.restrict_to_repacking()\n", + "\n", + "# disable packing everywhere except ROT_BLOCK\n", + "disable = torch.ones(\n", + " (pose_1ubq.n_poses, pose_1ubq.max_n_blocks), dtype=torch.bool, device=device\n", + ")\n", + "disable[0, ROT_BLOCK] = False\n", + "rot_task.disable_packing_by_block_mask(disable)\n", + "rot_task.add_conformer_sampler(rotset_sampler)\n", + "rot_task.add_conformer_sampler(FixedAAChiSampler())\n", + "\n", + "set_task = SetPackerTask.from_packer_task(rot_task)\n", + "rot_pose_stack, rotamer_set = build_rotamers(\n", + " pose_1ubq, set_task, pose_1ubq.packed_block_types.chem_db\n", + ")\n", + "\n", + "n_rots = int(rotamer_set.n_rots_for_block[0, ROT_BLOCK])\n", + "rot_offset = int(rotamer_set.rot_offset_for_block[0, ROT_BLOCK])\n", + "print(f\"built {n_rots} rotamers for block {ROT_BLOCK}\")\n", + "\n", + "pbt_rot = rot_pose_stack.packed_block_types\n", + "rot_coords = rotamer_set.coords.detach().cpu()\n", + "coord_offsets = rotamer_set.coord_offset_for_rot.cpu()\n", + "bt_for_rot = rotamer_set.block_type_ind_for_rot.cpu()\n", + "\n", + "with open(\"rotamers.pdb\", \"w\") as fid:\n", + " for model, rot in enumerate(range(rot_offset, rot_offset + n_rots), start=1):\n", + " bt = pbt_rot.active_block_types[int(bt_for_rot[rot])]\n", + " start = int(coord_offsets[rot])\n", + " fid.write(f\"MODEL {model:>4}\\n\")\n", + " for i, atom in enumerate(bt.atoms):\n", + " x, y, z = rot_coords[start + i].tolist()\n", + " fid.write(\n", + " f\"ATOM {i + 1:>5} {atom.name:<4}{bt.name3:>4} A\"\n", + " f\"{ROT_BLOCK + 1:>4} {x:8.3f}{y:8.3f}{z:8.3f} 1.00 0.00\\n\"\n", + " )\n", + " fid.write(\"ENDMDL\\n\")\n", + "\n", + "assert os.path.isfile(\"rotamers.pdb\")\n", + "with open(\"rotamers.pdb\") as fid:\n", + " n_models = sum(1 for line in fid if line.startswith(\"MODEL\"))\n", + "print(\"wrote rotamers.pdb with\", n_models, \"models\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "vucunYYYk2OX" + }, + "source": [ + "Kinematics" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": { + "id": "PYwjsJDvk4Ka" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "edges shape: (1, 2, 4) n_edges: [2]\n" + ] + } + ], + "source": [ + "# - Create an N->C fold tree (a fold *forest* in tmol) for a PoseStack\n", + "import numpy\n", + "\n", + "# Option 1: reasonable_fold_forest() produces a simple fold tree for\n", + "# each pose in the PoseStack of one N->C edge per chain.\n", + "ff1_1ubq = tmol.FoldForest.reasonable_fold_forest(pose_1ubq)\n", + "\n", + "# Option 2: specify the fold forest explicitly.\n", + "# Each edge is a 4-tuple of (edge_type, start_residue, end_residue, jump_index)\n", + "#\n", + "# tmol includes a new kind of jump edge: a root-jump. This connects\n", + "# the downstream residue to the root of the fold forest. Each\n", + "# fold tree must include at least one root-jump edge, but may\n", + "# include arbitrarily many. The start_residue of a root-jump edge\n", + "# is always the sentinel value of -1. The jump_index column is -1 for\n", + "# polymer edges and for root-jumps; ordinary jump edges must carry\n", + "# contiguous jump indices starting at 0.\n", + "#\n", + "# For this N->C fold tree, we define a root-jump to residue 0, and a polymer\n", + "# edge from 0 to the last residue in the pose.\n", + "#\n", + "# The edges array should be [n_poses, max_n_edges, 4] with the sentinel value\n", + "# of -1 in the edge-type column for unused edges.\n", + "edges = numpy.full((1, 2, 4), -1, dtype=int)\n", + "edges[0, 0] = [tmol.EdgeType.root_jump, -1, 0, -1]\n", + "edges[0, 1] = [tmol.EdgeType.polymer, 0, 75, -1]\n", + "ff2_1ubq = tmol.FoldForest.from_edges(edges)\n", + "print(\"edges shape:\", ff2_1ubq.edges.shape, \" n_edges:\", ff2_1ubq.n_edges)" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": { + "id": "LVS2vrVRZGbS" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n_edges per pose: [2 2 2]\n" + ] + } + ], + "source": [ + "# - Create a set of N->C fold trees for a multi-pose PoseStack of\n", + "# single-chain PDBs\n", + "ff1_ps3 = tmol.FoldForest.reasonable_fold_forest(pose_stack_3)\n", + "\n", + "# Option 2: build the edges array explicitly, as above.\n", + "edges = numpy.full((3, 2, 4), -1, dtype=int)\n", + "edges[:, 0] = numpy.array([tmol.EdgeType.root_jump, -1, 0, -1], dtype=int)[None, :]\n", + "n_res = pose_stack_3.n_res_per_pose.cpu().numpy()\n", + "edges[:, 1, 0] = tmol.EdgeType.polymer\n", + "edges[:, 1, 1] = 0\n", + "edges[:, 1, 2] = n_res - 1\n", + "edges[:, 1, 3] = -1\n", + "ff2_ps3 = tmol.FoldForest.from_edges(edges)\n", + "print(\"n_edges per pose:\", ff2_ps3.n_edges)" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": { + "id": "OBAw1Lw6lKKP" + }, + "outputs": [], + "source": [ + "# - Create a simple fold tree for a multi-chain PoseStack\n", + "ps_10vb = tmol.pose_stack_from_pdb(\"10VB.pdb\", device=device)\n", + "ff_10vb = tmol.FoldForest.reasonable_fold_forest(ps_10vb)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "metadata": { + "id": "e7m9yUztlKmg" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n edges: 6\n", + "edge kinds: {'polymer': 3, 'root_jump': 3}\n" + ] + } + ], + "source": [ + "# - Create a simple fold tree for a PoseStack with missing residues\n", + "#\n", + "# reasonable_fold_forest() understands chain breaks: a polymer gap *within*\n", + "# one chain becomes an ordinary jump edge, while a gap between two chains\n", + "# produces a separate root-jump. So a structure with missing residues just\n", + "# works.\n", + "ff_1s78 = tmol.FoldForest.reasonable_fold_forest(pose_1s78)\n", + "\n", + "edge_names = {int(e): e.name for e in tmol.EdgeType}\n", + "kinds = [\n", + " edge_names[int(ff_1s78.edges[0, i, 0])] for i in range(int(ff_1s78.n_edges[0]))\n", + "]\n", + "print(\"n edges:\", int(ff_1s78.n_edges[0]))\n", + "print(\"edge kinds:\", {k: kinds.count(k) for k in set(kinds)})" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": { + "id": "iblPqPjZlK8Z" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n edges per pose: [76 54 92]\n" + ] + } + ], + "source": [ + "# - Create a dandelion fold tree for a PoseStack\n", + "\n", + "# Many NNs, such as OpenFold, produce structures as one coordinate frame\n", + "# per residue and dihedrals for the side chains. Kinematically,\n", + "# we can reproduce that by describing a tree with one root-jump\n", + "# for every residue -- the side chains will still be described\n", + "# with internal geometries. This system is like a dandelion with\n", + "# tons of tiny stalks holding tiny seeds at the end all connected\n", + "# to a central hub. With such a kin_forest, e.g., it is possible to minimize\n", + "# the system in the same set of DOFs that the NN has access to.\n", + "\n", + "edges = numpy.full((3, pose_stack_3.max_n_blocks, 4), -1, dtype=int)\n", + "edges[:, :, 0] = tmol.EdgeType.root_jump\n", + "edges[:, :, 1] = -1\n", + "edges[:, :, 2] = numpy.arange(pose_stack_3.max_n_blocks, dtype=int)[None, :]\n", + "edges[:, :, 3] = -1\n", + "\n", + "# sentinel out the edges that are out-of-bounds for each pose\n", + "is_unreal_res = (pose_stack_3.block_type_ind == -1).cpu().numpy()\n", + "edges[is_unreal_res, :] = -1\n", + "\n", + "dandelion_ff = tmol.FoldForest.from_edges(edges)\n", + "print(\"n edges per pose:\", dandelion_ff.n_edges)" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": { + "id": "pUBbckP7lLRM" + }, + "outputs": [], + "source": [ + "# - Create a MoveMap that enables minimization for named torsions\n", + "mm_1ubq_all = tmol.MoveMap.from_pose_stack(pose_1ubq)\n", + "mm_1ubq_all.move_all_named_torsions = True" + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "metadata": { + "id": "fJUxnRqYlLnW" + }, + "outputs": [], + "source": [ + "# - Create a MoveMap that enables backbone minimization for some residues but not all\n", + "\n", + "# Let's turn on named-torsion minimization for residues 10-39\n", + "mm_1ubq_some = tmol.MoveMap.from_pose_stack(pose_1ubq)\n", + "mm_1ubq_some.set_move_all_named_torsions_for_blocks(0, torch.arange(30, dtype=torch.int64, device=device) + 10)" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "metadata": { + "id": "ZTRWFlPslL9b" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "194 residues; chain 2 starts at block 108\n", + "jump nodes in the kinematic forest: [2, 1821]\n", + "mean displacement, chain 1 block 0: 0.0000 A\n", + "mean displacement, chain 2 block 108: 1.0096 A\n" + ] + } + ], + "source": [ + "# - Apply a perturbation to the rigid-body DOFs between two chains\n", + "#\n", + "# 10VB has two chains. If we build a fold forest in which chain 2 hangs off\n", + "# chain 1 by an ordinary *jump* edge, the six rigid-body DOFs of that jump\n", + "# control the placement of chain 2 relative to chain 1. We read the current\n", + "# DOFs out with inverse kinematics, perturb the jump, and refold.\n", + "import attr\n", + "\n", + "# TEMP: the kinematic DOF machinery is not yet part of the top-level tmol API\n", + "from tmol.kinematics import JumpDOFTypes, NodeType, PoseStackKinematicsModule\n", + "from tmol.kinematics.compiled import inverse_kin\n", + "\n", + "n_res_10vb = int(ps_10vb.n_res_per_pose[0])\n", + "chain_ids = ps_10vb.chain_id[0, :n_res_10vb].cpu().numpy()\n", + "chain2_start = int(numpy.nonzero(chain_ids != chain_ids[0])[0][0])\n", + "print(f\"{n_res_10vb} residues; chain 2 starts at block {chain2_start}\")\n", + "\n", + "edges = numpy.full((1, 4, 4), -1, dtype=int)\n", + "edges[0, 0] = [tmol.EdgeType.root_jump, -1, 0, -1]\n", + "edges[0, 1] = [tmol.EdgeType.polymer, 0, chain2_start - 1, -1]\n", + "edges[0, 2] = [tmol.EdgeType.jump, 0, chain2_start, 0] # jump indices start at 0\n", + "edges[0, 3] = [tmol.EdgeType.polymer, chain2_start, n_res_10vb - 1, -1]\n", + "ff_10vb_jump = tmol.FoldForest.from_edges(edges)\n", + "\n", + "kin_module = PoseStackKinematicsModule(ps_10vb, ff_10vb_jump)\n", + "kmd = kin_module.kmd\n", + "ids = kmd.forest.id[1:].to(torch.int64)\n", + "\n", + "kincoords = torch.zeros(\n", + " (kmd.forest.id.shape[0], 3), dtype=torch.float32, device=device\n", + ")\n", + "kincoords[1:] = ps_10vb.coords.view(-1, 3)[ids]\n", + "dofs = inverse_kin(\n", + " kincoords,\n", + " kmd.forest.parent,\n", + " kmd.forest.frame_x,\n", + " kmd.forest.frame_y,\n", + " kmd.forest.frame_z,\n", + " kmd.forest.doftype,\n", + ")\n", + "\n", + "# the root-jump comes first; the inter-chain jump is the other jump node\n", + "jump_nodes = torch.nonzero(kmd.forest.doftype == NodeType.jump, as_tuple=True)[0]\n", + "print(\"jump nodes in the kinematic forest:\", jump_nodes.tolist())\n", + "inter_chain_jump = int(jump_nodes[-1])\n", + "\n", + "# translate chain 2 by 1 A along the jump frame's x axis and tilt it slightly\n", + "dofs[inter_chain_jump, JumpDOFTypes.RBx] += 1.0\n", + "dofs[inter_chain_jump, JumpDOFTypes.RBdel_alpha] += 0.05\n", + "\n", + "new_kincoords = kin_module(dofs)\n", + "flat = ps_10vb.coords.reshape(-1, 3).clone()\n", + "flat[ids] = new_kincoords[1:].to(flat.dtype)\n", + "ps_10vb_perturbed = attr.evolve(ps_10vb, coords=flat.view(ps_10vb.coords.shape))\n", + "\n", + "\n", + "def mean_block_displacement(a, b, block):\n", + " start = int(a.block_coord_offset64[0, block])\n", + " n_at = int(a.packed_block_types.n_atoms[int(a.block_type_ind64[0, block])])\n", + " d = a.coords[0, start : start + n_at] - b.coords[0, start : start + n_at]\n", + " return float(torch.linalg.norm(d, dim=-1).mean())\n", + "\n", + "\n", + "chain1_move = mean_block_displacement(ps_10vb, ps_10vb_perturbed, 0)\n", + "chain2_move = mean_block_displacement(ps_10vb, ps_10vb_perturbed, chain2_start)\n", + "print(f\"mean displacement, chain 1 block 0: {chain1_move:.4f} A\")\n", + "print(f\"mean displacement, chain 2 block {chain2_start}: {chain2_move:.4f} A\")\n", + "\n", + "# chain 1 is untouched; chain 2 moved as a rigid body\n", + "assert chain1_move < 1e-3\n", + "assert chain2_move > 0.5" + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "metadata": { + "id": "nM9g3oIwlMTy" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "block 10 torsions after assignment: {'phi': -57.0, 'psi': -47.0}\n", + "mean atom displacement vs. the crystal structure: 45.68 A\n" + ] + } + ], + "source": [ + "# - Assign dihedral values to all the residues in a PoseStack and\n", + "# calculate the coordinates\n", + "#\n", + "# set_named_torsions applies a batch of named-torsion assignments in one\n", + "# kinematic pass and returns a new PoseStack with the refolded coordinates.\n", + "# Which side of each bond stays fixed follows the fold forest.\n", + "\n", + "# Drive ubiquitin to ideal alpha-helical backbone torsions.\n", + "n_res_ubq = int(pose_1ubq.n_res_per_pose[0])\n", + "poses, blocks, names, values = [], [], [], []\n", + "for block in range(1, n_res_ubq - 1): # phi/psi are undefined at the termini\n", + " for name, value in ((\"phi\", -57.0), (\"psi\", -47.0)):\n", + " poses.append(0)\n", + " blocks.append(block)\n", + " names.append(name)\n", + " values.append(value)\n", + "\n", + "helical_1ubq = tmol.set_named_torsions(\n", + " pose_1ubq, poses, blocks, names, values, degrees=True\n", + ")\n", + "\n", + "# read the torsions back to confirm they took\n", + "measured = tmol.get_named_torsions(helical_1ubq, 0, 10, [\"phi\", \"psi\"], degrees=True)\n", + "print(\n", + " \"block 10 torsions after assignment:\",\n", + " {k: round(v, 2) for k, v in measured.items()},\n", + ")\n", + "\n", + "# the coordinates really did change\n", + "rmsd = float(\n", + " torch.linalg.norm(\n", + " helical_1ubq.coords[pose_1ubq.real_atoms]\n", + " - pose_1ubq.coords[pose_1ubq.real_atoms],\n", + " dim=-1,\n", + " ).mean()\n", + ")\n", + "print(f\"mean atom displacement vs. the crystal structure: {rmsd:.2f} A\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "klTMFJK6lQD1" + }, + "source": [ + "Scoring" + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "metadata": { + "id": "EnQIxSW4lWjZ" + }, + "outputs": [], + "source": [ + "# - Create the default score function\n", + "sfxn = tmol.beta2016_score_function(device)" + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "metadata": { + "id": "Yc8a4BnklZoh" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n terms: 11\n" + ] + } + ], + "source": [ + "# - Create the default score function from a custom Database\n", + "#\n", + "# Passing param_db bypasses the memoized default and builds a score function\n", + "# whose terms are parameterized from your database. Here we use the\n", + "# ligand-extended database built above.\n", + "sfxn_custom_db = tmol.beta2016_score_function(device, param_db=custom_param_db)\n", + "print(\"n terms:\", len(sfxn_custom_db.all_terms()))" + ] + }, + { + "cell_type": "code", + "execution_count": 42, + "metadata": { + "id": "5qErnZZSlbpA" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "available weight files: ['beta2016.sfxn', 'beta_soft.sfxn']\n", + "beta2016 score: tensor([235.4993], device='cuda:0')\n", + "soft-rep score: tensor([203.0768], device='cuda:0')\n", + "beta2016 weights + soft_rep option: tensor([125.0659], device='cuda:0')\n" + ] + } + ], + "source": [ + "# - Create the soft-rep version of the score function\n", + "#\n", + "# tmol ships weight files for its standard score functions; beta_soft.sfxn\n", + "# is the soft-repulsive variant (it softens the LJ repulsive shoulder and\n", + "# reweights accordingly). from_sfxn_file reads one of these YAML files.\n", + "from pathlib import Path\n", + "\n", + "sfxn_dir = Path(tmol.__file__).parent / \"database\" / \"score_functions\"\n", + "print(\"available weight files:\", sorted(p.name for p in sfxn_dir.glob(\"*.sfxn\")))\n", + "\n", + "sfxn_soft = tmol.ScoreFunction.from_sfxn_file(\n", + " sfxn_dir / \"beta_soft.sfxn\", default_param_db, device\n", + ")\n", + "\n", + "wpsm_hard = sfxn.render_whole_pose_scoring_module(pose_1ubq)\n", + "wpsm_soft = sfxn_soft.render_whole_pose_scoring_module(pose_1ubq)\n", + "print(\"beta2016 score: \", wpsm_hard(pose_1ubq.coords))\n", + "print(\"soft-rep score: \", wpsm_soft(pose_1ubq.coords))\n", + "\n", + "# The same softening can be requested on any score function via its options.\n", + "sfxn_soft_opt = tmol.beta2016_score_function(device, param_db=default_param_db)\n", + "sfxn_soft_opt.set_option(\"soft_rep\", True)\n", + "print(\n", + " \"beta2016 weights + soft_rep option:\",\n", + " sfxn_soft_opt.render_whole_pose_scoring_module(pose_1ubq)(pose_1ubq.coords),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 43, + "metadata": { + "id": "jLR7q9nXlc4H" + }, + "outputs": [], + "source": [ + "# - Create an empty score function\n", + "sfxn_empty = tmol.ScoreFunction(default_param_db, device)" + ] + }, + { + "cell_type": "code", + "execution_count": 44, + "metadata": { + "id": "r4vz-xXpmYV1" + }, + "outputs": [], + "source": [ + "# - Turn on a few terms in a score function\n", + "sfxn_rep_and_hbonds = tmol.ScoreFunction(default_param_db, device)\n", + "sfxn_rep_and_hbonds.set_weight(tmol.ScoreType.fa_ljrep, 0.55)\n", + "sfxn_rep_and_hbonds.set_weight(tmol.ScoreType.hbond, 1.0)" + ] + }, + { + "cell_type": "code", + "execution_count": 45, + "metadata": { + "id": "yMgEZFkQmYt_" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n terms to start: 2\n", + "after zeroing hbond: 1\n", + "after zeroing fa_ljrep: 0\n" + ] + } + ], + "source": [ + "# - Turn off a term in a score function\n", + "#\n", + "# Setting a weight to zero removes the term entirely.\n", + "print(\"n terms to start:\", len(sfxn_rep_and_hbonds.all_terms()))\n", + "\n", + "sfxn_rep_and_hbonds.set_weight(tmol.ScoreType.hbond, 0)\n", + "print(\"after zeroing hbond:\", len(sfxn_rep_and_hbonds.all_terms()))\n", + "assert len(sfxn_rep_and_hbonds.all_terms()) == 1\n", + "\n", + "sfxn_rep_and_hbonds.set_weight(tmol.ScoreType.fa_ljrep, 0)\n", + "print(\"after zeroing fa_ljrep:\", len(sfxn_rep_and_hbonds.all_terms()))\n", + "assert len(sfxn_rep_and_hbonds.all_terms()) == 0" + ] + }, + { + "cell_type": "code", + "execution_count": 46, + "metadata": { + "id": "EV89BMNEmZC-" + }, + "outputs": [], + "source": [ + "# - Score a PoseStack\n", + "wpsm = sfxn.render_whole_pose_scoring_module(pose_1ubq)\n", + "score = wpsm(pose_1ubq.coords)" + ] + }, + { + "cell_type": "code", + "execution_count": 47, + "metadata": { + "id": "XWLjbPsxmZYR" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "score: tensor([235.4993], device='cuda:0', grad_fn=)\n", + "grad shape: torch.Size([1, 1231, 3])\n", + "max |dE/dx|: 73.0069808959961\n" + ] + } + ], + "source": [ + "# - Score a PoseStack and back-propagate through the coordinates\n", + "#\n", + "# PoseStack coordinates do not require grad by default; opt in on a copy of\n", + "# the coordinate tensor and the whole score becomes differentiable.\n", + "coords_for_grad = pose_1ubq.coords.detach().clone().requires_grad_(True)\n", + "score2 = wpsm(coords_for_grad)\n", + "score2.sum().backward()\n", + "\n", + "print(\"score:\", score2)\n", + "print(\"grad shape:\", coords_for_grad.grad.shape)\n", + "print(\"max |dE/dx|:\", float(coords_for_grad.grad.abs().max()))" + ] + }, + { + "cell_type": "code", + "execution_count": 48, + "metadata": { + "id": "osRew_fxmZzF" + }, + "outputs": [], + "source": [ + "# - Score a PoseStack and return per-residue weighted energies\n", + "bpsm = sfxn.render_block_pair_scoring_module(pose_1ubq)\n", + "block_pair_scores = bpsm(pose_1ubq.coords)\n", + "assert block_pair_scores.shape == (1, pose_1ubq.max_n_blocks, pose_1ubq.max_n_blocks)" + ] + }, + { + "cell_type": "code", + "execution_count": 49, + "metadata": {}, + "outputs": [], + "source": [ + "# - Score a PoseStack and return per-residue un-weighted energies\n", + "unweighted_block_pair_scores = bpsm.unweighted_scores(pose_1ubq.coords)\n", + "assert unweighted_block_pair_scores.shape == (len(sfxn.all_score_types()), 1, pose_1ubq.max_n_blocks, pose_1ubq.max_n_blocks)" + ] + }, + { + "cell_type": "code", + "execution_count": 50, + "metadata": { + "id": "dYzEVvc7maRK" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "tot reg tensor(19602.5605, device='cuda:0', grad_fn=) tot_upweight_iface tensor(19503.5449, device='cuda:0', grad_fn=)\n", + "coord grad is populated: True\n" + ] + } + ], + "source": [ + "# - Score a PoseStack and return per-residue weighted energies, weight them\n", + "# according to some principle, and then back-propagate the total reweighted\n", + "# energy.\n", + "pose_1s78.coords.requires_grad_()\n", + "bpsm_1s78 = sfxn.render_block_pair_scoring_module(pose_1s78)\n", + "bps = bpsm_1s78(pose_1s78.coords)\n", + "tot_reg = bps.sum()\n", + "\n", + "# 1s78 is an antibody/antigen complex. Upweight every *inter-chain* residue\n", + "# pair by 2x, leaving the intra-chain pairs alone.\n", + "n_blocks_1s78 = pose_1s78.max_n_blocks\n", + "chain_of_block = pose_1s78.chain_id[0]\n", + "same_chain = chain_of_block[:, None] == chain_of_block[None, :]\n", + "\n", + "reweight = torch.ones(\n", + " (1, n_blocks_1s78, n_blocks_1s78), dtype=bps.dtype, device=device\n", + ")\n", + "reweight[0][~same_chain] = 2.0\n", + "\n", + "tot_upweight_iface = (bps * reweight).sum()\n", + "print(\"tot reg\", tot_reg, \"tot_upweight_iface\", tot_upweight_iface)\n", + "tot_upweight_iface.backward()\n", + "print(\"coord grad is populated:\", pose_1s78.coords.grad is not None)" + ] + }, + { + "cell_type": "code", + "execution_count": 51, + "metadata": { + "id": "-NW-sfNJmapX" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "start_ca_coords torch.Size([76, 3])\n", + "ca_dists torch.Size([76, 76])\n" + ] + } + ], + "source": [ + "# - Add constraints to a PoseStack\n", + "\n", + "# The PoseStack carries a ConstraintSet object that may be shared between\n", + "# multiple PoseStacks; thus ConstraintSet is immutable. Instead of\n", + "# being able to modify a ConstraintSet, the class makes it easy to\n", + "# create a new ConstraintSet with the contents you want.\n", + "# PoseStack is the same way: you cannot modify the ConstraintSet in \n", + "# an existing PoseStack, but using attr.evolve(...) you can readily\n", + "# construct a new PoseStack that holds your newly constructed \n", + "# ConstraintSet\n", + "import attr\n", + "\n", + "assert pose_1ubq.constraint_set is None\n", + "cst_set = tmol.ConstraintSet.create_empty(device=device, n_poses=pose_1ubq.n_poses)\n", + "\n", + "# let's create a constraint set with CA distances\n", + "# between all pairs of residues\n", + "start_ca_coords = []\n", + "res_inds = []\n", + "atom_inds = []\n", + "coords = pose_1ubq.coords.cpu()\n", + "for i in range(pose_1ubq.max_n_blocks):\n", + " i_bt = pose_1ubq.block_type(0, i)\n", + " # In general, not all block types will have a \"CA\" atom\n", + " # but in the case of 1ubq, they do happen to.\n", + " if \"CA\" in i_bt.atom_names_set:\n", + " res_inds.append(i)\n", + " at_ind = i_bt.atom_to_idx[\"CA\"]\n", + " atom_inds.append(at_ind)\n", + " start_ca_coords.append(coords[0, pose_1ubq.block_coord_offset64[0, i] + at_ind])\n", + "\n", + "start_ca_coords = torch.stack(start_ca_coords).to(device=device)\n", + "print(\"start_ca_coords\", start_ca_coords.shape)\n", + "res_inds = torch.tensor(res_inds, dtype=torch.int64, device=device)\n", + "atom_inds = torch.tensor(atom_inds, dtype=torch.int64, device=device)\n", + "\n", + "ca_dists = torch.linalg.norm(start_ca_coords[None, :, :] - start_ca_coords[:, None, :], dim=2)\n", + "print(\"ca_dists\", ca_dists.shape)\n", + "n_res_arange = torch.arange(len(start_ca_coords), dtype=torch.int64)\n", + "is_upper_triangle = n_res_arange[:, None] < n_res_arange[None, :]\n", + "nz_upper_triangle_r1, nz_upper_triangle_r2 = torch.nonzero(is_upper_triangle, as_tuple=True)\n", + "\n", + "n_csts = nz_upper_triangle_r1.shape[0]\n", + "cst_atoms = torch.zeros((n_csts, 2, 3), dtype=torch.int64, device=device)\n", + "cst_params = torch.zeros((n_csts, 4), dtype=torch.float32, device=device)\n", + "cst_atoms[:, :, 0] = 0 # pose index\n", + "cst_atoms[:, 0, 1] = res_inds[nz_upper_triangle_r1] # atom1 residue index\n", + "cst_atoms[:, 1, 1] = res_inds[nz_upper_triangle_r2] # atom2 residue index\n", + "cst_atoms[:, 0, 2] = atom_inds[nz_upper_triangle_r1] # atom1 atom index within its residue\n", + "cst_atoms[:, 1, 2] = atom_inds[nz_upper_triangle_r2] # atom2 atom index within its residue\n", + "\n", + "cst_params[:, 0] = ca_dists[nz_upper_triangle_r1, nz_upper_triangle_r2]\n", + "cst_params[:, 1] = 0.5 # 0.5A standard deviation\n", + "\n", + "cst_set = cst_set.add_constraints(\n", + " tmol.ConstraintEnergyTerm.harmonic,\n", + " cst_atoms,\n", + " cst_params\n", + ")\n", + "pose_1ubq_w_csts = attr.evolve(pose_1ubq, constraint_set=cst_set)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 52, + "metadata": {}, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "score w/ constraints tensor([235.4993], device='cuda:0')\n" + ] + } + ], + "source": [ + "wpsm_w_csts = sfxn.render_whole_pose_scoring_module(pose_1ubq_w_csts)\n", + "score = wpsm_w_csts(pose_1ubq_w_csts.coords)\n", + "print(\"score w/ constraints\", score)" + ] + }, + { + "cell_type": "code", + "execution_count": 53, + "metadata": { + "id": "swFqps8PmbCS" + }, + "outputs": [], + "source": [ + "# - Add the same constraints to all the Poses in a PoseStack\n", + "cst_set_10 = tmol.ConstraintSet.create_empty(device=device, n_poses=ten_1ubqs.n_poses)\n", + "cst_set_10 = cst_set_10.add_constraints_to_all_poses(tmol.ConstraintEnergyTerm.harmonic, cst_atoms, cst_params)\n", + "ten_1ubqs_w_csts = attr.evolve(ten_1ubqs, constraint_set=cst_set_10)" + ] + }, + { + "cell_type": "code", + "execution_count": 54, + "metadata": { + "id": "A26oZ-8ImbZf" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n CA constraints: 76\n", + "n mainchain constraints: 228\n" + ] + } + ], + "source": [ + "# - Add coordinate constraints to the current coordinates\n", + "\n", + "# TEMP: the constraint utilities are not yet part of the top-level tmol API\n", + "from tmol.score.constraint import constrain_all_ca\n", + "\n", + "# Option 1: constrain every CA to where it currently sits\n", + "pose_1ubq_w_ca_coord_csts1 = constrain_all_ca(pose_1ubq)\n", + "\n", + "# Option 2: constrain the declared main-chain atoms of every residue\n", + "pose_1ubq_w_ca_coord_csts2 = tmol.create_mainchain_coordinate_constraints(pose_1ubq)\n", + "\n", + "print(\n", + " \"n CA constraints: \",\n", + " len(pose_1ubq_w_ca_coord_csts1.constraint_set.constraint_function_inds),\n", + ")\n", + "print(\n", + " \"n mainchain constraints:\",\n", + " len(pose_1ubq_w_ca_coord_csts2.constraint_set.constraint_function_inds),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 55, + "metadata": { + "id": "0tOaC2Trmbvc" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "score_std tensor([235.4992], device='cuda:0') score_alt tensor([235.5190], device='cuda:0')\n" + ] + } + ], + "source": [ + "# - Alter the parameters for the cart-bonded energy function & rescore\n", + "\n", + "# The same immutable + shallow-copy + \"evolve\" strategy is how scoring\n", + "# parameters are controlled. In this case, the CartBondedDatabase\n", + "# keeps a hash of all of its parameters in order to ensure it uses\n", + "# the appropriate set of tensors, and so it provides its own\n", + "# evolve-like method for construction from a dictionary of parameters:\n", + "# CartBondedDatabase.from_cartres_dict.\n", + "#\n", + "# Let's imagine that we want to adjust the strength of only the\n", + "# bond distance between PRO CD and backbone N to make it stronger.\n", + "import copy\n", + "\n", + "# TEMP: CartBondedDatabase is not yet part of the top-level tmol API\n", + "from tmol.database.scoring import CartBondedDatabase\n", + "\n", + "cart_db = default_param_db.scoring.cartbonded\n", + "\n", + "# replace the default spring constant on the peptide bond with a stronger one\n", + "pro_params = cart_db.residue_params[\"PRO\"]\n", + "pro_length_params = pro_params.length_parameters\n", + "ind, cd_n_bond_params = next(\n", + " (i, p)\n", + " for (i, p) in enumerate(pro_length_params)\n", + " if p.atm1 == \"N\" and p.atm2 == \"CD\"\n", + ")\n", + "# double the default strength of ~122 kcal/mol*A\n", + "alt_cd_n_bond_params = attr.evolve(cd_n_bond_params, K=244)\n", + "alt_length_params = (\n", + " pro_length_params[0:ind]\n", + " + (alt_cd_n_bond_params,)\n", + " + pro_length_params[ind + 1 :]\n", + ")\n", + "alt_pro_params = attr.evolve(pro_params, length_parameters=alt_length_params)\n", + "alt_residue_params = copy.deepcopy(cart_db.residue_params)\n", + "alt_residue_params[\"PRO\"] = alt_pro_params\n", + "\n", + "alt_cart_db = CartBondedDatabase.from_cartres_dict(cartres_dict=alt_residue_params)\n", + "\n", + "alt_score_db = attr.evolve(default_param_db.scoring, cartbonded=alt_cart_db)\n", + "alt_param_db = attr.evolve(default_param_db, scoring=alt_score_db)\n", + "\n", + "sfxn_alt = tmol.beta2016_score_function(device, param_db=alt_param_db)\n", + "\n", + "wpsm_std = sfxn.render_whole_pose_scoring_module(pose_1ubq)\n", + "wpsm_alt = sfxn_alt.render_whole_pose_scoring_module(pose_1ubq)\n", + "score_std = wpsm_std(pose_1ubq.coords)\n", + "score_alt = wpsm_alt(pose_1ubq.coords)\n", + "print(\"score_std\", score_std, \"score_alt\", score_alt)\n", + "assert score_std < score_alt" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "dIPA0OfBmjhh" + }, + "source": [ + "Optimization" + ] + }, + { + "cell_type": "code", + "execution_count": 56, + "metadata": { + "id": "GffUmGNCmlIk" + }, + "outputs": [], + "source": [ + "# - Create a DunbrackSampler\n", + "\n", + "# constructing this object is a relatively heavy weight operation, so if you\n", + "# will be invoking the packer repeatedly, it's worthwhile to construct it\n", + "# once and then hold on to it\n", + "\n", + "# TEMP: the conformer samplers are not yet part of the top-level tmol API\n", + "from tmol.pack.rotamer.dunbrack import create_dunbrack_sampler_from_database\n", + "\n", + "dun_sampler = create_dunbrack_sampler_from_database(default_param_db, device)" + ] + }, + { + "cell_type": "code", + "execution_count": 57, + "metadata": {}, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "pack_rotamers 0.39 build rots: 0.09 calcRPEs: 0.10 build IG: 0.00 build IG part2: 0.00 run SA: 0.20 pose ctor: 0.00\n" + ] + } + ], + "source": [ + "# - Perform fixed-sequence side-chain optimization: repack\n", + "\n", + "# we will invoke \"the packer\": the module in tmol that optimizes the\n", + "# discrete side-chain conformation assignment.\n", + "# Steps:\n", + "# 1. Create a PackerPalette\n", + "# 2. Create a PackerTask using the palette\n", + "# 3. Configure the PackerTask to say \"optimize the sequence, but don't change it\"\n", + "# 4. Invoke pack_rotamers\n", + "\n", + "# TEMP: the packer is not yet part of the top-level tmol API\n", + "from tmol.pack import PackerPalette, PackerTask, pack_rotamers\n", + "from tmol.pack.rotamer import FixedAAChiSampler, IncludeCurrentSampler\n", + "\n", + "# 1.\n", + "# The default PackerPalette controls the initialization of the PackerTask\n", + "# to allow design from LCAAs to other LCAAs, DCAAs to other DCAAs, and\n", + "# otherwise only allows the original block types. The PackerPalette class is\n", + "# meant to be subclassed, so if you have more complex logic for which\n", + "# block types to consider at any given position, then you should write\n", + "# your own PackerPalette subclass.\n", + "palette = PackerPalette()\n", + "\n", + "# 2.\n", + "task = PackerTask(pose_1ubq, palette)\n", + "\n", + "# 3.\n", + "task.restrict_to_repacking()\n", + "task.add_conformer_sampler(dun_sampler)\n", + "task.add_conformer_sampler(FixedAAChiSampler())\n", + "# Beware: the native rotamer is often better than the naive rotamers\n", + "# and so if you add the IncludeCurrentSampler in e.g. a sequence\n", + "# recovery benchmark, you will be biasing the energies of the native\n", + "# sequence. For most modeling problems, though, you really do want\n", + "# to keep the input rotamer.\n", + "task.add_conformer_sampler(IncludeCurrentSampler())\n", + "\n", + "# 4.\n", + "pose_1ubq_repacked = pack_rotamers(pose_1ubq, sfxn, task, verbose=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 58, + "metadata": {}, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "score tensor([-100.0583], device='cuda:0')\n" + ] + } + ], + "source": [ + "wpsm = sfxn.render_whole_pose_scoring_module(pose_1ubq_repacked)\n", + "print(\"score\", wpsm(pose_1ubq_repacked.coords))" + ] + }, + { + "cell_type": "code", + "execution_count": 59, + "metadata": { + "id": "nDVbq3pDmqpf" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "pack_rotamers 0.36 build rots: 0.34 calcRPEs: 0.01 build IG: 0.00 build IG part2: 0.00 run SA: 0.01 pose ctor: 0.00\n", + "score after optH: tensor([189.3225], device='cuda:0')\n" + ] + } + ], + "source": [ + "# - Add hydrogens\n", + "\n", + "# Most NNs model only the heavy atoms; tmol's energy function requires Hs.\n", + "# The basic pathway of creating a PoseStack will place aliphatic hydrogens\n", + "# in their ideal geometries, however, the hydroxyl hydrogens will not be\n", + "# optimized. To optimize Hs, you can use the OptHSampler when invoking the\n", + "# packer; only the hydroxyl/proton-chi positions are repacked.\n", + "\n", + "# TEMP: the conformer samplers are not yet part of the top-level tmol API\n", + "from tmol.pack.rotamer import OptHSampler\n", + "\n", + "optH_task = PackerTask(pose_1ubq, palette)\n", + "optH_task.restrict_to_repacking()\n", + "optH_task.add_conformer_sampler(IncludeCurrentSampler())\n", + "optH_task.add_conformer_sampler(OptHSampler())\n", + "\n", + "pose_1ubq_optH_repacked = pack_rotamers(pose_1ubq, sfxn, optH_task, verbose=True)\n", + "\n", + "wpsm_optH = sfxn.render_whole_pose_scoring_module(pose_1ubq_optH_repacked)\n", + "print(\"score after optH:\", wpsm_optH(pose_1ubq_optH_repacked.coords))" + ] + }, + { + "cell_type": "code", + "execution_count": 60, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 61, + "metadata": { + "id": "e3a4AOs8mrDX" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "n atoms after stripping side chains: 374\n", + "n real atoms in the rebuilt PoseStack: 1231\n", + "n real atoms in the original PoseStack: 1231\n", + "score of the rebuilt structure: tensor([-92.7290], device='cuda:0')\n", + "build_missing_sidechains is available: True\n" + ] + } + ], + "source": [ + "# - Fill in side chains for a model that lacks them\n", + "#\n", + "# Sometimes your models will only include the backbone, or you are reading\n", + "# from the PDB where the electron density was absent for a side chain and the\n", + "# crystallographer did not model its conformation.\n", + "#\n", + "# The pose_stack_from_biotite pathway automatically builds in missing side\n", + "# chains (and hydrogens). To show that, strip ubiquitin back to its backbone\n", + "# + CB atoms with biotite and then rebuild.\n", + "backbone_only = bt_pdb_file.get_structure()\n", + "if isinstance(backbone_only, biotite.structure.AtomArrayStack):\n", + " backbone_only = backbone_only[0]\n", + "keep = numpy.isin(backbone_only.atom_name, [\"N\", \"CA\", \"C\", \"O\", \"CB\"])\n", + "backbone_only = backbone_only[keep]\n", + "print(\"n atoms after stripping side chains:\", backbone_only.array_length())\n", + "\n", + "pose_rebuilt = pose_stack_from_biotite(backbone_only, device)\n", + "print(\"n real atoms in the rebuilt PoseStack:\", int(pose_rebuilt.real_atoms.sum()))\n", + "print(\"n real atoms in the original PoseStack:\", int(pose_1ubq.real_atoms.sum()))\n", + "\n", + "# the rebuilt structure is complete and scorable\n", + "wpsm_rebuilt = sfxn.render_whole_pose_scoring_module(pose_rebuilt)\n", + "print(\"score of the rebuilt structure:\", wpsm_rebuilt(pose_rebuilt.coords))\n", + "\n", + "# build_missing_sidechains is the lower-level entry point that the IO\n", + "# pathway uses; it takes an explicit mask of which blocks need rebuilding.\n", + "from tmol.pack import build_missing_sidechains # TEMP: not yet top-level API\n", + "\n", + "print(\"build_missing_sidechains is available:\", callable(build_missing_sidechains))" + ] + }, + { + "cell_type": "code", + "execution_count": 62, + "metadata": { + "id": "Tgj8DwX8mrZ6" + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 63, + "metadata": { + "id": "zt23WbYtmru3" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "pack_rotamers 0.49 build rots: 0.02 calcRPEs: 0.32 build IG: 0.01 build IG part2: 0.00 run SA: 0.14 pose ctor: 0.00\n" + ] + } + ], + "source": [ + "# - Repack with extra rotamers\n", + "\n", + "# the \"extra chi\" flags in the PackerTask will trigger\n", + "# sampling at +/- 1 stdev. Chi indices are zero-based, so 0 and 1\n", + "# correspond conceptually to Rosetta's ex1 and ex2.\n", + "import gc\n", + "\n", + "gc.collect()\n", + "torch.cuda.empty_cache() if device.type == \"cuda\" else None\n", + "\n", + "palette = PackerPalette()\n", + "task = PackerTask(pose_1ubq, palette)\n", + "\n", + "task.restrict_to_repacking()\n", + "\n", + "# Expanding chi1 and chi2 at every one of the 76 positions multiplies the\n", + "# rotamer set enormously, so restrict the packing region to keep this example\n", + "# small enough to run anywhere.\n", + "ex_region = torch.zeros(\n", + " (pose_1ubq.n_poses, pose_1ubq.max_n_blocks), dtype=torch.bool, device=device\n", + ")\n", + "ex_region[:, 20:32] = True\n", + "task.disable_packing_by_block_mask(~ex_region)\n", + "\n", + "task.add_conformer_sampler(dun_sampler)\n", + "task.add_conformer_sampler(FixedAAChiSampler())\n", + "task.add_conformer_sampler(IncludeCurrentSampler())\n", + "task.or_expand_chi(0)\n", + "task.or_expand_chi(1)\n", + "\n", + "pose_1ubq_repacked_ex1ex2 = pack_rotamers(pose_1ubq, sfxn, task, verbose=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 64, + "metadata": {}, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "score tensor([211.5636], device='cuda:0')\n" + ] + } + ], + "source": [ + "wpsm = sfxn.render_whole_pose_scoring_module(pose_1ubq_repacked_ex1ex2)\n", + "print(\"score\", wpsm(pose_1ubq_repacked_ex1ex2.coords))" + ] + }, + { + "cell_type": "code", + "execution_count": 65, + "metadata": { + "id": "YmAY91mBmsCy" + }, + "outputs": [], + "source": [ + "# - Add hydrogens and back-propagate\n", + "\n", + "# The addition of missing alaphatic hydrogen atoms is differentiable\n", + "# and so it is possible to start from a model output from a NN\n", + "# which does not have hydrogens, build those hydrogens, score the\n", + "# structure, and then have the forces applied on those hydrogens\n", + "# backprop to the heavy-atoms they came from. All of that comes\n", + "# for free using the pose_stack_from_openfold utility.\n", + "\n", + "openfold_ubq_and_sumo_pred2 = torch.load(\"openfold_ubq_and_sumo.pt\", map_location=device)\n", + "# pretend an NN just gave us this\n", + "openfold_ubq_and_sumo_pred2[\"positions\"].requires_grad_()\n", + "ps_ubq_sumo2 = tmol.pose_stack_from_openfold(openfold_ubq_and_sumo_pred2)\n", + "wpsm = sfxn.render_whole_pose_scoring_module(ps_ubq_sumo2)\n", + "score = wpsm(ps_ubq_sumo2.coords)\n", + "score.sum().backward()" + ] + }, + { + "cell_type": "code", + "execution_count": 66, + "metadata": { + "id": "awPe0TtYmsaR" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "pack_rotamers 0.13 build rots: 0.02 calcRPEs: 0.05 build IG: 0.00 build IG part2: 0.00 run SA: 0.05 pose ctor: 0.00\n", + "n positions whose identity changed: 5\n", + "designed sequence (first 30): ['MET', 'GLN', 'ILE', 'PHE', 'VAL', 'LYS', 'THR', 'LEU', 'THR', 'GLY', 'LYS', 'THR', 'ILE', 'THR', 'LEU', 'GLU', 'VAL', 'GLU', 'PRO', 'SER', 'ASP', 'THR', 'ILE', 'SER', 'SER', 'VAL', 'LYS', 'SER', 'LYS', 'ILE']\n" + ] + } + ], + "source": [ + "# - Create a new PackerPalette subclass to handle logic of new block types\n", + "#\n", + "# PackerPalette decides which block types the packer may consider at each\n", + "# position. Subclass it when you have block types (or rules) the default\n", + "# palette does not know about. The contract of block_types_from_original is\n", + "# to return, for each pose/block: the number of allowed choices, a padded\n", + "# tensor of the allowed block-type indices (-1 padding), and a mask marking\n", + "# which of those choices is the original block type.\n", + "\n", + "from tmol.pack import PackerPalette # TEMP: not yet top-level API\n", + "\n", + "\n", + "class RestrictedPalette(PackerPalette):\n", + " \"\"\"Only allow design to a whitelist of residue names (plus the original).\"\"\"\n", + "\n", + " def __init__(self, allowed_name3s=(\"ALA\", \"GLY\", \"SER\")):\n", + " super().__init__()\n", + " self.allowed_name3s = set(allowed_name3s)\n", + "\n", + " def block_types_from_original(self, pbt, orig):\n", + " n_allowed, allowed, is_orig = super().block_types_from_original(pbt, orig)\n", + "\n", + " in_whitelist = torch.tensor(\n", + " [bt.name3 in self.allowed_name3s for bt in pbt.active_block_types],\n", + " dtype=torch.bool,\n", + " device=pbt.device,\n", + " )\n", + " is_real_choice = allowed >= 0\n", + " survives = torch.zeros_like(allowed, dtype=torch.bool)\n", + " survives[is_real_choice] = in_whitelist[allowed[is_real_choice]]\n", + " # always keep the original block type so every position has a choice\n", + " survives = (survives | is_orig) & is_real_choice\n", + "\n", + " # compact the survivors to the front of the padded dimension\n", + " order = torch.argsort((~survives).to(torch.int64), dim=2, stable=True)\n", + " allowed = torch.gather(allowed, 2, order)\n", + " is_orig = torch.gather(is_orig, 2, order)\n", + " survives = torch.gather(survives, 2, order)\n", + "\n", + " allowed = torch.where(survives, allowed, torch.full_like(allowed, -1))\n", + " is_orig = is_orig & survives\n", + " n_allowed = survives.sum(dim=2)\n", + " return n_allowed, allowed, is_orig\n", + "\n", + "\n", + "gc.collect()\n", + "torch.cuda.empty_cache() if device.type == \"cuda\" else None\n", + "\n", + "restricted_palette = RestrictedPalette((\"ALA\", \"GLY\", \"SER\"))\n", + "design_task = PackerTask(pose_1ubq, restricted_palette)\n", + "\n", + "# keep the design region small so the example is cheap\n", + "design_region = torch.zeros(\n", + " (pose_1ubq.n_poses, pose_1ubq.max_n_blocks), dtype=torch.bool, device=device\n", + ")\n", + "design_region[:, 20:32] = True\n", + "design_task.disable_packing_by_block_mask(~design_region)\n", + "\n", + "design_task.add_conformer_sampler(dun_sampler)\n", + "design_task.add_conformer_sampler(FixedAAChiSampler())\n", + "design_task.add_conformer_sampler(IncludeCurrentSampler())\n", + "\n", + "pose_1ubq_designed = pack_rotamers(pose_1ubq, sfxn, design_task, verbose=True)\n", + "\n", + "designed_names = [\n", + " pose_1ubq_designed.block_type(0, i).name3\n", + " for i in range(pose_1ubq_designed.max_n_blocks)\n", + "]\n", + "orig_names = [\n", + " pose_1ubq.block_type(0, i).name3 for i in range(pose_1ubq.max_n_blocks)\n", + "]\n", + "n_changed = sum(a != b for a, b in zip(orig_names, designed_names))\n", + "print(\"n positions whose identity changed:\", n_changed)\n", + "print(\"designed sequence (first 30):\", designed_names[:30])\n", + "\n", + "# every designed identity is either the original or in the whitelist\n", + "assert all(\n", + " new in (\"ALA\", \"GLY\", \"SER\") or new == old\n", + " for old, new in zip(orig_names, designed_names)\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 67, + "metadata": { + "id": "RTm3LjhAmsv4" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "score before: tensor([235.4993], device='cuda:0')\n", + "score after (float64 DOFs): tensor([-13.5674], device='cuda:0', grad_fn=)\n" + ] + } + ], + "source": [ + "# - Run minimization in double precision\n", + "#\n", + "# run_kin_min carries the kinematic degrees of freedom in the dtype given by\n", + "# kin_dtype. float64 costs more but gives a tighter-converging descent.\n", + "gc.collect()\n", + "torch.cuda.empty_cache() if device.type == \"cuda\" else None\n", + "\n", + "ff_dp = tmol.FoldForest.reasonable_fold_forest(pose_1ubq)\n", + "mm_dp = tmol.MoveMap.from_pose_stack(pose_1ubq)\n", + "mm_dp.move_all_named_torsions = True\n", + "\n", + "wpsm_dp = sfxn.render_whole_pose_scoring_module(pose_1ubq)\n", + "print(\"score before:\", wpsm_dp(pose_1ubq.coords))\n", + "\n", + "min_f64 = tmol.run_kin_min(\n", + " pose_1ubq,\n", + " sfxn,\n", + " ff_dp,\n", + " mm_dp,\n", + " kin_dtype=torch.float64,\n", + " optimizer_kwargs={\"max_iter\": 50},\n", + ")\n", + "print(\"score after (float64 DOFs):\", wpsm_dp(min_f64.coords))" + ] + }, + { + "cell_type": "code", + "execution_count": 68, + "metadata": { + "id": "Z2UFgON6mtFt" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "score before: tensor([235.4992], device='cuda:0')\n", + "score after: tensor([-182.1358], device='cuda:0')\n" + ] + } + ], + "source": [ + "# - Perform cartesian minimization\n", + "\n", + "gc.collect()\n", + "torch.cuda.empty_cache() if device.type == \"cuda\" else None\n", + "\n", + "cart_minimized_pose_1ubq = tmol.run_cart_min(\n", + " pose_1ubq, sfxn, optimizer_kwargs={\"max_iter\": 50}\n", + ")\n", + "\n", + "wpsm_cart = sfxn.render_whole_pose_scoring_module(cart_minimized_pose_1ubq)\n", + "print(\"score before:\", wpsm_cart(pose_1ubq.coords))\n", + "print(\"score after: \", wpsm_cart(cart_minimized_pose_1ubq.coords))" + ] + }, + { + "cell_type": "code", + "execution_count": 69, + "metadata": { + "id": "_Dwt2_lpmtbM" + }, + "outputs": [], + "source": [ + "# - Perform kinematic minimization\n", + "\n", + "# The run_kin_min function will perform minimization in\n", + "# internal degrees of freedom.\n", + "\n", + "# from the kinematics section above\n", + "ff1_1ubq = tmol.FoldForest.reasonable_fold_forest(pose_1ubq)\n", + "mm_1ubq_all = tmol.MoveMap.from_pose_stack(pose_1ubq)\n", + "mm_1ubq_all.move_all_named_torsions = True\n", + "\n", + "wpsm = sfxn.render_whole_pose_scoring_module(pose_1ubq)\n", + "score_before = wpsm(pose_1ubq.coords)\n", + "\n", + "minimized_pose_1ubq = tmol.run_kin_min(pose_1ubq, sfxn, ff1_1ubq, mm_1ubq_all)\n", + "\n", + "score_after = wpsm(minimized_pose_1ubq.coords)\n", + "assert score_after < score_before" + ] + }, + { + "cell_type": "code", + "execution_count": 70, + "metadata": { + "id": "oE7_fHfEmtwM" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "packing with fa_rep of 0.06 and constraint weight of 0.00\n", + "pack_rotamers 0.31 build rots: 0.02 calcRPEs: 0.10 build IG: 0.00 build IG part2: 0.00 run SA: 0.18 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.06 and constraint weight of 0.00\n", + "build_kinforest_network 0.00 s1: 0.00 s2: 0.00 s3: 0.00\n", + "run_min 0.32 setup: 0.00 opt 0.32 stack-ctor: 0.00\n", + "pack-min 0.63 task-init 0.00 packing 0.31 min 0.32\n", + "packing with fa_rep of 0.55 and constraint weight of 0.00\n", + "pack_rotamers 0.32 build rots: 0.02 calcRPEs: 0.10 build IG: 0.00 build IG part2: 0.00 run SA: 0.19 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.55 and constraint weight of 0.00\n", + "build_kinforest_network 0.01 s1: 0.00 s2: 0.00 s3: 0.00\n", + "run_min 0.32 setup: 0.00 opt 0.32 stack-ctor: 0.00\n", + "pack-min 0.64 task-init 0.00 packing 0.32 min 0.32\n", + "accepting new best scores\n", + " old best score: tensor([235.4993], device='cuda:0')\n", + " new best score: tensor([-151.4261], device='cuda:0', grad_fn=)\n", + "score after kinematic relax: tensor([-151.4264], device='cuda:0')\n" + ] + } + ], + "source": [ + "# - Relax a PoseStack w/ kinematic minimization\n", + "#\n", + "# fast_relax is the single relax protocol; whether it relaxes in internal or\n", + "# Cartesian DOFs is determined by the move map you hand it together with the\n", + "# matching min_fn. For kinematic relax, pass a MoveMap plus a min_fn built\n", + "# around run_kin_min.\n", + "\n", + "gc.collect()\n", + "torch.cuda.empty_cache() if device.type == \"cuda\" else None\n", + "\n", + "packer_palette = PackerPalette()\n", + "\n", + "# from the kinematics section above\n", + "ff1_1ubq = tmol.FoldForest.reasonable_fold_forest(pose_1ubq)\n", + "mm_1ubq_all = tmol.MoveMap.from_pose_stack(pose_1ubq)\n", + "mm_1ubq_all.move_all_named_torsions = True\n", + "\n", + "\n", + "def kinematic_min_fn(pose_stack, score_function, *, fold_forest, move_map, verbose):\n", + " return tmol.run_kin_min(\n", + " pose_stack,\n", + " score_function,\n", + " fold_forest,\n", + " move_map,\n", + " optimizer_kwargs={\"max_iter\": 50},\n", + " verbose=verbose,\n", + " )\n", + "\n", + "\n", + "# a short schedule / single repeat keeps the notebook quick; drop these\n", + "# arguments to get the full MonomerRelax2019 ramp.\n", + "short_schedule = [0.1, 1.0]\n", + "\n", + "kin_relaxed_1ubq = tmol.fast_relax(\n", + " pose_1ubq,\n", + " sfxn,\n", + " packer_palette,\n", + " mm_1ubq_all,\n", + " ff1_1ubq,\n", + " num_repeats=1,\n", + " schedule=short_schedule,\n", + " min_fn=kinematic_min_fn,\n", + " verbose=True,\n", + ")\n", + "\n", + "wpsm_kin = sfxn.render_whole_pose_scoring_module(kin_relaxed_1ubq)\n", + "print(\"score after kinematic relax:\", wpsm_kin(kin_relaxed_1ubq.coords))" + ] + }, + { + "cell_type": "code", + "execution_count": 71, + "metadata": { + "id": "9ZUjD_eYmuFw" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "packing with fa_rep of 0.06 and constraint weight of 0.00\n", + "pack_rotamers 0.31 build rots: 0.02 calcRPEs: 0.10 build IG: 0.00 build IG part2: 0.00 run SA: 0.19 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.06 and constraint weight of 0.00\n", + " iter 1 E=-203.217560 evals=6 done=0/1\n", + " iter 2 E=-216.729767 evals=1 done=0/1\n", + " iter 3 E=-224.710068 evals=1 done=0/1\n", + " iter 4 E=-231.657013 evals=2 done=0/1\n", + " iter 5 E=-237.187042 evals=1 done=0/1\n", + " iter 6 E=-241.635712 evals=2 done=0/1\n", + " iter 7 E=-243.062408 evals=1 done=0/1\n", + " iter 8 E=-245.966278 evals=2 done=0/1\n", + " iter 9 E=-248.399246 evals=1 done=0/1\n", + " iter 10 E=-250.893509 evals=2 done=0/1\n", + " iter 11 E=-252.440857 evals=1 done=0/1\n", + " iter 12 E=-255.059616 evals=2 done=0/1\n", + " iter 13 E=-256.729401 evals=1 done=0/1\n", + " iter 14 E=-258.787598 evals=2 done=0/1\n", + " iter 15 E=-260.711914 evals=1 done=0/1\n", + " iter 16 E=-262.557404 evals=2 done=0/1\n", + " iter 17 E=-264.623505 evals=1 done=0/1\n", + " iter 18 E=-266.267334 evals=2 done=0/1\n", + " iter 19 E=-268.098694 evals=1 done=0/1\n", + " iter 20 E=-269.639099 evals=2 done=0/1\n", + " iter 21 E=-271.829529 evals=1 done=0/1\n", + " iter 22 E=-273.494293 evals=2 done=0/1\n", + " iter 23 E=-275.363861 evals=1 done=0/1\n", + " iter 24 E=-277.106384 evals=2 done=0/1\n", + " iter 25 E=-279.277802 evals=1 done=0/1\n", + " iter 26 E=-281.111328 evals=2 done=0/1\n", + " iter 27 E=-283.141541 evals=1 done=0/1\n", + " iter 28 E=-284.699097 evals=2 done=0/1\n", + " iter 29 E=-286.568237 evals=1 done=0/1\n", + " iter 30 E=-288.128967 evals=2 done=0/1\n", + " iter 31 E=-290.021637 evals=1 done=0/1\n", + " iter 32 E=-291.425720 evals=2 done=0/1\n", + " iter 33 E=-293.087982 evals=1 done=0/1\n", + " iter 34 E=-294.431183 evals=2 done=0/1\n", + " iter 35 E=-296.057098 evals=1 done=0/1\n", + " iter 36 E=-297.286011 evals=2 done=0/1\n", + " iter 37 E=-298.724976 evals=1 done=0/1\n", + " iter 38 E=-299.779510 evals=2 done=0/1\n", + " iter 39 E=-301.027588 evals=1 done=0/1\n", + " iter 40 E=-301.988892 evals=2 done=0/1\n", + " iter 41 E=-303.251892 evals=1 done=0/1\n", + " iter 42 E=-304.208282 evals=2 done=0/1\n", + " iter 43 E=-305.287689 evals=1 done=0/1\n", + " iter 44 E=-306.065765 evals=2 done=0/1\n", + " iter 45 E=-306.941986 evals=1 done=0/1\n", + " iter 46 E=-307.616455 evals=2 done=0/1\n", + " iter 47 E=-308.370850 evals=1 done=0/1\n", + " iter 48 E=-308.952820 evals=2 done=0/1\n", + " iter 49 E=-309.641296 evals=1 done=0/1\n", + " iter 50 E=-310.195007 evals=2 done=0/1\n", + " iter 51 E=-310.925049 evals=1 done=0/1\n", + " iter 52 E=-311.556152 evals=2 done=0/1\n", + " iter 53 E=-312.320862 evals=1 done=0/1\n", + " iter 54 E=-312.934204 evals=2 done=0/1\n", + " iter 55 E=-313.712036 evals=1 done=0/1\n", + " iter 56 E=-314.329437 evals=2 done=0/1\n", + " iter 57 E=-315.116882 evals=1 done=0/1\n", + " iter 58 E=-315.766785 evals=2 done=0/1\n", + " iter 59 E=-316.460388 evals=1 done=0/1\n", + " iter 60 E=-317.026672 evals=2 done=0/1\n", + " iter 61 E=-317.716492 evals=1 done=0/1\n", + " iter 62 E=-318.303680 evals=2 done=0/1\n", + " iter 63 E=-318.974365 evals=1 done=0/1\n", + " iter 64 E=-319.478943 evals=2 done=0/1\n", + " iter 65 E=-320.085052 evals=1 done=0/1\n", + " iter 66 E=-320.585266 evals=2 done=0/1\n", + " iter 67 E=-321.208191 evals=1 done=0/1\n", + " iter 68 E=-321.708862 evals=2 done=0/1\n", + " iter 69 E=-322.346252 evals=1 done=0/1\n", + " iter 70 E=-322.848816 evals=2 done=0/1\n", + " iter 71 E=-323.487610 evals=1 done=0/1\n", + " iter 72 E=-323.994965 evals=2 done=0/1\n", + " iter 73 E=-324.591431 evals=1 done=0/1\n", + " iter 74 E=-325.033539 evals=2 done=0/1\n", + " iter 75 E=-325.600006 evals=1 done=0/1\n", + " iter 76 E=-326.051514 evals=2 done=0/1\n", + " iter 77 E=-326.597900 evals=1 done=0/1\n", + " iter 78 E=-327.035095 evals=2 done=0/1\n", + " iter 79 E=-327.535370 evals=1 done=0/1\n", + " iter 80 E=-327.937347 evals=2 done=0/1\n", + " iter 81 E=-328.430298 evals=1 done=0/1\n", + " iter 82 E=-328.803345 evals=2 done=0/1\n", + " iter 83 E=-329.266113 evals=1 done=0/1\n", + " iter 84 E=-329.618103 evals=2 done=0/1\n", + " iter 85 E=-330.042908 evals=1 done=0/1\n", + " iter 86 E=-330.355164 evals=2 done=0/1\n", + " iter 87 E=-330.737732 evals=1 done=0/1\n", + " iter 88 E=-331.044189 evals=2 done=0/1\n", + " iter 89 E=-331.433411 evals=1 done=0/1\n", + " iter 90 E=-331.745087 evals=2 done=0/1\n", + " iter 91 E=-332.142883 evals=1 done=0/1\n", + " iter 92 E=-332.454315 evals=2 done=0/1\n", + " iter 93 E=-332.856720 evals=1 done=0/1\n", + " iter 94 E=-333.174072 evals=2 done=0/1\n", + " iter 95 E=-333.603851 evals=1 done=0/1\n", + " iter 96 E=-333.925140 evals=2 done=0/1\n", + " iter 97 E=-334.301147 evals=1 done=0/1\n", + " iter 98 E=-334.584351 evals=2 done=0/1\n", + " iter 99 E=-334.914612 evals=1 done=0/1\n", + " iter 100 E=-335.171753 evals=2 done=0/1\n", + " iter 101 E=-335.502655 evals=1 done=0/1\n", + " iter 102 E=-335.773285 evals=2 done=0/1\n", + " iter 103 E=-336.107330 evals=1 done=0/1\n", + " iter 104 E=-336.371155 evals=2 done=0/1\n", + " iter 105 E=-336.723145 evals=1 done=0/1\n", + " iter 106 E=-336.977142 evals=2 done=0/1\n", + " iter 107 E=-337.289246 evals=1 done=0/1\n", + " iter 108 E=-337.512421 evals=2 done=0/1\n", + " iter 109 E=-337.769043 evals=1 done=0/1\n", + " iter 110 E=-337.960663 evals=2 done=0/1\n", + " iter 111 E=-338.240112 evals=1 done=0/1\n", + " iter 112 E=-338.446350 evals=2 done=0/1\n", + " iter 113 E=-338.714996 evals=1 done=0/1\n", + " iter 114 E=-338.908478 evals=2 done=0/1\n", + " iter 115 E=-339.160736 evals=1 done=0/1\n", + " iter 116 E=-339.361877 evals=2 done=0/1\n", + " iter 117 E=-339.611176 evals=1 done=0/1\n", + " iter 118 E=-339.817017 evals=2 done=0/1\n", + " iter 119 E=-340.075897 evals=1 done=0/1\n", + " iter 120 E=-340.261780 evals=2 done=0/1\n", + " iter 121 E=-340.495819 evals=1 done=0/1\n", + " iter 122 E=-340.665253 evals=2 done=0/1\n", + " iter 123 E=-340.895538 evals=1 done=0/1\n", + " iter 124 E=-341.072540 evals=2 done=0/1\n", + " iter 125 E=-341.285248 evals=1 done=0/1\n", + " iter 126 E=-341.432373 evals=2 done=0/1\n", + " iter 127 E=-341.640472 evals=1 done=0/1\n", + " iter 128 E=-341.811401 evals=2 done=0/1\n", + " iter 129 E=-342.033783 evals=1 done=0/1\n", + " iter 130 E=-342.205231 evals=2 done=0/1\n", + " iter 131 E=-342.445251 evals=1 done=0/1\n", + " iter 132 E=-342.624176 evals=2 done=0/1\n", + " iter 133 E=-342.734833 evals=1 done=1/1\n", + " finished: 1 converged, 0 stalled\n", + " LBFGS_Armijo done: 133 iters, 204 func evals, E=-342.7348\n", + "run_min 0.79 setup: 0.00 opt 0.78 stack-ctor: 0.00\n", + "pack-min 1.10 task-init 0.00 packing 0.31 min 0.79\n", + "packing with fa_rep of 0.55 and constraint weight of 0.00\n", + "pack_rotamers 0.32 build rots: 0.02 calcRPEs: 0.10 build IG: 0.00 build IG part2: 0.00 run SA: 0.19 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.55 and constraint weight of 0.00\n", + " iter 1 E=-169.317902 evals=5 done=0/1\n", + " iter 2 E=-186.513336 evals=2 done=0/1\n", + " iter 3 E=-193.982971 evals=1 done=0/1\n", + " iter 4 E=-199.485382 evals=2 done=0/1\n", + " iter 5 E=-202.543823 evals=1 done=0/1\n", + " iter 6 E=-205.944336 evals=2 done=0/1\n", + " iter 7 E=-210.085663 evals=1 done=0/1\n", + " iter 8 E=-213.694565 evals=2 done=0/1\n", + " iter 9 E=-217.090546 evals=1 done=0/1\n", + " iter 10 E=-219.411087 evals=2 done=0/1\n", + " iter 11 E=-222.017120 evals=1 done=0/1\n", + " iter 12 E=-224.229492 evals=2 done=0/1\n", + " iter 13 E=-226.560547 evals=1 done=0/1\n", + " iter 14 E=-228.325241 evals=2 done=0/1\n", + " iter 15 E=-230.155792 evals=1 done=0/1\n", + " iter 16 E=-231.912476 evals=2 done=0/1\n", + " iter 17 E=-233.826843 evals=1 done=0/1\n", + " iter 18 E=-235.395416 evals=2 done=0/1\n", + " iter 19 E=-237.198669 evals=1 done=0/1\n", + " iter 20 E=-238.633698 evals=2 done=0/1\n", + " iter 21 E=-240.455536 evals=1 done=0/1\n", + " iter 22 E=-241.857513 evals=2 done=0/1\n", + " iter 23 E=-243.238831 evals=1 done=0/1\n", + " iter 24 E=-244.424622 evals=2 done=0/1\n", + " iter 25 E=-245.646545 evals=1 done=0/1\n", + " iter 26 E=-246.727890 evals=2 done=0/1\n", + " iter 27 E=-247.886139 evals=1 done=0/1\n", + " iter 28 E=-248.730911 evals=2 done=0/1\n", + " iter 29 E=-249.865295 evals=1 done=0/1\n", + " iter 30 E=-250.747269 evals=2 done=0/1\n", + " iter 31 E=-251.612122 evals=1 done=0/1\n", + " iter 32 E=-252.344177 evals=2 done=0/1\n", + " iter 33 E=-253.088699 evals=1 done=0/1\n", + " iter 34 E=-253.768829 evals=2 done=0/1\n", + " iter 35 E=-254.626892 evals=1 done=0/1\n", + " iter 36 E=-255.301285 evals=2 done=0/1\n", + " iter 37 E=-256.018311 evals=1 done=0/1\n", + " iter 38 E=-256.595673 evals=2 done=0/1\n", + " iter 39 E=-257.257202 evals=1 done=0/1\n", + " iter 40 E=-257.802246 evals=2 done=0/1\n", + " iter 41 E=-258.446777 evals=1 done=0/1\n", + " iter 42 E=-258.908112 evals=2 done=0/1\n", + " iter 43 E=-259.448425 evals=1 done=0/1\n", + " iter 44 E=-259.881134 evals=2 done=0/1\n", + " iter 45 E=-260.391083 evals=1 done=0/1\n", + " iter 46 E=-260.764893 evals=2 done=0/1\n", + " iter 47 E=-261.208954 evals=1 done=0/1\n", + " iter 48 E=-261.548706 evals=2 done=0/1\n", + " iter 49 E=-261.942566 evals=1 done=0/1\n", + " iter 50 E=-262.264282 evals=2 done=0/1\n", + " iter 51 E=-262.632507 evals=1 done=0/1\n", + " iter 52 E=-262.908997 evals=2 done=0/1\n", + " iter 53 E=-263.248199 evals=1 done=0/1\n", + " iter 54 E=-263.519623 evals=2 done=0/1\n", + " iter 55 E=-263.851196 evals=1 done=0/1\n", + " iter 56 E=-264.088440 evals=2 done=0/1\n", + " iter 57 E=-264.356720 evals=1 done=0/1\n", + " iter 58 E=-264.583374 evals=2 done=0/1\n", + " iter 59 E=-264.891144 evals=1 done=0/1\n", + " iter 60 E=-265.120667 evals=2 done=0/1\n", + " iter 61 E=-265.360931 evals=1 done=0/1\n", + " iter 62 E=-265.558594 evals=2 done=0/1\n", + " iter 63 E=-265.791443 evals=1 done=0/1\n", + " iter 64 E=-265.995087 evals=2 done=0/1\n", + " iter 65 E=-266.225922 evals=1 done=0/1\n", + " iter 66 E=-266.401764 evals=2 done=0/1\n", + " iter 67 E=-266.594482 evals=1 done=0/1\n", + " iter 68 E=-266.751648 evals=2 done=0/1\n", + " iter 69 E=-266.959351 evals=1 done=0/1\n", + " iter 70 E=-267.130432 evals=2 done=0/1\n", + " iter 71 E=-267.319977 evals=1 done=0/1\n", + " iter 72 E=-267.474915 evals=2 done=0/1\n", + " iter 73 E=-267.650879 evals=1 done=0/1\n", + " iter 74 E=-267.790771 evals=2 done=0/1\n", + " iter 75 E=-267.959045 evals=1 done=0/1\n", + " iter 76 E=-268.084015 evals=2 done=0/1\n", + " iter 77 E=-268.239624 evals=1 done=0/1\n", + " iter 78 E=-268.360718 evals=2 done=0/1\n", + " iter 79 E=-268.501892 evals=1 done=0/1\n", + " iter 80 E=-268.618073 evals=2 done=0/1\n", + " iter 81 E=-268.765045 evals=1 done=0/1\n", + " iter 82 E=-268.885712 evals=2 done=0/1\n", + " iter 83 E=-269.036865 evals=1 done=0/1\n", + " iter 84 E=-269.152496 evals=2 done=0/1\n", + " iter 85 E=-269.284485 evals=1 done=0/1\n", + " iter 86 E=-269.388580 evals=2 done=0/1\n", + " iter 87 E=-269.519745 evals=1 done=0/1\n", + " iter 88 E=-269.613983 evals=2 done=0/1\n", + " iter 89 E=-269.713989 evals=1 done=0/1\n", + " iter 90 E=-269.791687 evals=2 done=1/1\n", + " finished: 1 converged, 0 stalled\n", + " LBFGS_Armijo done: 90 iters, 140 func evals, E=-269.7917\n", + "run_min 0.50 setup: 0.00 opt 0.50 stack-ctor: 0.00\n", + "pack-min 0.82 task-init 0.00 packing 0.32 min 0.50\n", + "accepting new best scores\n", + " old best score: tensor([235.4991], device='cuda:0')\n", + " new best score: tensor([-269.7917], device='cuda:0')\n", + "score after cartesian relax: tensor([-269.7915], device='cuda:0')\n" + ] + } + ], + "source": [ + "# - Relax a PoseStack w/ cartesian minimization\n", + "#\n", + "# The default min_fn is Cartesian, so passing a CartesianMoveMap is all it\n", + "# takes. The default CartesianMoveMap (coord_mask=None) lets every atom move.\n", + "gc.collect()\n", + "torch.cuda.empty_cache() if device.type == \"cuda\" else None\n", + "\n", + "packer_palette = PackerPalette()\n", + "cart_mm = tmol.CartesianMoveMap()\n", + "\n", + "cart_relaxed_1ubq = tmol.fast_relax(\n", + " pose_1ubq,\n", + " sfxn,\n", + " packer_palette,\n", + " cart_mm,\n", + " ff1_1ubq,\n", + " num_repeats=1,\n", + " schedule=short_schedule,\n", + " verbose=True,\n", + ")\n", + "\n", + "wpsm_cart_relax = sfxn.render_whole_pose_scoring_module(cart_relaxed_1ubq)\n", + "print(\"score after cartesian relax:\", wpsm_cart_relax(cart_relaxed_1ubq.coords))" + ] + }, + { + "cell_type": "code", + "execution_count": 72, + "metadata": { + "id": "2aUJd4X5mubN" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "batch of 3 poses; max_n_blocks = 92\n", + "scores before: tensor([ 235.4992, 1708.3586, 254.8842], device='cuda:0')\n", + "scores after: tensor([-269.6108, -116.8019, -318.5026], device='cuda:0')\n", + "wrote one PDB per relaxed pose\n" + ] + } + ], + "source": [ + "# - Relax structures generated one-at-a-time in batch format\n", + "#\n", + "# Though the two relax cells above were a) on single-pose PoseStacks, and\n", + "# b) on small structures, where the GPU really sings is with large structures\n", + "# and with many-pose PoseStacks. Structures that arrive one at a time can be\n", + "# collected with PoseStackBuilder.from_poses and then relaxed in a single\n", + "# batched call, which amortizes every kernel launch across the batch.\n", + "\n", + "gc.collect()\n", + "torch.cuda.empty_cache() if device.type == \"cuda\" else None\n", + "\n", + "incoming = [pose_1ubq, ps_3plc, ps_1qys] # e.g. arriving from a NN, one by one\n", + "batch = PoseStackBuilder.from_poses(incoming, device=device)\n", + "print(\"batch of\", batch.n_poses, \"poses; max_n_blocks =\", batch.max_n_blocks)\n", + "\n", + "batch_ff = tmol.FoldForest.reasonable_fold_forest(batch)\n", + "batch_mm = tmol.CartesianMoveMap()\n", + "\n", + "wpsm_batch = sfxn.render_whole_pose_scoring_module(batch)\n", + "print(\"scores before:\", wpsm_batch(batch.coords))\n", + "\n", + "relaxed_batch = tmol.fast_relax(\n", + " batch,\n", + " sfxn,\n", + " PackerPalette(),\n", + " batch_mm,\n", + " batch_ff,\n", + " num_repeats=1,\n", + " schedule=short_schedule,\n", + " verbose=False,\n", + ")\n", + "\n", + "print(\"scores after: \", wpsm_batch(relaxed_batch.coords))\n", + "\n", + "# split() recovers the individual poses again\n", + "for i in range(relaxed_batch.n_poses):\n", + " tmol.write_pose_stack_pdb(relaxed_batch.split(i), f\"relaxed_batch_{i}.pdb\")\n", + " assert os.path.isfile(f\"relaxed_batch_{i}.pdb\")\n", + "print(\"wrote one PDB per relaxed pose\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Mkj2U1Namuw7" + }, + "source": [ + "### Idealize a structure from the PDB\n", + "\n", + "**Not currently implemented in tmol.** There is no idealization protocol\n", + "(the Rosetta `idealize` equivalent) exposed by the current API \u2014 no\n", + "`idealize`-like function exists anywhere under `tmol/`. The closest available\n", + "building blocks are `tmol.run_kin_min` (which minimizes in internal DOFs and\n", + "therefore holds bond lengths and angles at whatever values the kinematic\n", + "description carries) and the `cart_lengths` / `cart_angles` score terms, which\n", + "penalize departures from ideal geometry during Cartesian minimization." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aWFR74xZmvFu" + }, + "source": [ + "### Idealize a structure from a dandelion\n", + "\n", + "**Not currently implemented in tmol.** See the note above: no idealization\n", + "protocol is exposed. A dandelion fold forest can be built (see the kinematics\n", + "section), but there is no routine that rebuilds its residues at ideal internal\n", + "geometry." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "NDZyRGtTmva3" + }, + "source": [ + "### Idealize just the backbone of a dandelion\n", + "\n", + "**Not currently implemented in tmol.** See the notes above." + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.11.13" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file