rc-foundry 0.1.1__py3-none-any.whl

rf3/data/paired_msa.py

@@ -0,0 +1,206 @@
+import os
+import socket
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from atomworks.common import exists
+from atomworks.enums import ChainType
+from atomworks.ml.datasets import StructuralDatasetWrapper, logger
+from atomworks.ml.datasets.parsers import (
+    MetadataRowParser,
+    load_example_from_metadata_row,
+)
+from atomworks.ml.transforms._checks import (
+    check_contains_keys,
+    check_is_instance,
+    check_nonzero_length,
+)
+from atomworks.ml.transforms.base import Transform, TransformedDict
+from atomworks.ml.transforms.msa._msa_loading_utils import load_msa_data_from_path
+from atomworks.ml.utils.rng import capture_rng_states
+from biotite.structure import AtomArray, concatenate
+
+
+# input data wrapper that allows multiple input files separated by ':'
+#   data is loaded as concatentation of all inputs
+class MultiInputDatasetWrapper(StructuralDatasetWrapper):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def __getitem__(self, idx: int) -> Any:
+        # Capture example ID & current rng state (for reproducibility & debugging)
+        if hasattr(self, "idx_to_id"):
+            # ...if the dataset has a custom idx_to_id method, use it (e.g., for a PandasDataset)
+            example_id = self.idx_to_id(idx)
+        else:
+            # ...otherwise, fallback to a the `id_column` or a string representation of the index
+            example_id = (
+                self.dataset[idx][self.id_column] if self.id_column else f"row_{idx}"
+            )
+
+        # Get process id and hostname (for debugging)
+        logger.debug(
+            f"({socket.gethostname()}:{os.getpid()}) Processing example ID: {example_id}"
+        )
+
+        # Load the row, using the __getitem__ method of the dataset
+        row = self.dataset[idx]
+        pdb_path = row["pdb_path"].split(":")
+
+        # Process the row into a transform-ready dictionary with the given CIF and dataset parsers
+        # We require the "data" dictionary output from `load_example_from_metadata_row` to contain, at a minimum:
+        #   (a) An "id" key, which uniquely identifies the example within the dataframe; and,
+        #   (b) The "path" key, which is the path to the CIF file
+        _start_parse_time = time.time()
+        data = None
+        assert len(pdb_path) <= 2
+
+        for pdb_i in pdb_path:
+            row_i = {"example_id": row["example_id"], "path": pdb_i}
+            data_i = load_example_from_metadata_row(
+                row_i, self.dataset_parser, cif_parser_args=self.cif_parser_args
+            )
+
+            if data is None:
+                data = data_i
+            else:
+                data_i["atom_array"].pn_unit_id = np.full(
+                    len(data_i["atom_array"]), "B_1"
+                )  # unique pn unit id
+                data_i["atom_array"].pn_unit_iid = np.full(
+                    len(data_i["atom_array"]), "B_1"
+                )  # unique pn unit iid
+                data_i["atom_array"].chain_id = np.full(
+                    len(data_i["atom_array"]), "B"
+                )  # unique chain id
+                data_i["atom_array"].chain_iid = np.full(
+                    len(data_i["atom_array"]), "B"
+                )  # unique chain iid
+                data["atom_array"] = concatenate(
+                    [data["atom_array"], data_i["atom_array"]]
+                )
+                data["atom_array_stack"] = concatenate(
+                    [data["atom_array_stack"], data_i["atom_array_stack"]]
+                )
+                data["chain_info"]["B"] = data_i["chain_info"]["A"]
+
+        # 'example_id', 'path', 'assembly_id', 'query_pn_unit_iids',
+        data["path"] = row["pdb_path"]
+        data["msa_path"] = Path(row["msa_path"])  # save msa
+        _stop_parse_time = time.time()
+
+        # Manually add timing for cif-parsing
+        data = TransformedDict(data)
+        data.__transform_history__.append(
+            dict(
+                name="load_example_from_metadata_row",
+                instance=hex(id(load_example_from_metadata_row)),
+                start_time=_start_parse_time,
+                end_time=_stop_parse_time,
+                processing_time=_stop_parse_time - _start_parse_time,
+            )
+        )
+
+        # Apply the transformation pipeline to the data
+        if exists(self.transform):
+            try:
+                rng_state_dict = capture_rng_states(include_cuda=False)
+                data = self.transform(data)
+            except KeyboardInterrupt as e:
+                raise e
+            except Exception as e:
+                # Log the error and save the failed example to disk (optional)
+                logger.info(f"Error processing row {idx} ({example_id}): {e}")
+
+                if exists(self.save_failed_examples_to_dir):
+                    save_failed_example_to_disk(
+                        example_id=example_id,
+                        error_msg=e,
+                        rng_state_dict=rng_state_dict,
+                        data={},  # We do not save the data, since it may be large.
+                        fail_dir=self.save_failed_examples_to_dir,
+                    )
+                raise e
+
+        return data
+
+
+class MultidomainDFParser(MetadataRowParser):
+    """Parser for Qian's multidomain data"""
+
+    def __init__(
+        self,
+        example_id_colname: str = "example_id",
+        path_colname: str = "path",
+    ):
+        self.example_id_colname = example_id_colname
+        self.path_colname = path_colname
+
+    def _parse(self, row: dict) -> dict[str, Any]:
+        query_pn_unit_iids = None
+        assembly_id = "1"
+
+        return {
+            "example_id": row[self.example_id_colname],
+            "path": Path(row[self.path_colname]),
+            "assembly_id": assembly_id,
+            "query_pn_unit_iids": query_pn_unit_iids,
+            "extra_info": row,
+        }
+
+
+class LoadPairedMSAs(Transform):
+    """
+    LoadPairedMSAs adds paired MSAs from disk, overwriting previously paired MSA data.
+    """
+
+    def check_input(self, data: dict[str, Any]):
+        check_contains_keys(data, ["atom_array", "msa_path"])
+        check_is_instance(data, "atom_array", AtomArray)
+        check_nonzero_length(data, "atom_array")
+
+    def forward(self, data: dict[str, Any]) -> dict[str, Any]:
+        atom_array = data["atom_array"]
+        msa_file_path = data["msa_path"]
+        chain_type = data["chain_info"]["A"]["chain_type"]
+        max_msa_sequences = 10000
+
+        msa_data = load_msa_data_from_path(
+            msa_file_path=msa_file_path,
+            chain_type=chain_type,
+            max_msa_sequences=max_msa_sequences,
+        )
+
+        # split into chains
+        start_idx = 0
+        allpolymerchains = np.unique(
+            atom_array.chain_id[
+                np.isin(atom_array.chain_type, ChainType.get_polymers())
+            ]
+        )
+
+        data["polymer_msas_by_chain_id"] = {}  # nuke old version
+        for chain_id in allpolymerchains:
+            sequence = data["chain_info"][chain_id][
+                "processed_entity_non_canonical_sequence"
+            ]
+            stop_idx = start_idx + len(sequence)
+
+            data["polymer_msas_by_chain_id"][chain_id] = {}
+
+            # trim all msa info to this chain only
+            for mkey in msa_data.keys():
+                data["polymer_msas_by_chain_id"][chain_id][mkey] = msa_data[mkey][
+                    ..., start_idx:stop_idx
+                ]
+
+            # mock msa_is_padded_mask (all 0s)
+            data["polymer_msas_by_chain_id"][chain_id]["msa_is_padded_mask"] = np.zeros(
+                data["polymer_msas_by_chain_id"][chain_id]["msa"].shape, dtype=bool
+            )
+
+            start_idx = stop_idx
+
+        return data

