rc-foundry 0.1.6__py3-none-any.whl → 0.1.9__py3-none-any.whl

rfd3/inference/symmetry/frames.py

@@ -10,12 +10,13 @@ def get_symmetry_frames_from_symmetry_id(symmetry_id):
     Returns:
         frames: list of rotation matrices
     """
+    from rfd3.inference.symmetry.symmetry_utils import SymmetryConfig
 
     # Get frames from symmetry id
     sym_conf = {}
-    if isinstance(symmetry_id, dict):
+    if isinstance(symmetry_id, SymmetryConfig):
         sym_conf = symmetry_id
-        symmetry_id = symmetry_id.get("id")
+        symmetry_id = symmetry_id.id
 
     if symmetry_id.lower().startswith("c"):
         order = int(symmetry_id[1:])
@@ -23,11 +24,17 @@ def get_symmetry_frames_from_symmetry_id(symmetry_id):
     elif symmetry_id.lower().startswith("d"):
         order = int(symmetry_id[1:])
         frames = get_dihedral_frames(order)
+    elif symmetry_id.lower() == "t":
+        frames = get_tetrahedral_frames()
+    elif symmetry_id.lower() == "o":
+        frames = get_octahedral_frames()
+    elif symmetry_id.lower() == "i":
+        frames = get_icosahedral_frames()
     elif symmetry_id.lower() == "input_defined":
         assert (
-            "symmetry_file" in sym_conf
+            sym_conf.symmetry_file is not None
         ), "symmetry_file is required for input_defined symmetry"
-        frames = get_frames_from_file(sym_conf.get("symmetry_file"))
+        frames = get_frames_from_file(sym_conf.symmetry_file)
     else:
         raise ValueError(f"Symmetry id {symmetry_id} not supported")
 
@@ -120,7 +127,9 @@ def get_symmetry_frames_from_atom_array(src_atom_array, input_frames):
     computed_frames = [(R, np.array([0, 0, 0])) for R in Rs]
 
     # check that the computed frames match the input frames
-    check_input_frames_match_symmetry_frames(computed_frames, input_frames)
+    check_input_frames_match_symmetry_frames(
+        computed_frames, input_frames, nids_by_entity
+    )
 
     return computed_frames
 
@@ -277,6 +286,248 @@ def get_dihedral_frames(order):
     return frames
 
 
+def get_tetrahedral_frames():
+    """
+    Get tetrahedral frames (T symmetry group, 12 elements).
+    Returns:
+        frames: list of rotation matrices
+    """
+
+    frames = []
+
+    # Identity
+    frames.append((np.eye(3), np.array([0, 0, 0])))
+
+    # 8 rotations by ±120° around body diagonals (±1, ±1, ±1)
+    diagonals = [
+        np.array([1, 1, 1]),
+        np.array([1, -1, -1]),
+        np.array([-1, 1, -1]),
+        np.array([-1, -1, 1]),
+    ]
+    for d in diagonals:
+        axis = d / np.linalg.norm(d)
+        for angle in [2 * np.pi / 3, 4 * np.pi / 3]:
+            R = _rotation_matrix_from_axis_angle(axis, angle)
+            frames.append((R, np.array([0, 0, 0])))
+
+    # 3 rotations by 180° around coordinate axes
+    for axis in [np.array([1, 0, 0]), np.array([0, 1, 0]), np.array([0, 0, 1])]:
+        R = _rotation_matrix_from_axis_angle(axis, np.pi)
+        frames.append((R, np.array([0, 0, 0])))
+
+    return frames
+
+
+def get_octahedral_frames():
+    """
+    Get octahedral frames (O symmetry group, 24 elements).
+    The axes are computed from the geometry of a cube with vertices at (±1, ±1, ±1).
+    Returns:
+        frames: list of rotation matrices
+    """
+
+    frames = []
+
+    # 8 vertices of the cube
+    vertices = []
+    for s1 in [1, -1]:
+        for s2 in [1, -1]:
+            for s3 in [1, -1]:
+                vertices.append(np.array([s1, s2, s3]))
+    vertices = np.array(vertices)
+
+    # 6 face centers of the cube (4-fold axes pass through these)
+    face_centers = [
+        np.array([1, 0, 0]),
+        np.array([-1, 0, 0]),
+        np.array([0, 1, 0]),
+        np.array([0, -1, 0]),
+        np.array([0, 0, 1]),
+        np.array([0, 0, -1]),
+    ]
+
+    # Find edges (pairs of vertices differing in exactly one coordinate)
+    edges = []
+    for i in range(len(vertices)):
+        for j in range(i + 1, len(vertices)):
+            diff = np.abs(vertices[i] - vertices[j])
+            if np.sum(diff > 0) == 1:  # Differ in exactly one coordinate
+                edges.append((i, j))
+
+    # Helper to get unique axis (normalize direction to avoid duplicates)
+    def normalize_axis(v):
+        axis = v / np.linalg.norm(v)
+        for c in axis:
+            if abs(c) > 1e-10:
+                if c < 0:
+                    axis = -axis
+                break
+        return tuple(np.round(axis, 10))
+
+    # Identity
+    frames.append((np.eye(3), np.array([0, 0, 0])))
+
+    # 4-fold axes (through opposite face centers) - 3 axes
+    # Each gives rotations at 90°, 180°, 270° (we skip 0° = identity)
+    fourfold_axes_set = set()
+    for fc in face_centers:
+        axis_tuple = normalize_axis(fc)
+        fourfold_axes_set.add(axis_tuple)
+
+    for axis_tuple in fourfold_axes_set:
+        axis = np.array(axis_tuple)
+        for k in [1, 2, 3]:  # 90°, 180°, 270°
+            angle = np.pi * k / 2
+            R = _rotation_matrix_from_axis_angle(axis, angle)
+            frames.append((R, np.array([0, 0, 0])))
+
+    # 3-fold axes (through opposite vertices) - 4 axes
+    # Each gives rotations at 120°, 240°
+    threefold_axes_set = set()
+    for v in vertices:
+        axis_tuple = normalize_axis(v)
+        threefold_axes_set.add(axis_tuple)
+
+    for axis_tuple in threefold_axes_set:
+        axis = np.array(axis_tuple)
+        for angle in [2 * np.pi / 3, 4 * np.pi / 3]:
+            R = _rotation_matrix_from_axis_angle(axis, angle)
+            frames.append((R, np.array([0, 0, 0])))
+
+    # 2-fold axes (through opposite edge midpoints) - 6 axes
+    # Each gives 1 rotation at 180°
+    twofold_axes_set = set()
+    for i, j in edges:
+        midpoint = (vertices[i] + vertices[j]) / 2
+        axis_tuple = normalize_axis(midpoint)
+        twofold_axes_set.add(axis_tuple)
+
+    for axis_tuple in twofold_axes_set:
+        axis = np.array(axis_tuple)
+        R = _rotation_matrix_from_axis_angle(axis, np.pi)
+        frames.append((R, np.array([0, 0, 0])))
+
+    return frames
+
+
+def get_icosahedral_frames():
+    """
+    Get icosahedral frames (I symmetry group, 60 elements).
+    The axes are computed from the geometry of a regular icosahedron with
+    vertices at (0, ±1, ±φ), (±1, ±φ, 0), (±φ, 0, ±1) where φ is the golden ratio.
+    Returns:
+        frames: list of rotation matrices
+    """
+
+    frames = []
+
+    # Golden ratio
+    phi = (1 + np.sqrt(5)) / 2
+
+    # 12 vertices of the icosahedron
+    vertices = []
+    for s1 in [1, -1]:
+        for s2 in [1, -1]:
+            vertices.append(np.array([0, s1 * 1, s2 * phi]))
+            vertices.append(np.array([s1 * 1, s2 * phi, 0]))
+            vertices.append(np.array([s2 * phi, 0, s1 * 1]))
+    vertices = np.array(vertices)
+
+    # Find edges (pairs of vertices at distance 2)
+    edges = []
+    for i in range(len(vertices)):
+        for j in range(i + 1, len(vertices)):
+            dist_sq = np.sum((vertices[i] - vertices[j]) ** 2)
+            if np.isclose(dist_sq, 4.0):
+                edges.append((i, j))
+
+    # Find faces (triangles of mutually adjacent vertices)
+    edge_set = set(edges)
+    faces = []
+    for i in range(len(vertices)):
+        for j in range(i + 1, len(vertices)):
+            for k in range(j + 1, len(vertices)):
+                if (i, j) in edge_set and (j, k) in edge_set and (i, k) in edge_set:
+                    faces.append((i, j, k))
+
+    # Helper to get unique axis (normalize direction to avoid duplicates)
+    def normalize_axis(v):
+        axis = v / np.linalg.norm(v)
+        # Make first significant component positive to avoid duplicate opposite axes
+        for c in axis:
+            if abs(c) > 1e-10:
+                if c < 0:
+                    axis = -axis
+                break
+        return tuple(np.round(axis, 10))
+
+    # Identity
+    frames.append((np.eye(3), np.array([0, 0, 0])))
+
+    # 5-fold axes (through opposite vertices) - 6 axes, 4 rotations each = 24
+    fivefold_axes_set = set()
+    for v in vertices:
+        axis_tuple = normalize_axis(v)
+        fivefold_axes_set.add(axis_tuple)
+
+    for axis_tuple in fivefold_axes_set:
+        axis = np.array(axis_tuple)
+        for k in [1, 2, 3, 4]:
+            angle = 2 * np.pi * k / 5
+            R = _rotation_matrix_from_axis_angle(axis, angle)
+            frames.append((R, np.array([0, 0, 0])))
+
+    # 3-fold axes (through opposite face centers) - 10 axes, 2 rotations each = 20
+    threefold_axes_set = set()
+    for i, j, k in faces:
+        center = (vertices[i] + vertices[j] + vertices[k]) / 3
+        axis_tuple = normalize_axis(center)
+        threefold_axes_set.add(axis_tuple)
+
+    for axis_tuple in threefold_axes_set:
+        axis = np.array(axis_tuple)
+        for angle in [2 * np.pi / 3, 4 * np.pi / 3]:
+            R = _rotation_matrix_from_axis_angle(axis, angle)
+            frames.append((R, np.array([0, 0, 0])))
+
+    # 2-fold axes (through opposite edge midpoints) - 15 axes, 1 rotation each = 15
+    twofold_axes_set = set()
+    for i, j in edges:
+        midpoint = (vertices[i] + vertices[j]) / 2
+        axis_tuple = normalize_axis(midpoint)
+        twofold_axes_set.add(axis_tuple)
+
+    for axis_tuple in twofold_axes_set:
+        axis = np.array(axis_tuple)
+        R = _rotation_matrix_from_axis_angle(axis, np.pi)
+        frames.append((R, np.array([0, 0, 0])))
+
+    return frames
+
+
+def _rotation_matrix_from_axis_angle(axis, angle):
+    """
+    Compute a rotation matrix from an axis and angle using Rodrigues' formula.
+    Arguments:
+        axis: unit vector of the rotation axis
+        angle: rotation angle in radians
+    Returns:
+        R: 3x3 rotation matrix
+    """
+
+    axis = axis / np.linalg.norm(axis)
+    K = np.array(
+        [
+            [0, -axis[2], axis[1]],
+            [axis[2], 0, -axis[0]],
+            [-axis[1], axis[0], 0],
+        ]
+    )
+    R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)
+    return R
+
+
 def get_frames_from_file(file_path):
     raise NotImplementedError("Input defined symmetry not implemented")
 

rfd3/inference/symmetry/symmetry_utils.py

@@ -39,18 +39,36 @@ ranked_logger = RankedLogger(__name__, rank_zero_only=True)
 
 
 class SymmetryConfig(BaseModel):
-    # AM / HE TODO: feel free to flesh this out and add validation as needed
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
         extra="allow",
     )
-    id: Optional[str] = Field(None)
-    # is_unsym_motif: Optional[np.ndarray[bool]] = Field(...)
-    # is_symmetric_motif: bool = Field(...)
+    id: Optional[str] = Field(
+        None,
+        description="Symmetry group ID. Supported types: Cyclic (C), Dihedral (D), Tetrahedral (T), Octahedral (O), Icosahedral (I). e.g. 'C3', 'D2', 'T', 'O', 'I'.",
+    )
+    is_unsym_motif: Optional[str] = Field(
+        None,
+        description="Comma separated list of contig/ligand names that should not be symmetrized such as DNA strands. \
+         e.g. 'HEM' or 'Y1-11,Z16-25'",
+    )
+    is_symmetric_motif: bool = Field(
+        True,
+        description="If True, the input motifs are expected to be already symmetric and won't be symmetrized. \
+        If False, the all input motifs are expected to be ASU and will be symmetrized.",
+    )
+
+
+def convery_sym_conf_to_symmetry_config(sym_conf: dict):
+    return SymmetryConfig(**sym_conf)
 
 
 def make_symmetric_atom_array(
-    asu_atom_array, sym_conf: SymmetryConfig, sm=None, has_2d=False, src_atom_array=None
+    asu_atom_array,
+    sym_conf: SymmetryConfig | dict,
+    sm=None,
+    has_dist_cond=False,
+    src_atom_array=None,
 ):
     """
     apply symmetry to an atom array.
