rc-foundry 0.1.1__py3-none-any.whl

mpnn/pipelines/mpnn.py

@@ -0,0 +1,162 @@
+from atomworks.constants import AF3_EXCLUDED_LIGANDS, STANDARD_AA, UNKNOWN_AA
+from atomworks.enums import ChainTypeInfo
+from atomworks.ml.transforms.atom_array import AddWithinChainInstanceResIdx
+from atomworks.ml.transforms.atomize import (
+    AtomizeByCCDName,
+    FlagNonPolymersForAtomization,
+)
+from atomworks.ml.transforms.base import (
+    AddData,
+    Compose,
+    ConditionalRoute,
+    ConvertToTorch,
+    Identity,
+    SubsetToKeys,
+)
+from atomworks.ml.transforms.covalent_modifications import (
+    FlagAndReassignCovalentModifications,
+)
+from atomworks.ml.transforms.featurize_unresolved_residues import (
+    MaskResiduesWithSpecificUnresolvedAtoms,
+)
+from atomworks.ml.transforms.filters import (
+    FilterToSpecifiedPNUnits,
+    HandleUndesiredResTokens,
+    RemoveHydrogens,
+    RemoveUnresolvedTokens,
+)
+from mpnn.transforms.feature_aggregation.mpnn import (
+    EncodeMPNNNonAtomizedTokens,
+    FeaturizeAtomizedTokens,
+    FeaturizeNonAtomizedTokens,
+)
+from mpnn.transforms.feature_aggregation.user_settings import (
+    FeaturizeUserSettings,
+)
+
+
+def TrainingRoute(transform):
+    return ConditionalRoute(
+        condition_func=lambda data: data["is_inference"],
+        transform_map={True: Identity(), False: transform},
+    )
+
+
+def InferenceRoute(transform):
+    return ConditionalRoute(
+        condition_func=lambda data: data["is_inference"],
+        transform_map={False: Identity(), True: transform},
+    )
+
+
+def ModelTypeRoute(transform, model_type: str):
+    return ConditionalRoute(
+        condition_func=lambda data: data["model_type"] == model_type,
+        transform_map={True: transform, False: Identity()},
+    )
+
+
+def build_mpnn_transform_pipeline(
+    *,
+    model_type: str = None,
+    occupancy_threshold_sidechain: float = 0.5,
+    occupancy_threshold_backbone: float = 0.8,
+    is_inference: bool = False,
+    minimal_return: bool = False,
+    train_structure_noise_default: float = 0.1,
+    undesired_res_names: list[str] = AF3_EXCLUDED_LIGANDS,
+    device=None,
+) -> Compose:
+    """Build the MPNN transform pipeline.
+    Args:
+        model_type (str): Model type identifier to include in data. Must be
+            provided. Defaults to None.
+        occupancy_threshold_sidechain (float): Minimum occupancy to consider
+            sidechain atoms as present in masks. Defaults to 0.5.
+        occupancy_threshold_backbone (float): Minimum occupancy to consider
+            backbone atoms as resolved. Residues with backbone atoms below this
+            threshold will be masked entirely. Defaults to 0.8.
+        train_structure_noise_default (float): Default standard deviation of
+            Gaussian noise to add to atomic coordinates during training for data
+            augmentation. Defaults to 0.1.
+        is_inference (bool): Whether this is inference mode. Defaults to
+            False (training mode).
+        minimal_return (bool): Whether to return minimal intermediate data.
+            Defaults to False.
+        undesired_res_names (list[str]): List of residue names to treat as
+            undesired and handle accordingly. Defaults to AF3_EXCLUDED_LIGANDS.
+        device (str | torch.device, optional): Device to move tensors to.
+            Defaults to None, which leads to default ConvertToTorch behavior.
+    """
+    if model_type not in ("protein_mpnn", "ligand_mpnn"):
+        raise ValueError(f"Unsupported model_type: {model_type}")
+
+    transforms = [
+        AddData({"model_type": model_type}),
+        AddData({"is_inference": is_inference}),
+        # + --------- Filters --------- +
+        RemoveHydrogens(),
+        # ... during training, filter to non-clashing chains (which are
+        # pre-computed and stored in the "extra_info" key)
+        TrainingRoute(
+            FilterToSpecifiedPNUnits(
+                extra_info_key_with_pn_unit_iids_to_keep="all_pn_unit_iids_after_processing"
+            ),
+        ),
+        # ... during training, remove undesired residues (e.g., non-biological
+        # crystallization artifacts), mapping to the closest canonical residue
+        # name where possible
+        TrainingRoute(
+            HandleUndesiredResTokens(undesired_res_tokens=undesired_res_names),
+        ),
+        # + --------- Atomization --------- +
+        # ... add within-chain instance res idx
+        AddWithinChainInstanceResIdx(),
+        FlagAndReassignCovalentModifications(),
+        # Atomization: keep standard AA + unknown AA as residues,
+        # atomize everything else
+        FlagNonPolymersForAtomization(),
+        AtomizeByCCDName(
+            atomize_by_default=True,
+            res_names_to_ignore=STANDARD_AA + (UNKNOWN_AA,),
+            move_atomized_part_to_end=False,
+            validate_atomize=False,
+        ),
+        # + --------- Occupancy filtering --------- +
+        MaskResiduesWithSpecificUnresolvedAtoms(
+            chain_type_to_atom_names={
+                ChainTypeInfo.PROTEINS: [
+                    "N",
+                    "CA",
+                    "C",
+                    "O",
+                ],  # MPNN needs backbone + oxygen
+            },
+            occupancy_threshold=occupancy_threshold_backbone,
+        ),
+        RemoveUnresolvedTokens(),
+        # +-------- Encoding and featurization --------- +
+        AddData({"input_features": dict()}),
+        # Encode and featurize non-atomized tokens
+        EncodeMPNNNonAtomizedTokens(occupancy_threshold=occupancy_threshold_sidechain),
+        FeaturizeNonAtomizedTokens(),
+        # LigandMPNN specific featurization: featurize atomized tokens
+        ModelTypeRoute(
+            transform=FeaturizeAtomizedTokens(),
+            model_type=model_type,
+        ),
+        # Featurize user settings
+        FeaturizeUserSettings(
+            is_inference=is_inference,
+            minimal_return=minimal_return,
+            train_structure_noise_default=train_structure_noise_default,
+        ),
+        # Convert to torch and subset keys
+        ConvertToTorch(
+            keys=["input_features"],
+            **({"device": device} if device is not None else {}),
+        ),
+        SubsetToKeys(keys=["input_features", "atom_array"]),
+    ]
+
+    return Compose(transforms)

