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

rfd3/engine.py

@@ -5,7 +5,7 @@ import time
 from dataclasses import dataclass, field
 from os import PathLike
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional
 
 import torch
 import yaml
@@ -48,9 +48,8 @@ class RFD3InferenceConfig:
     diffusion_batch_size: int = 16
 
     # RFD3 specific
-    skip_existing: bool = False
-    json_keys_subset: Optional[List[str]] = None
     skip_existing: bool = True
+    json_keys_subset: Optional[List[str]] = None
     specification: Optional[dict] = field(default_factory=dict)
     inference_sampler: SampleDiffusionConfig | dict = field(default_factory=dict)
 
@@ -216,6 +215,9 @@ class RFD3InferenceEngine(BaseInferenceEngine):
             inputs=inputs,
             n_batches=n_batches,
         )
+        if len(design_specifications) == 0:
+            ranked_logger.info("No design specifications to run. Skipping.")
+            return None
         ensure_inference_sampler_matches_design_spec(
             design_specifications, self.inference_sampler_overrides
         )
@@ -378,15 +380,21 @@ class RFD3InferenceEngine(BaseInferenceEngine):
 
     def _multiply_specifications(
         self, inputs: Dict[str, dict | DesignInputSpecification], n_batches=None
-    ) -> Dict[str, Dict[str, Any]]:
+    ) -> Dict[str, dict | DesignInputSpecification]:
         # Find existing example IDS in output directory
         if exists(self.out_dir):
