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

rfd3/inference/symmetry/frames.py

@@ -24,6 +24,12 @@ 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 (
             sym_conf.symmetry_file is not None
@@ -280,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

@@ -45,7 +45,7 @@ class SymmetryConfig(BaseModel):
     )
     id: Optional[str] = Field(
         None,
-        description="Symmetry group ID. e.g. 'C3', 'D2'. Only C and D symmetry types are supported currently.",
+        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,
@@ -101,10 +101,10 @@ def make_symmetric_atom_array(
         ), "Source atom array must be provided for symmetric motifs"
         # if symmetric motif is provided, get the frames from the src atom array.
         frames = get_symmetry_frames_from_atom_array(src_atom_array, frames)
-    elif (asu_atom_array._is_motif[~asu_atom_array._is_unsym_motif]).any():
-        # if the motifs that's not unsym motifs are present.
-        raise NotImplementedError(
-            "Asymmetric motif inputs are not implemented yet. please symmetrize the motif."
+    else:
+        # 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

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(

rfd3/model/layers/block_utils.py

@@ -118,14 +118,14 @@ def scatter_add_pair_features(P_LK_tgt, P_LK_indices, P_LA_src, P_LA_indices):
 
     Parameters
     ----------
-    P_LK_indices   : (D, L, k) LongTensor
+    P_LK_indices   : (B, L, k) LongTensor
         Key indices | P_LK_indices[d, i, k] = global atom index for which atom i attends to.
-    P_LK : (D, L, k, c) FloatTensor
+    P_LK : (B, L, k, c) FloatTensor
         Key features to scatter add into
 
-    P_LA_indices   : (D, L, a) LongTensor
+    P_LA_indices   : (B, L, a) LongTensor
         Additional feature indices to scatter into P_LK.
-    P_LA : (D, L, a, c) FloatTensor
+    P_LA : (B, L, a, c) FloatTensor
         Features corresponding to P_LA.
 
     Both index tensors contain indices representing D batch dim,
@@ -135,42 +135,42 @@ def scatter_add_pair_features(P_LK_tgt, P_LK_indices, P_LA_src, P_LA_indices):
 
     """
     # Handle case when indices and P_LA don't have batch dimensions
-    D, L, k = P_LK_indices.shape
+    B, L, k = P_LK_indices.shape
     if P_LA_indices.ndim == 2:
-        P_LA_indices = P_LA_indices.unsqueeze(0).expand(D, -1, -1)
+        P_LA_indices = P_LA_indices.unsqueeze(0).expand(B, -1, -1)
     if P_LA_src.ndim == 3:
-        P_LA_src = P_LA_src.unsqueeze(0).expand(D, -1, -1)
+        P_LA_src = P_LA_src.unsqueeze(0).expand(B, -1, -1)
     assert (
         P_LA_src.shape[-1] == P_LK_tgt.shape[-1]
     ), "Channel dims do not match, got: {} vs {}".format(
         P_LA_src.shape[-1], P_LK_tgt.shape[-1]
     )
 
-    matches = P_LA_indices.unsqueeze(-1) == P_LK_indices.unsqueeze(-2)  # (D, L, a, k)
+    matches = P_LA_indices.unsqueeze(-1) == P_LK_indices.unsqueeze(-2)  # (B, L, a, k)
     if not torch.all(matches.sum(dim=(-1, -2)) >= 1):
         raise ValueError("Found multiple scatter indices for some atoms")
     elif not torch.all(matches.sum(dim=-1) <= 1):
         raise ValueError("Did not find a scatter index for every atom")
-    k_indices = matches.long().argmax(dim=-1)  # (D, L, a)
+    k_indices = matches.long().argmax(dim=-1)  # (B, L, a)
     scatter_indices = k_indices.unsqueeze(-1).expand(
         -1, -1, -1, P_LK_tgt.shape[-1]
-    )  # (D, L, a, c)
+    )  # (B, L, a, c)
     P_LK_tgt = P_LK_tgt.scatter_add(dim=2, index=scatter_indices, src=P_LA_src)
     return P_LK_tgt
 
 
 def _batched_gather(values: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
     """
-    values : (D, L, C)
-    idx    : (D, L, k)
-    returns: (D, L, k, C)
+    values : (B, L, C)
+    idx    : (B, L, k)
+    returns: (B, L, k, C)
     """
-    D, L, C = values.shape
+    B, L, C = values.shape
     k = idx.shape[-1]
 
-    #   (D, L, 1, C)  → stride-0 along k  → (D, L, k, C)
+    #   (B, L, 1, C)  → stride-0 along k  → (B, L, k, C)
     src = values.unsqueeze(2).expand(-1, -1, k, -1)
-    idx = idx.unsqueeze(-1).expand(-1, -1, -1, C)  # (D, L, k, C)
+    idx = idx.unsqueeze(-1).expand(-1, -1, -1, C)  # (B, L, k, C)
 
     return torch.gather(src, 1, idx)  # dim=1 is the L-axis
 
@@ -196,7 +196,7 @@ def create_attention_indices(
         X_L = torch.randn(
             (1, L, 3), device=device, dtype=torch.float
         )  # [L, 3] - random
-    D_LL = torch.cdist(X_L, X_L, p=2)  # [D, L, L] - pairwise atom distances
+    D_LL = torch.cdist(X_L, X_L, p=2)  # [B, L, L] - pairwise atom distances
 
     # Create attention indices using neighbour distances
     base_mask = ~f["unindexing_pair_mask"][
@@ -231,7 +231,7 @@ def create_attention_indices(
             k_max=k_actual,
             chain_id=chain_ids,
             base_mask=base_mask,
-        )  # [D, L, k] | indices[b, i, j] = atom index for atom i to j-th attn query
+        )  # [B, L, k] | indices[b, i, j] = atom index for atom i to j-th attn query
 
     return attn_indices
 
@@ -245,7 +245,7 @@ def get_sparse_attention_indices_with_inter_chain(
 
     Args:
         tok_idx: atom to token mapping
-        D_LL: pairwise distances [D, L, L]
+        D_LL: pairwise distances [B, L, L]
         n_seq_neighbours: number of sequence neighbors
         k_intra: number of intra-chain attention keys
         k_inter: number of inter-chain attention keys
@@ -253,29 +253,29 @@ def get_sparse_attention_indices_with_inter_chain(
         base_mask: base mask for valid pairs
 
     Returns:
-        attn_indices: [D, L, k_total] where k_total = k_intra + k_inter
+        attn_indices: [B, L, k_total] where k_total = k_intra + k_inter
     """
-    D, L, _ = D_LL.shape
+    B, L, _ = D_LL.shape
 
     # Get regular intra-chain indices (limited to k_intra)
     intra_indices = get_sparse_attention_indices(
         tok_idx, D_LL, n_seq_neighbours, k_intra, chain_id, base_mask
-    )  # [D, L, k_intra]
+    )  # [B, L, k_intra]
 
     # Get inter-chain indices for clash avoidance
-    inter_indices = torch.zeros(D, L, k_inter, dtype=torch.long, device=D_LL.device)
-
-    for d in range(D):
-        for l in range(L):
-            query_chain = chain_id[l]
+    inter_indices = torch.zeros(B, L, k_inter, dtype=torch.long, device=D_LL.device)
+    unique_chains = torch.unique(chain_id)
+    for b in range(B):
+        for c in unique_chains:
+            query_chain = chain_id[c]
 
             # Find atoms from different chains
-            other_chain_mask = (chain_id != query_chain) & base_mask[l, :]
+            other_chain_mask = (chain_id != query_chain) & base_mask[c, :]
             other_chain_atoms = torch.where(other_chain_mask)[0]
 
             if len(other_chain_atoms) > 0:
                 # Get distances to other chains
-                distances_to_other = D_LL[d, l, other_chain_atoms]
+                distances_to_other = D_LL[b, c, other_chain_atoms]
 
                 # Select k_inter closest atoms from other chains
                 n_select = min(k_inter, len(other_chain_atoms))
@@ -283,23 +283,23 @@ def get_sparse_attention_indices_with_inter_chain(
                 selected_atoms = other_chain_atoms[closest_idx]
 
                 # Fill inter-chain indices
-                inter_indices[d, l, :n_select] = selected_atoms
+                inter_indices[b, c, :n_select] = selected_atoms
                 # Pad with random atoms if needed
                 if n_select < k_inter:
                     padding = torch.randint(
                         0, L, (k_inter - n_select,), device=D_LL.device
                     )
-                    inter_indices[d, l, n_select:] = padding
+                    inter_indices[b, c, n_select:] = padding
             else:
                 # No other chains found, fill with random indices
-                inter_indices[d, l, :] = torch.randint(
+                inter_indices[b, c, :] = torch.randint(
                     0, L, (k_inter,), device=D_LL.device
                 )
 
     # Combine intra and inter chain indices
     combined_indices = torch.cat(
         [intra_indices, inter_indices], dim=-1
-    )  # [D, L, k_total]
+    )  # [B, L, k_total]
 
     return combined_indices