rf3/data/pipeline_utils.py

@@ -0,0 +1,128 @@
+from functools import partial
+
+import torch
+from atomworks.enums import ChainType
+from atomworks.ml.transforms._checks import check_atom_array_annotation
+from atomworks.ml.transforms.crop import compute_local_hash
+from omegaconf import DictConfig
+from rf3.data.ground_truth_template import (
+    FeaturizeNoisedGroundTruthAsTemplateDistogram,
+    TokenGroupNoiseScaleSampler,
+    af3_noise_scale_distribution_wrapped,
+    af3_noise_scale_to_noise_level,
+)
+
+
+def annotate_pre_crop_hash(data: dict) -> dict:
+    hash_pre = compute_local_hash(data["atom_array"])
+    data["atom_array"].set_annotation("hash_pre", hash_pre)
+    return data
+
+
+def annotate_post_crop_hash(data: dict) -> dict:
+    hash_post = compute_local_hash(data["atom_array"])
+    data["atom_array"].set_annotation("hash_post", hash_post)
+    return data
+
+
+def set_to_occupancy_0_where_crop_hashes_differ(data: dict) -> dict:
+    check_atom_array_annotation(
+        data["atom_array"], ["hash_pre", "hash_post", "occupancy"]
+    )
+
+    # Create a mask of where hash_pre != hash_post
+    atom_array = data["atom_array"]
+    mask = atom_array.get_annotation("hash_pre") != atom_array.get_annotation(
+        "hash_post"
+    )
+
+    # Where the hashes differ, set occupancy to 0
+    atom_array.occupancy[mask] = 0
+
+    return data
+
+
+def build_ground_truth_distogram_transform(
+    *,
+    template_noise_scales: dict[str, float | None] | DictConfig,
+    allowed_chain_types_for_conditioning: list[ChainType] | None = None,
+    p_condition_per_token: float = 0.0,
+    p_provide_inter_molecule_distances: float = 0.0,
+    is_inference: bool = False,
+) -> FeaturizeNoisedGroundTruthAsTemplateDistogram:
+    """
+    Build a FeaturizeNoisedGroundTruthAsTemplateDistogram transform for either training or inference.
+
+    For inference, we must be deterministic, so we:
+        - Use constant noise scales (1.0)
+        - Always apply token-level conditioning
+
+    Args:
+        template_noise_scales (dict[str, float | None] | DictConfig):
+            Noise scales for 'atomized' and 'not_atomized' tokens. If is_inference=True, these are used as constants.
+            If is_inference=False, these are used as upper bounds for the noise scale distribution.
+        allowed_chain_types_for_conditioning (list[ChainType] | None):
+            List of allowed chain types for conditioning. None disables conditioning.
+        p_condition_per_token (float):
+            Probability of conditioning each eligible token.
+        p_provide_inter_molecule_distances (float):
+            Probability of providing inter-molecule (inter-chain) distances.
+        is_inference (bool):
+            If True, use constant noise scales for conditioning. If False, sample from provided distributions.
+
+    Returns:
+        FeaturizeNoisedGroundTruthAsTemplateDistogram: The configured transform.
+    """
+    mask_and_sampling_fns = []
+    if is_inference:
+        # Use constant noise scales for inference, rather than sampling (no stochasticity)
+        if template_noise_scales["atomized"] is not None:
+            mask_and_sampling_fns.append(
+                (
+                    lambda arr: arr.atomize,
+                    lambda size: torch.ones(size) * template_noise_scales["atomized"],
+                )
+            )
+        if template_noise_scales["not_atomized"] is not None:
+            mask_and_sampling_fns.append(
+                (
+                    lambda arr: ~arr.atomize,
+                    lambda size: torch.ones(size)
+                    * template_noise_scales["not_atomized"],
+                )
+            )
+    else:
+        # Use noise scale distributions for training
+        if template_noise_scales["atomized"] is not None:
+            mask_and_sampling_fns.append(
+                (
+                    lambda arr: arr.atomize,
+                    partial(
+                        af3_noise_scale_distribution_wrapped,
+                        upper_noise_level=af3_noise_scale_to_noise_level(
+                            template_noise_scales["atomized"]
+                        ).item(),
+                    ),
+                )
+            )
+        if template_noise_scales["not_atomized"] is not None:
+            mask_and_sampling_fns.append(
+                (
+                    lambda arr: ~arr.atomize,
+                    partial(
+                        af3_noise_scale_distribution_wrapped,
+                        upper_noise_level=af3_noise_scale_to_noise_level(
+                            template_noise_scales["not_atomized"]
+                        ).item(),
+                    ),
+                )
+            )
+
+    return FeaturizeNoisedGroundTruthAsTemplateDistogram(
+        noise_scale_distribution=TokenGroupNoiseScaleSampler(
+            mask_and_sampling_fns=tuple(mask_and_sampling_fns),
+        ),
+        allowed_chain_types=allowed_chain_types_for_conditioning,
+        p_condition_per_token=p_condition_per_token,
+        p_provide_inter_molecule_distances=p_provide_inter_molecule_distances,
+    )