diff --git a/notebooks/tmol_how_to_guide.ipynb b/notebooks/tmol_how_to_guide.ipynb new file mode 100644 index 000000000..4e8b37369 --- /dev/null +++ b/notebooks/tmol_how_to_guide.ipynb @@ -0,0 +1,1892 @@ +{ + "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": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "nU1-4FoBbApM", + "outputId": "68e50ec9-25eb-450e-f1b2-882c310165c2", + "scrolled": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Collecting tmol==0.1.36+cu128torch2.10\n", + " Downloading https://github.com/uw-ipd/tmol/releases/download/v0.1.36/tmol-0.1.36+cu128torch2.10-cp312-cp312-linux_x86_64.whl (95.5 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m95.5/95.5 MB\u001b[0m \u001b[31m6.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hRequirement already satisfied: torch>=2.5 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (2.11.0+cu128)\n", + "Requirement already satisfied: numpy>=1.24 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (2.0.2)\n", + "Requirement already satisfied: scipy>=1.11 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (1.16.3)\n", + "Requirement already satisfied: attrs>=21.0 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (26.1.0)\n", + "Collecting attrs_strict<2,>=1.0 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading attrs_strict-1.0.1-py3-none-any.whl.metadata (6.3 kB)\n", + "Collecting cattrs>=22.0 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading cattrs-26.1.0-py3-none-any.whl.metadata (8.5 kB)\n", + "Requirement already satisfied: frozendict>=2.4 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (2.4.7)\n", + "Requirement already satisfied: pandas<3,>=2.0 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (2.2.2)\n", + "Requirement already satisfied: pyarrow>=12 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (18.1.0)\n", + "Requirement already satisfied: pyyaml>=6.0 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (6.0.3)\n", + "Requirement already satisfied: networkx>=3.0 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (3.6.1)\n", + "Requirement already satisfied: toolz<2,>=0.12.0 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (0.12.1)\n", + "Requirement already satisfied: decorator>=4.0 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (4.4.2)\n", + "Requirement already satisfied: typing_extensions<5,>=4.0 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (4.16.0)\n", + "Collecting typing_inspect>=0.9 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading typing_inspect-0.9.0-py3-none-any.whl.metadata (1.5 kB)\n", + "Collecting pint<1,>=0.23 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading pint-0.25.3-py3-none-any.whl.metadata (10 kB)\n", + "Requirement already satisfied: psutil>=5.9 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (5.9.5)\n", + "Requirement already satisfied: requests<3,>=2.28 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (2.32.4)\n", + "Collecting sparse>=0.13 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading sparse-0.19.0-py2.py3-none-any.whl.metadata (5.3 kB)\n", + "Requirement already satisfied: numba>=0.56 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (0.60.0)\n", + "Requirement already satisfied: llvmlite>=0.39 in /usr/local/lib/python3.12/dist-packages (from tmol==0.1.36+cu128torch2.10) (0.43.0)\n", + "Collecting hypothesis>=5.35 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading hypothesis-6.156.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (5.7 kB)\n", + "Collecting astor>=0.8 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading astor-0.8.1-py2.py3-none-any.whl.metadata (4.2 kB)\n", + "Collecting biotite>=1.4.0 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading biotite-1.7.1-cp312-cp312-manylinux_2_28_x86_64.whl.metadata (5.7 kB)\n", + "Collecting rdkit>=2024.3 (from tmol==0.1.36+cu128torch2.10)\n", + " Downloading rdkit-2026.3.4-cp312-cp312-manylinux_2_28_x86_64.whl.metadata (3.8 kB)\n", + "Collecting biotraj<2.0,>=1.0 (from biotite>=1.4.0->tmol==0.1.36+cu128torch2.10)\n", + " Downloading biotraj-1.2.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (32 kB)\n", + "Requirement already satisfied: msgpack>=0.5.6 in /usr/local/lib/python3.12/dist-packages (from biotite>=1.4.0->tmol==0.1.36+cu128torch2.10) (1.2.1)\n", + "Requirement already satisfied: packaging>=24.0 in /usr/local/lib/python3.12/dist-packages (from biotite>=1.4.0->tmol==0.1.36+cu128torch2.10) (26.2)\n", + "Requirement already satisfied: sortedcontainers<3.0.0,>=2.1.0 in /usr/local/lib/python3.12/dist-packages (from hypothesis>=5.35->tmol==0.1.36+cu128torch2.10) (2.4.0)\n", + "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas<3,>=2.0->tmol==0.1.36+cu128torch2.10) (2.9.0.post0)\n", + "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas<3,>=2.0->tmol==0.1.36+cu128torch2.10) (2025.2)\n", + "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas<3,>=2.0->tmol==0.1.36+cu128torch2.10) (2026.3)\n", + "Collecting flexcache>=0.3 (from pint<1,>=0.23->tmol==0.1.36+cu128torch2.10)\n", + " Downloading flexcache-0.3-py3-none-any.whl.metadata (7.0 kB)\n", + "Collecting flexparser>=0.4 (from pint<1,>=0.23->tmol==0.1.36+cu128torch2.10)\n", + " Downloading flexparser-0.4-py3-none-any.whl.metadata (18 kB)\n", + "Requirement already satisfied: platformdirs>=2.1.0 in /usr/local/lib/python3.12/dist-packages (from pint<1,>=0.23->tmol==0.1.36+cu128torch2.10) (4.10.0)\n", + "Requirement already satisfied: Pillow in /usr/local/lib/python3.12/dist-packages (from rdkit>=2024.3->tmol==0.1.36+cu128torch2.10) (11.3.0)\n", + "Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests<3,>=2.28->tmol==0.1.36+cu128torch2.10) (3.4.9)\n", + "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests<3,>=2.28->tmol==0.1.36+cu128torch2.10) (3.18)\n", + "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests<3,>=2.28->tmol==0.1.36+cu128torch2.10) (2.5.0)\n", + "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests<3,>=2.28->tmol==0.1.36+cu128torch2.10) (2026.6.17)\n", + "Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (3.29.7)\n", + "Requirement already satisfied: setuptools<82 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (75.2.0)\n", + "Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (1.14.0)\n", + "Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (3.1.6)\n", + "Requirement already satisfied: fsspec>=0.8.5 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (2025.3.0)\n", + "Requirement already satisfied: cuda-toolkit==12.8.1 in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.8.1)\n", + "Requirement already satisfied: cuda-bindings<13,>=12.9.4 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.9.7)\n", + "Requirement already satisfied: nvidia-cudnn-cu12==9.19.0.56 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (9.19.0.56)\n", + "Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (0.7.1)\n", + "Requirement already satisfied: nvidia-nccl-cu12==2.28.9 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (2.28.9)\n", + "Requirement already satisfied: nvidia-nvshmem-cu12==3.4.5 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (3.4.5)\n", + "Requirement already satisfied: triton==3.6.0 in /usr/local/lib/python3.12/dist-packages (from torch>=2.5->tmol==0.1.36+cu128torch2.10) (3.6.0)\n", + "Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.8.4.1)\n", + "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.8.90)\n", + "Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (11.3.3.83)\n", + "Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (1.13.1.3)\n", + "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.8.90)\n", + "Requirement already satisfied: nvidia-curand-cu12==10.3.9.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (10.3.9.90)\n", + "Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (11.7.3.90)\n", + "Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.5.8.93)\n", + "Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.8.93)\n", + "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.8.93)\n", + "Requirement already satisfied: nvidia-nvtx-cu12==12.8.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch>=2.5->tmol==0.1.36+cu128torch2.10) (12.8.90)\n", + "Collecting mypy-extensions>=0.3.0 (from typing_inspect>=0.9->tmol==0.1.36+cu128torch2.10)\n", + " Downloading mypy_extensions-1.1.0-py3-none-any.whl.metadata (1.1 kB)\n", + "Requirement already satisfied: cuda-pathfinder~=1.1 in /usr/local/lib/python3.12/dist-packages (from cuda-bindings<13,>=12.9.4->torch>=2.5->tmol==0.1.36+cu128torch2.10) (1.5.6)\n", + "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas<3,>=2.0->tmol==0.1.36+cu128torch2.10) (1.17.0)\n", + "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch>=2.5->tmol==0.1.36+cu128torch2.10) (1.3.0)\n", + "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->torch>=2.5->tmol==0.1.36+cu128torch2.10) (3.0.3)\n", + "Downloading astor-0.8.1-py2.py3-none-any.whl (27 kB)\n", + "Downloading attrs_strict-1.0.1-py3-none-any.whl (14 kB)\n", + "Downloading biotite-1.7.1-cp312-cp312-manylinux_2_28_x86_64.whl (58.0 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.0/58.0 MB\u001b[0m \u001b[31m20.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading cattrs-26.1.0-py3-none-any.whl (73 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m73.1/73.1 kB\u001b[0m \u001b[31m9.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading hypothesis-6.156.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.1 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m78.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading pint-0.25.3-py3-none-any.whl (307 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m307.5/307.5 kB\u001b[0m \u001b[31m32.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading rdkit-2026.3.4-cp312-cp312-manylinux_2_28_x86_64.whl (37.4 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m37.4/37.4 MB\u001b[0m \u001b[31m17.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading sparse-0.19.0-py2.py3-none-any.whl (155 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m155.9/155.9 kB\u001b[0m \u001b[31m15.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading typing_inspect-0.9.0-py3-none-any.whl (8.8 kB)\n", + "Downloading biotraj-1.2.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.2 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m94.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading flexcache-0.3-py3-none-any.whl (13 kB)\n", + "Downloading flexparser-0.4-py3-none-any.whl (27 kB)\n", + "Downloading mypy_extensions-1.1.0-py3-none-any.whl (5.0 kB)\n", + "Installing collected packages: rdkit, mypy-extensions, hypothesis, flexparser, flexcache, cattrs, attrs_strict, astor, typing_inspect, sparse, pint, biotraj, biotite, tmol\n", + "Successfully installed astor-0.8.1 attrs_strict-1.0.1 biotite-1.7.1 biotraj-1.2.2 cattrs-26.1.0 flexcache-0.3 flexparser-0.4 hypothesis-6.156.6 mypy-extensions-1.1.0 pint-0.25.3 rdkit-2026.3.4 sparse-0.19.0 tmol-0.1.36 typing_inspect-0.9.0\n" + ] + } + ], + "source": [ + "# For Colab\n", + "# install tmol directly from the github wheel;\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": 35, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "y3f4dcA0mV8S", + "outputId": "29bc90e1-4fad-4e2a-a724-58d57032e1a2" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "--2026-07-22 11:47:22-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1ubq.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8000::154, 2606:50c0:8003::154, 2606:50c0:8001::154, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 99763 (97K) [text/plain]\n", + "Saving to: ‘1ubq.pdb’\n", + "\n", + "1ubq.pdb 100%[===================>] 97.42K --.-KB/s in 0.04s \n", + "\n", + "2026-07-22 11:47:22 (2.40 MB/s) - ‘1ubq.pdb’ saved [99763/99763]\n", + "\n", + "--2026-07-22 11:47:22-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1s78.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8001::154, 2606:50c0:8002::154, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 615762 (601K) [text/plain]\n", + "Saving to: ‘1s78.pdb’\n", + "\n", + "1s78.pdb 100%[===================>] 601.33K --.-KB/s in 0.1s \n", + "\n", + "2026-07-22 11:47:23 (4.44 MB/s) - ‘1s78.pdb’ saved [615762/615762]\n", + "\n", + "--2026-07-22 11:47:23-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1qys.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8001::154, 2606:50c0:8002::154, 2606:50c0:8000::154, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8001::154|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 119719 (117K) [text/plain]\n", + "Saving to: ‘1qys.pdb’\n", + "\n", + "1qys.pdb 100%[===================>] 116.91K --.-KB/s in 0.04s \n", + "\n", + "2026-07-22 11:47:23 (2.55 MB/s) - ‘1qys.pdb’ saved [119719/119719]\n", + "\n", + "--2026-07-22 11:47:23-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/cif/1BL8.cif\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8002::154, 2606:50c0:8000::154, 2606:50c0:8003::154, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8002::154|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 330009 (322K) [text/plain]\n", + "Saving to: ‘1BL8.cif’\n", + "\n", + "1BL8.cif 100%[===================>] 322.27K --.-KB/s in 0.1s \n", + "\n", + "2026-07-22 11:47:24 (3.31 MB/s) - ‘1BL8.cif’ saved [330009/330009]\n", + "\n", + "--2026-07-22 11:47:24-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/3plc.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8000::154, 2606:50c0:8003::154, 2606:50c0:8001::154, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 66502 (65K) [text/plain]\n", + "Saving to: ‘3plc.pdb’\n", + "\n", + "3plc.pdb 100%[===================>] 64.94K --.-KB/s in 0.04s \n", + "\n", + "2026-07-22 11:47:24 (1.72 MB/s) - ‘3plc.pdb’ saved [66502/66502]\n", + "\n", + "--2026-07-22 11:47:24-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/10VB.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8000::154, 2606:50c0:8003::154, 2606:50c0:8001::154, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 128304 (125K) [text/plain]\n", + "Saving to: ‘10VB.pdb’\n", + "\n", + "10VB.pdb 100%[===================>] 125.30K --.-KB/s in 0.05s \n", + "\n", + "2026-07-22 11:47:25 (2.31 MB/s) - ‘10VB.pdb’ saved [128304/128304]\n", + "\n", + "--2026-07-22 11:47:25-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/openfold/openfold_ubq_and_sumo.pt\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8000::154, 2606:50c0:8003::154, 2606:50c0:8001::154, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 30304474 (29M) [application/octet-stream]\n", + "Saving to: ‘openfold_ubq_and_sumo.pt’\n", + "\n", + "openfold_ubq_and_su 100%[===================>] 28.90M 17.9MB/s in 1.6s \n", + "\n", + "2026-07-22 11:47:26 (17.9 MB/s) - ‘openfold_ubq_and_sumo.pt’ saved [30304474/30304474]\n", + "\n" + ] + } + ], + "source": [ + "# download some structure files so we have something to work with\n", + "!wget -O 1ubq.pdb https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1ubq.pdb\n", + "!wget -O 1s78.pdb https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1s78.pdb\n", + "!wget -O 1qys.pdb https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1qys.pdb\n", + "!wget -O 1BL8.cif https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/cif/1BL8.cif\n", + "!wget -O 3plc.pdb https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/3plc.pdb\n", + "!wget -O 10VB.pdb https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/10VB.pdb\n", + "!wget -O openfold_ubq_and_sumo.pt https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/openfold/openfold_ubq_and_sumo.pt" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "A_hmNzYjZLH0", + "outputId": "3477568c-0a01-4d71-dad9-6385092be812" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ATOM 1 N MET A 1 27.340 24.430 2.614 1.00 9.67 N \n", + "ATOM 2 CA MET A 1 26.266 25.413 2.842 1.00 10.38 C \n", + "ATOM 3 C MET A 1 26.913 26.639 3.531 1.00 9.62 C \n", + "ATOM 4 O MET A 1 27.886 26.463 4.263 1.00 9.62 O \n", + "ATOM 5 CB MET A 1 25.112 24.880 3.649 1.00 13.77 C \n", + "ATOM 6 CG MET A 1 25.353 24.860 5.134 1.00 16.29 C \n", + "ATOM 7 SD MET A 1 23.930 23.959 5.904 1.00 17.17 S \n", + "ATOM 8 CE MET A 1 24.447 23.984 7.620 1.00 16.11 C \n", + "ATOM 9 1H MET A 1 26.961 23.619 2.168 1.00 0.00 H \n", + "ATOM 10 2H MET A 1 28.043 24.834 2.029 1.00 0.00 H \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": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "env: TMOL_USE_JIT=1\n" + ] + } + ], + "source": [ + "%env TMOL_USE_JIT=1" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import tmol" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "['/Users/leaverfa/GIT/tmol/tmol']\n" + ] + } + ], + "source": [ + "print(tmol.__path__)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "Mi9gshJmbVVC", + "outputId": "6a3b6001-415b-4570-a8b3-2b56561a709d" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "tensor([235.1642], grad_fn=)\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": 5, + "metadata": { + "id": "lBCg3XA8jMq8" + }, + "outputs": [], + "source": [ + "# - Initialize the default ParameterDatabase\n", + "default_param_db = tmol.ParameterDatabase.get_default()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9nyxfNgojaSB" + }, + "outputs": [], + "source": [ + "# - Load alternate block types\n", + "# TO DO: Kieran" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "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": null, + "metadata": { + "id": "PyhxBHq9pLMx" + }, + "outputs": [], + "source": [ + "# - Initialize a PackedBlockTypes object from a custom database\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "id": "OO32waIlp4-5" + }, + "outputs": [], + "source": [ + "# - Create a default CanonicalOrdering object\n", + "canonical_ordering = tmol.default_canonical_ordering()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "L-uVsXPLjiDC" + }, + "outputs": [], + "source": [ + "# - Create a CanonicalOrdering from a ParameterDatabase\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oXpe5H0Rjmfr" + }, + "outputs": [], + "source": [ + "# - Create a PackedBlockTypes object from a subset of the block types in a ParameterDatabase\n", + "# TO DO" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "42PMj1JkkTaD" + }, + "source": [ + "Input" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "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": 9, + "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": 9, + "metadata": { + "id": "N8RfPFCWYkfy" + }, + "outputs": [], + "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\n", + "from tmol.io.pose_stack_from_biotite 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)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "ysGCJuh8kXRc", + "outputId": "ea818011-19e4-4b37-c638-34ef10bef4d7" + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "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\n", + "from tmol.io.pose_stack_from_biotite 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", + "# TEMP: add pose_stack_from_biotite to API and modify this line\n", + "pose_1bl8 = pose_stack_from_biotite(bt_struct, device)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "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": 12, + "metadata": { + "id": "CgVHFNkVkYUZ" + }, + "outputs": [], + "source": [ + "# - Initialize a single pose PoseStack with missing residues\n", + "# This particular structure has a few regions where the backbone is missing\n", + "pose_1s78 = tmol.pose_stack_from_pdb(\"1s78.pdb\", device=device)\n", + "# TO DO: Assert that some residues are not bound to their upper / lower conns" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": { + "id": "HX9jTnCrkYry" + }, + "outputs": [], + "source": [ + "# - Load a bunch of PoseStacks from different sources and then concatenate 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: Add PoseStackBuilder to API\n", + "from tmol.pose.pose_stack_builder import PoseStackBuilder\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": 14, + "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": null, + "metadata": { + "id": "zIacYunXkZX4" + }, + "outputs": [], + "source": [ + "# - Add custom residue type\n", + "# TO DO\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "B9BJmHWCkZxQ" + }, + "outputs": [], + "source": [ + "# - Set the scoring parameters for a new residue type (elec, etc)\n", + "# TO DO\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "MTd5eVotkaJ9" + }, + "outputs": [], + "source": [ + "# - Build an extended pose from sequence\n", + "# TO DO" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "R1LJAlYJktGE" + }, + "source": [ + "Output" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "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": 17, + "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": 18, + "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": 80, + "metadata": { + "id": "wYa6C8EGk1GT" + }, + "outputs": [], + "source": [ + "# - Write a single pose PoseStack to an .mmcif file\n", + "# TO DO\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "gcNdsybLk1dR" + }, + "outputs": [], + "source": [ + "# - Write a rotamer set out as a multi-model PDB\n", + "# TO DO\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "vucunYYYk2OX" + }, + "source": [ + "Kinematics" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": { + "id": "PYwjsJDvk4Ka" + }, + "outputs": [], + "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 3-tuple of (edge_type, start_residue, end_residue)\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.\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, 3] with the sentinel value\n", + "# of -1 used for unused edges.\n", + "edges = numpy.full((1, 2, 3), -1, dtype=int)\n", + "edges[0, 0] = [tmol.EdgeType.root_jump, -1, 0]\n", + "edges[0, 1] = [tmol.EdgeType.polymer, 0, 75]\n", + "ff2_1ubq = tmol.FoldForest.from_edges(edges)" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": { + "id": "LVS2vrVRZGbS" + }, + "outputs": [], + "source": [ + "# - Create a set of N->C fold trees for a multi-pose PoseStack of 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, 3), -1, dtype=int)\n", + "edges[:, 0] = numpy.array([tmol.EdgeType.root_jump, -1, 0], 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", + "ff2_ps3 = tmol.FoldForest.from_edges(edges)" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "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": 23, + "metadata": { + "id": "e7m9yUztlKmg" + }, + "outputs": [], + "source": [ + "# - Create a simple fold tree for a PoseStack with missing residues\n", + "\n", + "# TO DO: Fix logic for reasonable_fold_forest so that it inserts jumps\n", + "# between blocks that are in the same chain but that are not connected.\n", + "ff_1s78 = tmol.FoldForest.reasonable_fold_forest(pose_1s78)" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": { + "id": "iblPqPjZlK8Z" + }, + "outputs": [], + "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", + "#\n", + "edges = numpy.full((3, pose_stack_3.max_n_blocks, 3), -1, dtype=int)\n", + "edges[:, :, 0] = tmol.EdgeType.root_jump\n", + "edges[:, :, 2] = numpy.arange(pose_stack_3.max_n_blocks, dtype=int)[None, :]\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)" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "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": 26, + "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": null, + "metadata": { + "id": "ZTRWFlPslL9b" + }, + "outputs": [], + "source": [ + "# - Apply a perturbation to the rigid-body DOFs between two chains\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "nM9g3oIwlMTy" + }, + "outputs": [], + "source": [ + "# - Assign dihedral values to all the residues in a PoseStack and calculate the coordinates\n", + "# TO DO" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "klTMFJK6lQD1" + }, + "source": [ + "Scoring" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": { + "id": "EnQIxSW4lWjZ" + }, + "outputs": [], + "source": [ + "# - Create the default score function\n", + "sfxn = tmol.beta2016_score_function(device)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Yc8a4BnklZoh" + }, + "outputs": [], + "source": [ + "# - Create the default score function from a custom Database\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5qErnZZSlbpA" + }, + "outputs": [], + "source": [ + "# - Create the soft-rep version of the score function\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": { + "id": "jLR7q9nXlc4H" + }, + "outputs": [], + "source": [ + "# - Create an empty score function\n", + "sfxn_empty = tmol.ScoreFunction(default_param_db, device)" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "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": 30, + "metadata": { + "id": "yMgEZFkQmYt_" + }, + "outputs": [], + "source": [ + "# - Turn off a term in a score function\n", + "# TO DO; fix bug\n", + "# print(len(sfxn_rep_and_hbonds.all_terms()))\n", + "\n", + "sfxn_rep_and_hbonds.set_weight(tmol.ScoreType.hbond, 0)\n", + "assert len(sfxn_rep_and_hbonds.all_terms()) == 1\n", + "\n", + "sfxn_rep_and_hbonds.set_weight(tmol.ScoreType.fa_ljrep, 0)\n", + "assert len(sfxn_rep_and_hbonds.all_terms()) == 0\n" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "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": 32, + "metadata": { + "id": "XWLjbPsxmZYR" + }, + "outputs": [], + "source": [ + "# - Score a PoseStack and back-propagate through the coordinates\n", + "score2 = wpsm(pose_1ubq.coords)\n", + "score2.sum().backward()" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "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": 34, + "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": 35, + "metadata": { + "id": "dYzEVvc7maRK" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "tot reg tensor(19602.2754, grad_fn=) tot_upweight_iface tensor(19575.9206, dtype=torch.float64, grad_fn=)\n" + ] + } + ], + "source": [ + "# - Score a PoseStack and return per-residue weighted energies, weight them according to some principle,\n", + "# and then back-propagate the total reweighted 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", + "# in 1s78, there are 555 residues in the antigen; the remaining residues are in the antibody\n", + "upweight_interchain_intxns = torch.ones((1, 991, 991), dtype=float, device=device)\n", + "# energies are written to the upper-triangle of the n-res x n-res table\n", + "upweight_interchain_intxns[0, 0:555, 555:991] = 2.0\n", + "\n", + "tot_upweight_iface = (bps * upweight_interchain_intxns).sum()\n", + "print(\"tot reg\", tot_reg, \"tot_upweight_iface\", tot_upweight_iface)\n", + "dcoords = tot_upweight_iface.backward()" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": { + "id": "-NW-sfNJmapX" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "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": 37, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "score w/ constraints tensor([235.1642], grad_fn=)\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": 38, + "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": 39, + "metadata": { + "id": "A26oZ-8ImbZf" + }, + "outputs": [], + "source": [ + "# - Add coordinate constraints to the current coordinates\n", + "\n", + "# Option 1: use the utility function\n", + "# TEMP!\n", + "pose_1ubq_w_ca_coord_csts1 = tmol.constrain_all_ca(pose_1ubq)\n", + "\n", + "# Option 2: use the other utility function\n", + "pose_1ubq_w_ca_coord_csts2 = tmol.create_mainchain_coordinate_constraints(pose_1ubq)" + ] + }, + { + "cell_type": "code", + "execution_count": 67, + "metadata": { + "id": "0tOaC2Trmbvc" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "score_std tensor([235.1642], grad_fn=) score_alt tensor([235.1841], grad_fn=)\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", + "# This is a somewhat advanced feature, so CartBondedDatabase is not imported\n", + "# with tmol by default, so we have to give its full scope\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", + "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((i, p) for (i, p) in enumerate(pro_length_params) if p.atm1 == \"N\" and p.atm2 == \"CD\")\n", + "alt_cd_n_bond_params = attr.evolve(cd_n_bond_params, K=244) # double the default strength of ~122 kcal/mol*A\n", + "alt_length_params = pro_length_params[0:ind] + (alt_cd_n_bond_params,) + pro_length_params[ind+1:]\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 = tmol.database.scoring.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", + "assert score_std < score_alt" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "dIPA0OfBmjhh" + }, + "source": [ + "Optimization" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "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", + "dun_sampler = tmol.create_dunbrack_sampler_from_database(default_param_db, device)" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pack_rotamers 4.89 build rots: 0.04 calcRPEs: 4.75 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 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 PackRotamers\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 = tmol.PackerPalette()\n", + "\n", + "# 2.\n", + "task = tmol.PackerTask(pose_1ubq, palette)\n", + "\n", + "# 3.\n", + "task.restrict_to_repacking()\n", + "task.add_conformer_sampler(dun_sampler)\n", + "task.add_conformer_sampler(tmol.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(tmol.IncludeCurrentSampler())\n", + "\n", + "# 4.\n", + "pose_1ubq_repacked = tmol.pack_rotamers(pose_1ubq, sfxn, task, verbose=True)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "score tensor([-100.3825], grad_fn=)\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": 13, + "metadata": { + "id": "nDVbq3pDmqpf" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pack_rotamers 4.89 build rots: 0.02 calcRPEs: 4.76 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\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 alaphatic 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\n", + "# packer; only repack the hydroxyl positions\n", + "\n", + "optH_task = tmol.PackerTask(pose_1ubq, palette)\n", + "optH_task.restrict_to_repacking()\n", + "optH_task.add_conformer_sampler(tmol.IncludeCurrentSampler())\n", + "optH_task.add_conformer_sampler(tmol.OptHSampler())\n", + "\n", + "pose_1ubq_optH_repacked = tmol.pack_rotamers(pose_1ubq, sfxn, task, verbose=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "e3a4AOs8mrDX" + }, + "outputs": [], + "source": [ + "# - Fill in side chains for a model that lacks them\n", + "# TO DO\n", + "# Sometimes your models will only include the backbone, or\n", + "# you are reading from the PDB where the electron density\n", + "# was absent for a side chain and the crystalographer did\n", + "# not add its conformation.\n", + "\n", + "# When working with a structure file, the\n", + "# pose_stack_from_biotite pathway will automatically \n", + "# build in missing side chains.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Tgj8DwX8mrZ6" + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": { + "id": "zt23WbYtmru3" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pack_rotamers 188.98 build rots: 0.13 calcRPEs: 185.85 build IG: 2.48 build IG part2: 0.00 run SA: 0.52 pose ctor: 0.00\n" + ] + } + ], + "source": [ + "# - Repack with extra rotamers\n", + "\n", + "# the \"extra chi\" flags in the PackerTask will trigger\n", + "# sampling at +/- 1stdev\n", + "palette = tmol.PackerPalette()\n", + "task = tmol.PackerTask(pose_1ubq, palette)\n", + "\n", + "# 3.\n", + "task.restrict_to_repacking()\n", + "task.add_conformer_sampler(dun_sampler)\n", + "task.add_conformer_sampler(tmol.FixedAAChiSampler())\n", + "task.add_conformer_sampler(tmol.IncludeCurrentSampler())\n", + "task.or_expand_chi(1)\n", + "task.or_expand_chi(2)\n", + "\n", + "# 4.\n", + "pose_1ubq_repacked_ex1ex2 = tmol.pack_rotamers(pose_1ubq, sfxn, task, verbose=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "score tensor([-111.3373], grad_fn=)\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": 24, + "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": null, + "metadata": { + "id": "awPe0TtYmsaR" + }, + "outputs": [], + "source": [ + "# - Create a new PackerPalette subclass to handle logic of new block types\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "RTm3LjhAmsv4" + }, + "outputs": [], + "source": [ + "# - Run minimization in double precision\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": { + "id": "Z2UFgON6mtFt" + }, + "outputs": [], + "source": [ + "# - Perform cartesian minimization\n", + "\n", + "cart_minimized_pose_1ubq = tmol.run_cart_min(pose_1ubq, sfxn)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "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": 9, + "metadata": { + "id": "oE7_fHfEmtwM" + }, + "outputs": [], + "source": [ + "# - Relax a PoseStack w/ kinematic minimization\n", + "\n", + "packer_palette = tmol.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", + "kin_relaxed_1ubq = tmol.kin_fast_relax(pose_1ubq, sfxn, packer_palette, mm_1ubq_all, ff1_1ubq, verbose=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "id": "9ZUjD_eYmuFw" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "packing with fa_rep of 0.02 and constraint weight of 0.00\n", + "pack_rotamers 4.89 build rots: 0.03 calcRPEs: 4.74 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.03 and constraint weight of 0.00\n", + " LBFGS_Armijo done: 200 iters, 305 func evals, E=-376.0524\n", + "run_min 5.20 setup: 0.00 opt 5.20 stack-ctor: 0.00\n", + "pack-min 10.10 task-init 0.00 packing 4.89 min 5.21\n", + "packing with fa_rep of 0.15 and constraint weight of 0.00\n", + "pack_rotamers 5.20 build rots: 0.02 calcRPEs: 5.07 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.15 and constraint weight of 0.00\n", + " converged: rel_dE 3.2835e-04 <= rtol 0.00034526698300124393\n", + " LBFGS_Armijo done: 69 iters, 108 func evals, E=-328.6090\n", + "run_min 1.84 setup: 0.00 opt 1.84 stack-ctor: 0.00\n", + "pack-min 7.04 task-init 0.00 packing 5.20 min 1.84\n", + "packing with fa_rep of 0.31 and constraint weight of 0.00\n", + "pack_rotamers 5.05 build rots: 0.02 calcRPEs: 4.92 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.32 and constraint weight of 0.00\n", + " converged: rel_dE 3.3045e-04 <= rtol 0.00034526698300124393\n", + " LBFGS_Armijo done: 52 iters, 83 func evals, E=-303.9801\n", + "run_min 1.38 setup: 0.00 opt 1.38 stack-ctor: 0.00\n", + "pack-min 6.43 task-init 0.00 packing 5.05 min 1.38\n", + "packing with fa_rep of 0.55 and constraint weight of 0.00\n", + "pack_rotamers 4.92 build rots: 0.02 calcRPEs: 4.79 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.55 and constraint weight of 0.00\n", + " converged: rel_dE 3.2141e-04 <= rtol 0.00034526698300124393\n", + " LBFGS_Armijo done: 10 iters, 21 func evals, E=-276.7285\n", + "run_min 0.34 setup: 0.00 opt 0.34 stack-ctor: 0.00\n", + "pack-min 5.27 task-init 0.00 packing 4.92 min 0.35\n", + "accepting new best scores\n", + " old best score: tensor([235.1642], grad_fn=)\n", + " new best score: tensor([-276.7285], grad_fn=)\n", + "packing with fa_rep of 0.02 and constraint weight of 0.00\n", + "pack_rotamers 4.89 build rots: 0.02 calcRPEs: 4.76 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.03 and constraint weight of 0.00\n", + " converged: rel_dE 3.1333e-04 <= rtol 0.00034526698300124393\n", + " LBFGS_Armijo done: 80 iters, 125 func evals, E=-380.3913\n", + "run_min 2.10 setup: 0.00 opt 2.10 stack-ctor: 0.00\n", + "pack-min 6.99 task-init 0.00 packing 4.89 min 2.10\n", + "packing with fa_rep of 0.15 and constraint weight of 0.00\n", + "pack_rotamers 4.98 build rots: 0.02 calcRPEs: 4.85 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.15 and constraint weight of 0.00\n", + " converged: rel_dE 2.9407e-04 <= rtol 0.00034526698300124393\n", + " LBFGS_Armijo done: 52 iters, 83 func evals, E=-333.4777\n", + "run_min 1.42 setup: 0.00 opt 1.42 stack-ctor: 0.00\n", + "pack-min 6.41 task-init 0.00 packing 4.98 min 1.42\n", + "packing with fa_rep of 0.31 and constraint weight of 0.00\n", + "pack_rotamers 4.96 build rots: 0.02 calcRPEs: 4.83 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.32 and constraint weight of 0.00\n", + " converged: rel_dE 3.0506e-04 <= rtol 0.00034526698300124393\n", + " LBFGS_Armijo done: 45 iters, 74 func evals, E=-306.8680\n", + "run_min 1.22 setup: 0.00 opt 1.22 stack-ctor: 0.00\n", + "pack-min 6.19 task-init 0.00 packing 4.97 min 1.22\n", + "packing with fa_rep of 0.55 and constraint weight of 0.00\n", + "pack_rotamers 4.91 build rots: 0.02 calcRPEs: 4.78 build IG: 0.07 build IG part2: 0.00 run SA: 0.04 pose ctor: 0.00\n", + "minimizing with fa_rep of 0.55 and constraint weight of 0.00\n", + " converged: rel_dE 3.1894e-04 <= rtol 0.00034526698300124393\n", + " LBFGS_Armijo done: 40 iters, 65 func evals, E=-283.6534\n", + "run_min 1.06 setup: 0.00 opt 1.06 stack-ctor: 0.00\n", + "pack-min 5.98 task-init 0.00 packing 4.91 min 1.07\n", + "accepting new best scores\n", + " old best score: tensor([-276.7285], grad_fn=)\n", + " new best score: tensor([-283.6534], grad_fn=)\n" + ] + } + ], + "source": [ + "# - Relax a PoseStack w/ cartesian minimization\n", + "packer_palette = tmol.PackerPalette()\n", + "\n", + "# default CartesianMoveMap allows all atoms to move\n", + "cart_mm = tmol.CartesianMoveMap()\n", + "\n", + "cart_relaxed_1ubq = tmol.cartesian_fast_relax(pose_1ubq, sfxn, packer_palette, cart_mm, verbose=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2aUJd4X5mubN" + }, + "outputs": [], + "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\n", + "# structures and with many-pose PoseStacks. Then, \n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Mkj2U1Namuw7" + }, + "outputs": [], + "source": [ + "# - Idealize a structure from the PDB\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "aWFR74xZmvFu" + }, + "outputs": [], + "source": [ + "# - Idealize a structure from a dandelion\n", + "# TO DO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "NDZyRGtTmva3" + }, + "outputs": [], + "source": [ + "# - Idealize just the backbone of a dandelion\n", + "# TO DO" + ] + } + ], + "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 +} diff --git a/tmol/__init__.py b/tmol/__init__.py index a49b21077..ce184da67 100644 --- a/tmol/__init__.py +++ b/tmol/__init__.py @@ -39,6 +39,7 @@ def include_paths(): default_canonical_ordering, default_packed_block_types, pose_stack_from_canonical_form, + pose_stack_from_biotite, canonical_form_from_openfold, canonical_ordering_for_openfold, packed_block_types_for_openfold, @@ -65,9 +66,11 @@ def include_paths(): run_kin_min, run_min, ) + from tmol.pose import ( PackedBlockTypes, PoseStack, + PoseStackBuilder, ConstraintSet, get_named_torsions, get_torsion_names, @@ -79,9 +82,29 @@ def include_paths(): ) from tmol.score.constraint import ( ConstraintEnergyTerm, + constrain_all_ca, create_mainchain_coordinate_constraints, ) -from tmol.relax import fast_relax +from tmol.pack import ( + pack_rotamers, + PackerPalette, + PackerTask, +) +from tmol.pack.rotamer import ( + FixedAAChiSampler, + IncludeCurrentSampler, + OptHSampler, +) +from tmol.pack.rotamer.dunbrack import ( + create_dunbrack_sampler_from_database, + DunbrackChiSampler, +) + +from tmol.relax import ( + fast_relax, + kin_fast_relax, + cartesian_fast_relax, +) try: __version__ = version("tmol") @@ -100,6 +123,7 @@ def include_paths(): "PackedBlockTypes", "ParameterDatabase", "PoseStack", + "PoseStackBuilder", "ScoreFunction", "ScoreType", "atom_records_from_pose_stack", @@ -110,6 +134,7 @@ def include_paths(): "canonical_form_from_rosettafold2", "canonical_ordering_for_openfold", "canonical_ordering_for_rosettafold2", + "constrain_all_ca", "create_mainchain_coordinate_constraints", "default_canonical_ordering", "default_packed_block_types", @@ -121,6 +146,7 @@ def include_paths(): "one2three", "packed_block_types_for_openfold", "packed_block_types_for_rosettafold2", + "pose_stack_from_biotite", "pose_stack_from_canonical_form", "pose_stack_from_openfold", "pose_stack_from_pdb", @@ -135,4 +161,14 @@ def include_paths(): "three2one", "view", "write_pose_stack_pdb", + "pack_rotamers", + "PackerPalette", + "PackerTask", + "FixedAAChiSampler", + "IncludeCurrentSampler", + "OptHSampler", + "create_dunbrack_sampler_from_database", + "DunbrackChiSampler", + "kin_fast_relax", + "cartesian_fast_relax", ] diff --git a/tmol/pose/_pose_stack.py b/tmol/pose/_pose_stack.py index 64cdf20c1..cabf27912 100644 --- a/tmol/pose/_pose_stack.py +++ b/tmol/pose/_pose_stack.py @@ -1,6 +1,6 @@ import attr import torch -from typing import Optional, TYPE_CHECKING +from typing import Optional, TYPE_CHECKING, Union from tmol.types import Tensor from tmol.chemical import RefinedResidueType @@ -288,6 +288,12 @@ def split(self, index) -> "PoseStack": ), ) + def to(self, dtype=Union[torch.float32, torch.float64]) -> "PoseStack": + """Create a new PoseStack with the dtype of the coords tensor changed to the requested dtype.""" + if self.coords.dtype == dtype: + return self + return attr.evolve(self, coords=self.coords.to(dtype=dtype)) + def expand_coords(self): """Load the coordinates into a 4D tensor: n_poses x max_n_blocks x max_n_atoms_per_block x 3 diff --git a/tmol/relax/__init__.py b/tmol/relax/__init__.py index 0579b0a4d..048a410c5 100644 --- a/tmol/relax/__init__.py +++ b/tmol/relax/__init__.py @@ -1,8 +1,11 @@ from ._fast_relax import ( # noqa: F401 DEFAULT_RELAX_SCHEDULE, - _default_cart_min_fn, + default_cart_min_fn, + default_kin_min_fn, accept_best, fast_relax, + kin_fast_relax, + cartesian_fast_relax, relax_pack_min_step, ) @@ -10,5 +13,7 @@ "DEFAULT_RELAX_SCHEDULE", "accept_best", "fast_relax", + "kin_fast_relax", + "cartesian_fast_relax", "relax_pack_min_step", ] diff --git a/tmol/relax/_fast_relax.py b/tmol/relax/_fast_relax.py index af97e4084..68f97306d 100644 --- a/tmol/relax/_fast_relax.py +++ b/tmol/relax/_fast_relax.py @@ -120,7 +120,7 @@ def constraint_fraction(step_index): return normalized -def _default_kin_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose): +def default_kin_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose): """Default minimization function: kinematic (torsion-space) LBFGS.""" return run_kin_min( pose_stack, @@ -132,7 +132,7 @@ def _default_kin_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose): ) -def _default_cart_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose): +def default_cart_min_fn(pose_stack, sfxn, *, fold_forest, move_map, verbose): """Default Cartesian minimization function for use as fast_relax min_fn. Extracts ``coord_mask`` from ``move_map`` if it is a @@ -159,7 +159,7 @@ def fast_relax( # noqa: C901 sfxn: ScoreFunction, packer_pallete: PackerPalette, move_map: Union[MoveMap, CartesianMoveMap], - fold_forest: FoldForest, + fold_forest: Optional[FoldForest], *, task_operations=None, num_repeats=2, @@ -176,7 +176,9 @@ def fast_relax( # noqa: C901 followed by an accept-to-best check. Args: - pose_stack: The input poses to relax. + pose_stack: The input poses to relax. Relax will use the precision of + the input coords tensor during minimization, but will only use + torch.float32 precision for packing. sfxn: Score function used for packing and minimization. If you wish to use constraints during relax, then the weight on the "constraint" score type must already have a non-zero value. @@ -251,7 +253,7 @@ def my_min(ps, sfxn, *, fold_forest, move_map, **kw): The relaxed PoseStack (best-scoring across all repeats). """ if min_fn is None: - min_fn = _default_cart_min_fn + min_fn = default_cart_min_fn if schedule is None: schedule = DEFAULT_RELAX_SCHEDULE @@ -352,6 +354,11 @@ def relax_pack_min_step( min_fn, verbose, ): + """Perform a single pack-min step of the FastRelax protocol. + + Convert the PoseStack to float32 for packing, then restore + it to the input dtype afterwards.""" + input_pose_dtype = pose_stack.coords.dtype if verbose and torch.cuda.is_available(): torch.cuda.synchronize() @@ -369,7 +376,12 @@ def relax_pack_min_step( if verbose and torch.cuda.is_available(): torch.cuda.synchronize() end_time1 = time.perf_counter() + + # convert pose_stack to float32 for packing, and restore it + # to the input dtype afterwards + pose_stack = pose_stack.to(torch.float32) packed_pose_stack = pack_rotamers(pose_stack, sfxn, task, verbose) + packed_pose_stack = packed_pose_stack.to(dtype=input_pose_dtype) sfxn.set_weight(ScoreType.fa_ljrep, fa_rep_min_weight) if verbose: @@ -442,3 +454,68 @@ def select_better(tensor_name): return new_best_pose_stack, new_best_pose_score else: # no change return best_pose_stack, best_pose_score + + +def kin_fast_relax( + pose_stack: PoseStack, + sfxn: ScoreFunction, + packer_pallete: PackerPalette, + move_map: MoveMap, + fold_forest: FoldForest, + *, + task_operations=None, + num_repeats=2, + ramp_constraints: Optional[bool] = None, # default True + schedule=None, + min_fn=default_kin_min_fn, + verbose: bool = False, +): + """Run the FastRelax protocol using kinematic (torsion-space) minimization. + + See documentation for fast_relax. + """ + return fast_relax( + pose_stack, + sfxn, + packer_pallete, + move_map, + fold_forest, + task_operations=task_operations, + num_repeats=num_repeats, + ramp_constraints=ramp_constraints, + schedule=schedule, + min_fn=min_fn or default_kin_min_fn, + verbose=verbose, + ) + + +def cartesian_fast_relax( + pose_stack: PoseStack, + sfxn: ScoreFunction, + packer_pallete: PackerPalette, + move_map: CartesianMoveMap, + *, + task_operations=None, + num_repeats=2, + ramp_constraints: Optional[bool] = None, # default True + schedule=None, + min_fn=default_cart_min_fn, + verbose: bool = False, +): + """Run the FastRelax protocol using Cartesian (coordinate-space) minimization. + + See documentation for fast_relax. + """ + return fast_relax( + pose_stack, + sfxn, + packer_pallete, + move_map, + None, + task_operations=task_operations, + num_repeats=num_repeats, + ramp_constraints=ramp_constraints, + schedule=schedule, + min_fn=min_fn or default_kin_min_fn, + verbose=verbose, + ) diff --git a/tmol/score/_score_function.py b/tmol/score/_score_function.py index 1a969ec2a..39eeac97e 100644 --- a/tmol/score/_score_function.py +++ b/tmol/score/_score_function.py @@ -53,99 +53,49 @@ def __init__(self, param_db: ParameterDatabase, device: torch.device): self.term_options = {} def set_weight(self, st: ScoreType, weight: float): - if not self.score_type_covered_by_contained_term(st): - self.retrieve_term_for_score_type(st) - if weight == 0 and self.term_for_st_has_no_other_non_zero_weights(st): - self.remove_term_for_score_type(st) # TO DO! + """Turn on or off a score type by setting its weight. + + This method will + """ + + if self._weights[st.value] == weight: + return + if not self._score_type_covered_by_contained_term(st): + self._retrieve_term_for_score_type(st) + elif weight == 0 and self._term_for_st_has_no_other_non_zero_weights(st): + self._remove_term_for_score_type(st) # TO DO! self._weights[st.value] = weight self._weights_tensor_out_of_date = True def get_weight(self, st: ScoreType): return self._weights[st.value] - def score_type_covered_by_contained_term(self, st: ScoreType): - for term in self._all_terms: - if st in term.score_types(): - return True - return False - - def retrieve_term_for_score_type(self, st: ScoreType): - term = ScoreTermFactory.create_term_for_score_type( - st, self._param_db, self._device - ) - # sanity check: if the ScoreTermFactory returns the wrong term, - # we want to know - assert st in term.score_types() - for tst in term.score_types(): - self._term_for_st[tst.value] = term - self._all_terms_unordered.append(term) - self._all_terms_out_of_date = True - if term.n_bodies() == 1: - self._one_body_terms_unordered.append(term) - self._one_body_terms_out_of_date = True - elif term.n_bodies() == 2: - self._two_body_terms_unordered.append(term) - self._two_body_terms_out_of_date = True - else: - self._multi_body_terms_unordered.append(term) - self._multi_body_terms_out_of_date = True - - def term_for_st_has_no_other_non_zero_weights(self, st: ScoreType): - term = self._term_for_st[st.value] - for st2 in term.score_types(): - if st2 == st: - continue - if self._weights[st2.value] != 0: - return True - return False - def all_terms(self): - """Grant read access to the list of terms. - - Do not modify this list directly - """ - if self._all_terms_out_of_date: - self._all_terms, self._all_score_types = self.get_sorted_terms( - self._all_terms_unordered - ) - self._all_terms_out_of_date = False - + """Grant read access to the list of terms.""" + self._update_all_terms() return self._all_terms def all_score_types(self): - if self._all_terms_out_of_date: - self._all_terms, self._all_score_types = self.get_sorted_terms( - self._all_terms_unordered - ) - self._all_terms_out_of_date = False + """Return the list of the currently-active score types + These are sorted by their order in the ScoreType enum, which + is also the order in which terms are ordered in the output + unweighted score tensor of the WholePoseScoringModule and + BlockPairScoringModule. + """ + self._update_all_terms() return self._all_score_types def one_body_terms(self): - if self._one_body_terms_out_of_date: - self._one_body_terms, _ = self.get_sorted_terms( - self._one_body_terms_unordered - ) - self._one_body_terms_out_of_date = False - + self._update_one_body_terms() return self._one_body_terms def two_body_terms(self): - if self._two_body_terms_out_of_date: - self._two_body_terms, _ = self.get_sorted_terms( - self._two_body_terms_unordered - ) - self._two_body_terms_out_of_date = False - + self._update_two_body_terms() return self._two_body_terms def multi_body_terms(self): - if self._multi_body_terms_out_of_date: - self._multi_body_terms, _ = self.get_sorted_terms( - self._multi_body_terms_unordered - ) - self._multi_body_terms_out_of_date = False - + self._update_multi_body_terms() return self._multi_body_terms def render_whole_pose_scoring_module(self, pose_stack: PoseStack): @@ -157,7 +107,7 @@ def render_whole_pose_scoring_module(self, pose_stack: PoseStack): object's __call__ will return a tensor of weighted energies of shape (n_poses,). """ - self.pre_work_initialization(pose_stack) + self._pre_work_initialization(pose_stack) term_modules = [ t.render_whole_pose_scoring_module(pose_stack) for t in self.all_terms() ] @@ -172,7 +122,7 @@ def render_block_pair_scoring_module(self, pose_stack: PoseStack): object's __call__ will return a tensor of weighted energies of shape (n_poses, max_n_blocks, max_n_blocks). """ - self.pre_work_initialization(pose_stack) + self._pre_work_initialization(pose_stack) term_modules = [ t.render_block_pair_scoring_module(pose_stack) for t in self.all_terms() ] @@ -189,32 +139,18 @@ def render_rotamer_scoring_module( object's __call__ will return a tensor of weighted energies of shape (n_poses, max_n_blocks, max_n_blocks). """ - self.pre_work_initialization(pose_stack) + self._pre_work_initialization(pose_stack) term_modules = [ t.render_rotamer_scoring_module(pose_stack, rotamer_set) for t in self.all_terms() ] return RotamerScoringModule(self.weights_tensor(), term_modules) - def pre_work_initialization(self, pose_stack: PoseStack): - # set_options must be first, since some of the logic that follows it - # may depend on the options - for energy_term in self.all_terms(): - energy_term.set_options(self.term_options) - - for block_type in pose_stack.packed_block_types.active_block_types: - for energy_term in self.all_terms(): - energy_term.setup_block_type(block_type) - for energy_term in self.all_terms(): - energy_term.setup_packed_block_types(pose_stack.packed_block_types) - for energy_term in self.all_terms(): - energy_term.setup_poses(pose_stack) - def set_option(self, key: str, value): """Set an option for all energy terms. Options are passed to each energy term's set_options method - as a dictionary during pre_work_initialization. + as a dictionary during _pre_work_initialization. """ self.term_options[key] = value @@ -291,8 +227,102 @@ def from_sfxn_file(cls, path, param_db, device): sfxn.set_options(data["options"]) return sfxn + def _score_type_covered_by_contained_term(self, st: ScoreType): + return self._term_for_st[st.value] is not None + + def _retrieve_term_for_score_type(self, st: ScoreType): + term = ScoreTermFactory.create_term_for_score_type( + st, self._param_db, self._device + ) + # sanity check: if the ScoreTermFactory returns the wrong term, + # we want to know + assert st in term.score_types() + for tst in term.score_types(): + self._term_for_st[tst.value] = term + self._all_terms_unordered.append(term) + self._all_terms_out_of_date = True + if term.n_bodies() == 1: + self._one_body_terms_unordered.append(term) + self._one_body_terms_out_of_date = True + elif term.n_bodies() == 2: + self._two_body_terms_unordered.append(term) + self._two_body_terms_out_of_date = True + else: + self._multi_body_terms_unordered.append(term) + self._multi_body_terms_out_of_date = True + + def _term_for_st_has_no_other_non_zero_weights(self, st: ScoreType): + term = self._term_for_st[st.value] + for st2 in term.score_types(): + if st2 == st: + continue + if self._weights[st2.value] != 0: + return False + return True + + def _remove_term_for_score_type(self, st: ScoreType): + term = self._term_for_st[st.value] + self._all_terms_unordered.remove(term) + self._all_terms_out_of_date = True + + if term.n_bodies() == 1: + self._one_body_terms_unordered.remove(term) + self._one_body_terms_out_of_date = True + elif term.n_bodies() == 2: + self._two_body_terms_unordered.remove(term) + self._two_body_terms_out_of_date = True + else: + self._multi_body_terms_unordered.remove(term) + self._multi_body_terms_out_of_date = True + + self._weights_tensor_out_of_date = True + for tst in term.score_types(): + self._term_for_st[tst.value] = None + + def _pre_work_initialization(self, pose_stack: PoseStack): + # set_options must be first, since some of the logic that follows it + # may depend on the options + for energy_term in self.all_terms(): + energy_term.set_options(self.term_options) + + for block_type in pose_stack.packed_block_types.active_block_types: + for energy_term in self.all_terms(): + energy_term.setup_block_type(block_type) + for energy_term in self.all_terms(): + energy_term.setup_packed_block_types(pose_stack.packed_block_types) + for energy_term in self.all_terms(): + energy_term.setup_poses(pose_stack) + + def _update_all_terms(self): + if self._all_terms_out_of_date: + self._all_terms, self._all_score_types = self._get_sorted_terms( + self._all_terms_unordered + ) + self._all_terms_out_of_date = False + + def _update_one_body_terms(self): + if self._one_body_terms_out_of_date: + self._one_body_terms, _ = self._get_sorted_terms( + self._one_body_terms_unordered + ) + self._one_body_terms_out_of_date = False + + def _update_two_body_terms(self): + if self._two_body_terms_out_of_date: + self._two_body_terms, _ = self._get_sorted_terms( + self._two_body_terms_unordered + ) + self._two_body_terms_out_of_date = False + + def _update_multi_body_terms(self): + if self._multi_body_terms_out_of_date: + self._multi_body_terms, _ = self._get_sorted_terms( + self._multi_body_terms_unordered + ) + self._multi_body_terms_out_of_date = False + @staticmethod - def get_sorted_terms(term_list): + def _get_sorted_terms(term_list): sorted_term_list = [] sorted_score_type_list = [] term_covered = [False] * ScoreType.n_score_types.value diff --git a/tmol/score/constraint/_utility.py b/tmol/score/constraint/_utility.py index bc82f1c38..e7023ccb9 100644 --- a/tmol/score/constraint/_utility.py +++ b/tmol/score/constraint/_utility.py @@ -2,12 +2,11 @@ import attrs from tmol.types import Tensor -from tmol.pose._packed_block_types import PackedBlockTypes -from tmol.pose._pose_stack import PoseStack +from tmol.pose import PackedBlockTypes, PoseStack def constrain_all_ca(pose_stack: PoseStack) -> PoseStack: - from tmol.pose._constraint_set import ConstraintSet + from tmol.pose import ConstraintSet from tmol.score.constraint._constraint_energy_term import ConstraintEnergyTerm constraint_set = pose_stack.constraint_set @@ -19,7 +18,7 @@ def constrain_all_ca(pose_stack: PoseStack) -> PoseStack: for block_ind in range(pose_stack.max_n_blocks): if pose_stack.is_real_block(pose_ind, block_ind): block_type = pose_stack.block_type(pose_ind, block_ind) - if "CA" not in block_type.atom_to_idx: + if "CA" not in block_type.atom_names_set: continue ca_ind = block_type.atom_to_idx["CA"] diff --git a/tmol/tests/relax/test_fast_relax.py b/tmol/tests/relax/test_fast_relax.py index 8741a1a42..139f4a0e7 100644 --- a/tmol/tests/relax/test_fast_relax.py +++ b/tmol/tests/relax/test_fast_relax.py @@ -2,7 +2,8 @@ import numpy import pytest -from tmol.relax import _default_cart_min_fn, fast_relax +from tmol.relax import default_cart_min_fn, default_kin_min_fn, fast_relax + import time from tmol.pose import ( @@ -93,6 +94,7 @@ def task_op(task): fold_forest, task_operations=[task_op], num_repeats=1, + min_fn=default_kin_min_fn, verbose=verbose, ) @@ -146,7 +148,7 @@ def task_op(task): cart_mm, fold_forest, task_operations=[task_op], - min_fn=_default_cart_min_fn, + min_fn=default_cart_min_fn, num_repeats=1, verbose=verbose, ) @@ -231,6 +233,7 @@ def task_op(task): fold_forest, task_operations=[task_op], num_repeats=1, + min_fn=default_kin_min_fn, verbose=verbose, ) @@ -305,6 +308,7 @@ def task_op(task): fold_forest, task_operations=[task_op], num_repeats=1, + min_fn=default_kin_min_fn, verbose=verbose, ) @@ -320,3 +324,46 @@ def task_op(task): print( f"Three differently-shaped PDBs relaxed; Execution time: {elapsed_time:.6f} seconds" ) + + +def test_fast_relax_with_f64(default_database, ubq_pdb, dun_sampler, torch_device): + # if torch_device == torch.device("cpu"): + # return + + pose_stack = pose_stack_from_pdb( + ubq_pdb, torch_device, residue_start=0, residue_end=76 + ).to(torch.float64) + + sfxn = get_relax_sfxn(default_database, torch_device) + + mm = MoveMap.from_pose_stack(pose_stack) + mm.move_all_jumps = True + mm.move_all_named_torsions = True + + palette = PackerPalette() + fold_forest = FoldForest.reasonable_fold_forest(pose_stack) + + def task_op(task): + task.restrict_to_repacking() + + fixed_sampler = FixedAAChiSampler() + task.add_conformer_sampler(dun_sampler) + task.add_conformer_sampler(fixed_sampler) + task.add_conformer_sampler(IncludeCurrentSampler()) + + # Now let's run fast_relax + verbose = True + new_pose_stack = fast_relax( + pose_stack, + sfxn, + palette, + mm, + fold_forest, + task_operations=[task_op], + num_repeats=1, + min_fn=default_kin_min_fn, + verbose=verbose, + ) + assert ( + new_pose_stack.coords.dtype == torch.float64 + ), "Output coords dtype should match input coords dtype"