diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..3791ec5e --- /dev/null +++ b/.dockerignore @@ -0,0 +1,13 @@ +.git +.gitignore +**/__pycache__ +**/*.pyc +*.egg-info +.eggs +.pytest_cache +.ruff_cache +.mypy_cache +.venv +venv +docs/_build +tests/ \ No newline at end of file diff --git a/.gitignore b/.gitignore index b1ed044f..5968716d 100644 --- a/.gitignore +++ b/.gitignore @@ -171,3 +171,6 @@ tinker/ **/.claude CLAUDE.md .github/workflows/claude*.yml + +# Local docker files +Dockerfile diff --git a/docs/conf.py b/docs/conf.py index 358c8db6..9755db2e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -44,8 +44,11 @@ "sphinx.ext.viewcode", # Add source code links "sphinx.ext.napoleon", # Google/NumPy style docstrings "sphinx_gallery.gen_gallery", # Generates auto_examples/ from examples/ + "myst_parser", # Support for Markdown files + "sphinx_design", # For better layout and design components ] + templates_path = ["_templates"] exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "examples/GALLERY_HEADER.rst", "ml/preprocessing.rst"] diff --git a/docs/docs_requirements.txt b/docs/docs_requirements.txt index e3661864..01f71eb3 100644 --- a/docs/docs_requirements.txt +++ b/docs/docs_requirements.txt @@ -5,4 +5,6 @@ sphinx-autodoc-typehints>=1.20.0,<2 nbsphinx>=0.8.9,<1 sphinx-gallery>=0.8.1,<1 ghp-import>=2.0.0,<3 -pandoc>=2.0.0,<3 \ No newline at end of file +pandoc>=2.0.0,<3 +myst-parser>=5.0.0 +sphinx_design?=>0.7.0 \ No newline at end of file diff --git a/docs/examples/annotate_and_save_structures.py b/docs/examples/annotate_and_save_structures.py index 7fb13648..9b3d3276 100644 --- a/docs/examples/annotate_and_save_structures.py +++ b/docs/examples/annotate_and_save_structures.py @@ -220,6 +220,6 @@ def fix_boolean_annotation(atom_array: struc.AtomArray, annotation_name: str) -> ######################################################################## # Related Examples -# --------------- +# ---------------- # # - :doc:`pocket_conditioning_transform` - Create custom transforms for ligand pocket identification and ML feature generation diff --git a/docs/examples/dataset_exploration.py b/docs/examples/dataset_exploration.py index 889cdeb5..2c2b01e7 100644 --- a/docs/examples/dataset_exploration.py +++ b/docs/examples/dataset_exploration.py @@ -239,7 +239,7 @@ def simple_loading_fn(raw_data: Any) -> dict: ######################################################################## # Related Examples -# --------------- +# ---------------- # # - :doc:`load_and_visualize_structures` - Learn how to load and explore protein structures # - :doc:`pocket_conditioning_transform` - Create custom transforms for ligand pocket identification and ML feature generation diff --git a/docs/examples/load_and_visualize_structures.py b/docs/examples/load_and_visualize_structures.py index 1421c200..5c02b258 100644 --- a/docs/examples/load_and_visualize_structures.py +++ b/docs/examples/load_and_visualize_structures.py @@ -170,7 +170,7 @@ ######################################################################## # Related Examples -# --------------- +# ---------------- # # - :doc:`annotate_and_save_structures` - Learn how to add custom annotations to structures and save them for later use # - :doc:`pocket_conditioning_transform` - Create custom transforms for ligand pocket identification and ML feature generation diff --git a/docs/examples/pocket_conditioning_transform.py b/docs/examples/pocket_conditioning_transform.py index 01b683cf..f7ebcfed 100644 --- a/docs/examples/pocket_conditioning_transform.py +++ b/docs/examples/pocket_conditioning_transform.py @@ -29,7 +29,7 @@ # Conventions # ----------- # **A.** Store information in ``AtomArray`` annotations, not in the state dictionary. -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # # This ensures robustness when atoms are added/removed downstream. # @@ -39,7 +39,7 @@ # - ❌ Store ``pocket_atom_indices`` in dictionary (which creates significant dependencies with operations that delete or re-order atoms) # # **B.** Within ``forward()``, call a stand-alone function with the same name as the transform class. -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # # We thus maintain an object-oriented and a functional API, making our core logic re-usable and testable outside of the ``Transform`` framework. # diff --git a/docs/how_to_build_a_model/how_to_build_a_model_part1.md b/docs/how_to_build_a_model/how_to_build_a_model_part1.md new file mode 100644 index 00000000..0fd877c9 --- /dev/null +++ b/docs/how_to_build_a_model/how_to_build_a_model_part1.md @@ -0,0 +1,337 @@ +# Part 1: Cleaning the Data + +## Table of Contents + +(aw_build_model_p1_intro)= +## Introduction +This is the first in a series of tutorials that walks you through how to use AtomWorks to build a machine learning model for protein design from start to finish. + +**In this installment, you will learn how to use the [IO functionalities in AtomWorks]() to prepare your data for use in a machine learning model.** + +By the end of this tutorial series you will have cleaned data and built a graph neural network to create plausible bound poses between a ligand and a protein pocket. + +```{important} +This tutorial will walk you through creating a script using the AtomWorks API. + +For those that want to use the tutorial text as structure and hints to write your own script, we have hidden the code in collapsible cells. + +If you would like to see the full script, it is provided in the tutorial files. +``` + +(aw_build_model_p1_prereq)= +## Prerequisites +Before starting this tutorial it is assumed that you have: +- An intermediate knowledge of Python and the Pandas library. +- A working installation of AtomWorks. Only the `io` side will be used for this part of the How to Build a Model tutorial series, however other parts will require the `ml` side. +- At least 1GB of space for storage of the parquet files +- Having the PDB mirror set up requires ~100GB of space + +```{note} +If you do not have over 100GB of space on your computing system, you can use a subset of the PDB instead of the full PDB mirror. See [Data Mirrors](../mirrors.rst) for how to only download specific PDB IDs. +``` + +(aw_build_model_p1_setup)= +## Setup +AtomWorks provides a few parquet files that already contains various pieces of metadata about the various structures included in the PDB. We will use these as our starting point. + +Download and decompress the parquet files via: +```bash +wget https://files.ipd.uw.edu/pub/atomworks/dfs/pdb/2026_01_06.tar.gz +tar -xvf 2026_01_06.tar.gz +``` +After decompressing the folder you should see three parquet files: +- **`assemblies.parquet`:** There are multiple bio assemblies stored in a single PDB, each assembly will have multiple chains, interfaces, etc. +- **`interfaces.parquet`:** Contains metadata for all binary interfaces in the PDB +- **`pn_units.parquet`:** Contains metadata for each PN unit in the [PDB](https://www.rcsb.org/) + +We will only use `interfaces.parquet` and `pn_units.parquet` in this tutorial. + + + +Later on you will also need a mirror to at least a subset of the PDB so that AtomWorks is able to find the related structure when training your model. You can learn how to set up a PDB mirror using AtomWorks [here](../mirrors.rst). This document also provides further information on the parquet files. + +(aw_build_model_p1_tutorial)= +## Creating Cleaned Parquet Files + +(aw_build_model_p1_load)= +### Loading the parquet file using Pandas +Let's first take a look at the information contained in the parquet file. Parquet files are not human parsable, but we can use [Pandas](https://pandas.pydata.org/) to inspect it. + +````{dropdown} Click to see the code +Load in the datasets: +```python +import pandas as pd + +interfaces = pd.read_parquet("2026_01_06/interfaces.parquet") +pn_units = pd.read_parquet("2026_01_06/pn_units.parquet") +``` +View the dataset columns: +```python +interfaces.columns +pn_units.columns +``` +```` + +Let's take a closer look at a few of the columns in the interfaces parquet: +- `pdb_id`: The identifier for the specific structure in the PDB. +- `assembly_id`: Integer label for what assembly the given interface belongs to. There may be multiple assemblies in a single PDB structure. +- `pn_unit_1_iid`: Label for the first PN unit in the interface. +- `pn_unit_2_iid`: Label for the second PN unit in the interface. +- `involves_loi`: Boolean for if the LOI (Ligand of Interest) is part of the interface. +- `is_inter_molecule`: Boolean for if the interfaces is between two molecules (True) or within the same molecule (False). +- `involves_metal`: Boolean for if the interface involves a metal atom. +- `involves_covalent_modification`: Boolean for if the interface involves a covalent modification (e.g. glycosylation). +- `num_contacts`: Number of contacts between the two PN units that create the interface. +- `min_distance`: Minimum distance between the contacts that create the interface. + +Let's also look at a few columns of interest from the PN units parquet file: +- `pn_unit_iid`: Label for the specific PN unit in the structure that the row corresponds to. +- `is_polymer`: Boolean for whether the PN unit is a polymer +- `num_resolved_residues`: Number of resolved residues in the structure. + +We encourage you to take a closer look at these datasets on your own. For some suggestions on what to do, see the collapsible group below: + +````{dropdown} Click to see the code +See the first 5 rows of specific columns in a dataset: +```python +interfaces[["pdb_id","assembly_id", "pn_unit_1_iid", "pn_unit_2_iid"]].head() +``` +See the unique values of a given column: +```python +interfaces["assembly_id"].unique() +``` +Determine the datatype of the data stored in a particular column: +```python +interfaces[["is_inter_molecule"]].dtype +``` +```` + +(aw_build_model_p1_merge)= +### Merging Datasets +While the interfaces parquet has most of the data we need to train our model to predict poses for ligands binding to protein pockets, we need the information stored in the `is_polymer` and `num_resolved_residues` columns in the PN units dataset as well. We will use the information in `is_polymer` to ensure that the interfaces we are looking at are between a protein (polymer) and ligand (non-polymer). We will use the information in `num_resolved_residues` to make sure our dataset remains small enough to train our model on a single GPU. + +We will need to merge this information with the interfaces dataset twice, one for each PN unit involved in each interface. Keep in mind that you need the information in `pdb_id`, `assembly_id`, and `pn_unit_iid` to uniquely identify a structure! + +````{dropdown} Click to see how to merge these datasets. +Create two copies of the relevant columns in the PN units dataset, one for each PN unit involved in the interface. +```python +u1_cols = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "is_polymer", "num_resolved_residues"]].copy() +u2_cols = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "is_polymer", "num_resolved_residues"]].copy() +``` +Rename the columns so that they match what is in the interfaces dataset and to ensure the new information is distinguishable between the two PN units. +```python +u1_cols = u1_cols.rename(columns={ + "pn_unit_iid": "pn_unit_1_iid", + "is_polymer": "u1_is_polymer", + "num_resolved_residues": "u1_num_resolved_residues" +}) +u2_cols = u2_cols.rename(columns={ + "pn_unit_iid": "pn_unit_2_iid", + "is_polymer": "u2_is_polymer", + "num_resolved_residues": "u2_num_resolved_residues" +}) +``` +Actually merge the three datasets together: +```python +df = interfaces.merge(u1_cols, on=["pdb_id", "assembly_id", "pn_unit_1_iid"], how="inner") +df = df.merge(u2_cols, on=["pdb_id", "assembly_id", "pn_unit_2_iid"], how="inner") +``` +```` + +Check that the merge occurred correctly by printing out the columns of the new data frame, inspecting the first few rows of the new data frame, etc. + +(aw_build_model_p1_clean)= +### Cleaning the Data +For the purposes of this tutorial, we want to remove rows where `involves_covalent_modification` is True and where the interface involves non-protein and non-ligand chains since we're looking for interfaces that are between a protein and a ligand. + +This means we only want to keep rows where: +- `involves_loi` is True +- `is_inter_molecule` is True +- `involves_covalent_modification` is False +- We want one PN unit involved in the interface to be a polymer and the other to be non-polymer +- The total number of resolved residues for the pocket/ligand combination should be less than 200 + +We also want to make sure we remove any duplicates. The `u1` and `u2` labels are arbitrary, so it's possible for two rows to be identical except for which PN unit was labeled as 1 or 2. + +````{dropdown} Click to see the code +Keep only rows where at least one PN unit is a ligand of interest: +```python +df = df[df["involves_loi"]==True] +``` +Keep only rows where the interfaces are intermolecular: +```python +df = df[df["is_inter_molecule"]==True] +``` +Remove metal mediated-interfaces: +```python +df = df[df["involves_metal"]!=True] +``` +Remove covalently modified residues +```python +df = df[df["involves_covalent_modification"] != True] +``` +Keep only interfaces where one PN unit is a polymer (protein) and one is not (ligand): +```python +df = df[df["u1_is_polymer"] != df["u2_is_polymer"]] +``` +Keep only small examples: +df = df[(df["u1_num_resolved_residues"] + df["u2_num_resolved_residues"]) < 200 ] + +Remove duplicates +df = df.drop_duplications(subset["pdb_id", "assembly_id", "pn_unit_1_iid", "pn_unit_2_iid"]) +```` + +You can check to make sure these filters are actually being applied to your data frame by checking the `len` of the data frame before and after applying each filter. + +To make sure the filter is doing what you expect, you can try running this procedure on a small subset of the data or locating specific rows in the larger dataset that should/should not be impacted by each filtering step. + +(aw_build_model_p1_new_cols)= +### Adding New Columns +It will be useful later on if one row contains unique labels for each remaining interface. Right now information from four rows (`pdb_id`, `assembly_id`, `pn_unit_1_iid`, and `pn_unit_2_iid`) are required to uniquely identify an interface in our dataset. Let's add a new column to our dataframe and store our custom lable there. + +````{dropdown} Click to see the code. +```python +df["example_id"] = (df["pdb_id"] + "_" + + df["assembly_id"] + "_" + + df["pn_unit_1_iid"] + "_" + + df["pn_unit_2_iid"]) +``` +```` + +It will also be useful during training if our dataset already contains the path to the structure file in our PDB mirror. Add this information as a new column to your dataset. +````{dropdown} Click to see the code. +```python +PDB_MIRROR_PATH = os.environ.get("PDB_MIRROR_PATH", "/PATH/TO/pdb_mirror") +df["path"] = df["pdb_id"].str.lower().map( + lambda x: f"{PDB_MIRROR_PATH}/{x[1:3]}/{x}.cif.gz" +) +``` +Replace `"/PATH/TO/pdb_mirror"` if this environment variable is not already set. If it is set, you can leave this second argument blank. +```` + +To check that your code works correctly, make sure these new columns are present when you run `df.columns` and check the first few values in each. Are they what you expected? + +You can also add a test to ensure that all the values in `example_id` are unique: + +````{dropdown} Click to see the code. +```python +assert df["example_id"].nunique() == len(df), "example_id is not unique!" +``` +```` + +(aw_build_model_p1_split)= +### Training, Testing, and Validation Sets +Now that we have the data, we need to split it up into three sets: `test`, `train`, and `val` (short for validation). There are many ways to do this and which is best will depend on your data and what you are trying to accomplish with your model. + +Here, we will use the `protein_cluster_30` column in the PN units data frame to split up our data. This column groups proteins by 30% sequence identity - it contains hash-based IDs that uniquely identify a cluster of proteins sharing more than 30% sequence identity. We will do an 80/10/10 split - 80% of the data will be in training, 10% in test, and 10% in validation. + +We will use this column to split the data by cluster, instead of individual rows. This will prevent the model from seeing near-identical protein pockets between the training and testing sets. + +Isolate this column and the data from the PN units data frame that uniquely identifies each row and merge it with the data frame. Remember that we will need to perform the merge twice, since there are two PN unites in each interface. + +````{dropdown} Click to see the code. +```python +protein_clusters = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "protein_cluster_30"]].copy() + +df = df.merge( + protein_clusters.rename(columns={ + "pn_unit_iid": "pn_unit_1_iid", + "protein_cluster_30": "u1_cluster" + }), + on=["pdb_id", "assembly_id", "pn_unit_1_iid"], + how="left" +) + +df = df.merge( + protein_clusters.rename(columns={ + "pn_unit_iid": "pn_unit_2_iid", + "protein_cluster_30": "u2_cluster" + }), + on=["pdb_id", "assembly_id", "pn_unit_2_iid"], + how="left" +) +``` +```` + +Only the side of the interface that corresponds to the polymer will have a value for `protein_cluster_30`, the ligand will always have a `Null` value. Instead of having to check both `u2_cluster` and `u1_cluster` to determine which cluster our interface belongs to, let's put the information all in one column: + +````{dropdown} Click to see the code. +```python +df["protein_cluster"] = np.where( + df["u1_is_polymer"], + df["u1_cluster"], + df["u2_cluster"] +) +``` +```` + +Check to see if there are any cases where no cluster was assigned. +````{dropdown} Click to see the code. +```python +# There are several ways to check, here we will just count the number of Null values: +len(df[df["protein_cluster"].isna()]) +``` +```` + +If any exist, we want to remove them from our dataset. They likely point to RNA/DNA chains, very short peptides, or low quality entries. +````{dropdown} Click to see the code. +```python +df = df[df["protein_cluster"].notna()].reset_index(drop=True) +``` +```` + +Before splitting the data up, let's shuffle the unique clusters. We use a seed of 42 for reproducibility. Use this seed if you want to exactly replicate what was produced in this segment of the tutorial. +````{dropdown} Click to see the code. +```python +unique_clusters = df["protein_cluster"].unique() +rng = np.random.default_rng(seed=42) +rng.shuffle(unique_clusters) +``` +```` + +Now we can finally spit the data into separate datasets and save them as parquet files for future use: +````{dropdown} Click to see the code. +```python +n = len(unique_clusters) +n_train = int(0.8 * n) +n_val = int(0.1 * n) +# test gets the remainder to avoid off-by-one gaps + +train_clusters = set(unique_clusters[:n_train]) +val_clusters = set(unique_clusters[n_train : n_train + n_val]) +test_clusters = set(unique_clusters[n_train + n_val :]) + +def assign_split(cluster): + if cluster in train_clusters: return "train" + if cluster in val_clusters: return "val" + if cluster in test_clusters: return "test" + return "unassigned" # rows where cluster was null + +df["split"] = df["protein_cluster"].map(assign_split) + +df_train = df[df["split"] == "train"].reset_index(drop=True) +df_val = df[df["split"] == "val"].reset_index(drop=True) +df_test = df[df["split"] == "test"].reset_index(drop=True) + +os.makedirs("splits", exist_ok=True) +df_train.to_parquet("splits/train.parquet", index=False) +df_val.to_parquet("splits/val.parquet", index=False) +df_test.to_parquet("splits/test.parquet", index=False) +``` +```` + +You now have created the datasets you need to train, test, and validate the machine learning model you'll create as you continue to go through the **How to Build a Model Using AtomWorks** tutorial series. + +(aw_build_model_p1_next)= +## What Next? + + + +(aw_build_model_p1_glossary)= +## Glossary + +parquet +PN units - this is actually in the docs [glossary](https://rosettacommons.github.io/atomworks/latest/glossary.html#chains-pn-units-and-molecules) + + + diff --git a/docs/how_to_build_a_model/how_to_build_a_model_part2.md b/docs/how_to_build_a_model/how_to_build_a_model_part2.md new file mode 100644 index 00000000..465e8832 --- /dev/null +++ b/docs/how_to_build_a_model/how_to_build_a_model_part2.md @@ -0,0 +1,714 @@ +# Part 2: The Transform Pipeline + +## Table of Contents + +(aw_build_model_p2_intro)= +## Introduction +This is the second tutorial in the **How to Build a Model Using AtomWorks** series. In [Part 1](how_to_build_a_model.md) you cleaned the PDB metadata and saved `train`/`val`/`test` parquet splits. + +**In this installment, you will learn how to turn each parquet row into model-ready tensors by wiring up a `PandasDataset`, a loader, and a pipeline of transforms — including two custom transforms of your own: `CropToPocket` and `FeaturizeForDocking`.** + +By the end of this part you will have a `smoke_test.py` that loads a single example, crops it to the binding pocket, and produces the five tensors your model will consume in Part 3. + +```{important} +This tutorial continues to build a script using the AtomWorks API. + +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. If you would like to see the full scripts, they are provided in the tutorial files. +``` + +(aw_build_model_p2_prereq)= +## Prerequisites +Before starting this part it is assumed that you have: +- Completed [Part 1](how_to_build_a_model.md) and saved `splits/train.parquet`, `splits/val.parquet`, and `splits/test.parquet`. +- A working installation of AtomWorks, including the `ml` side. +- A PDB mirror set up so the loader can find the structure file for each row (see [Data Mirrors](../mirrors.rst)). +- 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](https://www.biotite-python.org/) so the coordinate and annotation access patterns below feel familiar. +``` + +(aw_build_model_p2_goal)= +## The Goal of the Pipeline +Our task is pose generation: given a protein pocket and a small-molecule ligand, predict plausible bound cartesian coordinates for every atom. To get there, each raw structure needs to be reduced to just the binding pocket and converted into tensors. + +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 no resolved coordinates. +- **`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](how_to_build_a_model_part4.md) when we build the training loop; for now we work in NumPy so the outputs are easy to inspect. + +(aw_build_model_p2_wire)= +## 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_loader_with_query_pn_units()`** returns a picklable loader that parses the CIF file for each row and attaches the *query* PN unit IIDs to the example, so we know which chain(s) the interface of interest involves. + +It is worth reading the documentation for both 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. + +````{dropdown} Click to see how to inspect the documentation +```python +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units + +help(PandasDataset) +help(create_loader_with_query_pn_units) +``` +```` + +As a quick check that everything fits together, write a script that reads `train.parquet`, builds a `PandasDataset` from it, and uses `create_loader_with_query_pn_units` as the loader. + +````{dropdown} Click to see the code +```python +import pandas as pd +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units + +df_train = pd.read_parquet("splits/train.parquet") + +dataset = PandasDataset( + data=df_train, + name="docking_train", + id_column="example_id", + loader=create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"] + ), +) +``` +Note that `id_column="example_id"` reuses the unique identifier you created in Part 1, and the loader is told which columns hold the query PN unit IIDs. +```` + +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. + +````{dropdown} Click to see the code +```python +print(f"Dataset size: {len(dataset)}") + +# Load a single example +example = dataset[0] +print("\nLoaded one sample successfully.") +print("Sample type:", type(example)) + +if hasattr(example, "keys"): + print("Sample keys:", list(example.keys())) + for k, v in example.items(): + print(f"{k}: {type(v)}") +else: + print(example) +``` +```` + +Among the keys you should see `atom_array`, `query_pn_unit_iids`, and `chain_info` — the three inputs the `CropToPocket` transform will rely on. + +(aw_build_model_p2_compose)= +## Building the Transform Pipeline +Transforms are chained together with `Compose`. Let's start with the two transforms that already exist in AtomWorks and add them to the dataset via the `transform` argument. The *Dataset Exploration and Management in AtomWorks* example also shows how to feed a `Compose` pipeline into `PandasDataset`. + +````{dropdown} Click to see the code +```python +import pandas as pd +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import Compose + +df_train = pd.read_parquet("splits/train.parquet") + +pipe = Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), +]) + +dataset = PandasDataset( + data=df_train, + name="docking_train", + id_column="example_id", + loader=create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"] + ), + transform=pipe, +) +``` +```` + +Now we need the two transforms that do not yet exist: `CropToPocket` and `FeaturizeForDocking`. The *Creating Custom Transforms* 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 (snake_case) name. Keeping the logic in a standalone function makes it easy to test and reuse outside the transform machinery. +``` + +(aw_build_model_p2_crop)= +## Writing `CropToPocket` +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 `data` unchanged. Confirm it plugs into `Compose()`. + +````{dropdown} Click to see the code +In `transforms.py`: +```python +from atomworks.ml.transforms.base import Transform + +class CropToPocket(Transform): + def __init__(self, radius: float = 10.0): + super().__init__() + self.radius = radius + + def forward(self, data: dict) -> dict: + return data +``` +Back in your original script: +```python +from transforms import CropToPocket +... +pipe = Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=10.0), +]) +... +``` +```` + +### Add input validation +In `check_input()`, assert that `data` contains the three keys the transform needs: `atom_array`, `query_pn_unit_iids`, and `chain_info`. + +````{dropdown} Click to see the code +In the `CropToPocket` class: +```python + def check_input(self, data: dict) -> None: + assert "atom_array" in data, "Missing atom_array" + assert "query_pn_unit_iids" in data, "Missing query_pn_unit_iids" + assert "chain_info" in data, "Missing chain_info" +``` +```` + +### Add the standalone function +Add a standalone `crop_to_pocket()` function that `forward()` calls. For now, just copy the `AtomArray` and return it. + +````{dropdown} Click to see the code +In `transforms.py`: +```python +import numpy as np +from biotite.structure import AtomArray + +def crop_to_pocket( + atom_array: AtomArray, + query_pn_unit_iids: list, + chain_info: dict, + radius: float = 10.0, +) -> AtomArray: + atom_array = atom_array.copy() + return atom_array + +class CropToPocket(Transform): + ... + def forward(self, data: dict) -> dict: + data["atom_array"] = crop_to_pocket( + data["atom_array"], + query_pn_unit_iids=data["query_pn_unit_iids"], + chain_info=data["chain_info"], + radius=self.radius, + ) + return data +``` +```` + +### Decide which unit is the ligand and which is the protein +Each interface has two query PN units. Use `chain_info` to check which side is a polymer. If exactly one side is a polymer, the non-polymer side is the ligand. If the flags are ambiguous, fall back to a heuristic: the side with fewer atoms is treated as the ligand. + +````{dropdown} Click to see the code +```python +def crop_to_pocket( + atom_array: AtomArray, + query_pn_unit_iids: list, + chain_info: dict, + radius: float = 10.0, +) -> AtomArray: + atom_array = atom_array.copy() + + iid_a, iid_b = query_pn_unit_iids + chain_a = iid_a.split("_")[0] + chain_b = iid_b.split("_")[0] + + a_is_polymer = chain_info.get(chain_a, {}).get("is_polymer", True) + b_is_polymer = chain_info.get(chain_b, {}).get("is_polymer", True) + + if not a_is_polymer: + ligand_iid, protein_iid = iid_a, iid_b + elif not b_is_polymer: + ligand_iid, protein_iid = iid_b, iid_a + else: + mask_a = atom_array.pn_unit_iid == iid_a + mask_b = atom_array.pn_unit_iid == iid_b + if mask_a.sum() <= mask_b.sum(): + ligand_iid, protein_iid = iid_a, iid_b + else: + ligand_iid, protein_iid = iid_b, iid_a + + print("ligand_iid:", ligand_iid) + print("protein_iid:", protein_iid) + return atom_array +``` +```` + +### Build the atom masks and verify both sides have atoms +Create boolean masks for the ligand and protein atoms and raise a clear error if either side is empty. Failing early with a descriptive message makes debugging bad structures much easier later. + +````{dropdown} Click to see the code +```python + ligand_mask = atom_array.pn_unit_iid == ligand_iid + protein_mask = atom_array.pn_unit_iid == protein_iid + + print("ligand atoms:", int(ligand_mask.sum())) + print("protein atoms:", int(protein_mask.sum())) + + if ligand_mask.sum() == 0: + raise ValueError(f"Ligand {ligand_iid} has no atoms") + if protein_mask.sum() == 0: + raise ValueError(f"Protein {protein_iid} has no atoms") + return atom_array +``` +```` + +### Add the spatial crop with a KD-tree +Now for the actual crop. We keep every ligand atom plus every protein atom within `radius` 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`. +``` + +````{dropdown} Click to see the code +```python +from scipy.spatial import cKDTree + +def crop_to_pocket( + atom_array: AtomArray, + query_pn_unit_iids: list, + chain_info: dict, + radius: float = 10.0, +) -> AtomArray: + atom_array = atom_array.copy() + + iid_a, iid_b = query_pn_unit_iids + chain_a = iid_a.split("_")[0] + chain_b = iid_b.split("_")[0] + + a_is_polymer = chain_info.get(chain_a, {}).get("is_polymer", True) + b_is_polymer = chain_info.get(chain_b, {}).get("is_polymer", True) + + if not a_is_polymer: + ligand_iid, protein_iid = iid_a, iid_b + elif not b_is_polymer: + ligand_iid, protein_iid = iid_b, iid_a + else: + mask_a = atom_array.pn_unit_iid == iid_a + mask_b = atom_array.pn_unit_iid == iid_b + if mask_a.sum() <= mask_b.sum(): + ligand_iid, protein_iid = iid_a, iid_b + else: + ligand_iid, protein_iid = iid_b, iid_a + + ligand_mask = atom_array.pn_unit_iid == ligand_iid + protein_mask = atom_array.pn_unit_iid == protein_iid + + ligand_coords = atom_array.coord[ligand_mask] + protein_coords = atom_array.coord[protein_mask] + + if len(ligand_coords) == 0: + raise ValueError(f"Ligand {ligand_iid} has no atoms") + if len(protein_coords) == 0: + raise ValueError(f"Protein {protein_iid} has no atoms") + + tree = cKDTree(protein_coords) + neighbor_indices = tree.query_ball_point(ligand_coords, r=radius) + + total_neighbors = sum(len(n) for n in neighbor_indices) + if total_neighbors == 0: + raise ValueError(f"No protein atoms found within {radius}A of ligand {ligand_iid}") + + pocket_local_indices = np.unique(np.concatenate(neighbor_indices).astype(int)) + protein_global_indices = np.where(protein_mask)[0] + pocket_global_indices = protein_global_indices[pocket_local_indices] + ligand_global_indices = np.where(ligand_mask)[0] + + keep = np.sort(np.concatenate([pocket_global_indices, ligand_global_indices])) + cropped = atom_array[keep] + return cropped +``` +The `cKDTree` is built from the protein coordinates; `query_ball_point` returns indices *local* to `protein_coords`, so we map them back to global indices via `protein_global_indices` before slicing the `AtomArray`. +```` + +Confirm the crop runs in your test: + +````{dropdown} Click to see the code +```python +example = dataset[0] +cropped = example["atom_array"] +print("\nLoaded one sample successfully.") +print("Cropped atom count:", len(cropped)) +``` +```` + +### 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 `keep` and `ligand_global_indices`, 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. +``` + +````{dropdown} Click to see the code +Add these two lines just before returning, so the end of `crop_to_pocket` reads: +```python + keep = np.sort(np.concatenate([pocket_global_indices, ligand_global_indices])) + cropped = atom_array[keep] + + is_ligand = np.isin(keep, ligand_global_indices) + cropped.set_annotation("is_ligand", is_ligand) + return cropped +``` +```` + +Quick check in your test: + +````{dropdown} Click to see the code +```python +example = dataset[0] +cropped = example["atom_array"] +print("\nLoaded one sample successfully.") +print("Cropped atom count:", len(cropped)) +print("Has is_ligand annotation:", hasattr(cropped, "is_ligand")) +print("Ligand atoms kept:", int(cropped.is_ligand.sum())) +print("Pocket atoms kept:", int((~cropped.is_ligand).sum())) + +assert hasattr(cropped, "is_ligand") +assert cropped.is_ligand.any(), "No ligand atoms were kept" +assert (~cropped.is_ligand).any(), "No pocket atoms were kept" +print("CropToPocket sanity checks passed.") +``` +```` + +### Declare which transforms must run first +`CropToPocket` assumes hydrogens and unresolved atoms are already gone. Declare this ordering requirement so AtomWorks can enforce it. + +````{dropdown} Click to see the code +```python +class CropToPocket(Transform): + requires_previous_transforms = ["RemoveHydrogens", "RemoveUnresolvedAtoms"] + ... +``` +```` + +At this point you can assemble a full `smoke_test.py`: + +````{dropdown} Click to see the full Pyhton file: +```python +import pandas as pd +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import Compose +from transforms import CropToPocket + +df_train = pd.read_parquet("splits/train.parquet") + +pipe = Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=10.0), +]) + +dataset = PandasDataset( + data=df_train, + name="docking_train", + id_column="example_id", + loader=create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"] + ), + transform=pipe, +) + +print(f"Dataset size: {len(dataset)}") +example = dataset[0] +cropped = example["atom_array"] +print("\nLoaded one sample successfully.") +print("Sample keys:", list(example.keys())) +print("Cropped atom count:", len(cropped)) +print("Has is_ligand annotation:", hasattr(cropped, "is_ligand")) +print("Ligand atoms kept:", int(cropped.is_ligand.sum())) +print("Pocket atoms kept:", int((~cropped.is_ligand).sum())) + +assert hasattr(cropped, "is_ligand") +assert cropped.is_ligand.any(), "No ligand atoms were kept" +assert (~cropped.is_ligand).any(), "No pocket atoms were kept" +print("\nCropToPocket smoke test passed.") +``` +```` + +(aw_build_model_p2_featurize)= +## Writing `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 +````{dropdown} Click to see the code +In `transforms.py`: +```python +class FeaturizeForDocking(Transform): + def forward(self, data: dict) -> dict: + return data +``` +Add it to `Compose()` in `smoke_test.py`: +```python +from transforms import CropToPocket, FeaturizeForDocking + +pipe = Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=10.0), + FeaturizeForDocking(), +]) +``` +```` + +### 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()`. + +````{dropdown} Click to see the code +```python +from atomworks.ml.transforms._checks import check_atom_array_annotation + +class FeaturizeForDocking(Transform): + requires_previous_transforms = ["CropToPocket"] + + def check_input(self, data: dict) -> None: + check_atom_array_annotation(data, ["is_ligand"]) +``` +```` + +### Add the standalone function +````{dropdown} Click to see the code +```python +def featurize_for_docking(atom_array: AtomArray) -> dict: + return {} + +class FeaturizeForDocking(Transform): + ... + def forward(self, data: dict) -> dict: + features = featurize_for_docking(data["atom_array"]) + data.update(features) + return data +``` +```` + +### Add the `is_ligand` mask and target coordinates +`target_coords` are the ground-truth coordinates the model must reproduce. + +````{dropdown} Click to see the code +```python +def featurize_for_docking(atom_array: AtomArray) -> dict: + is_ligand = atom_array.is_ligand.astype(bool) + target_coords = atom_array.coord.astype(np.float32) + return { + "is_ligand": is_ligand, + "target_coords": target_coords, + } +``` +Check the new features in your smoke test: +```python +example = dataset[0] +print("Sample keys:", list(example.keys())) +print("is_ligand shape:", example["is_ligand"].shape) +print("ligand atoms:", int(example["is_ligand"].sum())) +assert "is_ligand" in example +assert example["is_ligand"].dtype == bool +print("target_coords shape:", example["target_coords"].shape) +print("target_coords dtype:", example["target_coords"].dtype) +assert example["target_coords"].ndim == 2 +assert example["target_coords"].shape[1] == 3 +``` +```` + +### 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. + +````{dropdown} Click to see the code +```python +def featurize_for_docking(atom_array: AtomArray) -> dict: + is_ligand = atom_array.is_ligand.astype(bool) + target_coords = atom_array.coord.astype(np.float32) + + input_coords = target_coords.copy() + input_coords[is_ligand] = 0.0 + + return { + "is_ligand": is_ligand, + "target_coords": target_coords, + "input_coords": input_coords, + } +``` +```` + +### 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. + +````{dropdown} Click to see the code +```python +from atomworks.constants import ELEMENT_NAME_TO_ATOMIC_NUMBER +... +def featurize_for_docking(atom_array: AtomArray) -> dict: + is_ligand = atom_array.is_ligand.astype(bool) + target_coords = atom_array.coord.astype(np.float32) + + input_coords = target_coords.copy() + input_coords[is_ligand] = 0.0 + + atomic_numbers = np.array( + [ELEMENT_NAME_TO_ATOMIC_NUMBER.get(e.upper(), 0) for e in atom_array.element], + dtype=np.int64, + ) + + return { + "atomic_numbers": atomic_numbers, + "is_ligand": is_ligand, + "target_coords": target_coords, + "input_coords": input_coords, + } +``` +Elements not found in the lookup table map to `0`, which acts as an "unknown atom" index in the embedding table. +```` + +### Add `edge_index` from the bond graph +`edge_index` is the bond graph in COO (coordinate) format: a `[2, E]` integer array where each column is a bonded pair of atoms. + +````{dropdown} Click to see the code +```python +def featurize_for_docking(atom_array: AtomArray) -> dict: + is_ligand = atom_array.is_ligand.astype(bool) + target_coords = atom_array.coord.astype(np.float32) + + input_coords = target_coords.copy() + input_coords[is_ligand] = 0.0 + + atomic_numbers = np.array( + [ELEMENT_NAME_TO_ATOMIC_NUMBER.get(e.upper(), 0) for e in atom_array.element], + dtype=np.int64, + ) + + bonds = atom_array.bonds.as_array() + edge_index = bonds[:, :2].T.astype(np.int64) + + return { + "atomic_numbers": atomic_numbers, + "input_coords": input_coords, + "target_coords": target_coords, + "edge_index": edge_index, + "is_ligand": is_ligand, + } +``` +`atom_array.bonds.as_array()` returns an array whose first two columns are the indices of the bonded atoms; we take those two columns and transpose to get the `[2, E]` shape. +```` + +Check the edge graph in your smoke test: + +````{dropdown} Click to see the code +```python +example = dataset[0] +print("edge_index shape:", example["edge_index"].shape) +print("edge_index dtype:", example["edge_index"].dtype) +assert example["edge_index"].ndim == 2 +assert example["edge_index"].shape[0] == 2 +assert example["edge_index"].max() < example["atomic_numbers"].shape[0] +``` +```` + +### The final `FeaturizeForDocking` class +````{dropdown} Click to see the code +```python +class FeaturizeForDocking(Transform): + requires_previous_transforms = ["CropToPocket"] + + def check_input(self, data: dict) -> None: + check_atom_array_annotation(data, ["is_ligand"]) + + def forward(self, data: dict) -> dict: + features = featurize_for_docking(data["atom_array"]) + data.update(features) + return data +``` +```` + +(aw_build_model_p2_smoke)= +## The Complete Smoke Test +Your final `smoke_test.py` should build the full pipeline and validate every tensor: + +````{dropdown} Click to see the complete file +```python +import pandas as pd +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import Compose +from transforms import CropToPocket, FeaturizeForDocking + +df_train = pd.read_parquet("splits/train.parquet") + +pipe = Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=10.0), + FeaturizeForDocking(), +]) + +dataset = PandasDataset( + data=df_train, + name="docking_train", + id_column="example_id", + loader=create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"] + ), + transform=pipe, +) + +print(f"Dataset size: {len(dataset)}") +example = dataset[0] +print("\nLoaded one sample successfully.") +print("Sample keys:", list(example.keys())) +print("atomic_numbers:", example["atomic_numbers"].shape, example["atomic_numbers"].dtype) +print("input_coords:", example["input_coords"].shape, example["input_coords"].dtype) +print("target_coords:", example["target_coords"].shape, example["target_coords"].dtype) +print("edge_index:", example["edge_index"].shape, example["edge_index"].dtype) +print("is_ligand:", example["is_ligand"].shape, example["is_ligand"].dtype) + +assert example["atomic_numbers"].ndim == 1 +assert example["target_coords"].ndim == 2 +assert example["target_coords"].shape[1] == 3 +assert example["input_coords"].shape == example["target_coords"].shape +assert example["edge_index"].ndim == 2 +assert example["edge_index"].shape[0] == 2 +assert example["edge_index"].max() < example["atomic_numbers"].shape[0] +assert (example["input_coords"][example["is_ligand"]] == 0).all(), \ + "Ligand coordinates should be zeroed" +assert (example["input_coords"][~example["is_ligand"]] != 0).any(), \ + "Pocket coordinates should not all be zero" +print("\nFeaturizeForDocking smoke test passed.") +``` +```` + +When this passes, every example your dataset yields carries five arrays: `atomic_numbers`, `input_coords`, `target_coords`, `edge_index`, and `is_ligand`. + +(aw_build_model_p2_next)= +## 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](how_to_build_a_model_part3.md). + +(aw_build_model_p2_glossary)= +## 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`. diff --git a/docs/how_to_build_a_model/how_to_build_a_model_part3.md b/docs/how_to_build_a_model/how_to_build_a_model_part3.md new file mode 100644 index 00000000..51b7fe9e --- /dev/null +++ b/docs/how_to_build_a_model/how_to_build_a_model_part3.md @@ -0,0 +1,323 @@ +# Part 3: Write a Model + +## Table of Contents + +(aw_build_model_p3_intro)= +## Introduction +This is the third tutorial in the **How to Build a Model Using AtomWorks** series. So far you have: + +- Parquet files for your `train`/`val`/`test` splits ([Part 1](how_to_build_a_model.md)). +- A loader that turns each parquet row into a protein-ligand structure ([Part 2](how_to_build_a_model_part2.md)). +- A transform pipeline that crops the pocket and converts it into tensors ([Part 2](how_to_build_a_model_part2.md)). + +**In this installment, you will write the neural network that consumes those tensors and predicts 3D coordinates for every atom.** + +By the end of this part you will have a `model.py` containing a trainable `PocketDockGNN` `LightningModule`. + +```{important} +This tutorial continues to build a script using the AtomWorks API and PyTorch. The full solution is available in the tutorial files; the code here is also hidden in collapsible cells so you can attempt each step yourself first. +``` + +(aw_build_model_p3_prereq)= +## Prerequisites +Before starting this part it is assumed that you have: +- Completed [Part 2](how_to_build_a_model_part2.md), with a working `transforms.py` and `smoke_test.py`. +- A working installation of PyTorch and PyTorch Lightning. +- Familiarity with basic neural-network building blocks (`Linear`, `Embedding`, `LayerNorm`) and the idea of message passing on a graph. + +(aw_build_model_p3_choice)= +## Choosing an Architecture +Before writing any code, we have to decide what architecture to use. For this tutorial we use a simple message-passing **graph neural network (GNN)**. A GNN is a natural choice because our data is already a graph: the bond graph is stored in `edge_index`, and each atom is a node with an atomic-number feature. + +```{warning} +This model is **not** rotation/translation invariant. It sees raw XYZ coordinates, which means it can learn to "cheat" based on absolute position rather than on geometry. Equivariant architectures (for example, those built on relative displacements or SE(3)-equivariant layers) address this, but they add substantial complexity. We keep things simple here so the pipeline is easy to follow; treat the resulting model as a teaching example rather than a production docking model. +``` + +A few other decisions to make before writing: + +- **How many GNN layers?** More layers propagate information further across the bond graph, at the cost of compute and the risk of over-smoothing. +- **What hidden dimension?** Wider layers can represent more, but use more memory. +- **What loss function?** This should reflect what "a good pose" means. + +To keep the example small and trainable on a single GPU, we use: + +- **3 GNN layers** — enough to propagate information a few hops through the bond graph. +- **128 hidden dimensions** — small enough to train on a single GPU. +- **Mean squared error (MSE)** between predicted and target coordinates. + +(aw_build_model_p3_goal)= +## The Goal of `model.py` +We will write the model in `model.py`. It needs to: + +- Take atom types, coordinates, and bond edges as input. +- Pass information along the atom graph. +- Predict 3D coordinates for each atom. +- Expose train/val/test steps through a `LightningModule`. + +(aw_build_model_p3_imports)= +## Imports and Class Definition +Import PyTorch, Lightning, and the basic building blocks. `torch` provides tensors and tensor ops; `torch.nn` provides layers like `Linear`, `Embedding`, and `LayerNorm`; `pytorch_lightning` provides `LightningModule`, which packages the model, loss, logging, and optimizer setup into one class. + +````{dropdown} Click to see the code +```python +import torch +import torch.nn as nn +import pytorch_lightning as pl + +class PocketDockGNN(pl.LightningModule): + ... +``` +```` + +(aw_build_model_p3_init)= +## `__init__` and Hyperparameters +Write `__init__` and save the hyperparameters so Lightning can restore them from a checkpoint. + +````{dropdown} Click to see the code +```python + def __init__( + self, + num_atom_types: int = 119, + hidden_dim: int = 128, + num_layers: int = 3, + learning_rate: float = 1e-3, + ): + super().__init__() + self.save_hyperparameters() + self.learning_rate = learning_rate +``` +```` + +These arguments mean: + +- **`num_atom_types=119`** — one embedding-table entry for each atomic number from 0 to 118. +- **`hidden_dim`** — the width of the learned per-atom representation. +- **`num_layers`** — how many rounds of message passing to run. +- **`learning_rate`** — the Adam step size. + +(aw_build_model_p3_layers)= +## Building the Layers +We build the layers in the order the data flows through them. + +### Embed atom identities +Each atom arrives as an integer atomic number. A neural network works better with learned vectors, so we add an embedding table. This lets the model learn different behavior for carbon, oxygen, nitrogen, and so on without us hand-coding any chemistry rules. + +````{dropdown} Click to see the code +```python + self.atom_embedding = nn.Embedding(num_atom_types, hidden_dim) +``` +```` + +### Combine atom type with input coordinates +The featurizer produces two especially important per-atom inputs: atom identity and input coordinates. Concatenate the atom embedding with the 3 coordinate values and project back to `hidden_dim`. + +````{dropdown} Click to see the code +```python + self.input_proj = nn.Sequential( + nn.Linear(hidden_dim + 3, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) +``` +```` + +### Build the message-passing blocks +The graph structure comes from `edge_index`, which says which atoms are bonded. To use that graph, add a stack of message and update blocks, plus layer norms. + +````{dropdown} Click to see the code +```python + self.conv_layers = nn.ModuleList([ + nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + for _ in range(num_layers) + ]) + + self.update_layers = nn.ModuleList([ + nn.Sequential( + nn.Linear(hidden_dim * 2, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + for _ in range(num_layers) + ]) + + self.layer_norms = nn.ModuleList([ + nn.LayerNorm(hidden_dim) + for _ in range(num_layers) + ]) +``` +```` + +These three lists work together: + +- **`conv_layers`** — transform each source atom into a message. +- **`update_layers`** — combine an atom's current state with the aggregated neighbor message. +- **`layer_norms`** — stabilize training after each update. + +### Project back to 3D coordinates +After message passing, each atom has a learned hidden representation. The output head maps that vector to an (x, y, z) prediction. + +````{dropdown} Click to see the code +```python + self.output_proj = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, 3), + ) +``` +```` + +### Define the loss +For the tutorial we use mean squared error between predicted and target coordinates. + +````{dropdown} Click to see the code +```python + self.loss_fn = nn.MSELoss() +``` +```` + +(aw_build_model_p3_forward)= +## The Forward Pass +Now write `forward()`. It does four things in order: + +1. Embed the atomic numbers. +2. Concatenate those embeddings with the input coordinates and project. +3. Run message passing over the bond graph. For each layer, compute a message for every atom, route each source message to its destination atom using `edge_index`, sum the incoming messages per destination, and update each atom with a residual connection. +4. Project to coordinates. + +````{dropdown} Click to see the code +```python + def forward( + self, + atomic_numbers: torch.Tensor, + input_coords: torch.Tensor, + edge_index: torch.Tensor, + ) -> torch.Tensor: + x = self.atom_embedding(atomic_numbers) + x = self.input_proj(torch.cat([x, input_coords], dim=-1)) + + src, dst = edge_index[0], edge_index[1] + + for conv, update, norm in zip(self.conv_layers, self.update_layers, self.layer_norms): + messages = conv(x) + agg = torch.zeros_like(x) + agg.scatter_add_(0, dst.unsqueeze(-1).expand(-1, x.size(-1)), messages[src]) + x = x + norm(update(torch.cat([x, agg], dim=-1))) + + pred_coords = self.output_proj(x) + return pred_coords +``` +`scatter_add_` sums each source atom's message into the row of `agg` belonging to its destination atom, which implements neighbor aggregation. The residual connection (`x = x + ...`) keeps gradients well-behaved across layers. +```` + +```{note} +`atom_array.bonds.as_array()` lists each bond once, so as written, messages flow from `src` to `dst` in a single direction per bond. If you want symmetric message passing (information flowing both ways along every bond), you can duplicate and flip the edges when you build `edge_index`, or concatenate `[src, dst]` with `[dst, src]`. This is optional for the tutorial but worth knowing. +``` + +(aw_build_model_p3_step)= +## One Shared Train/Val/Test Step +Lightning calls separate methods for training, validation, and test, but the logic is almost identical, so write it once in a helper. Skip `None` batches (these come from failed examples, which we handle in [Part 4](how_to_build_a_model_part4.md)), remove the leading batch dimension added by the `DataLoader` (we use `batch_size=1`), run the model, and compute the coordinate loss. + +The full-coordinate MSE is fine for optimization, but docking quality is best judged on the ligand atoms alone, so we also log a ligand RMSD. + +````{dropdown} Click to see the code +```python + def _shared_step(self, batch: dict, stage: str) -> torch.Tensor: + if batch is None: + return None + # Remove the batch dimension added by DataLoader (batch_size=1) + atomic_numbers = batch["atomic_numbers"].squeeze(0) # (N,) + input_coords = batch["input_coords"].squeeze(0) # (N, 3) + target_coords = batch["target_coords"].squeeze(0) # (N, 3) + edge_index = batch["edge_index"].squeeze(0) # (2, E) + is_ligand = batch["is_ligand"].squeeze(0) # (N,) + + pred_coords = self(atomic_numbers, input_coords, edge_index) + loss = self.loss_fn(pred_coords, target_coords) + + with torch.no_grad(): + ligand_rmsd = torch.sqrt( + ((pred_coords[is_ligand] - target_coords[is_ligand]) ** 2) + .sum(dim=-1).mean() + ) + + self.log(f"{stage}/loss", loss, prog_bar=True) + self.log(f"{stage}/ligand_rmsd", ligand_rmsd, prog_bar=True) + + return loss +``` +```` + +(aw_build_model_p3_hooks)= +## Connect Lightning's Step Methods +The stage-specific methods become tiny wrappers around the shared step. + +````{dropdown} Click to see the code +```python + def training_step(self, batch: dict, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, "train") + + def validation_step(self, batch: dict, batch_idx: int) -> None: + self._shared_step(batch, "val") + + def test_step(self, batch: dict, batch_idx: int) -> None: + self._shared_step(batch, "test") +``` +```` + +(aw_build_model_p3_optim)= +## Configure the Optimizer +Tell Lightning which optimizer to use. Here we use Adam. + +````{dropdown} Click to see the code +```python + def configure_optimizers(self): + return torch.optim.Adam(self.parameters(), lr=self.learning_rate) +``` +```` + +(aw_build_model_p3_shape)= +## The Shape of `model.py` +Your final `model.py` should have this shape: + +````{dropdown} Click to see the outline +```python +import torch +import torch.nn as nn +import pytorch_lightning as pl + +class PocketDockGNN(pl.LightningModule): + def __init__(...): + ... + def forward(...): + ... + def _shared_step(...): + ... + def training_step(...): + ... + def validation_step(...): + ... + def test_step(...): + ... + def configure_optimizers(...): + ... +``` +```` + +At this point you should have a trainable model class. + +(aw_build_model_p3_next)= +## What Next? +With `transforms.py` and `model.py` in place, you are ready to wire everything into a training loop. Continue to [Part 4: Train a Model](how_to_build_a_model_part4.md). + +(aw_build_model_p3_glossary)= +## Glossary + +GNN (graph neural network) — a network that operates on graph-structured data by passing messages between connected nodes. + +message passing — the process of computing a message from each node, aggregating messages at each destination node, and updating node states. + +RMSD (root-mean-square deviation) — the square root of the mean squared distance between predicted and true atom positions; a standard structural-accuracy metric. diff --git a/docs/how_to_build_a_model/how_to_build_a_model_part4.md b/docs/how_to_build_a_model/how_to_build_a_model_part4.md new file mode 100644 index 00000000..bc90702b --- /dev/null +++ b/docs/how_to_build_a_model/how_to_build_a_model_part4.md @@ -0,0 +1,395 @@ +# Part 4: Train a Model + +## Table of Contents + +(aw_build_model_p4_intro)= +## Introduction +This is the fourth and final tutorial in the **How to Build a Model Using AtomWorks** series. By now you have: + +- Parquet files for your `train`/`val`/`test` splits ([Part 1](how_to_build_a_model.md)). +- A loader that turns each parquet row into a protein-ligand structure ([Part 2](how_to_build_a_model_part2.md)). +- A transform pipeline that crops the pocket and converts it into tensors ([Part 2](how_to_build_a_model_part2.md)). +- A `PocketDockGNN` model in `model.py` ([Part 3](how_to_build_a_model_part3.md)). + +**In this installment, you will write `train.py`, which wires all of these pieces together to train, validate, checkpoint, and test the model.** + +```{important} +This tutorial completes the training script using the AtomWorks API, PyTorch, and PyTorch Lightning. The full solution is available in the tutorial files, and the code below is hidden in collapsible cells so you can attempt each step yourself first. +``` + +(aw_build_model_p4_prereq)= +## Prerequisites +Before starting this part it is assumed that you have: +- Completed [Parts 1–3](how_to_build_a_model.md), with `transforms.py` and `model.py` in place. +- A working installation of PyTorch and PyTorch Lightning. +- Access to a GPU (the example trainer is configured for a single GPU, but you can change the accelerator). + +(aw_build_model_p4_overview)= +## What the Training Script Does +`train.py` wires together four pieces: + +1. Read the split parquets. +2. Load and transform each example into model inputs. +3. Instantiate your model (`PocketDockGNN`). +4. Train, validate, checkpoint, and test. + +The model sees a pocket-centered graph with five tensors: + +- **`atomic_numbers`** — atomic identity for each atom. +- **`input_coords`** — protein pocket coordinates are kept, ligand coordinates are zeroed out. +- **`target_coords`** — the true coordinates the model should predict. +- **`edge_index`** — bond graph connectivity. +- **`is_ligand`** — a boolean mask for which atoms belong to the ligand. + +In other words, the task is: *given the protein pocket context and the ligand atoms, predict the ligand's 3D placement.* + +(aw_build_model_p4_imports)= +## Imports and Global Settings +Import PyTorch, PyTorch Lightning, the AtomWorks dataset utilities, the transforms you wrote in Part 2, and the model class from Part 3. + +````{dropdown} Click to see the code +```python +import torch +import pandas as pd +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping +from torch.utils.data import DataLoader, Dataset + +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import ConvertToTorch, Compose + +from transforms import CropToPocket, FeaturizeForDocking +from model import PocketDockGNN +``` +Note that we now also import `ConvertToTorch`, which turns the NumPy features from `FeaturizeForDocking` into `torch` tensors. +```` + +Add two small global settings. + +````{dropdown} Click to see the code +```python +torch.set_float32_matmul_precision("medium") +pl.seed_everything(42) +``` +`torch.set_float32_matmul_precision("medium")` is a practical speed/precision trade-off for training. `pl.seed_everything(42)` makes runs more reproducible. +```` + +(aw_build_model_p4_config)= +## Configuration +Put the tunable settings in one place near the top of the file. + +````{dropdown} Click to see the code +```python +CONFIG = { + "hidden_dim": 128, + "num_layers": 3, + "learning_rate": 1e-3, + "batch_size": 1, + "max_epochs": 5, + "pocket_radius": 10.0, + "num_workers": 0, + "max_train": 100, + "max_val": 20, + "max_test": 20, +} +``` +```` + +```{note} +The `max_*` values cap how many examples we use so the run stays small and fast to troubleshoot. For actual training, remove these caps (set them to `None`) and train on the full splits. +``` + +(aw_build_model_p4_keys)= +## List the Tensor Keys +Your featurization step produces five tensors. Write those keys down once so the rest of the file can reuse the same list. This list is used both to tell `ConvertToTorch` what to convert and to tell the collate function what to stack. + +````{dropdown} Click to see the code +```python +TENSOR_KEYS = [ + "atomic_numbers", + "input_coords", + "target_coords", + "edge_index", + "is_ligand", +] +``` +```` + +```{important} +`TENSOR_KEYS` must be defined before the functions that reference it (`collate_fn` and `build_pipeline`). Keep it near the top of the file with `CONFIG`. +``` + +(aw_build_model_p4_robust)= +## Make the Dataset Robust to Bad Examples +If you train on real structures — which we are — some examples will fail. A ligand may have no resolved coordinates, or the pocket crop may remove everything useful. You do not want one bad structure to crash the entire run. + +Wrap `PandasDataset` in a small `RobustDataset`. If a single example raises during loading or transform, it records the index and returns `None` instead of killing training. + +````{dropdown} Click to see the code +```python +class RobustDataset(Dataset): + def __init__(self, dataset: PandasDataset): + self.dataset = dataset + self.failed = [] + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, idx): + try: + return self.dataset[idx] + except Exception: + self.failed.append(idx) + return None +``` +```` + +(aw_build_model_p4_collate)= +## Write a Custom `collate_fn` +The default PyTorch collator is not a good fit here: the examples are variable-size graphs, the dataset may return `None`, and we only want to batch the specific tensors the model needs (ignoring any extra fields in the example dictionary). The collate function should do two things: drop failed examples, and stack only the tensor keys the model expects. + +````{dropdown} Click to see the code +```python +def collate_fn(batch): + batch = [b for b in batch if b is not None] + if len(batch) == 0: + return None + + return { + k: torch.stack([example[k] for example in batch]) + for k in TENSOR_KEYS + if k in batch[0] + } +``` +With `batch_size=1`, this stacks a single example and adds a leading dimension of size 1, which the model's `_shared_step` removes with `.squeeze(0)`. +```` + +(aw_build_model_p4_pipeline)= +## Rebuild the Transform Pipeline +Write a helper that reconstructs the same transform sequence you tested in Part 2, now with `ConvertToTorch` appended so the features come out as `torch` tensors. + +````{dropdown} Click to see the code +```python +def build_pipeline(radius: float) -> Compose: + return Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=radius), + FeaturizeForDocking(), + ConvertToTorch(keys=TENSOR_KEYS), + ]) +``` +```` + +(aw_build_model_p4_factories)= +## Dataset and Dataloader Factories +Write one function that takes a split parquet path and returns a ready-to-use dataset (this is essentially what you did in `smoke_test.py`, wrapped in `RobustDataset`). + +````{dropdown} Click to see the code +```python +def build_dataset(parquet_path: str, name: str, radius: float, max_examples: int = None): + df = pd.read_parquet(parquet_path) + if max_examples is not None: + df = df.head(max_examples).reset_index(drop=True) + + loader = create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"] + ) + + dataset = PandasDataset( + data=df, + name=name, + id_column="example_id", + loader=loader, + transform=build_pipeline(radius), + save_failed_examples_to_dir="failed_examples/", + ) + return RobustDataset(dataset) +``` +`save_failed_examples_to_dir` writes any example that fails inside the AtomWorks pipeline to disk so you can inspect it later. +```` + +Then write one small helper to build dataloaders consistently. This keeps the script tidy and reuses the same settings for train, validation, and test. + +````{dropdown} Click to see the code +```python +def build_dataloader(dataset, shuffle: bool) -> DataLoader: + return DataLoader( + dataset, + batch_size=CONFIG["batch_size"], + shuffle=shuffle, + num_workers=CONFIG["num_workers"], + collate_fn=collate_fn, + persistent_workers=False, + ) +``` +```` + +(aw_build_model_p4_build)= +## Build the Datasets and Dataloaders +Now the main script body. Start by building the three datasets. + +````{dropdown} Click to see the code +```python +print("Building datasets...") +train_dataset = build_dataset( + "splits/train.parquet", "docking_train", CONFIG["pocket_radius"], CONFIG["max_train"] +) +val_dataset = build_dataset( + "splits/val.parquet", "docking_val", CONFIG["pocket_radius"], CONFIG["max_val"] +) +test_dataset = build_dataset( + "splits/test.parquet", "docking_test", CONFIG["pocket_radius"], CONFIG["max_test"] +) + +print(f" Train: {len(train_dataset):,} examples") +print(f" Val: {len(val_dataset):,} examples") +print(f" Test: {len(test_dataset):,} examples") +``` +```` + +Then build the dataloaders. + +````{dropdown} Click to see the code +```python +train_loader = build_dataloader(train_dataset, shuffle=True) +val_loader = build_dataloader(val_dataset, shuffle=False) +test_loader = build_dataloader(test_dataset, shuffle=False) +``` +```` + +(aw_build_model_p4_model)= +## Instantiate the Model +Instantiate the network with the hyperparameters from `CONFIG`. + +````{dropdown} Click to see the code +```python +model = PocketDockGNN( + hidden_dim=CONFIG["hidden_dim"], + num_layers=CONFIG["num_layers"], + learning_rate=CONFIG["learning_rate"], +) +``` +Optionally, print the parameter count as a sanity check that you built the expected network: +```python +print(f"\nModel parameters: {sum(p.numel() for p in model.parameters()):,}") +``` +```` + +(aw_build_model_p4_callbacks)= +## Checkpointing and Early Stopping +Before creating the trainer, define the callbacks. Use `periodic_checkpoint` to give recovery points during training, and `best_checkpoint` to keep the best models by `val/loss`. + +````{dropdown} Click to see the code +```python +periodic_checkpoint = ModelCheckpoint( + dirpath="checkpoints/", + filename="pocketdockgnn-{epoch:02d}-{step}", + every_n_train_steps=50, + save_last=True, + verbose=True, +) + +best_checkpoint = ModelCheckpoint( + dirpath="checkpoints/", + filename="pocketdockgnn-best-{epoch:02d}-{val/loss:.4f}", + monitor="val/loss", + mode="min", + save_top_k=3, + verbose=True, +) + +callbacks = [ + best_checkpoint, + periodic_checkpoint, + EarlyStopping( + monitor="val/loss", + patience=10, + mode="min", + verbose=True, + ), +] +``` +The `EarlyStopping` callback stops training if validation loss stops improving for `patience` checks. +```` + +(aw_build_model_p4_trainer)= +## Create the Trainer +Now define the trainer itself with the PyTorch Lightning `Trainer`. + +````{dropdown} Click to see the code +```python +trainer = pl.Trainer( + max_epochs=CONFIG["max_epochs"], + accelerator="gpu", + devices=1, + callbacks=callbacks, + log_every_n_steps=10, + val_check_interval=50, + enable_progress_bar=True, +) +``` +```` + +```{tip} +If you do not have a GPU available, set `accelerator="cpu"` (and drop `devices=1`) to run the small capped example on CPU. +``` + +(aw_build_model_p4_fit)= +## Train and Evaluate +Call `fit()` to train. + +````{dropdown} Click to see the code +```python +print("\nStarting training...") +trainer.fit( + model, + train_dataloaders=train_loader, + val_dataloaders=val_loader, +) +``` +```` + +After training, print some diagnostics: how many examples were skipped, and which checkpoint Lightning considered best. + +````{dropdown} Click to see the code +```python +print(f"\nFailed examples during training: {len(train_dataset.failed)}") +print(f"Failed examples during val: {len(val_dataset.failed)}") +print(f"Best checkpoint: {best_checkpoint.best_model_path}") +``` +```` + +Finally, evaluate on the held-out test split using the best checkpoint. + +````{dropdown} Click to see the code +```python +print("\nEvaluating on test set...") +trainer.test( + model, + dataloaders=test_loader, + ckpt_path=best_checkpoint.best_model_path, +) + +print("\nDone.") +``` +```` + +(aw_build_model_p4_wrap)= +## Wrapping Up +You have now built a complete, if simplified, machine-learning pipeline with AtomWorks: from raw PDB metadata to a trained pose-generation model. As next steps you might: + +- Remove the `max_*` caps and train on the full splits. +- Add symmetric edges or an equivariant architecture to remove the model's dependence on absolute coordinates (see the warning in [Part 3](how_to_build_a_model_part3.md)). +- Track additional metrics, such as per-example ligand RMSD distributions, or visualize predicted poses. + +(aw_build_model_p4_glossary)= +## Glossary + +collate function — a function that assembles a list of examples into a single batch; here it also drops failed (`None`) examples. + +checkpoint — a saved snapshot of model weights (and optimizer state) that can be reloaded to resume training or run evaluation. + +early stopping — halting training once a monitored metric (here `val/loss`) stops improving, to save compute and reduce overfitting. diff --git a/docs/how_to_build_a_model/index.rst b/docs/how_to_build_a_model/index.rst new file mode 100644 index 00000000..6477ede1 --- /dev/null +++ b/docs/how_to_build_a_model/index.rst @@ -0,0 +1,13 @@ +How to Build a Model with AtomWorks +=================================== + +This tutorial series shows an example of how to build a model from scratch using AtomWorks. + +.. toctree:: + :maxdepth: 2 + :caption: Sections + + how_to_build_a_model_part1 + how_to_build_a_model_part2 + how_to_build_a_model_part3 + how_to_build_a_model_part4 diff --git a/docs/how_to_build_a_model/scripts/corrected/data_cleaning_script.py b/docs/how_to_build_a_model/scripts/corrected/data_cleaning_script.py new file mode 100644 index 00000000..0d029ad3 --- /dev/null +++ b/docs/how_to_build_a_model/scripts/corrected/data_cleaning_script.py @@ -0,0 +1,162 @@ +""" +Script to read in parquet files and clean the data needed for the +"How to Build a Model Using AtomWorks" tutorial. This script is meant to be run +from the command line and takes the path to the parquet files and the path to +your PDB mirror as arguments. It saves the cleaned data and the train/val/test +splits as parquet files for future use in the tutorial. + +Example usage: + python data_cleaning_script.py /path/to/parquet/files /path/to/pdb_mirror + +Last edited: July 9, 2026 +""" + +import os +import sys + +import numpy as np # BUGFIX: np was used (np.where, np.random) but never imported. +import pandas as pd + + +def read_in_parquet_file(path_to_parquets: str | os.PathLike, parquet_file: str) -> pd.DataFrame: + """ + Create a pandas DataFrame from a parquet file. + :param path_to_parquets: Directory containing the parquet files. + :param parquet_file: Name of the parquet file (without the .parquet extension). + :return: A pandas DataFrame with the parquet contents. + """ + file_path = os.path.join(path_to_parquets, (parquet_file + ".parquet")) + + if not os.path.exists(file_path): + raise FileNotFoundError(f"{file_path} not found") + + return pd.read_parquet(file_path) + + +def assign_split(cluster): + """Assign a protein cluster to a split, using the cluster sets defined below.""" + if cluster in train_clusters: + return "train" + if cluster in val_clusters: + return "val" + if cluster in test_clusters: + return "test" + return "unassigned" # Should already be filtered out; useful as a debug flag. + + +# ── Parse command-line arguments ───────────────────────────────────────────── +if len(sys.argv) > 2: + path_to_parquets = sys.argv[1] + PDB_MIRROR_PATH = sys.argv[2] +else: + print( + "To use this script, please provide the path to the parquet files and " + "the path to your PDB mirror as command line arguments." + ) + sys.exit(1) + +# ── Read the parquet files ─────────────────────────────────────────────────── +interfaces = read_in_parquet_file(path_to_parquets, "interfaces") +pn_units = read_in_parquet_file(path_to_parquets, "pn_units") + +# ── Merge PN-unit metadata onto each interface (once per side) ─────────────── +# pn_unit_iid is only unique within a (pdb_id, assembly_id) scope, so merge on +# all three keys, twice (once for each PN unit in the interface). +u1_cols = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "is_polymer", "num_resolved_residues"]].copy() +u2_cols = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "is_polymer", "num_resolved_residues"]].copy() + +u1_cols = u1_cols.rename(columns={ + "pn_unit_iid": "pn_unit_1_iid", + "is_polymer": "u1_is_polymer", + "num_resolved_residues": "u1_num_resolved_residues", +}) +u2_cols = u2_cols.rename(columns={ + "pn_unit_iid": "pn_unit_2_iid", + "is_polymer": "u2_is_polymer", + "num_resolved_residues": "u2_num_resolved_residues", +}) + +df = interfaces.merge(u1_cols, on=["pdb_id", "assembly_id", "pn_unit_1_iid"], how="inner") +df = df.merge(u2_cols, on=["pdb_id", "assembly_id", "pn_unit_2_iid"], how="inner") + +# ── Filter to drug-like protein-ligand interfaces ──────────────────────────── +df = df[df["involves_loi"]] +df = df[df["is_inter_molecule"]] +df = df[~df["involves_metal"]] +df = df[~df["involves_covalent_modification"]] +df = df[df["u1_is_polymer"] != df["u2_is_polymer"]] +df = df[(df["u1_num_resolved_residues"] + df["u2_num_resolved_residues"]) < 200] + +# Remove duplicate interface rows. +df = df.drop_duplicates(subset=["pdb_id", "assembly_id", "pn_unit_1_iid", "pn_unit_2_iid"]) + +# ── Add a unique example identifier ────────────────────────────────────────── +df["example_id"] = ( + df["pdb_id"] + "_" + + df["assembly_id"].astype(str) + "_" + + df["pn_unit_1_iid"] + "_" + + df["pn_unit_2_iid"] +) +assert df["example_id"].nunique() == len(df), "example_id is not unique!" + +# ── Add the path to each structure in the PDB mirror ───────────────────────── +# BUGFIX: the original overwrote PDB_MIRROR_PATH (from sys.argv[2]) with an +# os.environ lookup that fell back to "/PATH/TO/pdb_mirror", silently ignoring +# the command-line argument. Use the argument the user actually passed in. +df["path"] = df["pdb_id"].str.lower().map( + lambda x: f"{PDB_MIRROR_PATH}/{x[1:3]}/{x}.cif.gz" +) + +df.to_parquet("cleaned_data.parquet") + +# ── Split into train/val/test by protein cluster ───────────────────────────── +protein_clusters = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "protein_cluster_30"]].copy() + +df = df.merge( + protein_clusters.rename(columns={ + "pn_unit_iid": "pn_unit_1_iid", + "protein_cluster_30": "u1_cluster", + }), + on=["pdb_id", "assembly_id", "pn_unit_1_iid"], + how="left", +) +df = df.merge( + protein_clusters.rename(columns={ + "pn_unit_iid": "pn_unit_2_iid", + "protein_cluster_30": "u2_cluster", + }), + on=["pdb_id", "assembly_id", "pn_unit_2_iid"], + how="left", +) + +# The polymer side carries the cluster; the ligand side is null. Take whichever +# side is the polymer. +df["protein_cluster"] = np.where(df["u1_is_polymer"], df["u1_cluster"], df["u2_cluster"]) + +# Drop interfaces with no assigned cluster (RNA/DNA, short peptides, low quality). +df = df[df["protein_cluster"].notna()].reset_index(drop=True) + +# Shuffle the unique clusters (seeded for reproducibility) and split 80/10/10. +unique_clusters = df["protein_cluster"].unique() +rng = np.random.default_rng(seed=42) +rng.shuffle(unique_clusters) + +n = len(unique_clusters) +n_train = int(0.8 * n) +n_val = int(0.1 * n) +# test gets the remainder to avoid off-by-one gaps + +train_clusters = set(unique_clusters[:n_train]) +val_clusters = set(unique_clusters[n_train:n_train + n_val]) +test_clusters = set(unique_clusters[n_train + n_val:]) + +df["split"] = df["protein_cluster"].map(assign_split) + +df_train = df[df["split"] == "train"].reset_index(drop=True) +df_val = df[df["split"] == "val"].reset_index(drop=True) +df_test = df[df["split"] == "test"].reset_index(drop=True) + +os.makedirs("splits", exist_ok=True) +df_train.to_parquet("splits/train.parquet", index=False) +df_val.to_parquet("splits/val.parquet", index=False) +df_test.to_parquet("splits/test.parquet", index=False) diff --git a/docs/how_to_build_a_model/scripts/corrected/dataset_and_loader.py b/docs/how_to_build_a_model/scripts/corrected/dataset_and_loader.py new file mode 100644 index 00000000..ebeadc51 --- /dev/null +++ b/docs/how_to_build_a_model/scripts/corrected/dataset_and_loader.py @@ -0,0 +1,74 @@ +""" +Part of the "How to Build a Model Using AtomWorks" tutorial. +Applies a loader and transform pipeline to each example and runs a smoke test. + +Transforms applied: +- RemoveHydrogens +- RemoveUnresolvedAtoms +- CropToPocket +- FeaturizeForDocking + +Last updated: July 9, 2026 +""" + +import pandas as pd +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import Compose + +# NOTE: import matches the tutorial text and assumes you run this script from the +# directory containing transforms.py. (The PR used the absolute package path +# docs.how_to_build_a_model.scripts.transforms, which only resolves when run as a +# module from the repo root.) Keep this consistent with train.py. +from transforms import CropToPocket, FeaturizeForDocking + +# Load in the training data as a pandas DataFrame +df_train = pd.read_parquet("splits/train.parquet") + +# Define the transform pipeline +transforms_pipeline = Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=10.0), + FeaturizeForDocking(), +]) + +# Build the AtomWorks PandasDataset +dataset = PandasDataset( + data=df_train, + name="docking_train", + id_column="example_id", + loader=create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"]), + transform=transforms_pipeline, +) + +################################ +# Print statements for testing # +################################ +print(f"Dataset size: {len(dataset)}") + +example = dataset[0] + +print("\nLoaded one sample successfully.") +print("Sample keys:", list(example.keys())) +print("atomic_numbers:", example["atomic_numbers"].shape, example["atomic_numbers"].dtype) +print("input_coords:", example["input_coords"].shape, example["input_coords"].dtype) +print("target_coords:", example["target_coords"].shape, example["target_coords"].dtype) +print("edge_index:", example["edge_index"].shape, example["edge_index"].dtype) +print("is_ligand:", example["is_ligand"].shape, example["is_ligand"].dtype) + +assert example["atomic_numbers"].ndim == 1 +assert example["target_coords"].ndim == 2 +assert example["target_coords"].shape[1] == 3 +assert example["input_coords"].shape == example["target_coords"].shape +assert example["edge_index"].ndim == 2 +assert example["edge_index"].shape[0] == 2 +assert example["edge_index"].max() < example["atomic_numbers"].shape[0] +assert (example["input_coords"][example["is_ligand"]] == 0).all(), \ + "Ligand coordinates should be zeroed" +assert (example["input_coords"][~example["is_ligand"]] != 0).any(), \ + "Pocket coordinates should not all be zero" + +print("\nFeaturizeForDocking smoke test passed.") diff --git a/docs/how_to_build_a_model/scripts/corrected/model.py b/docs/how_to_build_a_model/scripts/corrected/model.py new file mode 100644 index 00000000..fab71381 --- /dev/null +++ b/docs/how_to_build_a_model/scripts/corrected/model.py @@ -0,0 +1,113 @@ +"""PocketDockGNN model. This file is unchanged from the PR — no bugs were found. +Included here so the corrected scripts run together as a set.""" + +import torch +import torch.nn as nn +import pytorch_lightning as pl + + +class PocketDockGNN(pl.LightningModule): + def __init__( + self, + num_atom_types: int = 119, + hidden_dim: int = 128, + num_layers: int = 3, + learning_rate: float = 1e-3, + ): + super().__init__() + self.save_hyperparameters() + self.learning_rate = learning_rate + self.atom_embedding = nn.Embedding(num_atom_types, hidden_dim) + + self.input_proj = nn.Sequential( + nn.Linear(hidden_dim + 3, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + + self.conv_layers = nn.ModuleList([ + nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + for _ in range(num_layers) + ]) + + self.update_layers = nn.ModuleList([ + nn.Sequential( + nn.Linear(hidden_dim * 2, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + for _ in range(num_layers) + ]) + + self.layer_norms = nn.ModuleList([ + nn.LayerNorm(hidden_dim) + for _ in range(num_layers) + ]) + + self.output_proj = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, 3), + ) + + self.loss_fn = nn.MSELoss() + + def forward( + self, + atomic_numbers: torch.Tensor, + input_coords: torch.Tensor, + edge_index: torch.Tensor, + ) -> torch.Tensor: + x = self.atom_embedding(atomic_numbers) + x = self.input_proj(torch.cat([x, input_coords], dim=-1)) + + src, dst = edge_index[0], edge_index[1] + + for conv, update, norm in zip(self.conv_layers, self.update_layers, self.layer_norms): + messages = conv(x) + agg = torch.zeros_like(x) + agg.scatter_add_(0, dst.unsqueeze(-1).expand(-1, x.size(-1)), messages[src]) + x = x + norm(update(torch.cat([x, agg], dim=-1))) + + pred_coords = self.output_proj(x) + return pred_coords + + def _shared_step(self, batch: dict, stage: str) -> torch.Tensor: + if batch is None: + return None + # Remove the batch dimension added by DataLoader (batch_size=1) + atomic_numbers = batch["atomic_numbers"].squeeze(0) # (N,) + input_coords = batch["input_coords"].squeeze(0) # (N, 3) + target_coords = batch["target_coords"].squeeze(0) # (N, 3) + edge_index = batch["edge_index"].squeeze(0) # (2, E) + is_ligand = batch["is_ligand"].squeeze(0) # (N,) + + pred_coords = self(atomic_numbers, input_coords, edge_index) + loss = self.loss_fn(pred_coords, target_coords) + + with torch.no_grad(): + ligand_rmsd = torch.sqrt( + ((pred_coords[is_ligand] - target_coords[is_ligand]) ** 2) + .sum(dim=-1).mean() + ) + + self.log(f"{stage}/loss", loss, prog_bar=True) + self.log(f"{stage}/ligand_rmsd", ligand_rmsd, prog_bar=True) + + return loss + + def training_step(self, batch: dict, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, "train") + + def validation_step(self, batch: dict, batch_idx: int) -> None: + self._shared_step(batch, "val") + + def test_step(self, batch: dict, batch_idx: int) -> None: + self._shared_step(batch, "test") + + def configure_optimizers(self): + return torch.optim.Adam(self.parameters(), lr=self.learning_rate) diff --git a/docs/how_to_build_a_model/scripts/corrected/train.py b/docs/how_to_build_a_model/scripts/corrected/train.py new file mode 100644 index 00000000..3d608b7b --- /dev/null +++ b/docs/how_to_build_a_model/scripts/corrected/train.py @@ -0,0 +1,203 @@ +import torch +import pandas as pd +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping +from torch.utils.data import DataLoader, Dataset + +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import ConvertToTorch, Compose + +# NOTE: These imports match the tutorial text and assume you run train.py from +# the directory that contains transforms.py and model.py. The PR used absolute +# package paths (docs.how_to_build_a_model.scripts.transforms), which only work +# when run as a module from the repo root. Pick whichever matches how you run it, +# but keep train.py, dataset_and_loader.py, and the tutorial consistent. +from transforms import CropToPocket, FeaturizeForDocking +from model import PocketDockGNN + +torch.set_float32_matmul_precision("medium") +pl.seed_everything(42) + +CONFIG = { + "hidden_dim": 128, + "num_layers": 3, + "learning_rate": 1e-3, + "batch_size": 1, + "max_epochs": 5, + "pocket_radius": 10.0, + "num_workers": 0, + "max_train": 100, + "max_val": 20, + "max_test": 20, +} + +# BUGFIX: TENSOR_KEYS was referenced by collate_fn and build_pipeline in the PR +# but never defined, which raised NameError at runtime. Define it here. +TENSOR_KEYS = [ + "atomic_numbers", + "input_coords", + "target_coords", + "edge_index", + "is_ligand", +] + + +class RobustDataset(Dataset): + def __init__(self, dataset: PandasDataset): + self.dataset = dataset + self.failed = [] + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, idx): + try: + return self.dataset[idx] + except Exception: + self.failed.append(idx) + return None + + +def collate_fn(batch): + batch = [b for b in batch if b is not None] + if len(batch) == 0: + return None + + return { + k: torch.stack([example[k] for example in batch]) + for k in TENSOR_KEYS + if k in batch[0] + } + + +def build_pipeline(radius: float) -> Compose: + return Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=radius), + FeaturizeForDocking(), + ConvertToTorch(keys=TENSOR_KEYS), + ]) + + +def build_dataset(parquet_path: str, name: str, radius: float, max_examples: int = None): + df = pd.read_parquet(parquet_path) + if max_examples is not None: + df = df.head(max_examples).reset_index(drop=True) + + loader = create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"] + ) + + dataset = PandasDataset( + data=df, + name=name, + id_column="example_id", + loader=loader, + transform=build_pipeline(radius), + save_failed_examples_to_dir="failed_examples/", + ) + return RobustDataset(dataset) + + +def build_dataloader(dataset, shuffle: bool) -> DataLoader: + return DataLoader( + dataset, + batch_size=CONFIG["batch_size"], + shuffle=shuffle, + num_workers=CONFIG["num_workers"], + collate_fn=collate_fn, + persistent_workers=False, + ) + + +def main(): + print("Building datasets...") + train_dataset = build_dataset( + "splits/train.parquet", "docking_train", CONFIG["pocket_radius"], CONFIG["max_train"] + ) + val_dataset = build_dataset( + "splits/val.parquet", "docking_val", CONFIG["pocket_radius"], CONFIG["max_val"] + ) + test_dataset = build_dataset( + "splits/test.parquet", "docking_test", CONFIG["pocket_radius"], CONFIG["max_test"] + ) + + print(f" Train: {len(train_dataset):,} examples") + print(f" Val: {len(val_dataset):,} examples") + print(f" Test: {len(test_dataset):,} examples") + + train_loader = build_dataloader(train_dataset, shuffle=True) + val_loader = build_dataloader(val_dataset, shuffle=False) + test_loader = build_dataloader(test_dataset, shuffle=False) + + model = PocketDockGNN( + hidden_dim=CONFIG["hidden_dim"], + num_layers=CONFIG["num_layers"], + learning_rate=CONFIG["learning_rate"], + ) + print(f"\nModel parameters: {sum(p.numel() for p in model.parameters()):,}") + + periodic_checkpoint = ModelCheckpoint( + dirpath="checkpoints/", + filename="pocketdockgnn-{epoch:02d}-{step}", + every_n_train_steps=50, + save_last=True, + verbose=True, + ) + + best_checkpoint = ModelCheckpoint( + dirpath="checkpoints/", + filename="pocketdockgnn-best-{epoch:02d}-{val/loss:.4f}", + monitor="val/loss", + mode="min", + save_top_k=3, + verbose=True, + ) + + callbacks = [ + best_checkpoint, + periodic_checkpoint, + EarlyStopping( + monitor="val/loss", + patience=10, + mode="min", + verbose=True, + ), + ] + + trainer = pl.Trainer( + max_epochs=CONFIG["max_epochs"], + accelerator="gpu", + devices=1, + callbacks=callbacks, + log_every_n_steps=10, + val_check_interval=50, + enable_progress_bar=True, + ) + + print("\nStarting training...") + trainer.fit( + model, + train_dataloaders=train_loader, + val_dataloaders=val_loader, + ) + + print(f"\nFailed examples during training: {len(train_dataset.failed)}") + print(f"Failed examples during val: {len(val_dataset.failed)}") + print(f"Best checkpoint: {best_checkpoint.best_model_path}") + + print("\nEvaluating on test set...") + trainer.test( + model, + dataloaders=test_loader, + ckpt_path=best_checkpoint.best_model_path, + ) + + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/docs/how_to_build_a_model/scripts/corrected/transforms.py b/docs/how_to_build_a_model/scripts/corrected/transforms.py new file mode 100644 index 00000000..60303c91 --- /dev/null +++ b/docs/how_to_build_a_model/scripts/corrected/transforms.py @@ -0,0 +1,126 @@ +from atomworks.ml.transforms.base import Transform +import numpy as np +from biotite.structure import AtomArray +from scipy.spatial import cKDTree +from atomworks.ml.transforms._checks import check_atom_array_annotation +from atomworks.constants import ELEMENT_NAME_TO_ATOMIC_NUMBER + + +class CropToPocket(Transform): + requires_previous_transforms = ["RemoveHydrogens", "RemoveUnresolvedAtoms"] + + def __init__(self, radius: float = 10.0) -> None: + super().__init__() + self.radius = radius + + def forward(self, data: dict) -> dict: + data["atom_array"] = crop_to_pocket( + data["atom_array"], + query_pn_unit_iids=data["query_pn_unit_iids"], + chain_info=data["chain_info"], + radius=self.radius, + ) + return data + + def check_input(self, data: dict) -> None: + assert "atom_array" in data, "Missing atom_array" + assert "query_pn_unit_iids" in data, "Missing query_pn_unit_iids" + assert "chain_info" in data, "Missing chain_info" + + +class FeaturizeForDocking(Transform): + requires_previous_transforms = ["CropToPocket"] + + def check_input(self, data: dict) -> None: + check_atom_array_annotation(data, ["is_ligand"]) + + def forward(self, data: dict) -> dict: + features = featurize_for_docking(data["atom_array"]) + data.update(features) + return data + + +def crop_to_pocket( + atom_array: AtomArray, + query_pn_unit_iids: list, + chain_info: dict, + radius: float = 10.0, +) -> AtomArray: + atom_array = atom_array.copy() + + iid_a, iid_b = query_pn_unit_iids + chain_a = iid_a.split("_")[0] + chain_b = iid_b.split("_")[0] + + a_is_polymer = chain_info.get(chain_a, {}).get("is_polymer", True) + b_is_polymer = chain_info.get(chain_b, {}).get("is_polymer", True) + + if not a_is_polymer: + ligand_iid, protein_iid = iid_a, iid_b + elif not b_is_polymer: + ligand_iid, protein_iid = iid_b, iid_a + else: + mask_a = atom_array.pn_unit_iid == iid_a + mask_b = atom_array.pn_unit_iid == iid_b + if mask_a.sum() <= mask_b.sum(): + ligand_iid, protein_iid = iid_a, iid_b + else: + ligand_iid, protein_iid = iid_b, iid_a + + ligand_mask = atom_array.pn_unit_iid == ligand_iid + protein_mask = atom_array.pn_unit_iid == protein_iid + + ligand_coords = atom_array.coord[ligand_mask] + protein_coords = atom_array.coord[protein_mask] + + if len(ligand_coords) == 0: + raise ValueError(f"Ligand {ligand_iid} has no atoms") + if len(protein_coords) == 0: + raise ValueError(f"Protein {protein_iid} has no atoms") + + tree = cKDTree(protein_coords) + neighbor_indices = tree.query_ball_point(ligand_coords, r=radius) + + total_neighbors = sum(len(n) for n in neighbor_indices) + if total_neighbors == 0: + raise ValueError(f"No protein atoms found within {radius}A of ligand {ligand_iid}") + + pocket_local_indices = np.unique(np.concatenate(neighbor_indices).astype(int)) + protein_global_indices = np.where(protein_mask)[0] + pocket_global_indices = protein_global_indices[pocket_local_indices] + ligand_global_indices = np.where(ligand_mask)[0] + + keep = np.sort(np.concatenate([pocket_global_indices, ligand_global_indices])) + cropped = atom_array[keep] + + # BUGFIX: the original PR computed `is_ligand` but never attached it to the + # cropped AtomArray, so FeaturizeForDocking.check_input / featurize_for_docking + # would fail (no `is_ligand` annotation). Set the annotation before returning. + is_ligand = np.isin(keep, ligand_global_indices) + cropped.set_annotation("is_ligand", is_ligand) + + return cropped + + +def featurize_for_docking(atom_array: AtomArray) -> dict: + is_ligand = atom_array.is_ligand.astype(bool) + target_coords = atom_array.coord.astype(np.float32) + + input_coords = target_coords.copy() + input_coords[is_ligand] = 0.0 + + atomic_numbers = np.array( + [ELEMENT_NAME_TO_ATOMIC_NUMBER.get(e.upper(), 0) for e in atom_array.element], + dtype=np.int64, + ) + + bonds = atom_array.bonds.as_array() + edge_index = bonds[:, :2].T.astype(np.int64) + + return { + "atomic_numbers": atomic_numbers, + "is_ligand": is_ligand, + "target_coords": target_coords, + "edge_index": edge_index, + "input_coords": input_coords, + } diff --git a/docs/how_to_build_a_model/scripts/data_cleaning_script.py b/docs/how_to_build_a_model/scripts/data_cleaning_script.py new file mode 100644 index 00000000..e1092c44 --- /dev/null +++ b/docs/how_to_build_a_model/scripts/data_cleaning_script.py @@ -0,0 +1,189 @@ +""" +Script to read in parquet files and clean the data needed for the "How to Build a Model Using AtomWorks" tutorial. This script is meant to be run from the command line and takes in the path to the parquet files as an argument. It saves the cleaned data as a parquet file for future use in the tutorial. + +Example usage: + python data_cleaning_script.py /path/to/parquet/files + +Last edited: April 14, 2026 +""" + +import pandas as pd +import os +import sys + +def read_in_parquet_file(path_to_parquets: str | os.PathLike, parquet_file: str) -> pd.DataFrame: + """ + Creates a pandas DataFrame from a parquet file. + :param path_to_parquets: The path to the directory containing the parquet files. + :param parquet_file: The name of the parquet file (without the .parquet extension). + :return: A pandas DataFrame containing the data from the parquet file. + """ + file_path = os.path.join(path_to_parquets, (parquet_file + ".parquet")) + + if not os.path.exists(file_path): + raise FileNotFoundError(f"{file_path} not found") + + mydf = pd.read_parquet(file_path) + + return mydf + +path_to_parquets = None +PDB_MIRROR_PATH = None + +if len(sys.argv) > 2: + path_to_parquets = sys.argv[1] + PDB_MIRROR_PATH = sys.argv[2] +else: + print("To use this script, please provide the path to the parquet files and the path to your PDB mirror as command line arguments.") + sys.exit(1) + +def assign_split(cluster): + """ + Helper function to assign each cluster to a split (train, val, or test) based on the sets of clusters we defined earlier. + :param cluster: The protein cluster to assign a split to. + :return: The split that the cluster belongs to (train, val, test, or unassigned). + """ + # TODO: shouldn't we already have removed any unassigned clusters in the previous step? + if cluster in train_clusters: return "train" + if cluster in val_clusters: return "val" + if cluster in test_clusters: return "test" + return "unassigned" # rows where cluster was null. For this tutorial, these should already be filtered out, but it's a good debugging tool! + + +#################################################### +# Reading in the parquet files and cleaning the data. +#################################################### + +interfaces = read_in_parquet_file(path_to_parquets, "interfaces") +pn_units = read_in_parquet_file(path_to_parquets, "pn_units") + +# We need to combine the information in the interfaces and pn_units datasets to get the data for each pn_unit involved in each interface. +# However, the label "pn_unit_iid" in the pn_units dataset is not unique - it is only unique for a given pdb_id and assembly_id. +# So when we merge these dataframes we need to merge on all three keys and we need to do it twice - once for the first pn_unit in each interface and once for the second pn_unit in each interface. + +# Saving the particular columns we want to merge for the first set of pn_units +# We also keep the "is_polymer" column to filter our data later. +u1_cols = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "is_polymer", "num_resolved_residues"]].copy() + +u2_cols = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "is_polymer", "num_resolved_residues"]].copy() + +# Renaming the columns in these new dataframes to merge on the pn_unit_iid labels (they are called "pn_unit_1_iid" and "pn_unit_2_iid" in the interfaces dataset) and to distinguish the other columns for the first and second pn_units in each interface. +u1_cols = u1_cols.rename(columns={"pn_unit_iid": "pn_unit_1_iid", + "is_polymer": "u1_is_polymer", + "num_resolved_residues": "u1_num_resolved_residues"}) + +u2_cols = u2_cols.rename(columns={"pn_unit_iid": "pn_unit_2_iid", + "is_polymer": "u2_is_polymer", + "num_resolved_residues": "u2_num_resolved_residues"}) + +# Merge these three dataframes together to get the data for each pn_unit involved in each interface. + +df = interfaces.merge(u1_cols, on=["pdb_id", "assembly_id", "pn_unit_1_iid"], how="inner") +df = df.merge(u2_cols, on=["pdb_id", "assembly_id", "pn_unit_2_iid"], how="inner") + +# Now we will filter the data to only keep drug-like protein-ligand pockets: +# involves_loi==True so that at least one of the pn_units is a "ligand of interest" +df = df[df["involves_loi"]] +# is_inter_molecule==True to ensure the two pn_units belong to different molecules +df = df[df["is_inter_molecule"]] +# involves_metal==False to avoid interfaces that are metal-mediated +df = df[~df["involves_metal"]] +# involves_covalent_modification==False to avoid covalently modified residues +df = df[~df["involves_covalent_modification"]] +# only keep interfaces where one pn_unit is a polymer and the other is not (should be a small molecule ligand). This wil avoid protein-protein, ligand-ligand, and RNA/DNA interfaces. +df = df[df["u1_is_polymer"] != df["u2_is_polymer"]] +# total of num_resolved_residues < 200 to keep the size tractable for training +df = df[(df["u1_num_resolved_residues"] + df["u2_num_resolved_residues"]) < 200] + +# Remove duplicate rows +df = df.drop_duplicates(subset=["pdb_id", "assembly_id", "pn_unit_1_iid", "pn_unit_2_iid"]) + +# Add a unique identifier for each interface. This will be useful later in the tutorial. +df["example_id"] = ( + df["pdb_id"] + "_" + + df["assembly_id"].astype(str) + "_" + + df["pn_unit_1_iid"] + "_" + + df["pn_unit_2_iid"] +) +assert df["example_id"].nunique() == len(df), "example_id is not unique!" + +# Add path information to match each entry to the location of the files in your PDB mirror. +PDB_MIRROR_PATH = os.environ.get("PDB_MIRROR_PATH", "/PATH/TO/pdb_mirror") +df["path"] = df["pdb_id"].str.lower().map( + lambda x: f"{PDB_MIRROR_PATH}/{x[1:3]}/{x}.cif.gz" +) + +# Save the cleaned data as a parquet file for future use in the tutorial. +df.to_parquet("cleaned_data.parquet") + +#################################################### +# Splitting the data into training, testing, and validation sets. +#################################################### +# There are many ways to do this, but for the tutorial we have decided to use the information stored in the column "protein_cluster_30" to split the data. This column groups proteins by 30% sequence identity. We will use this information to ensure that all interfaces involving proteins from the same cluster are in the same set (training, testing, or validation). This will help us to better evaluate the generalizability of our model to new proteins. + +protein_clusters = pn_units[["pdb_id", "assembly_id", "pn_unit_iid", "protein_cluster_30"]].copy() + +# merge this with u1 and u2 since either could be the protein side of the interface. We us a left merge for this to only keep the keys from the "left" (df) dataframe. We will keep all rows in df and just add the cluster information where it exists. +# Note: ligands will have a null cluster. + +# Merge cluster for u1 side +df = df.merge( + protein_clusters.rename(columns={ + "pn_unit_iid": "pn_unit_1_iid","protein_cluster_30": "u1_cluster"}), + on=["pdb_id", "assembly_id", "pn_unit_1_iid"], + how="left" +) +# Merge cluster for u2 side +df = df.merge( + protein_clusters.rename(columns={ + "pn_unit_iid": "pn_unit_2_iid", + "protein_cluster_30": "u2_cluster"}), + on=["pdb_id", "assembly_id", "pn_unit_2_iid"], + how="left" +) + +# Get the protein_cluster information from whichever side is the polymer: +df["protein_cluster"] = np.where( + df["u1_is_polymer"], + df["u1_cluster"], + df["u2_cluster"] +) +# TODO ask Hope about this step. It seems like the condition is only for if u1 is a polymer - what if u2 is the polymer? + +# There may be cases where no clusters were assigned, for example low-quality entries. Let's remove these from our dataset: +df = df[df["protein_cluster"].notna()].reset_index(drop=True) + +# Now we can finally split the data based on the protein_cluster information. We will use 80% of the clusters for training, 10% for testing, and 10% for validation. +# Shuffle the the clusters randomly. A seed is specified for reproducibility. +unique_clusters = df["protein_cluster"].dropna().unique() +# TODO: didn't we already drop these in the previous step? +rng = np.random.default_rng(seed=42) +rng.shuffle(unique_clusters) + +# split +n = len(unique_clusters) +n_train = int(0.8 * n) +n_val = int(0.1 * n) +# test gets the remainder to avoid off-by-one gaps + +train_clusters = set(unique_clusters[:n_train]) +val_clusters = set(unique_clusters[n_train : n_train + n_val]) +test_clusters = set(unique_clusters[n_train + n_val :]) + +# create datasets for each +df["split"] = df["protein_cluster"].map(assign_split) + +df_train = df[df["split"] == "train"].reset_index(drop=True) +df_val = df[df["split"] == "val"].reset_index(drop=True) +df_test = df[df["split"] == "test"].reset_index(drop=True) + +# Save datasets as parquet files in separate directories +os.makedirs("splits", exist_ok=True) +df_train.to_parquet("splits/train.parquet", index=False) +df_val.to_parquet("splits/val.parquet", index=False) +df_test.to_parquet("splits/test.parquet", index=False) + + + + + diff --git a/docs/how_to_build_a_model/scripts/dataset_and_loader.py b/docs/how_to_build_a_model/scripts/dataset_and_loader.py new file mode 100644 index 00000000..22864d34 --- /dev/null +++ b/docs/how_to_build_a_model/scripts/dataset_and_loader.py @@ -0,0 +1,71 @@ +""" +This script is part of the "How to Build a Model Using AtomWorks" tutorial. +It uses the AtomWorks API to apply a loader and transform to each 'example' (data point). + +Transforms applied: +- RemoveHydrogens +- RemoveUnresolvedAtoms +- CropToPocket +- FeaturizeFor +- ConvertToTorch + +Last updated: April 14, 2026 +""" + +import pandas as pd +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import Compose + +from docs.how_to_build_a_model.scripts.transforms import CropToPocket, FeaturizeForDocking + +# Load in the training data as a Pandas data frame +df_train = pd.read_parquet("splits/train.parquet") + +# Define our transforms pipeline with two transforms that are already defined in AtomWorks +transforms_pipeline = Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=10.0), + FeaturizeForDocking(), +]) + + +# Use the data to create an AtomWorks PandasDataset object +dataset = PandasDataset(data=df_train, + name="docking_train", + id_column="example_id", + loader=create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"]), + transform=transforms_pipeline, +) + +################################ +# Print statements for testing # +################################ +print(f"Dataset size: {len(dataset)}") + +example = dataset[0] + +print("\nLoaded one sample successfully.") +print("Sample keys:", list(example.keys())) +print("atomic_numbers:", example["atomic_numbers"].shape, example["atomic_numbers"].dtype) +print("input_coords:", example["input_coords"].shape, example["input_coords"].dtype) +print("target_coords:", example["target_coords"].shape, example["target_coords"].dtype) +print("edge_index:", example["edge_index"].shape, example["edge_index"].dtype) +print("is_ligand:", example["is_ligand"].shape, example["is_ligand"].dtype) + +assert example["atomic_numbers"].ndim == 1 +assert example["target_coords"].ndim == 2 +assert example["target_coords"].shape[1] == 3 +assert example["input_coords"].shape == example["target_coords"].shape +assert example["edge_index"].ndim == 2 +assert example["edge_index"].shape[0] == 2 +assert example["edge_index"].max() < example["atomic_numbers"].shape[0] +assert (example["input_coords"][example["is_ligand"]] == 0).all(), \ + "Ligand coordinates should be zeroed" +assert (example["input_coords"][~example["is_ligand"]] != 0).any(), \ + "Pocket coordinates should not all be zero" + +print("\nFeaturizeForDocking smoke test passed.") diff --git a/docs/how_to_build_a_model/scripts/model.py b/docs/how_to_build_a_model/scripts/model.py new file mode 100644 index 00000000..d5060e66 --- /dev/null +++ b/docs/how_to_build_a_model/scripts/model.py @@ -0,0 +1,119 @@ +import torch +import torch.nn as nn +import pytorch_lightning as pl + +class PocketDockGNN(pl.LightningModule): + def __init__( + self, + num_atom_types: int = 119, + hidden_dim: int = 128, + num_layers: int = 3, + learning_rate: float = 1e-3, + ): + super().__init__() + self.save_hyperparameters() + self.learning_rate = learning_rate + self.atom_embedding = nn.Embedding(num_atom_types, hidden_dim) + + self.input_proj = nn.Sequential( + nn.Linear(hidden_dim + 3, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + + self.conv_layers = nn.ModuleList([ + nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + for _ in range(num_layers) + ]) + + self.update_layers = nn.ModuleList([ + nn.Sequential( + nn.Linear(hidden_dim * 2, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + for _ in range(num_layers) + ]) + + self.layer_norms = nn.ModuleList([ + nn.LayerNorm(hidden_dim) + for _ in range(num_layers) + ]) + + self.output_proj = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, 3), + ) + + self.loss_fn = nn.MSELoss() + + def forward( + self, + atomic_numbers: torch.Tensor, + input_coords: torch.Tensor, + edge_index: torch.Tensor, + ) -> torch.Tensor: + x = self.atom_embedding(atomic_numbers) + + x = self.input_proj(torch.cat([x, input_coords], dim=-1)) + + src, dst = edge_index[0], edge_index[1] + + for conv, update, norm in zip(self.conv_layers, self.update_layers, self.layer_norms): + messages = conv(x) + + agg = torch.zeros_like(x) + agg.scatter_add_(0, dst.unsqueeze(-1).expand(-1, x.size(-1)), messages[src]) + + x = x + norm(update(torch.cat([x, agg], dim=-1))) + + pred_coords = self.output_proj(x) + return pred_coords + + def _shared_step(self, batch: dict, stage: str) -> torch.Tensor: + if batch is None: + return None + # Remove the batch dimension added by DataLoader (batch_size=1) + atomic_numbers = batch["atomic_numbers"].squeeze(0) # (N,) + input_coords = batch["input_coords"].squeeze(0) # (N, 3) + target_coords = batch["target_coords"].squeeze(0) # (N, 3) + edge_index = batch["edge_index"].squeeze(0) # (2, E) + is_ligand = batch["is_ligand"].squeeze(0) # (N,) + + pred_coords = self(atomic_numbers, input_coords, edge_index) + loss = self.loss_fn(pred_coords, target_coords) + + with torch.no_grad(): + ligand_rmsd = torch.sqrt( + ((pred_coords[is_ligand] - target_coords[is_ligand]) ** 2) + .sum(dim=-1).mean() + ) + + self.log(f"{stage}/loss", loss, prog_bar=True) + self.log(f"{stage}/ligand_rmsd", ligand_rmsd, prog_bar=True) + + return loss + + def training_step(self, batch: dict, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, "train") + + def validation_step(self, batch: dict, batch_idx: int) -> None: + self._shared_step(batch, "val") + + def test_step(self, batch: dict, batch_idx: int) -> None: + self._shared_step(batch, "test") + + def configure_optimizers(self): + return torch.optim.Adam(self.parameters(), lr=self.learning_rate) + + + + + + + diff --git a/docs/how_to_build_a_model/scripts/train.py b/docs/how_to_build_a_model/scripts/train.py new file mode 100644 index 00000000..86d6150f --- /dev/null +++ b/docs/how_to_build_a_model/scripts/train.py @@ -0,0 +1,180 @@ +import torch +import pandas as pd +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping +from torch.utils.data import DataLoader, Dataset + +from atomworks.ml.datasets import PandasDataset +from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units +from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms +from atomworks.ml.transforms.base import ConvertToTorch, Compose + +from docs.how_to_build_a_model.scripts.transforms import CropToPocket, FeaturizeForDocking +from docs.how_to_build_a_model.scripts.model import PocketDockGNN + +torch.set_float32_matmul_precision("medium") +pl.seed_everything(42) + +CONFIG = { + "hidden_dim": 128, + "num_layers": 3, + "learning_rate": 1e-3, + "batch_size": 1, + "max_epochs": 5, + "pocket_radius": 10.0, + "num_workers": 0, + "max_train": 100, + "max_val": 20, + "max_test": 20, +} + +class RobustDataset(Dataset): + def __init__(self, dataset: PandasDataset): + self.dataset = dataset + self.failed = [] + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, idx): + try: + return self.dataset[idx] + except Exception: + self.failed.append(idx) + return None + + +def collate_fn(batch): + batch = [b for b in batch if b is not None] + if len(batch) == 0: + return None + + return { + k: torch.stack([example[k] for example in batch]) + for k in TENSOR_KEYS + if k in batch[0] + } + + +def build_pipeline(radius: float) -> Compose: + return Compose([ + RemoveHydrogens(), + RemoveUnresolvedAtoms(), + CropToPocket(radius=radius), + FeaturizeForDocking(), + ConvertToTorch(keys=TENSOR_KEYS), + ]) + +def build_dataset(parquet_path: str, name: str, radius: float, max_examples: int = None): + df = pd.read_parquet(parquet_path) + if max_examples is not None: + df = df.head(max_examples).reset_index(drop=True) + + loader = create_loader_with_query_pn_units( + pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"] + ) + + dataset = PandasDataset( + data=df, + name=name, + id_column="example_id", + loader=loader, + transform=build_pipeline(radius), + save_failed_examples_to_dir="failed_examples/", + ) + return RobustDataset(dataset) + +def build_dataloader(dataset, shuffle: bool) -> DataLoader: + return DataLoader( + dataset, + batch_size=CONFIG["batch_size"], + shuffle=shuffle, + num_workers=CONFIG["num_workers"], + collate_fn=collate_fn, + persistent_workers=False, + ) + +print("Building datasets...") +train_dataset = build_dataset( + "splits/train.parquet", "docking_train", CONFIG["pocket_radius"], CONFIG["max_train"] +) +val_dataset = build_dataset( + "splits/val.parquet", "docking_val", CONFIG["pocket_radius"], CONFIG["max_val"] +) +test_dataset = build_dataset( + "splits/test.parquet", "docking_test", CONFIG["pocket_radius"], CONFIG["max_test"] +) + +print(f" Train: {len(train_dataset):,} examples") +print(f" Val: {len(val_dataset):,} examples") +print(f" Test: {len(test_dataset):,} examples") + +train_loader = build_dataloader(train_dataset, shuffle=True) +val_loader = build_dataloader(val_dataset, shuffle=False) +test_loader = build_dataloader(test_dataset, shuffle=False) + + +model = PocketDockGNN( + hidden_dim=CONFIG["hidden_dim"], + num_layers=CONFIG["num_layers"], + learning_rate=CONFIG["learning_rate"], +) + +periodic_checkpoint = ModelCheckpoint( + dirpath="checkpoints/", + filename="pocketdockgnn-{epoch:02d}-{step}", + every_n_train_steps=50, + save_last=True, + verbose=True, +) + +best_checkpoint = ModelCheckpoint( + dirpath="checkpoints/", + filename="pocketdockgnn-best-{epoch:02d}-{val/loss:.4f}", + monitor="val/loss", + mode="min", + save_top_k=3, + verbose=True, +) + +callbacks = [ + best_checkpoint, + periodic_checkpoint, + EarlyStopping( + monitor="val/loss", + patience=10, + mode="min", + verbose=True, + ), +] + +trainer = pl.Trainer( + max_epochs=CONFIG["max_epochs"], + accelerator="gpu", + devices=1, + callbacks=callbacks, + log_every_n_steps=10, + val_check_interval=50, + enable_progress_bar=True, +) +print("\nStarting training...") +trainer.fit( + model, + train_dataloaders=train_loader, + val_dataloaders=val_loader, +) + +print(f"\nFailed examples during training: {len(train_dataset.failed)}") +print(f"Failed examples during val: {len(val_dataset.failed)}") +print(f"Best checkpoint: {best_checkpoint.best_model_path}") + +print("\nEvaluating on test set...") +trainer.test( + model, + dataloaders=test_loader, + ckpt_path=best_checkpoint.best_model_path, +) + +print("\nDone.") + + diff --git a/docs/how_to_build_a_model/scripts/transforms.py b/docs/how_to_build_a_model/scripts/transforms.py new file mode 100644 index 00000000..e37b441c --- /dev/null +++ b/docs/how_to_build_a_model/scripts/transforms.py @@ -0,0 +1,117 @@ +from atomworks.ml.transforms.base import Transform +import numpy as np +from biotite.structure import AtomArray +from scipy.spatial import cKDTree +from atomworks.ml.transforms._checks import check_atom_array_annotation +from atomworks.constants import ELEMENT_NAME_TO_ATOMIC_NUMBER + +class CropToPocket(Transform): + requires_previous_transforms = ["RemoveHydrogens", "RemoveUnresolvedAtoms"] + def __init__(self, radius: float = 10.0) -> None: + super().__init__() + self.radius = radius + + def forward(self, data: dict) -> dict: + data["atom_array"] = crop_to_pocket( + data["atom_array"], + query_pn_unit_iids=data["query_pn_unit_iids"], + chain_info=data["chain_info"], + radius=self.radius, + ) + return data + + def check_input(self, data: dict) -> None: + assert "atom_array" in data, "Missing atom_array" + assert "query_pn_unit_iids" in data, "Missing query_pn_unit_iids" + assert "chain_info" in data, "Missing chain_info" + +class FeaturizeForDocking(Transform): + requires_previous_transforms = ["CropToPocket"] + + def check_input(self, data: dict) -> None: + check_atom_array_annotation(data, ["is_ligand"]) + + def forward(self, data: dict) -> dict: + features = featurize_for_docking(data["atom_array"]) + data.update(features) + return data + +def crop_to_pocket( + atom_array: AtomArray, + query_pn_unit_iids: list, + chain_info: dict, + radius: float = 10.0, +) -> AtomArray: + atom_array = atom_array.copy() + + iid_a, iid_b = query_pn_unit_iids + chain_a = iid_a.split("_")[0] + chain_b = iid_b.split("_")[0] + + a_is_polymer = chain_info.get(chain_a, {}).get("is_polymer", True) + b_is_polymer = chain_info.get(chain_b, {}).get("is_polymer", True) + + if not a_is_polymer: + ligand_iid, protein_iid = iid_a, iid_b + elif not b_is_polymer: + ligand_iid, protein_iid = iid_b, iid_a + else: + mask_a = atom_array.pn_unit_iid == iid_a + mask_b = atom_array.pn_unit_iid == iid_b + if mask_a.sum() <= mask_b.sum(): + ligand_iid, protein_iid = iid_a, iid_b + else: + ligand_iid, protein_iid = iid_b, iid_a + + ligand_mask = atom_array.pn_unit_iid == ligand_iid + protein_mask = atom_array.pn_unit_iid == protein_iid + + ligand_coords = atom_array.coord[ligand_mask] + protein_coords = atom_array.coord[protein_mask] + + if len(ligand_coords) == 0: + raise ValueError(f"Ligand {ligand_iid} has no atoms") + if len(protein_coords) == 0: + raise ValueError(f"Protein {protein_iid} has no atoms") + + tree = cKDTree(protein_coords) + neighbor_indices = tree.query_ball_point(ligand_coords, r=radius) + + total_neighbors = sum(len(n) for n in neighbor_indices) + if total_neighbors == 0: + raise ValueError(f"No protein atoms found within {radius}A of ligand {ligand_iid}") + + pocket_local_indices = np.unique(np.concatenate(neighbor_indices).astype(int)) + protein_global_indices = np.where(protein_mask)[0] + pocket_global_indices = protein_global_indices[pocket_local_indices] + ligand_global_indices = np.where(ligand_mask)[0] + + keep = np.sort(np.concatenate([pocket_global_indices, ligand_global_indices])) + is_ligand = np.isin(keep, ligand_global_indices) + cropped = atom_array[keep] + + return cropped + +def featurize_for_docking(atom_array: AtomArray) -> dict: + is_ligand = atom_array.is_ligand.astype(bool) + target_coords = atom_array.coord.astype(np.float32) + + input_coords = target_coords.copy() + input_coords[is_ligand] = 0.0 + + atomic_numbers = np.array( + [ELEMENT_NAME_TO_ATOMIC_NUMBER.get(e.upper(), 0) for e in atom_array.element], + dtype=np.int64, + ) + + bonds = atom_array.bonds.as_array() + edge_index = bonds[:, :2].T.astype(np.int64) + + return { + "atomic_numbers": atomic_numbers, + "is_ligand": is_ligand, + "target_coords": target_coords, + "edge_index": edge_index, + "input_coords": input_coords, + } + diff --git a/docs/index.rst b/docs/index.rst index 5bdf7632..d98c0596 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -4,7 +4,7 @@ contain the root `toctree` directive. atomworks documentation -====================== +======================= Welcome to **atomworks** — a toolkit for converting, parsing, and manipulating biological structure and sequence data, inspired by the Biotite library. Quickly convert between formats, extract features, and prepare data for machine learning or structural analysis. @@ -13,6 +13,12 @@ Welcome to **atomworks** — a toolkit for converting, parsing, and manipulating :alt: atomworks datapipelines +.. toctree:: + :maxdepth: 2 + :caption: Building a Model with AtomWorks + + how_to_build_a_model/index + .. toctree:: :maxdepth: 2 :caption: Navigation diff --git a/docs/io/transforms.rst b/docs/io/transforms.rst index 33dc6380..e38b74e7 100644 --- a/docs/io/transforms.rst +++ b/docs/io/transforms.rst @@ -1,5 +1,5 @@ Transforms -========= +========== .. toctree:: :maxdepth: 2 diff --git a/docs/io/transforms/atom_array.rst b/docs/io/transforms/atom_array.rst index e39345fd..16501695 100644 --- a/docs/io/transforms/atom_array.rst +++ b/docs/io/transforms/atom_array.rst @@ -1,5 +1,5 @@ Atom Array Transforms -==================== +===================== .. automodule:: atomworks.io.transforms.atom_array :members: diff --git a/docs/io/transforms/categories.rst b/docs/io/transforms/categories.rst index 65789f8d..15b73125 100644 --- a/docs/io/transforms/categories.rst +++ b/docs/io/transforms/categories.rst @@ -1,5 +1,5 @@ Category Transforms -================== +=================== .. automodule:: atomworks.io.transforms.categories :members: diff --git a/docs/io/utils/assembly.rst b/docs/io/utils/assembly.rst index 2627227e..8f707686 100644 --- a/docs/io/utils/assembly.rst +++ b/docs/io/utils/assembly.rst @@ -1,5 +1,5 @@ Assembly Utilities -================= +================== .. automodule:: atomworks.io.utils.assembly :members: diff --git a/docs/io/utils/bonds.rst b/docs/io/utils/bonds.rst index c51fcaf6..91bf7dc4 100644 --- a/docs/io/utils/bonds.rst +++ b/docs/io/utils/bonds.rst @@ -1,5 +1,5 @@ Bond Utilities -============= +============== .. automodule:: atomworks.io.utils.bonds :members: diff --git a/docs/io/utils/chain.rst b/docs/io/utils/chain.rst index ff797426..cb993611 100644 --- a/docs/io/utils/chain.rst +++ b/docs/io/utils/chain.rst @@ -1,5 +1,5 @@ Chain Utilities -============== +=============== .. automodule:: atomworks.io.utils.chain :members: diff --git a/docs/io/utils/non_rcsb.rst b/docs/io/utils/non_rcsb.rst index c874c7ee..771b813b 100644 --- a/docs/io/utils/non_rcsb.rst +++ b/docs/io/utils/non_rcsb.rst @@ -1,5 +1,5 @@ Non-RCSB Utilities -================ +================== .. automodule:: atomworks.io.utils.non_rcsb :members: diff --git a/docs/io/utils/selection.rst b/docs/io/utils/selection.rst index 8dafd0f1..70f8aee2 100644 --- a/docs/io/utils/selection.rst +++ b/docs/io/utils/selection.rst @@ -1,5 +1,5 @@ Selection Utilities -================= +=================== .. automodule:: atomworks.io.utils.selection :members: diff --git a/docs/io/utils/testing.rst b/docs/io/utils/testing.rst index b0584d68..9b3cb20e 100644 --- a/docs/io/utils/testing.rst +++ b/docs/io/utils/testing.rst @@ -1,5 +1,5 @@ Testing Utilities -=============== +================= .. automodule:: atomworks.io.utils.testing :members: diff --git a/docs/io/utils/visualize.rst b/docs/io/utils/visualize.rst index feb45a7c..d66dbccb 100644 --- a/docs/io/utils/visualize.rst +++ b/docs/io/utils/visualize.rst @@ -1,5 +1,5 @@ Visualization Utilities -==================== +======================= .. automodule:: atomworks.io.utils.visualize :members: diff --git a/docs/ml.rst b/docs/ml.rst index 4618eb95..d747f805 100644 --- a/docs/ml.rst +++ b/docs/ml.rst @@ -11,7 +11,7 @@ Core Modules ml/samplers Data Processing Modules ----------------------- +----------------------- .. toctree:: :maxdepth: 2 diff --git a/docs/ml/transforms/feature_aggregation.rst b/docs/ml/transforms/feature_aggregation.rst index 973d69a9..ec2fbcf9 100644 --- a/docs/ml/transforms/feature_aggregation.rst +++ b/docs/ml/transforms/feature_aggregation.rst @@ -1,5 +1,5 @@ Feature Aggregation Transforms -============================ +============================== This module contains transformations for feature aggregation tasks. diff --git a/docs/ml/transforms/msa.rst b/docs/ml/transforms/msa.rst index 36f40167..aad06b4c 100644 --- a/docs/ml/transforms/msa.rst +++ b/docs/ml/transforms/msa.rst @@ -1,5 +1,5 @@ MSA Transforms -============= +============== This module contains transformations for multiple sequence alignment (MSA) tasks.