{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Build your own model input interface\n", "\n", "[Open in Colab](https://colab.research.google.com/github/RosettaCommons/tmol/blob/main/docs/model_input_example.ipynb)\n", "\n", "Map OpenFold-style atom14 and RoseTTAFold2 prediction tensors to TMol, score them, and differentiate back to the original coordinates. The adapter functions below are editable application code using TMol's canonical construction and scoring APIs.\n", "\n", "Run cells in order. In Colab, select a GPU runtime; the shared setup installs published wheels matching PyTorch's CUDA major and downloads example predictions. No compilation is needed. Locally, follow the [installation guide](installation.md), then run from a source checkout for the fixtures. CPU and CUDA are supported. Set `TMOL_EXAMPLE_SOURCE` to another fixture checkout or `TMOL_EXAMPLE_DEVICE` to select a device explicitly.\n" ], "id": "model-inputs-00" }, { "cell_type": "code", "metadata": {}, "source": [ "import os\n", "from pathlib import Path\n", "\n", "try:\n", " import google.colab # noqa: F401\n", "except ImportError:\n", " IN_COLAB = False\n", "else:\n", " IN_COLAB = True\n", "\n", "if IN_COLAB:\n", " from urllib.request import urlretrieve\n", "\n", " urlretrieve(\n", " \"https://raw.githubusercontent.com/RosettaCommons/tmol/\"\n", " \"main/docs/tutorial/notebook_setup.py\",\n", " \"notebook_setup.py\",\n", " )\n", " import notebook_setup\n", "\n", " notebook_setup.setup_colab(\n", " [\n", " \"tmol/tests/data/openfold/openfold_ubq_and_sumo.pt\",\n", " \"tmol/tests/data/rosettafold2/ubiquitin.pt\",\n", " \"tmol/tests/data/rosettafold2/input_layout.json\",\n", " ]\n", " )\n", "\n", "repo = Path(os.environ.get(\"TMOL_EXAMPLE_SOURCE\", \".\")).resolve()\n", "os.environ.setdefault(\"OMP_NUM_THREADS\", \"2\")\n", "" ], "execution_count": null, "outputs": [], "id": "model-inputs-01" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Describe a coordinate layout once\n", "\n", "A slot count does not specify atom order or residue-token numbering. Supply ordered residue names and named slots for each residue. AtomWorks supplies common atom14/atom37 definitions; check the token order against your predictor version.\n", "\n", "This builder takes `[batch, residues, atoms, 3]` float32 tensors. `chain_id=-1` means a padded residue; NaN triplets or an observation mask mean missing atoms of a real residue. Torch indexed assignment preserves gradients from supplied coordinates. The example caches the mappings, while canonical pose construction still runs for each call. It expects the heavy atoms required for construction; it does not implement missing-sidechain packing." ], "id": "model-inputs-02" }, { "cell_type": "code", "metadata": {}, "source": [ "import torch\n", "import tmol\n", "from tmol.io import (\n", " CanonicalForm,\n", " default_canonical_ordering,\n", " default_packed_block_types,\n", " pose_stack_from_canonical_form,\n", ")\n", "from tmol.score import beta2016_score_function\n", "\n", "device = torch.device(\n", " os.environ.get(\n", " \"TMOL_EXAMPLE_DEVICE\",\n", " \"cuda\" if tmol.has_cuda() and torch.cuda.is_available() else \"cpu\",\n", " )\n", ")\n", "print(\"PyTorch:\", torch.__version__, \"device:\", device)\n", "\n", "\n", "def prepare_named_layout(sequence, chain_id, residue_names, atom_names):\n", " \"\"\"Prepare a model's fixed sequence/layout once; bind coordinates in Torch.\n", "\n", " This example uses the default protein database. General chemistry callers\n", " should use the ordering and packed types from their PoseBuildContext.\n", " Padding is explicitly identified by chain_id == -1.\n", " \"\"\"\n", " if sequence.ndim != 2 or chain_id.shape != sequence.shape:\n", " raise ValueError(\"sequence and chain_id must have shape [batch, residues]\")\n", " if sequence.dtype not in (torch.int32, torch.int64):\n", " raise TypeError(\"sequence must contain integer token IDs\")\n", " if chain_id.dtype not in (torch.int32, torch.int64):\n", " raise TypeError(\"chain_id must contain integer chain IDs\")\n", " if sequence.device != chain_id.device or torch.any(chain_id < -1):\n", " raise ValueError(\n", " \"chain IDs must share the sequence device and use -1 for padding\"\n", " )\n", " present = chain_id >= 0\n", " if torch.any(present & ((sequence < 0) | (sequence >= len(residue_names)))):\n", " raise ValueError(\"An observed residue has an unknown token ID\")\n", " sequence = sequence.clone().long()\n", " sequence[~present] = 0\n", " chain_id = chain_id.clone().to(torch.int32)\n", " ordering = default_canonical_ordering()\n", " packed_types = default_packed_block_types(sequence.device)\n", " restype_map, atom_map, real_atoms = ordering.create_src_2_tmol_mappings(\n", " residue_names, atom_names, sequence.device\n", " )\n", " restypes = restype_map[sequence].to(torch.int32)\n", " if torch.any(present & (restypes < 0)):\n", " raise ValueError(\"A residue needs a chemical definition before construction\")\n", " restypes[~present] = -1\n", " real = real_atoms[sequence] & present.unsqueeze(-1)\n", " pose, residue, slot = torch.nonzero(real, as_tuple=True)\n", " target_atom = atom_map[sequence][real]\n", " shape = (*sequence.shape, real.shape[-1], 3)\n", "\n", " def build(coords, *, observed=None):\n", " if tuple(coords.shape) != shape or coords.dtype != torch.float32:\n", " raise ValueError(f\"Expected float32 coordinates shaped {shape}\")\n", " if coords.device != sequence.device:\n", " raise ValueError(\"Coordinates and prepared mapping must share a device\")\n", " source = coords[pose, residue, slot]\n", " if observed is not None:\n", " if observed.shape != coords.shape[:-1] or observed.dtype != torch.bool:\n", " raise ValueError(\n", " \"observed must be a boolean mask over coordinate slots\"\n", " )\n", " source = torch.where(\n", " observed[pose, residue, slot, None], source, float(\"nan\")\n", " )\n", " valid = torch.isfinite(source).all(-1) | torch.isnan(source).all(-1)\n", " if not bool(valid.all()):\n", " raise ValueError(\n", " \"Supply finite coordinate triplets or all-NaN missing atoms\"\n", " )\n", " canonical = coords.new_full(\n", " (*sequence.shape, ordering.max_n_canonical_atoms, 3), float(\"nan\")\n", " )\n", " canonical[pose, residue, target_atom] = source\n", " form = CanonicalForm(\n", " chain_id=chain_id,\n", " res_types=restypes,\n", " coords=canonical,\n", " chain_labels=None,\n", " res_labels=None,\n", " residue_insertion_codes=None,\n", " atom_occupancy=None,\n", " atom_b_factor=None,\n", " disulfides=None,\n", " res_not_connected=None,\n", " )\n", " return pose_stack_from_canonical_form(ordering, packed_types, *form)\n", "\n", " return build\n" ], "execution_count": null, "outputs": [], "id": "model-inputs-03" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. OpenFold-style atom14 output\n", "\n", "Read the final structure-module coordinates, sequence tokens, chain IDs and atom-existence mask from a saved prediction. These dictionary keys describe this example's output format. Adapt them to your predictor; no OpenFold package or tmol-specific OpenFold interface is needed." ], "id": "model-inputs-04" }, { "cell_type": "code", "metadata": {}, "source": [ "def openfold_example(prediction):\n", " \"\"\"Adapt this predictor's token IDs, final atom14 coordinates and chain mask.\"\"\"\n", " from atomworks.ml.encoding_definitions import AF2_ATOM14_ENCODING\n", "\n", " # This token order is part of the predictor's contract, independent of shape.\n", " names = list(AF2_ATOM14_ENCODING.token_atoms)[:20]\n", " builder = prepare_named_layout(\n", " prediction[\"aatype\"],\n", " prediction[\"chain_index\"],\n", " names,\n", " AF2_ATOM14_ENCODING.token_atoms,\n", " )\n", " return builder(\n", " prediction[\"positions\"][-1], observed=prediction[\"atom14_atom_exists\"].bool()\n", " )\n" ], "execution_count": null, "outputs": [], "id": "model-inputs-05" }, { "cell_type": "code", "metadata": {}, "source": [ "from tmol.io import canonical_form_from_pose_stack\n", "\n", "\n", "def check_and_score(pose, input_coords, expected_ca):\n", " \"\"\"Check identity and differentiate the score into supplied coordinates.\"\"\"\n", " ordering = default_canonical_ordering()\n", " canonical = canonical_form_from_pose_stack(ordering, pose)\n", " ca_slot = torch.tensor(\n", " [\n", " ordering.restypes_atom_index_mapping[name].get(\"CA\", -1)\n", " for name in ordering.restype_io_equiv_classes\n", " ],\n", " device=device,\n", " )\n", " batch, residue = torch.nonzero(canonical.res_types >= 0, as_tuple=True)\n", " observed_ca = canonical.coords[\n", " batch, residue, ca_slot[canonical.res_types[batch, residue]]\n", " ]\n", " torch.testing.assert_close(observed_ca, expected_ca[batch, residue], rtol=0, atol=0)\n", " assert torch.isfinite(pose.coords[pose.real_atoms]).all()\n", " module = beta2016_score_function(device).render_whole_pose_scoring_module(pose)\n", " energy = module(pose.coords).sum()\n", " (gradient,) = torch.autograd.grad(energy, input_coords)\n", " assert torch.isfinite(energy) and torch.isfinite(gradient).all()\n", " assert torch.any(gradient != 0)\n", " return float(energy.detach()), gradient\n", "\n", "\n", "prediction = torch.load(\n", " repo / \"tmol/tests/data/openfold/openfold_ubq_and_sumo.pt\",\n", " map_location=device,\n", " weights_only=True,\n", ")\n", "openfold_coords = prediction[\"positions\"].detach().clone().requires_grad_()\n", "prediction[\"positions\"] = openfold_coords\n", "openfold_pose = openfold_example(prediction)\n", "assert openfold_pose.n_poses == 2\n", "openfold_energy, openfold_gradient = check_and_score(\n", " openfold_pose, openfold_coords, openfold_coords[-1, :, :, 1]\n", ")\n", "print(\"OpenFold-style batch energy:\", openfold_energy)\n", "print(\"Coordinate-gradient norm:\", float(openfold_gradient.norm()))\n" ], "execution_count": null, "outputs": [], "id": "model-inputs-06" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. RF2-style output and hydrogen policy\n", "\n", "Use `num2aa` and `aa2long` from the RF2 version that produced the tensor. The legacy prediction below has 27 slots; other RF2-family versions can use different layouts. The saved layout is example data, not a library interface.\n", "\n", "Choose `hydrogens=\"preserve\"` to keep supplied nonterminal H coordinates, or `\"rebuild\"` to mark all supplied H missing. Both rebuild a generic N-terminal amide H using the appropriate terminal model. Rebuilt slots have zero input gradient. Supplied coordinates keep their gradient connection." ], "id": "model-inputs-07" }, { "cell_type": "code", "metadata": {}, "source": [ "def rf2_example(prediction, residue_names, atom_names, *, hydrogens):\n", " \"\"\"Use num2aa/aa2long from the RF2 version that produced this prediction.\n", "\n", " ``hydrogens='preserve'`` retains supplied nonterminal H coordinates;\n", " ``'rebuild'`` marks all supplied H as missing. The generic amide H at each\n", " chain's N terminus is always rebuilt as the appropriate terminal hydrogens.\n", " \"\"\"\n", " if hydrogens not in (\"preserve\", \"rebuild\"):\n", " raise ValueError(\"Choose hydrogens='preserve' or 'rebuild'\")\n", " sequence = prediction[\"seq\"].unsqueeze(0)\n", " coords = prediction[\"xyz\"].unsqueeze(0)\n", " lengths = prediction[\"chainlens\"]\n", " if any(length <= 0 for length in lengths) or sum(lengths) != sequence.shape[1]:\n", " raise ValueError(\"Positive chain lengths must cover the entire sequence\")\n", " chain_id = torch.repeat_interleave(\n", " torch.arange(len(lengths), device=coords.device),\n", " torch.tensor(lengths, device=coords.device),\n", " ).unsqueeze(0)\n", " names = {\n", " name: [atom.strip() if atom else \"\" for atom in row]\n", " for name, row in zip(residue_names, atom_names)\n", " }\n", " builder = prepare_named_layout(sequence, chain_id, residue_names, names)\n", " hydrogen_slots = torch.tensor(\n", " [\n", " [atom.lstrip(\"123\").startswith(\"H\") for atom in names[name]]\n", " for name in residue_names\n", " ],\n", " device=coords.device,\n", " )\n", " amide_h = torch.tensor(\n", " [[atom == \"H\" for atom in names[name]] for name in residue_names],\n", " device=coords.device,\n", " )\n", " nterm = torch.ones_like(chain_id, dtype=torch.bool)\n", " nterm[:, 1:] = chain_id[:, 1:] != chain_id[:, :-1]\n", " suppressed = amide_h[sequence] & nterm.unsqueeze(-1)\n", " if hydrogens == \"rebuild\":\n", " suppressed |= hydrogen_slots[sequence]\n", " return builder(coords, observed=~suppressed)\n" ], "execution_count": null, "outputs": [], "id": "model-inputs-08" }, { "cell_type": "code", "metadata": {}, "source": [ "import json\n", "\n", "rf2_prediction = torch.load(\n", " repo / \"tmol/tests/data/rosettafold2/ubiquitin.pt\",\n", " map_location=device,\n", " weights_only=True,\n", ")\n", "layout = json.loads(\n", " (repo / \"tmol/tests/data/rosettafold2/input_layout.json\").read_text()\n", ")\n", "rf2_results = {}\n", "for policy in (\"preserve\", \"rebuild\"):\n", " coords = rf2_prediction[\"xyz\"].detach().clone().requires_grad_()\n", " rf2_prediction[\"xyz\"] = coords\n", " pose = rf2_example(\n", " rf2_prediction, layout[\"num2aa\"], layout[\"aa2long\"], hydrogens=policy\n", " )\n", " assert pose.n_poses == 1\n", " energy, gradient = check_and_score(pose, coords, coords[None, :, 1])\n", " if policy == \"rebuild\":\n", " # These saved legacy tables place hydrogens at slots 14 onward.\n", " assert torch.count_nonzero(gradient[:, 14:]) == 0\n", " rf2_results[policy] = {\"energy\": energy, \"gradient_norm\": float(gradient.norm())}\n", "print(rf2_results)\n" ], "execution_count": null, "outputs": [], "id": "model-inputs-09" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. Scoring and repeated guidance\n", "\n", "For a fixed sequence and topology, reuse the named mapping and score module. Bind new coordinate tensors entirely in Torch; do not detach or convert them to NumPy on the differentiable path. AtomWorks file parsing and chemical preparation belong before this loop. NaN completion preserves atom identity and does not itself impute finite coordinates.\n", "\n", "Discrete protonation choices, chemical topology changes and rotamer selection are outside this coordinate-gradient contract. A prepared-topology builder would also cache final topology; this small example only caches slot mappings." ], "id": "model-inputs-10" }, { "cell_type": "code", "metadata": {}, "source": [ "from atomworks.ml.encoding_definitions import AF2_ATOM14_ENCODING\n", "\n", "names = list(AF2_ATOM14_ENCODING.token_atoms)[:20]\n", "bind = prepare_named_layout(\n", " prediction[\"aatype\"],\n", " prediction[\"chain_index\"],\n", " names,\n", " AF2_ATOM14_ENCODING.token_atoms,\n", ")\n", "mask = prediction[\"atom14_atom_exists\"].bool()\n", "coordinates = openfold_coords[-1].detach().clone()\n", "score_module = None\n", "guidance_energies = []\n", "for step in range(3):\n", " coordinates.requires_grad_()\n", " pose = bind(coordinates, observed=mask)\n", " if score_module is None:\n", " score_module = beta2016_score_function(device).render_whole_pose_scoring_module(\n", " pose\n", " )\n", " topology = pose.block_type_ind.clone()\n", " torch.testing.assert_close(pose.block_type_ind, topology, rtol=0, atol=0)\n", " energy = score_module(pose.coords).sum()\n", " (gradient,) = torch.autograd.grad(energy, coordinates)\n", " assert torch.isfinite(energy) and torch.isfinite(gradient).all()\n", " guidance_energies.append(float(energy.detach()))\n", " # Stop the tape between iterations, after obtaining the current gradient.\n", " # An application can instead use this gradient in its diffusion/search update.\n", " with torch.no_grad():\n", " coordinates = coordinates - 0.001 * gradient / gradient.norm().clamp_min(1.0)\n", "print(\"Three guidance evaluations:\", guidance_energies)\n" ], "execution_count": null, "outputs": [], "id": "model-inputs-11" } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.12" }, "colab": { "name": "model_input_example.ipynb", "provenance": [] } }, "nbformat": 4, "nbformat_minor": 5 }