diff --git a/notebooks/tmol_how_to_guide.ipynb b/notebooks/tmol_how_to_guide.ipynb new file mode 100644 index 000000000..20f8c4417 --- /dev/null +++ b/notebooks/tmol_how_to_guide.ipynb @@ -0,0 +1,1288 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "markdown", + "source": [ + "# Contents:\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 single pose PoseStack from a .mmcif\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 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" + ], + "metadata": { + "id": "0NSgz0L8CeFA" + } + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "nU1-4FoBbApM", + "outputId": "68e50ec9-25eb-450e-f1b2-882c310165c2", + "collapsed": true + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "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": [ + "# 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", + "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" + ], + "metadata": { + "id": "y3f4dcA0mV8S", + "outputId": "29bc90e1-4fad-4e2a-a724-58d57032e1a2", + "colab": { + "base_uri": "https://localhost:8080/" + }, + "collapsed": true + }, + "execution_count": 4, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "--2026-07-16 19:46:02-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1ubq.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|: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.002s \n", + "\n", + "2026-07-16 19:46:03 (56.7 MB/s) - ‘1ubq.pdb’ saved [99763/99763]\n", + "\n", + "--2026-07-16 19:46:03-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1s78.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|: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.006s \n", + "\n", + "2026-07-16 19:46:03 (103 MB/s) - ‘1s78.pdb’ saved [615762/615762]\n", + "\n", + "--2026-07-16 19:46:03-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/pdb/1qys.pdb\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|: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.002s \n", + "\n", + "2026-07-16 19:46:03 (52.8 MB/s) - ‘1qys.pdb’ saved [119719/119719]\n", + "\n", + "--2026-07-16 19:46:03-- https://raw.githubusercontent.com/uw-ipd/tmol/refs/heads/master/tmol/tests/data/cif/1BL8.cif\n", + "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n", + "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|: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.004s \n", + "\n", + "2026-07-16 19:46:04 (86.6 MB/s) - ‘1BL8.cif’ saved [330009/330009]\n", + "\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "!head 1ubq.pdb" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "A_hmNzYjZLH0", + "outputId": "3477568c-0a01-4d71-dad9-6385092be812" + }, + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "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" + ] + } + ] + }, + { + "cell_type": "code", + "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())\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)" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "Mi9gshJmbVVC", + "outputId": "6a3b6001-415b-4570-a8b3-2b56561a709d" + }, + "execution_count": 3, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "tensor([231.5137], device='cuda:0', grad_fn=)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "Fundamentals" + ], + "metadata": { + "id": "Jp_iT-4DjP_x" + } + }, + { + "cell_type": "code", + "source": [ + "# - Initialize the default ParameterDatabase\n", + "param_db = tmol.ParameterDatabase.get_default()" + ], + "metadata": { + "id": "lBCg3XA8jMq8" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Load alternate block types\n", + "# TO DO" + ], + "metadata": { + "id": "9nyxfNgojaSB" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Initialize a PackedBlockTypes object from the default database\n", + "pbd = tmol.default_packed_block_types()" + ], + "metadata": { + "id": "xClcO442jcaO" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Initialize a PackedBlockTypes object from a custom database\n", + "# TO DO" + ], + "metadata": { + "id": "PyhxBHq9pLMx" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a default CanonicalOrdering object\n", + "canonical_ordering = tmol.default_canonical_ordering()" + ], + "metadata": { + "id": "OO32waIlp4-5" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a CanonicalOrdering from a ParameterDatabase\n", + "# TO DO" + ], + "metadata": { + "id": "L-uVsXPLjiDC" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a PackedBlockTypes object from a subset of the block types in a ParameterDatabase\n", + "# TO DO" + ], + "metadata": { + "id": "oXpe5H0Rjmfr" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Input" + ], + "metadata": { + "id": "42PMj1JkkTaD" + } + }, + { + "cell_type": "code", + "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" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "k_2XAIzBjqKn", + "outputId": "45d58612-7efc-4968-d77c-783e7404a70d" + }, + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "pose_1ubq 76\n" + ] + } + ] + }, + { + "cell_type": "code", + "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" + ], + "metadata": { + "id": "496P7xj4OCcr" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Initialize a single pose PoseStack from a PDB file using biotite\n", + "# TO DO" + ], + "metadata": { + "id": "N8RfPFCWYkfy" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "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", + " print(\"was atom array stack\")\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)" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "ysGCJuh8kXRc", + "outputId": "ea818011-19e4-4b37-c638-34ef10bef4d7" + }, + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stderr", + "text": [ + "WARNING:tmol.io.pose_stack_from_biotite:Unrecognized 3lc K\n", + "/usr/local/lib/python3.12/dist-packages/tmol/pack/rotamer/build_rotamers.py:920: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).\n", + " conf_dofs_kto[1:] = torch.tensor(\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# - Initialize a single pose PoseStack from an OpenFold set of tensors\n", + "# TO DO\n" + ], + "metadata": { + "id": "bNOuYq4QkX7A" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "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" + ], + "metadata": { + "id": "CgVHFNkVkYUZ" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Load a bunch of PoseStacks from different sources and then concatenate them to a single PoseStack\n", + "# TEMP: Add PoseStackBuilder to API\n", + "from tmol.pose.pose_stack_builder import PoseStackBuilder\n", + "pose_stack_3 = PoseStackBuilder.from_poses([pose_1ubq, pose_1s78, pose_1bl8], device=device)" + ], + "metadata": { + "id": "HX9jTnCrkYry" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "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" + ], + "metadata": { + "id": "IpP5VgmbkZBz" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Add custom residue type\n", + "# TO DO\n" + ], + "metadata": { + "id": "zIacYunXkZX4" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Set the scoring parameters for a new residue type (elec, etc)\n", + "# TO DO\n" + ], + "metadata": { + "id": "B9BJmHWCkZxQ" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Build an extended pose from sequence\n", + "# TO DO" + ], + "metadata": { + "id": "MTd5eVotkaJ9" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Output" + ], + "metadata": { + "id": "R1LJAlYJktGE" + } + }, + { + "cell_type": "code", + "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\")" + ], + "metadata": { + "id": "WAYlCkZ8kNh_" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "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" + ], + "metadata": { + "id": "a58GBcGrk0WW" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Write a multi-pose stack out to separate PDB files\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", + "for i in range(10):\n", + " assert os.path.isfile(f\"1ubq_{i:04}.pdb\")" + ], + "metadata": { + "id": "sPGfDWNrk0zL" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Write a single pose PoseStack to an .mmcif file\n", + "# TO DO\n" + ], + "metadata": { + "id": "wYa6C8EGk1GT" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Write a rotamer set out as a multi-model PDB\n", + "# TO DO\n" + ], + "metadata": { + "id": "gcNdsybLk1dR" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Kinematics" + ], + "metadata": { + "id": "vucunYYYk2OX" + } + }, + { + "cell_type": "code", + "source": [ + "# - Create an N->C fold tree (a fold *forest* in tmol) for a PoseStack\n", + "import numpy\n", + "\n", + "ff1_1ubq = tmol.FoldForest.reasonable_fold_forest(pose_1ubq)\n", + "\n", + "# also valid: specify a root-jump to residue 0, and a polymer\n", + "# edge from 0 to the last residue in the pose.\n", + "edges = numpy.zeros((1, 2, 3), 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)" + ], + "metadata": { + "id": "PYwjsJDvk4Ka" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a set of N->C fold trees for a multi-pose PoseStack\n", + "ff1_ps3 = tmol.FoldForest.reasonable_fold_forest(pose_stack_3)\n", + "\n", + "edges = numpy.zeros((3, 2, 3), 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] = -1\n", + "edges[:, 1, 2] = n_res\n", + "ff2_ps3 = tmol.FoldForest.from_edges" + ], + "metadata": { + "id": "LVS2vrVRZGbS" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a simple fold tree for a multi-chain PoseStack\n", + "# TO DO\n" + ], + "metadata": { + "id": "OBAw1Lw6lKKP" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a simple fold tree for a PoseStack with missing residues\n", + "# TO DO" + ], + "metadata": { + "id": "e7m9yUztlKmg" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a dandelion fold tree for a PoseStack\n", + "# TO DO" + ], + "metadata": { + "id": "iblPqPjZlK8Z" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a MoveMap that enables minimization for named torsions\n", + "# TO DO" + ], + "metadata": { + "id": "pUBbckP7lLRM" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a MoveMap that enables backbone minimization for some residues but not all\n", + "# TO DO" + ], + "metadata": { + "id": "fJUxnRqYlLnW" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Apply a perturbation to the rigid-body DOFs between two chains\n", + "# TO DO" + ], + "metadata": { + "id": "ZTRWFlPslL9b" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Assign dihedral values to all the residues in a PoseStack and calculate the coordinates\n", + "# TO DO" + ], + "metadata": { + "id": "nM9g3oIwlMTy" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Scoring" + ], + "metadata": { + "id": "klTMFJK6lQD1" + } + }, + { + "cell_type": "code", + "source": [ + "# - Create the default score function\n", + "# TO DO" + ], + "metadata": { + "id": "EnQIxSW4lWjZ" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create the default score function from a custom Database\n", + "# TO DO" + ], + "metadata": { + "id": "Yc8a4BnklZoh" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create the soft-rep version of the score function\n", + "# TO DO" + ], + "metadata": { + "id": "5qErnZZSlbpA" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create an empty score function\n", + "# TO DO" + ], + "metadata": { + "id": "jLR7q9nXlc4H" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Turn on a few terms in a score function\n", + "# TO DO" + ], + "metadata": { + "id": "r4vz-xXpmYV1" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Turn off a term in a score function\n", + "# TO DO" + ], + "metadata": { + "id": "yMgEZFkQmYt_" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Score a PoseStack\n", + "# TO DO" + ], + "metadata": { + "id": "EV89BMNEmZC-" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Score a PoseStack and back-propagate through the coordinates\n", + "# TO DO" + ], + "metadata": { + "id": "XWLjbPsxmZYR" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Score a PoseStack and return per-residue weighted energies\n", + "# TO DO" + ], + "metadata": { + "id": "osRew_fxmZzF" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Score a PoseStack and return per-residue weighted energies, weight them according to some principle, and then back-propagate the total energy.\n", + "# TO DO" + ], + "metadata": { + "id": "dYzEVvc7maRK" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Add constraints to a PoseStack\n", + "# TO DO" + ], + "metadata": { + "id": "-NW-sfNJmapX" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Add the same constraints to all the Poses in a PoseStack\n", + "# TO DO" + ], + "metadata": { + "id": "swFqps8PmbCS" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Add coordinate constraints to the current coordinates\n", + "# TO DO" + ], + "metadata": { + "id": "A26oZ-8ImbZf" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Alter the parameters for the cart-bonded energy function & rescore\n", + "# TO DO" + ], + "metadata": { + "id": "0tOaC2Trmbvc" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Optimization" + ], + "metadata": { + "id": "dIPA0OfBmjhh" + } + }, + { + "cell_type": "code", + "source": [ + "# - Create a DunbrackSampler\n", + "# TO DO" + ], + "metadata": { + "id": "GffUmGNCmlIk" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Add hydrogens\n", + "# TO DO" + ], + "metadata": { + "id": "nDVbq3pDmqpf" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Fill in side chains for a model that lacks them\n", + "# TO DO" + ], + "metadata": { + "id": "e3a4AOs8mrDX" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Perform fixed-sequence side-chain optimization: repack\n", + "# TO DO" + ], + "metadata": { + "id": "Tgj8DwX8mrZ6" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Repack with extra rotamers\n", + "# TO DO" + ], + "metadata": { + "id": "zt23WbYtmru3" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Add hydrogens and back-propagate\n", + "# TO DO" + ], + "metadata": { + "id": "YmAY91mBmsCy" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Create a new PackerPalette subclass to handle logic of new block types\n", + "# TO DO" + ], + "metadata": { + "id": "awPe0TtYmsaR" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Run minimization in double precision\n", + "# TO DO" + ], + "metadata": { + "id": "RTm3LjhAmsv4" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Perform cartesian minimization\n", + "# TO DO" + ], + "metadata": { + "id": "Z2UFgON6mtFt" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Perform kinematic minimization\n", + "# TO DO" + ], + "metadata": { + "id": "_Dwt2_lpmtbM" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Relax a PoseStack w/ kinematic minimization\n", + "# TO DO" + ], + "metadata": { + "id": "oE7_fHfEmtwM" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Relax a PoseStack w/ cartesian minimization\n", + "# TO DO" + ], + "metadata": { + "id": "9ZUjD_eYmuFw" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Relax structures generated one-at-a-time in batch format\n", + "# TO DO" + ], + "metadata": { + "id": "2aUJd4X5mubN" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Idealize a structure from the PDB\n", + "# TO DO" + ], + "metadata": { + "id": "Mkj2U1Namuw7" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Idealize a structure from a dandelion\n", + "# TO DO" + ], + "metadata": { + "id": "aWFR74xZmvFu" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# - Idealize just the backbone of a dandelion\n", + "# TO DO" + ], + "metadata": { + "id": "NDZyRGtTmva3" + }, + "execution_count": null, + "outputs": [] + } + ] +} \ No newline at end of file diff --git a/tmol/__init__.py b/tmol/__init__.py index e31abfd22..ff7a61a58 100644 --- a/tmol/__init__.py +++ b/tmol/__init__.py @@ -38,6 +38,9 @@ def include_paths(): from tmol.io.pose_stack_construction import ( pose_stack_from_canonical_form, ) +from tmol.io.pose_stack_from_biotite import ( + pose_stack_from_biotite, +) from tmol.io.pose_stack_from_openfold import ( canonical_form_from_openfold, canonical_ordering_for_openfold, @@ -65,6 +68,7 @@ def include_paths(): ) from tmol.pose.packed_block_types import PackedBlockTypes from tmol.pose.pose_stack import PoseStack +from tmol.pose.pose_stack_builder import PoseStackBuilder from tmol.score import beta2016_score_function from tmol.score.score_function import ScoreFunction from tmol.score.score_types import ScoreType diff --git a/tmol/io/pose_stack_from_biotite.py b/tmol/io/pose_stack_from_biotite.py index 1a74a6aae..f5690d808 100644 --- a/tmol/io/pose_stack_from_biotite.py +++ b/tmol/io/pose_stack_from_biotite.py @@ -17,8 +17,6 @@ from tmol.pose.pdb_info import DEFAULT_ATOM_B_FACTOR, DEFAULT_ATOM_OCCUPANCY from tmol.utility.biotite_util import get_all_residue_positions -from tmol import beta2016_score_function - logger = logging.getLogger(__name__) @@ -211,6 +209,7 @@ def pose_stack_from_biotite( from tmol.pack.rotamer.dunbrack.dunbrack_chi_sampler import ( create_dunbrack_sampler_from_database, ) + from tmol import beta2016_score_function if context is not None: if param_db is not None: diff --git a/tmol/pose/pose_stack.py b/tmol/pose/pose_stack.py index 6f9546d3b..cb857337c 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 +from typing import Optional, Union from tmol.types.torch import Tensor from tmol.chemical.restypes import RefinedResidueType @@ -276,6 +276,12 @@ def split(self, index) -> "PoseStack": device=self.device, ) + 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/fast_relax.py b/tmol/relax/fast_relax.py index 74c5700a1..554dd3203 100644 --- a/tmol/relax/fast_relax.py +++ b/tmol/relax/fast_relax.py @@ -166,7 +166,9 @@ def fast_relax( 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. @@ -340,6 +342,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() @@ -357,7 +364,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: diff --git a/tmol/tests/relax/test_fast_relax.py b/tmol/tests/relax/test_fast_relax.py index 0ec536c9b..e986c13a0 100644 --- a/tmol/tests/relax/test_fast_relax.py +++ b/tmol/tests/relax/test_fast_relax.py @@ -2,7 +2,7 @@ import numpy import pytest -from tmol.relax.fast_relax import _default_cart_min_fn, fast_relax +from tmol.relax.fast_relax import _default_cart_min_fn, _default_kin_min_fn, fast_relax import time from tmol.pose.pose_stack import PoseStack @@ -84,6 +84,7 @@ def task_op(task): fold_forest, task_operations=[task_op], num_repeats=1, + min_fn=_default_kin_min_fn, verbose=verbose, ) @@ -220,6 +221,7 @@ def task_op(task): fold_forest, task_operations=[task_op], num_repeats=1, + min_fn=_default_kin_min_fn, verbose=verbose, ) @@ -293,6 +295,7 @@ def task_op(task): fold_forest, task_operations=[task_op], num_repeats=1, + min_fn=_default_kin_min_fn, verbose=verbose, ) @@ -308,3 +311,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"