rc-foundry 0.1.1__py3-none-any.whl

rf3/utils/io.py

@@ -0,0 +1,198 @@
+from os import PathLike
+from pathlib import Path
+
+import numpy as np
+import torch
+from atomworks.io.utils.io_utils import to_cif_file
+from atomworks.ml.utils.io import apply_sharding_pattern
+from atomworks.ml.utils.misc import hash_sequence
+from beartype.typing import Literal
+from biotite.structure import AtomArray, AtomArrayStack, stack
+
+from foundry.utils.alignment import weighted_rigid_align
+from foundry.utils.ddp import RankedLogger
+
+ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+
+DICTIONARY_LIKE_EXTENSIONS = {".json", ".yaml", ".yml", ".pkl"}
+CIF_LIKE_EXTENSIONS = {".cif", ".pdb", ".bcif", ".cif.gz", ".pdb.gz", ".bcif.gz"}
+
+
+def get_sharded_output_path(
+    example_id: str,
+    base_dir: Path,
+    sharding_pattern: str | None = None,
+) -> Path:
+    """Get output directory path for an example with optional sharding.
+
+    Args:
+        example_id: Example identifier (used as final directory name).
+        base_dir: Base output directory.
+        sharding_pattern: Sharding pattern like ``/0:2/2:4/`` or ``None`` for no sharding.
+            Pattern defines how to split the hash of ``example_id`` into nested directories.
+
+    Returns:
+        Output directory path. If sharding is enabled, returns ``base_dir/shard1/shard2/.../example_id``.
+        Otherwise returns ``base_dir/example_id``.
+
+    Examples:
+        Without sharding::
+
+            get_sharded_output_path("entry_1", Path("/out"))
+            # Returns: /out/entry_1
+
+        With sharding pattern ``/0:2/2:4/``::
+
+            get_sharded_output_path("entry_1", Path("/out"), "/0:2/2:4/")
+            # Computes hash of "entry_1" (e.g., "a1b2c3d4e5f")
+            # Returns: /out/a1/b2/entry_1
+    """
+    if not sharding_pattern:
+        return base_dir / example_id
+
+    # Hash the example ID and apply sharding pattern
+    example_hash = hash_sequence(example_id)
+    sharded_path = apply_sharding_pattern(example_hash, sharding_pattern)
+
+    # Return base_dir / sharded_directories / example_id
+    return base_dir / sharded_path.parent / example_id
+
+
+def build_stack_from_atom_array_and_batched_coords(
+    coords: np.ndarray | torch.Tensor,
+    atom_array: AtomArray,
+) -> AtomArrayStack:
+    """Builds an AtomArrayStack from an AtomArray and a set of coordinates with a batch dimension.
+
+    Additionally, handles the case where the AtomArray contains multiple transformations and we must adjust the chain_id.
+
+    Args:
+        coords (np.array): The coordinates to be assigned to the AtomArrayStack. Must have shape (nbatch, n_atoms, 3).
+        atom_array (AtomArray): The AtomArray to be stacked. Must have shape (n_atoms,)
+    """
+    if isinstance(coords, torch.Tensor):
+        coords = coords.cpu().numpy()
+
+    # (Diffusion batch size will become the number of models)
+    n_batch = coords.shape[0]
+
+    # Build the stack and assign the coordinates
+    atom_array_stack = stack([atom_array for _ in range(n_batch)])
+    atom_array_stack.coord = coords
+
+    # Adjust chain_id if there are multiple transformations
+    # (Otherwise, we will have ambiguous bond annotations, since only `chain_id` is used for the bond annotations)
+    if (
+        "transformation_id" in atom_array.get_annotation_categories()
+        and len(np.unique(atom_array_stack.transformation_id)) > 1
+    ):
+        new_chain_ids = np.char.add(
+            atom_array_stack.chain_id, atom_array_stack.transformation_id
+        )
+        atom_array_stack.set_annotation("chain_id", new_chain_ids)
+
+    return atom_array_stack
+
+
+def dump_structures(
+    atom_arrays: AtomArrayStack | list[AtomArray] | AtomArray,
+    base_path: PathLike,
+    one_model_per_file: bool,
+    extra_fields: list[str] | Literal["all"] = [],
+    file_type: str = "cif.gz",
+) -> None:
+    """Dump structures to CIF files, given the coordinates and input AtomArray.
+
+    Args:
+        atom_arrays (AtomArrayStack | list[AtomArray] | AtomArray): Either an AtomArrayStack, a list of AtomArray objects,
+            or a single AtomArray object to be dumped to CIF file(s)
+        base_path (PathLike): Base path where the output files will be saved.
+        one_model_per_file (bool): Flag to determine if each model should be dumped into a separate file. Has no effect if
+            `atom_arrays` is a list of AtomArrays.
+        extra_fields (list[str] | Literal["all"]): List of extra fields to include in the CIF file.
+    """
+    base_path = Path(base_path)
+
+    if one_model_per_file:
+        assert (
+            isinstance(atom_arrays, AtomArrayStack) or isinstance(atom_arrays, list)
+        ), "AtomArrayStack or list of AtomArray required when one_model_per_file is True"
+        # One model per file —> loop over the diffusion batch
+        for i in range(len(atom_arrays)):
+            path = f"{base_path}_model_{i}"
+            to_cif_file(
+                atom_arrays[i],
+                path,
+                file_type=file_type,
+                include_entity_poly=False,
+                extra_fields=extra_fields,
+            )
+    else:
+        # Include all models in a single CIF file
+        to_cif_file(
+            atom_arrays,
+            base_path,
+            file_type=file_type,
+            include_entity_poly=False,
+            extra_fields=extra_fields,
+        )
+
+
+def dump_trajectories(
+    trajectory_list: list[torch.Tensor | np.ndarray],
+    atom_array: AtomArray,
+    base_path: Path,
+    align_structures: bool = True,
+    file_type: str = "cif.gz",
+) -> None:
+    """Write denoising trajectories to CIF files.
+
+    Args:
+        trajectory_list (List[torch.Tensor]): List of tensors of length n_steps representing the diffusion trajectory at each step.
+            Each tensor has shape [D, L, 3], where D is the diffusion batch size and L is the number of atoms.
+        atom_array (np.ndarray): Atom array corresponding to the coordinates.
+        base_path (Path): Base path where the output files will be saved.
+        align_structures (bool): Flag to determine if the structures should be aligned on the final prediction.
+            If False, each step may have a different alignment.
+        file_type (str): File type for output (e.g., "cif", "cif.gz", "pdb"). Defaults to ``"cif.gz"``.
+    """
+    n_steps = len(trajectory_list)
+
+    if align_structures:
+        # ... align the trajectories on the last prediction
+        w_L = torch.ones(*trajectory_list[0].shape[:2]).to(trajectory_list[0].device)
+        X_exists_L = torch.ones(trajectory_list[0].shape[1], dtype=torch.bool).to(
+            trajectory_list[0].device
+        )
+        for step in range(n_steps - 1):
+            trajectory_list[step] = weighted_rigid_align(
+                X_L=trajectory_list[-1],
+                X_gt_L=trajectory_list[step],
+                X_exists_L=X_exists_L,
+                w_L=w_L,
+            )
+
+    # ... invert the list, to make the trajectory compatible with PyMol (which builds the bond graph from the first frame)
+    trajectory_list = trajectory_list[::-1]
+
+    # ... iterate over the range of D (diffusion batch size; e.g., 5 during validation)
+    # (We want to convert `aligned_trajectory_list` to a list of length D where each item is a tensor of shape [n_steps, L, 3])
+    trajectories_split_by_model = []
+    for d in range(trajectory_list[0].shape[0]):
+        trajectory_for_single_model = torch.stack(
+            [trajectory_list[step][d] for step in range(n_steps)], dim=0
+        )
+        trajectories_split_by_model.append(trajectory_for_single_model)
+
+    #  ... write the trajectories to CIF files, named by epoch, dataset, example_id, and model index (within the diffusion batch)
+    for i, trajectory in enumerate(trajectories_split_by_model):
+        if isinstance(trajectory, torch.Tensor):
+            trajectory = trajectory.cpu().numpy()
+        atom_array_stack = build_stack_from_atom_array_and_batched_coords(
+            trajectory, atom_array
+        )
+
+        path = f"{base_path}_model_{i}"
+        to_cif_file(
+            atom_array_stack, path, file_type=file_type, include_entity_poly=False
+        )

