from collections import defaultdict from dataclasses import dataclass from pathlib import Path import traceback from typing import Dict, List, Optional import numpy as np import pandas as pd import pytorch_lightning as pl import torch from rdkit.Chem import Mol from torch import Tensor from torch.utils.data import DataLoader from boltzgen.data import const from boltzgen.data.crop.cropper import Cropper from boltzgen.data.select.selector import Selector from boltzgen.data.data import ( MSA, Input, Manifest, Record, Structure, ) from boltzgen.data.feature.featurizer import Featurizer from boltzgen.data.filter.dynamic.filter import DynamicFilter from boltzgen.data.mol import load_canonicals, load_molecules from boltzgen.data.pad import pad_to_max from boltzgen.data.sample.sampler import Sample, Sampler from boltzgen.data.template.features import load_dummy_templates from boltzgen.data.tokenize.tokenizer import Tokenizer from boltzgen.task.predict import data_ligands, data_protein_binder @dataclass class DatasetConfig: """Dataset configuration.""" target_dir: str msa_dir: str prob: Optional[float] sampler: Sampler cropper: Cropper selector: Optional[Selector] = None manifest_path: Optional[str] = None filters: Optional[list[DynamicFilter]] = None split: Optional[str] = None symmetry_correction: bool = True val_group: Optional[str] = "RCSB" use_train_subset: Optional[float] = None moldir: Optional[str] = None override_bfactor: Optional[bool] = False override_method: Optional[str] = None @dataclass class DataConfig: """Data configuration.""" datasets: List[DatasetConfig] featurizer: Featurizer tokenizer: Tokenizer selector: Selector max_atoms: int max_tokens: int max_seqs: int samples_per_epoch: int batch_size: int num_workers: int random_seed: int pin_memory: bool atoms_per_window_queries: int min_dist: float max_dist: float num_bins: int overfit: Optional[int] = None pad_to_max_tokens: bool = False pad_to_max_atoms: bool = False pad_to_max_seqs: bool = False return_train_symmetries: bool = False return_val_symmetries: bool = True val_batch_size: int = 1 single_sequence_prop_training: float = 0.0 msa_sampling_training: bool = False moldir: Optional[str] = None compute_frames: bool = True backbone_only: bool = False atom14: bool = False atom37: bool = False design: bool = False monomer_split: str = None monomer_target_dir: str = None monomer_seq_len: int = 100 monomer_target_structure_condition: bool = True inverse_fold: bool = False ligand_split: str = None ligand_target_dir: str = None ligand_seq_len: int = 100 use_msa: bool = True disulfide_prob: float = 1.0 disulfide_on: bool = False @dataclass class Dataset: """Data holder.""" samples: pd.DataFrame struct_dir: Path msa_dir: Path record_dir: Path prob: float cropper: Cropper tokenizer: Tokenizer featurizer: Featurizer val_group: str selector: Selector symmetry_correction: bool = True moldir: Optional[str] = None override_bfactor: Optional[bool] = False override_method: Optional[str] = None def load_record(record_id: str, record_dir: Path) -> Record: """Load the given record. Parameters ---------- record_id : str The record id to load. record_dir : Path The path to the record directory. Returns ------- Record The loaded record. """ return Record.load(record_dir / f"{record_id}.json") def load_structure(record: Record, struct_dir: Path) -> Structure: """Load the given input data. Parameters ---------- record : str The record to load. target_dir : Path The path to the data directory. Returns ------- Input The loaded input. """ if (struct_dir / f"{record.id}.npz").exists(): structure_path = struct_dir / f"{record.id}.npz" else: structure_path = struct_dir / f"{record.id}" / f"{record.id}_model_0.npz" return Structure.load(structure_path) def load_msas(chain_ids: set[int], record: Record, msa_dir: Path) -> Input: """Load the given input data. Parameters ---------- chain_ids : set[int] The chain ids to load. record : Record The record to load. msa_dir : Path The path to the MSA directory. Returns ------- Input The loaded input. """ msas = {} for chain in record.chains: if chain.chain_id not in chain_ids: continue msa_id = chain.msa_id if msa_id != -1: msa_path = msa_dir / f"{msa_id}.npz" msa = MSA.load(msa_path) msas[chain.chain_id] = msa return msas 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 [ "all_coords", "all_resolved_mask", "crop_to_all_atom_map", "chain_symmetries", "chain_swaps", "amino_acids_symmetries", "ligand_symmetries", "activity_name", "activity_qualifier", "sid", "cid", "normalized_protein_accession", "pair_id", "ligand_edge_index", "ligand_edge_lower_bounds", "ligand_edge_upper_bounds", "ligand_edge_bond_mask", "ligand_edge_angle_mask", "connections_edge_index", "ligand_chiral_atom_index", "ligand_chiral_check_mask", "ligand_chiral_atom_orientations", "ligand_stereo_bond_index", "ligand_stereo_check_mask", "ligand_stereo_bond_orientations", "ligand_aromatic_5_ring_index", "ligand_aromatic_6_ring_index", "ligand_planar_double_bond_index", "pdb_id", "id", "structure_bonds", "extra_mols", ]: if values[0] is not None: # 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 class TrainingDataset(torch.utils.data.Dataset): """Base iterable dataset.""" def __init__( self, datasets: List[Dataset], canonicals: dict[str, Mol], moldir: str, samples_per_epoch: int, max_atoms: int, max_tokens: int, max_seqs: int, pad_to_max_atoms: bool = False, pad_to_max_tokens: bool = False, pad_to_max_seqs: bool = False, atoms_per_window_queries: int = 32, min_dist: float = 2.0, max_dist: float = 22.0, num_bins: int = 64, overfit: Optional[int] = None, return_symmetries: Optional[bool] = False, single_sequence_prop: Optional[float] = 0.0, msa_sampling: bool = False, compute_frames: bool = True, backbone_only: bool = False, atom14: bool = False, atom37: bool = False, design: bool = False, disulfide_prob: float = 1.0, disulfide_on: bool = False, use_msa: bool = True, inverse_fold: bool = False, ) -> None: """Initialize the training dataset. Parameters ---------- datasets : List[Dataset] The datasets to sample from. samplers : List[Sampler] The samplers to sample from each dataset. probs : List[float] The probabilities to sample from each dataset. samples_per_epoch : int The number of samples per epoch. max_tokens : int The maximum number of tokens. """ super().__init__() self.datasets = datasets self.canonicals = canonicals self.moldir = moldir self.probs = [d.prob for d in datasets] self.samples_per_epoch = samples_per_epoch self.max_tokens = max_tokens self.max_seqs = max_seqs self.max_atoms = max_atoms self.pad_to_max_tokens = pad_to_max_tokens self.pad_to_max_atoms = pad_to_max_atoms self.pad_to_max_seqs = pad_to_max_seqs self.atoms_per_window_queries = atoms_per_window_queries self.min_dist = min_dist self.max_dist = max_dist self.num_bins = num_bins self.return_symmetries = return_symmetries self.backbone_only = backbone_only self.atom14 = atom14 self.atom37 = atom37 self.design = design self.disulfide_prob = disulfide_prob self.disulfide_on = disulfide_on self.single_sequence_prop = single_sequence_prop self.msa_sampling = msa_sampling self.use_msa = use_msa self.overfit = overfit self.compute_frames = compute_frames self.inverse_fold = inverse_fold self.samples: list[list[Dict]] = [] self.samples_weight: list[list[float]] = [] for d in self.datasets: if self.overfit: samples = d.samples[: self.overfit] else: samples = d.samples self.samples.append( [ samples.iloc[sample_idx].to_dict() for sample_idx in range(len(samples)) ] ) self.samples_weight.append(samples["weight"].tolist()) def __getitem__(self, idx: int) -> Dict[str, Tensor]: """Get an item from the dataset. Returns ------- Dict[str, Tensor] The sampled data features. """ # Set a random state random = np.random.default_rng() # Pick a random dataset dataset_idx = random.choice(len(self.datasets), p=self.probs) dataset = self.datasets[dataset_idx] # Get a sample from the dataset samples = self.samples[dataset_idx] sample_idx = random.choice( len(samples), p=( self.samples_weight[dataset_idx] / np.sum(self.samples_weight[dataset_idx]) if self.overfit else self.samples_weight[dataset_idx] ), ) sample = samples[sample_idx] sample: Sample = Sample( record_id=str(sample["record_id"]), chain_id=( int(sample["chain_id"]) if sample["chain_id"] is not None else None ), interface_id=( int(sample["interface_id"]) if sample["interface_id"] is not None else None ), weight=float(sample["weight"]), ) # Load record record = load_record(sample.record_id, dataset.record_dir) # Get the structure try: structure = load_structure(record, dataset.struct_dir) except Exception as e: # noqa: BLE001 print(f"Failed to load input for {record.id} with error {e}. Skipping.") return self.__getitem__(random.integers(0, len(self))) # Tokenize structure try: tokenized = dataset.tokenizer.tokenize( structure, inverse_fold=self.inverse_fold ) except Exception as e: # noqa: BLE001 print(f"Tokenizer failed on {record.id} with error {e}. Skipping.") traceback.print_exc() # noqa: T201 return self.__getitem__(random.integers(0, len(self))) # Compute crop try: if self.max_tokens is not None and len(tokenized.tokens) > self.max_tokens: tokenized = dataset.cropper.crop( tokenized, max_atoms=self.max_atoms, max_tokens=self.max_tokens, chain_id=sample.chain_id, interface_id=sample.interface_id, random=random, prefer_protein_queries=self.inverse_fold, ) if len(tokenized.tokens) == 0: msg = "No tokens in cropped structure." raise ValueError(msg) # noqa: TRY301 except Exception as e: # noqa: BLE001 print(f"Cropper failed on {record.id} with error {e}. Skipping.") traceback.print_exc() # noqa: T201 return self.__getitem__(random.integers(0, len(self))) # Select which tokens to design try: tokenized, design_task = dataset.selector.select( tokenized, random=random, ) except Exception as e: # noqa: BLE001 print(f"Selector failed on {record.id} with error {e}. Skipping.") # noqa: T201 traceback.print_exc() # noqa: T201 return self.__getitem__(random.integers(0, len(self))) structure = tokenized.structure # Get unique chain ids chain_ids = set(tokenized.tokens["asym_id"]) # Load msas and templates try: if self.use_msa: msas = load_msas( chain_ids=chain_ids, record=record, msa_dir=dataset.msa_dir, ) else: msas = {} except Exception as e: # noqa: BLE001 print(f"MSA loading failed for {record.id} with error {e}. Skipping.") return self.__getitem__(random.integers(0, len(self))) # Load molecules 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 dataset.moldir is not None: molecules.update(load_molecules(dataset.moldir, mol_names)) mol_names = mol_names - set(molecules.keys()) molecules.update(load_molecules(self.moldir, mol_names)) except Exception as e: # noqa: BLE001 print(f"Molecule loading failed for {record.id} with error {e}. Skipping.") return self.__getitem__(random.integers(0, len(self))) # Finalize input data input_data = Input( tokens=tokenized.tokens, bonds=tokenized.bonds, token_to_res=tokenized.token_to_res, structure=tokenized.structure, msa=msas, templates=None, record=record, ) # Compute features try: features: dict = dataset.featurizer.process( input_data, molecules=molecules, random=random, training=True, max_atoms=self.max_atoms if self.pad_to_max_atoms else None, max_tokens=self.max_tokens if self.pad_to_max_tokens else None, max_seqs=self.max_seqs, pad_to_max_seqs=self.pad_to_max_seqs, atoms_per_window_queries=self.atoms_per_window_queries, min_dist=self.min_dist, max_dist=self.max_dist, num_bins=self.num_bins, compute_symmetries=self.return_symmetries, single_sequence_prop=self.single_sequence_prop, msa_sampling=self.msa_sampling, override_bfactor=dataset.override_bfactor, override_method=dataset.override_method, compute_frames=self.compute_frames, backbone_only=self.backbone_only, atom14=self.atom14, atom37=self.atom37, design=self.design, disulfide_prob=self.disulfide_prob, inverse_fold=self.inverse_fold, ) except Exception as e: # noqa: BLE001 print(f"Featurizer failed on {record.id} with error {e}. Skipping.") traceback.print_exc() return self.__getitem__(random.integers(0, len(self))) # Check that there is enough stuff to design in the inverse folding case so we have no nan losses if self.inverse_fold and features["design_mask"].sum() < 3: print(f"Skipping {record.id}. Fewer than 3 design residues.") return self.__getitem__(random.integers(0, len(self))) # Set template features template_features = load_dummy_templates( tdim=1, num_tokens=len(features["res_type"]) ) features.update(template_features) features.update({"id": sample.record_id}) features["pdb_id"] = record.id # Assert that all design tokens make sense bad_protein_mask = ( (~features["is_standard"].bool()) & features["design_mask"].bool() & (features["mol_type"] == const.chain_type_ids["PROTEIN"]) ) assert not bad_protein_mask.any() return features def __len__(self) -> int: """Get the length of the dataset. Returns ------- int The length of the dataset. """ return self.samples_per_epoch class ValidationDataset(torch.utils.data.Dataset): """Base iterable dataset.""" def __init__( self, datasets: List[Dataset], canonicals: dict[str, Mol], moldir: str, seed: int, max_atoms: Optional[int] = None, max_tokens: Optional[int] = None, max_seqs: Optional[int] = None, pad_to_max_atoms: bool = False, pad_to_max_tokens: bool = False, pad_to_max_seqs: bool = False, atoms_per_window_queries: int = 32, min_dist: float = 2.0, max_dist: float = 22.0, num_bins: int = 64, overfit: Optional[int] = None, return_symmetries: Optional[bool] = False, compute_frames: bool = True, backbone_only: bool = False, atom14: bool = False, atom37: bool = False, design: bool = False, inverse_fold: bool = False, disulfide_prob: float = 1.0, disulfide_on: bool = False, ) -> None: """Initialize the training dataset. Parameters ---------- datasets : List[Dataset] The datasets to sample from. seed : int The random seed. max_tokens : int The maximum number of tokens. overfit : bool Whether to overfit the dataset """ super().__init__() self.datasets = datasets self.canonicals = canonicals self.moldir = moldir self.max_atoms = max_atoms self.max_tokens = max_tokens self.max_seqs = max_seqs self.seed = seed self.pad_to_max_tokens = pad_to_max_tokens self.pad_to_max_atoms = pad_to_max_atoms self.pad_to_max_seqs = pad_to_max_seqs self.overfit = overfit self.atoms_per_window_queries = atoms_per_window_queries self.min_dist = min_dist self.max_dist = max_dist self.num_bins = num_bins self.return_symmetries = return_symmetries self.compute_frames = compute_frames self.backbone_only = backbone_only self.atom14 = atom14 self.atom37 = atom37 self.design = design self.inverse_fold = inverse_fold self.disulfide_prob = disulfide_prob self.disulfide_on = disulfide_on def __getitem__(self, idx: int) -> Structure: """Get an item from the dataset. Returns ------- Dict[str, Tensor] The sampled data features. """ # Set random state seed = self.seed if self.overfit is None else None random = np.random.default_rng(seed) # Pick dataset based on idx for idx_dataset, dataset in enumerate(self.datasets): # noqa: B007 size = len(dataset.samples) if self.overfit is not None: size = min(size, self.overfit) if idx < size: break idx -= size # Get a sample from the dataset sample = Sample(**dataset.samples.iloc[idx].to_dict()) record = load_record(sample.record_id, dataset.record_dir) # Get the structure try: structure = load_structure(record, dataset.struct_dir) except Exception as e: # noqa: BLE001 print(f"Failed to load input for {record.id} with error {e}. Skipping.") return self.__getitem__(0) # Tokenize structure try: tokenized = dataset.tokenizer.tokenize(structure) except Exception as e: # noqa: BLE001 print(f"Tokenizer failed on {record.id} with error {e}. Skipping.") # noqa: T201 return self.__getitem__(0) # Compute crop try: if self.max_tokens is not None: tokenized = dataset.cropper.crop( tokenized, max_atoms=self.max_atoms, max_tokens=self.max_tokens, chain_id=sample.chain_id, interface_id=sample.interface_id, random=random, prefer_protein_queries=self.inverse_fold, ) if len(tokenized.tokens) == 0: msg = "No tokens in cropped structure." raise ValueError(msg) # noqa: TRY301 except Exception as e: # noqa: BLE001 print(f"Cropper failed on {record.id} with error {e}. Skipping.") return self.__getitem__(0) # Get unique chains chain_ids = set(np.unique(tokenized.tokens["asym_id"]).tolist()) # Load msas and templates try: msas = load_msas(chain_ids, record, dataset.msa_dir) except Exception as e: # noqa: BLE001 print(f"MSA loading failed for {record.id} with error {e}. Skipping.") return self.__getitem__(0) # Select which tokens to design try: tokenized, design_task = dataset.selector.select( tokenized, random=random, ) except Exception as e: # noqa: BLE001 print(f"Selector failed on {sample.record_id} with error {e}. Skipping.") # noqa: T201 traceback.print_exc() # noqa: T201 return self.__getitem__(0) structure = tokenized.structure 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 dataset.moldir is not None: molecules.update(load_molecules(dataset.moldir, mol_names)) mol_names = mol_names - set(molecules.keys()) molecules.update(load_molecules(self.moldir, mol_names)) except Exception as e: # noqa: BLE001 print(f"Molecule loading failed for {record.id} with error {e}. Skipping.") return self.__getitem__(0) # Finalize input data input_data = Input( tokens=tokenized.tokens, bonds=tokenized.bonds, token_to_res=tokenized.token_to_res, structure=tokenized.structure, msa=msas, templates=None, record=record, ) # Compute features try: features: dict = dataset.featurizer.process( input_data, molecules=molecules, random=random, training=False, max_atoms=None, max_tokens=None, max_seqs=self.max_seqs, pad_to_max_seqs=self.pad_to_max_seqs, atoms_per_window_queries=self.atoms_per_window_queries, min_dist=self.min_dist, max_dist=self.max_dist, num_bins=self.num_bins, compute_symmetries=self.return_symmetries, single_sequence_prop=0.0, override_method=dataset.override_method, compute_frames=self.compute_frames, backbone_only=self.backbone_only, atom14=self.atom14, atom37=self.atom37, design=self.design, inverse_fold=self.inverse_fold, disulfide_prob=self.disulfide_prob, disulfide_on=self.disulfide_on, ) except Exception as e: # noqa: BLE001 print(f"Featurizer failed on {record.id} with error {e}. Skipping.") return self.__getitem__(0) # Check that there is enough stuff to design in the inverse folding case so we have no nan losses if self.inverse_fold and features["design_mask"].sum() < 3: print(f"Skipping {record.id}. Fewer than 3 design residues.") return self.__getitem__(0) # Set template features template_features = load_dummy_templates( tdim=1, num_tokens=len(features["res_type"]) ) features.update(template_features) # Add dataset idx idx_dataset = torch.tensor([idx_dataset], dtype=torch.long) features.update({"idx_dataset": idx_dataset}) features.update({"id": record.id}) bad_protein_mask = ( (~features["is_standard"].bool()) & features["design_mask"].bool() & (features["mol_type"] == const.chain_type_ids["PROTEIN"]) ) assert not bad_protein_mask.any() return features def __len__(self) -> int: """Get the length of the dataset. Returns ------- int The length of the dataaset. """ if self.overfit is not None: length = sum(len(d.samples[: self.overfit]) for d in self.datasets) else: length = sum(len(d.samples) for d in self.datasets) return length class TrainingDataModule(pl.LightningDataModule): """DataModule for BoltzGen training.""" def __init__( self, cfg: DataConfig, ) -> None: """Initialize the DataModule. Parameters ---------- config : DataConfig The data configuration. """ super().__init__() self.cfg = cfg self.inverse_fold = cfg.inverse_fold assert self.cfg.val_batch_size == 1, "Validation only works with batch size=1." # Load datasets train: List[Dataset] = [] val: List[Dataset] = [] for data_config in cfg.datasets: # Get relevant directories if data_config.manifest_path is not None: manifest_path = Path(data_config.manifest_path) else: manifest_path = Path(data_config.target_dir) / "manifest.json" struct_dir = Path(data_config.target_dir) / "structures" record_dir = Path(data_config.target_dir) / "records" msa_dir = Path(data_config.msa_dir) # Get moldir, if any moldir = data_config.moldir moldir = Path(moldir) if moldir is not None else None # Load all records manifest: Manifest = Manifest.load(manifest_path) # Split records if givens if data_config.split is not None: with Path(data_config.split).open("r") as f: split = {x.lower() for x in f.read().splitlines()} train_records = [] val_records = [] for record in manifest.records: if record.id.lower() in split: val_records.append(record) else: train_records.append(record) else: train_records = manifest.records if cfg.overfit is None: val_records = [] else: print("Warning: modified overfit val behavior.") val_records = manifest.records[: cfg.overfit] print("train_records before filter", len(train_records)) # Apply dataset-specific filters if data_config.filters is not None: train_records = [ record for record in train_records if all(f.filter(record) for f in data_config.filters) ] # Train with subset of data if data_config.use_train_subset is not None: # Shuffle train_records list assert 0 < data_config.use_train_subset < 1.0 rng = np.random.default_rng(cfg.random_seed) rng.shuffle(train_records) train_records = train_records[ 0 : int(len(train_records) * data_config.use_train_subset) ] print("train_records after filter", len(train_records)) print("val_records after filter", len(val_records)) # Get samples train_samples: list[Sample] = data_config.sampler.sample(train_records) val_samples: list[Sample] = [Sample(r.id) for r in val_records] # Convert samples to pandas dataframe to avoid copy-on-write behavior train_samples = pd.DataFrame( [ ( r.record_id, r.chain_id, r.interface_id, r.weight, ) for r in train_samples ], columns=["record_id", "chain_id", "interface_id", "weight"], ) val_samples = pd.DataFrame( [s.record_id for s in val_samples], columns=["record_id"] ) # Use appropriate string type train_samples = train_samples.replace({np.nan: None}) val_samples = val_samples.replace({np.nan: None}) train_samples["record_id"] = train_samples["record_id"].astype("string") val_samples["record_id"] = val_samples["record_id"].astype("string") del manifest, train_records, val_records # Create train dataset if data_config.prob > 0: train.append( Dataset( samples=train_samples, record_dir=record_dir, struct_dir=struct_dir, msa_dir=msa_dir, moldir=moldir, prob=data_config.prob, cropper=data_config.cropper, tokenizer=cfg.tokenizer, featurizer=cfg.featurizer, val_group=data_config.val_group, symmetry_correction=data_config.symmetry_correction, override_bfactor=data_config.override_bfactor, override_method=data_config.override_method, selector=cfg.selector, ) ) # Create validation dataset if len(val_samples) > 0: val.append( Dataset( samples=val_samples, record_dir=record_dir, struct_dir=struct_dir, msa_dir=msa_dir, moldir=moldir, prob=data_config.prob, cropper=data_config.cropper, tokenizer=cfg.tokenizer, featurizer=cfg.featurizer, val_group=data_config.val_group, symmetry_correction=data_config.symmetry_correction, selector=cfg.selector, ) ) # Print dataset sizes for dataset in train: dataset: Dataset print(f"Training dataset size: {len(dataset.samples)}") self.val_group_mapper = defaultdict(dict) for i, dataset in enumerate(train if cfg.overfit is not None else val): dataset: Dataset print(f"Validation dataset size: {len(dataset.samples)}") self.val_group_mapper[i]["label"] = dataset.val_group self.val_group_mapper[i]["symmetry_correction"] = ( # If overfit, use symmetry_correction from val dataset instead of training dataset dataset.symmetry_correction if cfg.overfit is None else data_config.symmetry_correction ) # Load canonical molecules canonicals = load_canonicals(cfg.moldir) # Create wrapper datasets self._train_set = TrainingDataset( datasets=train, canonicals=canonicals, moldir=cfg.moldir, samples_per_epoch=cfg.samples_per_epoch, max_atoms=cfg.max_atoms, max_tokens=cfg.max_tokens, max_seqs=cfg.max_seqs, pad_to_max_atoms=cfg.pad_to_max_atoms, pad_to_max_tokens=cfg.pad_to_max_tokens, pad_to_max_seqs=cfg.pad_to_max_seqs, atoms_per_window_queries=cfg.atoms_per_window_queries, min_dist=cfg.min_dist, max_dist=cfg.max_dist, num_bins=cfg.num_bins, overfit=cfg.overfit, return_symmetries=cfg.return_train_symmetries, single_sequence_prop=cfg.single_sequence_prop_training, msa_sampling=cfg.msa_sampling_training, use_msa=cfg.use_msa, compute_frames=cfg.compute_frames, backbone_only=cfg.backbone_only, atom14=cfg.atom14, atom37=cfg.atom37, design=cfg.design, inverse_fold=cfg.inverse_fold, disulfide_prob=cfg.disulfide_prob, disulfide_on=cfg.disulfide_on, ) self._val_set = ValidationDataset( datasets=train if cfg.overfit is not None else val, canonicals=canonicals, moldir=cfg.moldir, seed=cfg.random_seed, max_atoms=cfg.max_atoms, max_tokens=cfg.max_tokens, max_seqs=cfg.max_seqs, pad_to_max_atoms=cfg.pad_to_max_atoms, pad_to_max_tokens=cfg.pad_to_max_tokens, pad_to_max_seqs=cfg.pad_to_max_seqs, atoms_per_window_queries=cfg.atoms_per_window_queries, min_dist=cfg.min_dist, max_dist=cfg.max_dist, num_bins=cfg.num_bins, overfit=cfg.overfit, return_symmetries=cfg.return_val_symmetries, compute_frames=cfg.compute_frames, backbone_only=cfg.backbone_only, atom14=cfg.atom14, atom37=cfg.atom37, design=cfg.design, inverse_fold=cfg.inverse_fold, disulfide_prob=cfg.disulfide_prob, disulfide_on=cfg.disulfide_on, ) self.monomer_split = cfg.monomer_split print("monomer_split", self.monomer_split) if self.monomer_split is not None: with Path(self.monomer_split).open("r") as f: monomer_ids = [x.lower() for x in f.read().splitlines()] print("monomer_split", monomer_ids) dataset = data_protein_binder.Dataset( struct_dir=Path(cfg.monomer_target_dir) / "structures", record_dir=Path(cfg.monomer_target_dir) / "records", target_ids=monomer_ids, seq_len=cfg.monomer_seq_len, tokenizer=cfg.tokenizer, featurizer=cfg.featurizer, ) # Load canonical molecules canonicals = load_canonicals(cfg.moldir) self.monomer_val_set = data_protein_binder.PredictionDataset( dataset=dataset, canonicals=canonicals, moldir=Path(cfg.moldir), backbone_only=cfg.backbone_only, atom14=cfg.atom14, atom37=cfg.atom37, design=cfg.design, target_structure_condition=cfg.monomer_target_structure_condition, inverse_fold=cfg.inverse_fold, disulfide_prob=cfg.disulfide_prob, disulfide_on=cfg.disulfide_on, ) self.ligand_split = cfg.ligand_split print("ligand_split", self.ligand_split) if self.ligand_split is not None: with Path(self.ligand_split).open("r") as f: ligand_ids = [x.lower() for x in f.read().splitlines()] print("ligand_split", ligand_ids) dataset = data_ligands.Dataset( struct_dir=Path(cfg.ligand_target_dir) / "structures", record_dir=Path(cfg.ligand_target_dir) / "records", target_ids=ligand_ids, min_len=cfg.ligand_seq_len, max_len=cfg.ligand_seq_len, tokenizer=cfg.tokenizer, featurizer=cfg.featurizer, ) # Load canonical molecules canonicals = load_canonicals(cfg.moldir) self.ligand_val_set = data_ligands.PredictionDataset( dataset=dataset, canonicals=canonicals, moldir=Path(cfg.moldir), backbone_only=cfg.backbone_only, atom14=cfg.atom14, atom37=cfg.atom37, design=cfg.design, disulfide_prob=cfg.disulfide_prob, disulfide_on=cfg.disulfide_on, ) def setup(self, stage: Optional[str] = None) -> None: # noqa: ARG002 (unused) """Run the setup for the DataModule. Parameters ---------- stage : str, optional The stage, one of 'fit', 'validate', 'test'. """ return def train_dataloader(self) -> DataLoader: """Get the training dataloader. Returns ------- DataLoader The training dataloader. """ return DataLoader( self._train_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 val_dataloader(self) -> DataLoader: """Get the validation dataloader. Returns ------- DataLoader The validation dataloader.s """ val_loaders = [] val_loaders.append( DataLoader( self._val_set, batch_size=self.cfg.val_batch_size, num_workers=self.cfg.num_workers if not self.inverse_fold else 1, pin_memory=self.cfg.num_workers if not self.inverse_fold else False, shuffle=False, collate_fn=collate, ) ) if self.monomer_split is not None: val_loaders.append( DataLoader( self.monomer_val_set, batch_size=self.cfg.val_batch_size, num_workers=self.cfg.num_workers if not self.inverse_fold else 1, pin_memory=self.cfg.pin_memory if not self.inverse_fold else False, shuffle=False, collate_fn=data_protein_binder.collate, ) ) if self.ligand_split is not None: val_loaders.append( DataLoader( self.ligand_val_set, batch_size=self.cfg.val_batch_size, num_workers=self.cfg.num_workers if not self.inverse_fold else 1, pin_memory=self.cfg.pin_memory if not self.inverse_fold else False, shuffle=False, collate_fn=data_ligands.collate, ) ) return val_loaders def predict_dataloader(self) -> DataLoader: return DataLoader( self._val_set, batch_size=self.cfg.val_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, # noqa: ARG002 ) -> Dict: """Transfer a batch to the given device. Parameters ---------- batch : Dict The batch to transfer. device : torch.device The device to transfer to. dataloader_idx : int The dataloader index. Returns ------- np.Any The transferred batch. """ for key in batch: if key not in [ "all_coords", "all_resolved_mask", "crop_to_all_atom_map", "chain_symmetries", "chain_swaps", "amino_acids_symmetries", "ligand_symmetries", "activity_name", "activity_qualifier", "sid", "cid", "normalized_protein_accession", "pair_id", "ligand_edge_index", "ligand_edge_lower_bounds", "ligand_edge_upper_bounds", "ligand_edge_bond_mask", "ligand_edge_angle_mask", "connections_edge_index", "ligand_chiral_atom_index", "ligand_chiral_check_mask", "ligand_chiral_atom_orientations", "ligand_stereo_bond_index", "ligand_stereo_check_mask", "ligand_stereo_bond_orientations", "ligand_aromatic_5_ring_index", "ligand_aromatic_6_ring_index", "ligand_planar_double_bond_index", "pdb_id", "id", "tokenized", "structure", "structure_bonds", "extra_mols", ]: if hasattr(batch[key], "to"): batch[key] = batch[key].to(device) return batch