rc-foundry 0.1.1__py3-none-any.whl

mpnn/trainers/mpnn.py

@@ -0,0 +1,193 @@
+import torch
+from beartype.typing import Any
+from lightning_utilities import apply_to_collection
+from mpnn.loss.nll_loss import LabelSmoothedNLLLoss
+from mpnn.metrics.nll import NLL, InterfaceNLL
+from mpnn.metrics.sequence_recovery import (
+    InterfaceSequenceRecovery,
+    SequenceRecovery,
+)
+from mpnn.model.mpnn import LigandMPNN, ProteinMPNN
+from omegaconf import DictConfig
+
+from foundry.metrics.metric import MetricManager
+from foundry.trainers.fabric import FabricTrainer
+from foundry.utils.ddp import RankedLogger
+from foundry.utils.torch import assert_no_nans
+
+ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+
+
+class MPNNTrainer(FabricTrainer):
+    """Standard Trainer for MPNN-style models"""
+
+    def __init__(
+        self,
+        *,
+        model_type: str,
+        loss: DictConfig | dict | None = None,
+        metrics: DictConfig | dict | None = None,
+        **kwargs,
+    ):
+        """
+        See `FabricTrainer` for the additional initialization arguments.
+
+        Args:
+            model_type (str): Type of model to use ("protein_mpnn" or
+                "ligand_mpnn").
+            loss (DictConfig | dict | None): Configuration for the loss
+                function. If None, default parameters will be used.
+            metrics (DictConfig | dict | None): Configuration for the metrics.
+                Ignored - metrics are hard-coded.
+        """
+        super().__init__(**kwargs)
+
+        self.model_type = model_type
+
+        # Metrics
+        metrics = {
+            "nll": NLL(),
+            "sequence_recovery": SequenceRecovery(),
+        }
+        if self.model_type == "ligand_mpnn":
+            metrics["interface_nll"] = InterfaceNLL()
+            metrics["interface_sequence_recovery"] = InterfaceSequenceRecovery()
+        self.metrics = MetricManager(metrics)
+
+        # Loss
+        loss_params = loss if loss else {}
+        self.loss = LabelSmoothedNLLLoss(**loss_params)
+
+    def construct_model(self):
+        """Construct the model with hard-coded parameters."""
+        with self.fabric.init_module():
+            ranked_logger.info(f"Instantiating {self.model_type} model...")
+
+            # Hard-coded model selection
+            if self.model_type == "protein_mpnn":
+                model = ProteinMPNN()
+            elif self.model_type == "ligand_mpnn":
+                model = LigandMPNN()
+            else:
+                raise ValueError(f"Invalid model type: {self.model_type}")
+
+            # Initialize model weights
+            model.apply(model.init_weights)
+
+        self.initialize_or_update_trainer_state({"model": model})
+
+    def training_step(
+        self,
+        batch: Any,
+        batch_idx: int,
+        is_accumulating: bool,
+    ) -> None:
+        """
+        Training step, running forward and backward passes.
+
+        Args:
+            batch (Any): The current batch; can be of any form.
+            batch_idx (int): The index of the current batch.
+            is_accumulating (bool): Whether we are accumulating gradients
+                (i.e., not yet calling optimizer.step()). If this is the case,
+                we should skip the synchronization during the backward pass.
+
+        Returns:
+            None; we call `loss.backward()` directly, and store the outputs in
+                `self._current_train_return`.
+        """
+        model = self.state["model"]
+        assert model.training, "Model must be training!"
+
+        network_input = batch
+
+        with self.fabric.no_backward_sync(model, enabled=is_accumulating):
+            # Forward pass
+            network_output = model.forward(network_input)
+            assert_no_nans(
+                network_output["decoder_features"],
+                msg="network_output['decoder_features'] "
+                + f"for batch_idx: {batch_idx}",
+            )
+
+            total_loss, loss_dict = self.loss(
+                network_input=batch,
+                network_output=network_output,
+                loss_input={},
+            )
+
+            # Backward pass
+            self.fabric.backward(total_loss)
+
+            # Optionally compute training metrics
+            train_return = {"total_loss": total_loss, "loss_dict": loss_dict}
+
+            # Store the outputs without gradients for use in logging,
+            # callbacks, learning rate schedulers, etc.
+            self._current_train_return = apply_to_collection(
+                train_return,
+                dtype=torch.Tensor,
+                function=lambda x: x.detach(),
+            )
+
+    def validation_step(
+        self,
+        batch: Any,
+        batch_idx: int,
+        compute_metrics: bool = True,
+    ) -> dict:
+        """
+        Validation step, running forward pass and computing validation
+        metrics.
+
+        Args:
+            batch (Any): The current batch; can be of any form.
+            batch_idx (int): The index of the current batch.
+            compute_metrics (bool): Whether to compute metrics. If False, we
+                will not compute metrics, and the output will be None. Set to
+                False during the inference pipeline, where we need the network
+                output but cannot compute metrics (since we do not have the
+                ground truth).
+
+        Returns:
+            dict: Output dictionary containing the validation metrics and
+                network output.
+        """
+        model = self.state["model"]
+        assert not model.training, "Model must be in evaluation mode during validation!"
+
+        network_input = batch
+
+        # Forward pass
+        network_output = model.forward(network_input)
+
+        assert_no_nans(
+            network_output["decoder_features"],
+            msg="network_output['decoder_features'] " + f"for batch_idx: {batch_idx}",
+        )
+
+        metrics_output = {}
+        if compute_metrics:
+            # Compute all metrics using MetricManager
+            metrics_output = self.metrics(
+                network_input=batch,
+                network_output=network_output,
+                extra_info={},
+            )
+
+            # Avoid gradients in stored values to prevent memory leaks
+            if metrics_output:
+                metrics_output = apply_to_collection(
+                    metrics_output, torch.Tensor, lambda x: x.detach()
+                )
+
+        network_output = apply_to_collection(
+            network_output, torch.Tensor, lambda x: x.detach()
+        )
+
+        validation_return = {
+            "metrics_output": metrics_output,
+            "network_output": network_output,
+        }
+
+        return validation_return

