model/boltzgen/task/predict/data_from_generated.py

from dataclasses import dataclass
from pathlib import Path
import random
import re
import warnings
from typing import Dict, List, Optional
from collections import defaultdict
from rdkit.Chem import Mol
import pickle

import numpy as np
import pytorch_lightning as pl
import torch
from torch import Tensor
from torch.utils.data import DataLoader

from boltzgen.data import const
from boltzgen.data.data import Input, Structure, Tokenized
from boltzgen.data.feature.featurizer import Featurizer
from boltzgen.data.mol import load_canonicals, load_molecules
from boltzgen.data.pad import pad_to_max
from boltzgen.data.parse import mmcif
from boltzgen.data.parse.pdb_parser import parse_pdb
from boltzgen.data.template.features import (
    load_dummy_templates,
)
from boltzgen.data.parse.schema import parse_redesign_yaml
from boltzgen.data.tokenize.tokenizer import Tokenizer


class DataFetchException(Exception):
    pass


@dataclass
class DataConfig:
    """Data configuration."""

    num_targets: int
    samples_per_target: int
    moldir: str
    tokenizer: Tokenizer
    featurizer: Featurizer
    batch_size: int
    num_workers: int
    pin_memory: bool
    suffix: str = ".cif"
    suffix_native: str = "_native.cif"
    suffix_metadata: str = ".npz"
    target_id_regex: str = (
        r"^(?:(?:sample\d+_|batch\d+_|rank\d+_)+)?([^_]+)(?:_[^_]+)*?(?:_(?:gen))*$"
    )
    design: bool = False
    # Featurizer args (if design is True these should match with training config):
    backbone_only: bool = False
    atom14: bool = True
    max_seqs: int = 1
    inverse_fold: bool = False
    extra_mol_dir: Optional[str] = None
    disulfide_prob: float = 1.0
    disulfide_on: bool = False
    design_mask_override: Optional[str] = None
    multiplicity: int = 1
    return_designfolding: bool = False


