diff --git a/fuson_plm/benchmarking/caid/README.md b/fuson_plm/benchmarking/caid/README.md
index 30f3a9ab3dfc5ff5be8fbe36b9785ba07940f828..5359e20798e3537ee5f00fbb817eefcb79b5ece9 100755
--- a/fuson_plm/benchmarking/caid/README.md
+++ b/fuson_plm/benchmarking/caid/README.md
@@ -120,47 +120,46 @@ benchmarking/
Here we describe what each script does and which files each script creates.
1. π **`scrape_fusionpdb.py`**
- 1. Scrapes metadata for FusionPDB Level 2 and Level 3
- 1. Pulls the online tables for [Level 2](https://compbio.uth.edu/FusionPDB/gene_search_result_0.cgi?type=chooseLevel&chooseLevel=level2) and [Level 3](https://compbio.uth.edu/FusionPDB/gene_search_result_0.cgi?type=chooseLevel&chooseLevel=level3), saving results to `raw_data/FusionPDB_level2_curated_09_05_2024.csv` and `raw_data/FusionPDB_level3_curated_09_05_2024.csv` respectively.
- 2. Retrieves structure links
- 1. Using the tables collected in step (i), visits the page for each fusion oncoprotein (FO) in FusionPDB Level 2 and 3, and downloads all AlphaFold2 structure links for each FO.
- 2. Saves results directly to `raw_data/FusionPDB_level2_fusion_structure_links.csv` and `raw_data/FusionPDB_level3_fusion_structure_links.csv`, respectively
- 3. Retrieves FO head gene and tail gene info
- 1. Using the tables collected in step (i), visits the page for each fusion oncoprotein (FO) in FusionPDB Level 2 and 3 to download head/tail info. Collects HGID and TGID (GeneIDs for head and tail) and UniProt accessions for each.
- 2. Saves results directly to `raw_data/level2_head_tail_info.txt` and `raw_data/level3_head_tail_info.txt`, respectively.
- 4. Combines Level 2 and 3 head/tail data
- 1. Merges `raw_data/level2_head_tail_info.txt` and `raw_data/level3_head_tail_info.txt` into a dataframe.
- 2. Saves result at `processed_data/fusionpdb/fusion_heads_and_tails.csv` (columns="FusionGID","HGID","TGID","HGUniProtAcc","TGUniProtAcc")
- 5. Combines Level 2 and 3 structure link data
- 1. Joins structure link data with metadata for each of levels 2 and 3, then combines the result.
- 2. Saves result at `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structure_links.csv`
- 6. Combines structure link data and metadata (result of step (v)) with head and tail data (result of step (iv)), and resolves any missing head/tail UniProt IDs.
- 1. Merges the data
- 2. Checks how many rows have either missing or wrong UniProt accessions for the head or tail gene, and compiles the gene symbols for online quering in the UniProt ID Mapping tool (`processed_data/fusionpdb/intermediates/unmapped_parts.txt`)
- 3. Reads the UniProt ID Mapping result. Combines this data with FusionPDB-scraped data by matching FusionPDB's HGID (GeneID for head) and TGID (GeneID for tail) with the GeneID returned by UniProt.
- 4. For any FO where FusionPDB lacked a UniProt ID for the head/tail, this ID is filled in from the UniProt ID Mapping result.
- 5. Saves result to `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_head_tail_info.csv`. Columns: "FusionGID","FusionGene","Hgene","Tgene","URL","HGID","TGID","HGUniProtAcc","TGUniProtAcc","HGUniProtAcc_Source","TGUniProtAcc_Source", where the "_Source" columns indicate whether the UniProt ID came from FusionPDB, or from the ID Map.
- 7. Downloads AlphaFold2 structures of FOs from FusionPDB.
- 1. Using structure links from `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structure_links.csv` (step (v)), directly downloads `.pdb` and `.cif` files.
- 2. Saves results in π`raw_data/fusionpdb/structures`
-
-
+ i. Scrapes metadata for FusionPDB Level 2 and Level 3
+ a. Pulls the online tables for [Level 2](https://compbio.uth.edu/FusionPDB/gene_search_result_0.cgi?type=chooseLevel&chooseLevel=level2) and [Level 3](https://compbio.uth.edu/FusionPDB/gene_search_result_0.cgi?type=chooseLevel&chooseLevel=level3), saving results to `raw_data/FusionPDB_level2_curated_09_05_2024.csv` and `raw_data/FusionPDB_level3_curated_09_05_2024.csv` respectively.
+ ii. Retrieves structure links
+ a. Using the tables collected in step (i), visits the page for each fusion oncoprotein (FO) in FusionPDB Level 2 and 3, and downloads all AlphaFold2 structure links for each FO.
+ b. Saves results directly to `raw_data/FusionPDB_level2_fusion_structure_links.csv` and `raw_data/FusionPDB_level3_fusion_structure_links.csv`, respectively
+ iii. Retrieves FO head gene and tail gene info
+ a. Using the tables collected in step (i), visits the page for each fusion oncoprotein (FO) in FusionPDB Level 2 and 3 to download head/tail info. Collects HGID and TGID (GeneIDs for head and tail) and UniProt accessions for each.
+ b. Saves results directly to `raw_data/level2_head_tail_info.txt` and `raw_data/level3_head_tail_info.txt`, respectively.
+ iv. Combines Level 2 and 3 head/tail data
+ a. Merges `raw_data/level2_head_tail_info.txt` and `raw_data/level3_head_tail_info.txt` into a dataframe.
+ b. Saves result at `processed_data/fusionpdb/fusion_heads_and_tails.csv` (columns="FusionGID","HGID","TGID","HGUniProtAcc","TGUniProtAcc")
+ v. Combines Level 2 and 3 structure link data
+ a. Joins structure link data with metadata for each of levels 2 and 3, then combines the result.
+ b. Saves result at `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structure_links.csv`
+ vi. Combines structure link data and metadata (result of step (v)) with head and tail data (result of step (iv)), and resolves any missing head/tail UniProt IDs.
+ a. Merges the data
+ b. Checks how many rows have either missing or wrong UniProt accessions for the head or tail gene, and compiles the gene symbols for online quering in the UniProt ID Mapping tool (`processed_data/fusionpdb/intermediates/unmapped_parts.txt`)
+ c. Reads the UniProt ID Mapping result. Combines this data with FusionPDB-scraped data by matching FusionPDB's HGID (GeneID for head) and TGID (GeneID for tail) with the GeneID returned by UniProt.
+ d. For any FO where FusionPDB lacked a UniProt ID for the head/tail, this ID is filled in from the UniProt ID Mapping result.
+ e. Saves result to `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_head_tail_info.csv`. Columns: "FusionGID","FusionGene","Hgene","Tgene","URL","HGID","TGID","HGUniProtAcc","TGUniProtAcc","HGUniProtAcc_Source","TGUniProtAcc_Source", where the "_Source" columns indicate whether the UniProt ID came from FusionPDB, or from the ID Map.
+ vii. Downloads AlphaFold2 structures of FOs from FusionPDB.
+ a. Using structure links from `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structure_links.csv` (step (v)), directly downloads `.pdb` and `.cif` files.
+ b. Saves results in π`raw_data/fusionpdb/structures`
+
2. π **`process_fusion_structures.py`**
- 1. Determines pLDDT(s) for each FO structure.
- 1. For each structure in π`raw_data/fusionpdb_structures/`, determines amino acid sequence, per-residue pLDDT, and average pLDDT from the AlphaFold2 structure.
- 2. Saves results in `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structures_processed.csv`.
- 2. Downloads AlphaFold2 structures for all head and tail proteins
- 1. Reads `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_head_tail_info.csv` and collects all unique UniProt IDs for all head/tail proteins.
- 2. For each UniProt ID, queries the AlphaFoldDB, downloads the AlphaFold2 structure (if available), and saves it to π`raw_data/fusionpdb/head_tail_af2db_structures/`. Saves files converted from PDB to CIF format in `mmcif_converted_files`. Then, extracts the sequence, per-residue pLDDT, and average pLDDT from the file.
- 3. Saves any UniProt IDs that did not have structures in the AlphaFoldDB to: `processed_data/fusionpdb/intermediates/uniprotids_not_in_afdb.txt`. Most of these were very long, but the shorter ones were folded and their average pLDDTs were manually inputted. These were put back into the AlphaFold ID map to look for alternative UniProt IDs, and their results are in `not_in_afdb_idmap.txt`.
- 4. Saves results to `processed_data/fusionpdb/heads_tails_structural_data.csv`
- 3. Cleans the dataase of level 2&3 structural info
- 1. Drops rows where no structure was successfully downloaded
- 2. Drops rows where the FO sequence from FusionPDB does not match the FO sequence from its own AlphaFold2 structure file
- 3. βοΈSaves **two final, cleaned databases**βοΈ:
- 1. βοΈ **`FusionPDB_level2-3_cleaned_FusionGID_info.csv`**: includes ful IDs and structural information for the Hgene and Tgene of each FO. Columns="FusionGID","FusionGene","Hgene","Tgene","URL","HGID","TGID","HGUniProtAcc","TGUniProtAcc","HGUniProtAcc_Source","TGUniProtAcc_Source","HG_pLDDT","HG_AA_pLDDTs","HG_Seq","TG_pLDDT","TG_AA_pLDDTs","TG_Seq".
- 2. βοΈ **`FusionPDB_level2-3_cleaned_structure_info.csv`**: includes full structural information for each FO. Columns= "FusionGID","FusionGene","Fusion_Seq","Fusion_Length","Hgene","Hchr","Hbp","Hstrand","Tgene","Tchr","Tbp","Tstrand","Level","Fusion_Structure_Link","Fusion_Structure_Type","Fusion_pLDDT","Fusion_AA_pLDDTs","Fusion_Seq_Source"
+ i. Determines pLDDT(s) for each FO structure.
+ a. For each structure in π`raw_data/fusionpdb_structures/`, determines amino acid sequence, per-residue pLDDT, and average pLDDT from the AlphaFold2 structure.
+ b. Saves results in `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structures_processed.csv`.
+ ii. Downloads AlphaFold2 structures for all head and tail proteins
+ a. Reads `processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_head_tail_info.csv` and collects all unique UniProt IDs for all head/tail proteins.
+ b. For each UniProt ID, queries the AlphaFoldDB, downloads the AlphaFold2 structure (if available), and saves it to π`raw_data/fusionpdb/head_tail_af2db_structures/`. Saves files converted from PDB to CIF format in `mmcif_converted_files`. Then, extracts the sequence, per-residue pLDDT, and average pLDDT from the file.
+ c. Saves any UniProt IDs that did not have structures in the AlphaFoldDB to: `processed_data/fusionpdb/intermediates/uniprotids_not_in_afdb.txt`. Most of these were very long, but the shorter ones were folded and their average pLDDTs were manually inputted. These were put back into the AlphaFold ID map to look for alternative UniProt IDs, and their results are in `not_in_afdb_idmap.txt`.
+ d. Saves results to `processed_data/fusionpdb/heads_tails_structural_data.csv`
+ iii. Cleans the dataase of level 2&3 structural info
+ a. Drops rows where no structure was successfully downloaded
+ b. Drops rows where the FO sequence from FusionPDB does not match the FO sequence from its own AlphaFold2 structure file
+ c. βοΈSaves **two final, cleaned databases**βοΈ:
+ a. βοΈ **`FusionPDB_level2-3_cleaned_FusionGID_info.csv`**: includes ful IDs and structural information for the Hgene and Tgene of each FO. Columns="FusionGID","FusionGene","Hgene","Tgene","URL","HGID","TGID","HGUniProtAcc","TGUniProtAcc","HGUniProtAcc_Source","TGUniProtAcc_Source","HG_pLDDT","HG_AA_pLDDTs","HG_Seq","TG_pLDDT","TG_AA_pLDDTs","TG_Seq".
+ b. βοΈ **`FusionPDB_level2-3_cleaned_structure_info.csv`**: includes full structural information for each FO. Columns= "FusionGID","FusionGene","Fusion_Seq","Fusion_Length","Hgene","Hchr","Hbp","Hstrand","Tgene","Tchr","Tbp","Tstrand","Level","Fusion_Structure_Link","Fusion_Structure_Type","Fusion_pLDDT","Fusion_AA_pLDDTs","Fusion_Seq_Source"
### Training
@@ -184,9 +183,12 @@ PERMISSION_TO_OVERWRITE_EMBEDDINGS = False # if False, scrip
PERMISSION_TO_OVERWRITE_MODELS = False # if False, script will halt if it believes these embeddings have already been made.
```
-
+`train.py` trains the models using embeddings indicated in `config.py`. It also performs a hyperparameter screen.
+- All **results** are stored in `caid/results/`, where `timestamp` is a unique string encoding the date and time when you started training.
+- All **raw outputs from models** are stored in `caid/trained_models/`, where `embedding_path` represents the embeddings used to build the disorder predictor.
+- All **embeddings** made for training will be stored in a new folder called `caid/embeddings/` with subfolders for each model. This allows you to use the same model multiple times without regenerating embeddings.
-`train.py` trains the models using embeddings indicated in `config.py`. It also performs a hyperparameter screen. Model raw outputs (probabilities) and performance metrics are saved in `trained_models`. For example, FusOn-pLM-Diso raw outputs (ESM-2-650M-Diso has a folder in the same format, and future trained models will as well):
+Below is the FusOn-pLM-Diso raw outputs folder, `trained_models/fuson_plm/best/'. (ESM-2-650M-Diso has a folder in the same format, and future trained models will as well):
```
benchmarking/
@@ -209,7 +211,7 @@ benchmarking/
- **`caid_train_losses.csv`**: train losses over the 2 training epochs for top-performing model
- **`params.txt`**: hyperparameters of top performing model
-The training script also populates the `results` directory. Results from the FusOn-pLM manuscript are found in `results/final`. A few extra data files and plots are added by `analyze_fus`
+Results from the FusOn-pLM manuscript are found in `results/final`. A few extra data files and plots are added by `analyze_fusion_preds.py`
```
benchmarking/
diff --git a/fuson_plm/benchmarking/idr_prediction/README.md b/fuson_plm/benchmarking/idr_prediction/README.md
new file mode 100755
index 0000000000000000000000000000000000000000..7bf9a02dc7cb0bceb7db54ed4bf0201d50ef6e63
--- /dev/null
+++ b/fuson_plm/benchmarking/idr_prediction/README.md
@@ -0,0 +1,211 @@
+## IDR Property Prediction Benchmark
+
+This folder contains all the data and code needed to perform the **IDR property prediction benchmark**, where FusOn-pLM-IDR (a regressor built on FusOn-pLM embeddings) is used to predict aggregate properties of intrinsically disordered regions (IDRs), specifically asphericity, end-to-end radius (Re), radius of gyration (Rg), and polymer scaling exponent (Figure 4A-B).
+
+### TL;DR
+The order in which to run the scripts, after downloading data:
+
+```
+python clean.py # clean the data
+python cluster.py # MMSeqs2 clustering
+python split.py # make cluster-based train/val/test splits
+python train.py # train the model
+python plot.py # if you want to remake r2 plots
+```
+
+### Downloading raw IDR data
+IDR properties from [Lotthammer et al. 2024](https://doi.org/10.1038/s41592-023-02159-5) (ALBATROSS model) were used to train FusOn-pLM-Diso. Sequences were downloaded from [this link](https://github.com/holehouse-lab/supportingdata/blob/master/2023/ALBATROSS_2023/simulations/data/all_sequences.tgz) and deposited in `raw_data`. All files in `raw_data` are from this direct download.
+
+```
+benchmarking/
+βββ idr_prediction/
+ βββ raw_data/
+ βββ asph_bio_synth_training_data_cleaned_05_09_2023.tsv
+ βββ asph_nat_meth_test.tsv
+ βββ scaled_re_bio_synth_training_data_cleaned_05_09_2023.tsv
+ βββ scaled_re_nat_meth_test.tsv
+ βββ scaled_rg_bio_synth_training_data_cleaned_05_09_2023.tsv
+ βββ scaled_rg_nat_meth_test.tsv
+ βββ scaling_exp_bio_synth_training_data_cleaned_05_09_2023.tsv
+ βββ scaling_exp_nat_meth_test.tsv
+```
+- **`asph`**=asphericity, **`scaled_re`**=scaled Re, **`scaled_rg`**=scaled Rg, **`scaling_exp`**=polymer scaling exponent
+- **`_bio_synth_training_data_cleaned_05_09_2023.tsv`** are ALBATROSS **training data** for the four properties, downloaded directly their GitHub
+- **`_nat_meth_test.tsv`** are ALBATROSS **testing data** for the four proeprties, downloaded directly from their GitHub
+
+### Cleaning raw IDR data
+`clean.py` cleans the raw training and testing data separately for each property. Any duplicates (in both train and test) are removed from train and kept in test. Finally, the four are combined into one file:
+
+```
+benchmarking/
+βββ idr_prediction/
+ βββ processed_data/
+ βββ all_albatross_seqs_and_properties.csv
+```
+
+- **`all_albatross_seqs_and_properties.csv`**: Columns = "Sequence","IDs","UniProt_IDs","UniProt_Names","Split","asph","scaled_re","scaled_rg","scaling_exp". All splits are either "Train" or "Test", indicating ALBATROSS model's usage of them
+
+To perform cleaning, run
+
+```
+python clean.py
+```
+
+### Using config.py for clustering, splitting, training
+
+This file has configurations for clustering, splitting, training.
+
+```
+# Clustering Parameters
+CLUSTER = CustomParams(
+ # MMSeqs2 parameters: see GitHub or MMSeqs2 Wiki for guidance
+ MIN_SEQ_ID = 0.3, # % identity
+ C = 0.5, # % sequence length overlap
+ COV_MODE = 1, # cov-mode: 0 = bidirectional, 1 = target coverage, 2 = query coverage, 3 = target-in-query length coverage.
+ CLUSTER_MODE = 2,
+ # File paths
+ INPUT_PATH = 'processed_data/all_albatross_seqs_and_properties.csv',
+ PATH_TO_MMSEQS = '../../mmseqs' # path to where you installed MMSeqs2
+)
+
+# Split config
+SPLIT = CustomParams(
+ IDR_DB_PATH = 'processed_data/all_albatross_seqs_and_properties.csv',
+ CLUSTER_OUTPUT_PATH = 'clustering/mmseqs_full_results.csv',
+ RANDOM_STATE_1 = 2, # random_state_1 = state for splitting all data into train & other
+ TEST_SIZE_1 = 0.21, # test size for data -> train/test split. e.g. 20 means 80% clusters in train, 20% clusters in other
+ RANDOM_STATE_2 = 6, # random_state_2 = state for splitting other from ^ into val and test
+ TEST_SIZE_2 = 0.50 # test size for train -> train/val split. e.g. 0.50 means 50% clusters in train, 50% clusters in test
+
+)
+
+# Which models to benchmark
+TRAIN = CustomParams(
+ BENCHMARK_FUSONPLM = True,
+ FUSONPLM_CKPTS= "FusOn-pLM", # Dictionary: key = run name, values = epochs, or string "FusOn-pLM"
+ BENCHMARK_ESM = True,
+
+ # GPU configs
+ CUDA_VISIBLE_DEVICES="0",
+
+ # Overwriting configs
+ PERMISSION_TO_OVERWRITE_EMBEDDINGS = False, # if False, script will halt if it believes these embeddings have already been made.
+ PERMISSION_TO_OVERWRITE_MODELS = False # if False, script will halt if it believes these embeddings have already been made.
+)
+```
+
+### Clustering
+Clustering of all sequences in `all_albatross_seqs_and_properties.csv` is performed by `cluster.py`.
+
+The clustering command entered by the script is:
+```
+mmseqs easy-cluster clustering/input.fasta clustering/raw_output/mmseqs clustering/raw_output --min-seq-id 0.3 -c 0.5 --cov-mode 1 --cluster-mode 2 --dbtype 1
+```
+The script will generate the following files:
+```
+benchmarking/
+βββ idr_prediction/
+ βββ clustering/
+ βββ input.fasta
+ βββ mmseqs_full_results.csv
+```
+- **`clustering/input.fasta`**: the input file used by MMSeqs2 to cluster the fusion oncoprotein sequences. Headers are our assigned sequence IDs (can be found in the `IDs` column of `processed_data/all_albatross_seqs_and_properties.csv`.)
+- **`clustering/mmseqs_full_results.csv`**: clustering results. Columns:
+ - `representative seq_id`: the seq_id of the sequence representing this cluster
+ - `member seq_id`: the seq_id of a member of the cluster
+ - `representative seq`: the amino acid sequence of the cluster representative (representative seq_id)
+ - `member seq`: the amino acid sequence of the cluster member
+
+### Splitting
+Cluster-based splitting is performed by `split.py`. Results are formatted as follows:
+
+```
+benchmarking/
+βββ idr_prediction/
+ βββ splits/
+ βββ asph/
+ βββ test_df.csv
+ βββ val_df.csv
+ βββ train_df.csv
+ βββ scaled_re/... # same format as splits/asph
+ βββ scaled_rg/... # same format as splits/asph
+ βββ scaling_exp/... # same format as splits/asph
+ βββ test_cluster_split.csv
+ βββ train_cluster_split.csv
+ βββ val_cluster_split.csv
+```
+
+- **`_cluster_split.csv`**: cluster information for the clusters in each split (train, val, test). Columns = "representative seq_id", "member seq_id", "representative seq", "member seq", "member length"
+- π **`asph/`**, **`scaled_re/`**, **`scaled_rg/`**, and **`scaling_exp/`** contain the train, val, and test sets for each property (`train_df.csv`, `val_df.csv`, and `test_df.csv`). The splits follow `_cluster_split.csv`, but not every property has a measurement for each of these sequences. The train-val-test ratio still remains 80-10-10 for each property, despite the sequence losses.
+
+### Training
+The model is defined in `model.py` and `utils.py`. The `train.py` script trains FusOn-pLM-IDR and ESM-2-650M-IDR models *separately for each property* (asphericity, Re, Rg, scaling exponent) with a hyperparameter screen, saves all results separated by property, and makes plots. `plot.py` can be used to regenerate the R2 plots.
+- All **results** are stored in `idr_prediction/results/`, where `timestamp` is a unique string encoding the date and time when you started training.
+- All **raw outputs from models** are stored in `idr_prediction/trained_models/`, where `embedding_path` represents the embeddings used to build the disorder predictor.
+- All **embeddings** made for training will be stored in a new folder called `idr_prediction/embeddings/` with subfolders for each model. This allows you to use the same model multiple times without regenerating embeddings.
+
+Below is the FusOn-pLM-Diso raw outputs folder, `trained_models/fuson_plm/best/`, and the results from the paper, `results/final/`...
+
+The outputs are structured as follows:
+
+```
+benchmarking/
+βββ idr_prediction/
+ βββ results/final/
+ βββ r2_plots
+ βββ asph/
+ βββ esm2_t33_650M_UR50D_asph_R2.png
+ βββ esm2_t33_650M_UR50D_asph_R2_source_data.csv
+ βββ fuson_plm_asph_R2.png
+ βββ fuson_plm_asph_R2_source_data.csv
+ βββ scaled_re/ # same format as r2_plots/asph/...
+ βββ scaled_rg/ # same format as r2_plots/asph/...
+ βββ scaling_exp/ # same format as r2_plots/asph/...
+ βββ asph_best_test_r2.csv
+ βββ asph_hyperparam_screen_test_r2.csv
+ βββ scaled_re_best_test_r2.csv
+ βββ scaled_re_hyperparam_screen_test_r2.csv
+ βββ scaled_rg_best_test_r2.csv
+ βββ scaled_rg_hyperparam_screen_test_r2.csv
+ βββ scaling_exp_best_test_r2.csv
+ βββ scaling_exp_hyperparam_screen_test_r2.csv
+ βββ trained_models/
+ βββ asph/
+ βββ fuson_plm/best/
+ βββ lr0.0001_bs32/
+ βββ asph_r2.csv
+ βββ train_val_losses.csv
+ βββ test_loss.csv
+ βββ asph_test_predictions.csv
+ βββ ... other hyperparameter folders with same format as lr0.001_bs32/
+ βββ esm2_t33_650M_UR50D # same format as asph/fuson_plm/best/
+ βββ scaled_re/ # same format as asph/
+ βββ scaled_rg/ # same format as asph/
+ βββ scaling_exp/ # same format as asph/
+
+```
+
+In both directories, results are organized by IDR property and by the type of embedding used to train FusOn-pLM-IDR.
+
+In the π `results/final` directory.
+- π **`r2_plots//`**: holds all R2 plots and source data (the formatted data used to make the R2 plots) for these properties.
+- **`_best_test_r2.csv`**: holds the R2 values for the top-performing models of each embedding type (e.g. ESM-2-650M and a specific checkpoint of FusOn-pLM)
+- **`_hyperparam_screen_test_r2.csv`**: holds the R2 values for all embedding types, for all screened hyperparaemters
+
+In the π `trained_models` directory:
+- π `/`: holds all results for all trained models predicting this property
+- π `asph/fuson_plm/best/`: holds all FusOn-pLM-IDR results on asphericity prediction for each set of hyperparameters screened when embeddings are made from "fuson_plm/best" (FusOn-pLM model). For example, π `lr0.0001_bs32/` holds results for learning rate of 0.001, batch size 32. If you were to retrain your own checkpoint of fuson_plm and run the IDR prediction benchmark, its results would be stored in a new subfolder of `trained_models/fuson_plm`.
+- **`asph/fuson_plm/best/lr0.0001_bs32/asph_r2.csv`**: R2 value for this set of hyperparameters with "fuson_plm/best" embeddings
+- **`asph/fuson_plm/best/lr0.0001_bs32/asph_test_predictions.csv`**: true asphericity values of the test set proteins, alongside FusOn-pLM-IDR's predictions of them.
+- **`asph/fuson_plm/best/lr0.0001_bs32/test_loss.csv`**: FusOn-pLM-IDR's asphericity test loss value
+- **`asph/fuson_plm/best/lr0.0001_bs32/train_val_losses.csv`**: FusOn-pLM-IDR's tarining and validation loss over each epoch while training on asphericity data
+
+To run the training script, enter:
+```
+nohup python train.py > train.out 2> train.err &
+```
+
+To run the plotting script, enter:
+```
+python plot.py
+```
diff --git a/fuson_plm/benchmarking/idr_prediction/__init__.py b/fuson_plm/benchmarking/idr_prediction/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/fuson_plm/benchmarking/idr_prediction/clean.py b/fuson_plm/benchmarking/idr_prediction/clean.py
new file mode 100755
index 0000000000000000000000000000000000000000..c31b1987aa2693daff14f59e41a075b7e2c78337
--- /dev/null
+++ b/fuson_plm/benchmarking/idr_prediction/clean.py
@@ -0,0 +1,289 @@
+import pandas as pd
+import numpy as np
+import os
+from fuson_plm.utils.logging import open_logfile, log_update
+from fuson_plm.utils.constants import DELIMITERS, VALID_AAS
+from fuson_plm.utils.data_cleaning import check_columns_for_listlike, find_invalid_chars
+from fuson_plm.benchmarking.idr_prediction.plot import plot_all_values_hist_grid, plot_all_train_val_test_values_hist_grid
+
+def process_raw_albatross(df):
+ # return a version of the df with first column split, duplicates cleaned,columns checked for weird characters and invalids
+
+ # first, look at the splits
+ split_str = df['Split'].value_counts().reset_index().rename(columns={'index': 'Split','Split': 'count'})
+ tot_prots = sum(split_str['count'])
+ split_str['pcnt'] = round(100*split_str['count']/tot_prots,2)
+ split_str = split_str.to_string(index=False)
+ split_str = "\t\t" + split_str.replace("\n","\n\t\t")
+ log_update(f"\tTotal proteins: {tot_prots}\n\tSplits:\n{split_str}")
+
+ # format: IDR_19076_tr___A0A8M9PNM5___A0A8M9PNM5_DANRE
+ # or: synth_test_sequence0
+ df['temp'] = df['ID'].str.split("_")
+ df['ID'] = df['temp'].apply(lambda x: f"{x[0]}" if len(x)==1 else f"{x[0]}_{x[1]}" if len(x)<3 else f"{x[0]}_{x[1]}_{x[2]}")
+ # Not ever column has UniProt IDs and Names, so we have to allow np.nan if this info is missing.
+ df['UniProt_ID'] = df['temp'].apply(lambda x: x[5].strip() if len(x)>=5 else np.nan)
+ df['UniProt_Name'] = df['temp'].apply(lambda x: f"{x[8].strip()}_{x[9].strip()}" if len(x)>=8 else np.nan)
+ df = df.drop(columns=['temp'])
+
+ cols_to_check = list(df.columns)
+ cols_to_check.remove('Value') # don't check this one because it shouldn't be string
+ # Investigate the colimns we just created and make sure they don't have any invalid features.
+ # make sure value is float type
+ assert df['Value'].dtype == 'float64'
+ check_columns_for_listlike(df, cols_of_interest=cols_to_check, delimiters=DELIMITERS)
+
+ # Check for invalid AAs
+ df['invalid_chars'] = df['Sequence'].apply(lambda x: find_invalid_chars(x, VALID_AAS))
+ df[df['invalid_chars'].str.len()>0].sort_values(by='Sequence')
+ all_invalid_chars = set().union(*df['invalid_chars'])
+ log_update(f"\tchecking for invalid characters...\n\t\tset of all invalid characters discovered within train_df: {all_invalid_chars}")
+
+ # Assert no invalid AAs
+ assert (df['invalid_chars'].str.len()==0).all()
+ df = df.drop(columns=['invalid_chars'])
+
+ # Check for duplicates - if we find any, REMOVE them from train and keep them in test
+ duplicates = df[df.duplicated('Sequence')]['Sequence'].unique().tolist()
+ n_rows_with_duplicates = len(df[df['Sequence'].isin(duplicates)])
+ log_update(f"\t{len(duplicates)} duplicated sequences, corresponding to {n_rows_with_duplicates} rows")
+
+ # Look for distribution of duplicates WITHIN train, WITHIN test, and BETWEEN train and test
+ # Train only
+ duplicates = df.loc[
+ (df['Split']=='Train')
+ ]
+ duplicates = duplicates[duplicates.duplicated('Sequence')]['Sequence'].unique().tolist()
+ n_rows_with_duplicates = len(df.loc[
+ (df['Sequence'].isin(duplicates)) &
+ (df['Split']=='Train')
+ ])
+ log_update(f"\t\twithin TRAIN only: {len(duplicates)} duplicated sequences, corresponding to {n_rows_with_duplicates} Train rows")
+
+ # Test only
+ duplicates = df.loc[
+ (df['Split']=='Test')
+ ]
+ duplicates = duplicates[duplicates.duplicated('Sequence')]['Sequence'].unique().tolist()
+ n_rows_with_duplicates = len(df.loc[
+ (df['Sequence'].isin(duplicates)) &
+ (df['Split']=='Test')
+ ])
+ log_update(f"\t\twithin TEST only: {len(duplicates)} duplicated sequences, corresponding to {n_rows_with_duplicates} Test rows")
+
+ # Between train and test
+ duplicates_df = df.groupby('Sequence').agg({
+ 'Split': lambda x: ','.join(set(x))
+ }).reset_index()
+ duplicates_df = duplicates_df.loc[duplicates_df['Split'].str.contains(',')].reset_index(drop=True)
+ duplicates = duplicates_df['Sequence'].unique().tolist()
+ n_rows_with_duplicates = len(df[df['Sequence'].isin(duplicates)])
+ log_update(f"\t\tduplicates in BOTH TRAIN AND TEST: {len(duplicates)} duplicated sequences, corresponding to {n_rows_with_duplicates} rows")
+ log_update(f"\t\tprinting portion of dataframe with train+test shared seqs:\n{duplicates_df.head(5)}")
+
+ log_update("\tGrouping by sequence, averaging values, and keeping any Train/Test duplicates in the Test set...")
+ df = df.replace(np.nan, '')
+ df = df.groupby('Sequence').agg(
+ Value=('Value', 'mean'),
+ Value_STD=('Value', 'std'),
+ IDs=('ID', lambda x: ','.join(x)),
+ UniProt_IDs=('UniProt_ID', lambda x: ','.join(x)),
+ UniProt_Names=('UniProt_Name', lambda x: ','.join(x)),
+ Split=('Split', lambda x: ','.join(x))
+ ).reset_index()
+ for col in ['IDs','UniProt_IDs','UniProt_Names','Split']:
+ df[col] = df[col].apply(lambda x: [y for y in x.split(',')])
+ df[col] = df[col].apply(lambda x: ','.join(x))
+ df[col] = df[col].str.strip(',')
+ # make sure there are no commas left
+ assert len(df[df[col].str.contains(',,')])==0
+ # set Split to Test if test is in it
+ df['Split'] = df['Split'].apply(lambda x: 'Test' if 'Test' in x else 'Train')
+
+ # For anything that wasn't duplicated, Value_STD is nan
+ log_update("\tChecking coefficients of variation for averaged rows")
+ # calculate coefficient of variation, should be < 10
+ df['Value_CV'] = 100*df['Value_STD']/df['Value']
+ log_update(f"\t\tTotal rows with coefficient of variation (CV)\n\t\t\t<=10%: {len(df[df['Value_CV']<=10])}\n\t\t\t>10%: {len(df[df['Value_CV']>10])}\n\t\t\t>20%: {len(df[df['Value_CV']>20])}")
+
+ # Ensure there are no duplicates
+ assert len(df[df['Sequence'].duplicated()])==0
+ log_update(f"\tNo remaining duplicates: {len(df[df['Sequence'].duplicated()])==0}")
+
+ # Print the final distribution of train and test values
+ split_str = df['Split'].value_counts().reset_index().rename(columns={'index': 'Split','Split': 'count'})
+ tot_prots = sum(split_str['count'])
+ split_str['pcnt'] = round(100*split_str['count']/tot_prots,2)
+ split_str = split_str.to_string(index=False)
+ split_str = "\t\t" + split_str.replace("\n","\n\t\t")
+ log_update(f"\tTotal proteins: {tot_prots}\n\tSplits:\n{split_str}")
+
+ return df
+
+def combine_albatross_seqs(asph, scaled_re, scaled_rg, scaling_exp):
+ log_update("\nCombining all four dataframes into one file of ALBATROSS sequences")
+
+ asph = asph[['Sequence','Value','IDs','UniProt_IDs','UniProt_Names','Split']].rename(columns={'Value':'asph'})
+ scaled_re = scaled_re[['Sequence','Value','IDs','UniProt_IDs','UniProt_Names','Split']].rename(columns={'Value':'scaled_re'})
+ scaled_rg = scaled_rg[['Sequence','Value','IDs','UniProt_IDs','UniProt_Names','Split']].rename(columns={'Value':'scaled_rg'})
+ scaling_exp = scaling_exp[['Sequence','Value','IDs','UniProt_IDs','UniProt_Names','Split']].rename(columns={'Value':'scaling_exp'})
+
+ combined = asph.merge(scaled_re, on='Sequence',how='outer',suffixes=('_asph', '_scaledre'))\
+ .merge(scaled_rg, on='Sequence',how='outer',suffixes=('_scaledre', '_scaledrg'))\
+ .merge(scaling_exp, on='Sequence',how='outer',suffixes=('_scaledrg', '_scalingexp')).fillna('')
+
+ # Make sure something that's in train for one is in train for all, and not test
+ combined['IDs'] = combined['IDs_asph']+','+combined['IDs_scaledre']+','+combined['IDs_scaledrg']+','+combined['IDs_scalingexp']
+ combined['UniProt_IDs'] = combined['UniProt_IDs_asph']+','+combined['UniProt_IDs_scaledre']+','+combined['UniProt_IDs_scaledrg']+','+combined['UniProt_IDs_scalingexp']
+ combined['UniProt_Names'] = combined['UniProt_Names_asph']+','+combined['UniProt_Names_scaledre']+','+combined['UniProt_Names_scaledrg']+','+combined['UniProt_Names_scalingexp']
+ combined['Split'] = combined['Split_asph']+','+combined['Split_scaledre']+','+combined['Split_scaledrg']+','+combined['Split_scalingexp']
+
+ # Make the lists clean
+ for col in ['IDs','UniProt_IDs','UniProt_Names','Split']:
+ combined[col] = combined[col].apply(lambda x: [y.strip() for y in x.split(',') if len(y)>0])
+ combined[col] = combined[col].apply(lambda x: ','.join(set(x)))
+ combined[col] = combined[col].str.strip(',')
+ # make sure there are no commas left
+ assert len(combined[combined[col].str.contains(',,')])==0
+ combined = combined[['Sequence','IDs','UniProt_IDs','UniProt_Names','Split','asph','scaled_re','scaled_rg','scaling_exp']] # drop unneeded merge relics
+ combined = combined.replace('',np.nan)
+ # Make sure there are no sequences where split is both train and test
+ log_update("\tChecking for any cases where a protein is Train for one IDR prediction task and Test for another (should NOT happen!)")
+ duplicates_df = combined.groupby('Sequence').agg({
+ 'Split': lambda x: ','.join(set(x))
+ }).reset_index()
+ duplicates_df = duplicates_df.loc[duplicates_df['Split'].str.contains(',')].reset_index(drop=True)
+ duplicates = duplicates_df['Sequence'].unique().tolist()
+ n_rows_with_duplicates = len(combined[combined['Sequence'].isin(duplicates)])
+ log_update(f"\t\tsequences in BOTH TRAIN AND TEST: {len(duplicates)} sequences, corresponding to {n_rows_with_duplicates} rows")
+ if len(duplicates)>0:
+ log_update(f"\t\tprinting portion of assert len(combined[combined['asph'].notna()])==len(asph)dataframe with train+test shared seqs:\n{duplicates_df.head(5)}")
+
+ # Now, get rid of duplicates
+ combined = combined.drop_duplicates().reset_index(drop=True)
+ duplicates = combined[combined.duplicated('Sequence')]['Sequence'].unique().tolist()
+ log_update(f"\tDropped duplicates.\n\tTotal duplicate sequences: {len(duplicates)}\n\tTotal sequences: {len(combined)}")
+ assert len(duplicates)==0
+
+ # See how many columns have multiple entries for each
+ log_update(f"\tChecking how many sequences have multiple of the following: ID, UniProt ID, UniProt Name")
+ for col in ['IDs','UniProt_IDs','UniProt_Names','Split']:
+ n_multiple = len(combined.loc[(combined[col].notna()) & (combined[col].str.contains(','))])
+ log_update(f"\t\t{col}: {n_multiple}")
+
+ # See how many entries there are of each cproperty (should match length of original database)
+ assert len(combined[combined['asph'].notna()])==len(asph)
+ assert len(combined[combined['scaled_re'].notna()])==len(scaled_re)
+ assert len(combined[combined['scaled_rg'].notna()])==len(scaled_rg)
+ assert len(combined[combined['scaling_exp'].notna()])==len(scaling_exp)
+ log_update("\tSequences with values for each property:")
+ for property in ['asph','scaled_re','scaled_rg','scaling_exp']:
+ log_update(f"\t\t{property}: {len(combined[combined[property].notna()])}")
+
+ log_update(f"\nPreview of combined database with columns: {combined.columns}\n{combined.head(10)}")
+ return combined
+
+def main():
+ with open_logfile("data_cleaning_log.txt"):
+ # Read in all of the raw data
+ raw_data_folder = 'raw_data'
+ dtype_dict = {0:str,1:str,2:float}
+ rename_dict = {0:'ID',1:'Sequence',2:'Value'}
+
+ # Read in the test data
+ asph_test = pd.read_csv(f"{raw_data_folder}/asph_nat_meth_test.tsv",sep=" ",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+ scaled_re_test = pd.read_csv(f"{raw_data_folder}/scaled_re_nat_meth_test.tsv",sep="\t",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+ scaled_rg_test = pd.read_csv(f"{raw_data_folder}/scaled_rg_nat_meth_test.tsv",sep="\t",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+ scaling_exp_test = pd.read_csv(f"{raw_data_folder}/scaling_exp_nat_meth_test.tsv",sep=" ",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+
+ # Read in the train data
+ asph_train = pd.read_csv(f"{raw_data_folder}/asph_bio_synth_training_data_cleaned_05_09_2023.tsv",sep=" ",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+ scaled_re_train = pd.read_csv(f"{raw_data_folder}/scaled_re_bio_synth_training_data_cleaned_05_09_2023.tsv",sep="\t",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+ scaled_rg_train = pd.read_csv(f"{raw_data_folder}/scaled_rg_bio_synth_training_data_cleaned_05_09_2023.tsv",sep="\t",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+ scaling_exp_train = pd.read_csv(f"{raw_data_folder}/scaling_exp_bio_synth_training_data_cleaned_05_09_2023.tsv",sep=" ",dtype=dtype_dict,header=None).rename(columns=rename_dict)
+
+ # Concatenate - include columns for split
+ asph_test['Split'] = ['Test']*len(asph_test)
+ scaled_re_test['Split'] = ['Test']*len(scaled_re_test)
+ scaled_rg_test['Split'] = ['Test']*len(scaled_rg_test)
+ scaling_exp_test['Split'] = ['Test']*len(scaling_exp_test)
+
+ asph_train['Split'] = ['Train']*len(asph_train)
+ scaled_re_train['Split'] = ['Train']*len(scaled_re_train)
+ scaled_rg_train['Split'] = ['Train']*len(scaled_rg_train)
+ scaling_exp_train['Split'] = ['Train']*len(scaling_exp_train)
+
+ asph = pd.concat([asph_test, asph_train])
+ scaled_re = pd.concat([scaled_re_test, scaled_re_train])
+ scaled_rg = pd.concat([scaled_rg_test, scaled_rg_train])
+ scaling_exp = pd.concat([scaling_exp_test, scaling_exp_train])
+
+ log_update("Initial counts:")
+ log_update(f"\tAsphericity: total entries={len(asph)}, not nan entries={len(asph.loc[asph['Value'].notna()])}")
+ log_update(f"\tScaled re: total entries={len(scaled_re)}, not nan entries={len(scaled_re.loc[scaled_re['Value'].notna()])}")
+ log_update(f"\tScaled rg: total entries={len(scaled_rg)}, not nan entries={len(scaled_rg.loc[scaled_rg['Value'].notna()])}")
+ # change any scaled_rg rows with values less than 1 to np.nan, as done in the paper
+ scaled_rg = scaled_rg.loc[
+ scaled_rg['Value']>=1].reset_index(drop=True)
+ log_update(f"\t\tAfter dropping Rg values <