mpnn/transforms/feature_aggregation/mpnn.py

@@ -0,0 +1,184 @@
+from typing import Any
+
+import numpy as np
+from atomworks.common import KeyToIntMapper
+from atomworks.ml.transforms._checks import (
+    check_atom_array_annotation,
+)
+from atomworks.ml.transforms.base import Transform
+from atomworks.ml.transforms.encoding import atom_array_to_encoding
+from atomworks.ml.utils.token import get_token_starts
+from mpnn.transforms.feature_aggregation.token_encodings import MPNN_TOKEN_ENCODING
+
+
+class EncodeMPNNNonAtomizedTokens(Transform):
+    """Encode non-atomized tokens for MPNN with X, X_m, and S features.
+
+    Creates:
+    - X: (L, 37, 3) coordinates for non-atomized tokens
+    - X_m: (L, 37) mask for atom existence and occupancy > occupancy_threshold
+    - S: (L) sequence encoding
+
+    Args:
+        occupancy_threshold (float): Minimum occupancy to consider atom as present. Defaults to 0.5.
+    """
+
+    def __init__(self, occupancy_threshold: float = 0.5):
+        self.occupancy_threshold = occupancy_threshold
+        self.encoding = MPNN_TOKEN_ENCODING
+
+    def check_input(self, data: dict[str, Any]) -> None:
+        check_atom_array_annotation(data, ["atomize", "res_name", "occupancy"])
+
+    def forward(self, data: dict[str, Any]) -> dict[str, Any]:
+        atom_array = data["atom_array"]
+
+        # Check that atom_array is not empty.
+        assert len(atom_array) > 0, "atom_array cannot be empty"
+
+        # Get non-atomized tokens only
+        non_atomized_mask = ~atom_array.atomize
+        non_atomized_array = atom_array[non_atomized_mask]
+
+        assert len(non_atomized_array) > 0, "No non-atomized atoms found"
+
+        if len(non_atomized_array) == 0:
+            # No non-atomized tokens, create empty arrays
+            data["input_features"].update(
+                {
+                    "X": np.zeros((0, 37, 3), dtype=np.float32),
+                    "X_m": np.zeros((0, 37), dtype=np.bool_),
+                    "S": np.zeros((0,), dtype=np.int64),
+                }
+            )
+            return data
+
+        # Encode using the MPNN token encoding
+        encoded = atom_array_to_encoding(
+            non_atomized_array,
+            encoding=self.encoding,
+            default_coord=0.0,  # Use 0.0 instead of NaN for MPNN
+            occupancy_threshold=self.occupancy_threshold,
+        )
+
+        # Create X: coordinates (L, 37, 3)
+        X = encoded["xyz"].astype(np.float32)
+
+        # Create X_m: mask for existence and occupancy > threshold (L, 37)
+        # encoded["mask"] already considers occupancy, we just need to check if atoms exist in encoding
+        X_m = encoded["mask"].astype(np.bool_)
+
+        # Create S: sequence encoding (L,)
+        S = encoded["seq"].astype(np.int64)
+
+        data["input_features"].update(
+            {
+                "X": X,
+                "X_m": X_m,
+                "S": S,
+            }
+        )
+
+        # Check that we have at least one non-atomized token.
+        L = X.shape[0]
+        assert L > 0, "At least one non-atomized token should be present"
+
+        return data
+
+
+class FeaturizeNonAtomizedTokens(Transform):
+    """Add additional features for non-atomized tokens: R_idx, chain_labels, residue_mask."""
+
+    def check_input(self, data: dict[str, Any]) -> None:
+        check_atom_array_annotation(
+            data, ["atomize", "within_chain_res_idx", "chain_iid"]
+        )
+
+    def forward(self, data: dict[str, Any]) -> dict[str, Any]:
+        atom_array = data["atom_array"]
+
+        # Subset to non-atomized.
+        non_atomized_array = atom_array[~atom_array.atomize]
+
+        # Get token starts for non-atomized tokens
+        non_atomized_token_starts = get_token_starts(non_atomized_array)
+        non_atomized_token_level = non_atomized_array[non_atomized_token_starts]
+
+        if len(non_atomized_token_level) == 0:
+            # No non-atomized tokens
+            data["input_features"].update(
+                {
+                    "R_idx": np.zeros((0,), dtype=np.int32),
+                    "chain_labels": np.zeros((0,), dtype=np.int64),
+                    "residue_mask": np.zeros((0,), dtype=np.bool_),
+                }
+            )
+            return data
+
+        # R_idx: residue indices within chains (0-indexed)
+        R_idx = non_atomized_token_level.within_chain_res_idx.astype(np.int32)
+
+        # chain_labels: convert chain_iid to unique integers
+        chain_mapper = KeyToIntMapper()
+        chain_labels = np.array(
+            [
+                chain_mapper(chain_iid)
+                for chain_iid in non_atomized_token_level.chain_iid
+            ],
+            dtype=np.int64,
+        )
+
+        # residue_mask: all 1's for non-atomized tokens
+        residue_mask = np.ones(len(non_atomized_token_level), dtype=np.bool_)
+
+        data["input_features"].update(
+            {
+                "R_idx": R_idx,
+                "chain_labels": chain_labels,
+                "residue_mask": residue_mask,
+            }
+        )
+        return data
+
+
+class FeaturizeAtomizedTokens(Transform):
+    """Add features for atomized tokens: Y, Y_t, Y_m."""
+
+    def check_input(self, data: dict[str, Any]) -> None:
+        check_atom_array_annotation(data, ["atomize", "atomic_number"])
+
+    def forward(self, data: dict[str, Any]) -> dict[str, Any]:
+        atom_array = data["atom_array"]
+
+        # Get atomized tokens only
+        atomized_mask = atom_array.atomize
+        atomized_array = atom_array[atomized_mask]
+
+        if len(atomized_array) == 0:
+            # No atomized tokens
+            data["input_features"].update(
+                {
+                    "Y": np.zeros((0, 3), dtype=np.float32),
+                    "Y_t": np.zeros((0,), dtype=np.int32),
+                    "Y_m": np.zeros((0,), dtype=np.bool_),
+                }
+            )
+            return data
+
+        # Y: coordinates of atomized tokens (n_atomized, 3)
+        Y = atomized_array.coord.astype(np.float32)
+
+        # Y_t: atomic numbers of atomized tokens (n_atomized,)
+        Y_t = atomized_array.atomic_number.astype(np.int32)
+
+        # Y_m: mask for atomized tokens (all 1's since they exist) (n_atomized,)
+        Y_m = np.ones(len(atomized_array), dtype=np.bool_)
+
+        data["input_features"].update(
+            {
+                "Y": Y,
+                "Y_t": Y_t,
+                "Y_m": Y_m,
+            }
+        )
+        return data

