feat: pep-design

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")
+

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

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]

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

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)