@@ -58,41 +76,35 @@ def make_symmetric_atom_array(
         asu_atom_array: atom array of the asymmetric unit
         sym_conf: symmetry configuration (dict, "id" key is required)
         sm: optional small molecule names (str, comma separated)
-        has_2d: whether to add 2d entity annotations
+        has_dist_cond: whether to add 2d entity annotations
     Returns:
         new_asu_atom_array: atom array with symmetry applied
     """
-    sym_conf = (
-        sym_conf.model_dump()
-    )  # TODO: JB: remove this line to keep as symmetry config for cleaner syntax(?)
-    ranked_logger.info(f"Symmetry Configs: {sym_conf}")
+    if not isinstance(sym_conf, SymmetryConfig):
+        sym_conf = convery_sym_conf_to_symmetry_config(sym_conf)
 
-    # Making sure that the symmetry config is valid
     check_symmetry_config(
-        asu_atom_array,
-        sym_conf,
-        sm,
-        has_dist_cond=has_2d,
-        src_atom_array=src_atom_array,
+        asu_atom_array, sym_conf, sm, has_dist_cond, src_atom_array=src_atom_array
     )
     # Adding utility annotations to the asu atom array
     asu_atom_array = _add_util_annotations(asu_atom_array, sym_conf, sm)
 
-    if has_2d:  # NB: this will only work for asymmetric motifs at the moment - need to add functionality for symmetric motifs
+    if has_dist_cond:  # NB: this will only work for asymmetric motifs at the moment - need to add functionality for symmetric motifs
         asu_atom_array = add_2d_entity_annotations(asu_atom_array)
 
     frames = get_symmetry_frames_from_symmetry_id(sym_conf)
 
     # If the motif is symmetric, we get the frames instead from the source atom array.
-    if sym_conf.get("is_symmetric_motif"):
+    if sym_conf.is_symmetric_motif:
         assert (
             src_atom_array is not None
         ), "Source atom array must be provided for symmetric motifs"
-        # if symmetric motif is provided, get the frames from the src atom array
+        # if symmetric motif is provided, get the frames from the src atom array.
         frames = get_symmetry_frames_from_atom_array(src_atom_array, frames)
     else:
-        raise NotImplementedError(
-            "Asymmetric motif inputs are not implemented yet. please symmetrize the motif."
+        # At this point, asym case would have been caught by the check_symmetry_config function.
+        ranked_logger.info(
+            "No motifs found in atom array. Generating unconditional symmetric proteins."
         )
 
     # Add symmetry annotations to the asu atom array
@@ -101,7 +113,7 @@ def make_symmetric_atom_array(
     # Extracting all things at this moment that we will not want to symmetrize.
     # This includes: 1) unsym motifs, 2) ligands
     unsym_atom_arrays = []
-    if sym_conf.get("is_unsym_motif"):
+    if sym_conf.is_unsym_motif:
         # unsym_motif_atom_array = get_unsym_motif(asu_atom_array, asu_atom_array._is_unsym_motif)
         # Now remove the unsym motifs from the asu atom array
         unsym_atom_arrays.append(asu_atom_array[asu_atom_array._is_unsym_motif])
@@ -128,7 +140,7 @@ def make_symmetric_atom_array(
     symmetrized_atom_array = struc.concatenate(symmetry_unit_list)
 
     # add 2D conditioning annotations
-    if has_2d:
+    if has_dist_cond:
         symmetrized_atom_array = reannotate_2d_conditions(symmetrized_atom_array)
 
     # set all motifs to not have any symmetrization applied to them
@@ -183,7 +195,7 @@ def make_symmetric_atom_array_for_partial_diffusion(atom_array, sym_conf):
     frames = get_symmetry_frames_from_symmetry_id(sym_conf)
 
     # Add symmetry ID
-    symmetry_ids = np.full(n, sym_conf.get("id"), dtype="U6")
+    symmetry_ids = np.full(n, sym_conf.id, dtype="U6")
     atom_array.set_annotation("symmetry_id", symmetry_ids)
 
     # Initialize transform annotations (use same format as original system)
@@ -244,7 +256,7 @@ def _add_util_annotations(asu_atom_array, sym_conf, sm):
     """
     n = asu_atom_array.shape[0]
     is_motif = get_motif_features(asu_atom_array)["is_motif_atom"].astype(np.bool_)
-    is_sm = np.zeros(asu_atom_array.shape[0], dtype=bool)
+    is_sm = np.zeros(n, dtype=bool)
     is_asu = np.ones(n, dtype=bool)
     is_unsym_motif = np.zeros(n, dtype=bool)
 
@@ -257,8 +269,8 @@ def _add_util_annotations(asu_atom_array, sym_conf, sm):
         )
 
     # assign unsym motifs