def collate(data: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
    """Collate the data.

    Parameters
    ----------
    data : List[Dict[str, Tensor]]
        The data to collate.

    Returns
    -------
    Dict[str, Tensor]
        The collated data.

    """
    # Get the keys
    keys = data[0].keys()

    # Collate the data
    collated = {}
    for key in keys:
        values = [d[key] for d in data]

        if key not in [
            "metadata",
            "str_gen",
            "id",
            "path",
            "native_metadata",
            "native_str_gen",
            "native_id",
            "native_path",
            "exception",
            "native_exception",
            "skip",
            "native_skip",
            "structure_bonds",
            "native_structure_bonds",
            "extra_mols",
            "native_extra_mols",
            "structure",
            "tokenized",
            "data_sample_idx",
        ]:
            # Check if all have the same shape
            shape = values[0].shape
            if not all(v.shape == shape for v in values):
                values = pad_to_max(values, 0)
            else:
                values = torch.stack(values, dim=0)

        # Stack the values
        collated[key] = values

    return collated


@dataclass(frozen=True)
class TemplateInfo:
    """TemplateInfo datatype."""

    name: str
    query_chain: str
    query_st: int
    query_en: int
    template_chain: str
    template_st: int
    template_en: int


def template_from_tokens(
    tokenized: Tokenized,
    token_mask: np.ndarray[bool],
    tdim: int = 1,
) -> dict[str, torch.Tensor]:
    """Get template features where the tokens specified in token_mask have their structure specified."""
    # Get num token
    num_tokens = len(tokenized.tokens)

    # Allocate features
    res_type = np.zeros((tdim, num_tokens), dtype=np.int64)
    frame_rot = np.zeros((tdim, num_tokens, 3, 3), dtype=np.float32)
    frame_t = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
    cb_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
    ca_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
    frame_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
    cb_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
    template_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
    query_to_template = np.zeros((tdim, num_tokens), dtype=np.int64)
    visibility_ids = np.zeros((tdim, num_tokens), dtype=np.float32)

    # Now create features per token
    template_indices = np.where(token_mask)[0]
    for token_idx in template_indices:
        token = tokenized.tokens[token_idx]
        res_type[:, token_idx] = token["res_type"]
        frame_rot[:, token_idx] = token["frame_rot"].reshape(3, 3)
        frame_t[:, token_idx] = token["frame_t"]
        cb_coords[:, token_idx] = token["disto_coords"]
        ca_coords[:, token_idx] = token["center_coords"]
        cb_mask[:, token_idx] = token["disto_mask"]
        frame_mask[:, token_idx] = token["frame_mask"]
        template_mask[:, token_idx] = 1.0
        visibility_ids[:, token_idx] = 1

    # Convert to one-hot
    res_type = torch.from_numpy(res_type)
    res_type = torch.nn.functional.one_hot(res_type, num_classes=const.num_tokens)

    return {
        "template_restype": res_type,
        "template_frame_rot": torch.from_numpy(frame_rot),
        "template_frame_t": torch.from_numpy(frame_t),
        "template_cb": torch.from_numpy(cb_coords),
        "template_ca": torch.from_numpy(ca_coords),
        "template_mask_cb": torch.from_numpy(cb_mask),
        "template_mask_frame": torch.from_numpy(frame_mask),
        "template_mask": torch.from_numpy(template_mask),
        "query_to_template": torch.from_numpy(query_to_template),
        "visibility_ids": torch.from_numpy(visibility_ids),
    }


class FromGeneratedDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        generated_paths: List[Path],
        metadata_paths: List[Path],
        native_paths: List[Path],
        moldir: Path,
        canonicals: dict[str, Mol],
        tokenizer: Tokenizer,
        featurizer: Featurizer,
        return_native: bool = False,
        reference_metadata_dir: Optional[Path] = None,
        target_templates: bool = False,
        design_mask_templates: bool = False,
        compute_affinity: bool = False,
        design: bool = False,
        backbone_only: bool = False,
        atom14: bool = True,
        max_seqs: int = 1,
        inverse_fold: bool = False,
        extra_mol_dir: Optional[Path] = None,
        extra_features: Optional[List[str]] = None,
        disulfide_prob: float = 1.0,
        disulfide_on: bool = False,
        design_mask_override: Optional[str] = None,
        use_new_design_mask: bool = False,
        multiplicity: int = 1,
        return_designfolding=False,
    ) -> None:
        """
        Parameters
        ----------
        design : bool
            Set to True if this dataset is used to make predictions over (i.e. design some parts
            of the structure). Set to False if this dataset is used to only evaluate the predictions
            under the paths (i.e. no design is done).
        """
        super().__init__()
        self.tokenizer = tokenizer
        self.moldir = moldir
        self.canonicals = canonicals
        self.featurizer = featurizer
        self.metadata_paths = metadata_paths
        self.generated_paths = generated_paths
        self.native_paths = native_paths
        self.return_native = return_native
        self.reference_metadata_dir = reference_metadata_dir
        self.target_templates = target_templates
        self.design_mask_templates = design_mask_templates
        self.compute_affinity = compute_affinity
        self.design = design
        self.backbone_only = backbone_only
        self.atom14 = atom14
        self.max_seqs = max_seqs
        self.inverse_fold = inverse_fold
        self.extra_mol_dir = extra_mol_dir
        self.extra_features = (
            set(extra_features) if extra_features is not None else set()
        )
        self.disulfide_prob = disulfide_prob
        self.disulfide_on = disulfide_on
        self.design_mask_override = design_mask_override
        self.use_new_design_mask = use_new_design_mask
        self.multiplicity = multiplicity
        self.return_designfolding = return_designfolding

    def __getitem__(self, idx: int) -> Dict:
        """Get an item from the dataset.

        Returns
        -------
        Dict[str, Tensor]

        """
        data_sample_idx = idx // len(self.generated_paths)
        idx = idx % len(self.generated_paths)

        try:
            feat = self.getitem_from_paths(
                self.metadata_paths[idx],
                self.generated_paths[idx],
                self.native_paths[idx],
            )
            if self.multiplicity > 1:
                feat["data_sample_idx"] = data_sample_idx
            return feat
        except DataFetchException:
            idx = random.randint(0, len(self) - 1)
            feat = self.getitem_from_paths(
                self.metadata_paths[idx],
                self.generated_paths[idx],
                self.native_paths[idx],
            )
            if self.multiplicity > 1:
                feat["data_sample_idx"] = data_sample_idx
            return feat

    def get_sample(self, design_dir: Path, sample_id: Optional[str] = None) -> Dict:
        metadata_path = design_dir / f"{sample_id}.npz"
        generated_path = design_dir / f"{sample_id}.cif"
        native_path = design_dir / f"{sample_id}_native.cif"
        return self.getitem_from_paths(metadata_path, generated_path, native_path)

    def getitem_from_paths(self, metadata_path, generated_path, native_path) -> Dict:
        """Get an item from the dataset.

        Returns
        -------
        Dict[str, Tensor]

        """
        # Get metadata

        if self.reference_metadata_dir:
            reference_metadata_path = self.reference_metadata_dir / metadata_path.name
            metadata = np.load(reference_metadata_path)
        else:
            metadata = np.load(metadata_path)

        # get conditioning information from metadata
        metadata_design_mask = metadata["design_mask"]
        if self.use_new_design_mask:
            design_mask = metadata["inverse_fold_design_mask"].astype(np.float32)
        else:
            design_mask = metadata_design_mask

        ss_type = None
        if "ss_type" in metadata:
            ss_type = metadata["ss_type"]

        binding_type = None
        if "binding_type" in metadata:
            binding_type = metadata["binding_type"]

        # Per-residue amino acid constraints for inverse folding
        aa_constraint_mask = None
        if "aa_constraint_mask" in metadata:
            loaded_mask = metadata["aa_constraint_mask"]
            # Validate the loaded mask is a proper array with expected shape
            if (
                isinstance(loaded_mask, np.ndarray)
                and loaded_mask.ndim == 2
                and loaded_mask.shape[1] == 20  # 20 canonical amino acids
            ):
                aa_constraint_mask = loaded_mask
            else:
                warnings.warn(
                    f"Invalid aa_constraint_mask in NPZ: "
                    f"type={type(loaded_mask)}, shape={getattr(loaded_mask, 'shape', 'N/A')}. "
                    f"Expected ndarray with shape (N, 20). Ignoring constraints.",
                    RuntimeWarning,
                    stacklevel=2,
                )

        # Get features
        feat = self.get_feat(generated_path, design_mask, ss_type, binding_type, aa_constraint_mask)

        # Get native features
        if self.return_native:
            if "native_design_mask" in metadata.keys():
                feat_native = self.get_feat(native_path, metadata["native_design_mask"])
            else:
                feat_native = self.get_feat(native_path, metadata_design_mask)

            for k, v in feat_native.items():
                feat[f"native_{k}"] = v

        return feat

    def get_feat(self, path, design_mask, ss_type=None, binding_type=None, aa_constraint_mask=None):
        # Load design
        if self.extra_mol_dir is not None:
            mols = {
                path.stem: pickle.load(path.open("rb"))
                for path in self.extra_mol_dir.glob("*.pkl")
            }
            for mol_name, mol in mols.items():
                element_counts = defaultdict(int)
                for i, atom in enumerate(mol.GetAtoms()):
                    symbol = atom.GetSymbol()
                    element_counts[symbol] += 1
                    atom_name = f"{symbol}{element_counts[symbol]}"
                    atom.SetProp("name", atom_name)

        try:
            if path.suffix == ".cif":
                structure = mmcif.parse_mmcif(
                    path, mols, moldir=self.moldir, use_original_res_idx=False
                ).data
            elif path.suffix == ".pdb":
                structure = parse_pdb(
                    path, moldir=self.moldir, use_original_res_idx=False
                ).data
            else:
                raise ValueError(f"Invalid path:{path}")  # noqa: T201
        except Exception as e:  # noqa: BLE001
            print(f"Failed to parse {path} with error {e}. Skipping.")  # noqa: T201
            raise DataFetchException() from e

        # Tokenize structure

        try:
            tokenized = self.tokenizer.tokenize(
                structure, inverse_fold=self.inverse_fold
            )
        except Exception as e:  # noqa: BLE001
            print(f"Tokenizer failed on {path} with error {e}. Skipping.")  # noqa: T201
            raise DataFetchException() from e

        # Propagate design mask to obtain chain_design_mask (True whenever something is covalently bound to any residue that is in a chain that contains a design residue).
        chain_design_mask = design_mask.astype(bool)
        asym_id = tokenized.tokens["asym_id"]
        while True:
            design_chains = np.unique(asym_id[chain_design_mask])
            chain_propagated = np.isin(asym_id, design_chains)
            for i, j, _ in tokenized.bonds:
                if any([chain_propagated[i], chain_propagated[j]]):
                    chain_propagated[i] = True
                    chain_propagated[j] = True
            if np.equal(chain_propagated, chain_design_mask).all():
                break
            chain_design_mask = chain_propagated.astype(bool)

        # Extract design for refolding the design only
        if self.return_designfolding:
            residue_design_mask = np.zeros(tokenized.token_to_res.max() + 1, dtype=bool)
            np.put_along_axis(
                residue_design_mask, tokenized.token_to_res, chain_design_mask, axis=0
            )
            structure = Structure.extract_residues(structure, residue_design_mask)
            tokenized = self.tokenizer.tokenize(structure)
            design_mask = design_mask[chain_design_mask]
            chain_design_mask = chain_design_mask[chain_design_mask]

        # For inverse folding, condition even on structure selected for design
        if self.inverse_fold:
            tokenized.tokens["structure_group"] = 1

        try:
            # Try to find molecules in the dataset moldir if provided
            # Find missing ones in global moldir and check if all found
            molecules = {}
            molecules.update(self.canonicals)
            mol_names = set(tokenized.tokens["res_name"].tolist())
            mol_names = mol_names - set(self.canonicals.keys())
            if mols is not None:
                molecules.update(mols)
            mol_names = mol_names - set(molecules.keys())
            if self.moldir is not None:
                molecules.update(load_molecules(self.moldir, mol_names))
            molecules.update(load_molecules(self.moldir, mol_names))
        except Exception as e:  # noqa: BLE001
            print(f"Molecule loading failed for {path} with error {e}. Skipping.")
            raise DataFetchException() from e

        # Set design mask for tokens. This will impact the featurization and add the atom14 features
        if self.design:
            tokenized.tokens["design_mask"] = torch.from_numpy(design_mask).bool()

        # Finalize input data
        input_data = Input(
            tokens=tokenized.tokens,
            bonds=tokenized.bonds,
            token_to_res=tokenized.token_to_res,
            structure=structure,
            msa={},
            templates=None,
        )

        # Compute features
        try:
            features = self.featurizer.process(
                input_data,
                molecules=molecules,
                random=np.random.default_rng(None),
                training=False,
                max_seqs=self.max_seqs,
                backbone_only=self.backbone_only,
                atom14=self.atom14,
                design=True,
                compute_affinity=self.compute_affinity,
                override_method="X-RAY DIFFRACTION",
                disulfide_prob=self.disulfide_prob,
                disulfide_on=self.disulfide_on,
            )
        except Exception as e:  # noqa: BLE001
            print(f"Featurizer failed on {path} with error {e}. Skipping.")  # noqa: T201
            raise DataFetchException() from e

        # Set chain design mask
        features["chain_design_mask"] = torch.from_numpy(chain_design_mask)

        # Set conditioning variables that were set during design
        if ss_type is not None:
            features["ss_type"] = torch.from_numpy(ss_type).long()
        if binding_type is not None:
            features["binding_type"] = torch.from_numpy(binding_type).long()
        # Per-residue amino acid constraints for inverse folding
        if aa_constraint_mask is not None:
            features["aa_constraint_mask"] = torch.from_numpy(aa_constraint_mask).float()

        # If we do not want the design mask to impact the featurizer (e.g. represent atoms as atom14), we set the design mask only here.
        if not self.design:
            features["design_mask"] = torch.from_numpy(design_mask).bool()

        # set chain_design_mask
        # Override design mask for inverse folding if the part that should be inverse folded differs from the previously designed part.
        if self.design and self.design_mask_override is not None:
            msg = f"design mask being overridden with user input: {self.design_mask_override}"
            print(msg)
            new_design_mask = parse_redesign_yaml(
                Path(self.design_mask_override), tokenized
            )
            features["inverse_fold_design_mask"] = torch.from_numpy(
                new_design_mask
            ).bool()

        # Perform assertions
        if len(tokenized.tokens) != len(design_mask):
            print(
                f"WARNING: len(tokenized.tokens) [{len(tokenized.tokens)}] != len(design_mask) "
                f"[{len(design_mask)}] for {path}"
            )
            features["exception"] = True
            return features
        else:
            features["exception"] = False

        # Set templates
        if self.target_templates:
            if self.design_mask_templates:
                template_mask = ~features["design_mask"].numpy()
            else:
                template_mask = ~features["chain_design_mask"].numpy()
            templates_features = template_from_tokens(tokenized, template_mask)
        else:
            # Compute template features
            templates_features = load_dummy_templates(
                tdim=1, num_tokens=len(features["res_type"])
            )
        features.update(templates_features)

        features["affinity_token_mask"] = (
            features["mol_type"] == const.chain_type_ids["NONPOLYMER"]
        )

        # Set additional features
        features["str_gen"] = structure
        features["path"] = path
        features["id"] = path.stem
        if "structure" in self.extra_features:
            features["structure"] = structure
        if "tokenized" in self.extra_features:
            features["tokenized"] = tokenized

        return features

    def __len__(self) -> int:
        return len(self.generated_paths) * self.multiplicity