-            existing_example_ids = set(
+            existing_example_ids_ = set(
                 extract_example_id_from_path(path, CIF_LIKE_EXTENSIONS)
                 for path in find_files_with_extension(self.out_dir, CIF_LIKE_EXTENSIONS)
             )
+            existing_example_ids = set(
+                [
+                    "_model_".join(eid.split("_model_")[:-1])
+                    for eid in existing_example_ids_
+                ]
+            )
             ranked_logger.info(
-                f"Found {len(existing_example_ids)} existing example IDs in the output directory."
+                f"Found {len(existing_example_ids)} existing example IDs in the output directory ({len(existing_example_ids_)} total)."
             )
 
         # Based on inputs, construct the specifications to loop through
@@ -405,7 +413,6 @@ class RFD3InferenceEngine(BaseInferenceEngine):
             for batch_id in range((n_batches) if exists(n_batches) else 1):
                 # ... Example ID
                 example_id = f"{prefix}_{batch_id}" if exists(n_batches) else prefix
-
                 if (
                     self.skip_existing
                     and exists(self.out_dir)

rfd3/inference/input_parsing.py

@@ -128,8 +128,10 @@ class DesignInputSpecification(BaseModel):
     # Motif selection from input file
     contig:  Optional[InputSelection] = Field(None, description="Contig specification string (e.g. 'A1-10,B1-5')")
     unindex: Optional[InputSelection] = Field(None, 
-        description="Unindexed components string (components must not overlap with contig). "\
-        "E.g. 'A15-20,B6-10' or dict. We recommend specifying")
+        description="Unindexed components selection. Components to fix in the generated structure without specifying sequence index. "\
+        "Components must not overlap with `contig` argument. "\
+        "E.g. 'A15-20,B6-10' or dict. We recommend specifying unindexed residues as a contig string, "\
+        "then using select_fixed_atoms will subset the atoms to the specified atoms")
     # Extra args:
     length:  Optional[str] = Field(None, description="Length range as 'min-max' or int. Constrains length of contig if provided")
     ligand:  Optional[str] = Field(None, description="Ligand name or index to include in design.")

rfd3/inference/symmetry/atom_array.py

@@ -1,5 +1,3 @@
-import string
-
 import numpy as np
 from rfd3.inference.symmetry.frames import (
     decompose_symmetry_frame,
@@ -9,68 +7,6 @@ from rfd3.inference.symmetry.frames import (
 FIXED_TRANSFORM_ID = -1
 FIXED_ENTITY_ID = -1
 
-# Alphabet for chain ID generation (uppercase letters only, per wwPDB convention)
-_CHAIN_ALPHABET = string.ascii_uppercase
-
-
-def index_to_chain_id(index: int) -> str:
-    """
-    Convert a zero-based index to a chain ID following wwPDB convention.
-
-    The naming follows the wwPDB-assigned chain ID system:
-    - 0-25: A-Z (single letter)
-    - 26-701: AA-ZZ (double letter)
-    - 702-18277: AAA-ZZZ (triple letter)
-    - And so on...
-
-    This is similar to Excel column naming (A, B, ..., Z, AA, AB, ...).
-
-    Arguments:
-        index: zero-based index (0 -> 'A', 25 -> 'Z', 26 -> 'AA', etc.)
-    Returns:
-        chain_id: string chain identifier
-    """
-    if index < 0:
-        raise ValueError(f"Chain index must be non-negative, got {index}")
-
-    result = ""
-    remaining = index
-
-    # Convert to bijective base-26 (like Excel columns)
-    while True:
-        result = _CHAIN_ALPHABET[remaining % 26] + result
-        remaining = remaining // 26 - 1
-        if remaining < 0:
-            break
-
-    return result
-
-
-def chain_id_to_index(chain_id: str) -> int:
-    """
-    Convert a chain ID back to a zero-based index.
-
-    Inverse of index_to_chain_id.
-
-    Arguments:
-        chain_id: string chain identifier (e.g., 'A', 'Z', 'AA', 'AB')
-    Returns:
-        index: zero-based index
-    """
-    if not chain_id or not all(c in _CHAIN_ALPHABET for c in chain_id):
-        raise ValueError(f"Invalid chain ID: {chain_id}")
-
-    # Offset for all shorter chain IDs (26 + 26^2 + ... + 26^(len-1))
-    offset = sum(26**k for k in range(1, len(chain_id)))
-
-    # Value within the current length group (standard base-26)
-    value = 0
-    for char in chain_id:
-        value = value * 26 + _CHAIN_ALPHABET.index(char)
-
-    return offset + value
-
-
 ########################################################
 # Symmetry annotations
 ########################################################
@@ -311,13 +247,11 @@ def reset_chain_ids(atom_array, start_id):
     Reset the chain ids and pn_unit_iids of an atom array to start from the given id.
     Arguments:
         atom_array: atom array with chain_ids and pn_unit_iids annotated
-        start_id: starting chain ID (e.g., 'A')
     """
     chain_ids = np.unique(atom_array.chain_id)
-    start_index = chain_id_to_index(start_id)
-    for i, old_id in enumerate(chain_ids):
-        new_id = index_to_chain_id(start_index + i)
-        atom_array.chain_id[atom_array.chain_id == old_id] = new_id
+    new_chain_range = range(ord(start_id), ord(start_id) + len(chain_ids))
+    for new_id, old_id in zip(new_chain_range, chain_ids):
+        atom_array.chain_id[atom_array.chain_id == old_id] = chr(new_id)
     atom_array.pn_unit_iid = atom_array.chain_id
     return atom_array
 
@@ -325,18 +259,15 @@ def reset_chain_ids(atom_array, start_id):
 def reannotate_chain_ids(atom_array, offset, multiplier=0):
     """
     Reannotate the chain ids and pn_unit_iids of an atom array.
-
-    Uses wwPDB-style chain IDs (A-Z, AA-ZZ, AAA-ZZZ, ...) to support
-    any number of chains.
-
     Arguments:
         atom_array: protein atom array with chain_ids and pn_unit_iids annotated
-        offset: offset to add to the chain ids (typically num_chains in ASU)
-        multiplier: multiplier for the offset (typically transform index)
+        offset: offset to add to the chain ids
+        multiplier: multiplier to add to the chain ids
     """
-    chain_ids_indices = np.array([chain_id_to_index(c) for c in atom_array.chain_id])
-    new_indices = chain_ids_indices + offset * multiplier
-    chain_ids = np.array([index_to_chain_id(idx) for idx in new_indices], dtype="U4")
+    chain_ids_int = (
+        np.array([ord(c) for c in atom_array.chain_id]) + offset * multiplier
+    )
+    chain_ids = np.array([chr(id) for id in chain_ids_int], dtype=str)
     atom_array.chain_id = chain_ids
     atom_array.pn_unit_iid = chain_ids
     return atom_array

rfd3/inference/symmetry/frames.py

@@ -24,12 +24,6 @@ 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
@@ -286,248 +280,6 @@ 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. Supported types: Cyclic (C), Dihedral (D), Tetrahedral (T), Octahedral (O), Icosahedral (I). e.g. 'C3', 'D2', 'T', 'O', 'I'.",
+        description="Symmetry group ID. e.g. 'C3', 'D2'. Only C and D symmetry types are supported currently.",
     )
     is_unsym_motif: Optional[str] = Field(
         None,
@@ -83,7 +83,7 @@ def make_symmetric_atom_array(
     if not isinstance(sym_conf, SymmetryConfig):
         sym_conf = convery_sym_conf_to_symmetry_config(sym_conf)
 
-    check_symmetry_config(
+    sym_conf = check_symmetry_config(
         asu_atom_array, sym_conf, sm, has_dist_cond, src_atom_array=src_atom_array
     )
     # Adding utility annotations to the asu atom array
@@ -99,7 +99,6 @@ def make_symmetric_atom_array(
         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.
         frames = get_symmetry_frames_from_atom_array(src_atom_array, frames)
     else:
         # At this point, asym case would have been caught by the check_symmetry_config function.

rfd3/model/inference_sampler.py

@@ -120,8 +120,10 @@ class SampleDiffusionWithMotif(SampleDiffusionConfig):
                 ranked_logger.info(
                     f"Using fallback: final step with t={noise_schedule[0].item():.6f}"
                 )
+        else:
+            noise_schedule = t_hat
 
-        return t_hat
+        return noise_schedule
 
     def _get_initial_structure(
         self,

rfd3/transforms/hbonds_hbplus.py

@@ -1,6 +1,7 @@
 import os
 import string
 import subprocess
+import tempfile
 from datetime import datetime
 from typing import Any, Tuple
 
@@ -66,10 +67,6 @@ def calculate_hbonds(
     cutoff_HA_dist: float = 3,
     cutoff_DA_distance: float = 3.5,
 ) -> Tuple[np.ndarray, np.ndarray, AtomArray]:
-    dtstr = datetime.now().strftime("%Y%m%d%H%M%S")
-    pdb_path = f"{dtstr}_{np.random.randint(10000)}.pdb"
-    atom_array, nan_mask, chain_map = save_atomarray_to_pdb(atom_array, pdb_path)
-
     hbplus_exe = os.environ.get("HBPLUS_PATH")
 
     if hbplus_exe is None or hbplus_exe == "":
@@ -78,49 +75,57 @@ def calculate_hbonds(
             "Please set it to the path of the hbplus executable in order to calculate hydrogen bonds."
         )
 
-    subprocess.call(
-        [
-            hbplus_exe,
-            "-h",
-            str(cutoff_HA_dist),
-            "-d",
-            str(cutoff_DA_distance),
-            pdb_path,
-            pdb_path,
-        ],
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-    )
-
-    HB = open(pdb_path.replace("pdb", "hb2"), "r").readlines()
-    hbonds = []
-    for i in range(8, len(HB)):
-        d_chain = HB[i][0]
-        d_resi = str(int(HB[i][1:5].strip()))
-        d_resn = HB[i][6:9].strip()
-        d_ins = HB[i][5].replace("-", " ")
-        d_atom = HB[i][9:13].strip()
-        a_chain = HB[i][14]
-        a_resi = str(int(HB[i][15:19].strip()))
-        a_ins = HB[i][19].replace("-", " ")
-        a_resn = HB[i][20:23].strip()
-        a_atom = HB[i][23:27].strip()
-        dist = float(HB[i][27:32].strip())
-
-        items = {
-            "d_chain": chain_map[d_chain],
-            "d_resi": d_resi,
-            "d_resn": d_resn,
-            "d_ins": d_ins,
-            "d_atom": d_atom,
-            "a_chain": chain_map[a_chain],
-            "a_resi": a_resi,
-            "a_resn": a_resn,
-            "a_ins": a_ins,
-            "a_atom": a_atom,
-            "dist": dist,
-        }
-        hbonds.append(items)
+    with tempfile.TemporaryDirectory() as tmpdir:
+        dtstr = datetime.now().strftime("%Y%m%d%H%M%S")
+        pdb_filename = f"{dtstr}_{np.random.randint(10000)}.pdb"
+        pdb_path = os.path.join(tmpdir, pdb_filename)
+        atom_array, _, chain_map = save_atomarray_to_pdb(atom_array, pdb_path)
+
+        subprocess.call(
+            [
+                hbplus_exe,
+                "-h",
+                str(cutoff_HA_dist),
+                "-d",
+                str(cutoff_DA_distance),
+                pdb_path,
+                pdb_path,
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+
+        hb2_path = pdb_path.replace(".pdb", ".hb2")
+        with open(hb2_path, "r") as hb_file:
+            HB = hb_file.readlines()
+        hbonds = []
+        for i in range(8, len(HB)):
+            d_chain = HB[i][0]
+            d_resi = str(int(HB[i][1:5].strip()))
+            d_resn = HB[i][6:9].strip()
+            d_ins = HB[i][5].replace("-", " ")
+            d_atom = HB[i][9:13].strip()
+            a_chain = HB[i][14]
+            a_resi = str(int(HB[i][15:19].strip()))
+            a_ins = HB[i][19].replace("-", " ")
+            a_resn = HB[i][20:23].strip()
+            a_atom = HB[i][23:27].strip()
+            dist = float(HB[i][27:32].strip())
+
+            items = {
+                "d_chain": chain_map[d_chain],
+                "d_resi": d_resi,
+                "d_resn": d_resn,
+                "d_ins": d_ins,
+                "d_atom": d_atom,
+                "a_chain": chain_map[a_chain],
+                "a_resi": a_resi,
+                "a_resn": a_resn,
+                "a_ins": a_ins,
+                "a_atom": a_atom,
+                "dist": dist,
+            }
+            hbonds.append(items)
 
     donor_array = np.zeros(len(atom_array))
     acceptor_array = np.zeros(len(atom_array))
@@ -162,8 +167,6 @@ def calculate_hbonds(
     donor_array[donor_mask] = 1
     acceptor_array[acceptor_mask] = 1
 
-    os.remove(pdb_path)
-    os.remove(pdb_path.replace("pdb", "hb2"))
     atom_array.set_annotation("active_donor", donor_array)
     atom_array.set_annotation("active_acceptor", acceptor_array)
 

rfd3/transforms/symmetry.py

@@ -60,22 +60,13 @@ class AddSymmetryFeats(Transform):
         )
         TIDs = torch.from_numpy(atom_array.get_annotation("sym_transform_id"))
 
-        # Get unique transforms by TID (more robust than unique_consecutive on each array)
-        unique_TIDs, inverse_indices = torch.unique(TIDs, return_inverse=True)
-
-        # Get the first occurrence of each unique TID
-        first_occurrence = torch.zeros(len(unique_TIDs), dtype=torch.long)
-        for i in range(len(TIDs)):
-            tid_idx = inverse_indices[i]
-            if first_occurrence[tid_idx] == 0 or i < first_occurrence[tid_idx]:
-                first_occurrence[tid_idx] = i
-
-        # Extract Ori, X, Y for each unique transform
-        Oris = Oris[first_occurrence]
-        Xs = Xs[first_occurrence]
-        Ys = Ys[first_occurrence]
-        TIDs = unique_TIDs
-
+        Oris = torch.unique_consecutive(Oris, dim=0)
+        Xs = torch.unique_consecutive(Xs, dim=0)
+        Ys = torch.unique_consecutive(Ys, dim=0)
+        TIDs = torch.unique_consecutive(TIDs, dim=0)
+        # the case in which there is only rotation (no translation), Ori = [0,0,0]
+        if len(Oris) == 1 and (Oris == 0).all():
+            Oris = Oris.repeat(len(Xs), 1)
         Rs, Ts = framecoords_to_RTs(Oris, Xs, Ys)
 
         for R, T, transform_id in zip(Rs, Ts, TIDs):

rfd3/utils/inference.py

@@ -373,12 +373,13 @@ def ensure_inference_sampler_matches_design_spec(
         design_spec: Design specification dictionary
         inference_sampler: Inference sampler dictionary
     """
-    has_symmetry_specification = [
-        True
-        if "symmetry" in item.keys() and item.get("symmetry") is not None
-        else False
-        for item in design_spec.values()
-    ]
+    has_symmetry_specification = []
+    for item in design_spec.values():
+        if hasattr(item, "symmetry"):
+            has_symmetry = item.symmetry is not None
+        else:
+            has_symmetry = "symmetry" in item and item.get("symmetry") is not None
+        has_symmetry_specification.append(has_symmetry)
     if any(has_symmetry_specification):
         if (
             inference_sampler is None