-    if sym_conf.get("is_unsym_motif"):
-        unsym_motif_names = sym_conf["is_unsym_motif"].split(",")
+    if sym_conf.is_unsym_motif:
+        unsym_motif_names = sym_conf.is_unsym_motif.split(",")
         unsym_motif_names = expand_contig_unsym_motif(unsym_motif_names)
         is_unsym_motif = get_unsym_motif_mask(asu_atom_array, unsym_motif_names)
 
@@ -361,38 +373,4 @@ def apply_symmetry_to_xyz_atomwise(X_L, sym_feats, partial_diffusion=False):
                 "blc,cd->bld", asu_xyz, sym_transforms[target_id][0].to(asu_xyz.dtype)
             ) + sym_transforms[target_id][1].to(asu_xyz.dtype)
 
-    # Log inter-chain distances for debugging - use actual chain annotations
-    if sym_X_L.shape[1] > 100:  # Only for large structures
-        # Use symmetry entity annotations to find different chains
-        sym_entity_id = sym_feats["sym_entity_id"]
-        unique_entities = torch.unique(sym_entity_id)
-
-        if len(unique_entities) >= 2:
-            # Get atoms from first two different entities
-            entity_0_mask = sym_entity_id == unique_entities[0]
-            entity_1_mask = sym_entity_id == unique_entities[1]
-
-            if entity_0_mask.sum() > 0 and entity_1_mask.sum() > 0:
-                entity_0_atoms = sym_X_L[0, entity_0_mask, :]
-                entity_1_atoms = sym_X_L[0, entity_1_mask, :]
-
-                # Sample subset to avoid memory issues
-                entity_0_sample = entity_0_atoms[: min(50, entity_0_atoms.shape[0]), :]
-                entity_1_sample = entity_1_atoms[: min(50, entity_1_atoms.shape[0]), :]
-
-                min_distance = (
-                    torch.cdist(entity_0_sample, entity_1_sample).min().item()
-                )
-                ranked_logger.info(
-                    f"Min inter-chain distance after symmetry: {min_distance:.2f} Å"
-                )
-
-                # Also log the centers of each entity
-                entity_0_center = entity_0_atoms.mean(dim=0)
-                entity_1_center = entity_1_atoms.mean(dim=0)
-                center_distance = torch.norm(entity_0_center - entity_1_center).item()
-                ranked_logger.info(
-                    f"Distance between chain centers: {center_distance:.2f} Å"
-                )
-
     return sym_X_L

