From cbcb2569cbbfce9bfa52cb93c0929f92a0b3bd98 Mon Sep 17 00:00:00 2001 From: chungongyu Date: Fri, 25 Jul 2025 14:49:59 +0800 Subject: [PATCH] feat: pep-design --- configs/datasets_config/pdb/pep_train.yaml | 44 +++++++++ .../training_ca_motif_pep.yaml | 75 ++++++++++++++ proteinfoundation/datasets/pdb_data.py | 27 ++--- proteinfoundation/utils/cluster_utils.py | 2 +- script_utils/pdb_dataset_from_list.py | 99 +++++++++++++++++++ 5 files changed, 234 insertions(+), 13 deletions(-) create mode 100644 configs/datasets_config/pdb/pep_train.yaml create mode 100644 configs/experiment_config/training_ca_motif_pep.yaml create mode 100644 script_utils/pdb_dataset_from_list.py diff --git a/configs/datasets_config/pdb/pep_train.yaml b/configs/datasets_config/pdb/pep_train.yaml new file mode 100644 index 0000000..cc79403 --- /dev/null +++ b/configs/datasets_config/pdb/pep_train.yaml @@ -0,0 +1,44 @@ +datamodule: + _target_: "proteinfoundation.datasets.pdb_data.PDBLightningDataModule" + data_dir: ${oc.env:DATA_PATH}/pep_train/ # Directory where the dataset is stored + in_memory: False + format: "cif" # format for file download + overwrite: False # Whether to overwrite existing dataset files and reprocess the raw data + # arguments for BaseLightningDataModule class + batch_padding: True # whether we want a sparse PyG batch or a padded dense batch + sampling_mode: "cluster-random" # sample randomly inside each sequence similarity cluster during training + transforms: + - _target_: "proteinfoundation.datasets.transforms.GlobalRotationTransform" # Transforms to apply to dataset examples + - _target_: "proteinfoundation.datasets.transforms.ChainBreakPerResidueTransform" + # - _target_: "proteinfoundation.datasets.transforms.CATHLabelTransform" # activate for fold-conditional training + # root_dir: ${oc.env:DATA_PATH}/cathdata/ # Root directory for CATH labels + batch_size: 4 # Batch size for dataloader + num_workers: 16 # Number of workers for dataloader + pin_memory: True # Pin memory for dataloader + + # dataselector: + # _target_: "proteinfoundation.datasets.pdb_data.PDBDataSelector" + # data_dir: ${oc.env:DATA_PATH}/pep_train/ # Directory where the dataset is stored + # fraction: 0.95 # Fraction of dataset to use + # molecule_type: "protein" # Type of molecule for which to select + # experiment_types: ["diffraction", "EM"] # other options are "NMR" and "other" + # min_length: 8 # Exclude peptides of length 50 + # max_length: 512 # Exclude polypeptides greater than length 500 + # oligomeric_min: null + # oligomeric_max: null + # best_resolution: 0.0 # Include only proteins with resolution >= 0.0 + # worst_resolution: 5.0 # Include only proteins with resolution <= 5.0 + # has_ligands: [] # Include only proteins containing the ligand `ZN` + # remove_ligands: [] # Exclude specific ligands from any available protein-ligand complexes + # remove_non_standard_residues: True # Include only proteins containing standard amino acid residues + # remove_pdb_unavailable: True # Include only proteins that are available to download + # exclude_ids: [] # IDs can be excluded here like ["9b57", "1crr"] + + datasplitter: + _target_: "proteinfoundation.datasets.pdb_data.PDBDataSplitter" + data_dir: ${oc.env:DATA_PATH}/pep_train/ # Directory where the dataset is stored + train_val_test: [0.98, 0.019, 0.001] # Cross-validation ratios to use for train, val, and test splits + split_type: "sequence_similarity" # Split sequences by sequence similarity clustering, other option is "random" + split_sequence_similarity: 0.75 # Clustering at 50% sequence similarity (argument is ignored if split_type!="sequence_similarity") + overwrite_sequence_clusters: False # Previous clusterings at same sequence similarity are reused and not overwritten (argument is ignored if split_type!="sequence_similarity") + diff --git a/configs/experiment_config/training_ca_motif_pep.yaml b/configs/experiment_config/training_ca_motif_pep.yaml new file mode 100644 index 0000000..9d61805 --- /dev/null +++ b/configs/experiment_config/training_ca_motif_pep.yaml @@ -0,0 +1,75 @@ +run_name_: train_motif_1gpu + + +hardware: + ncpus_per_task_train_: 24 # Number of CPUs per tast during training + ncpus_per_task_prepro_: 32 # Number of CPUs used for data preprocessing run + accelerator: gpu + ngpus_per_node_: 1 # Number of GPUs per node + nnodes_: 1 # Number of nodes + + +# Below, for t_distribution, options are +# - name: uniform. p2 is the maximum time that we can sample (<=1, p1 is ignored). +# - name: logit-normal (normal + sigmoid). p1 is the mean of the normal, p2 the std (>0). +# - name: beta. This is beta(p1, p2). +loss: + t_distribution: + name: mix_up02_beta + p1: 1.9 + p2: 1.0 + loss_t_clamp: 0.9 # Used for loss stability in frameflow, 1. for no clamping + use_aux_loss: True # Whether to use auxiliary loss + aux_loss_t_lim: 0.3 # Time limit to apply auxiliary loss + thres_aux_2d_loss: 0.6 # This is nm not Å + aux_loss_weight: 1.0 + num_dist_buckets: 64 # Number of buckets to discretize the pairwise distance + max_dist_boundary: 1.0 # Given by nanometer + motif_aux_loss_weight: 5.0 + scaffold_aux_loss_weight: 0 + + +defaults: + - model: caflow_motif # caflow or frameflow + - _self_ + +dataset: pep_train +dataset_config_subdir: pdb + +force_precision_f32: False # If false will use bf16-mixed precision + +training: + motif_conditioning: True + motif_conditioning_sequence_rep: True + self_cond: True + fold_cond: False + mask_T_prob: 0.5 + mask_A_prob: 0.5 + mask_C_prob: 0.5 + fold_label_sample_ratio: [0.5, 0.1, 0.15, 0.25] # Training proportion for [None, C, CA, CAT]. If specified, will override mask_{C,A,T}_prob + +opt: + lr: 0.0001 + max_epochs: 10000000 + log_every_n_steps: 1 # For wandb + accumulate_grad_batches: 1 + val_check_interval: 5000 # Number of training steps after which we check validation loss + skip_nan_grad: False # Skip updates with nan gradient + grad_and_weight_analysis: False # Log some statistics of gradients and weights + dist_strategy: ddp # For multi GPU training, do not change + + +log: + wandb_project: protein_transformer_big_runs # Leave this so we can compare runs easily + log_wandb: True # whether to log to wandb + checkpoint: True # whether to store checkpoints + checkpoint_every_n_steps: 10000 # How often we store a checkpoint, should be greater than val_check_interval above + last_ckpt_every_n_steps: 3500 # How often do we update our last ckpt, needed for requeuing without losing progress + +seed: 42 + +ema: + decay: 0.999 # 0 means no EMA, so all the EMA machinery is unused and no EMA checkpoints are stored + validate_original_weights: False # Whether to run validation on regular or EMA weights + every_n_steps: 1 # Frequency of EMA updates + cpu_offload: False # Whether to offload EMA weights to cpu diff --git a/proteinfoundation/datasets/pdb_data.py b/proteinfoundation/datasets/pdb_data.py index 7c7ff71..052eb62 100644 --- a/proteinfoundation/datasets/pdb_data.py +++ b/proteinfoundation/datasets/pdb_data.py @@ -681,23 +681,23 @@ def _load_and_process_pdb( fill_value_coords=fill_value_coords, ) + coord_mask = graph.coords != fill_value_coords + graph.coord_mask = coord_mask[..., 0] + graph.residue_type = torch.tensor( + [resname_to_idx[residue] for residue in graph.residues] + ).long() + graph.database = "pdb" + graph.bfactor_avg = torch.mean(graph.bfactor, dim=-1) + graph.residue_pdb_idx = torch.tensor( + [int(s.split(":")[2]) for s in graph.residue_id], dtype=torch.long + ) + graph.seq_pos = torch.arange(graph.coords.shape[0]).unsqueeze(-1) except Exception as e: logger.warning(f"Error processing {pdb} {chains}: {e}") return None fname = f"{pdb}.pt" if chains == "all" else f"{pdb}_{chains}.pt" graph.id = fname.split(".")[0] - coord_mask = graph.coords != fill_value_coords - graph.coord_mask = coord_mask[..., 0] - graph.residue_type = torch.tensor( - [resname_to_idx[residue] for residue in graph.residues] - ).long() - graph.database = "pdb" - graph.bfactor_avg = torch.mean(graph.bfactor, dim=-1) - graph.residue_pdb_idx = torch.tensor( - [int(s.split(":")[2]) for s in graph.residue_id], dtype=torch.long - ) - graph.seq_pos = torch.arange(graph.coords.shape[0]).unsqueeze(-1) if self.pre_transform: graph = self.pre_transform(graph) @@ -760,10 +760,13 @@ def _get_dataset(self, split: Literal["train", "val", "test"]) -> PDBDataset: df_split = self.dfs_splits[split] self.clusterid_to_seqid_mappings = self.clusterid_to_seqid_mappings pdb_codes = df_split["pdb"].tolist() + # Check if 'chain' column exists in the DataFrame if 'chain' in df_split.columns: chains = df_split["chain"].tolist() - file_names = [f"{pdb}_{chain}" for pdb, chain in zip(pdb_codes, chains)] + def _file_name(pdb, chain): + return f"{pdb}_{chain}" if chain else f"{pdb}" + file_names = [_file_name(pdb, chain) for pdb, chain in zip(pdb_codes, chains)] else: chains = None file_names = [f"{pdb}" for pdb in pdb_codes] diff --git a/proteinfoundation/utils/cluster_utils.py b/proteinfoundation/utils/cluster_utils.py index 134cd2b..b9bc0e9 100644 --- a/proteinfoundation/utils/cluster_utils.py +++ b/proteinfoundation/utils/cluster_utils.py @@ -375,7 +375,7 @@ def expand_cluster_splits( if len(split_cluster_members_df) > 0: split_cluster_members_df[["pdb", "chain"]] = split_cluster_members_df[ "id" - ].str.split("_", n=1, expand=True) + ].str.split("_|$", n=1, expand=True) # Add the expanded DataFrame to the dictionary full_cluster_splits[split_name] = split_cluster_members_df # Add the split-specific cluster_dict to the dictionary diff --git a/script_utils/pdb_dataset_from_list.py b/script_utils/pdb_dataset_from_list.py new file mode 100644 index 0000000..84a3652 --- /dev/null +++ b/script_utils/pdb_dataset_from_list.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +import os +import sys +import contextlib +from dotenv import load_dotenv +import hydra +from loguru import logger + +@contextlib.contextmanager +def _open_file(fname): + if fname == "-": + yield sys.stdin + else: + with open(fname, "r") as f: + yield f + +def _read_pid(f): + for line in filter(lambda x: x, map(lambda x: x.strip(), f)): + if not line.startswith("#"): + yield line + +def main(args): + load_dotenv() + + # Load experiment config + config_path = "../configs/experiment_config" + with hydra.initialize(config_path, version_base=hydra.__version__): + cfg_exp = hydra.compose(config_name=args.config_name) + logger.info(f"Exp config {cfg_exp}") + + # Load data config + dataset_config_subdir = cfg_exp.get("dataset_config_subdir", None) + if dataset_config_subdir is not None: + # if args.dataset_config_subdir: + config_path = f"../configs/datasets_config/{dataset_config_subdir}" + else: + config_path = "../configs/datasets_config/" + with hydra.initialize(config_path, version_base=hydra.__version__): + cfg_data = hydra.compose(config_name=cfg_exp["dataset"]) + logger.info(f"Data config {cfg_data}") + + # create datamodule containing default train and val dataloader + datamodule = hydra.utils.instantiate(cfg_data.datamodule) + assert datamodule.dataselector is not None + df_data = datamodule.dataselector.create_dataset() + # filter by pdb_id + if args.pid_list_file: + with _open_file(args.pid_list_file) as f: + pid_list = list(_read_pid(f)) + df_data = df_data.loc[df_data["pdb"].isin(pid_list)] + logger.info(f"{len(df_data)} chains remaining") + + logger.info( + f"Dataset created with {len(df_data)} entries. Now downloading structure data..." + ) + + datamodule._download_structure_data(df_data["pdb"].tolist()) + + # process pdb files into seperate chains and save processed objects as .pt files + datamodule._process_structure_data( + df_data["pdb"].tolist(), df_data["chain"].tolist() + ) + # save df_data to disk for later use (in splitting, dataloading etc) + file_identifier = datamodule._get_file_identifier(datamodule.dataselector) + df_data_name = f"{file_identifier}.csv" + logger.info(f"Saving dataset csv to {df_data_name}") + df_data.to_csv(datamodule.dataselector.data_dir / df_data_name, index=False) + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--config_name", + type=str, + default="training_ca", + help="Name of the config yaml file.", + ) + parser.add_argument( + "--pid_list_file", + type=str, + default=None, + help="Filter pdb id by list file.", + ) + + args = parser.parse_args() + + main(args)