rc-foundry 0.1.1__py3-none-any.whl

rf3/symmetry/resolve.py

@@ -0,0 +1,284 @@
+"""Generalized symmetry resolution implementation, operating on the outputs of AtomWorks.io `parse` function."""
+
+import logging
+from typing import Any, Dict
+
+import numpy as np
+import torch
+from atomworks.io.transforms.atom_array import ensure_atom_array_stack
+from atomworks.ml.transforms.atom_array import AddGlobalTokenIdAnnotation
+from atomworks.ml.transforms.atomize import AtomizeByCCDName
+from atomworks.ml.transforms.base import Compose, convert_to_torch
+from atomworks.ml.transforms.symmetry import FindAutomorphismsWithNetworkX
+from biotite.structure import AtomArray, AtomArrayStack
+from jaxtyping import Bool, Float, Int
+from rf3.loss.af3_losses import (
+    ResidueSymmetryResolution,
+    SubunitSymmetryResolution,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_symmetries(
+    predicted_atom_array: AtomArray | AtomArrayStack,
+    ground_truth_atom_array: AtomArray | AtomArrayStack,
+    resolve_residue_symmetries: bool = True,
+    resolve_subunit_symmetries: bool = True,
+) -> AtomArrayStack:
+    """
+    Generalized symmetry resolution for both residue- and subunit-level symmetries.
+
+    Returns updated ground truth AtomArray with coordinates that minimize RMSD with the predicted structure.
+
+    Args:
+        predicted_atom_array: Predicted structure as AtomArray or AtomArrayStack
+        ground_truth_atom_array: Ground truth structure as AtomArray or AtomArrayStack
+        resolve_residue_symmetries: Whether to resolve residue-level symmetries
+        resolve_subunit_symmetries: Whether to resolve subunit-level symmetries
+
+    Returns:
+        Updated ground truth AtomArray or AtomArrayStack with resolved coordinates
+    """
+    predicted_stack = ensure_atom_array_stack(predicted_atom_array)
+    ground_truth_stack = ensure_atom_array_stack(ground_truth_atom_array)
+
+    # Set ground truth coordinates to nan if they are nan in the predicted coordinates...
+    ground_truth_stack.coord[np.isnan(predicted_stack.coord)] = np.nan
+
+    # ... then nan-to-num the pred_aa coordinates (otherwise, the symmetry resolution may fail)
+    # TODO: Update the symmetry resolution to handle NaNs in the predicted coordinates
+    predicted_stack.coord = np.nan_to_num(predicted_stack.coord)
+
+    # Extract predicted and ground truth coordinates
+    X_pred: Float[torch.Tensor, "D L 3"] = torch.tensor(
+        predicted_stack.coord, dtype=torch.float32
+    )
+    X_gt: Float[torch.Tensor, "D L 3"] = torch.tensor(
+        ground_truth_stack.coord, dtype=torch.float32
+    )
+
+    # (Match dimensions)
+    D_pred, L_pred = X_pred.shape[:2]
+    D_gt, L_gt = X_gt.shape[:2]
+
+    if D_pred != D_gt:
+        if D_gt == 1:
+            X_gt = X_gt.expand(D_pred, -1, -1)
+        else:
+            raise ValueError(
+                f"Cannot broadcast ground truth of shape ({D_gt}) to prediction of shape ({D_pred})"
+            )
+    assert L_pred == L_gt, "Length mismatch: predicted {L_pred}, ground truth {L_gt}"
+
+    # Generate symmetric features (e.g., automorphisms, entity information, etc.) inputs from ground truth
+    symmetry_data = generate_symmetry_resolution_inputs_from_atom_array(
+        ground_truth_stack
+    )
+
+    # Extract coordinate mask from ground truth stack
+    crd_mask: Bool[torch.Tensor, "D L"]
+    if "occupancy" in ground_truth_stack.get_annotation_categories():
+        crd_mask = torch.tensor(ground_truth_stack.occupancy > 0.0, dtype=torch.bool)
+    else:
+        logger.warning(
+            "No occupancy annotation found in ground truth, using coordinate validity mask (not NaN)"
+        )
+        crd_mask = ~torch.isnan(torch.tensor(ground_truth_stack.coord)).any(dim=-1)
+
+    assert not torch.isnan(
+        X_pred
+    ).any(), "NaN coordinates found in predicted structure!"
+
+    # Apply symmetry resolution (returns updated ground truth coordinates)
+    X_gt_resolved: Float[torch.Tensor, "D L 3"] = apply_symmetry_resolution(
+        X_pred=X_pred,
+        X_gt=X_gt,
+        crd_mask=crd_mask,
+        automorphisms=symmetry_data["automorphisms"],
+        molecule_entity=symmetry_data["molecule_entity"],
+        molecule_iid=symmetry_data["molecule_iid"],
+        crop_mask=symmetry_data["crop_mask"],
+        coord_atom_lvl=symmetry_data["coord_atom_lvl"],
+        mask_atom_lvl=symmetry_data["mask_atom_lvl"],
+        resolve_residue=resolve_residue_symmetries,
+        resolve_subunit=resolve_subunit_symmetries,
+    )
+
+    # Update the ground truth AtomArray with resolved coordinates
+    result_stack = ground_truth_stack.copy()
+    result_stack.coord = X_gt_resolved.cpu().numpy()
+
+    return result_stack
+
+
+def generate_symmetry_resolution_inputs_from_atom_array(
+    atom_array: AtomArray | AtomArrayStack,
+) -> Dict[str, Any]:
+    """
+    Generate all inputs needed for symmetry resolution from an AtomArray.
+
+    Args:
+        atom_array: Input AtomArray or AtomArrayStack
+
+    Returns:
+        Dictionary containing:
+            - automorphisms: List[np.ndarray]
+            - molecule_entity: torch.Tensor [N_atoms]
+            - molecule_iid: torch.Tensor [N_atoms]
+            - coord_atom_lvl: torch.Tensor [N_atoms, 3]
+            - mask_atom_lvl: torch.Tensor [N_atoms]
+            - atom_to_token_map: torch.Tensor [N_atoms]
+            - crop_mask: torch.Tensor [N_atoms]
+    """
+    # (Take first model)
+    atom_array_stack = ensure_atom_array_stack(atom_array)
+    atom_array = atom_array_stack[0]
+
+    # (Avoid modifying the original)
+    atom_array = atom_array.copy()
+
+    # Prepare transform pipeline to generate features
+    transforms = [AtomizeByCCDName(atomize_by_default=True)]
+
+    if "token_id" not in atom_array.get_annotation_categories():
+        transforms.append(AddGlobalTokenIdAnnotation())
+
+    transforms.append(FindAutomorphismsWithNetworkX())
+
+    pipeline = Compose(transforms)
+    data = pipeline({"atom_array": atom_array})
+    atom_array = data["atom_array"]
+
+    result: Dict[str, Any] = {}
+    # Extract automorphisms
+    result["automorphisms"] = data.get("automorphisms", [])
+    # Extract molecule annotations (assert they exist)
+    assert (
+        "molecule_entity" in atom_array.get_annotation_categories()
+    ), "molecule_entity annotation required"
+    assert (
+        "molecule_iid" in atom_array.get_annotation_categories()
+    ), "molecule_iid annotation required"
+
+    result["molecule_entity"] = atom_array.molecule_entity
+    result["molecule_iid"] = atom_array.molecule_iid
+
+    # Extract coordinates
+    coords: np.ndarray = atom_array.coord
+    result["coord_atom_lvl"] = coords
+
+    # Extract mask from occupancy (like in lddt.py) - no batch dimension for SubunitSymmetryResolution
+    mask: np.ndarray
+    if "occupancy" in atom_array.get_annotation_categories():
+        mask = atom_array.occupancy > 0.0
+    else:
+        # Fallback to coordinate validity
+        mask = ~np.isnan(atom_array.coord).any(axis=-1)
+
+    # Keep mask as [N_atoms] for SubunitSymmetryResolution compatibility
+    result["mask_atom_lvl"] = mask
+
+    # Extract atom to token map (like in lddt.py)
+    if "token_id" in atom_array.get_annotation_categories():
+        result["atom_to_token_map"] = atom_array.token_id.astype(np.int32)
+    else:
+        # This should not happen since AddGlobalTokenIdAnnotation was applied
+        raise ValueError(
+            "token_id annotation not found after AddGlobalTokenIdAnnotation"
+        )
+
+    # Create crop_mask (full range)
+    result["crop_mask"] = np.arange(len(atom_array), dtype=np.int32)
+
+    # Step 3: Convert all numpy arrays to torch tensors using convert_to_torch
+    # First, create a temporary dict with the keys we want to convert
+    torch_data = {
+        "molecule_entity": result["molecule_entity"],
+        "molecule_iid": result["molecule_iid"],
+        "coord_atom_lvl": result["coord_atom_lvl"],
+        "mask_atom_lvl": result["mask_atom_lvl"],
+        "atom_to_token_map": result["atom_to_token_map"],
+        "crop_mask": result["crop_mask"],
+    }
+
+    # Convert to torch tensors
+    torch_data = convert_to_torch(torch_data, list(torch_data.keys()))
+
+    # Update result with torch tensors
+    result.update(torch_data)
+
+    return result
+
+
+def apply_symmetry_resolution(
+    X_pred: Float[torch.Tensor, "D L 3"],
+    X_gt: Float[torch.Tensor, "D L 3"],
+    crd_mask: Bool[torch.Tensor, "D L"],
+    automorphisms: list,
+    molecule_entity: Int[torch.Tensor, "N_atoms"],
+    molecule_iid: Int[torch.Tensor, "N_atoms"],
+    crop_mask: Int[torch.Tensor, "N_atoms"],
+    coord_atom_lvl: Float[torch.Tensor, "N_atoms 3"],
+    mask_atom_lvl: Bool[torch.Tensor, "N_atoms"],
+    resolve_residue: bool = True,
+    resolve_subunit: bool = True,
+) -> Float[torch.Tensor, "D L 3"]:
+    """
+    Apply the actual symmetry resolution using the existing classes and return updated coordinates.
+
+    Args:
+        X_pred: Predicted coordinates [D, L, 3]
+        X_gt: Ground truth coordinates [D, L, 3]
+        crd_mask: Coordinate mask [D, L]
+        automorphisms: List of automorphism groups
+        molecule_entity: Molecule entity IDs [N_atoms]
+        molecule_iid: Molecule instance IDs [N_atoms]
+        crop_mask: Crop mask indices [N_atoms]
+        coord_atom_lvl: Atom-level coordinates [N_atoms, 3]
+        mask_atom_lvl: Atom-level mask [N_atoms]
+        resolve_residue: Whether to resolve residue symmetries
+        resolve_subunit: Whether to resolve subunit symmetries
+
+    Returns:
+        Updated ground truth coordinates [D, L, 3]
+    """
+    # Prepare loss_input dictionary for existing classes
+    loss_input: Dict[str, torch.Tensor] = {
+        "X_gt_L": X_gt.clone(),
+        "crd_mask_L": crd_mask.clone(),
+    }
+
+    # Apply subunit symmetry resolution
+    if resolve_subunit:
+        subunit_resolver = SubunitSymmetryResolution()
+
+        # Create symmetry resolution input
+        symmetry_resolution: Dict[str, torch.Tensor] = {
+            "molecule_entity": molecule_entity,
+            "molecule_iid": molecule_iid,
+            "crop_mask": crop_mask,
+            "coord_atom_lvl": coord_atom_lvl,
+            "mask_atom_lvl": mask_atom_lvl,
+        }
+
+        # Create network output dict
+        network_output: Dict[str, torch.Tensor] = {"X_L": X_pred}
+
+        # Apply subunit resolution
+        logger.info("Applying subunit symmetry resolution")
+        loss_input = subunit_resolver(network_output, loss_input, symmetry_resolution)
+
+    # Apply residue symmetry resolution
+    if resolve_residue and automorphisms:
+        logger.info("Applying residue symmetry resolution")
+        residue_resolver = ResidueSymmetryResolution()
+
+        # Create network output dict
+        network_output: Dict[str, torch.Tensor] = {"X_L": X_pred}
+
+        # Apply residue resolution
+        loss_input = residue_resolver(network_output, loss_input, automorphisms)
+
+    # Return the updated ground truth coordinates
+    return loss_input["X_gt_L"]

rf3/train.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env -S /bin/sh -c '"$(dirname "$0")/../../../../.ipd/shebang/rf3_exec.sh" "$0" "$@"'
+
+import logging
+import os
+
+import hydra
+import rootutils
+from dotenv import load_dotenv
+from omegaconf import DictConfig
+
+from foundry.utils.logging import suppress_warnings
+from foundry.utils.weights import CheckpointConfig
+
+# Setup root dir and environment variables (more info: https://github.com/ashleve/rootutils)
+# NOTE: Sets the `PROJECT_ROOT` environment variable to the root directory of the project (where `.project-root` is located)
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+load_dotenv(override=True)
+
+_config_path = os.path.join(os.environ["PROJECT_ROOT"], "models/rf3/configs")
+
+_spawning_process_logger = logging.getLogger(__name__)
+
+
+@hydra.main(config_path=_config_path, config_name="train", version_base="1.3")
+def train(cfg: DictConfig) -> None:
+    # ==============================================================================
+    # Import dependencies and resolve Hydra configuration
+    # ==============================================================================
+
+    _spawning_process_logger.info("Importing dependencies...")
+
+    # Lazy imports to make config generation fast
+    import torch
+    from lightning.fabric import seed_everything
+    from lightning.fabric.loggers import Logger
+
+    # If training on DIGS L40, set precision of matrix multiplication to balance speed and accuracy
+    # Reference: https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
+    torch.set_float32_matmul_precision("medium")
+
+    from foundry.callbacks.callback import BaseCallback  # noqa
+    from foundry.utils.instantiators import instantiate_loggers, instantiate_callbacks  # noqa
+    from foundry.utils.logging import (
+        print_config_tree,
+        log_hyperparameters_with_all_loggers,
+    )  # noqa
+    from foundry.utils.ddp import RankedLogger  # noqa
+    from foundry.utils.ddp import is_rank_zero, set_accelerator_based_on_availability  # noqa
+    from foundry.utils.datasets import (
+        recursively_instantiate_datasets_and_samplers,
+        assemble_distributed_loader,
+        subset_dataset_to_example_ids,
+        assemble_val_loader_dict,
+    )  # noqa
+
+    set_accelerator_based_on_availability(cfg)
+
+    ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+    _spawning_process_logger.info("Completed dependency imports ...")
+
+    # ... print the configuration tree (NOTE: Only prints for rank 0)
+    print_config_tree(cfg, resolve=True)
+
+    # ==============================================================================
+    # Logging and Callback instantiation
+    # ==============================================================================
+
+    # Reduce the logging level for all dataset and sampler loggers (unless rank 0)
+    # We will still see messages from Rank 0; they are identical, since all ranks load and sample from the same datasets
+    if not is_rank_zero():
+        dataset_logger = logging.getLogger("datasets")
+        sampler_logger = logging.getLogger("atomworks.ml.samplers")
+        dataset_logger.setLevel(logging.WARNING)
+        sampler_logger.setLevel(logging.ERROR)
+
+    # ... seed everything (NOTE: By setting `workers=True`, we ensure that the dataloaders are seeded as well)
+    # (`PL_GLOBAL_SEED` environment varaible will be passed to the spawned subprocessed; e.g., through `ddp_spawn` backend)
+    if cfg.get("seed"):
+        ranked_logger.info(f"Seeding everything with seed={cfg.seed}...")
+        seed_everything(cfg.seed, workers=True, verbose=True)
+    else:
+        ranked_logger.warning("No seed provided - Not seeding anything!")
+
+    ranked_logger.info("Instantiating loggers...")
+    loggers: list[Logger] = instantiate_loggers(cfg.get("logger"))
+
+    ranked_logger.info("Instantiating callbacks...")
+    callbacks: list[BaseCallback] = instantiate_callbacks(cfg.get("callbacks"))
+
+    # ==============================================================================
+    # Trainer and model instantiation
+    # ==============================================================================
+
+    # ... instantiate the trainer
+    ranked_logger.info("Instantiating trainer...")
+    trainer = hydra.utils.instantiate(
+        cfg.trainer,
+        loggers=loggers or None,
+        callbacks=callbacks or None,
+        _convert_="partial",
+        _recursive_=False,
+    )
+    # (Store the Hydra configuration in the trainer state)
+    trainer.initialize_or_update_trainer_state({"train_cfg": cfg})
+
+    # ... spawn processes for distributed training
+    # (We spawn here, rather than within `fit`, so we can use Fabric's `init_module` to efficiently initialize the model on the appropriate device)
+    ranked_logger.info(
+        f"Spawning {trainer.fabric.world_size} processes from {trainer.fabric.global_rank}..."
+    )
+    trainer.fabric.launch()
+
+    # ... construct the model
+    trainer.construct_model()
+
+    # ... construct the optimizer and schedule (which requires the model to be constructed)
+    trainer.construct_optimizer()
+    trainer.construct_scheduler()
+
+    # ==============================================================================
+    # Dataset instantiation
+    # ==============================================================================
+
+    # Number of examples per epoch (accross all GPUs)
+    # (We must sample this many indices from our sampler)
+    n_examples_per_epoch = cfg.trainer.n_examples_per_epoch
+
+    # ... build the train dataset
+    assert (
+        "train" in cfg.datasets and cfg.datasets.train
+    ), "No 'train' dataloader configuration provided! If only performing validation, use `validate.py` instead."
+    dataset_and_sampler = recursively_instantiate_datasets_and_samplers(
+        cfg.datasets.train
+    )
+    train_dataset, train_sampler = (
+        dataset_and_sampler["dataset"],
+        dataset_and_sampler["sampler"],
+    )
+
+    # ... compose the train loader
+    if "subset_to_example_ids" in cfg.datasets:
+        # Backdoor for debugging and overfitting: subset the dataset to a specific set of example IDs
+        train_dataset = subset_dataset_to_example_ids(
+            train_dataset, cfg.datasets.subset_to_example_ids
+        )
+        train_sampler = None  # Sampler is no longer valid, since we are using a subset of the dataset
+
+    train_loader = assemble_distributed_loader(
+        dataset=train_dataset,
+        sampler=train_sampler,
+        rank=trainer.fabric.global_rank,
+        world_size=trainer.fabric.world_size,
+        n_examples_per_epoch=n_examples_per_epoch,
+        loader_cfg=cfg.dataloader["train"],
+    )
+
+    # ... compose the validation loader(s)
+    if "val" in cfg.datasets and cfg.datasets.val:
+        val_loaders = assemble_val_loader_dict(
+            cfg=cfg.datasets.val,
+            rank=trainer.fabric.global_rank,
+            world_size=trainer.fabric.world_size,
+            loader_cfg=cfg.dataloader["val"],
+        )
+    else:
+        ranked_logger.warning("No validation datasets provided! Skipping validation...")
+        val_loaders = None
+
+    ranked_logger.info("Logging hyperparameters...")
+    log_hyperparameters_with_all_loggers(
+        trainer=trainer, cfg=cfg, model=trainer.state["model"]
+    )
+
+    # ... load the checkpoint configuration
+    ckpt_config = None
+    if "ckpt_config" in cfg and cfg.ckpt_config:
+        ckpt_config = hydra.utils.instantiate(cfg.ckpt_config)
+    elif "ckpt_path" in cfg and cfg.ckpt_path:
+        # Just a checkpoint path
+        if cfg.ckpt_path is not None:
+            ckpt_config = CheckpointConfig(path=cfg.ckpt_path)
+
+    # ... train the model
+    ranked_logger.info("Training model...")
+
+    with suppress_warnings():
+        trainer.fit(
+            train_loader=train_loader, val_loaders=val_loaders, ckpt_config=ckpt_config
+        )
+
+
+if __name__ == "__main__":
+    train()