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

rfd3/inference/legacy_input_parsing.py

@@ -139,13 +139,18 @@ def fetch_motif_residue_(
         subarray, motif=True, unindexed=False, dtype=int
     )  # all values init to True (fix all)
 
+    to_unindex = f"{src_chain}{src_resid}" in unindexed_components
+    to_index = f"{src_chain}{src_resid}" in components
+
     # Assign is motif atom and sequence
     if exists(atoms := fixed_atoms.get(f"{src_chain}{src_resid}")):
+        # If specified, we set fixed atoms in the residue to be motif atoms
         atom_mask = get_name_mask(subarray.atom_name, atoms, res_name)
         subarray.set_annotation("is_motif_atom", atom_mask)
         # subarray.set_annotation("is_motif_atom_with_fixed_coord", atom_mask)  # BUGFIX: uncomment
 
     elif redesign_motif_sidechains and res_name in STANDARD_AA:
+        # If redesign_motif_sidechains is True, we only make the backbone atoms to be motif atoms
         n_atoms = subarray.shape[0]
         diffuse_oxygen = False
         if n_atoms < 3:
@@ -178,6 +183,18 @@ def fetch_motif_residue_(
         subarray.set_annotation(
             "is_motif_atom_with_fixed_seq", np.zeros(subarray.shape[0], dtype=int)
         )
+    elif to_index or to_unindex:
+        # If the residue is in the contig or unindexed components,
+        # we set all atoms in the residue to be motif atoms
+        subarray.set_annotation("is_motif_atom", np.ones(subarray.shape[0], dtype=int))
+    else:
+        if to_unindex and not (
+            unfix_all or f"{src_chain}{src_resid}" in unfix_residues
+        ):
+            raise ValueError(
+                f"{src_chain}{src_resid} is not found in fixed_atoms, contig or unindex contig."
+                "Please check your input and contig specification."
+            )
     if unfix_all or f"{src_chain}{src_resid}" in unfix_residues:
         subarray.set_annotation(
             "is_motif_atom_with_fixed_coord", np.zeros(subarray.shape[0], dtype=int)
@@ -197,7 +214,6 @@ def fetch_motif_residue_(
         subarray.set_annotation(
             "is_flexible_motif_atom", np.zeros(subarray.shape[0], dtype=bool)
         )
-    to_unindex = f"{src_chain}{src_resid}" in unindexed_components
     if to_unindex:
         subarray.set_annotation(
             "is_motif_atom_unindexed", subarray.is_motif_atom.copy()

rfd3/inference/parsing.py

@@ -117,6 +117,7 @@ def from_any_(v: Any, atom_array: AtomArray):
 
         # Split to atom names
         data_split[idx] = token.atom_name[comp_mask_subset].tolist()
+        # TODO: there is a bug where when you select specifc atoms within a ligand, output ligand is fragmented
 
         # Update mask & token dictionary
         mask[comp_mask] = comp_mask_subset

rfd3/inference/symmetry/atom_array.py

@@ -4,12 +4,8 @@ from rfd3.inference.symmetry.frames import (
     get_symmetry_frames_from_symmetry_id,
 )
 
-from foundry.utils.ddp import RankedLogger
-
 FIXED_TRANSFORM_ID = -1
 FIXED_ENTITY_ID = -1
-ranked_logger = RankedLogger(__name__, rank_zero_only=True)
-
 
 ########################################################
 # Symmetry annotations
@@ -28,7 +24,7 @@ def add_sym_annotations(atom_array, sym_conf):
     is_asu = np.full(n, True, dtype=np.bool_)
     atom_array.set_annotation("is_sym_asu", is_asu)
     # 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)
     return atom_array
 

rfd3/inference/symmetry/checks.py

@@ -1,10 +1,13 @@
 import numpy as np
-from rfd3.inference.symmetry.contigs import expand_contig_unsym_motif
+from rfd3.inference.symmetry.contigs import (
+    expand_contig_unsym_motif,
+    get_unsym_motif_mask,
+)
 from rfd3.transforms.conditioning_base import get_motif_features
 
 from foundry.utils.ddp import RankedLogger
 
-MIN_ATOMS_ALIGN = 100
+MIN_ATOMS_ALIGN = 30
 MAX_TRANSFORMS = 10
 RMSD_CUT = 1.0  # Angstroms
 
@@ -18,29 +21,33 @@ def check_symmetry_config(
     Check if the symmetry configuration is valid. Add all basic checks here.
     """
 
-    assert sym_conf.get("id"), "symmetry_id is required. e.g. {'id': 'C2'}"
+    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
-    if sym_conf.get("is_unsym_motif"):
+
+    is_unsym_motif = np.zeros(atom_array.shape[0], dtype=bool)
+    if sym_conf.is_unsym_motif:
         assert (
             src_atom_array is not None
         ), "Source atom array must be provided for symmetric motifs"
-        unsym_motif_names = sym_conf["is_unsym_motif"].split(",")
+        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(atom_array, unsym_motif_names)
         for n in unsym_motif_names:
             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 (
-        get_motif_features(atom_array)["is_motif_token"].any()
-        and not sym_conf.get("is_symmetric_motif")
+        is_motif_token[~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. "
+            "Asymmetric motif inputs should be distance constrained."
             "Use atomwise_fixed_dist to constrain the distance between the motif atoms."
         )
-    # else: if unconditional symmetry, no need to have symmetric input motif
 
-    if partial and not sym_conf.get("is_symmetric_motif"):
+    if partial and not sym_conf.is_symmetric_motif:
         raise ValueError(
             "Partial diffusion with symmetry is only supported for symmetric inputs."
         )
@@ -54,9 +61,6 @@ def check_atom_array_is_symmetric(atom_array):
     Returns:
         bool: True if the atom array is symmetric, False otherwise
     """
-    # TODO: Implement something like this https://github.com/baker-laboratory/ipd/blob/main/ipd/sym/sym_detect.py#L303
-    #       and maybe this https://github.com/baker-laboratory/ipd/blob/main/ipd/sym/sym_detect.py#L231
-
     import biotite.structure as struc
     from rfd3.inference.symmetry.atom_array import (
         apply_symmetry_to_atomarray_coord,
@@ -68,8 +72,10 @@ def check_atom_array_is_symmetric(atom_array):
     # remove hetero atoms
     atom_array = atom_array[~atom_array.hetero]
     if len(atom_array) == 0:
-        ranked_logger.info("Atom array has no protein chains. Please check your input.")
-        return False
+        ranked_logger.warning(
+            "Atom array has no protein chains. Please check your input."
+        )
+        return True
 
     chains = np.unique(atom_array.chain_id)
     asu_mask = atom_array.chain_id == chains[0]
@@ -162,16 +168,22 @@ def find_optimal_rotation(coords1, coords2, max_points=1000):
         return None
 
 
-def check_input_frames_match_symmetry_frames(computed_frames, original_frames) -> None:
+def check_input_frames_match_symmetry_frames(
+    computed_frames, original_frames, nids_by_entity
+) -> None:
     """
     Check if the atom array matches the symmetry_id.
     Arguments:
         computed_frames: list of computed frames
         original_frames: list of original frames
     """
-    assert len(computed_frames) == len(
-        original_frames
-    ), "Number of computed frames does not match number of original frames"
+    assert len(computed_frames) == len(original_frames), (
+        "Number of computed frames does not match number of original frames.\n"
+        f"Computed Frames: {len(computed_frames)}. Original Frames: {len(original_frames)}.\n"
+        "If the computed frames are not as expected, please check if you have one-to-one mapping "
+        "(size, sequence, folding) of an entity across all chains.\n"
+        f"Computed Entity Mapping: {nids_by_entity}."
+    )
 
 
 def check_valid_multiplicity(nids_by_entity) -> None:
@@ -184,25 +196,35 @@ def check_valid_multiplicity(nids_by_entity) -> None:
     multiplicity = min([len(i) for i in nids_by_entity.values()])
     if multiplicity == 1:  # no possible symmetry
         raise ValueError(
-            "Input has no possible symmetry. If asymmetric motif, please use 2D conditioning inference instead."
+            "Input has no possible symmetry. If asymmetric motif, please use 2D conditioning inference instead.\n"
+            "Multiplicity: 1"
         )
 
     # Check that the input is not asymmetric
     multiplicity_good = [len(i) % multiplicity == 0 for i in nids_by_entity.values()]
     if not all(multiplicity_good):
-        raise ValueError("Invalid multiplicities of subunits. Please check your input.")
+        raise ValueError(
+            "Expected multiplicity does not match for some entities.\n"
+            "Please modify your input to have one-to-one mapping (size, sequence, folding) of an entity across all chains.\n"
+            f"Expected Multiplicity: {multiplicity}.\n"
+            f"Computed Entity Mapping: {nids_by_entity}."
+        )
 
 
 def check_valid_subunit_size(nids_by_entity, pn_unit_id) -> None:
     """
     Check that the subunits in the input are of the same size.
     Arguments:
-        nids_by_entity: dict mapping entity to ids
+        nids_by_entity: dict mapping entity to ids. e.g. {0: (['A_1', 'B_1', 'C_1']), 1: (['A_2', 'B_2', 'C_2'])}
+        pn_unit_id: array of ids. e.g. ['A_1', 'B_1', 'C_1', 'A_2', 'B_2', 'C_2']
     """
-    for i, js in nids_by_entity.items():
-        for j in js[1:]:
-            if (pn_unit_id == js[0]).sum() != (pn_unit_id == j).sum():
-                raise ValueError("Size mismatch in the input. Please check your file.")
+    for js in nids_by_entity.values():
+        for js_i in js[1:]:
+            if (pn_unit_id == js[0]).sum() != (pn_unit_id == js_i).sum():
+                raise ValueError(
+                    f"Size mismatch between chain {js[0]} ({(pn_unit_id == js[0]).sum()} atoms) "
+                    f"and chain {js_i} ({(pn_unit_id == js_i).sum()} atoms). Please check your input file."
+                )
 
 
 def check_min_atoms_to_align(natm_per_unique, reference_entity) -> None:
@@ -212,7 +234,10 @@ def check_min_atoms_to_align(natm_per_unique, reference_entity) -> None:
         nids_by_entity: dict mapping entity to ids
     """
     if natm_per_unique[reference_entity] < MIN_ATOMS_ALIGN:
-        raise ValueError("Not enough atoms to align. Please check your input.")
+        raise ValueError(
+            f"Not enough atoms to align < {MIN_ATOMS_ALIGN} atoms."
+            f"Please provide a input with at least {MIN_ATOMS_ALIGN} atoms."
+        )
 
 
 def check_max_transforms(chains_to_consider) -> None:
@@ -224,7 +249,7 @@ def check_max_transforms(chains_to_consider) -> None:
     """
     if len(chains_to_consider) > MAX_TRANSFORMS:
         raise ValueError(
-            "Number of transforms exceeds the max number of transforms (10)"
+            f"Number of transforms exceeds the max number of transforms ({MAX_TRANSFORMS})."
         )
 
 

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:])
@@ -25,9 +26,9 @@ def get_symmetry_frames_from_symmetry_id(symmetry_id):
         frames = get_dihedral_frames(order)
     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 +121,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
 

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. e.g. 'C3', 'D2'. Only C and D symmetry types are supported currently.",
+    )
+    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,39 +76,33 @@ 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:
+    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."
         )
@@ -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/run_inference.py

@@ -12,7 +12,9 @@ load_dotenv(override=True)
 
 # For pip-installed package, configs should be relative to this file
 # Adjust this path based on where configs are bundled in the package
-_config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "configs")
+_config_path = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "configs"
+)
 
 
 @hydra.main(

rfd3/utils/inference.py

@@ -391,6 +391,29 @@ def ensure_input_is_abspath(args: dict, path: PathLike | None):
     return args
 
 
+def ensure_inference_sampler_matches_design_spec(
+    design_spec: dict, inference_sampler: dict | None = None
+):
+    """
+    Ensure the inference sampler is set to the correct sampler for the design specification.
+    Args:
+        design_spec: Design specification dictionary
+        inference_sampler: Inference sampler dictionary
+    """
+    has_symmetry_specification = [
+        True if "symmetry" in item.keys() else False for item in design_spec.values()
+    ]
+    if any(has_symmetry_specification):
+        if (
+            inference_sampler is None
+            or inference_sampler.get("kind", "default") != "symmetry"
+        ):
+            raise ValueError(
+                "You requested for symmetric designs, but inference sampler is not set to symmetry. "
+                "Please add inference_sampler.kind='symmetry' to your command."
+            )
+
+
 #################################################################################
 # Custom infer_ori functions
 #################################################################################