Part 2: Creating a Transform Pipeline#

Table of Contents#

Introduction
Prerequisites
The Goal of the Pipeline
Wiring Up the Dataset and Loader
Building the Transform Pipeline
- CropToPocket
- FeaturizeForDocking
The Complete Smoke Test
What Next?
Glossary

Introduction#

This is the second tutorial in the How to Build a Model with AtomWorks series. In Part 1 you cleaned the PDB metadata and saved train/validation/test parquet splits.

In this installment, you will learn how to transform each parquet file into tensors that can be used in a machine learning model.

More specifically, you will learn how to:

create custom transforms
create a transform pipeline
utilize the PandasDataset class to go from parquets to trainable tensors

By the end of this part of the tutorial series you will have two python files:

transforms.py: Definition of custom transforms
smoke_test.py: Testing the pipeline by loading a single example

Important

This tutorial will walk you through the creation of these two scripts.

For those who want to use the tutorial text as structure and hints to write your own code, the solutions are hidden in collapsible cells. The new code in cells with repeated code is highlighted.

If you would like to see the full scripts, they are provided in the tutorial files.

Prerequisites#

Before starting this part it is assumed that you have:

Completed Part 1 and saved splits/train.parquet, splits/validation.parquet, and splits/test.parquet.
A working installation of AtomWorks, including the ml side.
A (partial) PDB mirror set up so the loader can find the structure file for each row (see Data Mirrors).
Familiarity with biotite.structure.AtomArray, NumPy, and SciPy.

Note

The transforms you write here operate on a biotite AtomArray. If you have not worked with AtomArray objects before, it helps to skim the Biotite structure documentation so the coordinate and annotation access patterns below feel familiar.

The Goal of the Pipeline#

Our task is pose generation: given a protein pocket and a small-molecule ligand (from the parquet files we generated in Part 1), predict plausible bound cartesian coordinates for every ligand atom. To get there, each raw structure needs to be reduced to just the binding pocket and converted into tensors.

We will create a transform pipeline that can take protein/ligand pairs from the parquet files we created in Part 1, reduce the information to just the binding pocket, and convert them into tensors that can be used in the training of our machine learning model.

We will apply four transforms in order, two that already exist in AtomWorks and two that we will write ourselves:

RemoveHydrogens (exists already): drops hydrogen atoms.
RemoveUnresolvedAtoms (exists already): drops atoms with occupancy below the default threshold of 0.5.
CropToPocket (new): spatially crops the structure to the atoms within a radius of the ligand.
FeaturizeForDocking (new): converts the AtomArray into the tensors our model needs.

A fifth transform, ConvertToTorch, converts the NumPy features into torch tensors. We add it in Part 4 when we build the training loop; for now we work in NumPy so the outputs are easy to inspect.

Wiring Up the Dataset and Loader#

Before writing any transforms, let’s confirm we can load a single example. AtomWorks provides two pieces we need:

PandasDataset wraps a parquet/DataFrame and applies a loader and transform to each row.
create_structure_loader() returns a picklable loader that parses the CIF file for each row and maps metadata columns to example keys, so we know which PN units the interface of interest involves.

Read the documentation for both (linked above) before using them. You can search the API docs, or use help() at a Python prompt. There is also more detail on PandasDataset in the Dataset Exploration and Management in AtomWorks example.

Let’s create our PandasDataset. We need to supply it with

data: let’s use train.parquet
name: however you want to label this structure, in the tutorial we’ll call it docking_train
id_column: The column that has the unique identifier for each of piece of data, example_id
loader: we will use create_structure_loader here

As a quick check that everything fits together, write a script (smoke_test.py) that reads train.parquet, builds a PandasDataset from it, and uses create_structure_loader as the loader.

Now load one example and inspect it. This tells us what keys the loader attaches and confirms the parquet, the loader, and your PDB mirror are all talking to each other.

You should see that the example type is <class 'atomworks.ml.transforms.base.TransformedDict'>.

Among the keys you should see atom_array, query_pn_unit_iids, and query_is_polymer. These are the pieces of information that AtomWorks has pulled from the parquet files that we will use in our transforms pipeline:

atom_array is the AtomArray for this structure. Every atom carries a pn_unit_iid annotation (atom_array.pn_unit_iid), tagging it with the PN unit it belongs to, using the same string format as query_pn_unit_iids below.
query_pn_unit_iids is a list of PN unit IID strings, one per column mapped to query_pn_unit_iids. Each IID has the form "{chain_id}_{transformation_id}" (e.g. "A_1"), or a comma-joined list of such tokens for a covalently-linked, multi-chain unit (e.g. "G_1,R_1"). Since we asked for two columns (pn_unit_1_iid and pn_unit_2_iid), this list has two entries: the two sides of the interface we are studying.
query_is_polymer is a list of polymer flags, in the same order as query_pn_unit_iids. Part 1 keeps exactly one polymer side, so CropToPocket can use these flags to work out which side is the protein and which is the ligand.