mpnn/samplers/samplers.py

@@ -0,0 +1,167 @@
+from typing import Any, Callable, Iterator, List
+
+import numpy as np
+import torch
+from atomworks.ml.samplers import set_sampler_epoch
+from torch.utils.data import BatchSampler, Sampler
+
+
+class PaddedTokenBudgetBatchSampler(BatchSampler):
+    """
+    Token-based batch sampler that wraps existing samplers and creates batches
+    of similar-token length samples, respecting a maximum token count
+    constraint (considering that the batches will be padded to the maximum
+    length in the batch).
+
+    Args:
+        sampler: The underlying sampler to wrap around.
+        get_num_tokens: Function that takes an index from the previous sampler
+            and returns the number of tokens for that sample.
+        max_tokens_with_padding: Maximum number of tokens allowed per batch,
+            including padding. The constraint is
+            max(batch_lengths) * len(batch) <= max_tokens.
+        shuffle_batches: Whether to randomize the order of batches after
+            grouping by length. Defaults to True.
+    """
+
+    def __init__(
+        self,
+        sampler: Sampler,
+        get_num_tokens: Callable[[Any], int],
+        max_tokens_with_padding: int = 6000,
+        shuffle_batches: bool = True,
+    ):
+        # Initialize BatchSampler with a dummy batch_size (we don't use it).
+        super().__init__(sampler, batch_size=1, drop_last=False)
+
+        self.sampler = sampler
+        self.get_num_tokens = get_num_tokens
+        self.max_tokens_with_padding = max_tokens_with_padding
+        self.shuffle_batches = shuffle_batches
+
+        # Add drop_last attribute for DataLoader compatibility
+        self.drop_last = False
+
+        self.epoch = 0
+
+        # Will hold our "one epoch" cache.
+        self._batches: List[List[Any]] | None = None
+
+        # Validate inputs
+        if max_tokens_with_padding <= 0:
+            raise ValueError("max_tokens_with_padding must be greater than 0")
+        if get_num_tokens is None:
+            raise ValueError("get_num_tokens function must be provided")
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+        set_sampler_epoch(self.sampler, epoch)
+
+    def _build_batches(self) -> List[List[Any]]:
+        """
+        Compute all batches and cache them for the current epoch.
+
+        Returns:
+            List of batches, where each batch is a list of indices.
+        Raises:
+            TypeError: If get_num_tokens returns invalid types.
+            ValueError: If get_num_tokens returns invalid values.
+        """
+        # Extract all indices and their token counts
+        sample_indices_and_lengths = []
+        for idx in self.sampler:
+            num_tokens = self.get_num_tokens(idx)
+
+            # Validate num_tokens type
+            if not isinstance(num_tokens, (int, float, np.integer, np.floating)):
+                raise TypeError(
+                    f"get_num_tokens returned invalid type {type(num_tokens)} "
+                    f"for index {idx}. Expected numeric type."
+                )
+
+            # Validate num_tokens value.
+            num_tokens = int(num_tokens)
+            if num_tokens <= 0:
+                raise ValueError(
+                    f"get_num_tokens returned invalid value {num_tokens} "
+                    f"for index {idx}. Expected positive integer."
+                )
+            if num_tokens > self.max_tokens_with_padding:
+                raise ValueError(
+                    f"Index {idx} has {num_tokens} tokens, exceeding "
+                    f"max_tokens_with_padding={self.max_tokens_with_padding}."
+                )
+
+            sample_indices_and_lengths.append((idx, num_tokens))
+
+        # Sort by token length (ascending order).
+        sample_indices_and_lengths.sort(key=lambda x: x[1])
+
+        # Batch by length
+        batches = []
+        current_batch = []
+        current_max_length = 0
+        for idx, length in sample_indices_and_lengths:
+            # Check if adding this sample would violate max_tokens_with_padding
+            # constraint.
+            if current_batch:
+                potential_max_length = max(length, current_max_length)
+                new_batch_size = len(current_batch) + 1
+
+                if potential_max_length * new_batch_size > self.max_tokens_with_padding:
+                    # Current batch is full, start a new batch
+                    batches.append(current_batch)
+
+                    current_batch = [idx]
+                    current_max_length = length
+                else:
+                    # Add to current batch
+                    current_batch.append(idx)
+                    current_max_length = potential_max_length
+            else:
+                # First sample in batch
+                current_batch = [idx]
+                current_max_length = length
+
+        # Add the last batch if it's not empty
+        if current_batch:
+            batches.append(current_batch)
+
+        # Randomize batch order if requested
+        if self.shuffle_batches:
+            # Set the seed based on the epoch
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+
+            perm = torch.randperm(len(batches), generator=g).tolist()
+
+            batches = [batches[i] for i in perm]
+
+        return batches
+
+    def __iter__(self) -> Iterator[List[Any]]:
+        """
+        Generate batches of indices grouped by token length while respecting
+        the max_tokens_with_padding constraint.
+
+        Returns:
+            Iterator[List[Any]]: Iterator over batches of dataset indices.
+        """
+        # Build/reference cached batches if not already done.
+        if self._batches is None:
+            self._batches = self._build_batches()
+
+        for batch in self._batches:
+            yield batch
+
+        # End of __iter__ — clear cache so next call recomputes.
+        self._batches = None
+
+    def __len__(self) -> int:
+        """
+        Return the exact number of batches that will be produced.
+        """
+        # Build/reference cached batches if not already done.
+        if self._batches is None:
+            self._batches = self._build_batches()
+        return len(self._batches)

