Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -textmodels/nn/__pycache__/feature_factory.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,260 @@
+---
+datasets:
+- OneScience/la-proteina
+frameworks:
+- ""
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- protein structure generation
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">La-Proteina</span>
+  </strong>
+</p>
+
+# Model Overview
+
+La-Proteina is a protein structure generation model based on **Partially Latent Flow Matching**. It can directly generate all-atom protein structures together with their corresponding amino acid sequences.
+
+Paper: _La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching_ (arXiv 2025).
+
+- [Paper link](https://arxiv.org/abs/2507.09466)
+
+# Model Description
+
+La-Proteina explicitly models the protein backbone (backbone CA), while sequence and atom-level details are captured through fixed-dimensional latent variables for each residue. This effectively avoids the challenges introduced by explicit side-chain representations.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :---: |
+| Protein structure generation | Generates protein backbones (backbone CA) and local latents based on flow matching. |
+| Motif-constrained generation | Supports motif position and sequence constraints for backbone design around functional motifs. |
+| Diffusion model training | Trains the main La-Proteina model on PDB datasets. |
+| Autoencoder training and inference | Trains the local-latent autoencoder and performs encoding, decoding, and reconstruction evaluation on PDB structures. |
+| Generated result evaluation | Computes metrics such as RMSD, sequence recovery, and (co-)designability. |
+
+# Usage
+
+## 1. Using OneCode
+
+You can try intelligent one-click AI4S programming through the OneCode online environment:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- Running on a GPU or DCU is recommended.
+- CPU can be used for connectivity checks, but it is relatively slow.
+- DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster.
+
+**Software Requirements**
+
+DCU users who want to learn more about adaptation details can contact liubiao@sugon.com.
+
+**Environment Checks**
+
+- NVIDIA GPU:
+
+```bash
+nvidia-smi
+```
+
+- Hygon DCU:
+
+```bash
+hy-smi
+```
+
+### Environment Preparation
+
+Optional: override the default paths through environment variables:
+
+```bash
+export LAPROTEINA_ROOT=/path/to/la-proteina
+export LAPROTEINA_DATASET_DIR=/path/to/dataset
+export LAPROTEINA_CHECKPOINTS_DIR=/path/to/checkpoints_laproteina
+export DATA_PATH=/path/to/dataset
+```
+
+## Quick Start
+
+### 1. Install the Runtime Environment
+
+```bash
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+If the following code cannot find required libraries at runtime, activate CUDA as shown below.
+
+```bash
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
+```
+
+### 2. Download the Model Package
+
+```bash
+# By default, the package is downloaded to the model folder under the current path. To change this, adjust the path after local_dir.
+hf download --model OneScience-Sugon/La-Proteina --local-dir ./model
+cd model
+```
+
+### Training Weights and Datasets
+
+They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.
+
+### 3. Usage Modes
+
+It is recommended to run the scripts under the `laproteina` directory so that outputs are managed in one place.
+
+#### 1. Train the Main La-Proteina Model (`run_train.sh`)
+
+```bash
+bash scripts/run_train.sh
+```
+
+**Common Hydra parameter overrides:**
+
+```bash
+# Single-device debugging
+bash scripts/run_train.sh hardware.ngpus_per_node_=1 single=true
+
+# Specify a run name
+bash scripts/run_train.sh run_name=my_laproteina_run
+
+# Override dataset or network configuration (for motif training scenarios)
+bash scripts/run_train.sh dataset=pdb/pdb_train_motif_aa nn=local_latents_score_nn_160M_motif_idx_aa
+```
+
+Output:
+
+- `./store/<run_name>/`: training logs, checkpoints, and Hydra configurations
+
+#### 2. Protein Structure Generation (`run_generate.sh`)
+
+```bash
+bash scripts/run_generate.sh
+```
+
+By default, the `inference_ucond_tri` configuration is used for unconditional generation (LD2 model + AE1).
+
+**Switching generation configurations:**
+
+```bash
+# Unconditional generation without triangle attention
+bash scripts/run_generate.sh --config_name inference_ucond_notri
+
+# Unconditional generation for long chains (300-800 residues)
+bash scripts/run_generate.sh --config_name inference_ucond_notri_long
+
+# Indexed all-atom motif scaffolding
+bash scripts/run_generate.sh --config_name inference_motif_idx_aa
+
+# Indexed tip-atom motif scaffolding
+bash scripts/run_generate.sh --config_name inference_motif_idx_tip
+
+# Unindexed all-atom motif scaffolding
+bash scripts/run_generate.sh --config_name inference_motif_uidx_aa
+
+# Unindexed tip-atom motif scaffolding
+bash scripts/run_generate.sh --config_name inference_motif_uidx_tip
+```
+
+Output:
+
+- `./inference/<config_name>/`: generated protein structure files and metadata
+
+---
+
+#### 3. Generated Result Evaluation (`run_evaluate.sh`)
+
+```bash
+bash scripts/run_evaluate.sh
+```
+
+By default, this evaluates the generated results corresponding to the `inference_ucond_tri` configuration.
+
+**Preparing ProteinMPNN weights:**
+
+ProteinMPNN weights must be downloaded before evaluation. They can be downloaded from the Hugging Face community:
+
+```bash
+hf download --model OneScience-Sugon/ProteinMPNN --local-dir ./weight
+```
+
+Output:
+
+- `./inference/<config_name>/evaluation/`: evaluation result files
+
+---
+
+#### 4. Autoencoder Inference (`run_ae_infer.sh`)
+
+```bash
+bash scripts/run_ae_infer.sh
+```
+
+This performs encode-decode reconstruction on the PDB dataset and evaluates reconstruction metrics, such as all-atom RMSD and sequence recovery.
+
+The script checks whether `DATA_PATH/pdb_train` and `AE1_ucond_512.ckpt` exist.
+
+The script internally calls:
+
+```bash
+python infer_laproteina_ae.py "$@"
+```
+
+Common parameter overrides:
+
+| Environment variable | Default value | Description |
+|----------|--------|------|
+| `LAPROTEINA_ROOT` | `${ONESCIENCE_DATASETS_DIR}/la-proteina` | Root directory for data and weights |
+| `LAPROTEINA_CHECKPOINTS_DIR` | `${LAPROTEINA_ROOT}/checkpoints_laproteina` | Autoencoder weight directory |
+| `DATA_PATH` | `${LAPROTEINA_ROOT}/dataset` | Dataset directory |
+
+Output:
+
+- `./inference_ae/`: reconstructed structures and evaluation metrics
+
+### Notes
+
+- Make sure the `ONESCIENCE_DATASETS_DIR` environment variable is correctly set before running the scripts.
+- The training script checks whether `DATA_PATH/pdb_train` and `AE1_ucond_512.ckpt` exist by default. If either is missing, it will report an error and exit.
+- In the current integration, `dataset=pdb` is available. `dataset=genie2` and `dataset=pdb_multimer` are not packaged in the current repository snapshot, and running them will produce an explicit error.
+- The scripts automatically set ROCm/DCU-related `LD_LIBRARY_PATH` values and can run directly on Hygon DCU platforms. On CUDA platforms, these settings can be ignored or adjusted as needed.
+- When generating motif scaffold structures, make sure the LD model is correctly paired with the corresponding AE model. Otherwise, generation may fail due to length or task mismatch.
+- Evaluating (co-)designability requires ProteinMPNN weights. Run `script_utils/download_pmpnn_weights.sh` in advance to download them.
+- It is recommended to run all scripts under `examples/biosciences/laproteina` so that output directories remain unified.
+- When overriding configurations through Hydra, you can use `+CK_PATH=...` to specify the checkpoint root path. If it is not provided, the script automatically sets it to `LAPROTEINA_ROOT`.
+
+# Official OneScience Information
+
+| Platform | OneScience main repository | Skills repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+The example code is licensed under Apache 2.0. La-Proteina model weights are licensed under the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/), and other materials are licensed under [CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode).
+
+If you use La-Proteina in your research, please cite the original paper:
+
+```bibtex
+@article{geffner2025laproteina,
+  title={La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching},
+  author={Geffner, Tomas and Didi, Kieran and Cao, Zhonglin and Reidenbach, Danny and Zhang, Zuobai and Dallago, Christian and Kucukbenli, Emine and Kreis, Karsten and Vahdat, Arash},
+  journal={arXiv preprint arXiv:2507.09466},
+  year={2025}
+}
+```

configs/__init__.py

@@ -0,0 +1 @@
+

configs/dataset/__init__.py

@@ -0,0 +1 @@
+

configs/dataset/pdb/__init__.py

@@ -0,0 +1 @@
+

configs/dataset/pdb/pdb_train_motif_aa.yaml

@@ -0,0 +1,54 @@
+datamodule:
+  _target_: "models.datasets.pdb_data.PDBLightningDataModule"
+  data_dir: ${oc.env:DATA_PATH}/pdb_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_: "models.datasets.transforms.CoordsToNanometers"
+    - _target_: "models.datasets.transforms.CenterStructureTransform"
+    - _target_: "models.datasets.transforms.GlobalRotationTransform"
+    - _target_: "models.datasets.transforms.ChainBreakPerResidueTransform"
+    - _target_: "models.datasets.transforms.MotifMaskTransform"
+      atom_selection_mode: "all"
+      motif_max_pct_res: 0.3
+      motif_prob: 1.0
+    - _target_: "models.datasets.transforms.CenteringTransform"
+      center_mode: "motif"
+      data_mode: "all-atom"
+    - _target_: "models.datasets.transforms.ExtractMotifCoordinatesTransform"
+  batch_size: 2 # Batch size for dataloader
+  num_workers: 16 # Number of workers for dataloader
+  pin_memory: True # Pin memory for dataloader
+
+  dataselector:
+    _target_: "models.datasets.pdb_data.PDBDataSelector"
+    data_dir: ${oc.env:DATA_PATH}/pdb_train/ # Directory where the dataset is stored
+    fraction: 0.001 # 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: 50 # Exclude peptides of length 50
+    max_length: 256 # 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: ["9b57", "9b5p", "9b5s", "9b5n", "9b5a", "9b5k", "9b5v", "9b59", "9b5w", "9b5g", "9b5i", "9b5x", "9b5d", "9b5e", "9b5t", "9b5f", 
+                  "9b5o", "9b58", "9b5u", "9b5c", "9b5j", "9b5b", "9b5h", "9b5q", "9b5l", "9b56", "9b5m", "9b5l", "9b55", "9b5r", "9ij9", "9iix",
+                  "2ezq", "1kld", "1crr", "2ezs", "2ezr", "1vve", "7ll9_G"]   # These ids are failed to be processed
+
+
+  datasplitter:
+    _target_: "models.datasets.pdb_data.PDBDataSplitter"
+    data_dir: ${oc.env:DATA_PATH}/pdb_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.5 # 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/dataset/pdb/pdb_train_motif_tip.yaml

@@ -0,0 +1,54 @@
+datamodule:
+  _target_: "models.datasets.pdb_data.PDBLightningDataModule"
+  data_dir: ${oc.env:DATA_PATH}/pdb_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_: "models.datasets.transforms.CoordsToNanometers"
+    - _target_: "models.datasets.transforms.CenterStructureTransform"
+    - _target_: "models.datasets.transforms.GlobalRotationTransform"
+    - _target_: "models.datasets.transforms.ChainBreakPerResidueTransform"
+    - _target_: "models.datasets.transforms.MotifMaskTransform"
+      atom_selection_mode: "tip_atoms"
+      motif_max_pct_res: 0.3
+      motif_prob: 1.0
+    - _target_: "models.datasets.transforms.CenteringTransform"
+      center_mode: "motif"
+      data_mode: "all-atom"
+    - _target_: "models.datasets.transforms.ExtractMotifCoordinatesTransform"
+  batch_size: 2 # Batch size for dataloader
+  num_workers: 16 # Number of workers for dataloader
+  pin_memory: True # Pin memory for dataloader
+
+  dataselector:
+    _target_: "models.datasets.pdb_data.PDBDataSelector"
+    data_dir: ${oc.env:DATA_PATH}/pdb_train/ # Directory where the dataset is stored
+    fraction: 0.001 # 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: 50 # Exclude peptides of length 50
+    max_length: 256 # 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: ["9b57", "9b5p", "9b5s", "9b5n", "9b5a", "9b5k", "9b5v", "9b59", "9b5w", "9b5g", "9b5i", "9b5x", "9b5d", "9b5e", "9b5t", "9b5f", 
+                  "9b5o", "9b58", "9b5u", "9b5c", "9b5j", "9b5b", "9b5h", "9b5q", "9b5l", "9b56", "9b5m", "9b5l", "9b55", "9b5r", "9ij9", "9iix",
+                  "2ezq", "1kld", "1crr", "2ezs", "2ezr", "1vve", "7ll9_G"]   # These ids are failed to be processed
+
+
+  datasplitter:
+    _target_: "models.datasets.pdb_data.PDBDataSplitter"
+    data_dir: ${oc.env:DATA_PATH}/pdb_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.5 # 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/dataset/pdb/pdb_train_ucond.yaml

@@ -0,0 +1,46 @@
+datamodule:
+  _target_: "models.datasets.pdb_data.PDBLightningDataModule"
+  data_dir: ${oc.env:DATA_PATH}/pdb_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_: "models.datasets.transforms.CoordsToNanometers"
+    - _target_: "models.datasets.transforms.CenterStructureTransform"
+    - _target_: "models.datasets.transforms.GlobalRotationTransform" # Transforms to apply to dataset examples
+    - _target_: "models.datasets.transforms.ChainBreakPerResidueTransform"
+  batch_size: 4 # Batch size for dataloader
+  num_workers: 16 # Number of workers for dataloader
+  pin_memory: True # Pin memory for dataloader
+
+  dataselector:
+    _target_: "models.datasets.pdb_data.PDBDataSelector"
+    data_dir: ${oc.env:DATA_PATH}/pdb_train/ # Directory where the dataset is stored
+    fraction: 0.001 # 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: 50 # Exclude peptides of length 50
+    max_length: 256 # 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: ["9b57", "9b5p", "9b5s", "9b5n", "9b5a", "9b5k", "9b5v", "9b59", "9b5w", "9b5g", "9b5i", "9b5x", "9b5d", "9b5e", "9b5t", "9b5f", 
+                  "9b5o", "9b58", "9b5u", "9b5c", "9b5j", "9b5b", "9b5h", "9b5q", "9b5l", "9b56", "9b5m", "9b5l", "9b55", "9b5r", "9ij9", "9iix",
+                  "2ezq", "1kld", "1crr", "2ezs", "2ezr", "1vve", "7ll9_G"]   # These ids are failed to be processed
+
+
+  datasplitter:
+    _target_: "models.datasets.pdb_data.PDBDataSplitter"
+    data_dir: ${oc.env:DATA_PATH}/pdb_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.5 # 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/__init__.py

@@ -0,0 +1 @@
+

configs/experiment_config/inference_my.yaml

@@ -0,0 +1,9 @@
+defaults:
+  - inference_base
+  - generation: uncond_codes
+  - _self_
+
+run_name_: laproteina_ucond_notri
+ckpt_path: ./store/test_release_diffusion/checkpoints
+ckpt_name: last-EMA.ckpt
+autoencoder_ckpt_path: ./store/test_release_ae/checkpoints/last-EMA.ckpt

configs/generation/__init__.py

@@ -0,0 +1 @@
+

configs/generation/motif.yaml

@@ -0,0 +1,76 @@
+defaults:
+  - motif_dict
+  - _self_
+
+# General sampling params
+args:
+  fold_cond: False     # not supported
+
+
+# Sampling dataset
+dataset:
+  # Lengths to sample. If not None, then use length-centric sampling
+  # Uses `nres_lens` if specified, otherwise arange(min_len, max_len+1, step)
+  nlens_cfg:
+    nres_lens: null
+    min_len:
+    max_len:
+    step_len:
+
+  cath_codes: null  # Not supported
+
+  # Number of samples per protein length or per cath code
+  nsamples: 200
+
+  # Maximum number of samples in each batch
+  max_nsamples_per_batch: 5
+
+  # Empirical (length, cath code) distribution
+  empirical_distribution_cfg:
+    # Path to the empirical (len, cath_code) distribution file
+    len_cath_code_path: null
+
+    # Level of cath codes to generate
+    cath_code_level: null
+
+    # Arguments for bucketizing the length distribution
+    bucket_min_len: 50
+    bucket_max_len: 274
+    bucket_step_size: 25
+
+
+# Sampling metrics
+metric:
+  # Designability metrics
+  compute_designability: True
+  designability_modes: ["ca", "bb3o"]  # Options: ca, bb3o
+  designability_folding_models: ["esmfold"]  # List of folding models to use
+  designability_motif_eval: True  # Enable motif-specific designability evaluation
+  
+  # Codesignability metrics
+  compute_codesignability: True
+  codesignability_modes: ["ca", "bb3o", "all_atom"]  # Options: ca, bb3o, all_atom
+  codesignability_folding_models: ["esmfold"]  # List of folding models to use
+  codesignability_motif_eval: True  # Enable motif-specific codesignability evaluation
+  compute_co_sequence_recovery: False
+
+  # Direct motif RMSD evaluation (generated structure vs ground truth motif)
+  compute_motif_rmsd: True
+  motif_rmsd_modes: ["ca", "bb3o", "all_atom"]  # Which atom models to compute for direct motif RMSD
+
+  # Output management
+  keep_folding_outputs: True  # Whether to keep folding model outputs (PDB files) after evaluation.
+                               # If false (default), temporary directories containing predicted structures
+                               # are deleted to save disk space. If true, all folding outputs are preserved
+                               # in the temporary directories for debugging and further analysis.
+                               # WARNING: Keeping outputs can consume significant disk space!
+
+  # Novelty
+  compute_novelty_pdb: False
+  compute_novelty_afdb: False
+  compute_novelty_afdb_rep_v4: False
+  compute_novelty_afdb_rep_v4_geniefilters_maxlen512: False
+
+  # FID-related metrics
+  compute_fid: False
+  metric_factory: null

configs/generation/motif_dict.yaml

@@ -0,0 +1,418 @@
+dataset:
+  motif_dict_cfg:
+    1PRW_AA:
+      contig_string: 5-20/A1-20/10-25/B1-20/5-20
+      motif_pdb_path: motif_benchmark_pdb_files/1prw_aa.pdb
+      motif_only: True
+      motif_min_length: 60
+      motif_max_length: 105
+      segment_order: "A;B"
+      atom_selection_mode: "all_atom"
+    1PRW_AA_TIP:
+      contig_string: 5-20/A1-20/10-25/B1-20/5-20
+      motif_pdb_path: motif_benchmark_pdb_files/1prw_aa.pdb
+      motif_only: True
+      motif_min_length: 60
+      motif_max_length: 105
+      segment_order: "A;B"
+      atom_selection_mode: "tip_atoms"
+    1BCF_AA:
+      contig_string: 8-15/A92-99/16-30/A123-130/16-30/A47-53/17-31/A18-25/8-15
+      motif_pdb_path: motif_benchmark_pdb_files/1bcf_aa.pdb
+      motif_only: True
+      motif_min_length: 96
+      motif_max_length: 152
+      segment_order: "A;B;C;D"
+      atom_selection_mode: "all_atom"
+    1BCF_AA_TIP:
+      contig_string: 8-15/A92-99/16-30/A123-130/16-30/A47-53/17-31/A18-25/8-15
+      motif_pdb_path: motif_benchmark_pdb_files/1bcf_aa.pdb
+      motif_only: True
+      motif_min_length: 96
+      motif_max_length: 152
+      segment_order: "A;B;C;D"
+      atom_selection_mode: "tip_atoms"
+    5TPN_AA:
+      contig_string: 10-40/A163-181/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/5tpn_aa.pdb
+      motif_only: True
+      motif_min_length: 50
+      motif_max_length: 75
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    5TPN_AA_TIP:
+      contig_string: 10-40/A163-181/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/5tpn_aa.pdb
+      motif_only: True
+      motif_min_length: 50
+      motif_max_length: 75
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    5IUS_AA:
+      contig_string: 0-30/B119-130/1/A132-140/15-40/A63-82/0-30
+      motif_pdb_path: motif_benchmark_pdb_files/5ius_aa.pdb
+      motif_only: True
+      motif_min_length: 57
+      motif_max_length: 142
+      segment_order: "B;A"
+      atom_selection_mode: "all_atom"
+    5IUS_AA_TIP:
+      contig_string: 0-30/B119-130/1/A132-140/15-40/A63-82/0-30
+      motif_pdb_path: motif_benchmark_pdb_files/5ius_aa.pdb
+      motif_only: True
+      motif_min_length: 57
+      motif_max_length: 142
+      segment_order: "B;A"
+      atom_selection_mode: "tip_atoms"
+    3IXT_AA:
+      contig_string: 10-40/P254-277/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/3ixt_aa.pdb
+      motif_only: True
+      motif_min_length: 50
+      motif_max_length: 75
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    3IXT_AA_TIP:
+      contig_string: 10-40/P254-277/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/3ixt_aa.pdb
+      motif_only: True
+      motif_min_length: 50
+      motif_max_length: 75
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    5YUI_AA:
+      contig_string: 5-30/A93-97/5-20/B118-120/10-35/C198-200/10-30
+      motif_pdb_path: motif_benchmark_pdb_files/5yui_aa.pdb
+      motif_only: True
+      motif_min_length: 50
+      motif_max_length: 100
+      segment_order: "A;B;C"
+      atom_selection_mode: "all_atom"
+    5YUI_AA_TIP:
+      contig_string: 5-30/A93-97/5-20/B118-120/10-35/C198-200/10-30
+      motif_pdb_path: motif_benchmark_pdb_files/5yui_aa.pdb
+      motif_only: True
+      motif_min_length: 50
+      motif_max_length: 100
+      segment_order: "A;B;C"
+      atom_selection_mode: "tip_atoms"
+    5AOU_AA:
+      contig_string: 40-60/A1051/20-40/A2083/20-35/A2110/100-140
+      motif_pdb_path: motif_benchmark_pdb_files/5aou_aa.pdb
+      motif_only: True
+      motif_min_length: 230
+      motif_max_length: 270
+      segment_order: "A;B;C"
+      atom_selection_mode: "all_atom"
+    5AOU_AA_TIP:
+      contig_string: 40-60/A1051/20-40/A2083/20-35/A2110/100-140
+      motif_pdb_path: motif_benchmark_pdb_files/5aou_aa.pdb
+      motif_only: True
+      motif_min_length: 230
+      motif_max_length: 270
+      segment_order: "A;B;C"
+      atom_selection_mode: "tip_atoms"
+    5AOU_QUAD_AA:
+      contig_string: 40-60/A1051/20-40/A2083/20-35/A2110/60-80/A2180/40-60
+      motif_pdb_path: motif_benchmark_pdb_files/5aou_quad_aa.pdb
+      motif_only: True
+      motif_min_length: 230
+      motif_max_length: 270
+      segment_order: "A;B;C;D"
+      atom_selection_mode: "all_atom"
+    5AOU_QUAD_AA_TIP:
+      contig_string: 40-60/A1051/20-40/A2083/20-35/A2110/60-80/A2180/40-60
+      motif_pdb_path: motif_benchmark_pdb_files/5aou_quad_aa.pdb
+      motif_only: True
+      motif_min_length: 230
+      motif_max_length: 270
+      segment_order: "A;B;C;D"
+      atom_selection_mode: "tip_atoms"
+    7K4V_AA:
+      contig_string: 40-50/A44/3-8/A50/70-85/A127/150-200
+      motif_pdb_path: motif_benchmark_pdb_files/7k4v_aa.pdb
+      motif_only: True
+      motif_min_length: 280
+      motif_max_length: 320
+      segment_order: "A;B;C"
+      atom_selection_mode: "all_atom"
+    7K4V_AA_TIP:
+      contig_string: 40-50/A44/3-8/A50/70-85/A127/150-200
+      motif_pdb_path: motif_benchmark_pdb_files/7k4v_aa.pdb
+      motif_only: True
+      motif_min_length: 280
+      motif_max_length: 320
+      segment_order: "A;B;C"
+      atom_selection_mode: "tip_atoms"
+    1YCR_AA:
+      contig_string: 10-40/B19-27/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/1ycr_aa.pdb
+      motif_only: True
+      motif_min_length: 40
+      motif_max_length: 100
+      segment_order: "A"
+      atom_selection_mode: "all_atom"
+    1YCR_AA_TIP:
+      contig_string: 10-40/B19-27/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/1ycr_aa.pdb
+      motif_only: True
+      motif_min_length: 40
+      motif_max_length: 100
+      segment_order: "A"
+      atom_selection_mode: "tip_atoms"
+    4JHW_AA:
+      contig_string: 10-25/F196-212/15-30/F63-69/10-25
+      motif_pdb_path: motif_benchmark_pdb_files/4jhw_aa.pdb
+      motif_only: True
+      motif_min_length: 60
+      motif_max_length: 90
+      segment_order: "B;A"
+      atom_selection_mode: "all_atom"
+    4JHW_AA_TIP:
+      contig_string: 10-25/F196-212/15-30/F63-69/10-25
+      motif_pdb_path: motif_benchmark_pdb_files/4jhw_aa.pdb
+      motif_only: True
+      motif_min_length: 60
+      motif_max_length: 90
+      segment_order: "B;A"
+      atom_selection_mode: "tip_atoms"
+    5WN9_AA:
+      contig_string: 10-40/A170-186/1/A188-189/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/5wn9_aa.pdb
+      motif_only: True
+      motif_min_length: 35
+      motif_max_length: 50
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    5WN9_AA_TIP:
+      contig_string: 10-40/A170-186/1/A188-189/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/5wn9_aa.pdb
+      motif_only: True
+      motif_min_length: 35
+      motif_max_length: 50
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    4ZYP_AA:
+      contig_string: 10-40/A422-436/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/4zyp_aa.pdb
+      motif_only: True
+      motif_min_length: 30
+      motif_max_length: 50
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    4ZYP_AA_TIP:
+      contig_string: 10-40/A422-436/10-40
+      motif_pdb_path: motif_benchmark_pdb_files/4zyp_aa.pdb
+      motif_only: True
+      motif_min_length: 30
+      motif_max_length: 50
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    6VW1_AA:
+      contig_string: 20-30/A24-42/4-10/A64-82/0-5
+      motif_pdb_path: motif_benchmark_pdb_files/6vw1_aa.pdb
+      motif_only: True
+      motif_min_length: 62
+      motif_max_length: 83
+      segment_order: A;B
+      atom_selection_mode: "all_atom"
+    6VW1_AA_TIP:
+      contig_string: 20-30/A24-42/4-10/A64-82/0-5
+      motif_pdb_path: motif_benchmark_pdb_files/6vw1_aa.pdb
+      motif_only: True
+      motif_min_length: 62
+      motif_max_length: 83
+      segment_order: A;B
+      atom_selection_mode: "tip_atoms"
+    1QJG_AA:
+      contig_string: 10-20/A38/15-30/A14/15-30/A99/10-20
+      motif_pdb_path: motif_benchmark_pdb_files/1qjg_aa.pdb
+      motif_only: False
+      motif_min_length: 53
+      motif_max_length: 103
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    1QJG_AA_TIP:
+      contig_string: 10-20/A38/15-30/A14/15-30/A99/10-20
+      motif_pdb_path: motif_benchmark_pdb_files/1qjg_aa.pdb
+      motif_only: False
+      motif_min_length: 53
+      motif_max_length: 103
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    1QJG_AA_NATIVE:
+      contig_string: 10-20/A14/15-30/A38/50-70/A99/25-30
+      motif_pdb_path: motif_benchmark_pdb_files/1qjg_aa.pdb
+      motif_only: False
+      motif_min_length: 115
+      motif_max_length: 135
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    1QJG_AA_NATIVE_TIP:
+      contig_string: 10-20/A14/15-30/A38/50-70/A99/25-30
+      motif_pdb_path: motif_benchmark_pdb_files/1qjg_aa.pdb
+      motif_only: False
+      motif_min_length: 115
+      motif_max_length: 135
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    2KL8_AA:
+      contig_string: A1-7/20/A28-79
+      motif_pdb_path: motif_benchmark_pdb_files/2kl8_aa.pdb
+      motif_only: False
+      motif_min_length: 79
+      motif_max_length: 79
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    2KL8_AA_TIP:
+      contig_string: A1-7/20/A28-79
+      motif_pdb_path: motif_benchmark_pdb_files/2kl8_aa.pdb
+      motif_only: False
+      motif_min_length: 79
+      motif_max_length: 79
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    7MRX_AA_60:
+      contig_string: 0-38/B25-46/0-38
+      motif_pdb_path: motif_benchmark_pdb_files/7mrx_aa.pdb
+      motif_only: False
+      motif_min_length: 60
+      motif_max_length: 60
+      segment_order: B
+      atom_selection_mode: "all_atom"
+    7MRX_AA_60_TIP:
+      contig_string: 0-38/B25-46/0-38
+      motif_pdb_path: motif_benchmark_pdb_files/7mrx_aa.pdb
+      motif_only: False
+      motif_min_length: 60
+      motif_max_length: 60
+      segment_order: B
+      atom_selection_mode: "tip_atoms"
+    7MRX_AA_85:
+      contig_string: 0-63/B25-46/0-63
+      motif_pdb_path: motif_benchmark_pdb_files/7mrx_aa.pdb
+      motif_only: False
+      motif_min_length: 85
+      motif_max_length: 85
+      segment_order: B
+      atom_selection_mode: "all_atom"
+    7MRX_AA_85_TIP:
+      contig_string: 0-63/B25-46/0-63
+      motif_pdb_path: motif_benchmark_pdb_files/7mrx_aa.pdb
+      motif_only: False
+      motif_min_length: 85
+      motif_max_length: 85
+      segment_order: B
+      atom_selection_mode: "tip_atoms"
+    7MRX_AA_128:
+      contig_string: 0-122/B25-46/0-122
+      motif_pdb_path: motif_benchmark_pdb_files/7mrx_aa.pdb
+      motif_only: False
+      motif_min_length: 128
+      motif_max_length: 128
+      segment_order: B
+      atom_selection_mode: "all_atom"
+    7MRX_AA_128_TIP:
+      contig_string: 0-122/B25-46/0-122
+      motif_pdb_path: motif_benchmark_pdb_files/7mrx_aa.pdb
+      motif_only: False
+      motif_min_length: 128
+      motif_max_length: 128
+      segment_order: B
+      atom_selection_mode: "tip_atoms"
+    5TRV_AA_SHORT:
+      contig_string: 0-35/A46-48/1/A50-55/1/A57-58/2/A61-65/0-35
+      motif_pdb_path: motif_benchmark_pdb_files/5trv_aa.pdb
+      motif_only: False
+      motif_min_length: 56
+      motif_max_length: 56
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    5TRV_AA_SHORT_TIP:
+      contig_string: 0-35/A46-48/1/A50-55/1/A57-58/2/A61-65/0-35
+      motif_pdb_path: motif_benchmark_pdb_files/5trv_aa.pdb
+      motif_only: False
+      motif_min_length: 56
+      motif_max_length: 56
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    5TRV_AA_MED:
+      contig_string: 0-65/A46-48/1/A50-55/1/A57-58/2/A61-65/0-65
+      motif_pdb_path: motif_benchmark_pdb_files/5trv_aa.pdb
+      motif_only: False
+      motif_min_length: 86
+      motif_max_length: 86
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    5TRV_AA_MED_TIP:
+      contig_string: 0-65/A46-48/1/A50-55/1/A57-58/2/A61-65/0-65
+      motif_pdb_path: motif_benchmark_pdb_files/5trv_aa.pdb
+      motif_only: False
+      motif_min_length: 86
+      motif_max_length: 86
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    5TRV_AA_LONG:
+      contig_string: 0-95/A46-48/1/A50-55/1/A57-58/2/A61-65/0-95
+      motif_pdb_path: motif_benchmark_pdb_files/5trv_aa.pdb
+      motif_only: False
+      motif_min_length: 116
+      motif_max_length: 116
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    5TRV_AA_LONG_TIP:
+      contig_string: 0-95/A46-48/1/A50-55/1/A57-58/2/A61-65/0-95
+      motif_pdb_path: motif_benchmark_pdb_files/5trv_aa.pdb
+      motif_only: False
+      motif_min_length: 116
+      motif_max_length: 116
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    6E6R_AA_SHORT:
+      contig_string: 0-35/A23-29/1/A31-35/0-35
+      motif_pdb_path: motif_benchmark_pdb_files/6e6r_aa.pdb
+      motif_only: False
+      motif_min_length: 48
+      motif_max_length: 48
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    6E6R_AA_SHORT_TIP:
+      contig_string: 0-35/A23-29/1/A31-35/0-35
+      motif_pdb_path: motif_benchmark_pdb_files/6e6r_aa.pdb
+      motif_only: False
+      motif_min_length: 48
+      motif_max_length: 48
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    6E6R_AA_MED:
+      contig_string: 0-65/A23-29/1/A31-35/0-65
+      motif_pdb_path: motif_benchmark_pdb_files/6e6r_aa.pdb
+      motif_only: False
+      motif_min_length: 78
+      motif_max_length: 78
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    6E6R_AA_MED_TIP:
+      contig_string: 0-65/A23-29/1/A31-35/0-65
+      motif_pdb_path: motif_benchmark_pdb_files/6e6r_aa.pdb
+      motif_only: False
+      motif_min_length: 78
+      motif_max_length: 78
+      segment_order: A
+      atom_selection_mode: "tip_atoms"
+    6E6R_AA_LONG:
+      contig_string: 0-95/A23-29/1/A31-35/0-95
+      motif_pdb_path: motif_benchmark_pdb_files/6e6r_aa.pdb
+      motif_only: False
+      motif_min_length: 108
+      motif_max_length: 108
+      segment_order: A
+      atom_selection_mode: "all_atom"
+    6E6R_AA_LONG_TIP:
+      contig_string: 0-95/A23-29/1/A31-35/0-95
+      motif_pdb_path: motif_benchmark_pdb_files/6e6r_aa.pdb
+      motif_only: False
+      motif_min_length: 108
+      motif_max_length: 108
+      segment_order: A
+      atom_selection_mode: "tip_atoms"

configs/generation/uncond_codes.yaml

@@ -0,0 +1,76 @@
+# General sampling params
+args:
+  fold_cond: False     # not supported
+
+
+# Sampling dataset
+dataset:
+  # Lengths to sample. If not None, then use length-centric sampling
+  # Uses `nres_lens` if specified, otherwise arange(min_len, max_len+1, step)
+  nlens_cfg:
+    # nres_lens: [100, 200, 300, 400, 500]
+    nres_lens: [100]
+    min_len:
+    max_len:
+    step_len:
+
+  cath_codes: null
+
+  # Number of samples per protein length or per cath code
+  # nsamples: 100
+  nsamples: 2
+
+  # Maximum number of samples in each batch
+  # max_nsamples_per_batch: 10
+  max_nsamples_per_batch: 1
+
+
+  # Empirical (length, cath code) distribution
+  empirical_distribution_cfg:
+    # Path to the empirical (len, cath_code) distribution file
+    len_cath_code_path: null
+
+    # Level of cath codes to generate
+    cath_code_level: null
+
+    # Arguments for bucketizing the length distribution
+    bucket_min_len: 50
+    bucket_max_len: 274
+    bucket_step_size: 25
+
+
+# Sampling metrics
+metric:
+  # Designability
+  compute_designability: True
+  designability_folding_models: ["esmfold"]  # List of folding models to use
+  
+  # codesign metrics
+  compute_codesignability: True
+  codesignability_folding_models: ["esmfold"]  # List of folding models to use
+  compute_co_sequence_recovery: False
+
+  # Novelty
+  compute_novelty_pdb: False
+  compute_novelty_afdb: False
+  compute_novelty_afdb_rep_v4: False
+  compute_novelty_afdb_rep_v4_geniefilters_maxlen512: False
+
+  # motif scaffolding
+  compute_motif_scaffolding: False
+  compute_motif_backbone_designability: False  # backbone scRMSD
+  compute_motif_aa_designability: False    # all-atom scRMSD
+  compute_motif_rmsd: False
+
+  # Output management
+  keep_folding_outputs: True
+
+  # binder design
+  compute_binder_metrics: False
+
+  # multimer generation
+  compute_multimer_metrics: False
+
+  # FID-related metrics
+  compute_fid: False
+  metric_factory: null

configs/generation/uncond_codes_800.yaml

@@ -0,0 +1,74 @@
+# General sampling params
+args:
+  fold_cond: False     # not supported
+
+
+# Sampling dataset
+dataset:
+  # Lengths to sample. If not None, then use length-centric sampling
+  # Uses `nres_lens` if specified, otherwise arange(min_len, max_len+1, step)
+  nlens_cfg:
+    nres_lens: [300, 400, 500, 600, 700, 800]
+    min_len:
+    max_len:
+    step_len:
+
+  # CATH codes to sample. If not None and nlens_cfg is None, then use fold-centric sampling
+  cath_codes: null       # ["1.x.x.x", "2.x.x.x", "3.x.x.x"]
+
+  # Number of samples per protein length or per cath code
+  nsamples: 100
+
+  # Maximum number of samples in each batch
+  max_nsamples_per_batch: 10
+
+
+  # Empirical (length, cath code) distribution
+  empirical_distribution_cfg:
+    # Path to the empirical (len, cath_code) distribution file
+    len_cath_code_path: null
+
+    # Level of cath codes to generate
+    cath_code_level: null
+
+    # Arguments for bucketizing the length distribution
+    bucket_min_len: 50
+    bucket_max_len: 274
+    bucket_step_size: 25
+
+
+# Sampling metrics
+metric:
+  # Designability
+  compute_designability: True
+  designability_folding_models: ["esmfold"]  # List of folding models to use
+  
+  # codesign metrics
+  compute_codesignability: True
+  codesignability_folding_models: ["esmfold"]  # List of folding models to use
+  compute_co_sequence_recovery: False
+
+  # Novelty
+  compute_novelty_pdb: False
+  compute_novelty_afdb: False
+  compute_novelty_afdb_rep_v4: False
+  compute_novelty_afdb_rep_v4_geniefilters_maxlen512: False
+
+  # motif scaffolding
+  compute_motif_scaffolding: False
+  compute_motif_backbone_designability: False  # backbone scRMSD
+  compute_motif_aa_designability: False    # all-atom scRMSD
+  compute_motif_rmsd: False
+
+  # Output management
+  keep_folding_outputs: True
+
+  # binder design
+  compute_binder_metrics: False
+
+  # multimer generation
+  compute_multimer_metrics: False
+
+  # FID-related metrics
+  compute_fid: False
+  metric_factory: null

configs/generation/validation_local_latents.yaml

@@ -0,0 +1,167 @@
+# General sampling params
+args:
+  # some params for sampling
+  nsteps: 200
+  self_cond: False
+  fold_cond: False
+  guidance_w: 1.0   # conditional model weights, 1.0 for conditional generation, 0.0 for unconditional generation. We typically set this value greater than 1
+  # Always set autoguidance as 0.0 during validation
+  ag_ratio: 0.0   # a value between 0 and 1, determining the proportion of autoguidance v.s. classifier-free guidance, 1.0 for all autoguidance, 0.0 for all CFG
+  ag_ckpt_path: null    # path to bad model checkpoint
+
+
+# Sampling dataset
+dataset:
+  # Validation sampling dataset
+  # Lengths to sample. If not None, then use length-centric sampling
+  # Uses `nres_lens` if specified, otherwise arange(min_len, max_len+1, step)
+  nlens_cfg:
+    nres_lens: [50, 100, 150, 200, 250]
+    min_len:
+    max_len:
+    step_len:
+
+  # CATH codes to sample. If not None and nlens_cfg is None, then use fold-centric sampling
+  cath_codes: null       # ["1.x.x.x", "2.x.x.x", "3.x.x.x"]
+
+  # Number of samples per protein length or per cath code
+  nsamples: 20
+
+  # Maximum number of samples in each batch
+  max_nsamples_per_batch: 5
+
+  # Empirical (length, cath code) distribution
+  empirical_distribution_cfg:
+    # Path to the empirical (len, cath_code) distribution file
+    len_cath_code_path: null
+
+    # Level of cath codes to generate
+    cath_code_level: null
+
+    # Arguments for bucketizing the length distribution
+    bucket_min_len: 50
+    bucket_max_len: 274
+    bucket_step_size: 25
+
+model:
+  ode:
+    bb_ca:
+      schedule:
+        mode: uniform
+        p: 1.0
+      gt:
+        mode: "tan"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: vf
+        sc_scale_noise: 0.0
+        sc_scale_score: 0.0
+        t_lim_ode: 0.98
+        t_lim_ode_below: 0.02
+        center_every_step: True
+    local_latents:
+      schedule:
+        mode: uniform
+        p: 1.0
+      gt:
+        mode: "tan"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: vf
+        sc_scale_noise: 0.0
+        sc_scale_score: 0.0
+        t_lim_ode: 0.98
+        t_lim_ode_below: 0.02
+        center_every_step: False
+  
+  temp_035_only_ca:
+    bb_ca:
+      schedule:
+        mode: uniform
+        p: 1.0
+      gt:
+        mode: "tan"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: sc
+        sc_scale_noise: 0.35
+        sc_scale_score: 1.0
+        t_lim_ode: 0.98
+        t_lim_ode_below: 0.02
+        center_every_step: True
+    local_latents:
+      schedule:
+        mode: uniform
+        p: 1.0
+      gt:
+        mode: "tan"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: vf
+        sc_scale_noise: 0.0
+        sc_scale_score: 0.0
+        t_lim_ode: 0.98
+        t_lim_ode_below: 0.02
+        center_every_step: False
+  
+  temp_035_all:
+    bb_ca:
+      schedule:
+        mode: uniform
+        p: 1.0
+      gt:
+        mode: "tan"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: sc
+        sc_scale_noise: 0.35
+        sc_scale_score: 1.0
+        t_lim_ode: 0.98
+        t_lim_ode_below: 0.02
+        center_every_step: True
+    local_latents:
+      schedule:
+        mode: uniform
+        p: 1.0
+      gt:
+        mode: "tan"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: sc
+        sc_scale_noise: 0.35
+        sc_scale_score: 1.0
+        t_lim_ode: 0.98
+        t_lim_ode_below: 0.02
+        center_every_step: False
+
+# Sampling metrics
+metric:
+  # Designability
+  compute_designability: False
+
+  # Novelty
+  compute_novelty_pdb: False
+  compute_novelty_afdb: False
+
+  # FID-related metrics
+  compute_fid: True
+  metric_factory:
+    - metrics: ["FID", "fJSD_C", "fJSD_A", "fJSD_T"]
+      ckpt_path: ${oc.env:DATA_PATH}/metric_factory/model_weights/gearnet_ca.pth
+      real_features_path: ${oc.env:DATA_PATH}/metric_factory/features/pdb_eval_ca_features.pth
+      ca_only: True
+      prefix: PDB_
+    - metrics: ["FID", "fJSD_C", "fJSD_A", "fJSD_T"]
+      ckpt_path: ${oc.env:DATA_PATH}/metric_factory/model_weights/gearnet_ca.pth
+      real_features_path: ${oc.env:DATA_PATH}/metric_factory/features/afdb_eval_ca_features.pth
+      ca_only: True
+      prefix: AFDB_
+    - metrics: ["IS_C", "IS_A", "IS_T"]
+      ckpt_path: ${oc.env:DATA_PATH}/metric_factory/model_weights/gearnet_ca.pth
+      ca_only: True

configs/inference_ae.yaml

@@ -0,0 +1,24 @@
+# Checkpoint
+#ckpt_file: ./checkpoints_laproteina/AE1_ucond_512.ckpt
+ckpt_file:  ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE1_ucond_512.ckpt
+# Job management
+ncpus_: 24
+seed: 5
+
+# Which data to test autoencoder on, and how many structures
+dataset: pdb
+n_structs: 1000
+bs: 20
+
+# Stuff to compute
+metrics:
+  bb_rmsd: false
+  all_atom_rmsd: true
+  seq_rec_rate: true
+  seq_cross_entropy: false
+  kl_latent_dist: true
+  std_latent_dist: false
+
+plot:
+  pca_latents: false
+  pca_latents_per_residue_type: true

configs/inference_base.yaml

@@ -0,0 +1,61 @@
+run_name_:
+ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina
+ckpt_name: 
+autoencoder_ckpt_path: 
+
+# Job management
+ncpus_: 24
+seed: 5
+gen_njobs: 1  # Leave as 1
+eval_njobs: 1  # Leave as 1
+
+# All sampling arugments that we want to change frequently
+generation:
+  n_recycle: 0  # Leave as 0
+
+  # General sampling params
+  args:
+    # nsteps: 400
+    # self_cond: True
+    nsteps: 20
+    self_cond: False
+
+    # Guidance
+    guidance_w: 1.0   # guidance model weights, 1.0 for w/o CFG and autoguidance, 0.0 for excluding the main model
+    ag_ratio: 0.0   # a value between 0 and 1, determining the proportion of autoguidance v.s. classifier-free guidance, 1.0 for all autoguidance, 0.0 for all CFG -- ag not supported for now
+    ag_ckpt_path: null  # path to bad model checkpoint
+
+    save_trajectory_every: 0 # at which step interval to save trajectory snapshots of generation, 0 for no saving
+
+  # Model-specific sampling arguments
+  model:
+    bb_ca:
+      schedule:
+        mode: log
+        p: 2.0
+      gt:
+        mode: "1/t"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: sc  # Options are: vf (plain fow matching) or sc (using score, where parameters below matter)
+        sc_scale_noise: 0.1  # scale used to multiply noise if mode == sc
+        sc_scale_score: 1.0  # scale used to multiply score if mode == sc
+        t_lim_ode: 0.98  # At what point between 0 and 1 we switch to ODE with score scaling (for sc mode)
+        t_lim_ode_below: 0.02  # Before what point between 0 and 1 we use sc mode (for vf_ss mode)
+        center_every_step: True  # Whether to center the sample every step
+    local_latents:
+      schedule:
+        mode: power
+        p: 2.0
+      gt:
+        mode: "tan"
+        p: 1.0
+        clamp_val: null
+      simulation_step_params:
+        sampling_mode: sc
+        sc_scale_noise: 0.1
+        sc_scale_score: 1.0
+        t_lim_ode: 0.98
+        t_lim_ode_below: 0.02
+        center_every_step: False

configs/inference_motif_idx_aa.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - inference_base
+  - generation: motif
+  - _self_
+
+run_name_: laproteina_motif_idx_aa
+ckpt_name: LD4_motif_idx_aa.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE3_motif.ckpt
+
+generation:
+  model:
+    bb_ca:
+      simulation_step_params:
+        center_every_step: False  # Whether to center the sample every step
+  dataset:
+    motif_task_name: 1YCR_AA

configs/inference_motif_idx_tip.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - inference_base
+  - generation: motif
+  - _self_
+
+run_name_: laproteina_motif_idx_tip
+ckpt_name: LD5_motif_idx_tip.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE3_motif.ckpt
+
+generation:
+  model:
+    bb_ca:
+      simulation_step_params:
+        center_every_step: False  # Whether to center the sample every step
+  dataset:
+    motif_task_name: 1YCR_AA_TIP

configs/inference_motif_uidx_aa.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - inference_base
+  - generation: motif
+  - _self_
+
+run_name_: laproteina_motif_uidx_aa
+ckpt_name: LD6_motif_uidx_aa.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE3_motif.ckpt
+
+generation:
+  model:
+    bb_ca:
+      simulation_step_params:
+        center_every_step: False  # Whether to center the sample every step
+  dataset:
+    motif_task_name: 1YCR_AA

configs/inference_motif_uidx_tip.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - inference_base
+  - generation: motif
+  - _self_
+
+run_name_: laproteina_motif_uidx_tip
+ckpt_name: LD7_motif_uidx_tip.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE3_motif.ckpt
+
+generation:
+  model:
+    bb_ca:
+      simulation_step_params:
+        center_every_step: False  # Whether to center the sample every step
+  dataset:
+    motif_task_name: 1YCR_AA_TIP

configs/inference_my.yaml

@@ -0,0 +1,13 @@
+defaults:
+  - inference_base
+  - generation: uncond_codes
+  - _self_
+
+run_name_: laproteina_ucond_notri
+ckpt_path: ./store/test_release_diffusion/checkpoints
+ckpt_name: last-EMA.ckpt
+autoencoder_ckpt_path: ./store/test_release_ae/checkpoints/last-EMA.ckpt
+
+generation:
+  dataset:
+    max_nsamples_per_batch: 2  # 原来是 10，改成 2 或 1

configs/inference_ucond_notri.yaml

@@ -0,0 +1,9 @@
+defaults:
+  - inference_base
+  - generation: uncond_codes
+  - _self_
+
+run_name_: laproteina_ucond_notri
+ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina
+ckpt_name: LD1_ucond_notri_512.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE1_ucond_512.ckpt

configs/inference_ucond_notri_long.yaml

@@ -0,0 +1,17 @@
+defaults:
+  - inference_base
+  - generation: uncond_codes_800
+  - _self_
+
+run_name_: laproteina_ucond_notri_long
+ckpt_name: LD3_ucond_notri_800.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/checkpoints_laproteina/AE2_ucond_800.ckpt
+
+generation:
+  model:
+    bb_ca:
+      simulation_step_params:
+        sc_scale_noise: 0.15  # scale used to multiply noise if mode == sc
+    local_latents:
+      simulation_step_params:
+        sc_scale_noise: 0.05  # scale used to multiply noise if mode == sc

configs/inference_ucond_tri.yaml

@@ -0,0 +1,9 @@
+defaults:
+  - inference_base
+  - generation: uncond_codes
+  - _self_
+
+run_name_: laproteina_ucond_tri
+ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina
+ckpt_name: LD2_ucond_tri_512.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE1_ucond_512.ckpt

configs/nn/__init__.py

@@ -0,0 +1 @@
+

configs/nn/local_latents_score_nn_160M.yaml

@@ -0,0 +1,47 @@
+name: local_latents_transformer  # name of architecture
+
+output_parameterization:
+  bb_ca: v  # or x_1
+  local_latents: v  # or x_1
+
+
+nlayers: 14  # number of transformer layers
+token_dim: 768  # dimension of the tokens in the sequence
+nheads: 12  # number of attn heads  12
+parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
+
+strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
+# If True, if some feature is not provided, then it raises an error
+
+# feats_seq: ["res_seq_pdb_idx", "xt_bb_ca", "xt_local_latents", "x_sc_bb_ca", "x_sc_local_latents", "x_recycle_bb_ca", "x_recycle_local_latents"]  # Sequence features to include in initial representation
+feats_seq: ["xt_bb_ca", "xt_local_latents", "x_sc_bb_ca", "x_sc_local_latents", "optional_ca_coors_nm_seq_feat", "optional_res_type_seq_feat"]  # Sequence features to include in initial representation
+feats_cond_seq: ["time_emb_bb_ca", "time_emb_local_latents"]  # Sequence features to include in conditioning variables
+
+# Parameters for the features we extract
+dim_cond: 256  # dimension of conditioning vector
+idx_emb_dim: 256  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
+t_emb_dim: 256  # dimension of the time embedding
+
+# feats_pair_repr: ["rel_seq_sep", "xt_bb_ca_pair_dists", "x_sc_bb_ca_pair_dists", "optional_ca_pair_dist", "x_recycle_bb_ca_pair_dists"]  # Features to include in the pair representation
+feats_pair_repr: ["rel_seq_sep", "xt_bb_ca_pair_dists", "x_sc_bb_ca_pair_dists", "optional_ca_pair_dist"]  # Features to include in the pair representation
+feats_pair_cond: ["time_emb_bb_ca", "time_emb_local_latents"]  # Features to include in the pair representation conditioning
+
+# Parameters for the pair features we extract
+# Binning for the pair distances of noisy xt
+xt_pair_dist_dim: 30
+xt_pair_dist_min: 0.1  # in nm (not Å)
+xt_pair_dist_max: 3  # in nm (not Å)
+# Binning for the pair distances for self conditioning
+x_sc_pair_dist_dim: 30
+x_sc_pair_dist_min: 0.1  # in nm (not Å)
+x_sc_pair_dist_max: 3  # in nm (not Å)
+# Relative sequence separation
+seq_sep_dim: 127  # should be odd >= 5
+# Dimension of final pair representation
+pair_repr_dim: 256
+
+update_pair_repr: False  # whether to update pair representation
+update_pair_repr_every_n: 3  # Update the pair representation every n layers
+use_tri_mult: False  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation
+
+use_qkln: True

configs/nn/local_latents_score_nn_160M_motif_idx_aa.yaml

@@ -0,0 +1,45 @@
+name: local_latents_transformer  # name of architecture
+
+output_parameterization:
+  bb_ca: v  # or x_1
+  local_latents: v  # or x_1
+
+
+nlayers: 14  # number of transformer layers
+token_dim: 768  # dimension of the tokens in the sequence
+nheads: 12  # number of attn heads  12
+parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
+
+strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
+# If True, if some feature is not provided, then it raises an error
+
+feats_seq: ["xt_bb_ca", "xt_local_latents", "x_sc_bb_ca", "x_sc_local_latents", "optional_ca_coors_nm_seq_feat", "optional_res_type_seq_feat", "bulk_all_atom_xmotif"]  # Sequence features to include in initial representation
+feats_cond_seq: ["time_emb_bb_ca", "time_emb_local_latents"]  # Sequence features to include in conditioning variables
+
+# Parameters for the features we extract
+dim_cond: 256  # dimension of conditioning vector
+idx_emb_dim: 256  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
+t_emb_dim: 256  # dimension of the time embedding
+
+feats_pair_repr: ["rel_seq_sep", "xt_bb_ca_pair_dists", "x_sc_bb_ca_pair_dists", "optional_ca_pair_dist", "x_motif_pair_dists"] # Features to include in the pair representation
+feats_pair_cond: ["time_emb_bb_ca", "time_emb_local_latents"]  # Features to include in the pair representation conditioning
+
+# Parameters for the pair features we extract
+# Binning for the pair distances of noisy xt
+xt_pair_dist_dim: 30
+xt_pair_dist_min: 0.1  # in nm (not Å)
+xt_pair_dist_max: 3  # in nm (not Å)
+# Binning for the pair distances for self conditioning
+x_sc_pair_dist_dim: 30
+x_sc_pair_dist_min: 0.1  # in nm (not Å)
+x_sc_pair_dist_max: 3  # in nm (not Å)
+# Relative sequence separation
+seq_sep_dim: 127  # should be odd >= 5
+# Dimension of final pair representation
+pair_repr_dim: 256
+
+update_pair_repr: False  # whether to update pair representation
+update_pair_repr_every_n: 3  # Update the pair representation every n layers
+use_tri_mult: False  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation
+
+use_qkln: True

configs/nn/local_latents_score_nn_160M_motif_idx_tip.yaml

@@ -0,0 +1,45 @@
+name: local_latents_transformer  # name of architecture
+
+output_parameterization:
+  bb_ca: v  # or x_1
+  local_latents: v  # or x_1
+
+
+nlayers: 14  # number of transformer layers
+token_dim: 768  # dimension of the tokens in the sequence
+nheads: 12  # number of attn heads  12
+parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
+
+strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
+# If True, if some feature is not provided, then it raises an error
+
+feats_seq: ["xt_bb_ca", "xt_local_latents", "x_sc_bb_ca", "x_sc_local_latents", "optional_ca_coors_nm_seq_feat", "optional_res_type_seq_feat", "motif_mask", "motif_abs_coords", "motif_seq"]  # Sequence features to include in initial representation
+feats_cond_seq: ["time_emb_bb_ca", "time_emb_local_latents"]  # Sequence features to include in conditioning variables
+
+# Parameters for the features we extract
+dim_cond: 256  # dimension of conditioning vector
+idx_emb_dim: 256  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
+t_emb_dim: 256  # dimension of the time embedding
+
+feats_pair_repr: ["rel_seq_sep", "xt_bb_ca_pair_dists", "x_sc_bb_ca_pair_dists", "optional_ca_pair_dist", "x_motif_pair_dists"] # Features to include in the pair representation  #! x_motif_pair_dists was in old not in new
+feats_pair_cond: ["time_emb_bb_ca", "time_emb_local_latents"]  # Features to include in the pair representation conditioning
+
+# Parameters for the pair features we extract
+# Binning for the pair distances of noisy xt
+xt_pair_dist_dim: 30
+xt_pair_dist_min: 0.1  # in nm (not Å)
+xt_pair_dist_max: 3  # in nm (not Å)
+# Binning for the pair distances for self conditioning
+x_sc_pair_dist_dim: 30
+x_sc_pair_dist_min: 0.1  # in nm (not Å)
+x_sc_pair_dist_max: 3  # in nm (not Å)
+# Relative sequence separation
+seq_sep_dim: 127  # should be odd >= 5
+# Dimension of final pair representation
+pair_repr_dim: 256
+
+update_pair_repr: False  # whether to update pair representation
+update_pair_repr_every_n: 3  # Update the pair representation every n layers
+use_tri_mult: False  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation
+
+use_qkln: True

configs/nn/local_latents_score_nn_160M_motif_uidx.yaml

@@ -0,0 +1,45 @@
+name: local_latents_transformer_motif_uidx  # name of architecture
+
+output_parameterization:
+  bb_ca: v  # or x_1
+  local_latents: v  # or x_1
+
+
+nlayers: 14  # number of transformer layers
+token_dim: 768  # dimension of the tokens in the sequence
+nheads: 12  # number of attn heads  12
+parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
+
+strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
+# If True, if some feature is not provided, then it raises an error
+
+feats_seq: ["xt_bb_ca", "xt_local_latents", "x_sc_bb_ca", "x_sc_local_latents", "optional_ca_coors_nm_seq_feat", "optional_res_type_seq_feat", "motif_mask", "motif_abs_coords", "motif_seq"]  # Sequence features to include in initial representation
+feats_cond_seq: ["time_emb_bb_ca", "time_emb_local_latents"]  # Sequence features to include in conditioning variables
+
+# Parameters for the features we extract
+dim_cond: 256  # dimension of conditioning vector
+idx_emb_dim: 256  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
+t_emb_dim: 256  # dimension of the time embedding
+
+feats_pair_repr: ["rel_seq_sep", "xt_bb_ca_pair_dists", "x_sc_bb_ca_pair_dists", "optional_ca_pair_dist", "x_motif_pair_dists"] # Features to include in the pair representation  #! x_motif_pair_dists was in old not in new
+feats_pair_cond: ["time_emb_bb_ca", "time_emb_local_latents"]  # Features to include in the pair representation conditioning
+
+# Parameters for the pair features we extract
+# Binning for the pair distances of noisy xt
+xt_pair_dist_dim: 30
+xt_pair_dist_min: 0.1  # in nm (not Å)
+xt_pair_dist_max: 3  # in nm (not Å)
+# Binning for the pair distances for self conditioning
+x_sc_pair_dist_dim: 30
+x_sc_pair_dist_min: 0.1  # in nm (not Å)
+x_sc_pair_dist_max: 3  # in nm (not Å)
+# Relative sequence separation
+seq_sep_dim: 127  # should be odd >= 5
+# Dimension of final pair representation
+pair_repr_dim: 256
+
+update_pair_repr: False  # whether to update pair representation
+update_pair_repr_every_n: 3  # Update the pair representation every n layers
+use_tri_mult: False  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation
+
+use_qkln: True

configs/nn/local_latents_score_nn_160M_tri.yaml

@@ -0,0 +1,48 @@
+name: local_latents_transformer  # name of architecture
+
+output_parameterization:
+  bb_ca: v  # or x_1
+  local_latents: v  # or x_1
+
+
+nlayers: 14  # number of transformer layers
+token_dim: 768  # dimension of the tokens in the sequence
+nheads: 12  # number of attn heads  12
+parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
+
+strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
+# If True, if some feature is not provided, then it raises an error
+
+# feats_seq: ["xt_bb_ca", "xt_local_latents", "x_sc_bb_ca", "x_sc_local_latents", "optional_ca_coors_nm_seq_feat", "optional_res_type_seq_feat"]  # Sequence features to include in initial representation
+# feats_cond_seq: ["time_emb_bb_ca", "time_emb_local_latents"]  # Sequence features to include in conditioning variables
+
+feats_seq: ["xt_bb_ca", "xt_local_latents", "x_sc_bb_ca", "x_sc_local_latents", "optional_ca_coors_nm_seq_feat", "optional_res_type_seq_feat"]  # Sequence features to include in initial representation
+feats_cond_seq: ["time_emb_bb_ca", "time_emb_local_latents"]  # Sequence features to include in conditioning variables
+
+# Parameters for the features we extract
+dim_cond: 256  # dimension of conditioning vector
+idx_emb_dim: 256  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
+t_emb_dim: 256  # dimension of the time embedding
+
+feats_pair_repr: ["rel_seq_sep", "xt_bb_ca_pair_dists", "x_sc_bb_ca_pair_dists", "optional_ca_pair_dist"]  # Features to include in the pair representation
+feats_pair_cond: ["time_emb_bb_ca", "time_emb_local_latents"]  # Features to include in the pair representation conditioning
+
+# Parameters for the pair features we extract
+# Binning for the pair distances of noisy xt
+xt_pair_dist_dim: 30
+xt_pair_dist_min: 0.1  # in nm (not Å)
+xt_pair_dist_max: 3  # in nm (not Å)
+# Binning for the pair distances for self conditioning
+x_sc_pair_dist_dim: 30
+x_sc_pair_dist_min: 0.1  # in nm (not Å)
+x_sc_pair_dist_max: 3  # in nm (not Å)
+# Relative sequence separation
+seq_sep_dim: 127  # should be odd >= 5
+# Dimension of final pair representation
+pair_repr_dim: 256
+
+update_pair_repr: True  # whether to update pair representation
+update_pair_repr_every_n: 2  # Update the pair representation every n layers
+use_tri_mult: True  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation
+
+use_qkln: True

configs/nn_ae/__init__.py

@@ -0,0 +1 @@
+

configs/nn_ae/nn_130m.yaml

@@ -0,0 +1,73 @@
+name: transformer_ae  # name of architecture
+
+latent_z_dim: &latent_z_dim 8
+
+encoder:
+  nlayers: 12  # number of transformer layers
+  token_dim: 768  # dimension of the tokens in the sequence
+  nheads: 12  # number of attn heads
+  parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
+
+  strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
+  # If True, if some feature is not provided, then it raises an error
+
+  # feats_seq: ["res_seq_pdb_idx", "chain_break_per_res", "x1_aatype", "x1_a37coors_nm", "x1_a37coors_nm_rel", "x1_bb_angles", "x1_sidechain_angles"]  # Sequence features to include in initial representation
+  feats_seq: ["chain_break_per_res", "x1_aatype", "x1_a37coors_nm", "x1_a37coors_nm_rel", "x1_bb_angles", "x1_sidechain_angles", "chain_idx_seq"]  # Sequence features to include in initial representation
+  # feats_seq: ["chain_break_per_res", "x1_aatype", "x1_bb_angles", "x1_sidechain_angles"]  # Sequence features to include in initial representation
+  feats_cond_seq:  # Sequence features to include in conditioning variables
+
+  # Parameters for the features we extract
+  dim_cond: 128  # dimension of conditioning vector
+  idx_emb_dim: 128  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
+
+  feats_pair_repr: ["rel_seq_sep", "x1_bb_pair_dists_nm", "x1_bb_pair_orientation", "chain_idx_pair"]  # Features to include in the pair representation
+  # feats_pair_repr: ["rel_seq_sep", "x1_bb_pair_dists_nm"]  # Features to include in the pair representation
+
+  seq_sep_dim: 127  # should be odd >= 5
+  pair_repr_dim: 256
+
+  update_pair_repr: False  # whether to update pair representation
+  update_pair_repr_every_n: 3  # Update the pair representation every n layers
+  use_tri_mult: False  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation
+
+  use_qkln: True
+
+  normalize_latent: false
+  latent_z_dim: *latent_z_dim
+
+decoder:
+  nlayers: 12  # number of transformer layers
+  token_dim: 768  # dimension of the tokens in the sequence
+
+  abs_coors: False  # Compute absolute coors or use relative to CA
+
+  # type: ff_local
+  # # This for 7m params
+  # nlayers: 12  # number of transformer layers
+  # token_dim: 768  # dimension of the tokens dim in the sequence
+
+  nheads: 12  # number of attn heads
+  parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
+
+  strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
+  # If True, if some feature is not provided, then it raises an error
+
+  feats_seq: ["ca_coors_nm", "z_latent_seq"]  # Sequence features to include in initial representation
+  feats_cond_seq:  # Sequence features to include in conditioning variables
+
+  # Parameters for the features we extract
+  dim_cond: 128  # dimension of conditioning vector
+  idx_emb_dim: 128  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
+
+  feats_pair_repr: ["rel_seq_sep", "ca_coors_nm_pair_dists"]  # Features to include in the pair representation
+
+  seq_sep_dim: 127  # should be odd >= 5
+  pair_repr_dim: 256
+
+  update_pair_repr: False  # whether to update pair representation
+  update_pair_repr_every_n: 3  # Update the pair representation every n layers
+  use_tri_mult: False  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation
+
+  use_qkln: True
+  latent_z_dim: *latent_z_dim
+

configs/training_ae.yaml

@@ -0,0 +1,67 @@
+run_name_: test_release_ae
+
+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
+
+defaults:
+  - nn_ae: nn_130m
+  - dataset: pdb/pdb_train_ucond
+  - _self_
+
+# Pre-training checkpoint to start from there
+pretrain_ckpt_path:
+
+# Loss
+loss:
+  kl:
+    weight: 1e-4
+    anneal: false
+    patience: 10000
+  struct:
+    type: l2
+    # type: l2cut_0.01
+    weight: 1.0
+  seq:
+    weight: 1.0
+
+
+# Precision
+force_precision_f32: False  # If false will use bf16-mixed precision
+
+# Optimization
+opt:
+  lr: 0.0001
+  max_epochs: 10000000
+  accumulate_grad_batches: 1
+  skip_nan_grad: False  # Skip updates with nan gradient
+  grad_and_weight_analysis: True  # Log some statistics of gradients and weights
+  dist_strategy: ddp  # For multi GPU training, do not change
+  val_check_interval: 5000  # Number of training steps after which we check validation loss
+  skip_large_grad_updates:
+    use: False  # Do not skip updates large gradients
+    moving_avg_size: 100
+    factor_threshold: 10
+    min_opt_steps: 1500
+
+# wandb logging and checkpoints
+log:
+  wandb_project: protein_ae
+  log_wandb: False  # whether to log to wandb
+  log_every_n_steps: 1  # wandb logging frequency
+  checkpoint: True  # whether to store checkpoints
+  checkpoint_every_n_steps: 10000  # How often we store a checkpoint, should be greater than val_check_interval in `opt`
+  last_ckpt_every_n_steps: 1500  # 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
+
+

configs/training_local_latents.yaml

@@ -0,0 +1,96 @@
+run_name_: test_release_diffusion
+
+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  ##2026  UPDATE
+  nnodes_: 1
+
+defaults:
+  - dataset: pdb/pdb_train_ucond
+  - nn: local_latents_score_nn_160M
+  - generation: validation_local_latents    # specifies all arguments for validation
+  - _self_
+
+
+# Pre-training checkpoint to start from there
+pretrain_ckpt_path:
+
+
+# Modalities to use
+product_flowmatcher:
+  bb_ca:
+    zero_com_noise: True  # whether to center noise samples
+    guidance_enabled: True  # whether this modality can use guidance
+    dim: 3  # Only CA atoms in bb_ca
+  local_latents:
+    zero_com_noise: False  # whether to center noise samples
+    guidance_enabled: True
+    dim: null  # Overwritten in the code by the correct value loaded from the autoencoder (likely 4 or 8)
+    
+
+    #autoencoder_ckpt_path: ./checkpoints_laproteina/AE1_ucond_512.ckpt
+autoencoder_ckpt_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/la-proteina/checkpoints_laproteina/AE1_ucond_512.ckpt
+
+# Loss
+loss:
+  t_distribution:
+    bb_ca:
+      name: mix_unif_beta
+      p1: 1.9
+      p2: 1.0
+      p3: 0.02
+    local_latents:
+      name: mix_unif_beta
+      p1: 1.0
+      p2: 1.5
+      p3: 0.2
+
+    shared_groups:
+
+# Precision
+force_precision_f32: False  # If false will use bf16-mixed precision
+
+# Training
+training:
+  n_recycle: 0
+  motif_conditioning: False
+  self_cond: True
+  p_folding_n_inv_folding_iters: 0.1
+
+# LoRA
+lora:
+  r: null   # Rank of low-rank adpation of matrices. Set r as null to turn off lora
+  lora_alpha: 32.0     # Scaling of lora matrices
+  lora_dropout: 0.0   # Dropout of lora inputs
+  train_bias: none  # Fine-tune bias or not. ['none', 'all', 'lora_only']
+
+# Optimization
+opt:
+  lr: 0.0001
+  max_epochs: 10000000
+  accumulate_grad_batches: 1
+  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
+  val_check_interval: 5000  # Number of training steps after which we check validation loss
+
+# wandb logging and checkpoints
+log:
+  wandb_project: laproteina
+  log_wandb: False  # whether to log to wandb
+  log_every_n_steps: 10  # wandb logging frequency
+  checkpoint: True  # whether to store checkpoints
+  checkpoint_every_n_steps: 10000  # How often we store a checkpoint, should be greater than val_check_interval in `opt`
+  last_ckpt_every_n_steps: 1500  # 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
+
+

models/__init__.py

@@ -0,0 +1,5 @@
+"""Embedded La-Proteina model package for OneScience."""
+
+from models import datasets, utils
+
+__all__ = ["datasets", "utils"]

models/datasets/base_data.py

@@ -0,0 +1,205 @@
+from abc import ABC, abstractmethod
+from typing import Callable, Dict, Iterable, List, Literal, Optional
+
+import lightning as L
+from loguru import logger
+from torch_geometric import transforms as T
+from torch_geometric.data import Dataset
+from torch_geometric.loader import DataLoader
+
+from models.utils.cluster_utils import ClusterSampler
+from models.utils.dense_padding_data_loader import DensePaddingDataLoader
+
+
+class BaseLightningDataModule(L.LightningDataModule, ABC):
+    """Base class for all datamodules"""
+
+    def __init__(
+        self,
+        batch_padding: bool = True,
+        sampling_mode: Literal["random", "cluster-random", "cluster-reps"] = "random",
+        transforms: Optional[List[Callable]] = None,
+        pre_transforms: Optional[List[Callable]] = None,
+        pre_filters: Optional[List[Callable]] = None,
+        batch_size: int = 32,
+        num_workers: int = 32,
+        pin_memory: bool = False,
+    ):
+        """Initialising the base data module class.
+
+        Args:
+            batch_padding (bool, optional): Whether batches should be padded to a dense representation
+                with the length being either a pre-specified max length or the maximum length of the
+                sample in the batch (base PyTorch batch) or whether a sparse representation should be
+                used (PyG batch). Defaults to True (base PyTorch batch).
+            sampling_mode (Literal["random", "cluster-random", "cluster-reps"], optional): How the data should be
+                sampled from the dataset later on:
+                - "random": Select a random sequence and ignore clusters.
+                - "cluster-random": Select a random sequence from each cluster. Keep all samples for each cluster.
+                - "cluster-reps": Select the cluster representative from each cluster. Only keep the representative for each cluster.
+                  Defaults to "random".
+            transforms (List[Callable]): List of transforms applied to each example.
+            pre_transforms (List[Callable]): List of transforms applied to each example before processing.
+            pre_filters (List[Callable]): List of filters applied to each example before processing.
+            batch_size (int, optional): Batch size used for dataloaders. Defaults to 32.
+            num_workers (int, optional): Number of workers used for dataloading. Defaults to 32.
+            pin_memory (bool, optional): Whether memory should be pinned. Defaults to False.
+        """
+        super().__init__()
+        self.batch_padding = batch_padding
+        self.sampling_mode = sampling_mode
+        self.transform = (
+            self._compose_transforms(transforms) if transforms is not None else None
+        )
+        self.pre_transform = (
+            self._compose_transforms(pre_transforms)
+            if pre_transforms is not None
+            else None
+        )
+        self.pre_filter = (
+            self._compose_filters(pre_filters) if pre_filters is not None else None
+        )
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+        self.train_ds = None
+        self.val_ds = None
+        self.test_ds = None
+        self.clusterid_to_seqid_mappings = None  # for cluster sampling
+
+    def setup(self, stage: Optional[str] = None):
+        if stage == "fit" or stage is None:
+            self.train_ds = self.train_dataset()
+        elif stage == "validation":
+            self.val_ds = self.val_dataset()
+        elif stage == "test":
+            self.test_ds = self.test_dataset()
+
+    def _compose_transforms(self, transforms: Iterable[Callable]) -> T.Compose:
+        try:
+            return T.Compose(list(transforms.values()))
+        except Exception:
+            return T.Compose(transforms)
+
+    def _compose_filters(self, filters: Iterable[Callable]) -> T.ComposeFilters:
+        try:
+            return T.ComposeFilters(list(filters.values()))
+        except Exception:
+            return T.ComposeFilters(filters)
+
+    @abstractmethod
+    def _get_dataset(self, split: str) -> Dataset:
+        """Creates a dataset given a split.
+
+        Args:
+            split (str): Split for which to get the dataset, with options "train", "val" or "test"
+
+        Returns:
+            Dataset: Dataset created for the respective split
+        """
+        ...
+
+    def train_dataset(self) -> Dataset:
+        return self._get_dataset("train")
+
+    def val_dataset(self) -> Dataset:
+        return self._get_dataset("val")
+
+    def test_dataset(self) -> Dataset:
+        return self._get_dataset("test")
+
+    def _get_dataloader(
+        self,
+        dataset: Dataset,
+        shuffle: bool = False,
+        clusterid_to_seqid_mapping: Dict[str, List[str]] = None,
+    ) -> DataLoader:
+        """Returns the dataloader for the corresponding dataset.
+
+        Args:
+            dataset (Dataset): PyG dataset for which the dataloader will be created.
+            shuffle (bool, optional): Whether the dataloader should be shuffled. Defaults to False. False when cluster_id mapping is given.
+            clusterid_to_seqid_mapping (Dict[str, List[str]], optional): Maps cluster ids to sequence ids. Defaults to None.
+
+        Returns:
+            DataLoader: Dataloader to be used by model.
+        """
+        if self.sampling_mode is None:
+            raise ValueError(
+                "Sampling mode not set, should be one of 'random', 'cluster-random' or 'cluster-reps'"
+            )
+        if clusterid_to_seqid_mapping and self.sampling_mode != "random":
+            sampler = ClusterSampler(
+                dataset=dataset,
+                clusterid_to_seqid_mapping=clusterid_to_seqid_mapping,
+                sampling_mode=self.sampling_mode,
+            )
+            shuffle = False
+        elif self.sampling_mode == "random":
+            sampler = None
+            shuffle = shuffle
+        else:
+            raise ValueError(
+                f"Sampling mode is {self.sampling_mode}, but clusterid_to_seqid_mapping is {clusterid_to_seqid_mapping}"
+            )
+
+        dataloader_class = DensePaddingDataLoader if self.batch_padding else DataLoader
+
+        return dataloader_class(
+            dataset,
+            batch_size=self.batch_size,
+            sampler=sampler,
+            shuffle=shuffle,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            drop_last=True,
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        if self.train_ds is None:
+            self.train_ds = self.train_dataset()
+        clusterid_to_seqid_mapping = (
+            self.clusterid_to_seqid_mappings["train"]
+            if self.clusterid_to_seqid_mappings
+            else None
+        )
+        shuffle = True
+        train_dl = self._get_dataloader(
+            dataset=self.train_ds,
+            shuffle=shuffle,
+            clusterid_to_seqid_mapping=clusterid_to_seqid_mapping,
+        )
+        return train_dl
+
+    def val_dataloader(self) -> DataLoader:
+        if self.val_ds is None:
+            self.val_ds = self.val_dataset()
+        clusterid_to_seqid_mapping = (
+            self.clusterid_to_seqid_mappings["val"]
+            if self.clusterid_to_seqid_mappings
+            else None
+        )
+        shuffle = False
+        logger.info(f"Length of validation set: {len(self.val_ds)}")
+        val_dl = self._get_dataloader(
+            dataset=self.val_ds,
+            shuffle=shuffle,
+            clusterid_to_seqid_mapping=clusterid_to_seqid_mapping,
+        )
+        return val_dl
+
+    def test_dataloader(self) -> DataLoader:
+        if self.test_ds is None:
+            self.test_ds = self.test_dataset()
+        clusterid_to_seqid_mapping = (
+            self.clusterid_to_seqid_mappings["test"]
+            if self.clusterid_to_seqid_mappings
+            else None
+        )
+        shuffle = False
+        test_dl = self._get_dataloader(
+            dataset=self.test_ds,
+            shuffle=shuffle,
+            clusterid_to_seqid_mapping=clusterid_to_seqid_mapping,
+        )
+        return test_dl

models/datasets/gen_dataset.py

@@ -0,0 +1,427 @@
+import random
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from loguru import logger
+from torch.utils.data import Dataset
+
+from models.utils.align_utils import mean_w_mask
+from models.utils.fold_utils import mask_cath_code_by_level
+from models.utils.motif_utils import parse_motif, save_motif_csv
+from models.utils.coors_utils import ang_to_nm
+
+
+class GenDataset(Dataset):
+    """
+    This class provides length-centric and fold-centric sampling for unconditional
+    and conditional protein structure generation. Each returned item is a dictionary
+    with key information for generation, which contains the length of proteins, the
+    number of proteins and cath codes if conditional sampling is used.
+
+    If length distribution is specified, sample `nsamples` proteins for each length,
+    cath codes are randomly sampled based on empirical distribution.
+    Otherwise, if cath code set is specified, sample `nsamples` proteins for each cath code,
+    lengths are randomly sampled based on empirical distribution.
+
+    Each sample returned by this dataset is a 2-tuple (L, nsamples) or 3-tuple (L, nsamples, cath_code) where
+      - nres (int) is the number of residues in the proteins to be samples
+      - nsamples (int) is the number of proteins to generate (happens in parallel),
+        so if nsamples=10 it means that it will produce 10 proteins of length L (all sampled in parallel)
+      - cath_code (List[str]) is the cath code for the nsamples if conditional generation is used
+    """
+
+    def __init__(
+        self,
+        nlens_cfg: Optional[Dict] = None,
+        cath_codes: Optional[List[str]] = None,
+        empirical_distribution_cfg: Optional[Dict] = None,
+        motif_dict_cfg: Optional[Dict] = None,
+        motif_task_name: Optional[str] = None,
+        motif_csv_path: Optional[str] = None,
+        target_as_features: bool = False,
+        nsamples: Optional[int] = 1,
+        max_nsamples_per_batch: Optional[int] = 1,
+        n_replicas: int = 1,
+    ):
+        """
+        Args:
+            nlens_cfg (Optional[Dict]): Config dict for specifying length distribution. If not None, use length-centric sampling.
+                Defaults to None.
+            cath_codes (Optional[List[str]]): List of cath codes to sample.. If not None and nlens_cfg is None, use fold-centric sampling.
+                Defaults to None.
+            empirical_distribution_cfg (Optional[Dict]): Config dict for specifying (length, cath code) distribution.
+                Defaults to None.
+            motif_dict_cfg (Optional[Dict]): Config dict for all motif configs.
+                Defaults to None.
+            motif_task_name (Optional[str]): Name of the motif task to sample.
+                Defaults to None.
+            motif_csv_path (Optional[str]): Path to the motif csv file.
+                Defaults to None.
+
+            nsamples (Optional[int]): Number of samples to generate for each length or each cath code.
+                Defaults to 1.
+            max_nsamples_per_batch (Optional[int]): Maximum number of samples for each batch.
+                Defaults to 1.
+            n_replicas (Optional[int]): Number of devices. Used for validation on multiple devices.
+                Defaults to 1.
+        """
+        super(GenDataset, self).__init__()
+        ##################################################################################
+        ################### 1. Parse length and cath codes ###############################
+        ##################################################################################
+        nres = self.parse_nlens_cfg(nlens_cfg)
+        self.target_as_features = target_as_features
+        self.motif_task_name = motif_task_name
+        if nres is not None:
+            logger.info("Use length-centric sampling.")
+            nsamples = [nsamples] * len(nres)
+        elif motif_task_name:
+            logger.info("Use motif-conditioned sampling.")
+            if motif_task_name in motif_dict_cfg:
+                motif_cfg = motif_dict_cfg[motif_task_name]
+            else:
+                raise ValueError(
+                    f"Motif task name {motif_task_name} not found in motif_dict_cfg"
+                )
+            nsamples = [nsamples]
+        else:
+            raise ValueError("Error in GenDataset init.")
+
+        ##################################################################################
+        ################### 2. Parse and bucketize empirical distribution ################
+        ##################################################################################
+        if empirical_distribution_cfg:
+            self.parse_empirical_distribution_cfg(empirical_distribution_cfg)
+            self.bucketize()
+
+        ##################################################################################
+        ################### 3. Generate data points ######################################
+        ##################################################################################
+        self.motif_masks = [None] * len(nsamples)
+        self.x_motifs = [None] * len(nsamples)
+        self.masks = [None] * len(nsamples)
+        if nres is not None:
+            # Length-centric generation
+            self.nres, self.cath_codes, self.nsamples = (
+                self.generate_cath_code_given_len(nres, nsamples)
+            )
+        elif cath_codes:
+            # Fold-centric generation
+            self.nres, self.cath_codes, self.nsamples = (
+                self.generate_len_given_cath_code(cath_codes, nsamples)
+            )
+        else:
+            self.nsamples = nsamples
+            self.cath_codes = [None] * len(nsamples)
+            self.motif_masks, self.x_motifs, self.residue_types = (
+                self.generate_motif_info(motif_cfg, nsamples[0], motif_csv_path)
+            )
+
+        ##################################################################################
+        # 4. Make sure the nsamples for each data point is not greater than max_nsamples #
+        ##################################################################################
+        if max_nsamples_per_batch:
+            if nres is not None or cath_codes:
+                self.nres, self.cath_codes, self.nsamples = self.flatten(
+                    max_nsamples_per_batch
+                )
+            else:
+                (
+                    self.nres,
+                    self.cath_codes,
+                    self.nsamples,
+                    self.masks,
+                    self.motif_masks,
+                    self.x_motifs,
+                    self.residue_types,
+                ) = self.flatten_motif(max_nsamples_per_batch)
+
+        ##################################################################################
+        # 5. Make sure this won't cause an error during validation on multiple devices ###
+        ##################################################################################
+        if n_replicas > 1:
+            self.pad_nlens(n_replicas)
+        assert all(
+            [n <= max_nsamples_per_batch for n in self.nsamples]
+        ), f"The nsamples for each len shouldn't be greater than {max_nsamples_per_batch}"
+        assert (
+            len(self.nsamples) % n_replicas == 0
+        ), f"Should be evenly splitable over {n_replicas} devices"
+
+        logger.info(
+            f"Adding generation dataset to sample {self.nsamples} sequences of length {self.nres}."
+        )
+
+    def bucketize(self):
+        """Build length buckets for cath_codes. Record the cath_code distribution given length bucket and the reverse"""
+        if self.len_cath_codes is None:
+            self.cath_codes_given_len_bucket = None
+            self.len_bucket_given_cath_codes = None
+            return
+
+        bucket = list(
+            range(self.bucket_min_len, self.bucket_max_len, self.bucket_step_size)
+        )
+        cath_codes_given_len_bucket = [[] for _ in range(len(bucket))]
+        len_bucket_given_cath_codes = defaultdict(set)
+        for _len, codes in self.len_cath_codes:
+            if len(codes) == 0:
+                continue
+            bucket_idx = (_len - self.bucket_min_len) // self.bucket_step_size
+            bucket_idx = min(bucket_idx, self.bucket_size - 1)  # Boundary cutoff
+            bucket_idx = max(bucket_idx, 0)
+
+            # Record all possible cath codes for each bucket
+            cath_codes_given_len_bucket[bucket_idx].append(codes)
+
+            # Record all possible len bucket for each cath code
+            for code in codes:
+                for level in ["C", "A", "T"]:
+                    ns = {"C": 3, "A": 2, "T": 1}
+                    level_code = code.rsplit(".", ns[level])[0] + ".x" * ns[level]
+                    len_bucket_given_cath_codes[level_code].add(bucket_idx)
+
+        for k, v in len_bucket_given_cath_codes.items():
+            len_bucket_given_cath_codes[k] = tuple(v)
+
+        self.cath_codes_given_len_bucket = cath_codes_given_len_bucket
+        self.len_bucket_given_cath_codes = len_bucket_given_cath_codes
+
+    def generate_cath_code_given_len(self, nres: List[int], nsamples: List[int]):
+        """Pre-generate corresponding cath codes for each length"""
+        cath_codes = []
+        for i in range(len(nres)):
+            if self.cath_codes_given_len_bucket is None:
+                cath_code = None
+            else:
+                if nres[i] <= self.bucket_max_len:
+                    bucket_idx = (
+                        nres[i] - self.bucket_min_len
+                    ) // self.bucket_step_size
+                else:
+                    bucket_idx = -1
+                cath_code = random.choices(
+                    self.cath_codes_given_len_bucket[bucket_idx], k=nsamples[i]
+                )
+            cath_codes.append(cath_code)
+        return nres, cath_codes, nsamples
+
+    def generate_len_given_cath_code(self, cath_codes: List[str], nsamples: List[int]):
+        """Pre-generate corresponding lengths for each cath code, then gather proteins of the same length as one batch"""
+        assert (
+            self.len_bucket_given_cath_codes is not None
+        ), "Need len_cath_code distribution for fold-centric generation"
+        tmp_nres = []
+        tmp_cath_codes = []
+        for i in range(len(cath_codes)):
+            for _ in range(nsamples[i]):
+                if cath_codes[i] not in self.len_bucket_given_cath_codes:
+                    raise ValueError(
+                        f"CATH code {cath_codes[i]} not in the empirical distribution"
+                    )
+                bucket_idx = random.choices(
+                    self.len_bucket_given_cath_codes[cath_codes[i]], k=1
+                )[0]
+                _len = self.bucket_min_len + bucket_idx * self.bucket_step_size
+
+                tmp_nres.append(_len)
+                tmp_cath_codes.append([cath_codes[i]])
+
+        # Gather the same lengths, as we need to generate proteins of the same length together
+        len_bucket = defaultdict(list)
+        out_nres, out_cath_codes, out_nsamples = [], [], []
+        for n, code in zip(tmp_nres, tmp_cath_codes):
+            len_bucket[n].append(code)
+
+        for n, code in len_bucket.items():
+            out_nres.append(n)
+            out_cath_codes.append(code)
+            out_nsamples.append(len(code))
+
+        return out_nres, out_cath_codes, out_nsamples
+
+    def generate_motif_info(self, motif_cfg, nsamples, motif_csv_path):
+        # Always return motif_masks, x_motifs, residue_types as lists of tensors, regardless of input type
+        lengths, motif_masks, x_motifs, residue_types, outstrs = parse_motif(
+            nsamples=nsamples, **motif_cfg
+        )
+        idx = np.argsort(lengths)
+        motif_masks = [motif_masks[i] for i in idx]
+        x_motifs = [x_motifs[i] for i in idx]
+        residue_types = [residue_types[i] for i in idx]
+        # center motifs to origin
+        for i in range(len(x_motifs)):
+            motif_center = mean_w_mask(
+                x_motifs[i].flatten(0, 1), motif_masks[i].flatten(0, 1)
+            ).unsqueeze(0)
+            x_motifs[i] = x_motifs[i] - motif_center
+            x_motifs[i] = x_motifs[i] * motif_masks[i][..., None]
+        # Only save CSV for contig_string (residue/range) case
+        if "motif_atom_spec" not in motif_cfg or motif_cfg["motif_atom_spec"] is None:
+            outstrs = [outstrs[i] for i in idx]
+            save_motif_csv(
+                motif_cfg["motif_pdb_path"],
+                self.motif_task_name,
+                outstrs,
+                outpath=motif_csv_path,
+                segment_order=motif_cfg["segment_order"],
+            )
+        return motif_masks, x_motifs, residue_types
+
+    def flatten(self, max_nsamples: int):
+        """Flatten the list to make sure each data point have no more than max_nsamples"""
+        nres, cath_codes, nsamples = [], [], []
+        for i in range(len(self.nsamples)):
+            for j in range(0, self.nsamples[i], max_nsamples):
+                nres.append(self.nres[i])
+                if self.cath_codes[i] is not None:
+                    cath_codes.append(self.cath_codes[i][j : j + max_nsamples])
+                else:
+                    cath_codes.append(None)
+                if j + max_nsamples <= self.nsamples[i]:
+                    nsamples.append(max_nsamples)
+                else:
+                    nsamples.append(self.nsamples[i] - j)
+        return nres, cath_codes, nsamples
+
+    def flatten_motif(self, max_nsamples: int):
+        """Flatten the list to make sure each data point have no more than max_nsamples"""
+        nres, cath_codes, nsamples = [], [], []
+        masks, motif_masks = [], []
+        x_motifs, residue_types = [], []
+        for i in range(len(self.nsamples)):
+            for j in range(0, self.nsamples[i], max_nsamples):
+
+                if self.cath_codes[i] is not None:
+                    cath_codes.append(self.cath_codes[i][j : j + max_nsamples])
+                else:
+                    cath_codes.append(None)
+                if j + max_nsamples <= self.nsamples[i]:
+                    nsamples.append(max_nsamples)
+                    motif_mask = self.motif_masks[j : j + max_nsamples]
+                    x_motif = self.x_motifs[j : j + max_nsamples]
+                    residue_type = self.residue_types[j : j + max_nsamples]
+                else:
+                    nsamples.append(self.nsamples[i] - j)
+                    motif_mask = self.motif_masks[j : self.nsamples[i]]
+                    x_motif = self.x_motifs[j : self.nsamples[i]]
+                    residue_type = self.residue_types[j : self.nsamples[i]]
+                mask = [torch.Tensor([True] * x.shape[0]) for x in motif_mask]
+                padded_mask = torch.nn.utils.rnn.pad_sequence(
+                    mask, batch_first=True, padding_value=False
+                )
+                padded_motif_mask = torch.nn.utils.rnn.pad_sequence(
+                    motif_mask, batch_first=True, padding_value=False
+                )
+                padded_x_motif = torch.nn.utils.rnn.pad_sequence(
+                    x_motif, batch_first=True, padding_value=0
+                )
+                padded_residue_type = torch.nn.utils.rnn.pad_sequence(
+                    residue_type, batch_first=True, padding_value=0
+                )
+                masks.append(padded_mask)
+                motif_masks.append(padded_motif_mask)
+                x_motifs.append(padded_x_motif)
+                residue_types.append(padded_residue_type)
+                nres.append(padded_mask.shape[1])
+        return nres, cath_codes, nsamples, masks, motif_masks, x_motifs, residue_types
+
+    def pad_nlens(self, n_replicas: int):
+        """Split nlens into data points (len, nsample) as val dataset and guarantee that
+        1. len(val_dataset) should be a multiple of n_replica, to ensure that we don't introduce additional samples for multi-gpu validation
+        2. nsample should be the same for all data points if n_replica > 1 (multi-gpu)
+        """
+        # Add samples to the small bins
+        max_nsamples = max(self.nsamples)
+        for i in range(len(self.nsamples)):
+            while self.cath_codes[i] != None and len(self.cath_codes[i]) < max_nsamples:
+                self.cath_codes[i] += self.cath_codes[i][
+                    : (max_nsamples - len(self.cath_codes[i]))
+                ]
+            self.nsamples[i] += max_nsamples - self.nsamples[i]
+
+        # Keep adding lengths in the dataset to make it a multiple of n_replica
+        while len(self.nres) % n_replicas != 0:
+            self.nres.append(self.nres[-1])
+            self.nsamples.append(max_nsamples)
+            self.cath_codes.append(self.cath_codes[-1])
+            self.cath_codes.append(self.cath_codes[-1])
+            if hasattr(self, "chain_masks"):
+                self.chain_masks.append(self.chain_masks[-1])
+            if hasattr(self, "general_masks"):
+                self.general_masks.append(self.general_masks[-1])
+            if hasattr(self, "structures"):
+                self.structures.append(self.structures[-1])
+
+    def parse_empirical_distribution_cfg(self, cfg: Dict):
+        """Load empirical (len, cath_codes) joint distribution. Apply mask according to the guidance cath code level"""
+        if cfg.len_cath_code_path is not None:
+            logger.info(
+                f"Loading empirical (length, cath_code) distribution from {cfg.len_cath_code_path}"
+            )
+            raw_len_cath_codes = torch.load(cfg.len_cath_code_path)
+
+            # By applying mask to the cath code distribution, we can control the level we want to sample
+            level = cfg.cath_code_level
+            self.len_cath_codes = []
+            for i in range(len(raw_len_cath_codes)):
+                _len, code = raw_len_cath_codes[i]
+                code = mask_cath_code_by_level(code, level="H")
+                if level == "A" or level == "C":
+                    code = mask_cath_code_by_level(code, level="T")
+                    if level == "C":
+                        code = mask_cath_code_by_level(code, level="A")
+                self.len_cath_codes.append((_len, code))
+
+            self.bucket_min_len = cfg.bucket_min_len
+            self.bucket_max_len = cfg.bucket_max_len
+            self.bucket_step_size = cfg.bucket_step_size
+            self.bucket_size = (
+                self.bucket_max_len - self.bucket_min_len
+            ) // self.bucket_step_size + 1
+        else:
+            logger.info(
+                "No empirical (length, cath_code) distribution provided. Use unconditional training."
+            )
+            self.len_cath_codes = None
+
+    def parse_nlens_cfg(self, cfg: Dict):
+        """Load nlens config."""
+        if cfg is None:
+            return None
+        if cfg.nres_lens:
+            nres = [int(n) for n in cfg.nres_lens]
+        elif cfg.min_len:
+            nres = np.arange(cfg.min_len, cfg.max_len + 1, cfg.step_len).tolist()
+        else:
+            nres = None
+        return nres
+
+    def __len__(self):
+        return len(self.nres)
+
+    def __getitem__(self, index: int):
+        result = {
+            "nres": self.nres[index],
+            "nsamples": self.nsamples[index],
+        }
+
+        # Add CATH codes if available
+        if self.cath_codes[index] is not None:
+            result["cath_code"] = self.cath_codes[index]
+        # Motif-conditioned
+        if self.motif_task_name is not None:
+            # Assume motif_mask, x_motif, seq_motif_mask, seq_motif are available
+            result["motif_mask"] = self.motif_masks[index].bool()  # [bs, num_res, 37]
+            result["x_motif"] = self.x_motifs[index]  # [bs, num_res, 37, 3]
+            result["seq_motif_mask"] = (
+                self.motif_masks[index].sum(dim=-1).bool()
+            )  # [bs, num_res]
+            result["seq_motif"] = self.residue_types[index]  # [bs, num_res]
+            result["mask"] = self.masks[index].bool()  # [bs, num_res]
+            return result
+
+        # Fallback: unconditional
+        return result