Tip

The loader also attaches other keys to example, such as example_id, metadata, and ligand_info, however the three listed above are the ones this tutorial’s transforms rely on.

Building the Transform Pipeline#

We can start creating our pipeline by adding to the code in our smoke_test.py file. We will use the Compose class to store the transforms, in an object we’ll call pipe, in the order we want them to be applied and then add that object to our dataset.

Let’s start with the two transforms that are native to AtomWorks, RemoveHydrogens and RemoveUnresolvedAtoms:

Creating Custom Transforms#

Now we need the two transforms that do not yet exist: CropToPocket and FeaturizeForDocking. The Creating Custom Transforms: Ligand Pocket Conditioning example is a great place to start.

Tip

Write your transforms in a separate transforms.py file. Each transform is a class whose forward() method calls a standalone function of the same name in snake_case. Keeping the logic in a standalone function makes it easy to test and reuse outside the transform machinery.

`CropToPocket`#

The CropToPocket transform will take the protein-ligand interfaces from our dataset and crop them to just the pocket. We will build CropToPocket up incrementally, confirming it plugs into Compose at each stage.

Start with a Shell#

Begin with a class that has an __init__(radius=10.0) and a forward() that just returns the input data unchanged. Confirm it plugs into Compose().

Input Validation#

We need to add input validation to our class, add a check_input() method. In check_input(), check that data contains atom_array, query_pn_unit_iids, and query_is_polymer, and that the atom array has a pn_unit_iid annotation.

Write the `crop_to_pocket` Function#

Add a standalone crop_to_pocket() function that forward() calls. It should take atom_array, query_pn_unit_iids, query_is_polymer, and radius as inputs. For now, just return the AtomArray.

Each interface has two query PN units, a ligand and a protein (polymer). Add to the crop_to_pocket function to use query_is_polymer to assign the IIDs to ligand_iid and protein_iid. Exactly one side must be a polymer; raise an error if the flags are ambiguous.

Create boolean masks to make sure that the PN unit actually contains atoms whose identity match the ligand_iid or protein_iid. Raise a clear error if either side is empty. Failing early with a descriptive message makes debugging bad structures much easier later.

Add the Spatial Crop with a KD-tree#

Now for the actual crop. We keep every ligand atom plus every protein atom within a given radius (in angstroms) of any ligand atom.

Note

A KD-tree (k-dimensional tree) is a data structure that partitions points in space so you can answer “which points are within radius r of this query point?” efficiently, without comparing every pair of atoms. SciPy’s cKDTree.query_ball_point returns, for each ligand atom, the indices of all protein atoms within radius.

Confirm the crop runs in your test. The crop keeps individual nearby protein atoms; it does not expand the selection to whole residues:

Add the `is_ligand` annotation#

The featurizer (and later analysis) needs to know which cropped atoms are ligand atoms. Annotate the cropped AtomArray with a boolean is_ligand array.

Important

This annotation is easy to forget: it is computed from ligand_mask[keep], but you must actually attach it to the array with set_annotation. If you skip this line, FeaturizeForDocking will fail because atom_array.is_ligand will not exist.

Quick check in your test:

Declare which transforms must run first#

CropToPocket assumes hydrogens and unresolved atoms are already gone. Declare this ordering requirement so AtomWorks can enforce it.

At this point you can assemble a full smoke_test.py:

`FeaturizeForDocking`#

We follow the same incremental procedure. FeaturizeForDocking reads the cropped AtomArray and returns a dictionary of tensors that we merge into the example.

Start with a Shell#

Declare the dependency on `CropToPocket`#

FeaturizeForDocking depends on the is_ligand annotation created by CropToPocket, so declare it as a required previous transform and validate the annotation in check_input().

Write the `featurize_for_docking` Function#

Start the function by returning an empty dictionary.

Also update the forward() function. Have it update the data dictionary with the dictionary features the featurize_for_docking function will eventually create.

Add the `is_ligand` Mask and Target Coordinates#

target_coords are the ground-truth coordinates the model must reproduce.

Add `input_coords` with the ligand zeroed out#

This is the heart of the task. We keep the real pocket coordinates but zero out the ligand coordinates: the model must learn to place the ligand.

Add `atomic_numbers`#

Encode each atom’s element as an integer atomic number using the AtomWorks lookup table. These integers will feed an embedding layer in the model.

Add Bonds Via an `edge_index`#

edge_index is the bond graph in COO (coordinate) format: a [2, <number of bonds>] integer array where each column is a bonded pair of atoms.