rf3/utils/loss.py

@@ -0,0 +1,72 @@
+import torch
+
+
+def convert_batched_losses_to_list_of_dicts(loss_dict: dict[str, torch.Tensor]):
+    """Converts a dictionary of batched and non-batched loss tensors into a list of dictionaries.
+
+    Args:
+        loss_dict (dict): A dictionary where keys are loss names and values are PyTorch tensors.
+                          Some values may be batched (1D tensors), while others are not (0D tensors).
+
+    Returns:
+        list: A list of dictionaries, each representing a batch or non-batched losses.
+
+    Example:
+        >>> outputs = {
+        ...     "loss_dict": {
+        ...         "diffusion_loss": torch.tensor([0.0509, 0.0062]),
+        ...         "smoothed_lddt_loss": torch.tensor([0.2507, 0.2797]),
+        ...         "t": torch.tensor([1.7329, 9.3498]),
+        ...         "distogram_loss": torch.tensor(1.7663),
+        ...         "total_loss": torch.tensor(1.2281),
+        ...     }
+        ... }
+        >>> convert_batched_losses_to_list_of_dicts(outputs["loss_dict"])
+        [{'batch_idx': 0, 'diffusion_loss': 0.0509, 'smoothed_lddt_loss': 0.2507, 't': 1.7329},
+         {'batch_idx': 1, 'diffusion_loss': 0.0062, 'smoothed_lddt_loss': 0.2797, 't': 9.3498},
+         {'distogram_loss': 1.7663, 'total_loss': 1.2281}]
+    """
+    result = []
+    batch_size = next((v.size(0) for v in loss_dict.values() if v.dim() == 1), 1)
+
+    # Create a dictionary for each batch index
+    for batch_idx in range(batch_size):
+        batch_dict = {"batch_idx": batch_idx}
+
+        for key, value in loss_dict.items():
+            if value.dim() == 1:  # Check if the tensor is batched
+                batch_dict[key] = value[batch_idx].item()
+
+        result.append(batch_dict)
+
+    # Create a dictionary for non-batched losses
+    non_batched_dict = {}
+    for key, value in loss_dict.items():
+        if value.dim() == 0:  # Check if the tensor is not batched
+            non_batched_dict[key] = value.item()
+
+    result.append(non_batched_dict)
+
+    return result
+
+
+def mean_losses(loss_dict_batched: dict[str, torch.Tensor]) -> dict:
+    """Compute the mean of each tensor in a dictionary of batched losses.
+
+    Args:
+        loss_dict_batched (Dict[str, torch.Tensor]): A dictionary where each key maps to a tensor of losses.
+
+    Returns:
+        dict: A dictionary with the mean loss for each key (as a tensor).
+
+    Example:
+        >>> loss_dict_batched = {"loss1": torch.tensor([0.5, 0.7]), "loss2": torch.tensor([1.0])}
+        >>> mean_losses(loss_dict_batched)
+        {'loss1': 0.6, 'loss2': 1.0}
+    """
+    loss_dict = {}
+    for key, batched_loss in loss_dict_batched.items():
+        # Compute the mean of the tensor and store it in the dictionary
+        loss_dict[key] = batched_loss.mean().item()
+
+    return loss_dict

