# Setting up the OpenFold PDB training set from RODA

The multiple sequence alignments of OpenProteinSet and mmCIF structure files required to train OpenFold are freely available at the [Registry of Open Data on AWS (RODA)](https://registry.opendata.aws/openfold/). Additionally, OpenFold requires some postprocessing and [auxiliary files](Aux_seq_files.md) for training that need to be generated from the AWS data manually. This documentation is intended to give a full overview of those steps starting from the data download, assuming that the OpenFold codebase has already been set up on your system previously at the path `$OF_DIR` and the `openfold` environment is activated.

## 1. Downloading alignments and structure files
To fetch all the alignments corresponding to the original PDB training set of OpenFold alongside their mmCIF 3D structures, you can run the following commands:

```bash
mkdir -p alignment_data/alignment_dir_roda --recursive --no-sign-request
aws s3 cp s3://openfold/pdb/ alignment_data/alignment_dir_roda/ --recursive --no-sign-request

mkdir pdb_data
aws s3 cp s3://openfold/pdb_mmcif.zip pdb_data/ --no-sign-request
aws s3 cp s3://openfold/duplicate_pdb_chains.txt pdb_data/ --no-sign-request
```

The nested alignment directory structure is not yet exactly what OpenFold expects, so you can run the `flatten_roda.sh` script to convert them to the correct format:

```bash
bash $OF_DIR/scripts/flatten_roda.sh alignment_data/alignment_dir_roda alignment_data/
```

Afterwards, the old directory can be safely removed:

```bash
rm -r alignment_data/alignment_dir_roda
```

## 2. Creating alignment DBs (optional)
As further explained in [Auxiliary Sequence Files in OpenFold](Aux_seq_files.md), OpenFold supports an alternate format for storing alignments that can increase training performance in I/O bottlenecked systems. These so-called `alignment_db` files can be generated with the following script:

```bash
python $OF_DIR/scripts/alignment_db_scripts/create_alignment_db_sharded.py \
    alignment_data/alignments \
    alignment_data/alignment_dbs \
    alignment_db \
    --n_shards 10 \
    --duplicate_chains_file pdb_data/duplicate_pdb_chains.txt
```

We recommend creating 10 total `alignment_db` files (= "shards") for better
filesystem health and fast preprocessing, but note that this script will only run
optimally if the number of CPUs on your machine is at least as big as the number
of shards you are creating.

As an optional check, you can run the following command which should return 634,434:

```bash
grep "files" alignment_data/alignment_dbs/alignment_db.index | wc -l
```

## 3. Adding duplicate chains to alignments
To save space, the OpenProteinSet alignment database is stored without duplicates, meaning that only one representative alignment is stored for all chains with identical sequences in the PDB and duplicate instances are tracked with a [`duplicate_chains.txt`](Aux_seq_files.md#duplicate-pdb-chain-files) file. As OpenFold will select chains during training based on the chains in the alignment directory (or `alignment_db`), we therefore need to add those duplicate chains back in in order to train on the full conformational diversity of chains in the PDB.

If you've followed the optional Step 2, the `.index` file of your `alignment_db` files will have already been adjusted for duplicates and you can proceed to the next step. Otherwise, the standard alignment directory can be expanded to accommodate duplicates by inserting symlinked directories for the duplicate chains that point to their representative alignments:

```bash
python $OF_DIR/scripts/expand_alignment_duplicates.py \
    alignment_data/alignments \
    pdb_data/duplicate_pdb_chains.txt
```

As an optional check, the following command should return 634,434:

```bash
ls alignment_data/alignments/ | wc -l
```

## 4. Generating cluster-files
[JENNIFER: We could simply upload this cluster file as well and having it set up by the user is pretty unnecessary, but the scripts would still be useful to have for setup of new datasets]\
The AlphaFold dataloader adjusts the sampling probability of chains by their inverse cluster size, so we need to generate these sequence clusters for our training set.

As a first step, we'll need a `.fasta` file of all sequences in the training set. This can be generated with the following scripts, depending on how you set up your alignment data in the previous steps:

**Use this if you set up the duplicate-expanded alignment directory (faster):**
```bash
python $OF_DIR/scripts/alignment_data_to_fasta.py \
    alignment_data/all-seqs.fasta \
    --alignment_dir alignment_data/alignments
```

[JENNIFER: These scripts replace `data_dir_to_fasta.py` which is horribly slow as it reparses all mmCIF structure files. What those scripts do instead is fetch the sequence information from the >query line in the MSA files which is a lot faster. Technically this means that the chains in the generated .fasta may not 100% mirror the available mmCIF files in the PDB directory if there were some MSA generation errors for some of those structures. However as OF only trains on the structures with alignments available this is fine in practice, but I just wanted to note that the output of `data_dir_to_fasta.py` may technically not be 100% identical.]

**Use this if you set up the `alignment_db` files:**
```bash
python $OF_DIR/scripts/alignment_data_to_fasta.py \
    alignment_data/all-seqs.fasta \
    --alignment_db_index alignment_data/alignment_dbs/alignment_db.index
```

Next, we need to generate a cluster file at 40% sequence identity, which will contain all chains in a particular cluster on the same line. You'll need [MMSeqs2](https://github.com/soedinglab/MMseqs2?tab=readme-ov-file#installation) for this as well, which can be set up either in a conda environment or as a binary. (JENNIFER: do we want to add that as a dependency?)

```bash
python $OF_DIR/scripts/fasta_to_clusterfile.py \
    alignment_data/all-seqs.fasta \
    alignment_data/all-seqs_clusters-40.txt \
    /path/to/mmseqs \
    --seq-id 0.4
```

## 5. Generating cluster-files
As a last step, OpenFold requires "cache" files (JENNIFER: insert cross-link here) with metadata information for each chain that are used for choosing templates and samples during training.

The mmCIF-cache is used for filtering templates and can be generated with the following script:

```bash
mkdir pdb_data/data_caches

python $OF_DIR/scripts/generate_mmcif_cache.py \
    pdb_data/mmcif_files \
    pdb_data/data_caches/mmcif_cache.json \
    --no_workers 16
```

The chain-data-cache is used for filtering training samples and adjusting per-chain sampling probabilities and can be generated with the following script:

```bash
python $OF_DIR/scripts/generate_chain_data_cache.py \
    pdb_data/mmcif_files \
    pdb_data/data_caches/chain_data_cache.json \
    --cluster_file alignment_data/all-seqs_clusters-40.txt \
    --no_workers 16
```
