rc-foundry 0.1.1__py3-none-any.whl

rfd3/utils/io.py

@@ -0,0 +1,245 @@
+import json
+import re
+from os import PathLike
+from pathlib import Path
+from typing import Literal
+
+import numpy as np
+import torch
+from atomworks.io.utils.io_utils import to_cif_file
+from biotite.structure import AtomArray, AtomArrayStack, stack
+
+from foundry.utils.alignment import weighted_rigid_align
+
+DICTIONARY_LIKE_EXTENSIONS = {".json", ".yaml", ".yml", ".pkl"}
+CIF_LIKE_EXTENSIONS = {".cif", ".pdb", ".bcif", ".cif.gz", ".pdb.gz", ".bcif.gz"}
+
+
+def dump_structures(
+    atom_arrays: AtomArrayStack | list[AtomArray] | AtomArray,
+    base_path: PathLike,
+    one_model_per_file: bool,
+    extra_fields: list[str] | Literal["all"] = [],
+    **kwargs,
+) -> 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="cif.gz",
+                include_entity_poly=False,
+                extra_fields=extra_fields,
+                **kwargs,
+            )
+    else:
+        # Include all models in a single CIF file
+        to_cif_file(
+            atom_arrays,
+            base_path,
+            file_type="cif.gz",
+            include_entity_poly=False,
+            extra_fields=extra_fields,
+            **kwargs,
+        )
+
+
+def dump_metadata(
+    prediction_metadata: dict,
+    base_path: PathLike,
+    one_model_per_file: bool,
+):
+    """
+    Dump JSONs of prediction metadata to disk.
+
+    Args:
+        prediction_metadata (dict): Dictionary containing metadata for the predictions.
+            Instantiated after the models' predictions in trainer. Keys are model indices (from 0)
+        base_path (PathLike): Base path where the output files will be saved.
+        one_model_per_file (bool): If True, save each model's metadata in a separate file.
+    """
+    if one_model_per_file:
+        # One model per file —> loop over the diffusion batch
+        for i in prediction_metadata:
+            path = f"{base_path}_model_{i}"
+            with open(f"{path}.json", "w") as f:
+                json.dump(prediction_metadata[i], f, indent=4)
+    else:
+        # Include all models in a single JSON file
+        with open(f"{base_path}.json", "w") as f:
+            json.dump(prediction_metadata, f, indent=4)
+
+
+def dump_trajectories(
+    trajectory_list: list[torch.Tensor | np.ndarray],
+    atom_array: AtomArray,
+    base_path: Path,
+    align_structures: bool = True,
+    coord_atom_lvl_to_be_noised: torch.Tensor | None = None,
+    is_motif_atom_with_fixed_pos: torch.Tensor | None = None,
+    max_frames: int = 100,
+) -> 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.
+    """
+    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]
+
+    # ... Select subset of frames if necessary
+    if n_steps > max_frames:
+        selected_indices = torch.linspace(0, n_steps - 1, max_frames).long().tolist()
+        trajectory_list = [trajectory_list[i] for i in selected_indices]
+        n_steps = len(trajectory_list)
+
+    # ... 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="cif.gz", include_entity_poly=False
+        )
+
+
+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()
+
+    assert (
+        coords.shape[-2] == atom_array.array_length()
+    ), f"N batched coordinates {coords.shape} != {atom_array.array_length()}"
+
+    # (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 find_files_with_extension(path: PathLike, supported_file_types: list) -> list[Path]:
+    """Recursively find all files with the given extensions in the specified path.
+
+    Args:
+        path (PathLike): Path to the directory containing the files.
+        supported_file_types (list): List of supported file extensions.
+
+    Returns:
+        list[Path]: List of files with the given extensions.
+    """
+    files_with_supported_types = []
+    path = Path(path)
+
+    # Check if the path is a directory
+    if path.is_dir():
+        # Search for files with each supported extension
+        for file_type in supported_file_types:
+            files_with_supported_types.extend(path.glob(f"*{file_type}"))
+    elif path.is_file() and path.suffix in supported_file_types:
+        # If it's a file and has a supported extension, add to the list
+        files_with_supported_types.append(path)
+
+    return files_with_supported_types
+
+
+def create_example_id_extractor(extensions: set | list = CIF_LIKE_EXTENSIONS) -> str:
+    """Create a function with closure that extracts example_ids from file paths with specified extensions.
+
+    Example:
+        >>> extractor = create_example_id_extractor({".cif", ".cif.gz"})
+        >>> extractor("example.path.example_id.cif.gz")
+        'example_id'
+    """
+    pattern = re.compile(
+        "(" + "|".join(re.escape(ext) + "$" for ext in extensions) + ")"
+    )
+
+    def extract_id(file_path: PathLike) -> str:
+        """Extract example_id from file path."""
+        # Remove extension and get last part after splitting by dots
+        without_ext = pattern.sub("", Path(file_path).name)
+        return without_ext.split(".")[-1]
+
+    return extract_id
+
+
+def extract_example_id_from_path(file_path: PathLike, extensions: set | list) -> str:
+    """Extract example_id from file path with specified extensions."""
+    extractor = create_example_id_extractor(extensions)
+    return extractor(file_path)

rfd3/utils/vizualize.py

@@ -0,0 +1,276 @@
+#!/usr/bin/env -S /bin/sh -c '"$(dirname "$0")/../../../scripts/shebang/modelhub_exec.sh" "$0" "$@"'
+"""
+If you add the `/path/to/visualize.py` to your .bashrc/.zshrc like this:
+
+```bash
+viz () {
+    /path/to/visualize.py "$@"
+}
+```
+
+Then you can run `viz /path/to/file.cif` to visualize the structures via pymol-remote.
+"""
+
+import pathlib
+import sys
+
+# NOTE: This is needed here to enable `viz` to be used as script
+if (project_dir := str(pathlib.Path(__file__).parents[3])) not in sys.path:
+    sys.path.append(project_dir)
+
+import logging
+
+import biotite.structure as struc
+import numpy as np
+from atomworks.io.utils.visualize import get_pymol_session, view_pymol
+from atomworks.ml.conditions import C_DIS, Condition
+from atomworks.ml.conditions.base import Level
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+def get_atom_array_style_cmd(
+    atom_array: struc.AtomArray | struc.AtomArrayStack,
+    obj: str,
+    label: bool = False,
+    max_distances: int = 100,
+    grid_slot: int | None = None,
+) -> str:
+    """Generate PyMOL commands to style an atom array visualization.
+
+    Creates a series of PyMOL commands that style different parts of a molecular structure, including:
+    - Applying a color spectrum to polymer chains
+    - Showing backbone atoms as sticks and CA atoms as spheres
+    - Styling non-polymer atoms with different colors and representation
+    - Highlighting specially annotated atoms with different colors and visualizations
+
+    Args:
+        atom_array: The biotite AtomArray or AtomArrayStack to be styled
+        obj: PyMOL object name to apply the styling to
+        label: Whether to label all atoms with their 0-indexed atom_id
+        max_num_dist_lines: Maximum number of distance lines to show (pymol hangs when there's too many distance objects)
+
+    Returns:
+        str: A PyMOL command string that styles the atom array
+    """
+    grid_slot = grid_slot or np.random.randint(0, 10_000)
+    commands = [f"hide everything, {obj}"]
+    annotations = atom_array.get_annotation_categories()
+
+    offset = 1  # ... default offset since pymol 1-indexes atom id's
+    atom_ids = np.arange(1, atom_array.array_length() + 1)
+    if "atom_id" in annotations:
+        offset = 0
+        atom_ids = atom_array.get_annotation("atom_id")
+
+    # Convert MAS residues to ALA for compatibility
+    atom_array = atom_array.copy()
+    # atom_array.res_name[atom_array.res_name == "MAS"] = "ALA"
+
+    # Style the backbone for each polymer chain with a color spectrum
+    for chain_id in struc.get_chains(atom_array):
+        if (~atom_array.hetero[atom_array.chain_id == chain_id]).any():
+            commands.append(
+                f"spectrum resi, RFd_darkblue RFd_blue RFd_lightblue RFd_purple RFd_pink RFd_melon RFd_navaho, "
+                f"{obj} and chain {chain_id} and elem C"
+            )
+
+    # Add basic styling commands for protein backbone and non-polymer components
+    commands.extend(
+        [
+            f"show sticks, model {obj} and name n+c+ca+cb",
+            f"show spheres, model {obj} and name ca",
+            f"set sphere_scale, 0.23, model {obj} and name ca",
+            f"set sphere_transparency, 0, model {obj} and name ca",
+            f"color grey60, model {obj} and not polymer and elem C",
+            f"show nb_spheres, model {obj} and not polymer",
+            f"show sticks, model {obj} and not polymer",
+        ]
+    )
+    if label:
+        # label 0-indexed for correspondence with biotite
+        commands.append("label all, ID") if offset == 0 else commands.append(
+            f"label all, ID-{offset}"
+        )
+
+    # Style atoms marked for "atomize" if present
+    if "atomize" in annotations and atom_array.atomize.any():
+        atomize_ids = np.where(atom_array.atomize)[0]
+        atomize_ids = atom_ids[atomize_ids]
+        commands.extend(
+            [
+                f"select {obj}_atomize, model {obj} and id {'+'.join(str(id) for id in atomize_ids)}",
+                f"show sticks, {obj}_atomize and byres {obj}_atomize",
+            ]
+        )
+
+    # Style constraints, if present:
+    # 2-body:
+    # ... add distance lines between atoms if specified in annotations
+    if hasattr(atom_array, "_annot_2d"):
+        distance_commands = []
+        if C_DIS.full_name in atom_array._annot_2d:
+            constraint_data = C_DIS.annotation(atom_array).as_array()
+            if len(constraint_data) > 0:
+                _atom_idxs = np.unique(constraint_data[:, :2].flatten()).astype(int)
+                _atom_ids = atom_ids[_atom_idxs]
+                _selection = f"{obj} and id {'+'.join(str(id) for id in _atom_ids)}"
+                commands.extend(
+                    [
+                        f"delete m2d_{obj}",
+                        f"select m2d_{obj}, {_selection}",
+                        f"show spheres, m2d_{obj}",
+                        f"set sphere_color, lime, m2d_{obj}",
+                        f"set sphere_scale, 0.25, m2d_{obj}",
+                        f"set sphere_transparency, 0.5, m2d_{obj}",
+                        f"show sticks, byres m2d_{obj}",
+                    ]
+                )
+
+                if len(constraint_data) > max_distances:
+                    logger.warning(
+                        f"Too many distance conditions ({len(constraint_data)}), sampling {max_distances}."
+                    )
+                    constraint_idxs = np.arange(len(constraint_data))
+                    constraint_idxs = np.random.choice(
+                        constraint_idxs, max_distances, replace=False
+                    )
+                    constraint_data = constraint_data[constraint_idxs]
+
+                for row in constraint_data:
+                    idx_i, idx_j, value = row
+                    if (idx_i > idx_j) or (value == 0):
+                        continue
+
+                    i, j = atom_ids[idx_i], atom_ids[idx_j]
+                    # ... if we have a stack, we grab the last frame for the distance computation
+                    if isinstance(atom_array, struc.AtomArrayStack):
+                        distance = struc.distance(
+                            atom_array[0, idx_i], atom_array[0, idx_j]
+                        )
+                    else:
+                        distance = struc.distance(atom_array[idx_i], atom_array[idx_j])
+
+                    distance_name = f"d{idx_i}-{idx_j}_{value:.2f}_{distance:.2f}"
+                    distance_commands.extend(
+                        [
+                            f"distance {distance_name}, model {obj} and id {i}, model {obj} and id {j}",
+                            f"set grid_slot, {grid_slot}, {distance_name}",
+                        ]
+                    )
+
+        commands.extend(distance_commands)
+
+    # Handle 1-D conditions (only display masks)
+    for cond in Condition:
+        if cond.n_body == 1 and cond.mask(atom_array, default="generate").any():
+            _atom_ids = atom_ids[np.where(cond.mask(atom_array, default="generate"))[0]]
+            if cond.level == Level.ATOM:
+                _selection = (
+                    f"model {obj} and id {'+'.join(str(id) for id in _atom_ids)}"
+                )
+                commands.extend([f"select {cond.mask_name}_{obj}, {_selection}"])
+            elif cond.level == Level.RESIDUE or cond.level == Level.TOKEN:
+                _selection = f"model {obj} and byres (id {'+'.join(str(id) for id in _atom_ids)})"
+                commands.extend([f"select {cond.mask_name}_{obj}, {_selection}"])
+
+    return "\n".join(commands)
+
+
+def viz(
+    atom_array: struc.AtomArray | struc.AtomArrayStack,
+    id: str = "obj",
+    clear: bool = True,
+    label: bool = True,
+    max_distances: int = 100,
+    view_ori_token: bool = False,
+) -> None:
+    """Quickly visualize a molecular structure in PyMOL with predefined styling.
+
+    This function creates a PyMOL session, loads the atom array structure, and applies
+    a set of styling commands to make the visualization informative and aesthetically pleasing.
+    The styling highlights different structural components and annotated features.
+
+    Args:
+        atom_array: The biotite AtomArray or AtomArrayStack to visualize
+        id: PyMOL object identifier (default: "obj")
+        clear: Whether to clear existing PyMOL objects before visualization (default: True)
+        label: Whether to label all atoms with their 0-indexed atom_id (default: True)
+        view_ori_token: Whether to view the ori token (default: False)
+
+    Example:
+        >>> from biotite.structure import AtomArray
+        >>> # Create or load an atom array
+        >>> structure = AtomArray(...)
+        >>> # Visualize the structure in PyMOL
+        >>> viz(structure)
+    """
+    atom_array = atom_array.copy()
+    # pymol only considers chain_id annotation, which can make weird looking artifacts if we have two different chains with the same chain_id
+    # We always disambiguate by using the chain_iid annotation, so we need to have pymol use that to do the same
+    if "chain_iid" in atom_array.get_annotation_categories():
+        atom_array.chain_id = atom_array.get_annotation("chain_iid")
+        atom_array.chain_id = atom_array.chain_id.astype(str)
+
+    pymol = get_pymol_session()
+    pymol.do("set valence, 1; set connect_mode, 2;")
+    if clear:
+        pymol.do("delete d*")
+        pymol.delete("all")
+    slot = np.random.randint(0, 10_000)
+    obj_name = view_pymol(atom_array, id=id, grid_slot=slot)
+    cmd = get_atom_array_style_cmd(
+        atom_array, obj_name, label=label, grid_slot=slot, max_distances=max_distances
+    )
+    pymol.do(cmd)
+
+    if view_ori_token:
+        pymol.do(f"pseudoatom ori_{obj_name}, pos=[0,0,0]")
+        pymol.do(
+            [
+                f"set grid_slot, {slot}, ori_{obj_name}",
+                f"show spheres, ori_{obj_name}",
+                f"set sphere_color, white, ori_{obj_name}",
+                f"set sphere_scale, 0.5, ori_{obj_name}",
+                f"set sphere_transparency, 0.5, ori_{obj_name}",
+            ]
+        )
+
+
+def _viz_from_file(
+    file_path: str,
+    id: str = "obj",
+    clear: bool = True,
+    label: bool = True,
+    max_distances: int = 100,
+):
+    if file_path.endswith(".pkl.gz"):
+        import gzip
+        import pickle
+
+        with gzip.open(file_path, "rb") as f:
+            atom_array = pickle.load(f)
+    elif file_path.endswith(".pkl"):
+        import pickle
+
+        with open(file_path, "rb") as f:
+            atom_array = pickle.load(f)
+    elif file_path.endswith((".cif", ".cif.gz", ".bcif", ".bcif.gz")):
+        from atomworks.io.utils.io_utils import get_structure, read_any
+        from rfd3.utils.inference import (
+            _add_design_annotations_from_cif_block_metadata,
+        )
+
+        cif_file = read_any(file_path)
+        atom_array = get_structure(cif_file, include_bonds=True, extra_fields="all")
+        atom_array = _add_design_annotations_from_cif_block_metadata(
+            atom_array, cif_file.block
+        )
+    viz(atom_array, id=id, clear=clear, label=label, max_distances=max_distances)
+
+
+if __name__ == "__main__":
+    import fire
+
+    fire.Fire(_viz_from_file)