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

rf3/configs/validate.yaml

@@ -2,6 +2,11 @@
 # ^ The "package" determines where the content of the config is placed in the output config
 # For more information about overriding configs, see: https://hydra.cc/docs/advanced/overriding_packages/#overriding-packages-using-the-defaults-list
 
+hydra:
+  searchpath:
+    - pkg://rf3.configs
+    - pkg://configs
+
 # NOTE: order of defaults determines the order in which configs override each other (higher up items are overridden by lower items)
 defaults:
   - callbacks: default

rf3/inference.py

@@ -16,7 +16,9 @@ rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
 
 load_dotenv(override=True)
 
-_config_path = os.path.join(os.environ["PROJECT_ROOT"], "models/rf3/configs")
+_config_path = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "configs"
+)
 
 
 @hydra.main(

rfd3/configs/dev.yaml

@@ -2,6 +2,7 @@
 # Inference engine config for development
 hydra:
   searchpath:
+    - pkg://rfd3.configs
     - pkg://configs
 
 defaults:

rfd3/configs/inference.yaml

@@ -2,6 +2,7 @@
 
 hydra:
   searchpath:
+    - pkg://rfd3.configs
     - pkg://configs
 
 defaults:

rfd3/configs/train.yaml

@@ -1,6 +1,7 @@
 # @package _global_
 hydra:
   searchpath:
+    - pkg://rfd3.configs
     - pkg://configs
 
 defaults:
@@ -25,4 +26,4 @@ ckpt_path: null
 
 # Placeholders
 name: aa_design
-tags: [aa_design]
+tags: [aa_design]

rfd3/configs/trainer/xpu.yaml

@@ -0,0 +1,6 @@
+strategy: xpu_single
+
+accelerator: xpu
+devices_per_node: 1
+num_nodes: 1
+

rfd3/configs/validate.yaml

@@ -2,6 +2,7 @@
 
 hydra:
   searchpath:
+    - pkg://rfd3.configs
     - pkg://configs
 
 defaults:

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
@@ -21,11 +21,13 @@ from rfd3.constants import SAVED_CONDITIONING_ANNOTATIONS
 from rfd3.inference.datasets import (
     assemble_distributed_inference_loader_from_json,
 )
-from rfd3.inference.input_parsing import DesignInputSpecification
+from rfd3.inference.input_parsing import (
+    DesignInputSpecification,
+    ensure_input_is_abspath,
+)
 from rfd3.model.inference_sampler import SampleDiffusionConfig
 from rfd3.utils.inference import (
     ensure_inference_sampler_matches_design_spec,
-    ensure_input_is_abspath,
 )
 from rfd3.utils.io import (
     CIF_LIKE_EXTENSIONS,
@@ -46,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)
 
@@ -214,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
         )
@@ -376,28 +380,39 @@ 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
         design_specifications = {}
         for prefix, example_spec in inputs.items():
             # Record task name in the specification
-            example_spec["extra"]["task_name"] = prefix
+            if isinstance(example_spec, DesignInputSpecification):
+                example_spec.extra = example_spec.extra or {}
+                example_spec.extra["task_name"] = prefix
+            else:
+                if "extra" not in example_spec:
+                    example_spec["extra"] = {}
+                example_spec["extra"]["task_name"] = prefix
 
             # ... Create n_batches for example
             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/datasets.py

@@ -14,8 +14,8 @@ from atomworks.ml.transforms.base import Compose, Transform
 from omegaconf import DictConfig, OmegaConf
 from rfd3.inference.input_parsing import (
     DesignInputSpecification,
+    ensure_input_is_abspath,
 )
-from rfd3.utils.inference import ensure_input_is_abspath
 from torch.utils.data import (
     DataLoader,
     SequentialSampler,

rfd3/inference/input_parsing.py

@@ -5,6 +5,7 @@ import os
 import time
 import warnings
 from contextlib import contextmanager
+from os import PathLike
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
@@ -127,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.")
@@ -1121,3 +1124,33 @@ def accumulate_components(
     if atom_array_accum.bonds is None:
         atom_array_accum.bonds = BondList(atom_array_accum.array_length())
     return atom_array_accum
+
+
+def ensure_input_is_abspath(args: Dict[str, Any], path: PathLike | None):
+    """
+    Ensures the input source is an absolute path if exists, if not it will convert
+
+    args:
+        args: Inference specification for atom array
+        path: None or file to which the input is relative to.
+    """
+    if isinstance(args, str):
+        raise ValueError(
+            "Expected args to be a dictionary, got a string: {}. If you are using an input JSON ensure it contains dictionaries of arguments".format(
+                args
+            )
+        )
+    if "input" not in args or not exists(args["input"]):
+        return args
+    input = str(args["input"])
+    if not os.path.isabs(input):
+        if path is None:
+            raise ValueError(
+                "Input path is relative, but no base path was provided to resolve it against."
+            )
+        input = os.path.abspath(os.path.join(os.path.dirname(str(path)), input))
+        logger.info(
+            f"Input source path is relative, converted to absolute path: {input}"
+        )
+        args["input"] = input
+    return args

rfd3/inference/symmetry/checks.py

@@ -24,7 +24,16 @@ def check_symmetry_config(
     assert sym_conf.id, "symmetry_id is required. e.g. {'id': 'C2'}"
     # if unsym motif is provided, check that each motif name is in the atom array
 
+    is_motif_atom = get_motif_features(atom_array)["is_motif_atom"]
     is_unsym_motif = np.zeros(atom_array.shape[0], dtype=bool)
+
+    if not is_motif_atom.any():
+        sym_conf.is_symmetric_motif = None
+        ranked_logger.warning(
+            "No motifs found in atom array. Setting is_symmetric_motif to None."
+        )
+        return sym_conf
+
     if sym_conf.is_unsym_motif:
         assert (
             src_atom_array is not None
@@ -36,21 +45,20 @@ def check_symmetry_config(
             if (sm and n not in sm.split(",")) and (n not in atom_array.src_component):
                 raise ValueError(f"Unsym motif {n} not found in atom_array")
 
-    is_motif_token = get_motif_features(atom_array)["is_motif_token"]
     if (
-        is_motif_token[~is_unsym_motif].any()
+        is_motif_atom[~is_unsym_motif].any()
         and not sym_conf.is_symmetric_motif
         and not has_dist_cond
     ):
         raise ValueError(
-            "Asymmetric motif inputs should be distance constrained."
-            "Use atomwise_fixed_dist to constrain the distance between the motif atoms."
+            "Asymmetric motif inputs are not supported yet. Please provide a symmetric motif."
         )
 
     if partial and not sym_conf.is_symmetric_motif:
         raise ValueError(
             "Partial diffusion with symmetry is only supported for symmetric inputs."
         )
+    return sym_conf
 
 
 def check_atom_array_is_symmetric(atom_array):

rfd3/inference/symmetry/symmetry_utils.py

@@ -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,12 +99,11 @@ 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)
-    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,14 +114,16 @@ 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(
                     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,
@@ -221,6 +227,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 +245,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 +454,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 +472,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