mpnn/transforms/feature_aggregation/polymer_ligand_interface.py

@@ -0,0 +1,76 @@
+"""
+Feature aggregation for polymer-ligand interface masks.
+
+This module provides transforms to compute interface masks for polymer residues
+that are at the interface with ligand molecules.
+"""
+
+from typing import Any
+
+import numpy as np
+from atomworks.ml.transforms._checks import check_atom_array_annotation
+from atomworks.ml.transforms.base import Transform
+from atomworks.ml.utils.token import get_token_starts
+
+
+class FeaturizePolymerLigandInterfaceMask(Transform):
+    """
+    Compute a polymer mask indicating which residues are at the polymer-ligand
+    interface.
+
+    This transform processes an atom array to identify polymer residues that
+    have any atoms within the specified distance threshold of ligand atoms.
+    It expects that the atom array already has the
+    'at_polymer_ligand_interface' annotation computed by the
+    ComputePolymerLigandInterface transform.
+    """
+
+    def check_input(self, data: dict[str, Any]) -> None:
+        """Check that required annotations are present."""
+        check_atom_array_annotation(
+            {"atom_array": data["atom_array"]},
+            required=["element", "atomize", "at_polymer_ligand_interface"],
+        )
+
+    def forward(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Compute polymer-ligand interface mask and add to input_features."""
+        atom_array = data["atom_array"]
+
+        # Get interface annotation that should already be computed
+        interface_atoms = atom_array.at_polymer_ligand_interface
+
+        # Get token starts to map atoms to residues
+        token_starts = get_token_starts(atom_array)
+
+        # Create residue-level interface mask for all tokens
+        all_residue_interface_mask = np.zeros(len(token_starts), dtype=bool)
+
+        # For each token (residue), check if any of its atoms are at the
+        # interface.
+        for i, start_idx in enumerate(token_starts):
+            if i < len(token_starts) - 1:
+                end_idx = token_starts[i + 1]
+            else:
+                end_idx = len(atom_array)
+
+            # Check if any atom in this residue is at the interface
+            residue_atoms = interface_atoms[start_idx:end_idx]
+            all_residue_interface_mask[i] = np.any(residue_atoms)
+
+        # Get token-level atomize annotation
+        token_level_array = atom_array[token_starts]
+        non_atomized_mask = ~token_level_array.atomize
+
+        # Get interface mask for non-atomized residues only
+        polymer_interface_mask = all_residue_interface_mask[non_atomized_mask]
+
+        # Initialize input_features if it doesn't exist.
+        if "input_features" not in data:
+            data["input_features"] = {}
+
+        # Add the interface mask to input_features
+        data["input_features"]["polymer_ligand_interface_mask"] = (
+            polymer_interface_mask.astype(np.bool_)
+        )
+
+        return data

mpnn/transforms/feature_aggregation/token_encodings.py

@@ -0,0 +1,132 @@
+from atomworks.constants import AA_LIKE_CHEM_TYPES, STANDARD_AA, UNKNOWN_AA
+from atomworks.ml.encoding_definitions import TokenEncoding
+
+# Token ordering for MPNN.
+token_order = STANDARD_AA + (UNKNOWN_AA,)
+
+# Token ordering for old versions of MPNN.
+legacy_token_order = (
+    "ALA",
+    "CYS",
+    "ASP",
+    "GLU",
+    "PHE",
+    "GLY",
+    "HIS",
+    "ILE",
+    "LYS",
+    "LEU",
+    "MET",
+    "ASN",
+    "PRO",
+    "GLN",
+    "ARG",
+    "SER",
+    "THR",
+    "VAL",
+    "TRP",
+    "TYR",
+    "UNK",
+)
+
+# Atom ordering for new versions of MPNN.
+atom_order = (
+    "N",
+    "CA",
+    "C",
+    "O",
+    "CB",
+    "CG",
+    "CG1",
+    "CG2",
+    "OG",
+    "OG1",
+    "SG",
+    "CD",
+    "CD1",
+    "CD2",
+    "ND1",
+    "ND2",
+    "OD1",
+    "OD2",
+    "SD",
+    "CE",
+    "CE1",
+    "CE2",
+    "CE3",
+    "NE",
+    "NE1",
+    "NE2",
+    "OE1",
+    "OE2",
+    "CH2",
+    "NH1",
+    "NH2",
+    "OH",
+    "CZ",
+    "CZ2",
+    "CZ3",
+    "NZ",
+    "OXT",
+)
+
+# Atom ordering for old versions of MPNN.
+legacy_atom_order = (
+    "N",
+    "CA",
+    "C",
+    "CB",
+    "O",
+    "CG",
+    "CG1",
+    "CG2",
+    "OG",
+    "OG1",
+    "SG",
+    "CD",
+    "CD1",
+    "CD2",
+    "ND1",
+    "ND2",
+    "OD1",
+    "OD2",
+    "SD",
+    "CE",
+    "CE1",
+    "CE2",
+    "CE3",
+    "NE",
+    "NE1",
+    "NE2",
+    "OE1",
+    "OE2",
+    "CH2",
+    "NH1",
+    "NH2",
+    "OH",
+    "CZ",
+    "CZ2",
+    "CZ3",
+    "NZ",
+    "OXT",
+)
+
+# Token encoding for MPNN.
+MPNN_TOKEN_ENCODING = TokenEncoding(
+    token_atoms={token: atom_order for token in token_order},
+    chemcomp_type_to_unknown={chem_type: "UNK" for chem_type in AA_LIKE_CHEM_TYPES},
+)
+
+# Token encoding for versions of MPNN using the legacy token order and
+# new atom order.
+MPNN_LEGACY_TOKEN_ENCODING = TokenEncoding(
+    token_atoms={token: atom_order for token in legacy_token_order},
+    chemcomp_type_to_unknown={chem_type: "UNK" for chem_type in AA_LIKE_CHEM_TYPES},
+)
+
+# Token encoding for versions of MPNN using the legacy token order and
+# legacy atom order.
+MPNN_LEGACY_TOKEN_LEGACY_ATOM_ENCODING = TokenEncoding(
+    token_atoms={token: legacy_atom_order for token in legacy_token_order},
+    chemcomp_type_to_unknown={chem_type: "UNK" for chem_type in AA_LIKE_CHEM_TYPES},
+)