rf3/utils/predict_and_score.py

@@ -0,0 +1,165 @@
+"""Utility to run RF3 predictions and then a set of metrics on those predictions."""
+
+from pathlib import Path
+
+import biotite.structure as struc
+from atomworks.ml.transforms.filters import remove_protein_terminal_oxygen
+from atomworks.ml.utils.rng import create_rng_state_from_seeds, rng_state
+from beartype.typing import Any
+from rf3.inference_engines.rf3 import RF3InferenceEngine
+from rf3.utils.inference import InferenceInput
+
+from foundry.metrics.metric import MetricManager
+from foundry.utils.ddp import RankedLogger
+
+ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+
+
+def _clean_atom_array_for_rf3(atom_array: struc.AtomArray) -> struc.AtomArray:
+    """Preprocess atom array by removing terminal oxygen and hydrogens."""
+    original_count = len(atom_array)
+
+    # Remove terminal oxygen atoms
+    atom_array = remove_protein_terminal_oxygen(atom_array)
+    if len(atom_array) < original_count:
+        ranked_logger.warning(
+            f"Removed {original_count - len(atom_array)} terminal oxygen atoms. "
+            f"Atom count changed from {original_count} to {len(atom_array)}."
+        )
+        original_count = len(atom_array)
+
+    # Filter to heavy atoms only (no hydrogen)
+    atom_array = atom_array[atom_array.element != "H"]
+    if len(atom_array) < original_count:
+        ranked_logger.warning(
+            f"Removed {original_count - len(atom_array)} hydrogen atoms. "
+            f"Atom count changed from {original_count} to {len(atom_array)}."
+        )
+
+    return atom_array
+
+
+def predict_and_score_with_rf3(
+    atom_arrays: list[struc.AtomArray],
+    ckpt_path: str | Path,
+    metrics=None,
+    n_recycles: int = 10,
+    diffusion_batch_size: int = 5,
+    num_steps: int = 50,
+    example_ids: list[str] | None = None,
+    annotate_b_factor_with_plddt: bool = True,
+    rng_seed: int = 1,
+) -> dict[str, dict[str, Any]]:
+    """Predict structures with RF3 and evaluate against inputs.
+
+    Metrics are computed using the RF3 inference engine's internal trainer,
+    which automatically handles symmetry resolution during validation.
+
+    Args:
+        atom_arrays: List of input structures (ground truth).
+        ckpt_path: Path to RF3 checkpoint file.
+        metrics: Metrics to compute. Can be:
+            - Dict mapping names to Metric objects
+            - List of (name, Metric) tuples
+            - None (no metrics)
+        n_recycles: Number of recycles. Defaults to ``10``.
+        diffusion_batch_size: Number of structures per input. Defaults to ``5``.
+        num_steps: Number of diffusion steps. Defaults to ``50``.
+        example_ids: Optional IDs for each structure. Defaults to "example_0", "example_1", etc.
+        annotate_b_factor_with_plddt: Whether to write pLDDT to B-factor. Defaults to ``True``.
+        rng_seed: RNG seed for reproducibility. Defaults to ``1``.
+
+    Returns:
+        Dict mapping example_id to::
+
+            {
+                "predicted_structures": list[AtomArray] | AtomArrayStack,
+                "metrics": dict[str, float],
+            }
+
+    Example:
+        ```python
+        metrics = [
+            ("all_atom_lddt", AllAtomLDDT()),
+            ("by_type_lddt", ByTypeLDDT()),
+        ]
+
+        results = predict_and_score_with_rf3(
+            atom_arrays=structures,
+            ckpt_path="rf3_latest.pt",
+            metrics=metrics,
+        )
+        ```
+    """
+    # Generate example IDs if not provided
+    if example_ids is None:
+        example_ids = [f"example_{i}" for i in range(len(atom_arrays))]
+
+    # Preprocess atom arrays (remove terminal oxygen and hydrogens so that atom counts match)
+    ranked_logger.info("Preprocessing atom arrays...")
+    preprocessed_arrays = [_clean_atom_array_for_rf3(arr.copy()) for arr in atom_arrays]
+
+    # Convert metrics to MetricManager if provided
+    if metrics is not None:
+        metric_manager = MetricManager.from_metrics(metrics)
+    else:
+        # (Prediction only, no metrics)
+        metric_manager = None
+
+    # Initialize RF3 engine (one-time) with custom metrics
+    ranked_logger.info("Initializing RF3 inference engine...")
+    inference_engine = RF3InferenceEngine(
+        ckpt_path=ckpt_path,
+        n_recycles=n_recycles,
+        diffusion_batch_size=diffusion_batch_size,
+        num_steps=num_steps,
+        seed=None,  # We'll use external RNG state (set below)
+        metrics_cfg=metric_manager,  # Pass MetricManager to engine
+    )
+
+    with rng_state(create_rng_state_from_seeds(rng_seed, rng_seed, rng_seed)):
+        results = {}
+
+        # Loop over each example
+        for example_id, ground_truth_array in zip(example_ids, preprocessed_arrays):
+            ranked_logger.info(f"Predicting structure for {example_id}...")
+
+            # Create InferenceInput from AtomArray
+            inference_input = InferenceInput.from_atom_array(
+                ground_truth_array, example_id=example_id
+            )
+
+            # Run inference in-memory
+            # The engine's trainer.validation_step() automatically:
+            # 1. Runs inference
+            # 2. Applies symmetry resolution
+            # 3. Computes configured metrics
+            inference_results = inference_engine.run(
+                inputs=inference_input,
+                out_dir=None,  # Return in-memory
+                annotate_b_factor_with_plddt=annotate_b_factor_with_plddt,
+            )
+
+            # Extract results for this example
+            result = inference_results[example_id]
+
+            # Check for early stopping
+            if result.get("early_stopped", False):
+                ranked_logger.warning(
+                    f"Early stopping triggered for {example_id} "
+                    f"(mean pLDDT = {result.get('mean_plddt', 'N/A'):.2f})"
+                )
+                results[example_id] = {
+                    "predicted_structures": None,
+                    "metrics": result.get("metrics", {}),
+                    "early_stopped": True,
+                }
+                continue
+
+            # Store results
+            results[example_id] = {
+                "predicted_structures": result["predicted_structures"],
+                "metrics": result.get("metrics", {}),
+            }
+
+    return results