mpnn/train.py

@@ -0,0 +1,341 @@
+#!/usr/bin/env -S /bin/sh -c '"$(dirname "$0")/../../../../.ipd/shebang/mpnn_exec.sh" "$0" "$@"'
+
+import sys
+from pathlib import Path
+
+import pandas as pd
+import torch
+from atomworks.io.parser import STANDARD_PARSER_ARGS
+from atomworks.ml.datasets.pandas_dataset import PandasDataset, StructuralDatasetWrapper
+from atomworks.ml.datasets.parsers.default_metadata_row_parsers import GenericDFParser
+from atomworks.ml.samplers import (
+    DistributedMixedSampler,
+    calculate_weights_for_pdb_dataset_df,
+)
+from omegaconf import DictConfig
+from torch.utils.data import DataLoader, WeightedRandomSampler
+
+from foundry.callbacks.metrics_logging import StoreValidationMetricsInDFCallback
+from foundry.utils.datasets import wrap_dataset_and_sampler_with_fallbacks
+from mpnn.collate.feature_collator import TokenBudgetAwareFeatureCollator
+from mpnn.pipelines.mpnn import build_mpnn_transform_pipeline
+from mpnn.samplers.samplers import PaddedTokenBudgetBatchSampler
+from mpnn.trainers.mpnn import MPNNTrainer
+
+model_type = sys.argv[1]
+
+if model_type == "protein_mpnn":
+    batch_size = 10000
+    train_date_cutoff = "2021-08-02"
+    clip_grad_max_norm = None
+    train_structure_noise_default = 0.2
+elif model_type == "ligand_mpnn":
+    batch_size = 6000
+    train_date_cutoff = "2022-12-16"
+    clip_grad_max_norm = 1.0
+    train_structure_noise_default = 0.1
+else:
+    raise ValueError(f"Unknown model_type: {model_type}")
+
+
+def create_noam_scheduler(optimizer, d_model, warmup_steps=4000, factor=2):
+    """
+    Create a NoamOpt-style scheduler using standard PyTorch components.
+
+    Args:
+        optimizer: PyTorch optimizer
+        d_model: Model dimension (for scaling)
+        warmup_steps: Number of warmup steps
+        factor: Scaling factor
+
+    Returns:
+        LambdaLR scheduler that implements NoamOpt schedule
+    """
+
+    def noam_lambda(step):
+        # NoamOpt formula: factor * (d_model ** (-0.5)) * min(step ** (-0.5), step * warmup ** (-1.5))
+        base_lr = factor * (d_model ** (-0.5))
+        if step == 0:
+            return 0.0  # Start with zero learning rate
+
+        # Calculate the schedule component
+        schedule = min(step ** (-0.5), step * warmup_steps ** (-1.5))
+        return base_lr * schedule
+
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=noam_lambda)
+
+
+def get_num_tokens(df, idx):
+    """
+    Extract the number of non-atomized tokens for a given index.
+
+    Args:
+        df: DataFrame containing the dataset
+        idx: Index to extract token count from
+
+    Returns:
+        Number of non-atomized tokens for the sample at idx
+    """
+    if isinstance(idx, (list, tuple)):
+        # If idx is a list/tuple, return the first element's token count
+        idx = idx[0]
+    return df.iloc[idx]["n_non_atomized_tokens"]
+
+
+# Common filters for MPNN datasets
+MPNN_FILTERS = [
+    "resolution < 3.5 and ~method.str.contains('NMR')",
+    "n_non_atomized_tokens >= 30",
+    "cluster.notnull() and cluster != 'nan'",
+    "method in ['X-RAY_DIFFRACTION', 'ELECTRON_MICROSCOPY']",
+    f"n_non_atomized_tokens <= {batch_size}",
+    "n_prot == 1",
+]
+
+MPNN_TRAIN_FILTERS = [
+    f"deposition_date < '{train_date_cutoff}'",
+] + MPNN_FILTERS
+
+# Data loading setup
+train_path = "/projects/ml/datahub/dfs/mpnn/splits/2025_07_13/pn_units_df_train.parquet"
+val_path = "/projects/ml/datahub/dfs/mpnn/splits/2025_07_13/pn_units_df_val.parquet"
+
+# Load datasets
+train_df = pd.read_parquet(train_path)
+val_df = pd.read_parquet(val_path)
+
+# Create different pipelines for training and inference
+train_pipeline = build_mpnn_transform_pipeline(
+    model_type=model_type,
+    is_inference=False,
+    minimal_return=True,
+    train_structure_noise_default=train_structure_noise_default,
+)
+inference_pipeline = build_mpnn_transform_pipeline(
+    model_type=model_type, is_inference=True, minimal_return=True
+)
+
+# Create train dataset with fallback
+train_structural_dataset = StructuralDatasetWrapper(
+    dataset=PandasDataset(
+        data=train_df,
+        id_column="example_id",
+        name="pn_units_df_train",
+        filters=MPNN_TRAIN_FILTERS,
+    ),
+    dataset_parser=GenericDFParser(
+        example_id_colname="example_id",
+        path_colname="path",
+        assembly_id_colname="assembly_id",
+    ),
+    transform=train_pipeline,
+    cif_parser_args={
+        **STANDARD_PARSER_ARGS,
+        "load_from_cache": True,
+        "save_to_cache": True,
+        "cache_dir": "/net/tukwila/akubaney/cifutils/cache",
+    },
+)
+
+train_fallback_dataset = StructuralDatasetWrapper(
+    dataset=PandasDataset(
+        data=train_df,
+        id_column="example_id",
+        name="pn_units_df_train",
+        filters=MPNN_TRAIN_FILTERS,
+    ),
+    dataset_parser=GenericDFParser(
+        example_id_colname="example_id",
+        path_colname="path",
+        assembly_id_colname="assembly_id",
+    ),
+    transform=train_pipeline,
+    cif_parser_args={
+        **STANDARD_PARSER_ARGS,
+        "load_from_cache": True,
+        "save_to_cache": True,
+        "cache_dir": "/net/tukwila/akubaney/cifutils/cache",
+    },
+)
+
+# Calculate weights for train dataset
+train_weights = calculate_weights_for_pdb_dataset_df(
+    dataset_df=train_structural_dataset.data,
+    beta=1.0,  # For chains
+    alphas={"a_prot": 1.0, "a_nuc": 0, "a_ligand": 0, "a_loi": 0},
+)
+
+# Create train sampler with fallback
+train_sampler = DistributedMixedSampler(
+    datasets_info=[
+        {
+            "sampler": WeightedRandomSampler(
+                train_weights, len(train_structural_dataset)
+            ),
+            "dataset": train_structural_dataset,
+            "probability": 1.0,
+        }
+    ],
+    num_replicas=1,
+    rank=0,
+    n_examples_per_epoch=20000,
+)
+
+train_fallback_sampler = WeightedRandomSampler(
+    train_weights, len(train_structural_dataset)
+)
+
+train_dataset_with_fallback, train_sampler_with_fallback = (
+    wrap_dataset_and_sampler_with_fallbacks(
+        dataset_to_be_wrapped=train_structural_dataset,
+        sampler_to_be_wrapped=train_sampler,
+        dataset_to_fallback_to=train_fallback_dataset,
+        sampler_to_fallback_to=train_fallback_sampler,
+        n_fallback_retries=5,
+    )
+)
+
+batched_train_sampler = PaddedTokenBudgetBatchSampler(
+    sampler=train_sampler_with_fallback,
+    get_num_tokens=lambda idx: get_num_tokens(train_structural_dataset.data, idx),
+    max_tokens_with_padding=batch_size,
+    shuffle_batches=True,
+)
+
+# Create val dataset with fallback
+val_structural_dataset = StructuralDatasetWrapper(
+    dataset=PandasDataset(
+        data=val_df,
+        id_column="example_id",
+        name="pn_units_df_val",
+        filters=MPNN_FILTERS,
+    ),
+    dataset_parser=GenericDFParser(
+        example_id_colname="example_id",
+        path_colname="path",
+        assembly_id_colname="assembly_id",
+    ),
+    transform=inference_pipeline,
+    cif_parser_args={
+        **STANDARD_PARSER_ARGS,
+        "load_from_cache": True,
+        "save_to_cache": True,
+        "cache_dir": "/net/tukwila/akubaney/cifutils/cache",
+    },
+)
+
+# Create val sampler with fallback
+val_weights = calculate_weights_for_pdb_dataset_df(
+    dataset_df=val_structural_dataset.data,
+    beta=1.0,  # For chains
+    alphas={"a_prot": 1.0, "a_nuc": 0, "a_ligand": 0, "a_loi": 0},
+)
+
+val_sampler = DistributedMixedSampler(
+    datasets_info=[
+        {
+            "sampler": WeightedRandomSampler(val_weights, len(val_structural_dataset)),
+            "dataset": val_structural_dataset,
+            "probability": 1.0,
+        }
+    ],
+    num_replicas=1,
+    rank=0,
+    n_examples_per_epoch=100,
+)
+
+# Create collator
+collator = TokenBudgetAwareFeatureCollator(max_tokens_with_padding=batch_size)
+
+# Create DataLoaders
+train_loader = DataLoader(
+    train_dataset_with_fallback,
+    batch_sampler=batched_train_sampler,
+    num_workers=12,
+    collate_fn=collator,
+)
+
+val_loaders = {
+    "test_val": DataLoader(
+        val_structural_dataset,
+        sampler=val_sampler,
+        num_workers=12,
+        collate_fn=collator,
+    )
+}
+
+# Create output directory for logs and checkpoints
+output_dir = Path(f"./mpnn_output_{model_type}")
+output_dir.mkdir(exist_ok=True)
+
+# Create CSV logging callback
+csv_callback = StoreValidationMetricsInDFCallback(
+    save_dir=output_dir / "val_metrics", metrics_to_save="all"
+)
+
+# Create trainer with minimal configuration for testing
+trainer = MPNNTrainer(
+    model_type=model_type,
+    accelerator="gpu",
+    devices_per_node=1,
+    max_epochs=500,
+    output_dir=output_dir,
+    callbacks=[csv_callback],
+    precision="bf16-mixed",
+    clip_grad_max_norm=clip_grad_max_norm,
+)
+
+# Create minimal train_cfg for optimizer and scheduler construction
+train_cfg = DictConfig(
+    {
+        "model": {
+            "optimizer": {
+                "_target_": "torch.optim.Adam",
+                "lr": 1.0,  # This will be overridden by the NoamOpt scheduler
+                "betas": [0.9, 0.98],  # NoamOpt uses (0.9, 0.98)
+                "eps": 1e-9,  # NoamOpt uses 1e-9
+                "weight_decay": 0.0,
+            },
+            "lr_scheduler": {
+                "_target_": "__main__.create_noam_scheduler",
+                "d_model": 128,  # Adjust based on your model's hidden dimension
+                "warmup_steps": 4000,
+                "factor": 2,
+            },
+        }
+    }
+)
+
+# Initialize trainer state with train_cfg
+trainer.initialize_or_update_trainer_state({"train_cfg": train_cfg})
+
+# Launch Fabric (this sets up the distributed environment)
+trainer.fabric.launch()
+
+# Construct model
+trainer.construct_model()
+
+# Construct optimizer and scheduler
+trainer.construct_optimizer()
+trainer.construct_scheduler()
+
+
+class CkptConfig:
+    def __init__(self, path, weight_loading_config=None, reset_optimizer=False):
+        self.path = path
+        self.weight_loading_config = weight_loading_config
+        self.reset_optimizer = reset_optimizer
+
+
+ckpt_dir = output_dir / "ckpt"
+if ckpt_dir.exists():
+    ckpt_config = CkptConfig(
+        path=ckpt_dir, weight_loading_config=None, reset_optimizer=False
+    )
+else:
+    ckpt_config = None
+
+# Run the full training using fit method
+print("Starting training...")
+trainer.fit(train_loader=train_loader, val_loaders=val_loaders, ckpt_config=ckpt_config)
+print("Training completed!")