class FromGeneratedDataModule(pl.LightningDataModule):
    def __init__(
        self,
        cfg: DataConfig,
        return_native: bool = False,
        compute_affinity: bool = False,
        target_templates: bool = False,
        design_mask_templates: bool = False,
        skip_existing: bool = False,
        skip_existing_kind: str = None,
        legacy_gen_suffix: str = "_gen.cif",
        legacy_metadata_suffix: str = "_metadata.npz",
        reference_metadata_dir: Optional[Path] = None,
        design_dir: Optional[str] = None,
        extra_features: Optional[List[str]] = None,
        design_mask_override: Optional[str] = None,
        subset_target_ids: Optional[str] = None,
        skip_specific_ids: Optional[List[str]] = None,
        use_new_design_mask: bool = False,
        fail_if_no_designs: bool = False,
        output_dir: Optional[str] = None,
    ) -> None:
        super().__init__()
        self.cfg = cfg
        self.return_native = return_native
        self.skip_existing = skip_existing
        self.skip_existing_kind = skip_existing_kind
        self.reference_metadata_dir = (
            Path(reference_metadata_dir) if reference_metadata_dir else None
        )

        self.legacy_gen_suffix = legacy_gen_suffix
        self.legacy_metadata_suffix = legacy_metadata_suffix
        self.compute_affinity = compute_affinity
        self.target_templates = target_templates
        self.design_mask_templates = design_mask_templates
        self.extra_features = extra_features
        self.disulfide_prob = cfg.disulfide_prob
        self.disulfide_on = cfg.disulfide_on
        self.design_mask_override = cfg.design_mask_override
        self.collate = collate
        self.fail_if_no_designs = fail_if_no_designs
        self.subset_target_ids = subset_target_ids
        self.output_dir = Path(output_dir) if output_dir else None

        if design_dir is not None:
            self.init_dataset(
                design_dir,
                skip_specific_ids=skip_specific_ids,
                extra_features=extra_features,
                use_new_design_mask=use_new_design_mask,
            )
        else:
            # Load canonical molecules
            canonicals = load_canonicals(self.cfg.moldir)

            self.predict_set = FromGeneratedDataset(
                generated_paths=[],
                metadata_paths=[],
                native_paths=[],
                canonicals=canonicals,
                moldir=Path(self.cfg.moldir),
                tokenizer=self.cfg.tokenizer,
                featurizer=self.cfg.featurizer,
                return_native=self.return_native,
                reference_metadata_dir=self.reference_metadata_dir,
                target_templates=self.target_templates,
                design_mask_templates=self.design_mask_templates,
                compute_affinity=self.compute_affinity,
                design=self.cfg.design,
                backbone_only=self.cfg.backbone_only,
                atom14=self.cfg.atom14,
                max_seqs=self.cfg.max_seqs,
                inverse_fold=self.cfg.inverse_fold,
                extra_features=self.extra_features,
                disulfide_prob=self.disulfide_prob,
                disulfide_on=self.disulfide_on,
                design_mask_override=self.design_mask_override,
                use_new_design_mask=use_new_design_mask,
                multiplicity=self.cfg.multiplicity,
                return_designfolding=self.cfg.return_designfolding,
            )

    def init_dataset(
        self,
        design_dir,
        skip_specific_ids: Optional[List[str]] = None,
        extra_features: Optional[List[str]] = None,
        use_new_design_mask: bool = False,
    ):
        print(f"Initializing FromGeneratedDataModule datasets for {design_dir}")
        design_dir = Path(design_dir)
        assert design_dir.exists(), f"Path does not exist design_dir: {design_dir}"

        # Aggregate generated structure files (.cif or .pdb) while skipping companion native/metadata files.
        generated_paths = sorted(
            p
            for p in design_dir.iterdir()
            if p.suffix in {".cif", ".pdb"}
            and "_native.cif" not in p.name
            and "_metadata.npz" not in p.name
        )
        if self.fail_if_no_designs and len(generated_paths) == 0:
            raise ValueError(f"No designs found in {design_dir}")

        # skip certain ids
        num_files_before = len(generated_paths)
        print(
            f"[Info] Number of files to process (including already processed ones): {num_files_before}"
        )

        if skip_specific_ids:
            filtered_generated_paths = [
                p
                for p in generated_paths
                if not any(prob_id in p.name for prob_id in skip_specific_ids)
            ]
            num_files_after = len(filtered_generated_paths)
            print(f"[Info] Skipped specific IDs: {skip_specific_ids}")
            print(f"[Info] Number of files after filtering: {num_files_after}")
            generated_paths = filtered_generated_paths

        if self.skip_existing:
            # Functions to map an input path to a list of output paths.
            # If all output paths exist, the input path is skipped.
            def output_path_inverse_fold(input_path):
                assert self.output_dir is not None
                return [
                    self.output_dir / f"{input_path.stem}.cif",
                    self.output_dir / f"{input_path.stem}.npz",
                ]

            def output_path_folded(input_path):
                output_dir = (
                    design_dir / const.folding_dirname
                    if self.output_dir is None
                    else self.output_dir
                )
                return [
                    output_dir / f"{input_path.stem}.npz",
                    output_dir / f"{input_path.stem}.npz",
                ]

            def output_path_design_folded(input_path):
                output_dir = (
                    design_dir / const.refold_design_cif_dirname
                    if self.output_dir is None
                    else self.output_dir
                )
                return [
                    output_dir / f"{input_path.stem}.cif",
                ]

            def output_path_affinity(input_path):
                output_dir = (
                    design_dir / const.affinity_dirname
                    if self.output_dir is None
                    else self.output_dir
                )
                return [
                    output_dir / f"{input_path.stem}.npz",
                ]

            def output_path_analyzed(input_path):
                output_dir = (
                    design_dir / const.metrics_dirname
                    if self.output_dir is None
                    else self.output_dir
                )
                return [
                    output_dir / f"data_{input_path.stem}.npz",
                    output_dir / f"metrics_{input_path.stem}.npz",
                ]

            mappings = {
                "inverse_fold": output_path_inverse_fold,
                "folded": output_path_folded,
                "design_folded": output_path_design_folded,
                "affinity": output_path_affinity,
                "analyzed": output_path_analyzed,
            }
            if self.skip_existing_kind not in mappings:
                msg = f"Invalid skip_existing_kind: {self.skip_existing_kind}. Available kinds: {list(mappings.keys())}"
                raise ValueError(msg)
            selected_mapping = mappings[self.skip_existing_kind]

            generated_paths = [
                p
                for p in generated_paths
                if not all(output_path.exists() for output_path in selected_mapping(p))
            ]
            msg = f"[Info] Skipped already {self.skip_existing_kind} IDs. Number of files after filtering: {len(generated_paths)}"
            print(msg)

        target_ids = [
            re.search(rf"{self.cfg.target_id_regex}", p.stem).group(1)
            for p in generated_paths
        ]
        target_ids = list(set(target_ids))

        if self.cfg.num_targets is not None:
            target_ids = target_ids[: self.cfg.num_targets]
            generated_paths = [
                p
                for p in generated_paths
                if re.search(rf"{self.cfg.target_id_regex}", p.stem).group(1)
                in target_ids
            ]

        filtered_paths = []
        for target_id in target_ids:
            paths_of_target = [
                p
                for p in generated_paths
                if re.search(rf"{self.cfg.target_id_regex}", p.stem).group(1)
                == target_id
            ]

            filtered_paths.extend(paths_of_target[: self.cfg.samples_per_target])

        filtered_paths2 = []
        if self.subset_target_ids is not None:
            subset_ids = [
                l.strip() for l in open(self.subset_target_ids, "r").readlines()
            ]
            for path in filtered_paths:
                if any([sid in str(path) for sid in subset_ids]):
                    filtered_paths2.append(path)
            filtered_paths = filtered_paths2

        metadata_paths = []
        native_paths = []
        # Sort the paths to make sure each subprocess (when using multiple GPUs) has the same order and the index distribution when fetching from the dataset fetches the correct paths instead of fetching the same paths multiple times.
        filtered_paths = sorted(filtered_paths)
        for path in filtered_paths:
            ext = path.suffix

            # Legacy files contain "_gen" before the extension.
            if path.stem.endswith("_gen"):
                metadata_path = path.with_name(
                    path.name.replace(f"_gen{ext}", "_metadata.npz")
                )
                native_path = path.with_name(
                    path.name.replace(f"_gen{ext}", "_native.cif")
                )
            else:
                metadata_path = path.with_suffix(".npz")
                native_path = path.with_name(f"{path.stem}_native.cif")

            if not metadata_path.exists():
                print(f"[WARNING] Path does not exist: {metadata_path}")
            metadata_paths.append(metadata_path)
            if self.return_native:
                if not native_path.exists():
                    print(f"[WARNING] Path does not exist: {native_path}")
                native_paths.append(native_path)
            else:
                native_paths.append(None)
        msg = f"Found {len(target_ids)} targets and {len(filtered_paths)} remaining designs that still need to be processed in this step."
        print(msg)

        # Load canonical molecules
        canonicals = load_canonicals(self.cfg.moldir)

        self.predict_set = FromGeneratedDataset(
            generated_paths=filtered_paths,
            metadata_paths=metadata_paths,
            native_paths=native_paths,
            canonicals=canonicals,
            moldir=Path(self.cfg.moldir),
            tokenizer=self.cfg.tokenizer,
            featurizer=self.cfg.featurizer,
            return_native=self.return_native,
            reference_metadata_dir=self.reference_metadata_dir,
            target_templates=self.target_templates,
            design_mask_templates=self.design_mask_templates,
            compute_affinity=self.compute_affinity,
            design=self.cfg.design,
            backbone_only=self.cfg.backbone_only,
            atom14=self.cfg.atom14,
            max_seqs=self.cfg.max_seqs,
            inverse_fold=self.cfg.inverse_fold,
            extra_mol_dir=design_dir / const.molecules_dirname,
            extra_features=self.extra_features,
            disulfide_prob=self.disulfide_prob,
            disulfide_on=self.disulfide_on,
            design_mask_override=self.design_mask_override,
            use_new_design_mask=use_new_design_mask,
            multiplicity=self.cfg.multiplicity,
            return_designfolding=self.cfg.return_designfolding,
        )

    def predict_dataloader(self) -> DataLoader:
        return DataLoader(
            self.predict_set,
            batch_size=self.cfg.batch_size,
            num_workers=self.cfg.num_workers,
            pin_memory=self.cfg.pin_memory,
            shuffle=False,
            collate_fn=collate,
        )

    def transfer_batch_to_device(
        self,
        batch: Dict,
        device: torch.device,
        dataloader_idx: int = 0,
    ) -> Dict:
        for key in batch:
            if key not in [
                "metadata",
                "str_gen",
                "id",
                "path",
                "native_metadata",
                "native_str_gen",
                "native_id",
                "native_path",
                "exception",
                "native_exception",
                "skip",
                "native_skip",
                "structure_bonds",
                "native_structure_bonds",
                "extra_mols",
                "native_extra_mols",
                "structure",
                "tokenized",
                "data_sample_idx",
            ]:
                batch[key] = batch[key].to(device)

        return batch