rfd3/model/inference_sampler.py

@@ -1,11 +1,15 @@
 import inspect
+import time
 from dataclasses import dataclass
 from typing import Any, Literal
 
 import torch
 from jaxtyping import Float
+from rfd3.inference.symmetry.symmetry_utils import apply_symmetry_to_xyz_atomwise
+from rfd3.model.cfg_utils import strip_X
 
 from foundry.common import exists
+from foundry.utils.alignment import weighted_rigid_align
 from foundry.utils.ddp import RankedLogger
 from foundry.utils.rotation_augmentation import (
     rot_vec_mul,
@@ -110,7 +114,7 @@ class SampleDiffusionWithMotif(SampleDiffusionConfig):
                 )
                 # Fallback to smallest available step
                 noise_schedule_original = self._construct_inference_noise_schedule(
-                    device=coord_atom_lvl_to_be_noised.device
+                    device=device
                 )
                 noise_schedule = noise_schedule_original[-1:]  # Just use the final step
                 ranked_logger.info(
@@ -221,6 +225,7 @@ class SampleDiffusionWithMotif(SampleDiffusionConfig):
             # Handle chunked mode vs standard mode
             if "chunked_pairwise_embedder" in initializer_outputs:
                 # Chunked mode: explicitly provide P_LL=None
+                tic = time.time()
                 chunked_embedder = initializer_outputs[
                     "chunked_pairwise_embedder"
                 ]  # Don't pop, just get
@@ -238,6 +243,8 @@ class SampleDiffusionWithMotif(SampleDiffusionConfig):
                     initializer_outputs=other_outputs,
                     **other_outputs,
                 )
+                toc = time.time()
+                ranked_logger.info(f"Chunked mode time: {toc - tic} seconds")
             else:
                 # Standard mode: P_LL is included in initializer_outputs
                 outs = diffusion_module(
@@ -445,6 +452,7 @@ class SampleDiffusionWithSymmetry(SampleDiffusionWithMotif):
             # Handle chunked mode vs standard mode (same as default sampler)
             if "chunked_pairwise_embedder" in initializer_outputs:
                 # Chunked mode: explicitly provide P_LL=None
+                tic = time.time()
                 chunked_embedder = initializer_outputs[
                     "chunked_pairwise_embedder"
                 ]  # Don't pop, just get
@@ -462,6 +470,8 @@ class SampleDiffusionWithSymmetry(SampleDiffusionWithMotif):
                     initializer_outputs=other_outputs,
                     **other_outputs,
                 )
+                toc = time.time()
+                ranked_logger.info(f"Chunked mode time: {toc - tic} seconds")
             else:
                 # Standard mode: P_LL is included in initializer_outputs
                 outs = diffusion_module(