An empty bond list produces an edge_index with shape [2, 0]. Check the edge graph in your smoke test:

The final `FeaturizeForDocking` class#

The Complete Smoke Test#

Your final smoke_test.py should build the full pipeline and validate every array. Run it from the repository root with python docs/how_to_build_a_model/scripts/smoke_test.py splits/train.parquet. Use --index to select a row whose structure exists in a partial mirror:

Click to see the complete file

"""Load one Part 1 example through the Part 2 transform pipeline."""

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from transforms import CropToPocket, FeaturizeForDocking

from atomworks.ml.datasets import PandasDataset
from atomworks.ml.datasets.loaders import create_structure_loader
from atomworks.ml.transforms.base import Compose
from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms


def build_dataset(split_path: Path) -> PandasDataset:
    """Build the tutorial dataset from a Part 1 Parquet split."""
    pipeline = Compose(
        [
            RemoveHydrogens(),
            RemoveUnresolvedAtoms(),
            CropToPocket(radius=10.0),
            FeaturizeForDocking(),
        ]
    )
    return PandasDataset(
        data=pd.read_parquet(split_path),
        name="docking_train",
        id_column="example_id",
        loader=create_structure_loader(
            altloc_seed_colname="altloc_seed",
            column_mapping={
                "query_pn_unit_iids": ["pn_unit_1_iid", "pn_unit_2_iid"],
                "query_is_polymer": ["pn_unit_1_is_polymer", "pn_unit_2_is_polymer"],
            },
        ),
        transform=pipeline,
    )


def main() -> None:
    """Print shapes and check the basic feature contract for one example."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("split", type=Path, nargs="?", default=Path("splits/train.parquet"))
    parser.add_argument("--index", type=int, default=0, help="Row to load from the split.")
    args = parser.parse_args()

    dataset = build_dataset(args.split)
    example = dataset[args.index]
    atom_count = len(example["atomic_numbers"])
    assert example["target_coords"].shape == (atom_count, 3)
    assert example["input_coords"].shape == (atom_count, 3)
    assert example["is_ligand"].shape == (atom_count,)
    assert example["edge_index"].shape[0] == 2
    assert not example["edge_index"].size or example["edge_index"].max() < atom_count
    assert np.all(example["input_coords"][example["is_ligand"]] == 0)
    assert np.array_equal(
        example["input_coords"][~example["is_ligand"]],
        example["target_coords"][~example["is_ligand"]],
    )
    assert example["is_ligand"].any() and (~example["is_ligand"]).any()

    print(f"Loaded {example['example_id']} from {len(dataset)} examples")
    for name in ("atomic_numbers", "input_coords", "target_coords", "edge_index", "is_ligand"):
        value = example[name]
        print(f"{name}: shape={value.shape}, dtype={value.dtype}")


if __name__ == "__main__":
    main()

When this passes, every example your dataset yields carries five arrays: atomic_numbers, input_coords, target_coords, edge_index, and is_ligand.

What Next?#

With a working transform pipeline, you are ready to write the model that consumes these tensors. Continue to Part 3: Write a Model.

Glossary#

KD-tree: a space-partitioning data structure for fast nearest-neighbor and radius queries.

COO format: “coordinate” sparse format; here, a [2, E] array listing the endpoints of each bond.

Transform: a reusable step that reads and rewrites the example dictionary; chained together with Compose.

Part 2: Creating a Transform Pipeline#

Table of Contents#

Introduction#

Prerequisites#

The Goal of the Pipeline#

Wiring Up the Dataset and Loader#

Building the Transform Pipeline#

Creating Custom Transforms#

CropToPocket#

Start with a Shell#

Input Validation#

Write the crop_to_pocket Function#

Add the Spatial Crop with a KD-tree#

Add the is_ligand annotation#

Declare which transforms must run first#

FeaturizeForDocking#

Start with a Shell#

Declare the dependency on CropToPocket#

Write the featurize_for_docking Function#

Add the is_ligand Mask and Target Coordinates#

Add input_coords with the ligand zeroed out#

Add atomic_numbers#

Add Bonds Via an edge_index#

The final FeaturizeForDocking class#

The Complete Smoke Test#

What Next?#

Glossary#

`CropToPocket`#

Write the `crop_to_pocket` Function#

Add the `is_ligand` annotation#

`FeaturizeForDocking`#

Declare the dependency on `CropToPocket`#

Write the `featurize_for_docking` Function#

Add the `is_ligand` Mask and Target Coordinates#

Add `input_coords` with the ligand zeroed out#

Add `atomic_numbers`#

Add Bonds Via an `edge_index`#

The final `FeaturizeForDocking` class#