caid benchmark

fuson_plm/benchmarking/caid/README.md

@@ -0,0 +1,294 @@
+## CAID Benchmark
+
+This folder contains all the data and code needed to perform the **CAID benchmark**, where FusOn-pLM-Diso (a classifier built on FusOn-pLM embeddings) is used to predict per-residue disorder propensities (Figure 4C-F) and plot disorder properties (Figure 1C-1D, S1)
+
+### TL;DR
+The order in which to run the scripts:
+
+```
+python scrape_fusionpdb.py              # pull FusionPDB structures
+python process_fusion_structures.py     # process FusionPDB structures, and head/tail protein structures
+python clean.py                         # clean disorder data and structure data. Assemble train/test/benchmark splits
+python train.py                         # train models 
+python analyze_fusion_preds.py          # make box chart and line plot of model performance on fusion proteins
+python plot.py                          # plot AUROC of model performance, and additional figures based on disorder data
+```
+
+Additional notes: 
+
+* `color_disorder_residues.ipynb` is used to plot fusion structures with pLDDT or disorder prediction color overlays. 
+* We recommend using `nohup` to run longer scripts like `scrape_fusionpdb.py`, `process_fusion_structures.py`, `clean.py`, and `train.py`
+
+### Downloading raw disorder data
+Per-residue disorder predictions were used to train and test FusOn-pLM-Diso.
+
+1. **flDPnn** ([Hu et al. 2021](https://doi.org/10.1038/s41467-021-24773-7))
+    1. At this [link](http://biomine.cs.vcu.edu/servers/flDPnn/?fbclid=IwZXh0bgNhZW0CMTEAAR0KO5CkNdkGC9e5O32S0QoG3BWOw6_egbnioXQNBSv3UC-m_b_dxh70Nnk_aem_z285WFCHdBLw3vOj7LL37A), scroll down to the bottom to find links to the [training](http://biomine.cs.vcu.edu/servers/flDPnn/data/flDPnn_Training_Annotation.txt) and [validation](http://biomine.cs.vcu.edu/servers/flDPnn/data/flDPnn_Validation_Annotation.txt) sets. 
+2. **IDP-CRF** ([Liu et al. 2018](https://doi.org/10.3390/ijms19092483))
+    1. Download zipped data from [this link](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6164615/bin/ijms-19-02483-s001.zip), remove header and footer, and save as a FASTA file 
+3. **CAID2-Disorder-NOX** ([Del Conte et al. 2023](https://doi.org/10.1002/prot.26582))
+    1. Go to [CAID Round 2 Results](https://caid.idpcentral.org/challenge/results?fbclid=IwZXh0bgNhZW0CMTEAAR12dKaA0KywcT71FnyXIrrNS91pwGREsLiq5c2RmfdYl7L0VdUNG7jYai8_aem_tW6Wm9_11ZuiI_GKzbNZjA). Scroll to "Here you can download the references used in the CAID-2 challenge" and you'll find the following links.
+        1. [disorder_nox.fasta](https://caid.idpcentral.org/assets/sections/challenge/static/references/2/disorder_nox.fasta)
+        2. [predictions](https://caid.idpcentral.org/assets/sections/challenge/static/predictions/2/predictions.zip) made by all CAID2 participants; AUROC curves can be reconstructed from these
+
+Raw disorder data are stored in `caid/raw_data`
+```
+benchmarking/
+└── caid/ 
+    └── raw_data/ 
+        └── caid2_competition_results/...
+        └── caid2_train_and_test_data/
+            ├── CAID-2_Disorder_NOX_Testing_Sequences.fasta
+            ├── flDPnn_Training_Dataset.txt
+            ├── flDPnn_Validation_Annotation.txt
+            ├── IDP-CRF_Training_Dataset.txt
+```
+- 📁 **`raw_data/caid2_competition_results/`**: folder containing raw predictions from CAID2 competitors, downloaded directly from the CAID2 website. Models: AlphaFold-disorder, AlphaFOld-rsa, DeepIDP-2L, disomine, DisoPred, DISOPRED3-diso, Dispredict3, ESpritz-D, flDPlr2, flDPnn, flDPnn2, flDPtr, IDP-Fusion, IUPred3.
+- **`raw_data/caid2_train_and_test_data/CAID-2_Disorder_NOX_Testing_Sequences.fasta`**: Disorder-NOX dataset (used as the test set in this benchmark)
+- **`raw_data/caid2_train_and_test_data/flDPnn_Training_Dataset.txt`**: training set for flDPnn
+- **`raw_data/caid2_train_and_test_dataflDPnn_Validation_Dataset.txt`**: validation set for flDPnn
+- **`raw_data/IDP-CRF_Training_Dataset.txt`**: training set for IDP-CRF
+
+### Processing disorder data
+
+```
+benchmarking/
+└── caid/ 
+    └── processed_data/
+        └── caid2_competition_results/...
+        ├── CAID-2_Disorder_NOX_Processed.csv
+        ├── flDPnn_Training_Dataset.csv
+        ├── flDPnn_Validation_Dataset.csv
+        ├── IDP-CRF_Training_Dataset.csv
+    └── splits/
+        ├── splits.csv
+        ├── train_df.csv 
+        ├── test_df.csv
+        ├── fusion_bench_df.csv
+```
+
+The   **`clean.py`** processes and combines the raw data files, generating the following files in 📁`processed_data/`:
+- 📁 **`caid2_competition_results/`**: a folder with table versions of all the files in 📁 `raw_data/caid2_competition_results/`
+- **`CAID-2_Disorder_NOX_Processed.csv`**: a table of test data, made by parsing `raw_data/caid2_train_and_test_data/CAID-2_Disorder_NOX_Testing_Sequences.fasta`
+- **`flDPnn_Training_Dataset.csv`**: a table of flDPnn's training data, made by parsing `raw_data/caid2_train_and_test_data/flDPnn_Training_Dataset.txt`
+- **`flDPnn_Validation_Dataset.csv`**: a table of flDPnn's validation data, made by parsing `raw_data/caid2_train_and_test_data/flDPnn_Validation_Dataset.txt`
+- **`IDP-CRF_Training_Dataset.csv`**: a table of IDP-CRF's training data, made by parsing `raw_data/caid2_train_and_test_data/CRF_Training_Dataset.txt`
+
+`clean.py` also generates **the final train-test splits and fusion oncoprotein benchmarking file used to train and evaluate the disorder predictors.** These are stored in 📁`splits/`
+- **`splits.csv`**: sequences, IDs, split (either "Train", "Test", or "Fusion_Benchmark"), andpper-residue disorder labels based on AlphaFold-pLDDT (1 (disordered) if pLDDT< 68.8, 0 (ordered) if >=68.8)
+- **`train_df.csv`**: just the Train set portion of `splits.csv`
+- **`test_df.csv`**: just the Test set portion of `splits.csv`
+- **`fusion_bench_df.csv`**: just the Fusion_Benchmark portion of `splits.csv`. Includes 524 fusion oncoproteins from the FusOn-pLM test set whose structures were collected from FusionPDB (see "Downloading and Processing FusionPDB data
+
+
+### Downloading and Processing FusionPDB data
+The structures of fusion oncoproteins from the FusionPDB database were used to evaluate FusOn-pLM-Diso's performance on fusion oncoproteins. This data was collected by running `scrape_fusionpdb.py`, followed by `process_fusion_structures.py`. These scripts populated the `raw_data` and `processed_data` files simultaneously.
+
+Listed below are all the relevant files: 
+
+```
+benchmarking/
+└── caid/ 
+    └── raw_data/ 
+        └── fusionpdb/
+            └── structures/... # created by scrape_fusionpdb.py (folder not included in repo)
+            └── head_tail_af2db_structures/... # created by process_fusion_structures.py (folder not included in repo)
+            ├── FusionPDB_level2_curated_09_05_2024.csv
+            ├── FusionPDB_level2_fusion_structure_links.csv
+            ├── FusionPDB_level3_curated_09_05_2024.csv
+            ├── FusionPDB_level3_fusion_structure_links.csv
+            ├── fusionpdb_structureless_ids.txt
+            ├── hgene_tgene_uniprot_idmap_07_10_2024.txt
+            ├── level2_head_tail_info.txt
+            ├── level3_head_tail_info.txt
+            ├── not_in_afdb_idmap.txt
+    └── processed_data/
+        └── fusion_pdb/
+            └── intermediates/
+                ├── giant_level_2-3_fusion_protein_head_tail_info.csv
+                ├── giant_level2-3_fusion_protein_structure_links.csv
+                ├── giant_level2-3_fusion_protein_structures_processed.csv
+                ├── uniprotids_not_in_afdb.txt
+                ├── unmapped_parts.tt
+            ├── fusion_heads_and_tails.csv
+            ├── FusionPDB_level2-3_cleaned_FusionGID_info.csv
+            ├── FusionPDB_level2-3_cleaned_structure_info.csv
+            ├── heads_tails_structural_data.csv
+```
+
+#### ⚙️ Pipeline
+
+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`
+
+<br> 
+
+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"
+
+
+### Training
+
+The model is defined in `model.py` and `utils.py`. Training configs can be provided in `config.py`:
+
+```
+# Which models to benchmark
+BENCHMARK_FUSONPLM = True
+# FUSONPLM_CKPTS. If you've traiend your own model, this is a dictionary: key = run name, values = epochs
+# If you want to use the trained FusOn-pLM, instead FUSONPLM_CKPTS="FusOn-pLM"
+FUSONPLM_CKPTS= "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.
+```
+
+<br>
+
+`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):
+
+```
+benchmarking/
+└── caid/ 
+    └── trained_models/ 
+        └── esm2_t33_650M_UR50D/best/
+        └── fuson_plm/best/
+            ├── caid_hyperparam_screen_fusion_benchmark_metrics.csv
+            ├── caid_hyperparam_screen_fusion_benchmark_probs.csv
+            ├── caid_hyperparam_screen_test_metrics.csv
+            ├── caid_hyperparam_screen_test_probs.csv
+            ├── caid_train_losses.csv
+            ├── params.txt
+```
+
+- **`caid_hyperparam_screen_fusion_benchmark_metrics.csv`**: performance metrics (Accuracy, Precision, Recall, F1 Score, AUROC) for the top model on the fusion benchmark set (`splits/fusion_bench_df.csv`)
+- **`caid_hyperparam_screen_fusion_benchmark_probs.csv`**: for the fusion benchmark, raw probabilities of class 1 (disorder), threshold used to assign 0/1 based on maximized F1 score, prediction labels based on probabilities and threshold 
+- **`caid_hyperparam_screen_test_metrics.csv`**: same as `caid_hyperparam_screen_fusion_benchmark_metrics.csv`, but for CAID2 Disorder-NOX (`splits/test_df.csv`)
+- **`caid_hyperparam_screen_test_probs.csv`**: same as `caid_hyperparam_screen_fusion_benchmark_probs`, but for CAID2 Disorder-NOX
+- **`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`
+
+```
+benchmarking/
+└── caid/ 
+    └── results/final
+        ├── best_caid_model_results.csv 
+        ├── caid_hyperparam_screen_test_metrics.csv
+        ├── caid_hyperparam_screen_fusion_benchmark_metrics.csv
+        ├── caid_hyperparam_screen_train_losses.csv
+        ├── fusion_disorder_boxplots.png
+        ├── fusion_pred_disorder_r2.png
+        ├── fusion_disorder_boxplots_source_data.csv
+        ├── fusion_pred_disorder_r2_source_data.csv
+        ├── CAID2_FusOn-pLM-Diso_with_ESM_AUROC_curve.png   
+        ├── CAID_fpr_tpr_source_data.csv 
+        ├── CAID_prediction_source_data.csv
+```
+
+- **`best_caid_model_results.csv`**: Summary file of hyperparameters, test set statistics, and fusion benchmark statistics for the best model of each type screened (ESM-2-650M, FusOn-pLM)
+- **`caid_hyperparam_screen_fusion_benchmark_metrics.csv`**: Fusion benchmark set statistics for full hyperparameter screen
+- **`caid_hyperparam_screen_fusion_benchmark_metrics.csv`**: Test set statistics for full hyperparameter screen 
+- **`caid_hyperparam_screen_train_losses.csv`**: Train losses for full hyperparameter screen
+- 📊 **`fusion_disorder_boxplots.png`**: Fig. 4E, left (data directly used to produce the plot at `fusion_disorder_boxplots_source_data.csv`)
+- 📊 **`fusion_pred_disorder_r2_source_data.csv`**: Fig. 4E, right (data directly used to produce the plot at `fusion_pred_disorder_r2_source_data.csv`)
+- 📊 **`CAID2_FusOn-pLM-Diso_with_ESM_AUROC_curve.png`**: Fig. 4D (probabilities used at `CAID_prediction_source_data.csv`, FPR/TPR relationships directly used to make the plot at `CAID_fpr_tpr_source_data.csv`)
+
+To run the training script, use
+
+```
+nohup python train.py > train.out 2> train.err &
+```
+
+### Plotting
+
+The `plot.py` script generates many figures from the paper, alongside the formatted data directly used for plotting. 
+
+```
+benchmarking/
+└── caid/ 
+    └── results/final/
+        ├── CAID2_FusOn-pLM-Diso_with_ESM_AUROC_curve.png
+    └── processed_data/ 
+        └── figures/ 
+            └── fusion_disorder/ 
+                ├── plddt_sequence_EML4-ALK.png 
+                ├── plddt_sequence_EML4::ALK_source_data.csv
+                ├── plddt_sequence_EWSR1-FLI1.png 
+                ├── plddt_sequence_EWSR1::FLI1_source_data.csv
+                ├── plddt_sequence_PAX3-FOXO1.png 
+                ├── plddt_sequence_PAX3::FOXO1_source_data.csv
+                ├── plddt_sequence_SS18-SSX1.png
+                ├── plddt_sequence_SS18::SSX1_source_data.csv
+            └── histograms/ 
+                ├── disorder_nox_histogram.png 
+                ├── disorder_nox_histogram_source_data.csv
+                ├── fusions_histogram.png  
+                ├── fusions_histogram_source_data.csv
+                ├── heads_histogram.png  
+                ├── heads_histogram_source_data.csv
+                ├── tails_histogram.png
+                ├── tails_histogram_source_data.csv
+```
+
+- Plots in `fusion_disorder` are from Fig. 1C
+- Plots in `hisograms` are from Fig. 1D and Fig. S1
+
+To regenerate these plots and source data, run: 
+
+```
+python plot.py
+```
+
+### Colored structure images
+`color_disorder_residues.ipynb` is used to plot fusion structures with pLDDT or disorder prediction color overlays. By running certain (or all) of its cells, you will recreate images from Fig. 1C and 4F, as well as the following file: 
+
+```
+benchmarking/
+└── caid/ 
+    └── disorder_coloring_data
+        ├── normalized_disorder_propensities_source_data.csv
+```
+- **`normalized_disorder_propensities_source_data.csv`**: the normalized disorder propensities that were visualized on fusion structures in Fig. 4F

fuson_plm/benchmarking/caid/analyze_fusion_preds.py

@@ -0,0 +1,158 @@
+
+import torch
+import torch.nn as nn
+import os
+import pickle
+import pandas as pd
+import numpy as np
+
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, precision_recall_curve, average_precision_score
+from fuson_plm.utils.logging import log_update, open_logfile
+from fuson_plm.benchmarking.caid.plot import plot_fusion_stats_boxplots, plot_fusion_frac_disorder_r2
+
+# calculate AUROC and AUPRC for each sequence
+def calc_metrics(row):
+    probs = row['prob_1']
+    probs = [float(y) for y in probs.split(',')]
+    true_labels = row['Label']
+    true_labels = [int(y) for y in list(true_labels)]
+    pred_labels = row['pred_labels']
+    pred_labels = [int(y) for y in list(pred_labels)]
+    
+    # Calculate AUROC
+    # Calculate AUPRC
+    # Calculate all the other stats based on the predicted labels
+    
+    flat_binary_preds = np.array(pred_labels)
+    flat_prob_preds = np.array(probs)
+    flat_labels = np.array(true_labels)
+
+    accuracy = accuracy_score(flat_labels, flat_binary_preds)
+    precision = precision_score(flat_labels, flat_binary_preds)
+    recall = recall_score(flat_labels, flat_binary_preds)
+    f1 = f1_score(flat_labels, flat_binary_preds)
+    try:
+        roc_auc = roc_auc_score(flat_labels, flat_prob_preds)
+    except:
+        roc_auc = np.nan
+        
+    try:
+        auprc = average_precision_score(flat_labels, flat_prob_preds)
+    except: 
+        auprc = np.nan
+    
+    return pd.Series({
+        'Accuracy': round(accuracy,3),
+        'Precision': round(precision,3),
+        'Recall': round(recall,3),
+        'F1': round(f1,3),
+        'AUROC': round(roc_auc,3) if not(np.isnan(roc_auc)) else roc_auc,
+        'AUPRC': round(auprc,3) if not(np.isnan(auprc)) else auprc,
+    })
+
+def get_model_preds_with_metrics(path_to_model_predictions):
+    # Define paths and dataframes that we will need 
+    fusion_benchmark_set = pd.read_csv('splits/fusion_bench_df.csv')
+    model_predictions = pd.read_csv(path_to_model_predictions)
+    fusion_structure_data = pd.read_csv('processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv')
+    fusion_structure_data['Fusion_Structure_Link'] = fusion_structure_data['Fusion_Structure_Link'].apply(lambda x: x.split('/')[-1])
+    
+    # merge fusion data with seq ids 
+    fuson_db = pd.read_csv('../../data/fuson_db.csv')
+    fuson_db = fuson_db[['aa_seq','seq_id']].rename(columns={'aa_seq':'Fusion_Seq'})
+    fusion_structure_data = pd.merge(
+        fusion_structure_data,
+        fuson_db,
+        on='Fusion_Seq',
+        how='inner'
+    )
+
+    # merge fusion structure data with top swissprot alignments
+    swissprot_top_alignments = pd.read_csv("../../data/blast/blast_outputs/swissprot_top_alignments.csv")
+    fusion_structure_data = pd.merge(
+        fusion_structure_data,
+        swissprot_top_alignments,
+        on="seq_id",
+        how="left"
+    )
+
+    model_predictions_labeled = pd.merge(model_predictions,fusion_benchmark_set.rename(columns={'Sequence':'sequence'}),on='sequence',how='inner')
+    model_predictions_labeled = pd.merge(model_predictions_labeled, 
+                                        fusion_structure_data[['FusionGene','Fusion_Seq','Fusion_Structure_Link','Fusion_pLDDT','Fusion_AA_pLDDTs',
+                                                                'top_hg_UniProtID', 'top_hg_UniProt_isoform', 'top_hg_UniProt_fus_indices', 'top_tg_UniProtID', 'top_tg_UniProt_isoform', 
+                                                                'top_tg_UniProt_fus_indices', 'top_UniProtID', 'top_UniProt_isoform', 'top_UniProt_fus_indices', 'top_UniProt_nIdentities', 
+                                                                'top_UniProt_nPositives']].rename(
+                                            columns={'Fusion_Seq': 'sequence'}
+                                        ),
+                                        on='sequence',
+                                        how='left')
+    model_predictions_labeled['length'] = model_predictions_labeled['sequence'].str.len()
+    model_predictions_labeled['Fusion_Structure_Link'] = model_predictions_labeled['Fusion_Structure_Link'].apply(lambda x: x.split('/')[-1])
+
+    model_predictions_labeled[['Accuracy','Precision','Recall','F1','AUROC','AUPRC']] = model_predictions_labeled.apply(lambda row: calc_metrics(row),axis=1)
+    model_predictions_labeled = model_predictions_labeled.sort_values(by=['AUROC','F1','AUPRC','Accuracy','Precision','Recall'],ascending=[False,False,False,False,False,False]).reset_index(drop=True)
+    model_predictions_labeled['pcnt_disordered'] = round(100*model_predictions_labeled['Label'].apply(lambda x: sum([int(y) for y in list(x)]))/model_predictions_labeled['sequence'].str.len(),2)
+    model_predictions_labeled['pred_pcnt_disordered'] = round(100*model_predictions_labeled['pred_labels'].apply(lambda x: sum([int(y) for y in list(x)]))/model_predictions_labeled['sequence'].str.len(),2)
+    log_update(f"Model predictions for {len(model_predictions_labeled)} fusion oncoproteins. Preview:")
+    log_update(
+        model_predictions_labeled[['sequence','length','FusionGene','Fusion_pLDDT','pcnt_disordered','pred_pcnt_disordered','AUROC','F1','AUPRC','Accuracy','Precision','Recall']].head()
+    )
+    cols_str = '\n\t'+ '\n\t'.join(list(model_predictions_labeled.columns))
+    log_update(f"Columns in model predictions stats database: {cols_str}")
+    
+    # There is one duplicate row
+    duplicate_sequences = model_predictions_labeled.loc[model_predictions_labeled['sequence'].duplicated()]['sequence'].unique().tolist()
+    log_update(f"\nTotal duplicate sequences: {len(duplicate_sequences)}")
+    gb = model_predictions_labeled.groupby('sequence').agg(
+        pred_labels=("pred_labels", list),
+    )
+    gb['pred_labels'] = gb['pred_labels'].apply(lambda x: list(set(x)))
+    gb['unique_label_vectors'] = gb['pred_labels'].apply(lambda x: len(x))
+    log_update(f"Duplicate entries for sequences have the exact same label vector: {(gb['unique_label_vectors']==1).all()}")
+    log_update("\tSince above statement is true, randomly dropping duplicate sequence rows - should make no difference to prediction.")
+    
+    model_predictions_labeled = model_predictions_labeled.drop_duplicates('sequence').reset_index(drop=True)
+    #os.makedirs("processed_data/fusion_predictions",exist_ok=True)
+    return model_predictions_labeled
+
+def calc_average_stats(model_pred_stats):
+    # cols: Accuracy	Precision	Recall	F1	AUROC	AUPRC	pcnt_disordered	pred_pcnt_disordered
+    averages = model_pred_stats[[
+        'Accuracy', 'Precision', 'Recall', 'F1', 'AUROC', 'AUPRC'
+    ]].mean()
+    averages
+
+def main():
+    with open_logfile("analyze_fusion_preds.txt"):
+        ## Put path to model predictions you'd like to use for benchmarking
+        path_to_model_predictions = "trained_models/fuson_plm/best/caid_hyperparam_screen_fusion_benchmark_probs.csv"
+        save_dir = "results/final"
+        preds_with_metrics_save_path = f"{save_dir}/model_preds_with_metrics.csv"
+        boxplot_save_path = f"{save_dir}/fusion_disorder_boxplots.png"
+        r2_save_path = "results/final/fusion_pred_disorder_r2.png"
+        
+        # Additional benchmarking on fusion predictions
+        fuson_db = pd.read_csv("../../data/fuson_db.csv")
+        seq_id_dict = dict(zip(fuson_db['aa_seq'],fuson_db['seq_id']))
+        model_preds_with_metrics = get_model_preds_with_metrics(path_to_model_predictions)
+        model_preds_with_metrics['seq_id'] = model_preds_with_metrics['sequence'].map(seq_id_dict)
+        model_preds_with_metrics.to_csv(preds_with_metrics_save_path,index=False)
+        print(model_preds_with_metrics.columns)
+        
+        # Box plot
+        boxplot_model_preds = model_preds_with_metrics[['seq_id','FusionGene',
+        'Accuracy', 'Precision', 'Recall', 'F1', 'AUROC'
+        ]]
+        
+        boxplot_model_preds.to_csv(boxplot_save_path.replace(".png","_source_data.csv"),index=False)
+        plot_fusion_stats_boxplots(boxplot_model_preds[['Accuracy', 'Precision', 'Recall', 'F1', 'AUROC'
+        ]], save_path=boxplot_save_path)
+        
+        # R2 plot
+        r2_model_preds = model_preds_with_metrics[['seq_id','FusionGene','pcnt_disordered','pred_pcnt_disordered']]
+        r2_model_preds.to_csv(r2_save_path.replace(".png","_source_data.csv"),index=False)
+        plot_fusion_frac_disorder_r2(r2_model_preds['pcnt_disordered'], r2_model_preds['pred_pcnt_disordered'], save_path=r2_save_path)
+        calc_average_stats(model_preds_with_metrics)
+
+if __name__ == "__main__":
+    main()

fuson_plm/benchmarking/caid/clean.py

@@ -0,0 +1,671 @@
+import os
+import numpy as np
+import re
+import pandas as pd
+import requests
+
+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.caid.scrape_fusionpdb import scrape_fusionpdb_level_2_3
+from fuson_plm.benchmarking.caid.process_fusion_structures import process_fusions_and_hts
+            
+def download_fasta(uniprotid, includeIsoform, output_file):
+    try:
+        url = f"https://rest.uniprot.org/uniprotkb/search?format=fasta&includeIsoform={includeIsoform}&query=accession%3A{uniprotid}&size=500&sort=accession+asc"
+        # Send a GET request to the URL
+        response = requests.get(url)
+        
+        # Raise an exception if the request was unsuccessful
+        response.raise_for_status()
+        
+        # Write the content to a file in text mode
+        with open(output_file, 'a+') as file:
+            file.write(response.text)
+        
+        log_update(f"FASTA file for {uniprotid} successfully downloaded and added to '{output_file}'")
+    
+    except requests.exceptions.RequestException as e:
+        log_update(f"An error occurred: {e}")
+
+# Test Sequences (CAID-2 Disorder-NOX)
+def parse_caid_txt(fast_file):
+    '''
+    Parses correctly fasta-formatted text file with conditions:
+    Line 1: ID
+    Line 2: Sequence
+    Line 3: Label
+    '''
+
+    seq_to_label = {}
+    id_to_sequence = {}
+
+    with open(fast_file, 'r') as file:
+        label = None
+        sequence = ""
+        seq_id = None
+        reading_sequence = False
+        for line in file:
+            line = line.strip()
+            if line.startswith(">"):
+                if label is not None and sequence:
+                    seq_to_label[sequence] = (label, seq_id)
+                seq_id = line[1:]  # Capture the ID without the '>'
+                label = None
+                sequence = ""
+                reading_sequence = True
+            elif reading_sequence:
+                if all(c in "01-" for c in line):
+                    label = line
+                    reading_sequence = False
+                else:
+                    sequence += line
+        if label is not None and sequence:
+            seq_to_label[sequence] = (label, seq_id)
+
+    return seq_to_label
+
+def check_df_for_mismatched_labels(sd):
+    log_update("\tChecking dataframe for mismatched sequences and labels...")
+    counter=0
+    for idx, row in sd.iterrows():
+        seq = row['Sequence']
+        label = row['Label']
+
+        if len(seq) != len(label):
+            counter+=1
+            log_update(f"\t\tLength mismatch at index {idx}: sequence length = {len(seq)}, label length = {len(label)}")
+            
+    log_update(f"\t\tTotal mismatched lengths/labels: {counter}")
+
+
+def process_caid2_disorder_nox_test(caid_path):
+    """
+    Processes the CAID-2_Disorder_NOX_Testing_Sequences.fasta file
+    """
+    log_update("Processing CAID-2-Disorder-NOX Testing Dataset")
+    # Parse the fasta file
+    caid_dict = parse_caid_txt(caid_path)
+    
+    # Gather the sequences
+    caid_seqs = {}
+    for k, (v, seq_id) in caid_dict.items():
+        caid_seqs[seq_id] = (k, v)
+    log_update(f"\tTotal sequences: {len(caid_seqs)}")
+
+    # Form dataframe from processed data
+    caid_df = pd.DataFrame({
+        'ID': list(caid_seqs.keys()),
+        'Sequence': [seq for seq, _ in caid_seqs.values()],
+        'Label': [lbl for _, lbl in caid_seqs.values()],
+        'Split': 'Test'
+    })
+    
+    check_df_for_mismatched_labels(caid_df)
+    return caid_df
+    
+# Training Squences (fldpnn and IDP-CRF)
+# fldpnn Training Sequences
+def parse_fldpnn_fasta(file_path):
+    """
+    Parse flDPnn_Training_Dataset.txt, where there are 5 sequence lines. We only want the first 
+
+    >Disprot ID
+    Amino acid sequence
+    Experimental annotation for intrinsic disorder
+    Experimental annotation for disordered protein binding
+    Experimental annotation for disordered DNA binding
+    Experimental annotation for disordered RNA binding
+    Experimental annotation for disordered flexible linkers
+    """
+    sequences = []
+    labels = []
+    ids = []
+
+    with open(file_path, 'r') as file:
+        lines = file.readlines()
+
+        seq_id = ""
+        current_sequence = ""
+        seen_label_lines = 0    # should go up to 5 for each
+        current_labels = []
+        is_label = False
+
+        for line in lines:
+            line = line.strip()
+            if line.startswith('>'):
+                if current_sequence and current_labels:
+                    assert seen_label_lines==5  # we should've seen 5 labels, otherwise something is wrong
+                    ids.append(seq_id)
+                    sequences.append(current_sequence)
+                    labels.append(''.join(current_labels))
+                seq_id = line[1:]  # Capture the ID without the '>'
+                current_sequence = ""
+                current_labels = []
+                is_label = False
+                seen_label_lines = 0
+            elif re.match('^[A-Z]+$', line):  # Sequence lines
+                current_sequence += line
+            else:  # Label lines
+                seen_label_lines+=1
+                if seen_label_lines==1:
+                    current_labels.append(line)
+                is_label = True
+
+        # Add the last sequence and labels
+        if current_sequence and current_labels:
+            sequences.append(current_sequence)
+            labels.append(''.join(current_labels))
+            ids.append(seq_id)
+
+    return ids, sequences, labels
+
+def parse_idp_crf_fasta(file_path):
+    sequences = []
+    labels = []
+    ids = []
+
+    with open(file_path, 'r') as file:
+        lines = file.readlines()
+
+        seq_id = ""
+        current_sequence = ""
+        current_labels = []
+        is_label = False
+
+        for line in lines:
+            line = line.strip()
+            if line.startswith('>'):
+                if current_sequence and current_labels:
+                    ids.append(seq_id)
+                    sequences.append(current_sequence)
+                    labels.append(''.join(current_labels))
+                seq_id = line[1:]  # Capture the ID without the '>'
+                current_sequence = ""
+                current_labels = []
+                is_label = False
+            elif re.match('^[A-Z]+$', line):  # Sequence lines
+                current_sequence += line
+            else:  # Label lines
+                current_labels.append(line)
+                is_label = True
+
+        # Add the last sequence and labels
+        if current_sequence and current_labels:
+            sequences.append(current_sequence)
+            labels.append(''.join(current_labels))
+            ids.append(seq_id)
+
+    return ids, sequences, labels
+
+def process_fldpnn(fldpnn_path, split="training"):
+    """
+    Process the fldpnn_Training_Dataset
+    """
+    log_update(f"\nProcessing flDPnn {split} dataset")
+    # Parse fasta
+    fldpnn_ids, fldpnn_seqs, fldpnn_labels = parse_fldpnn_fasta(fldpnn_path)
+
+    # Collect cleaned labels
+    cleaned_fldpnn_ids = []
+    cleaned_fldpnn_labels = []
+    for i in range(len(fldpnn_seqs)):
+        seq_len = len(fldpnn_seqs[i])
+        label = fldpnn_labels[i]      # Should only be the first set of labels
+        id = fldpnn_ids[i]
+        cleaned_fldpnn_labels.append(label)
+        
+    log_update(f"\tTotal labels: {len(cleaned_fldpnn_labels)}, total sequences: {len(fldpnn_seqs)},total IDs: {len(fldpnn_ids)}")
+
+    fldpnn_df = pd.DataFrame({'Sequence': fldpnn_seqs, 
+                              'Label': cleaned_fldpnn_labels, 
+                              "Split": "Train" if split=="training" else "Val",
+                              "ID": fldpnn_ids})
+    check_df_for_mismatched_labels(fldpnn_df)
+    
+    return fldpnn_df
+
+def combine_fldpnn_train_val(fldpnn_train_df, fldpnn_val_df):
+    log_update("\nJoining flDPnn train and val sets into one training set for CAID predictor")
+    combined = pd.concat([fldpnn_train_df,fldpnn_val_df])
+    
+    # check for duplicates
+    duplicates = combined[combined['Sequence'].duplicated()]['Sequence'].unique().tolist()
+    n_rows_with_duplicates = len(combined[combined['Sequence'].isin(duplicates)])
+    log_update(f"\t{len(duplicates)} sequences in both train and val datasets, corresponding to {n_rows_with_duplicates} rows")
+    for dup in duplicates:
+        train_id = combined.loc[(combined['Sequence']==dup) & (combined['Split']=='Train')]['ID'].reset_index(drop=True).iloc[0]
+        val_id = combined.loc[(combined['Sequence']==dup) & (combined['Split']=='Val')]['ID'].reset_index(drop=True).iloc[0]
+        train_label = combined.loc[(combined['Sequence']==dup) & (combined['Split']=='Train')]['Label'].reset_index(drop=True).iloc[0]
+        val_label = combined.loc[(combined['Sequence']==dup) & (combined['Split']=='Val')]['Label'].reset_index(drop=True).iloc[0]
+        log_update(f"\t\tTrain ID: {train_id}\tVal ID: {val_id}\tSame labels: {train_label==val_label}\tSequence: {dup}")
+        
+        # if the labels are not equal, get rid of it completely. Otherwise just get rid of the val case
+        if not(train_label==val_label):
+            log_update(f"\t\t\tSince labels are not equal, removing sequence completely")
+            combined = combined[combined['Sequence']!=dup].reset_index(drop=True)
+        else:
+            log_update(f"\t\t\tSince labels are equal, removing validation copy")
+            combined = combined.loc[(combined['Sequence']!=dup) |
+                                    ((combined['Sequence']==dup) & (combined['Split']=='Train'))]
+    # drop duplicates
+    log_update(f"\tLength of joined flDPnn data: {len(combined)}")
+    
+    return combined
+    
+def process_idp_crf_train(idp_crf_train_path):
+    """
+    Process IDP-CRF_Training_Dataset
+    
+    Args:
+        idp_crf_train_path
+    """
+    log_update("\nProcessing IDP-CRF training dataset")
+    # Parse the fasta, get sequences and labels
+    idp_crf_ids, idp_crf_seqs, idp_crf_labels = parse_idp_crf_fasta(idp_crf_train_path)
+    log_update(f"\tTotal labels: {len(idp_crf_labels)}, total sequences: {len(idp_crf_seqs)}, total IDs: {len(idp_crf_ids)}")
+
+    # Clean the labels
+    cleaned_idp_ids, cleaned_idp_seqs, cleaned_idp_labels = [], [], []
+    counter = 0
+    log_update("\tCleaning labels and counting length-mismatched examples...")
+    for i, label in enumerate(idp_crf_labels):
+        # If length of sequence and labels doesn't match, log it
+        if len(idp_crf_seqs[i]) != len(idp_crf_labels[i]):
+            log_update(f"\t\tLength mismatch at index {i}: sequence length = {len(idp_crf_seqs[i])}, label length = {len(idp_crf_labels[i])}")
+            
+            counter += 1
+        # Else, "clean" the labels by mapping them to ints and converting them to a list 
+        else:
+            cleaned_idp_ids.append(idp_crf_ids[i])
+            cleaned_idp_labels.append(label)
+            cleaned_idp_seqs.append(idp_crf_seqs[i])
+
+    log_update(f"\t\tMismatched lengths/labels: {counter}")
+
+    # Confirm that final database has no mismatched labels
+    idp_crf_df = pd.DataFrame({'Sequence': cleaned_idp_seqs, 
+                               'Label': cleaned_idp_labels, 
+                               "Split": "Train",
+                               "ID": cleaned_idp_ids})
+    check_df_for_mismatched_labels(idp_crf_df)
+            
+    return idp_crf_df
+
+def find_agreeing_labels(row, lab1="", lab2=""):
+    """
+    If there's only one possible label, return that label. If the two labels disagree, return np.nan
+    """
+    val1 = row[lab1]
+    val2 = row[lab2]
+    
+    # If one of them is nan, then they won't match anyway, so return True because there is no conflict
+    if type(val1)==float and np.isnan(val1):
+        return val2
+    elif type(val2)==float and np.isnan(val2):
+        return val1
+    else:
+        if val1==val2:
+            return val1
+        else:
+            return np.nan
+
+def get_unique_ids(row):
+    source_to_id = {
+        "IDP-CRF": row["IDP-CRF ID"],
+        "flDPnn": row["flDPnn ID"],
+        "CAID-2_Disorder_NOX": row["CAID-2_Disorder_NOX ID"]
+    }
+    
+    all_sources = row["Source"].split(",")
+    all_ids = []
+    # they are already in the desired order so just iterate through them
+    for source in all_sources:
+        candidate_id = source_to_id[source]
+        if not(candidate_id in all_ids):
+            all_ids.append(candidate_id)
+    
+    return ",".join(all_ids)
+
+def parse_caid2_results(processed_caid2_df,lines):
+    # iterate through the lines
+    all_caid2_disorder_nox_ids = processed_caid2_df['ID'].tolist()
+    all_caid2_disorder_nox_sequences = processed_caid2_df['Sequence'].tolist()
+    
+    cur_id = None
+    results = {
+        }
+    for i, line in enumerate(lines):
+        # If line starts with >, that means we have a new ID
+        if line[0]==">":
+            # If we are currently on a different cur_id, finish that one out
+            if not(cur_id is None):
+                results[cur_id]['prob_1'] = ",".join(results[cur_id]['prob_1'])
+                results[cur_id]['pred_labels'] = ",".join(results[cur_id]['pred_labels'])
+                sequence = results[cur_id]['sequence']
+                # get the true labels from the CAID2 dataset - IF POSSIBLE
+                if (cur_id not in all_caid2_disorder_nox_ids) and (sequence not in all_caid2_disorder_nox_sequences):
+                    results[cur_id]['labels'] = np.nan
+                else: 
+                    true_labels = processed_caid2_df.loc[
+                        processed_caid2_df['ID']==cur_id,'Label'
+                    ].item()
+                    true_labels = ",".join(list(true_labels))
+                    results[cur_id]['labels'] = true_labels
+            # Now process the new one
+            cur_id = line[1::].strip('\t').strip('\n')
+            results[cur_id] = {
+                'sequence': '',
+                'prob_1': [],
+                'pred_labels': []
+            }
+        # if cur id is not None
+        else:
+            # if we have a cur id to process, process it
+            if not(cur_id is None):
+                # Extract the information - not every .caid file as predicted labels!!
+                lsplit =  line.strip('\n').split('\t')
+                label=''
+                idx, aa, prob = lsplit[0], lsplit[1], lsplit[2]
+                if len(lsplit)==4: label=lsplit[3]
+                # Add to dict
+                results[cur_id]['sequence']+=aa
+                results[cur_id]['prob_1'].append(prob)
+                results[cur_id]['pred_labels'].append(label)
+
+            # if we're on the last line, combine
+            if i==len(lines)-1:
+                results[cur_id]['prob_1'] = ",".join(results[cur_id]['prob_1'])
+                results[cur_id]['pred_labels'] = ",".join(results[cur_id]['pred_labels'])
+                sequence = results[cur_id]['sequence']
+                # get the true labels from the CAID2 dataset - IF POSSIBLE
+                if (cur_id not in all_caid2_disorder_nox_ids) and (sequence not in all_caid2_disorder_nox_sequences):
+                    results[cur_id]['labels'] = np.nan
+                else: 
+                    true_labels = processed_caid2_df.loc[
+                        processed_caid2_df['ID']==cur_id,'Label'
+                    ].item()
+                    true_labels = ",".join(list(true_labels))
+                    results[cur_id]['labels'] = true_labels
+                
+    df = pd.DataFrame.from_dict(results,orient='index').reset_index().rename(columns={'index':'seq_id'})
+    df = df.loc[df['labels'].notna()].reset_index(drop=True)
+    # drop pred_labels if it's empty
+    if set(','.join(df['pred_labels'].tolist()))=={','}: 
+        df = df.drop(columns=['pred_labels'])
+        log_update(f"\t\tno predicted labels provided for this dataset; only probabilities")
+    log_update(f"\t\t{len(df)}/{len(all_caid2_disorder_nox_sequences)} total CAID2-Nox sequences")
+    return df
+    
+def parse_all_caid2_results(processed_caid2_df, caid_raw_folder="raw_data/caid2_competition_results"):
+    save_dir ="processed_data/caid2_competition_results" 
+    os.makedirs(save_dir,exist_ok=True)
+    
+    log_update(f"\nExtracting all CAID-2_Disorder_NOX results from CAID2 competition results files...")
+    all_caid_files = os.listdir(caid_raw_folder)
+    for caid_file in all_caid_files:
+        # figure out how to parse .caid files 
+        with open(f"{caid_raw_folder}/{caid_file}", "r") as f:
+            lines = f.readlines()
+            log_update(f"\t{caid_file}:")
+            results_df = parse_caid2_results(processed_caid2_df,lines)
+            # save it
+            competitor_name = caid_file.split('.caid')[0]
+            results_df.to_csv(f"{save_dir}/{competitor_name}_CAID-2_Disorder_NOX.csv",index=False)
+    
+def make_train_df(fldpnn_df, idp_crf_df):
+    """
+    Make training dataframe by concatenating the two processed training sets. 
+    """
+    # Add source columns so we can track where each sequence came from
+    idp_crf_df = idp_crf_df.rename(columns={'Label':'IDP-CRF Label', 'ID': 'IDP-CRF ID'}).drop(columns=['Split'])
+    fldpnn_df = fldpnn_df.rename(columns={'Label':'flDPnn Label', 'ID': 'flDPnn ID'}).drop(columns=['Split'])
+    ########### Combine fldpnn and idp crf
+    # Join
+    log_update("\nJoining flDPnn and IDP-CRF data by sequence make unified training set")
+    train_df = pd.merge(idp_crf_df, 
+                        fldpnn_df,
+                        on='Sequence',
+                        how='outer',
+                        indicator=True)
+    train_df["Split"] = ["Train"]*len(train_df)
+    # Map _merge column to desired labels
+    train_df['Source'] = train_df['_merge'].map({
+        'left_only': 'IDP-CRF',
+        'right_only': 'flDPnn',
+        'both': 'IDP-CRF,flDPnn'
+    })
+    train_df = train_df.drop(columns=["_merge"])
+    log_update(f"\tIDP-CRF dataset size: {len(idp_crf_df)}\n\tfLDpnn dataset size: {len(fldpnn_df)}\n\tinitial train dataset size: {len(train_df)}")
+
+    # Check for duplicate sequences
+    log_update(f"\tChecking for sequences in both datasets...")
+    duplicates = train_df[train_df["Source"].str.contains(",")]['Sequence'].unique().tolist()
+    n_rows_with_duplicates = len(train_df[train_df['Sequence'].isin(duplicates)])
+    log_update(f"\t\t{len(duplicates)} sequences in both datasets, corresponding to {n_rows_with_duplicates} rows")
+
+    # Check for consistency between IDP-CRF Label and flDPnn label
+    train_df["Label"] = train_df.apply(lambda row: find_agreeing_labels(row,lab1="IDP-CRF Label",lab2="flDPnn Label"),axis=1)
+    train_df["No Label Conflicts"]= ~train_df["Label"].isna()
+    log_update(f"\tChecked for label inconsistencies between IDP-CRF and flDPnn on the same sequence:")
+    match_str = train_df['No Label Conflicts'].value_counts().reset_index().rename(columns={'index': 'No Label Conflicts','No Label Conflicts': 'count'}).to_string(index=False)
+    match_str = "\t\t" + match_str.replace("\n","\n\t\t")
+    log_update(match_str)
+    
+    # Dropping rows where labels don't match 
+    #train_df[train_df['No Label Conflicts']==False][['Sequence','Split','IDP-CRF ID','flDPnn ID','IDP-CRF Label','flDPnn Label','No Label Conflicts']].to_csv('mismatch.csv',index=False)
+    # Drop row with known conflict with disprot
+    conflict_seq="MASREEEQRETTPERGRGAARRPPTMEDVSSPSPSPPPPRAPPKKRMRRRIESEDEEDSSQDALVPRTPSPRPSTSAADLAIAPKKKKKRPSPKPERPPSPEVIVDSEEEREDVALQMVGFSNPPVLIKHGKGGKRTVRRLNEDDPVARGMRTQEEEEEPSEAESEITVMNPLSVPIVSAWEKGMEAARALMDKYHVDNDLKANFKLLPDQVEALAAVCKTWLNEEHRGLQLTFTSKKTFVTMMGRFLQAYLQSFAEVTYKHHEPTGCALWLHRCAEIEGELKCLHGSIMINKEHVIEMDVTSENGQRALKEQSSKAKIVKNRWGRNVVQISNTDARCCVHDAACPANQFSGKSCGMFFSEGAKAQVAFKQIKAFMQALYPNAQTGHGHLLMPLRCECNSKPGHAPFLGRQLPKLTPFALSNAEDLDADLISDKSVLASVHHPALIVFQCCNPVYRNSRAQGGGPNCDFKISAPDLLNALVMVRSLWSENFTELPRMVVPEFKWSTKHQYRNVSLPVAHSDARQNPFDF"
+    train_df = train_df.loc[train_df['Sequence']!=conflict_seq].reset_index(drop=True)
+    log_update(f"\tDropping rows with label mismatch or known error (total={len(train_df[train_df['No Label Conflicts']==False])+1})")
+    train_df = train_df.loc[train_df['No Label Conflicts']].reset_index(drop=True)
+    
+    # Make a new label column
+    train_df = train_df.drop(columns=["IDP-CRF Label","flDPnn Label"])
+    log_update(f"\t\tNew dataset size: {len(train_df)}")
+    
+    ######## Final checks
+    # Check for any invalid sequences or invalid characters
+    cols_of_interest =  ['Sequence','Split','Label','IDP-CRF ID','flDPnn ID']
+    listlike_dict = check_columns_for_listlike(train_df, cols_of_interest, DELIMITERS)
+    
+    # Check for invalid characters
+    train_df['invalid_chars'] = train_df['Sequence'].apply(lambda x: find_invalid_chars(x, VALID_AAS))
+    train_df[train_df['invalid_chars'].str.len()>0].sort_values(by='Sequence')
+    all_invalid_chars = set().union(*train_df['invalid_chars'])
+    log_update(f"\tchecking for invalid characters...\n\t\tset of all invalid characters discovered within train_df: {all_invalid_chars}")
+
+    # Dropping rows where invalid characters(should be none)
+    log_update(f"\tDropping rows with invalid characters (total={len(train_df[train_df['invalid_chars'].str.len()>0])})")
+    train_df = train_df.loc[train_df['invalid_chars'].str.len()==0].reset_index(drop=True)
+    train_df = train_df.drop(columns=['invalid_chars'])
+    log_update(f"\t\tNew dataset size: {len(train_df)}")
+    
+    source_str = train_df['Source'].value_counts().reset_index().rename(columns={'index': 'Source','Source': 'count'}).to_string(index=False)
+    source_str = "\t\t" + source_str.replace("\n","\n\t\t")
+    log_update(f"\tSources:\n{source_str}")
+    return train_df
+
+def make_train_and_test_df(train_df, test_df):
+    """
+    Combine the training and testing dataframe into one
+    """
+    log_update("\nMaking final dataframe with train and test splits")
+    # Concatenate proposed train and test
+    test_df["Source"] = ["CAID-2_Disorder_NOX"]*len(test_df)
+    splits_df = pd.concat([train_df.drop(columns=['No Label Conflicts']),
+                           test_df.rename(columns={'ID':'CAID-2_Disorder_NOX ID', 'Label': 'CAID-2_Disorder_NOX Label'})])
+    split_str = splits_df['Split'].value_counts().reset_index().rename(columns={'index': 'Split','Split': 'count'}).to_string(index=False)
+    split_str = "\t\t" + split_str.replace("\n","\n\t\t")
+    log_update(f"\tTrain dataset size: {len(train_df)}\n\tTest dataset size: {len(test_df)}\n\tinitial combined dataset size: {len(splits_df)}")
+    
+    # Check for duplicates - if we find any, REMOVE them from train and keep them in test
+    duplicates = splits_df[splits_df.duplicated('Sequence')]['Sequence'].unique().tolist()
+    n_rows_with_duplicates = len(splits_df[splits_df['Sequence'].isin(duplicates)])
+    log_update(f"\t\t{len(duplicates)} duplicated sequences, corresponding to {n_rows_with_duplicates} rows")
+    for i, dup in enumerate(duplicates):
+        fldpnn_id = splits_df.loc[(splits_df['Sequence']==dup)&(splits_df['Split']=='Train')]['flDPnn ID'].item()
+        idp_crf_id = splits_df.loc[(splits_df['Sequence']==dup)&(splits_df['Split']=='Train')]['IDP-CRF ID'].item()
+        caid2_disorder_nox_id = splits_df.loc[(splits_df['Sequence']==dup)&(splits_df['Split']=='Test')]['CAID-2_Disorder_NOX ID'].item()
+        log_update(f"\t\t\t{i+1}: flDPnn ID: {fldpnn_id}\tIDP-CRF ID: {idp_crf_id}\tCAID-2_Disorder_NOX ID: {caid2_disorder_nox_id}\n\t\t\t\tSequence: {dup}")
+    # remove from train and keep in test
+    splits_df = splits_df.loc[
+        (~splits_df['Sequence'].isin(duplicates)) |       # Either the sequence is NOT duplicated, or
+        ((splits_df['Sequence'].isin(duplicates)) & (splits_df['Split']=='Test'))     # Sequence is duplicated, and it's in test set
+    ].reset_index(drop=True)
+    split_str = splits_df['Split'].value_counts().reset_index().rename(columns={'index': 'Split','Split': 'count'}).to_string(index=False)
+    split_str = "\t\t" + split_str.replace("\n","\n\t\t")
+    log_update(f"\tRemoved duplicate sequences from training split, kept in test split\n\t\tNew dataset size: {len(splits_df)}\n\n{split_str}")
+    
+    # Everything in the train set should have a label; nothing in the test set should
+    assert splits_df[splits_df["Label"].isna()]["Split"].value_counts().reset_index()['index'].tolist()==['Test']
+    splits_df.loc[
+        splits_df["Split"]=="Test","Label"
+    ] = splits_df.loc[
+        splits_df["Split"]=="Test","CAID-2_Disorder_NOX Label"
+    ] 
+    splits_df = splits_df.drop(columns=["CAID-2_Disorder_NOX Label"])
+    # Make sure there are no na's in label
+    assert len(splits_df[splits_df["Label"].isna()])==0
+    
+    # Print out distribution of sources
+    source_str = splits_df['Source'].value_counts().reset_index().rename(columns={'index': 'Source','Source': 'count'}).to_string(index=False)
+    source_str = "\t\t" + source_str.replace("\n","\n\t\t")
+    total_sources = sum(splits_df['Source'].value_counts().reset_index()['Source'])
+    assert total_sources == len(splits_df)
+    log_update(f"\n\tSource distribution:\n{source_str}\n\n\t\t\t\t\t\tSum:  {total_sources}")
+    
+    # Print largest and smallest seq len in each set
+    longest_train = max(splits_df[splits_df['Split']=='Train']['Sequence'].apply(lambda x: len(x)).tolist())
+    shortest_train = min(splits_df[splits_df['Split']=='Train']['Sequence'].apply(lambda x: len(x)).tolist())
+    longest_test = max(splits_df[splits_df['Split']=='Test']['Sequence'].apply(lambda x: len(x)).tolist())
+    shortest_test = min(splits_df[splits_df['Split']=='Test']['Sequence'].apply(lambda x: len(x)).tolist())
+    log_update(f"\n\tLength distributions...\n\t\tTrain: max={longest_train}\tmin={shortest_train}\n\t\tTest: max={longest_test}\tmin={shortest_test}")
+    
+    # Consolidate the IDs a bit 
+    splits_df["IDs"] = splits_df.apply(lambda row: get_unique_ids(row),axis=1)
+    assert len(splits_df[splits_df["IDs"].isna()])==0
+    n_different_ids = len(splits_df.loc[splits_df["IDs"].str.contains(",")])
+    log_update(f"\n\tProvided comma-separated IDs in same listed order as Source\n\t\t- train: IDP-CRF first, flDPnn second ({n_different_ids} seqs have multiple distinct IDs)\n\t\t- test: CAID-2_Disorder_NOX")
+    
+    # Keep only desired columns
+    splits_df = splits_df[[
+        'Sequence','IDs','Split','Source','Label'
+    ]]
+    
+    return splits_df
+
+def main():
+    with open_logfile("data_cleaning_log.txt"):
+        rawdata_train_test_path = "raw_data/caid2_train_and_test_data"
+        # make directory to save processed data
+        processeddata_path = "processed_data"
+        splits_path = "splits"
+        os.makedirs(processeddata_path,exist_ok=True)
+        os.makedirs(splits_path,exist_ok=True)
+        
+        # Process CAID-2_Disorder_NOX_Testing_Sequences dataset from fasta file
+        caid_path = f"{rawdata_train_test_path}/CAID-2_Disorder_NOX_Testing_Sequences.fasta"
+        caid_df = process_caid2_disorder_nox_test(caid_path)
+        caid_df.to_csv(f"{processeddata_path}/CAID-2_Disorder_NOX_Processed.csv", index=False)
+        
+        # Process fldpnn Training and Validation Datasets
+        fldpnn_train_path = f"{rawdata_train_test_path}/flDPnn_Training_Dataset.txt"
+        fldpnn_val_path = f"{rawdata_train_test_path}/flDPnn_Validation_Annotation.txt"
+        fldpnn_train_df = process_fldpnn(fldpnn_train_path, split="training")
+        fldpnn_val_df = process_fldpnn(fldpnn_val_path, split="validation")
+        fldpnn_train_df.to_csv(f"{processeddata_path}/flDPnn_Training_Dataset.csv", index=False)
+        fldpnn_val_df.to_csv(f"{processeddata_path}/flDPnn_Validation_Dataset.csv", index=False)
+        # Combine train and val
+        fldpnn_df = combine_fldpnn_train_val(fldpnn_train_df, fldpnn_val_df)
+        
+        # Process IDP-CRF_Training_Dataset
+        idp_crf_train_path = f"{rawdata_train_test_path}/IDP-CRF_Training_Dataset.txt"
+        idp_crf_df= process_idp_crf_train(idp_crf_train_path)
+        idp_crf_df.to_csv(f"{processeddata_path}/IDP-CRF_Training_Dataset.csv", index=False)
+        
+        # Merge
+        train_df = make_train_df(fldpnn_df, idp_crf_df)
+        
+        # Make a full splits file 
+        splits_df = make_train_and_test_df(train_df, caid_df)
+        final_train_df = splits_df.loc[splits_df['Split']=='Train'].reset_index(drop=True)
+        final_test_df = splits_df.loc[splits_df['Split']=='Test'].reset_index(drop=True)
+        
+        # Save final files
+        final_train_df.to_csv(f"{splits_path}/train_df.csv", index=False)
+        final_test_df.to_csv(f"{splits_path}/test_df.csv", index=False)
+    
+        # Process the caid competition results and save them in a more accessible format
+        processed_caid2_df = pd.read_csv(f"{processeddata_path}/CAID-2_Disorder_NOX_Processed.csv")
+        parse_all_caid2_results(processed_caid2_df)
+        
+        # Process data for visualizing fusion oncoproteins
+        # Scrape FusionPDB
+        scrape_fusionpdb_level_2_3()
+        # Process the structures that we downloaded from scraping
+        process_fusions_and_hts()
+        
+        # Now, figure out which structures are in the test set and isolate those for benchmarking in splits/fusion_bench_df.csv
+        fusion_test_set = pd.read_csv("../../data/splits/test_df.csv")
+        # columns are sequence, member length, snp_probabilities
+        fusion_test_set = set(fusion_test_set['sequence'].tolist())
+        log_update(f"\nFinding level 2 and 3 fusion structures that are in the FusOn-pLM test set...\n\tTest set size: {len(fusion_test_set)} seqs")
+        level_2_3_info = pd.read_csv('processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv')
+        # there are duplicate sequences in here so drop duplicates randomly
+        level_2_3_seqs = level_2_3_info.drop_duplicates('Fusion_Seq').reset_index(drop=True)
+        level_2_3_seqs = set(level_2_3_seqs.loc[
+            level_2_3_info['Fusion_pLDDT'].notna()  # make sure we've got a structure
+            ]['Fusion_Seq'].tolist())
+        # if it has a structure, it's in the test set, and it's not in the caid train set, we can benchmark with it
+        test_benchmark_seqs = fusion_test_set.intersection(level_2_3_seqs)   
+        log_update(f"\tTotal fusion proteins in the FusOn-pLM test set: {len(test_benchmark_seqs)}")
+        caid_train_set = set(pd.read_csv('splits/train_df.csv')['Sequence'].tolist())
+        test_benchmark_seqs = test_benchmark_seqs.difference(caid_train_set)    # subtract off the caid train set to be sure
+        log_update(f"\tTotal fusion proteins in the FusOn-pLM test set and NOT in the CAID train set: {len(test_benchmark_seqs)}")
+        
+        # Finally, make a dataframe structured like train_df and test_df. Columns are: Sequence,IDs,Split,Source,Label
+        # Let's make the IDs FusionGID
+        test_benchmark_df = pd.DataFrame(
+            data = {'Sequence': list(test_benchmark_seqs)}
+
+        )
+        seq_id_dict = dict(zip(level_2_3_info['Fusion_Seq'],level_2_3_info['FusionGID']))
+        seq_plddts_dict = dict(zip(level_2_3_info['Fusion_Seq'],level_2_3_info['Fusion_AA_pLDDTs']))
+        test_benchmark_df['IDs'] = test_benchmark_df['Sequence'].map(seq_id_dict)
+        test_benchmark_df['Split'] = ['Fusion_Benchmark']*len(test_benchmark_df)
+        test_benchmark_df['Source'] = ['FusionPDB_AlphaFold2']*len(test_benchmark_df)
+        test_benchmark_df['Label'] = test_benchmark_df['Sequence'].map(seq_plddts_dict)
+        # convert label to 1 or 0
+        test_benchmark_df['Label'] = test_benchmark_df['Label'].apply(lambda x: x.split(","))
+        test_benchmark_df['Label'] = test_benchmark_df['Label'].apply(lambda x: [float(y) for y in x])    # make it a float list of pLDDTs
+        test_benchmark_df['Label'] = test_benchmark_df['Label'].apply(lambda x: ['1' if y < 68.8 else '0' for y in x])    # disordered if pLDDT is < 68.8, accoridng to AlphaFold-pLDDT published threshold
+        test_benchmark_df['Label'] = test_benchmark_df['Label'].apply(lambda x: ''.join(x)) # change ['1','1','0''] to '110'
+        
+        # check lengths
+        test_benchmark_df['SeqLen'] = test_benchmark_df['Sequence'].apply(lambda x: len(x))
+        test_benchmark_df['LabelLen'] = test_benchmark_df['Label'].apply(lambda x: len(x))
+        log_update(f"\tAll seq lengths and label lengths match: {(test_benchmark_df['SeqLen']==test_benchmark_df['LabelLen']).all()}")
+        test_benchmark_df = test_benchmark_df.drop(columns=['SeqLen','LabelLen'])
+        
+        # convert to string
+        test_benchmark_df_str = test_benchmark_df.head(10)
+        test_benchmark_df_str['Sequence'] = test_benchmark_df_str['Sequence'].apply(lambda x: x[0:10]+'...')
+        test_benchmark_df_str['Label'] = test_benchmark_df_str['Label'].apply(lambda x: x[0:10]+'...')
+        test_benchmark_df_str = test_benchmark_df_str.to_string(index=False)
+        test_benchmark_df_str = "\t" + test_benchmark_df_str.replace("\n","\n\t")
+        log_update(f"\nPreview of benchmarking set:\n{test_benchmark_df_str}")
+        test_benchmark_df.to_csv('splits/fusion_bench_df.csv',index=False)
+        
+        # Add the benchmarking sequences to split 
+        log_update(f"\nAdding benchmarking sequences to splits_df.csv:\n\tLength before adding bench seqs: {len(splits_df)}")
+        splits_df = pd.concat([splits_df,test_benchmark_df])
+        log_update(f"\tLength after adding bench seqs: {len(splits_df)}")
+        split_str = splits_df['Split'].value_counts().reset_index().rename(columns={'index': 'Split','Split': 'count'}).to_string(index=False)
+        split_str = "\t" + split_str.replace("\n","\n\t")
+        log_update(f"Distribution among splits:\n{split_str}")
+        splits_df.to_csv(f"{splits_path}/splits.csv",index=False)
+
+if __name__ == "__main__":
+    main()

fuson_plm/benchmarking/caid/color_disordered_residues.ipynb

@@ -0,0 +1,849 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "FJd6a9gdZNjG"
+   },
+   "source": [
+    "### Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Put path to model predictions you'd like to use for benchmarking\n",
+    "path_to_model_predictions = \"trained_models/fuson_plm/best/caid_hyperparam_screen_fusion_benchmark_probs.csv\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install torch pandas numpy py3Dmol scikit-learn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "id": "ZEWZVc9lUxjI"
+   },
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "import pickle\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "\n",
+    "import py3Dmol\n",
+    "\n",
+    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, precision_recall_curve, average_precision_score"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "id": "Uu0z_JYYUzMC"
+   },
+   "outputs": [],
+   "source": [
+    "# Define paths and dataframes that we will need \n",
+    "fusion_benchmark_set = pd.read_csv('splits/fusion_bench_df.csv')\n",
+    "model_predictions = pd.read_csv(path_to_model_predictions)\n",
+    "fusion_structure_data = pd.read_csv('processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv')\n",
+    "fusion_structure_data['Fusion_Structure_Link'] = fusion_structure_data['Fusion_Structure_Link'].apply(lambda x: x.split('/')[-1])\n",
+    "fusion_structure_folder = \"raw_data/fusionpdb/structures\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# merge fusion data with seq ids \n",
+    "fuson_db = pd.read_csv('../../data/fuson_db.csv')\n",
+    "print(fuson_db.columns)\n",
+    "fuson_db = fuson_db[['aa_seq','seq_id']].rename(columns={'aa_seq':'Fusion_Seq'})\n",
+    "print(f\"Length of fusion structure data before merge on seqid: {len(fusion_structure_data)}\")\n",
+    "fusion_structure_data = pd.merge(\n",
+    "    fusion_structure_data,\n",
+    "    fuson_db,\n",
+    "    on='Fusion_Seq',\n",
+    "    how='inner'\n",
+    ")\n",
+    "print(f\"Length of fusion structure data after merge on seqid: {len(fusion_structure_data)}\")\n",
+    "fusion_structure_data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# merge fusion structure data with top swissprot alignments\n",
+    "swissprot_top_alignments = pd.read_csv(\"../../data/blast/blast_outputs/swissprot_top_alignments.csv\")\n",
+    "fusion_structure_data = pd.merge(\n",
+    "    fusion_structure_data,\n",
+    "    swissprot_top_alignments,\n",
+    "    on=\"seq_id\",\n",
+    "    how=\"left\"\n",
+    ")\n",
+    "fusion_structure_data.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(list(fusion_structure_data.columns))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "zORkLJztZWp9"
+   },
+   "source": [
+    "### Prepare data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "id": "JKbEOWP390ba"
+   },
+   "outputs": [],
+   "source": [
+    "def interpolate_color(value, start_color, end_color):\n",
+    "    return [(1 - value) * start + value * end for start, end in zip(start_color, end_color)]\n",
+    "\n",
+    "# Define the color points for the gradient\n",
+    "\n",
+    "def get_color(value):\n",
+    "    colors = [\n",
+    "        (0, 0, 0.545),  # Dark Blue\n",
+    "        (0.678, 0.847, 0.902),  # Light Blue\n",
+    "        (1, 1, 0),  # Yellow\n",
+    "        (1, 0.65, 0)  # Orange\n",
+    "    ]\n",
+    "\n",
+    "    if value <= 10:\n",
+    "        return interpolate_color(value / 10, colors[0], colors[1])\n",
+    "    elif value <= 30:\n",
+    "        return interpolate_color((value - 10) / 20, colors[1], colors[2])\n",
+    "    elif value <= 50:\n",
+    "        return interpolate_color((value - 30) / 20, colors[2], colors[3])\n",
+    "    else:\n",
+    "        return interpolate_color((value - 50) / 50, colors[3], colors[3])\n",
+    "\n",
+    "def color_by_disorder(cif_file, disorder_values):\n",
+    "    # Create viewer\n",
+    "    viewer = py3Dmol.view()\n",
+    "\n",
+    "    # Load CIF file\n",
+    "    with open(cif_file, 'r') as f:\n",
+    "        cif_data = f.read()\n",
+    "\n",
+    "    # Add structure\n",
+    "    viewer.addModel(cif_data, 'cif')\n",
+    "\n",
+    "    # Normalize disorder values to be between 0 and 100\n",
+    "    min_disorder = min(disorder_values)\n",
+    "    max_disorder = max(disorder_values)\n",
+    "    normalized_disorder = [(val - min_disorder) / (max_disorder - min_disorder) * 100 for val in disorder_values]\n",
+    "\n",
+    "    # Apply colors based on normalized disorder values\n",
+    "    for i, value in enumerate(normalized_disorder, start=1):\n",
+    "        mix_color = get_color(value)\n",
+    "        rgb_color = f'rgb({int(mix_color[0]*255)}, {int(mix_color[1]*255)}, {int(mix_color[2]*255)})'\n",
+    "        viewer.setStyle({'resi': i}, {'cartoon': {'color': rgb_color}})\n",
+    "\n",
+    "    # Show viewer\n",
+    "    viewer.zoomTo()\n",
+    "    return viewer.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# AlphaFold2 Database colors\n",
+    "def get_alphafold2_color(pLDDT):\n",
+    "  if pLDDT > 90: # darkblue\n",
+    "    return (13, 87, 211)\n",
+    "  elif pLDDT > 70: # lightblue\n",
+    "    return (106, 203, 241)\n",
+    "  elif pLDDT > 50: # yellow\n",
+    "    return (254, 217, 54)\n",
+    "  else: # orange\n",
+    "    return (253, 125, 77)\n",
+    "\n",
+    "def color_by_pLDDT(cif_file, plddt_values):\n",
+    "    viewer = py3Dmol.view(width=400, height=400)\n",
+    "    viewer.addModel(open(cif_file, 'r').read(), 'cif')\n",
+    "\n",
+    "    # Apply colors based on normalized disorder values\n",
+    "    for i, value in enumerate(plddt_values, start=1):\n",
+    "        mix_color = get_alphafold2_color(value)\n",
+    "        rgb_color = f'rgb({int(mix_color[0])}, {int(mix_color[1])}, {int(mix_color[2])})'\n",
+    "        viewer.setStyle({'resi': i}, {'cartoon': {'color': rgb_color}})\n",
+    "\n",
+    "    viewer.zoomTo()\n",
+    "    viewer.show()\n",
+    "  \n",
+    "def color_by_head_tail(cif_file, head_coords, tail_coords):\n",
+    "    \"\"\"\n",
+    "    head_coords and tail_coords are list of [start_psn, end_psn] inclusive. 1-indexed.\n",
+    "    \"\"\"\n",
+    "    head_resi_nos = list(range(head_coords[0], head_coords[1]+1))\n",
+    "    tail_resi_nos = list(range(tail_coords[0], tail_coords[1]+1))\n",
+    "    \n",
+    "    viewer = py3Dmol.view(width=400, height=400)\n",
+    "    viewer.addModel(open(cif_file, 'r').read(), 'cif')\n",
+    "\n",
+    "    viewer.setStyle({'cartoon': {'color': '#cf9dfa'}})  # fusion color purple\n",
+    "    viewer.setStyle({'resi': head_resi_nos}, {'cartoon': {'color': '#eb8888'}})\n",
+    "    viewer.setStyle({'resi': tail_resi_nos}, {'cartoon': {'color': '#5fa3e3'}})\n",
+    "\n",
+    "    viewer.zoomTo()\n",
+    "    viewer.show()\n",
+    "\n",
+    "def two_coloring_viewer_grid(cif_file, plddt_values, head_coords, tail_coords):\n",
+    "  head_resi_nos = list(range(head_coords[0], head_coords[1]+1))\n",
+    "  tail_resi_nos = list(range(tail_coords[0], tail_coords[1]+1))\n",
+    "    \n",
+    "  view = py3Dmol.view(width=800, height=400, viewergrid=(1,2))\n",
+    "  view.addModel(open(cif_file, 'r').read(), 'cif')\n",
+    "\n",
+    "  # Apply colors based on normalized disorder values\n",
+    "  for i, value in enumerate(plddt_values, start=1):\n",
+    "        mix_color = get_alphafold2_color(value)\n",
+    "        rgb_color = f'rgb({int(mix_color[0])}, {int(mix_color[1])}, {int(mix_color[2])})'\n",
+    "        view.setStyle({'resi': i}, {'cartoon': {'color': rgb_color}},viewer=(0,0))\n",
+    "\n",
+    "  view.setStyle({'cartoon': {'color': '#cf9dfa'}},viewer=(0,1))  # fusion color purple\n",
+    "  view.setStyle({'resi': head_resi_nos}, {'cartoon': {'color': '#eb8888'}},viewer=(0,1))\n",
+    "  view.setStyle({'resi': tail_resi_nos}, {'cartoon': {'color': '#5fa3e3'}},viewer=(0,1))\n",
+    "\n",
+    "  view.zoomTo()\n",
+    "  view.render()\n",
+    "  view.show()\n",
+    "\n",
+    "import os\n",
+    "def four_coloring_viewer_grid(cif_file, plddt_values, disorder_values, head_coords, tail_coords, seq_id=None, save_normalized_propensities=False):\n",
+    "    head_resi_nos = list(range(head_coords[0], head_coords[1]+1))\n",
+    "    tail_resi_nos = list(range(tail_coords[0], tail_coords[1]+1))\n",
+    "        \n",
+    "    view = py3Dmol.view(width=800, height=800, viewergrid=(2,2))\n",
+    "    view.addModel(open(cif_file, 'r').read(), 'cif')\n",
+    "\n",
+    "    # (0,0): apply AF2 colors based on normalized disorder values\n",
+    "    for i, value in enumerate(plddt_values, start=1):\n",
+    "        mix_color = get_alphafold2_color(value)\n",
+    "        rgb_color = f'rgb({int(mix_color[0])}, {int(mix_color[1])}, {int(mix_color[2])})'\n",
+    "        view.setStyle({'resi': i}, {'cartoon': {'color': rgb_color}},viewer=(0,0))\n",
+    "\n",
+    "    # (0,1): apply our colors based on normalized disorder values\n",
+    "    min_plddt = min(plddt_values)\n",
+    "    max_plddt = max(plddt_values)\n",
+    "    normalized_inverse_plddts = [(val - min_plddt) / (max_plddt - min_plddt) * 100 for val in plddt_values]\n",
+    "    normalized_inverse_plddts = [100 - val for val in normalized_inverse_plddts]\n",
+    "\n",
+    "    for i, value in enumerate(normalized_inverse_plddts, start=1):\n",
+    "        mix_color = get_color(value)\n",
+    "        rgb_color = f'rgb({int(mix_color[0]*255)}, {int(mix_color[1]*255)}, {int(mix_color[2]*255)})'\n",
+    "        view.setStyle({'resi': i}, {'cartoon': {'color': rgb_color}},viewer=(0,1))\n",
+    "\n",
+    "    # (1,0): apply colors based on head and tail segments\n",
+    "    view.setStyle({'cartoon': {'color': '#cf9dfa'}},viewer=(1,0))  # fusion color purple\n",
+    "    view.setStyle({'resi': head_resi_nos}, {'cartoon': {'color': '#eb8888'}},viewer=(1,0))\n",
+    "    view.setStyle({'resi': tail_resi_nos}, {'cartoon': {'color': '#5fa3e3'}},viewer=(1,0))\n",
+    "    \n",
+    "    ## (1,1): apply colors based on disorder predictions\n",
+    "    # Normalize disorder values to be between 0 and 100\n",
+    "    min_disorder = min(disorder_values)\n",
+    "    max_disorder = max(disorder_values)\n",
+    "    normalized_disorder = [(val - min_disorder) / (max_disorder - min_disorder) * 100 for val in disorder_values]\n",
+    "\n",
+    "    # if we're saving normalized propensities, then save them here\n",
+    "    if save_normalized_propensities:\n",
+    "      normalized_disorder_df = pd.DataFrame(\n",
+    "        data={\n",
+    "          'seq_id': seq_id,\n",
+    "          'normalized_disorder': [normalized_disorder]\n",
+    "        }\n",
+    "      )\n",
+    "      normalized_disorder_df['normalized_disorder'] = normalized_disorder_df['normalized_disorder'].apply(lambda x: [round(y,6) for y in x])\n",
+    "      os.makedirs(\"disorder_coloring_data\",exist_ok=True)\n",
+    "      fname = \"disorder_coloring_data/normalized_disorder_propensities_source_data.csv\"\n",
+    "      if os.path.isfile(fname):\n",
+    "        normalized_disorder_df.to_csv(fname,index=False,mode='a',header=False)\n",
+    "      else:\n",
+    "        normalized_disorder_df.to_csv(fname,index=False)\n",
+    "    # Apply colors based on normalized disorder values\n",
+    "    for i, value in enumerate(normalized_disorder, start=1):\n",
+    "        mix_color = get_color(value)\n",
+    "        rgb_color = f'rgb({int(mix_color[0]*255)}, {int(mix_color[1]*255)}, {int(mix_color[2]*255)})'\n",
+    "        view.setStyle({'resi': i}, {'cartoon': {'color': rgb_color}},viewer=(1,1))\n",
+    "\n",
+    "    view.zoomTo()\n",
+    "    view.render()\n",
+    "    view.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "SiTAKvW9aQvg"
+   },
+   "source": [
+    "### Save disorder prediction results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Y6Pxxyk5U_un"
+   },
+   "outputs": [],
+   "source": [
+    "model_predictions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_predictions_labeled = pd.merge(model_predictions,fusion_benchmark_set.rename(columns={'Sequence':'sequence'}),on='sequence',how='inner')\n",
+    "model_predictions_labeled = pd.merge(model_predictions_labeled, \n",
+    "                                     fusion_structure_data[['FusionGene','Fusion_Seq','Fusion_Structure_Link','Fusion_pLDDT','Fusion_AA_pLDDTs',\n",
+    "                                                            'top_hg_UniProtID', 'top_hg_UniProt_isoform', 'top_hg_UniProt_fus_indices', 'top_tg_UniProtID', 'top_tg_UniProt_isoform', \n",
+    "                                                            'top_tg_UniProt_fus_indices', 'top_UniProtID', 'top_UniProt_isoform', 'top_UniProt_fus_indices', 'top_UniProt_nIdentities', \n",
+    "                                                            'top_UniProt_nPositives']].rename(\n",
+    "                                         columns={'Fusion_Seq': 'sequence'}\n",
+    "                                     ),\n",
+    "                                     on='sequence',\n",
+    "                                     how='left')\n",
+    "model_predictions_labeled['length'] = model_predictions_labeled['sequence'].str.len()\n",
+    "seq_id_dict = dict(zip(fuson_db['Fusion_Seq'],fuson_db['seq_id']))\n",
+    "model_predictions_labeled['seq_id'] = model_predictions_labeled['sequence'].map(seq_id_dict)\n",
+    "model_predictions_labeled['Fusion_Structure_Link'] = model_predictions_labeled['Fusion_Structure_Link'].apply(lambda x: x.split('/')[-1])\n",
+    "model_predictions_labeled"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# calculate AUROC and AUPRC for each sequence\n",
+    "def calc_metrics(row):\n",
+    "    probs = row['prob_1']\n",
+    "    probs = [float(y) for y in probs.split(',')]\n",
+    "    true_labels = row['Label']\n",
+    "    true_labels = [int(y) for y in list(true_labels)]\n",
+    "    pred_labels = row['pred_labels']\n",
+    "    pred_labels = [int(y) for y in list(pred_labels)]\n",
+    "    \n",
+    "    # Calculate AUROC\n",
+    "    # Calculate AUPRC\n",
+    "    # Calculate all the other stats based on the predicted labels\n",
+    "    \n",
+    "    flat_binary_preds = np.array(pred_labels)\n",
+    "    flat_prob_preds = np.array(probs)\n",
+    "    flat_labels = np.array(true_labels)\n",
+    "\n",
+    "    accuracy = accuracy_score(flat_labels, flat_binary_preds)\n",
+    "    precision = precision_score(flat_labels, flat_binary_preds)\n",
+    "    recall = recall_score(flat_labels, flat_binary_preds)\n",
+    "    f1 = f1_score(flat_labels, flat_binary_preds)\n",
+    "    try:\n",
+    "        roc_auc = roc_auc_score(flat_labels, flat_prob_preds)\n",
+    "    except:\n",
+    "        roc_auc = np.nan\n",
+    "        \n",
+    "    try:\n",
+    "        auprc = average_precision_score(flat_labels, flat_prob_preds)\n",
+    "    except: \n",
+    "        auprc = np.nan\n",
+    "    \n",
+    "    return pd.Series({\n",
+    "        'Accuracy': round(accuracy,3),\n",
+    "        'Precision': round(precision,3),\n",
+    "        'Recall': round(recall,3),\n",
+    "        'F1': round(f1,3),\n",
+    "        'AUROC': round(roc_auc,3) if not(np.isnan(roc_auc)) else roc_auc,\n",
+    "        'AUPRC': round(auprc,3) if not(np.isnan(auprc)) else auprc,\n",
+    "    })\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_predictions_labeled[['Accuracy','Precision','Recall','F1','AUROC','AUPRC']] = model_predictions_labeled.apply(lambda row: calc_metrics(row),axis=1)\n",
+    "model_predictions_labeled = model_predictions_labeled.sort_values(by=['AUROC','F1','AUPRC','Accuracy','Precision','Recall'],ascending=[False,False,False,False,False,False]).reset_index(drop=True)\n",
+    "model_predictions_labeled['pcnt_disordered'] = round(100*model_predictions_labeled['Label'].apply(lambda x: sum([int(y) for y in list(x)]))/model_predictions_labeled['sequence'].str.len(),2)\n",
+    "model_predictions_labeled['pred_pcnt_disordered'] = round(100*model_predictions_labeled['pred_labels'].apply(lambda x: sum([int(y) for y in list(x)]))/model_predictions_labeled['sequence'].str.len(),2)\n",
+    "model_predictions_labeled[['sequence','length','FusionGene','Fusion_pLDDT','pcnt_disordered','pred_pcnt_disordered','AUROC','F1','AUPRC','Accuracy','Precision','Recall']]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.chdir('../../..')\n",
+    "!pip install -e .\n",
+    "import sys\n",
+    "sys.path.append('fuson_plm')\n",
+    "os.chdir('fuson_plm/benchmarking/caid')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Metrics calculation\n",
+    "def calculate_metrics(preds, labels, threshold=0.5):\n",
+    "    \"\"\"\n",
+    "    Calculates metrics to assess model performance\n",
+    "    Args:\n",
+    "        preds (list): model's predictions (probabilities)\n",
+    "        labels (list): ground truth labels\n",
+    "        threshold (float): minimum threshold a prediction must be met to be considered disordered\n",
+    "    Returns:\n",
+    "        accuracy (float): accuracy\n",
+    "        precision (float): precision\n",
+    "        recall (float): recall\n",
+    "        f1 (float): F1 score\n",
+    "        roc_auc (float): AUROC score\n",
+    "    \"\"\"\n",
+    "    flat_binary_preds, flat_prob_preds, flat_labels = [], [], []\n",
+    "\n",
+    "    for pred, label in zip(preds, labels):\n",
+    "        flat_binary_preds.extend((pred > threshold).astype(int).flatten())   # binary preds are 1 or 0; 1 if the prob > threshold\n",
+    "        flat_prob_preds.extend(pred.flatten())\n",
+    "        flat_labels.extend(label.flatten())\n",
+    "\n",
+    "    flat_binary_preds = np.array(flat_binary_preds)\n",
+    "    flat_prob_preds = np.array(flat_prob_preds)\n",
+    "    flat_labels = np.array(flat_labels)\n",
+    "\n",
+    "    accuracy = accuracy_score(flat_labels, flat_binary_preds)\n",
+    "    precision = precision_score(flat_labels, flat_binary_preds)\n",
+    "    recall = recall_score(flat_labels, flat_binary_preds)\n",
+    "    f1 = f1_score(flat_labels, flat_binary_preds)\n",
+    "    roc_auc = roc_auc_score(flat_labels, flat_prob_preds)\n",
+    "    \n",
+    "    # make dictionary of the results and return it\n",
+    "    metrics_dict = {\n",
+    "        'Accuracy': accuracy, \n",
+    "        'Precision': precision, \n",
+    "        'Recall': recall, \n",
+    "        'F1 Score': f1, \n",
+    "        'AUROC': roc_auc\n",
+    "    }\n",
+    "\n",
+    "    return metrics_dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Recalculate the true statistics\n",
+    "prob_and_label_df = model_predictions_labeled[['prob_1','Label']]\n",
+    "probs = ','.join(prob_and_label_df['prob_1'].tolist())\n",
+    "probs = [float(x) for x in probs.split(\",\")]\n",
+    "true_labels = ''.join(prob_and_label_df['Label'].tolist())\n",
+    "true_labels = [int(x) for x in list(true_labels)]\n",
+    "\n",
+    "calculate_metrics(np.array(probs), np.array(true_labels), threshold=0.022)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "w25hagtZaV65"
+   },
+   "source": [
+    "### Visualize on Favorite Fusions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "l = model_predictions_labeled['FusionGene'].tolist()\n",
+    "for i in range(len(l)):\n",
+    "    print(i, l[i])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def visualize_fusion_protein_fourgrid(selected_row, save_normalized_propensities=False):\n",
+    "    cif_file = f\"{fusion_structure_folder}/{selected_row['Fusion_Structure_Link']}\"\n",
+    "    plddt_values = [float(x) for x in selected_row['Fusion_AA_pLDDTs'].split(\",\")]\n",
+    "    disorder_values = [float(x) for x in selected_row['prob_1'].split(\",\")]\n",
+    "    head_indices = selected_row[\"top_hg_UniProt_fus_indices\"]\n",
+    "    if len(head_indices)>0:\n",
+    "        head_indices = [int(x) for x in head_indices.split(',')]\n",
+    "    else:\n",
+    "        head_indices = []\n",
+    "    tail_indices = selected_row[\"top_tg_UniProt_fus_indices\"]\n",
+    "    if len(tail_indices)>0:\n",
+    "        tail_indices = [int(x) for x in tail_indices.split(',')]\n",
+    "    else:\n",
+    "        tail_indices = []\n",
+    "    print(f\"Visualizing file: {cif_file}\")\n",
+    "    print(f\"Fusion protein: {selected_row['FusionGene']}\")\n",
+    "    print(f\"Fusion pLDDT: {selected_row['Fusion_pLDDT']}\")\n",
+    "    print(f\"Fusion pcnt_disordered: {selected_row['pcnt_disordered']}\\tpredicted pcnt_disordered: {selected_row['pred_pcnt_disordered']}\")\n",
+    "    print(f\"Sequence start: {selected_row['sequence'][0:10]}...\")\n",
+    "    print(f\"Sequence length: {len(selected_row['sequence'])}...\")\n",
+    "    print(f\"seq_id: {len(selected_row['seq_id'])}...\")\n",
+    "    print(f\"Fusion indices of best head alignment: {head_indices}\")\n",
+    "    print(f\"Fusion indices of best tail alignment: {tail_indices}\")\n",
+    "    print(f\"Disorder values start: {disorder_values[0:10]}...\")\n",
+    "    \n",
+    "    print(f\"\\nQuadrants:\\n\\t(0,0): pLDDT values in AF2 coloring\\n\\t(0,1): normalized 100-pLDDT values in AF2 coloring\\n\\t(1,0): head and tail segment coloring\\n\\t(1, 1): FusOn-pLM Disorder Model predictions, normalized, in AF2 coloring\")\n",
+    "\n",
+    "    #color_by_head_tail(cif_file, head_indices, tail_indices)\n",
+    "    #color_by_pLDDT(cif_file, plddt_values)\n",
+    "    four_coloring_viewer_grid(cif_file, plddt_values, disorder_values, head_indices, tail_indices, \n",
+    "                              seq_id=selected_row['seq_id'],\n",
+    "                              save_normalized_propensities=save_normalized_propensities)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = model_predictions_labeled.iloc[103,:]\n",
+    "visualize_fusion_protein_fourgrid(selected_row,save_normalized_propensities=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 497
+    },
+    "id": "Vr9ZECa8sMUg",
+    "outputId": "30b38391-9482-4cd5-8997-bb245fe5e6e5"
+   },
+   "outputs": [],
+   "source": [
+    "# Example usage:\n",
+    "### Used for FIgure 4D \n",
+    "selected_row = model_predictions_labeled.iloc[103,:]\n",
+    "cif_file = f\"{fusion_structure_folder}/{selected_row['Fusion_Structure_Link']}\"\n",
+    "disorder_values = [float(x) for x in selected_row['prob_1'].split(\",\")]\n",
+    "print(f\"Visualizing file: {cif_file}\")\n",
+    "print(f\"Fusion protein: {selected_row['FusionGene']}\")\n",
+    "print(f\"Fusion pLDDT: {selected_row['Fusion_pLDDT']}\")\n",
+    "print(f\"Fusion pcnt_disordered: {selected_row['pcnt_disordered']}\\tpredicted pcnt_disordered: {selected_row['pred_pcnt_disordered']}\")\n",
+    "print(f\"Sequence start: {selected_row['sequence'][0:10]}...\")\n",
+    "print(f\"Sequence length: {len(selected_row['sequence'])}...\")\n",
+    "print(f\"Disorder values start: {disorder_values[0:10]}...\")\n",
+    "color_by_disorder(cif_file, disorder_values)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#### Used for disorder figure\n",
+    "selected_row = model_predictions_labeled.loc[\n",
+    "    model_predictions_labeled['FusionGene'].str.contains('EWSR1::FLI1')\n",
+    "].reset_index(drop=True).iloc[0,:]\n",
+    "\n",
+    "visualize_fusion_protein_fourgrid(selected_row,save_normalized_propensities=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = model_predictions_labeled.loc[\n",
+    "    model_predictions_labeled['FusionGene'].str.contains('BCR::ABL')\n",
+    "].reset_index(drop=True).iloc[-1,:]\n",
+    "\n",
+    "visualize_fusion_protein_fourgrid(selected_row,save_normalized_propensities=True)\n",
+    "cif_file = cif_file = f\"{fusion_structure_folder}/{selected_row['Fusion_Structure_Link']}\"\n",
+    "disorder_values = [float(x) for x in selected_row['prob_1'].split(\",\")]\n",
+    "color_by_disorder(cif_file, disorder_values)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = model_predictions_labeled.loc[\n",
+    "    model_predictions_labeled['FusionGene'].str.contains('EML4::ALK')\n",
+    "].reset_index(drop=True).iloc[0,:]\n",
+    "\n",
+    "visualize_fusion_protein_fourgrid(selected_row,save_normalized_propensities=True)\n",
+    "cif_file = f\"{fusion_structure_folder}/{selected_row['Fusion_Structure_Link']}\"\n",
+    "disorder_values = [float(x) for x in selected_row['prob_1'].split(\",\")]\n",
+    "color_by_disorder(cif_file, disorder_values)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = model_predictions_labeled.loc[\n",
+    "    model_predictions_labeled['FusionGene'].str.contains('SS18::SSX1')\n",
+    "].reset_index(drop=True).iloc[0,:]\n",
+    "\n",
+    "visualize_fusion_protein_fourgrid(selected_row,save_normalized_propensities=True)\n",
+    "cif_file = f\"{fusion_structure_folder}/{selected_row['Fusion_Structure_Link']}\"\n",
+    "disorder_values = [float(x) for x in selected_row['prob_1'].split(\",\")]\n",
+    "color_by_disorder(cif_file, disorder_values)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visualize important fusions, colored by head/tail and by actual pLDDT"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def visualize_fusion_protein_twogrid(selected_row):\n",
+    "    cif_file = f\"{fusion_structure_folder}/{selected_row['Fusion_Structure_Link']}\"\n",
+    "    plddt_values = [float(x) for x in selected_row['Fusion_AA_pLDDTs'].split(\",\")]\n",
+    "    head_indices = selected_row[\"top_hg_UniProt_fus_indices\"]\n",
+    "    if len(head_indices)>0:\n",
+    "        head_indices = [int(x) for x in head_indices.split(',')]\n",
+    "    else:\n",
+    "        head_indices = []\n",
+    "    tail_indices = selected_row[\"top_tg_UniProt_fus_indices\"]\n",
+    "    if len(tail_indices)>0:\n",
+    "        tail_indices = [int(x) for x in tail_indices.split(',')]\n",
+    "    else:\n",
+    "        tail_indices = []\n",
+    "    print(f\"Visualizing file: {cif_file}\")\n",
+    "    print(f\"Fusion protein: {selected_row['FusionGene']}\")\n",
+    "    print(f\"Fusion pLDDT: {selected_row['Fusion_pLDDT']}\")\n",
+    "    print(f\"Sequence start: {selected_row['Fusion_Seq'][0:10]}...\")\n",
+    "    print(f\"Sequence length: {len(selected_row['Fusion_Seq'])}...\")\n",
+    "    print(f\"Fusion indices of best head alignment: {head_indices}\")\n",
+    "    print(f\"Fusion indices of best tail alignment: {tail_indices}\")\n",
+    "    #print(f\"Disorder values start: {disorder_values[0:10]}...\")\n",
+    "\n",
+    "    #color_by_head_tail(cif_file, head_indices, tail_indices)\n",
+    "    #color_by_pLDDT(cif_file, plddt_values)\n",
+    "    two_coloring_viewer_grid(cif_file, plddt_values, head_indices, tail_indices)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = pd.read_csv('../../data/fuson_db.csv')\n",
+    "df.loc[df['fusiongenes'].str.contains('EML4::ALK')]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fusion_structure_data.loc[\n",
+    "    fusion_structure_data['FusionGene'].str.contains('PAX3::FOXO1')\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def apply_thresh(y):\n",
+    "    if y < 68.8:\n",
+    "        return 1 #disorder\n",
+    "    return 0\n",
+    "\n",
+    "fusion_structure_data['n_disorder'] = fusion_structure_data['Fusion_AA_pLDDTs'].apply(lambda x: x.split(','))\n",
+    "fusion_structure_data['n_disorder'] = fusion_structure_data['n_disorder'].apply(lambda x: [float(y) for y in x])\n",
+    "#68.8\n",
+    "fusion_structure_data['n_disorder'] = fusion_structure_data['n_disorder'].apply(lambda x: [apply_thresh(y) for y in x])\n",
+    "fusion_structure_data['n_disorder'] = fusion_structure_data['n_disorder'].apply(lambda x: sum(x))\n",
+    "fusion_structure_data['pcnt_disorder'] = round(100*fusion_structure_data['n_disorder']/fusion_structure_data['Fusion_Length'],2)\n",
+    "fusion_structure_data['pcnt_disorder'] "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 40,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.makedirs(\"disorder_coloring_data\",exist_ok=True)\n",
+    "fusion_structure_data.loc[\n",
+    "    fusion_structure_data['FusionGene'].str.contains('PAX3::FOXO1') |\n",
+    "    fusion_structure_data['FusionGene'].str.contains('EWSR1::FLI1') |\n",
+    "    fusion_structure_data['FusionGene'].str.contains('EML4::ALK')|\n",
+    "    fusion_structure_data['FusionGene'].str.contains('SS18::SSX1')\n",
+    "][['FusionGID','FusionGene','pcnt_disorder','Fusion_Length','Fusion_pLDDT']].reset_index(drop=True)#.to_csv(\"disorder_coloring_data/favorite_fusions_all_structures.csv\",index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = fusion_structure_data.loc[\n",
+    "    fusion_structure_data['FusionGene'].str.contains('PAX3::FOXO1')\n",
+    "].sort_values(by='Fusion_pLDDT',ascending=True).reset_index(drop=True).iloc[0,:]\n",
+    "print('GID', selected_row['FusionGID'])\n",
+    "print('percent disorder', selected_row['pcnt_disorder'])\n",
+    "visualize_fusion_protein_twogrid(selected_row)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = fusion_structure_data.loc[\n",
+    "    fusion_structure_data['FusionGene'].str.contains('EWSR1::FLI1')\n",
+    "].sort_values(by='Fusion_pLDDT',ascending=True).reset_index(drop=True).iloc[0,:]\n",
+    "\n",
+    "visualize_fusion_protein_twogrid(selected_row)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = fusion_structure_data.loc[\n",
+    "    fusion_structure_data['FusionGene'].str.contains('EML4::ALK')\n",
+    "].sort_values(by='Fusion_pLDDT',ascending=True).reset_index(drop=True).iloc[0,:]\n",
+    "\n",
+    "#visualize_fusion_protein(selected_row)\n",
+    "#selected_row[['Fusion_Length','top_hg_UniProt_fus_indices','top_tg_UniProt_fus_indices']]\n",
+    "visualize_fusion_protein_twogrid(selected_row)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_row = fusion_structure_data.loc[\n",
+    "    fusion_structure_data['FusionGene'].str.contains('SS18::SSX1')\n",
+    "].sort_values(by='Fusion_pLDDT',ascending=True).reset_index(drop=True).iloc[1,:]\n",
+    "\n",
+    "visualize_fusion_protein_twogrid(selected_row)\n",
+    "#selected_row[['Fusion_Length','top_hg_UniProt_fus_indices','top_tg_UniProt_fus_indices']]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_predictions_labeled.loc[\n",
+    "    model_predictions_labeled['seq_id'].isin(['seq26176','seq12668','seq4546','seq11520','seq36572'])\n",
+    "][['seq_id','FusionGene']]"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "collapsed_sections": [
+    "FJd6a9gdZNjG",
+    "zORkLJztZWp9",
+    "w25hagtZaV65",
+    "IbyqxlvAFUAK",
+    "0n5PSprbhLk7"
+   ],
+   "machine_shape": "hm",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

fuson_plm/benchmarking/caid/config.py

@@ -0,0 +1,14 @@
+# Which models to benchmark
+BENCHMARK_FUSONPLM = True
+# FUSONPLM_CKPTS. If you've traiend your own model, this is a dictionary: key = run name, values = epochs
+# If you want to use the trained FusOn-pLM, instead FUSONPLM_CKPTS="FusOn-pLM"
+FUSONPLM_CKPTS= "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.

fuson_plm/benchmarking/caid/disorder_coloring_data/normalized_disorder_propensities_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4874f8b7e228a71038c231eca02abf06cee45f63f9bed106f5561b5cfa45c3b3
+size 43282

fuson_plm/benchmarking/caid/model.py

@@ -0,0 +1,26 @@
+import torch.nn as nn
+
+# Transformer model class
+class DisorderPredictor(nn.Module):
+    def __init__(self, input_dim, hidden_dim, num_heads, num_layers, dropout):
+        super(DisorderPredictor, self).__init__()
+        self.embedding_dim = input_dim
+        self.self_attention = nn.MultiheadAttention(embed_dim=input_dim, num_heads=num_heads, dropout=dropout)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dim,
+            nhead=num_heads,
+            dropout=dropout,
+            batch_first=True
+        )
+        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.classifier = nn.Linear(input_dim, 1)       # juts produce probabilities for 1 class
+        #self.softmax = nn.Softmax(dim=-1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, embeddings):
+        attn_out, _ = self.self_attention(embeddings, embeddings, embeddings) # Start with embeddings as random Q, K, V vectors
+        transformer_out = self.transformer_encoder(attn_out) # Get outputs from encoder layers
+        logits = self.classifier(transformer_out) # Linear classifier
+        probs = self.sigmoid(logits.squeeze(-1)) # sigmoid for probabilities; remove the last dimension of size 1 (since we only predicted 1 class)
+        return probs  # Get probabilities of dimension seq_len
+    

fuson_plm/benchmarking/caid/plot.py

@@ -0,0 +1,1030 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+import pandas as pd
+import numpy as np
+import os
+import matplotlib.colors as mcolors
+import matplotlib.patches as mpatches
+from matplotlib import font_manager
+import matplotlib.patches as patches
+from sklearn.metrics import roc_curve, auc, r2_score
+
+from fuson_plm.utils.visualizing import set_font
+
+global caid2_winners, caid2_model_rankings
+caid2_winners = pd.DataFrame(data=
+        {
+        'Model Name': ['Dispredict3','flDPnn2','flDPnn','flDPlr','flDPlr2','DisoPred',
+                       'IDP-Fusion','ESpritz-D','DeepIDP-2L','disomine','DISOPRED3-diso','IUPred3',
+                       'AlphaFold-rsa','AlphaFold-pLDDT'],    # do the top 6 models, and IUPred because it's well-known
+        'AUROC': [0.838,0.836,0.833,0.827,0.821,0.821,
+                  0.818,0.802,0.800,0.797,0.692,0.755,0.747,0.695],
+})
+caid2_winners['Model Type'] = ['caid2_competition']*len(caid2_winners)
+caid2_winners['Model Epoch'] = [np.nan]*len(caid2_winners)
+
+caid2_model_rankings = {
+    'Dispredict3': 1,
+    'flDPnn2': 2,
+    'flDPnn': 3,
+    'flDPlr': 4,
+    'flDPlr2': 5,
+    'DisoPred': 6,
+    'IDP-Fusion': 7,
+    'ESpritz-D': 8,
+    'DeepIDP-2L': 9,
+    'disomine': 10,
+    'DISOPRED3-diso': 35,
+    'IUPred3': 21,
+    'AlphaFold-rsa': 24,
+    'AlphaFold-pLDDT': 34
+}
+
+# Method for lengthening the model name
+def lengthen_model_name(row):
+    model_type = row['Model Type']
+    name = row['Model Name']
+    epoch = row['Model Epoch']
+    
+    if 'esm' in name:
+        return name
+    if 'puncta' in name:
+        return name
+    if model_type=='caid2_competition':
+        return name
+    
+    return f'{name}_e{epoch}'
+
+# Method for shortening the model name for display
+def shorten_model_name(row):
+    model_type = row['Model Type']
+    name = row['Model Name']
+    epoch = row['Model Epoch']
+    
+    if 'esm' in name:
+        return 'ESM-2-650M'
+    if model_type=='caid2_competition':
+        return name
+    
+    if 'snp_' in name:
+        prob_type = 'snp'
+    elif 'uniform_' in name:
+        prob_type = 'uni'
+    
+    layers = name.split('layers')[0].split('_')[-1]
+    maskrate = name.split('mask')[1].split('-', 1)[0]
+    kqv_tag = name.split('layers_')[1].split('_')[0]
+    dt = name.split('mask')[1].split('-', 1)[1]
+    
+    return f'{prob_type}_{layers}L_{kqv_tag}_mask{maskrate}_{dt}_e{epoch}'
+
+def make_heatmap(df, results_dir='.', gold_standard_model_name="esm2_t33_650M_UR50D",split="test",thresh=None,ax=None):
+    # Set font to Ubuntu
+    set_font()
+    
+    # Declare columns to compare: metrics
+    columns_to_compare = ['AUROC']
+    
+    # Define the literature-reported values for CAID competition winners - only IF the split is not "benchmark"
+    if not(split=="benchmark"):
+        df = pd.concat([df,caid2_winners])
+    
+    # Create Short Model Name and Full Model Name columns for later use
+    df['Model Epoch'] = df['Model Epoch'].apply(lambda x: str(int(x)) if not(np.isnan(x)) else '')
+    df['Short Model Name'] = df.apply(lambda row: shorten_model_name(row),axis=1)
+    df['Full Model Name'] = df.apply(lambda row: lengthen_model_name(row), axis=1)
+    
+    # Isolate gold standard row for later comparison
+    gold_standard = df[df['Full Model Name'] == gold_standard_model_name].reset_index(drop=True).iloc[0]
+    gold_standard_short_model_name = df[df['Full Model Name'] == gold_standard_model_name]['Short Model Name'].item()
+
+    # Create a new dataframe for the heatmap; sort by model type and place gold standard on top
+    heatmap_data = df[['Model Type','Short Model Name','Full Model Name'] + columns_to_compare].copy()
+    heatmap_data['is_gold_standard'] = (heatmap_data['Full Model Name'] == gold_standard_model_name).astype(int)
+    heatmap_data = heatmap_data.sort_values(by=['is_gold_standard','Model Type','AUROC'], ascending=[False,True,False]).reset_index(drop=True).drop(columns=['is_gold_standard'])
+    # Save the original values before calculating differences so we can use them for annotation
+    original_values = heatmap_data[columns_to_compare].copy()
+    
+    # Calculate differences from the gold standard
+    for col in columns_to_compare:
+        heatmap_data[col] = heatmap_data[col] - gold_standard[col]
+
+    # Create a color map where values equal to 0 are white, above are red, and below are blue
+    cmap = sns.color_palette("coolwarm", as_cmap=True)  # other option is diverging_palette(220, 20, as_cmap=True)
+
+    ### Make the plot
+    # can plot on a bigger plot, or make it an individual plot
+    if ax is None:
+        tallsize = max(8, 8 +.25*(len(heatmap_data)-26))
+        fig, ax = plt.subplots(1, 1, figsize=(8, tallsize), dpi=300)
+        
+    # Plot the heatmap with original values as annotations
+    hm = sns.heatmap(heatmap_data.set_index('Short Model Name').drop(columns=['Model Type','Full Model Name']),
+                    annot=False, fmt='', cmap=cmap, center=0, 
+                    cbar_kws={'label': 'Difference from Gold Standard'})
+    
+    # Explicitly set tick labels to prevent them from being messed up
+    ax.set_yticklabels(heatmap_data['Short Model Name'], rotation=0, fontsize=12)
+    # Add padding to the y-axis label
+    ax.set_ylabel("Short Model Name", labelpad=20)  # Increase the labelpad value to add more padding
+
+    # Bold any values values that exceed the gold standard
+    for i in range(original_values.shape[0]):
+        for j in range(original_values.shape[1]):
+            value = original_values.iloc[i, j]
+            if value > gold_standard[columns_to_compare[j]]:
+                ax.text(j + 0.5, i + 0.5, f'{value:.3f}', ha='center', va='center', fontweight='bold', color='black')
+            else:
+                ax.text(j + 0.5, i + 0.5, f'{value:.3f}', ha='center', va='center', color='black')
+                
+    # Add horizontal lines between different model types
+    model_type_series = heatmap_data['Model Type'].values
+    last_index = 0
+    labels_positions = []  # To store the positions for labels
+    for i in range(1, len(model_type_series)):
+        if model_type_series[i] != model_type_series[i - 1]:
+            hm.axhline(i, color='white', linewidth=8)  # Draw a thick white line between groups
+            labels_positions.append((last_index + i) / 2)  # Store the midpoint for labeling
+            last_index = i
+
+    # Add label for the last group
+    labels_positions.append((last_index + len(model_type_series)) / 2)
+    
+    # Italic and bold models that win AUROC; apply yellow coloring to gold standard model
+    for ytick, model_name in enumerate(heatmap_data['Short Model Name']):
+        if model_name == gold_standard_short_model_name:
+            # color yellow
+            label = ax.get_yticklabels()[ytick]
+            #label.set_color('gold')
+            label.set_bbox(dict(facecolor='gold', alpha=0.5, edgecolor='gold'))
+        if model_name != gold_standard_short_model_name:
+            auroc_value = original_values.loc[ytick, 'AUROC']
+            
+            # Apply bold and italic for wins on either AUROC or F1 Score
+            if (auroc_value > gold_standard['AUROC']):
+                label = ax.get_yticklabels()[ytick]
+                #label.set_style('italic')
+                #label.set_weight('bold')
+                label.set_bbox(dict(facecolor='red', alpha=0.3, edgecolor='red'))
+    
+    # Make legend
+    gold_patch = mpatches.Patch(color='gold', alpha=0.5, label='Gold Standard')
+    red_patch = mpatches.Patch(color='red', alpha=0.5, label='Winner')
+    plt.legend(handles=[gold_patch, red_patch], loc='best', bbox_to_anchor=(0, 0))  # You can change loc to position the legend
+
+    split_fname_dict = {
+        "testing": "CAID2_test",
+        "training": "CAID2_train",
+        "benchmark": "FusionPDB_pLDDT_disorder"
+    }
+    split_title_dict = {
+        "testing": "CAID-2 Disorder Prediction",
+        "training": "CAID-2 Disorder Prediction",
+        "benchmark": "FusionPDB_pLDDT Disorder Prediction"
+    }
+    ax.set_title(split_title_dict[split])
+    
+    # Rotate the color bar label
+    cbar = hm.collections[0].colorbar
+    cbar.ax.yaxis.set_label_position('right')
+    cbar.ax.yaxis.set_ticks_position('right')
+    cbar.set_label('Difference from Gold Standard', rotation=270, labelpad=20)  # Rotate 270 degrees and add some padding
+    
+    # Set tight layout using fig
+    fig.tight_layout(rect=[0, 0, 0.95, 1])  # Add extra padding on the right side to fit the label
+
+    plt.savefig(f"{results_dir}/{split_fname_dict[split]}_heatmap_vs_{gold_standard_model_name}.png")
+
+# Plot AUROC curve of ONE model of interest on its fusion pdb performance
+def make_benchmark_auroc_curve(results_dir='.', seq_label_dict=None, path_to_results_of_interest='', model_alias=None):
+    # Isolate the information for the model we'll be plotting
+    benchmark_model = path_to_results_of_interest.split('trained_models/')[1].split('/')
+    benchmark_model_type = benchmark_model[0]
+    benchmark_model_epoch = np.nan
+    benchmark_model_hyperparams = None
+    if len(benchmark_model)==5: 
+        benchmark_model_name = benchmark_model[1]
+        benchmark_model_epoch = benchmark_model[2].split('epoch')[1]
+        benchmark_model_hyperparams = benchmark_model[3]
+    else:
+        benchmark_model_name = benchmark_model[0]
+        benchmark_model_hyperparams = benchmark_model[1]
+    benchmark_model_info = pd.DataFrame(data={
+        'Model Type': [benchmark_model_type], 'Model Name': [benchmark_model_name], 'Model Epoch': [benchmark_model_epoch]
+    })
+    if model_alias is None:
+        model_alias = benchmark_model_info.apply(lambda row: shorten_model_name(row),axis=1).iloc[0]
+        
+    color_map = {
+        model_alias: 'black'
+    }
+    method_results = {model_alias: path_to_results_of_interest}
+    method_results = {k:v for k,v in method_results.items() if v not in [None, '']}
+    
+    set_font()
+    plt.figure(figsize=(10,6),dpi=300)
+    
+    # To store AUROC values and corresponding labels for sorting
+    roc_data = []
+    # Read each result file and plot the metrics
+    for method, path in method_results.items():
+        df = pd.read_csv(path) # columns = prob_1,labels
+
+        # Extract probabilities and labels
+        prob_1 = ",".join(df['prob_1'].tolist())
+        df['labels'] = df['sequence'].apply(lambda x: seq_label_dict[x])
+        labels = "".join(df['labels'].tolist())
+        prob_1 = [float(x) for x in prob_1.split(",")]
+        labels = [int(x) for x in list(labels)]
+        sequences = "".join(df['sequence'].tolist())
+        assert len(prob_1)==len(labels)==len(sequences)
+
+        # Compute ROC curve and ROC area
+        fpr, tpr, thresholds = roc_curve(labels, prob_1)
+        roc_auc = auc(fpr, tpr)
+        
+        # Store data for sorting later
+        roc_data.append((method, fpr, tpr, roc_auc))
+        
+    # Sort the methods by AUROC values
+    roc_data = sorted(roc_data, key=lambda x: x[3], reverse=True)
+
+    # Plot sorted ROC curves
+    for method, fpr, tpr, roc_auc in roc_data:
+        if method == model_alias:
+            plt.plot(fpr, tpr, color=color_map[method], lw=2, label=f'{method} ({roc_auc:0.3f})')
+        else:
+            plt.plot(fpr, tpr, color=color_map[method], lw=1, alpha=0.7, label=f'{method} ({roc_auc:0.3f})')
+
+    # Set other stylistic elements
+    plt.xlim([0.0, 1.0])
+    plt.ylim([0.0, 1.05])
+    plt.plot([0, 1], [0, 1], color='darkgrey', lw=2, linestyle='--')
+    plt.xlabel('False Positive Rate')
+    plt.ylabel('True Positive Rate')
+    plt.title('Receiver Operating Characteristic (ROC) Curve')
+    
+    # After plotting the ROC curves, customize the legend
+    handles, labels = plt.gca().get_legend_handles_labels()
+
+    # Create the legend first
+    legend = plt.legend(handles, labels, loc="center left", bbox_to_anchor=(1, 0.5))
+
+    # Iterate through the legend's text labels
+    for text in legend.get_texts():
+        if model_alias in text.get_text():
+            text.set_fontweight('bold')  # Bold the alias model
+        
+    plt.tight_layout()
+    plt.savefig(f'{results_dir}/FusionPDB_pLDDT_disorder_{model_alias}_AUROC_curve.png')
+    
+# Plot AUROC curve of ONE model of interest with all the CAID models
+def make_auroc_curve(results_dir='.', seq_label_dict=None, seq_ids_dict=None, path_to_results_of_interest='', model_alias=None, path_to_esm_results=None, with_rankings=False):
+    # Isolate the information for the model we'll be plotting
+    benchmark_model = path_to_results_of_interest.split('trained_models/')[1].split('/')
+    benchmark_model_type = benchmark_model[0]
+    benchmark_model_epoch = np.nan
+    benchmark_model_hyperparams = None
+    if len(benchmark_model)==5: 
+        benchmark_model_name = benchmark_model[1]
+        benchmark_model_epoch = benchmark_model[2].split('epoch')[1]
+        benchmark_model_hyperparams = benchmark_model[3]
+    else:
+        benchmark_model_name = benchmark_model[0]
+        benchmark_model_hyperparams = benchmark_model[1]
+    benchmark_model_info = pd.DataFrame(data={
+        'Model Type': [benchmark_model_type], 'Model Name': [benchmark_model_name], 'Model Epoch': [benchmark_model_epoch]
+    })
+    if model_alias is None:
+        model_alias = benchmark_model_info.apply(lambda row: shorten_model_name(row),axis=1).iloc[0]
+        
+    color_map = {
+        'Dispredict3': '#d62727',           #1
+        'flDPnn2': '#ff7f0f',               #2
+        'flDPnn': '#1f77b4',                #3
+        'flDPlr': '#bcbd21',                #4
+        'flDPlr2': '#16becf',               #5
+        'DisoPred': '#1f77b4',              #6
+        'IDP-Fusion': '#d62727',            #7
+        'ESpritz-D': '#8b564c',             #8
+        'DeepIDP-2L': '#e377c2',            #9
+        'disomine': '#e377c2',                #10
+        'DISOPRED3-diso': '#ff892d',             
+        'IUPred3': '#8b564c',
+        'AlphaFold-rsa': '#2ba02b',
+        'AlphaFold-pLDDT': '#ff892d',
+        model_alias: 'black'
+    }
+    method_results = {'Dispredict3': 'processed_data/caid2_competition_results/Dispredict3_CAID-2_Disorder_NOX.csv',
+                    'flDPnn2': 'processed_data/caid2_competition_results/flDPnn2_CAID-2_Disorder_NOX.csv',
+                    'flDPnn': 'processed_data/caid2_competition_results/flDPnn_CAID-2_Disorder_NOX.csv',
+                    'flDPlr': 'processed_data/caid2_competition_results/flDPtr_CAID-2_Disorder_NOX.csv',   # name doesn't match but this is what it is in raw download
+                    'flDPlr2': 'processed_data/caid2_competition_results/flDPlr2_CAID-2_Disorder_NOX.csv',
+                    'DisoPred': 'processed_data/caid2_competition_results/DisoPred_CAID-2_Disorder_NOX.csv',
+                    'IDP-Fusion': 'processed_data/caid2_competition_results/IDP-Fusion_CAID-2_Disorder_NOX.csv',        
+                    'ESpritz-D': 'processed_data/caid2_competition_results/ESpritz-D_CAID-2_Disorder_NOX.csv',           
+                    'DeepIDP-2L': 'processed_data/caid2_competition_results/DeepIDP-2L_CAID-2_Disorder_NOX.csv',          
+                    'disomine': 'processed_data/caid2_competition_results/disomine_CAID-2_Disorder_NOX.csv',              
+                    'DISOPRED3-diso': 'processed_data/caid2_competition_results/DISOPRED3-diso_CAID-2_Disorder_NOX.csv',             
+                    'AlphaFold-rsa': 'processed_data/caid2_competition_results/AlphaFold-rsa_CAID-2_Disorder_NOX.csv',
+                    'AlphaFold-pLDDT': 'processed_data/caid2_competition_results/AlphaFold-disorder_CAID-2_Disorder_NOX.csv',        # name doesn't match but this is what it is in raw download
+                    'IUPred3': 'processed_data/caid2_competition_results/IUPred3_CAID-2_Disorder_NOX.csv',
+                    model_alias: path_to_results_of_interest
+                }
+    if path_to_esm_results is not None:
+        method_results['ESM-2-650M'] = path_to_esm_results
+        color_map['ESM-2-650M'] = 'black'
+        
+    method_results = {k:v for k,v in method_results.items() if v not in [None, '']}
+    
+    set_font()
+    plt.figure(figsize=(12,6),dpi=300)
+    
+    # To store AUROC values and corresponding labels for sorting
+    merged_preds = pd.DataFrame(data={'sequence':[]})
+    merged_tpr_fpr = pd.DataFrame(data={'model': [],'fpr':[],'tpr':[]})
+    roc_data = []
+    # Read each result file and plot the metrics
+    for method, path in method_results.items():
+        df = pd.read_csv(path) # columns = prob_1,labels
+        merged_preds = pd.merge(merged_preds, 
+                                df.rename(columns={'prob_1':f"{method}_prob_1"})[['sequence',f"{method}_prob_1",]],
+                                on=['sequence'],how='outer')
+        
+        # Extract probabilities and labels
+        prob_1 = ",".join(df['prob_1'].tolist())
+        df['labels'] = df['sequence'].apply(lambda x: seq_label_dict[x])
+        labels = "".join(df['labels'].tolist())
+        prob_1 = [float(x) for x in prob_1.split(",")]
+        labels = [int(x) for x in list(labels)]
+        sequences = "".join(df['sequence'].tolist())
+        assert len(prob_1)==len(labels)==len(sequences)
+
+        # Compute ROC curve and ROC area
+        fpr, tpr, thresholds = roc_curve(labels, prob_1)
+        new_tpr_fpr = pd.DataFrame(data={
+            'model': [method]*len(fpr),
+            'fpr': fpr, 'tpr': tpr
+        })
+        merged_tpr_fpr = pd.concat([merged_tpr_fpr,new_tpr_fpr])
+        roc_auc = auc(fpr, tpr)
+        
+        if method==model_alias:
+            path_to_og_metrics = path_to_results_of_interest.rsplit('/',1)[0]+'/caid_hyperparam_screen_test_metrics.csv'
+            og_metrics = pd.read_csv(path_to_og_metrics)
+            roc_auc = og_metrics['AUROC'][0]
+        
+        # Store data for sorting later
+        roc_data.append((method, fpr, tpr, roc_auc))
+       
+    # Save the merged dataframe as source data
+    merged_preds['labels'] = merged_preds['sequence'].apply(lambda x: seq_label_dict[x])
+    merged_preds['labels'] = merged_preds['labels'].apply(lambda x: ",".join([str(y) for y in x]))
+    merged_preds['ids'] = merged_preds['sequence'].apply(lambda x: seq_ids_dict[x])
+    merged_preds.drop(columns={'sequence'}).to_csv(f"{results_dir}/CAID_prediction_source_data.csv",index=False)
+    merged_tpr_fpr.to_csv(f"{results_dir}/CAID_fpr_tpr_source_data.csv",index=False)
+    # Sort the methods by AUROC values
+    roc_data = sorted(roc_data, key=lambda x: x[3], reverse=True)
+    
+    # figure out the labels
+    labels = {method: method for method in method_results}
+    if with_rankings:
+        for method in labels:
+            if method in caid2_model_rankings:
+                labels[method] = f"{caid2_model_rankings[method]}. {method}"
+
+    # Plot sorted ROC curves
+    for method, fpr, tpr, roc_auc in roc_data:
+        if method=='ESM-2-650M' and path_to_esm_results is not None:
+            plt.plot(fpr, tpr, color=color_map[method], lw=2, linestyle='--', label=f'{labels[method]} ({roc_auc:0.3f})')
+        elif method == model_alias:
+            plt.plot(fpr, tpr, color=color_map[method], lw=2, label=f'{labels[method]} ({roc_auc:0.3f})')
+        else:
+            plt.plot(fpr, tpr, color=color_map[method], lw=1, alpha=0.7, label=f'{labels[method]} ({roc_auc:0.3f})')
+
+    # Set other stylistic elements
+    plt.xlim([0.0, 1.0])
+    plt.ylim([0.0, 1.05])
+    plt.xticks(fontsize=20)
+    plt.yticks(fontsize=20)
+    plt.plot([0, 1], [0, 1], color='darkgrey', lw=2, linestyle='--')
+    plt.xlabel('False Positive Rate', fontsize=22)
+    plt.ylabel('True Positive Rate', fontsize=22)
+    plt.title('CAID2 Disorder NOX Dataset: ROC Curve', fontsize=22)
+    
+    # After plotting the ROC curves, customize the legend
+    handles, labels = plt.gca().get_legend_handles_labels()
+
+    # Create the legend first
+    legend = plt.legend(handles, labels, loc="center left", bbox_to_anchor=(1.1, 0.5), fontsize=16)
+
+    # Iterate through the legend's text labels
+    for text in legend.get_texts():
+        if model_alias in text.get_text():
+            text.set_fontweight('bold')  # Bold the alias model
+        elif (path_to_esm_results is not None) and "ESM-2-650M" in text.get_text():
+            text.set_fontweight('bold') # Bold ESM if we're comparing to it
+        
+    plt.tight_layout()
+    figpath = f'{results_dir}/CAID2_{model_alias}_AUROC_curve.png'
+    if path_to_esm_results is not None:
+        figpath = f'{results_dir}/CAID2_{model_alias}_with_ESM_AUROC_curve.png'
+    plt.savefig(figpath)
+
+    
+def plot_disorder_content_scatter(train_labels, test_labels, benchmark_labels, savepath='splits/disorder_content_scatter.png'):
+    """
+    Compare disorder content between the train, test, and fusion benchmark sets based on the TRUE labels.
+    Each labels vector should have ['11110000','0001110',...] format.  
+    """
+    
+    # Get train disorder distribution
+    train_lengths = []
+    train_frac_disorder = []
+    for vec in train_labels:
+        veclist = [int(x) for x in vec]
+        train_lengths.append(len(veclist))
+        train_frac_disorder.append(sum(veclist)/len(veclist))
+    
+    # Get test disorder distribution
+    test_lengths = []
+    test_frac_disorder = []
+    for vec in test_labels:
+        veclist = [int(x) for x in vec]
+        test_lengths.append(len(veclist))
+        test_frac_disorder.append(sum(veclist)/len(veclist))
+        
+    # Get benchmark disorder distribution
+    benchmark_lengths = []
+    benchmark_frac_disorder = []
+    for vec in benchmark_labels:
+        veclist = [int(x) for x in vec]
+        benchmark_lengths.append(len(veclist))
+        benchmark_frac_disorder.append(sum(veclist)/len(veclist))
+    
+    # make a plot
+    set_font()
+    color_map = {
+    'train': '#0072B2',
+    'test': '#E69F00',
+    'fusion': 'purple'
+    }
+    
+    # Plotting
+    fig, ax = plt.subplots(figsize=(10, 6))
+
+    ax.scatter(train_lengths, train_frac_disorder, color=color_map['train'], label='Train', alpha=0.7)
+    ax.scatter(test_lengths, test_frac_disorder, color=color_map['test'], label='Test', alpha=0.7)
+    ax.scatter(benchmark_lengths, benchmark_frac_disorder, color=color_map['fusion'], label='Fusion', alpha=0.7)
+
+    # Labels and title
+    ax.set_xlabel('Length')
+    ax.set_ylabel('Fraction of Disorder')
+    ax.set_title('Length vs. Fraction of Disorder for Train, Test, and Benchmark Datasets')
+    ax.legend()
+    plt.tight_layout()
+    plt.savefig(savepath)
+
+def plot_disorder_content_hist(labels, ids, title="data", color="black", savepath='splits/disorder_content_histograms.png'):
+    """
+    Compare disorder content between the train, test, and fusion benchmark sets based on the TRUE labels.
+    Each labels vector should have ['11110000','0001110',...] format.  
+    """
+    set_font()
+    
+    # Get disorder distribution
+    lengths = []
+    frac_disorder = []
+    for vec in labels:
+        veclist = [int(x) for x in vec]
+        lengths.append(len(veclist))
+        frac_disorder.append(100*sum(veclist)/len(veclist)) # make it a percent, i like this better
+    
+    # save the source data
+    source_data = pd.DataFrame(data={
+        'ID': ids,
+        'Percent_Disordered': frac_disorder 
+    })
+    source_data['Percent_Disordered'] = source_data['Percent_Disordered'].round(3)
+    source_data.to_csv(savepath.replace(".png","_source_data.csv"),index=False)
+    
+    fig, ax = plt.subplots(1, 1, figsize=(20, 12))
+
+    # Plot histogram for train data
+    title_fontsize = 70
+    axislabel_fontsize = 70
+    tick_fontsize = 50
+    ax.hist(frac_disorder, bins=20, color=color, alpha=0.7)
+    ax.set_title(title, fontsize=title_fontsize)
+    ax.set_xlabel('% Disordered', fontsize=axislabel_fontsize)
+    ax.set_ylabel('Count', fontsize=axislabel_fontsize)
+    ax.grid(True)
+    ax.set_axisbelow(True)
+    ax.tick_params(axis='both', which='major', labelsize=tick_fontsize)
+    
+    # Calculate the mean and median of the percent coverage
+    mean_coverage = np.mean(frac_disorder)
+    median_coverage = np.median(frac_disorder)
+
+    # Add vertical line for the mean
+    ax.axvline(mean_coverage, color='black', linestyle='--', linewidth=2, label=f'Mean: {mean_coverage:.1f}%')
+    
+    # Add vertical line for the median
+    ax.axvline(median_coverage, color='black', linestyle='-', linewidth=2, label=f'Median: {median_coverage:.1f}%')
+
+    ax.legend(fontsize=50, title_fontsize=50)
+
+    plt.tight_layout()
+    plt.savefig(savepath)
+
+def plot_group_disorder_content_hist(train_labels, test_labels, benchmark_labels, savepath='splits/disorder_content_histograms.png',orient='horizontal'):
+    """
+    Compare disorder content between the train, test, and fusion benchmark sets based on the TRUE labels.
+    Each labels vector should have ['11110000','0001110',...] format.  
+    """
+    
+    # Get train disorder distribution
+    train_lengths = []
+    train_frac_disorder = []
+    for vec in train_labels:
+        veclist = [int(x) for x in vec]
+        train_lengths.append(len(veclist))
+        train_frac_disorder.append(sum(veclist)/len(veclist))
+    
+    # Get test disorder distribution
+    test_lengths = []
+    test_frac_disorder = []
+    for vec in test_labels:
+        veclist = [int(x) for x in vec]
+        test_lengths.append(len(veclist))
+        test_frac_disorder.append(sum(veclist)/len(veclist))
+        
+    # Get benchmark disorder distribution
+    benchmark_lengths = []
+    benchmark_frac_disorder = []
+    for vec in benchmark_labels:
+        veclist = [int(x) for x in vec]
+        benchmark_lengths.append(len(veclist))
+        benchmark_frac_disorder.append(sum(veclist)/len(veclist))
+    
+    # make a plot
+    set_font()
+    color_map = {
+    'train': '#0072B2',
+    'test': '#E69F00',
+    'fusion': 'mediumpurple'
+    }
+    
+    # Create a 1x3 subplot (1 row, 3 columns) or 3x1 
+    if orient=='horizontal':
+        fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=False)
+    if orient=='vertical':
+        fig, axes = plt.subplots(3, 1, figsize=(5, 15), sharey=False)
+
+    # Plot histogram for train data
+    title_fontsize = 26
+    axislabel_fontsize = 26
+    tick_fontsize = 16
+    axes[0].hist(train_frac_disorder, bins=20, color=color_map['train'], alpha=0.7)
+    axes[0].set_title('CAID2 Train', fontsize=title_fontsize)
+    if orient=="horizontal":
+        axes[0].set_xlabel('Fraction of Disorder', fontsize=axislabel_fontsize)
+    axes[0].set_ylabel('Frequency', fontsize=axislabel_fontsize)
+    axes[0].grid(True)
+    axes[0].set_axisbelow(True)
+    axes[0].tick_params(axis='both', which='major', labelsize=tick_fontsize)
+
+
+    # Plot histogram for test data
+    axes[1].hist(test_frac_disorder, bins=20, color=color_map['test'], alpha=0.7)
+    axes[1].set_title('CAID2 Test',fontsize=title_fontsize)
+    if orient=="horizontal":
+        axes[1].set_xlabel('Fraction of Disorder', fontsize=axislabel_fontsize)
+    if orient=="vertical":
+        axes[1].set_ylabel('Frequency', fontsize=axislabel_fontsize)
+    axes[1].grid(True)
+    axes[1].set_axisbelow(True)
+    axes[1].tick_params(axis='both', which='major', labelsize=tick_fontsize)
+
+    # Plot histogram for benchmark (fusion) data
+    axes[2].hist(benchmark_frac_disorder, bins=20, color=color_map['fusion'], alpha=0.7)
+    axes[2].set_title('Fusion Oncoproteins',fontsize=title_fontsize)
+    axes[2].set_xlabel('Fraction of Disorder', fontsize=axislabel_fontsize)
+    if orient=="vertical":
+        axes[2].set_ylabel('Frequency', fontsize=axislabel_fontsize)
+    axes[2].grid(True)
+    axes[2].set_axisbelow(True)
+    axes[2].tick_params(axis='both', which='major', labelsize=tick_fontsize)
+    plt.tight_layout()
+    plt.savefig(savepath)
+    
+def categorize_plddt(values):
+    categories = {
+        "<= 50": sum(1 for x in values if x <= 50),
+        "50-70": sum(1 for x in values if 50 < x <= 70),
+        "70-90": sum(1 for x in values if 70 < x <= 90),
+        "> 90": sum(1 for x in values if x > 90)
+    }
+    return categories
+
+
+def plot_fusion_sequence_pLDDT_left_to_right(fusion_structure_data, fusiongene, save_path=''):
+    """
+    Plot each amino acid in the sequence as a separate colored bar based on pLDDT values.
+    """
+    set_font()
+    # Filter for specific fusion data and preprocess
+    df_of_interest = fusion_structure_data[fusion_structure_data['FusionGene'] == fusiongene].copy()
+    df_of_interest['Fusion_AA_pLDDTs'] = df_of_interest['Fusion_AA_pLDDTs'].apply(lambda x: [float(i) for i in x.split(',')])
+    df_of_interest['Label'] = df_of_interest['Fusion_Length'].astype(str) + 'AAs'
+    
+    # Sort data by Fusion_Length
+    df_of_interest = df_of_interest.sort_values(by='Fusion_Length', ascending=True).reset_index(drop=True)
+    
+    # Define colors for each pLDDT range
+    category_colors = {"<= 50": "#f27842", "50-70": "#f8d514", "70-90": "#60c1e8", "> 90": "#004ecb"}
+    
+    # Helper function to get color based on pLDDT
+    def get_color(pLDDT):
+        if pLDDT > 90:
+            return category_colors["> 90"]
+        elif pLDDT > 70:
+            return category_colors["70-90"]
+        elif pLDDT > 50:
+            return category_colors["50-70"]
+        else:
+            return category_colors["<= 50"]
+    
+    # Start plotting each sequence with colored bars
+    fig, ax = plt.subplots(figsize=(10, 6))
+    if len(df_of_interest)<3:
+        fig, ax = plt.subplots(figsize=(10, 2))
+
+    average_plddt = dict(zip(df_of_interest['Label'], df_of_interest['Fusion_pLDDT']))
+    df_of_interest['Fusion_AA_colors'] = df_of_interest['Fusion_AA_pLDDTs'].apply(lambda x: [get_color(plddt) for plddt in x])
+    df_of_interest['Fusion_pLDDT_color'] = df_of_interest['Fusion_pLDDT'].apply(lambda plddt: get_color(plddt))
+    # just save the columns needed for the plot 
+    df_of_interest[['FusionGene','seq_id','Fusion_Length','Fusion_pLDDT','Fusion_AA_pLDDTs','Fusion_AA_colors','Fusion_pLDDT_color',
+                    'top_hg_UniProtID','top_hg_UniProt_isoform','top_hg_UniProt_fus_indices',
+                    'top_tg_UniProtID','top_tg_UniProt_isoform','top_tg_UniProt_fus_indices']].to_csv(f"{save_path}/plddt_sequence_{fusiongene}_source_data.csv",index=False)
+    
+    for idx, row in df_of_interest.iterrows():
+        pLDDT_values = row['Fusion_AA_pLDDTs']
+        colors = [get_color(plddt) for plddt in pLDDT_values]
+        
+        # Plot each amino acid in the sequence with the respective color
+        ax.bar(range(len(pLDDT_values)), 
+               [0.7] * len(pLDDT_values), color=colors, edgecolor='none', 
+               bottom=idx - 0.7 / 2)  # Centering each row at idx
+        
+    labels = df_of_interest['Label'].tolist()
+    # Annotate each bar with the Fusion_pLDDT value on the right, colored by PLDDT category
+    for idx, label in enumerate(labels):
+        avg_plddt_value = average_plddt[label]
+        
+        # Determine color based on the PLDDT category
+        if avg_plddt_value > 90:
+            color = '#004ecb'
+        elif avg_plddt_value > 70:
+            color = "#60c1e8"
+        elif avg_plddt_value > 50:
+            color = '#f8d514'
+        else:
+            color = '#f27842'
+            
+        # Annotate with the determined color
+        if len(df_of_interest)>10:
+            markersize = 10
+        elif len(df_of_interest)>5:
+            markersize = 16
+        else:
+            markersize=12
+        ax.plot(1.02*max(df_of_interest['Fusion_Length']), 
+                idx, marker='o', color="black", markersize=markersize, markerfacecolor=color, markeredgewidth=2)
+
+        # Add breakpoint box - make sure we actually HAVE one of each
+        hg_indices, tg_indices = None, None 
+        if not(type(df_of_interest['top_hg_UniProt_fus_indices'][idx])==float):
+            hg_indices = [int(x) for x in df_of_interest['top_hg_UniProt_fus_indices'][idx].split(',')]
+        if not(type(df_of_interest['top_tg_UniProt_fus_indices'][idx])==float):
+            tg_indices = [int(x) for x in df_of_interest['top_tg_UniProt_fus_indices'][idx].split(',')]
+        print(hg_indices, tg_indices)
+        
+        if (hg_indices is not None) and (tg_indices is not None):
+            box_start = min(hg_indices[-1],tg_indices[0])
+            box_end = max(hg_indices[-1],tg_indices[0])
+        elif hg_indices is not None:
+            box_start, box_end = hg_indices[-1], hg_indices[-1]
+        elif tg_indices is not None:
+            box_start, box_end = tg_indices[0], tg_indices[0]
+            
+        print(f"box indices for structure {idx}, fusion gene {fusiongene}", box_start, box_end)
+        
+        # Plot the rectangle, making it slightly larger than the rest of the bar
+        rect = patches.Rectangle((box_start, idx - 0.7 / 2), box_end-box_start, 0.7, linewidth=2, edgecolor='black', facecolor='none')
+        ax.add_patch(rect)
+    
+    # Customize plot
+    ax.set_yticks([])  # Hide y-axis ticks
+    ax.set_yticklabels([])  # Hide y-axis labels
+    ax.set_ylim(-0.5, len(df_of_interest) - 0.5) # reduce white space at top
+    ax.set_xlabel("Amino Acid Sequence (ordered)", fontsize=14)
+    # Customize x-axis for labeling
+    ax.set_xlim(left=0)  # Start x-axis at 0 to make bars flush left
+    ax.set_xlabel("Amino Acid Sequence (ordered)", fontsize=14)
+    ax.tick_params(axis='x', labelsize=30) 
+
+    
+    plt.title(f"{fusiongene} pLDDT Distribution by Amino Acid Sequence", fontsize=16)
+    plt.tight_layout()
+
+    # Save figure
+    fusiongene_savename = fusiongene.replace("::","-")
+    plt.savefig(f"{save_path}/plddt_sequence_{fusiongene_savename}.png", dpi=300)
+    plt.show()
+    
+def plot_favorite_fusion_pLDDT_distribution(fusion_structure_data, fusiongene, save_path=''):
+    """
+    Make a stacked bar chart of the pLDDT distribution 
+    """
+    set_font()
+    # Filter for EWSR1::FLI1 fusion data and preprocess
+    df_of_interest = fusion_structure_data[fusion_structure_data['FusionGene'] == fusiongene].copy()
+    df_of_interest['Fusion_AA_pLDDTs'] = df_of_interest['Fusion_AA_pLDDTs'].apply(lambda x: [float(i) for i in x.split(',')])
+    df_of_interest['Label'] = df_of_interest['Fusion_Length'].astype(str) + 'AAs'
+    # Sort data by Fusion_Length
+    df_of_interest = df_of_interest.sort_values(by='Fusion_Length', ascending=True).reset_index(drop=True)
+    # Convert to dictionary format
+    data_dict = dict(zip(df_of_interest['Label'], df_of_interest['Fusion_AA_pLDDTs']))
+    average_plddt = dict(zip(df_of_interest['Label'], df_of_interest['Fusion_pLDDT']))
+    
+    # Categorize each structure
+    categorized_data = {structure: categorize_plddt(plddt_values) for structure, plddt_values in data_dict.items()}
+
+    # Extract counts for each category
+    labels = list(categorized_data.keys())
+    categories = ["<= 50", "50-70", "70-90", "> 90"]
+    counts = {cat: [categorized_data[structure][cat] for structure in labels] for cat in categories}
+
+    # Define colors for each category
+    category_colors = {"<= 50": "#f27842", "50-70": "#f8d514", "70-90": "#60c1e8", "> 90": "#004ecb"}
+
+    # Re-categorize PLDDT values for the bar chart
+    categorized_data = {structure: categorize_plddt(plddt_values) for structure, plddt_values in data_dict.items()}
+    labels = list(categorized_data.keys())
+    counts = {cat: [categorized_data[structure][cat] for structure in labels] for cat in categories}
+
+    # Plotting the horizontal stacked bar chart with annotations for 'Fusion_pLDDT' values
+    fig, ax = plt.subplots(figsize=(10, 6))
+    if len(data_dict)<3:
+        fig, ax = plt.subplots(figsize=(10, 2))
+    bottom = np.zeros(len(labels))
+
+    # Stack each category horizontally
+    for cat in categories:
+        ax.barh(labels, counts[cat], label=cat, color=category_colors[cat], left=bottom)
+        bottom += counts[cat]  # Update the left position for the next stack
+
+    # Annotate each bar with the Fusion_pLDDT value on the right, colored by PLDDT category
+    for idx, label in enumerate(labels):
+        avg_plddt_value = average_plddt[label]
+        
+        # Determine color based on the PLDDT category
+        if avg_plddt_value > 90:
+            color = '#004ecb'
+        elif avg_plddt_value > 70:
+            color = "#60c1e8"
+        elif avg_plddt_value > 50:
+            color = '#f8d514'
+        else:
+            color = '#f27842'
+            
+        # Annotate with the determined color
+        #ax.text(bottom[idx] + 1, idx, f"{avg_plddt_value:.2f}", va='center', ha='left', color="black", fontsize=18, fontweight='bold')
+        if len(df_of_interest)>10:
+            markersize = 10
+        elif len(df_of_interest)>5:
+            markersize = 16
+        else:
+            markersize=12
+        ax.plot(bottom[idx] + .02*max(df_of_interest['Fusion_Length']), idx, marker='s', color="black", markersize=markersize, markerfacecolor=color, markeredgewidth=2)
+
+
+    # Add labels and legend
+    #ax.set_xlim([0,max(df_of_interest['Fusion_Length'])*1.0])
+    #ax.set_ylabel("Structures")
+    # Save original ticks before changing label size
+    #ax.tick_params(axis='x', labelsize=16) 
+    #original_xticks = ax.get_xticks()
+    # Set ticks explicitly to avoid automatic adjustment
+    #ax.set_xticks(original_xticks)
+
+    #ax.set_xlabel("Length",fontsize=40)
+    ax.tick_params(axis='x', labelsize=30) 
+    #ax.tick_params(axis='y', labelsize=16) 
+    ax.tick_params(axis='y', left=False, labelleft=False)
+    #ax.set_title(f"{fusiongene} pLDDT Distribution")
+    #ax.legend(title="pLDDT Ranges", fontsize=16, bbox_to_anchor=(1, 1), title_fontsize=16)
+
+    plt.tight_layout()
+    fusiongene_savename = fusiongene.replace("::","-")
+    plt.savefig(f"{save_path}/plddt_dist_{fusiongene_savename}.png",dpi=300)
+
+def make_all_favorite_fusion_pLDDT_plots(favorite_fusions,left_to_right=True):
+    fusion_structure_data = pd.read_csv('processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv')
+    swissprot_top_alignments = pd.read_csv("../../data/blast/blast_outputs/swissprot_top_alignments.csv")
+    fuson_db = pd.read_csv("../../data/fuson_db.csv")
+    seq_id_dict = dict(zip(fuson_db['aa_seq'],fuson_db['seq_id']))
+    fusion_structure_data['seq_id'] = fusion_structure_data['Fusion_Seq'].map(seq_id_dict)
+    fusion_structure_data = pd.merge(
+        fusion_structure_data,
+        swissprot_top_alignments,
+        on="seq_id",
+        how="left"
+    )
+    for x in favorite_fusions:
+        if left_to_right:
+            plot_fusion_sequence_pLDDT_left_to_right(fusion_structure_data, x, save_path='processed_data/figures/fusion_disorder')
+        else:
+            plot_favorite_fusion_pLDDT_distribution(fusion_structure_data, x, save_path='processed_data/figures/fusion_disorder')
+    
+def prep_data_for_ht_disorder_comparison():
+    ht_structure_data = pd.read_csv('processed_data/fusionpdb/heads_tails_structural_data.csv')
+    fusion_structure_data = pd.read_csv('processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv')
+    fusion_heads_and_tails = pd.read_csv('processed_data/fusionpdb/fusion_heads_and_tails.csv')
+
+    all_hts_with_structures = ht_structure_data['UniProtID'].unique().tolist()
+
+    fuson_ht_db = pd.read_csv('../../data/blast/fuson_ht_db.csv')[['seq_id','aa_seq','fusiongenes','hgUniProt','tgUniProt']]
+
+    merge = pd.merge(
+        fuson_ht_db.rename(columns={'aa_seq':'Fusion_Seq'}),
+        fusion_structure_data[['FusionGID', 'Fusion_Seq','Fusion_pLDDT','Fusion_AA_pLDDTs']],
+        on='Fusion_Seq',
+        how='right'
+    )
+    # now merge again
+    merge['hgUniProt'] = merge['hgUniProt'].apply(lambda x: x.split(','))
+    merge['tgUniProt'] = merge['tgUniProt'].apply(lambda x: x.split(','))
+    merge = merge.explode('hgUniProt')
+    merge = merge.explode('tgUniProt')
+    merge = merge.loc[
+        merge['hgUniProt'].isin(all_hts_with_structures) &
+        merge['tgUniProt'].isin(all_hts_with_structures)
+    ].reset_index(drop=True)
+
+    merge = pd.merge(
+        merge,
+        ht_structure_data.rename(columns=
+            {'UniProtID':'hgUniProt',
+            'Avg pLDDT': 'hg_pLDDT',
+            'All pLDDTs': 'hg_AA_pLDDTs',
+            'Seq': 'hg_seq'}),
+        on='hgUniProt',
+        how='inner'
+    )
+
+    merge = pd.merge(
+        merge,
+        ht_structure_data.rename(columns=
+            {'UniProtID':'tgUniProt',
+            'Avg pLDDT': 'tg_pLDDT',
+            'All pLDDTs': 'tg_AA_pLDDTs',
+            'Seq': 'tg_seq'}),
+        on='tgUniProt',
+        how='inner'
+    )
+    merge = merge.loc[merge['hg_AA_pLDDTs'].notna()]
+    merge = merge.loc[merge['tg_AA_pLDDTs'].notna()].reset_index(drop=True)
+
+    # finally, calcualte label
+    merge['hg_label'] = merge['hg_AA_pLDDTs'].apply(lambda x: x.split(','))
+    merge['hg_label'] = merge['hg_label'].apply(lambda x: [float(y) for y in x])
+    merge['hg_label'] = merge['hg_label'].apply(lambda x: [apply_plddt_thresh(y) for y in x])
+    merge['hg_label'] = merge['hg_label'].apply(lambda x: ''.join(x))
+
+    merge['tg_label'] = merge['tg_AA_pLDDTs'].apply(lambda x: x.split(','))
+    merge['tg_label'] = merge['tg_label'].apply(lambda x: [float(y) for y in x])
+    merge['tg_label'] = merge['tg_label'].apply(lambda x: [apply_plddt_thresh(y) for y in x])
+    merge['tg_label'] = merge['tg_label'].apply(lambda x: ''.join(x))
+
+    merge['fusion_label'] = merge['Fusion_AA_pLDDTs'].apply(lambda x: x.split(','))
+    merge['fusion_label'] = merge['fusion_label'].apply(lambda x: [float(y) for y in x])
+    merge['fusion_label'] = merge['fusion_label'].apply(lambda x: [apply_plddt_thresh(y) for y in x])
+    merge['fusion_label'] = merge['fusion_label'].apply(lambda x: ''.join(x))
+
+    return merge
+
+def apply_plddt_thresh(y):
+    if y < 68.8:
+        return '1'
+    else: 
+        return '0'
+
+def plot_fusion_stats_boxplots(data, save_path="fusion_disorder_boxplots.png"):
+    set_font()
+    # Create box plots
+    plt.figure(figsize=(6, 5))
+    # for ones that are 100% disordered, AUROC was NaN, so drop these
+    box = plt.boxplot([data[col].dropna() for col in data.columns], labels=data.columns, patch_artist=True)
+
+    # Set color of each box plot
+    for patch in box['boxes']:
+        patch.set_facecolor('#ff68b4')
+        patch.set_edgecolor('#ff68b4')
+    
+    # Customize other elements if needed
+    #for whisker in box['whiskers']:
+        #whisker.set_color('#ff68b4')
+    #for cap in box['caps']:
+        #cap.set_color('#ff68b4')
+    for median in box['medians']:
+        median.set_color('black')
+    # Add labels and title
+    #plt.xlabel('Metrics')
+    #plt.ylabel('Values')
+    plt.title(f"Per-Residue Disorder (n={len(data)})",fontsize=22)
+    plt.xticks(rotation=20,fontsize=22)
+    plt.yticks(fontsize=22)
+
+    # Show plot
+    plt.tight_layout()
+    plt.show()
+    plt.savefig(save_path,dpi=300)
+
+def plot_fusion_frac_disorder_r2(actual_values, predicted_values, save_path="fusion_pred_disorder_r2.png"):
+    set_font()
+    plt.figure(figsize=(6, 6))
+    r2 = r2_score(actual_values, predicted_values)
+    #sns.kdeplot(actual_values, label="Actual Values", shade=True)
+    #sns.kdeplot(predicted_values, label="Predicted Values", shade=True)
+    plt.scatter(actual_values, predicted_values, alpha=0.5, label=f"Predictions", color="#ff68b4")
+    plt.plot([min(actual_values), max(actual_values)], [min(actual_values), max(actual_values)], 'k--', label='Ideal Fit')
+    plt.text(0, 92, f"$R^2$={r2:.2f}", fontsize=32)
+    # Adjusting font sizes and setting font properties
+    plt.xlabel(f'AlphaFold-pLDDT',size=32)
+    plt.ylabel(f'FusOn-pLM-Diso',size=32)
+    plt.title(f"% Disordered (n={len(actual_values)})",size=32)
+    plt.xticks(fontsize=24)
+    plt.yticks(fontsize=24)
+    #plt.xlabel("Values")
+    #plt.ylabel("Density")
+    #plt.title(f"Density Plot of Actual vs Predicted Values (R^2 = {r2:.2f})")
+    plt.legend(prop={'size': 16})
+    plt.tight_layout()
+    plt.show()
+    plt.savefig(save_path, dpi=300)
+    
+def main():
+    set_font()
+    #output_dir = "results/test"
+    output_dir = "results/final"
+    seq_label_dict = pd.read_csv('splits/test_df.csv')
+    seq_ids_dict = dict(zip(seq_label_dict['Sequence'],seq_label_dict['IDs']))
+    seq_label_dict = dict(zip(seq_label_dict['Sequence'],seq_label_dict['Label']))
+    best_caid_model_results = pd.read_csv(f"{output_dir}/best_caid_model_results.csv")
+    make_auroc_curve(results_dir=output_dir, 
+                    seq_label_dict=seq_label_dict,
+                    seq_ids_dict=seq_ids_dict,
+                    path_to_results_of_interest="trained_models/fuson_plm/best/caid_hyperparam_screen_test_probs.csv", 
+                    model_alias="FusOn-pLM", 
+                    path_to_esm_results="trained_models/esm2_t33_650M_UR50D/best/caid_hyperparam_screen_test_probs.csv",
+                    with_rankings=True)
+        
+    caid2_test_data = pd.read_csv(f"splits/splits.csv")
+    caid2_test_data = caid2_test_data.loc[caid2_test_data['Split']=='Test']
+    caid2_test_labels = caid2_test_data['Label'].tolist()
+    caid2_test_ids = caid2_test_data['IDs'].tolist()
+    # fusions, heads, and tails
+    fusion_ht_data = prep_data_for_ht_disorder_comparison()
+    os.makedirs("processed_data/figures",exist_ok=True)
+
+    head_data = fusion_ht_data.drop_duplicates(['hg_seq']).reset_index(drop=True)
+    head_labels = head_data['hg_label'].tolist()
+    head_ids = head_data['hgUniProt'].tolist()
+    tail_data = fusion_ht_data.drop_duplicates(['tg_seq']).reset_index(drop=True)
+    tail_labels = tail_data['tg_label'].tolist()
+    tail_ids = tail_data['tgUniProt'].tolist()
+    fusion_data = fusion_ht_data.drop_duplicates(['Fusion_Seq']).reset_index(drop=True)
+    fusion_labels = fusion_data['fusion_label'].tolist()
+    fusion_ids = fusion_data['seq_id'].tolist()
+    
+    plt.rc('text', usetex=False)
+    math_part = r"$n$"
+
+    os.makedirs("processed_data/figures/histograms",exist_ok=True)
+    plot_disorder_content_hist(caid2_test_labels, caid2_test_ids, title=f"CAID2 Disorder-NOX ({math_part}={len(caid2_test_labels):,})", color="black", savepath='processed_data/figures/histograms/disorder_nox_histogram.png')
+    plot_disorder_content_hist(head_labels, head_ids, title=f"Head Proteins ({math_part}={len(head_labels):,})", color="#df8385", savepath='processed_data/figures/histograms/heads_histogram.png')
+    plot_disorder_content_hist(tail_labels, tail_ids, title=f"Tail Proteins ({math_part}={len(tail_labels):,})", color="#6ea4da", savepath='processed_data/figures/histograms/tails_histogram.png')
+    plot_disorder_content_hist(fusion_labels, fusion_ids, title=f"Fusion Oncoproteins ({math_part}={len(fusion_labels):,})", color="mediumpurple", savepath='processed_data/figures/histograms/fusions_histogram.png')
+    
+    os.makedirs("processed_data/figures/fusion_disorder",exist_ok=True)
+    make_all_favorite_fusion_pLDDT_plots([
+                                        "EWSR1::FLI1",
+                                        "PAX3::FOXO1",
+                                        "EML4::ALK",
+                                        "SS18::SSX1"],
+                                            left_to_right=True)
+
+if __name__ == "__main__":
+    main()

fuson_plm/benchmarking/caid/process_fusion_structures.py

@@ -0,0 +1,799 @@
+# Process fusion structures and the structures of their head and tail proteins for pLDDTs 
+
+import requests
+import json
+import pandas as pd
+import numpy as np
+
+import requests
+import re
+import os
+import shutil
+
+from Bio.PDB import MMCIFParser
+import Bio.PDB as PDB
+from Bio import pairwise2
+from Bio.pairwise2 import format_alignment
+from bs4 import BeautifulSoup
+import pdb
+
+from fuson_plm.utils.logging import log_update, open_logfile
+
+#@markdown Define AlphaFoldStructure class
+class AlphaFoldStructure:
+    '''
+    This class processes an mmCIF file, either uploaded or downloaded from the AlphaFold2 database, to provide comprehensive information.
+    '''
+    def __init__(self, fold_path=None, uniprot_to_download=None, uniprot_output_dir= None, secondary_structure_types=None):
+        # If the user provided a PDB path, convert their file to mmcif. Isolate the suffix
+        if fold_path is not None:
+          fold_fname = fold_path.split('/')[-1]
+          prefix, suffix = fold_fname.split('.')
+
+          if suffix == 'pdb': # convert to cif
+            # make a directory for converted cif files
+            conversion_path = 'mmcif_converted_files'
+            if not(os.path.exists(conversion_path)):
+              os.makedirs(conversion_path)
+
+            fold_path = self.__convert_pdb_to_mmcif__(fold_path, f'{conversion_path}/{prefix}.cif')
+
+        self.file_path = fold_path
+
+        # If user provided a uniprot ID to download, download it and save it as the file path so it can be processed
+        if uniprot_to_download is not None:
+            if fold_path is not None:
+              log_update("WARNING: both a fold_path and a uniprot_to_download were provided. Running default: downloading the CIF file for provided UniProt ID.")
+            self.file_path = self.__download_mmCIF(uniprot_to_download, output_path=uniprot_output_dir)
+
+        # Either they provide acceptable secondary structure types, or query the internet for them
+        if secondary_structure_types is None:
+          self.secondary_structure_types = self.__pull_secondary_structure_types()
+        else:
+          self.secondary_structure_types = secondary_structure_types
+
+        # If there's a CIF file, initialize the object
+        if self.file_path:
+          self.cif_lines = self.__parse_cif()
+          self.secondary_structures = self.__extract_secondary_structures()
+          self.structure_dict = self.__calc_pLDDTs()
+          self.sequence = self.structure_dict['seq']
+          self.plddts = self.structure_dict['res_pLDDTs']
+          self.avg_pLDDT = self.structure_dict['avg_pLDDT']
+          self.residues_df = self.__create_residues_summary_dataframe()
+          self.secondary_structures_df = self.__create_secondary_structures_summary_dataframe()
+        # Otherwise, print an error.
+        else:
+          log_update("ERROR: structure could not be created. No CIF file found.")
+
+    def __convert_pdb_to_mmcif__(self, pdb_filename, mmcif_filename):
+      parser = PDB.PDBParser()
+      structure = parser.get_structure('structure', pdb_filename)
+
+      io = PDB.MMCIFIO()
+      io.set_structure(structure)
+      io.save(mmcif_filename)
+      return mmcif_filename
+
+    def __download_mmCIF(self, uniprot_id, output_path=None):
+        '''
+        Download mmCIF file with provided uniprot_id and optional output_path for the downloaded file.
+
+        Return: path to downloaded file if successful, None otherwise
+        '''
+        full_file_name = f"AF-{uniprot_id}-F1-model_v4.cif"     # define file name that will be found on the AlphaFold2 database.
+        # if output path not provided, just save locally under full_file_name
+        if output_path is None:
+            output_path = full_file_name
+        else:
+            output_path = f"{output_path}/{full_file_name}"
+
+        # request the URL for the file
+        url = f"https://alphafold.ebi.ac.uk/files/{full_file_name}"
+        response = requests.get(url)
+
+        if response.status_code == 200:
+            with open(output_path, 'wb') as file:
+                file.write(response.content)
+            #log_update(f"File downloaded successfully and saved as {output_path}")
+        else:
+            log_update(f"Failed to download file. Status code: {response.status_code}")
+            return None
+
+        return output_path
+
+    def __pull_secondary_structure_types(self):
+        '''
+        Pull a dictionary of secondary structure types and their descriptions from the PDB mmCIF website (necessary for annotating the CIF file)
+        Only called if the user does not provide such a dictionary themselves.
+        '''
+
+        # request the .html tree from the website with all secondary structure terms
+        url = "https://mmcif.wwpdb.org/dictionaries/mmcif_pdbx_v50.dic/Items/_struct_conf_type.id.html"
+        response = requests.get(url)
+
+        if response.status_code != 200:
+            raise Exception("Failed to retrieve mmCIF dictionary")
+
+        # Parse the response content
+        soup = BeautifulSoup(response.content, 'html.parser')
+
+        # Debug: Print the soup to understand the structure
+        # log_update(soup.prettify())
+        # write the prettified soup to a txt file
+        #with open('mmcif_dictionary.txt', 'w') as f:
+        #    f.write(soup.prettify())
+
+        # Find the h4 header with the class "panel-title" and text "Controlled Vocabulary"
+        header = soup.find('h4', class_='panel-title')
+        if header is None or 'Controlled Vocabulary' not in header.text:
+            raise Exception("Could not find the 'Controlled Vocabulary' header")
+
+        # Debug: Print the found header
+        #log_update(f"Found header: {header}")
+
+        # The table should be the next sibling of the header
+        table = header.find_next('table')
+        if table is None:
+            raise Exception("Could not find the table following the 'Controlled Vocabulary' header")
+
+        # Debug: Print the found table (only the opening <table> tag)
+        #log_update(f"Found table (showing header line): {str(table).split('<thead')[0]}")
+
+        # Iterate through rows in the table and process each entry
+        secondary_structure_types = {}
+        rows = table.find_all('tr')
+        for row in rows[1:]:  # Skip the header row
+            cols = row.find_all('td')
+            if len(cols) > 1:
+                type_id = cols[0].text.strip()
+                description = cols[1].text.replace('\t', ' ').strip()
+
+                # Replace multiple spaces with a single space
+                description = re.sub(' +', ' ', description)
+
+                # If this is a protein secondary structure (the table also contains nucleic acid structures), add it to teh dictionary
+                if '(protein)' in description:
+                  secondary_structure_types[type_id] = description
+
+        return secondary_structure_types
+
+    def get_secondary_structure_types(self):
+      '''
+      Display secondary structure types
+      '''
+      log_update("Secondary Structure Types in mmCIF files:")
+      for ss_type, description in self.secondary_structure_types.items():
+          log_update(f"{ss_type}: {description}")
+
+      return self.secondary_structure_types
+
+    def __parse_cif(self):
+        '''
+        Read cif file lines from self.file_path
+        '''
+        with open(self.file_path, 'r') as file:
+            lines = file.readlines()
+        return lines
+
+    def __extract_secondary_structures(self):
+        '''
+        Iterate through the lines of the cif files to find each secondary structure.
+        Returns a tuple for each amino acid that has a secondary structure annotation. Tuple contains:
+          1. Structure Type (e.g. STRN)
+          2. Structure ID (e.g. STRN1)
+          3. Description (e.g. beta strand)
+          4. Position (e.g. 3)
+        '''
+        secondary_structures = []
+        parsing_secondary_structure = False
+
+        # iterate throhugh cif lines
+        for line in self.cif_lines:
+            # hone in on the right section of the cif file
+            if line.startswith("_struct_conf.conf_type_id"):
+                parsing_secondary_structure = True
+                continue
+            # if we're in the right section...
+            if parsing_secondary_structure:
+                if line.startswith("#"):
+                    parsing_secondary_structure = False   # no longer in the right section
+                    continue
+                # still in the right section
+                columns = line.split()
+                # iterate through columns to find each piece of info we need
+                if len(columns) >= 7:
+                    sec_struc_type = columns[6]
+                    sec_struc_id = columns[13]
+                    start_res = int(columns[2])
+                    end_res = int(columns[9])
+                    sec_struc_name = self.secondary_structure_types.get(sec_struc_type, 'Unknown')
+                    # make tuple for this position in the sequence
+                    for pos in range(start_res, end_res + 1):
+                        secondary_structures.append((sec_struc_type, sec_struc_id, sec_struc_name, pos))
+
+        return secondary_structures
+
+    def __calc_pLDDTs(self):
+        '''
+        This method iterates through the cif file to return a dictionary with a few key pieces of info:
+          1. Sequence
+          2. pLDDTs for each residue
+          3. Average pLDDT
+        '''
+
+        # define dictionary needed to translate into single-letter AA code
+        aa_dict = {
+            "ALA": "A", "CYS": "C", "ASP": "D", "GLU": "E", "PHE": "F",
+            "GLY": "G", "HIS": "H", "ILE": "I", "LYS": "K", "LEU": "L",
+            "MET": "M", "ASN": "N", "PRO": "P", "GLN": "Q", "ARG": "R",
+            "SER": "S", "THR": "T", "VAL": "V", "TRP": "W", "TYR": "Y"
+        }
+
+        parser = MMCIFParser(QUIET=True)    # create a parser
+        data = parser.get_structure("structure", self.file_path)    # parse structure
+
+        # count models and chains (should be 1 model and 1 chain; don't use this class to parse a complex)
+        model = data.get_models()
+        models = list(model)
+        chains = list(models[0].get_chains())
+
+        # iterate through the chains and get amino acid letters and pLDDTs
+        all_pLDDTs = []
+        for n in range(len(chains)):
+            chainname = chr(n + 65)     # turn chain number into letter (e.g. 1 --> "A" so we have Chain A instead of Chain 1)
+            residues = list(chains[n].get_residues())   # extract all residues
+            seq = ''
+            pLDDTs = [0] * len(residues)    # initialize empty pLDDT array for this chain
+
+            # iterate through all residues in this chain
+            for i in range(len(residues)):
+                r = residues[i]
+                # which amino acid is here?
+                try:
+                    seq += aa_dict[r.get_resname()]
+                # error if it's not a real amino acid
+                except KeyError:
+                    log_update('residue name invalid')
+                    break
+
+                # look at each atom. Get its pLDDT (bfactor). make sure bfactor for all atoms within one residue are equal.
+                atoms = list(r.get_atoms())
+                bfactor = atoms[0].get_bfactor()
+                for a in range(len(atoms)):
+                    # if not all atoms within an AA have the same pLDDT, error.
+                    if atoms[a].get_bfactor() != bfactor:
+                        break
+
+                pLDDTs[i] = bfactor   # add pLDDT for this residue to the list.
+
+            all_pLDDTs.extend(pLDDTs) # add pLDDTs for this chain to list of all pLDDTs
+
+        avg_pLDDT = np.mean(all_pLDDTs) # average pLDDTs across all chains
+        return_dict = {
+            'avg_pLDDT': round(avg_pLDDT, 2),
+            'res_pLDDTs': all_pLDDTs,
+            'seq': seq
+        }
+        return return_dict
+
+    def __create_residues_summary_dataframe(self):
+        '''
+        Create a dataframe that summarizes the secondary structure information for each residue.
+        Columns:
+          1. Position: amino acid position (e.g. 3)
+          2. Residue: amino acid 1-letter code (e.g. A)
+          3. pLDDT: alphafold2's pLDDT score for this residue to 2 decimal places (e.g. 77.54)
+          4. Structure Type: type of secondary structure (e.g. STRN)
+          5. Structure ID: ID of this secondary structure (e.g. STRN1)
+          5. Description: description of this secondary structure (e.g. beta strand)
+          6. Disordered: is this residue disordered or not? A residue is not disordered if it's in a HELX or STRN. (True/False)
+
+        '''
+        # Convert the secondary structures to a dataframe
+        df_secondary_structures = pd.DataFrame(self.secondary_structures, columns=['Structure Type', 'Structure ID', 'Description', 'Position'])
+
+        # Add Residue and pLDDT columns to the dataframe
+        df_temp = pd.DataFrame(
+            data={
+                'Position': list(range(1, len(self.sequence) + 1)),
+                'Residue': list(self.sequence),
+                'pLDDT': self.plddts
+            })
+
+        df_secondary_structures = pd.merge(df_secondary_structures, df_temp, on='Position', how='right')
+        # Determine if each residue is disordered or not based on what Structure Type it's in. If helix or strand, it's ordered. If anything else or NaN, it's disordered.
+        df_secondary_structures['Disordered'] = df_secondary_structures['Structure Type'].apply(
+            lambda x: False if (type(x)==str and (('HELX' in x) or ('STRN' in x))) else True
+        )
+
+        return df_secondary_structures
+
+    def __create_secondary_structures_summary_dataframe(self):
+        '''
+        Create a dataframe grouped by each Structure ID, providing a summary of each secondary structure in the chain.
+        Columns:
+          1. Structure ID: ID of this secondary structure (e.g. STRN1)
+          2. Start: start position of this secondary structure (e.g. 3)
+          3. End: end position of this secondary structure (e.g. 12)
+          4. Start Residue: amino acid 1-letter code of the start position (e.g. A)
+          5. End Residue: amino acid 1-letter code of the end position (e.g. L)
+          6. Disordered: is this residue disordered or not? A residue is not disordered if it's in a HELX or STRN. (True/False)
+          7. Description: description of this secondary structure (e.g. beta strand)
+          8. Structure Type: type of secondary structure (e.g. STRN)
+          9. avg_pLDDT: average pLDDT for this secondary structure (e.g. 77.54)
+        '''
+
+        # Apply groupby on self.residues_df to reorganize it by Structure ID
+        secondary_structures_df = self.residues_df.groupby('Structure ID').agg({
+            'Position': ['first', 'last'],
+            'Residue': ['first','last'],
+            'Disordered': 'first',
+            'Description': 'first',
+            'Structure Type': 'first',
+            'pLDDT': 'mean'
+        }).reset_index()
+
+        # Flatten the multi-level columns
+        secondary_structures_df.columns = ['Structure ID', 'Start', 'End', 'Start Residue', 'End Residue', 'Disordered', 'Description', 'Structure Type', 'avg_pLDDT']
+        secondary_structures_df['avg_pLDDT'] = secondary_structures_df['avg_pLDDT'].round(2)
+
+        # Display the summarized DataFrame
+        return secondary_structures_df
+
+    def get_residues_df(self):
+        return self.residues_df
+
+    def get_secondary_structures_df(self):
+        return self.secondary_structures_df
+
+    def get_full_sequence(self):
+        return ''.join([res for res in self.residues_df['Residue']])
+
+    def get_average_plddt(self):
+        plddt_values = [plddt for plddt in self.residues_df['pLDDT'] if plddt is not None]
+        return sum(plddt_values) / len(plddt_values) if plddt_values else None
+
+def pull_secondary_structure_types():
+    url = "https://mmcif.wwpdb.org/dictionaries/mmcif_pdbx_v50.dic/Items/_struct_conf_type.id.html"
+    response = requests.get(url)
+
+    if response.status_code != 200:
+        raise Exception("Failed to retrieve mmCIF dictionary")
+
+    soup = BeautifulSoup(response.content, 'html.parser')
+
+    # Debug: Print the soup to understand the structure
+    # log_update(soup.prettify())
+    # write the prettified soup to a txt file
+    with open('mmcif_dictionary.txt', 'w') as f:
+        f.write(soup.prettify())
+
+    # Find the h4 header with the class "panel-title" and text "Controlled Vocabulary"
+    header = soup.find('h4', class_='panel-title')
+    if header is None or 'Controlled Vocabulary' not in header.text:
+        raise Exception("Could not find the 'Controlled Vocabulary' header")
+
+    # Debug: Print the found header
+    #log_update(f"Found header: {header}")
+
+    # The table should be the next sibling of the header
+    table = header.find_next('table')
+    if table is None:
+        raise Exception("Could not find the table following the 'Controlled Vocabulary' header")
+
+    # Debug: Print the found table (only the opening <table> tag)
+    #log_update(f"Found table (showing header line): {str(table).split('<thead')[0]}")
+
+    secondary_structure_types = {}
+    rows = table.find_all('tr')
+    for row in rows[1:]:  # Skip the header row
+        cols = row.find_all('td')
+        if len(cols) > 1:
+            type_id = cols[0].text.strip()
+            description = cols[1].text.replace('\t', ' ').strip()
+
+            # Replace multiple spaces with a single space
+            description = re.sub(' +', ' ', description)
+
+            if '(protein)' in description:
+              secondary_structure_types[type_id] = description
+
+    return secondary_structure_types
+
+# Process structures downloaded from FusionPDB 
+def process_fusionpdb_fusion_files(files, level_2_3_structure_info, folder, save_path=None):
+    # get secondary structure types so we can process PDBs
+    secondary_structure_types = pull_secondary_structure_types()
+
+    # Initialize 3 columns to store structural info - the AA seq in the fold (should match), the Avg pLDDT, and the per-residue pLDDTs (comma-separated, 2 decimal pts.)
+    level_2_3_structure_info['Fold AA seq'] = ['']*len(level_2_3_structure_info)
+    level_2_3_structure_info['Avg pLDDT'] = [0]*len(level_2_3_structure_info)
+    level_2_3_structure_info['pLDDTs'] = ['']*len(level_2_3_structure_info)
+    
+    # pre-loop processed
+    pre_loop_processed = []
+    if os.path.exists(save_path):
+        pre_loop_processed = pd.read_csv(save_path)
+        pre_loop_processed = pre_loop_processed['Structure Link'].tolist()
+        pre_loop_processed = [x.split('/')[-1] for x in pre_loop_processed]
+        log_update(f"Total structures already processed: {len(pre_loop_processed)}")
+    
+    log_update("\nProcessing fusion structures...")
+    # only process structures we haven't processed yet
+    for i, structure in enumerate(files):
+        log_update(f'\tProcessing #{i+1}: {structure}')
+        
+        # make sure we haven't already processed it and aren't wasting time
+        if structure in pre_loop_processed:
+            log_update(f"\t\tAlready processed. Continuing...")
+            continue
+        
+        # create AlphaFoldStructure object
+        obj = AlphaFoldStructure(fold_path=f'{folder}/{structure}', secondary_structure_types=secondary_structure_types)
+        aa_seq = obj.get_full_sequence()
+        avg_plddt = obj.get_average_plddt()
+        residues_df = obj.get_residues_df()
+        all_plddts = ",".join(residues_df['pLDDT'].astype(str).tolist())
+        
+        log_update(f"\t\tAvg pLDDT: {round(avg_plddt,2)}\tFold AA seq: {aa_seq}\tFirst 5 pLDDTs: {','.join(all_plddts.split(',')[0:5])}")
+
+        level_2_3_structure_info.loc[level_2_3_structure_info['Structure Link'].str.contains(f"/{structure}"), 'Fold AA seq'] = aa_seq
+        level_2_3_structure_info.loc[level_2_3_structure_info['Structure Link'].str.contains(f"/{structure}"), 'Avg pLDDT'] = avg_plddt
+        level_2_3_structure_info.loc[level_2_3_structure_info['Structure Link'].str.contains(f"/{structure}"), 'pLDDTs'] = all_plddts
+        
+        # write level_2_3_structure_info to csv
+        cur_df = level_2_3_structure_info.loc[level_2_3_structure_info['Structure Link'].str.contains(f"/{structure}")].reset_index(drop=True)
+        if os.path.exists(save_path):
+            cur_df.to_csv(save_path,mode='a',header=False,index=False)
+        else:
+            cur_df.to_csv(save_path,index=False)
+            
+    # now reload the completed dataframe
+    level_2_3_structure_info = pd.read_csv(save_path)
+    return level_2_3_structure_info
+
+def process_fusionpdb_head_tail_files(ht, save_path='heads_and_tails_structures_processed.csv'):
+    # ht is a list of head and tail proteins we have to process.  
+    log_update("\nProcessing head and tail structures...")
+    
+    # get secondary structure types so we can process PDBs
+    secondary_structure_types = pull_secondary_structure_types()
+    
+    # make directory to save alphafold DB structures of heads and tails
+    os.makedirs('raw_data/fusionpdb/head_tail_af2db_structures',exist_ok=True)
+    
+    # pre-loop processed
+    pre_loop_processed = []
+    if os.path.exists(save_path):
+        pre_loop_processed = pd.read_csv(save_path)
+        pre_loop_processed = pre_loop_processed['UniProtID'].tolist()
+        log_update(f"Heads and tails already processed: {len(pre_loop_processed)}")
+        
+    ht_structures_df = pd.DataFrame(
+        data = {
+            'UniProtID': ['']*len(ht),
+            'Avg pLDDT': ['']*len(ht),
+            'All pLDDTs': ['']*len(ht),
+            'Seq': ['']*len(ht)
+        }
+    )
+    
+    for i, uniprotid in enumerate(ht):
+        log_update(f'\tProcessing #{i+1}: {uniprotid}')
+        aa_seq, avg_plddt, all_plddts = None, None, None
+        
+        # make sure we haven't processed it yet!
+        if uniprotid in pre_loop_processed:
+            log_update(f"\t\tAlready processed. Continuing")
+            continue
+
+        try:
+            obj = AlphaFoldStructure(uniprot_to_download=uniprotid, secondary_structure_types=secondary_structure_types,
+                                     uniprot_output_dir='raw_data/fusionpdb/head_tail_af2db_structures')
+            aa_seq = obj.get_full_sequence()
+            avg_plddt = obj.get_average_plddt()
+            residues_df = obj.get_residues_df()
+            all_plddts = ",".join(residues_df['pLDDT'].astype(str).tolist())
+
+            log_update(f"\t\tAvg pLDDT: {round(avg_plddt,2)}\tFold AA seq: {aa_seq}\tFirst 5 pLDDTs: {','.join(all_plddts.split(',')[0:5])}")
+
+        except:
+            log_update(f"\t\tAvg pLDDT: {None}\tFold AA seq: {None}\tFirst 5 pLDDTs: {None}")
+            
+        # Fill in info for combined ht df
+        ht_structures_df.loc[i, 'UniProtID'] = uniprotid
+        ht_structures_df.loc[i, 'Avg pLDDT'] = avg_plddt
+        ht_structures_df.loc[i, 'All pLDDTs'] = all_plddts
+        ht_structures_df.loc[i, 'Seq'] = aa_seq
+        
+        # write level_2_3_structure_info to csv
+        cur_df = pd.DataFrame(ht_structures_df.iloc[i,:]).T.reset_index(drop=True)
+        if os.path.exists(save_path):
+            cur_df.to_csv(save_path,mode='a',header=False,index=False)
+        else:
+            cur_df.to_csv(save_path,index=False)
+    
+    # ensure we got everything
+    ht_structures_df = pd.read_csv(save_path)
+    level_2_3 = pd.read_csv(f'processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_head_tail_info.csv')
+    level_2_3['FusionGene'] = level_2_3['FusionGene'].str.replace('-','::') 
+    heads = level_2_3['HGUniProtAcc'].tolist()
+    tails = level_2_3['TGUniProtAcc'].tolist()
+    ht = heads + tails
+    ht = set([x for x in ht if type(x)==str])
+    ht = set(','.join(ht).split(','))
+
+    log_update(f"total heads and tails: {len(ht)}")
+    log_update(f"total processed: {len(ht_structures_df)}\t{len(ht_structures_df['UniProtID'].unique())}")
+
+    # which ones are missing?
+    missing = set(ht) - set(ht_structures_df['UniProtID'].unique())
+    log_update(f"missing: {len(missing)}")
+    log_update(missing)
+    
+    # Some heads and tails are not in the alphxwafold database. I folded these myself, externally. 
+    ht_structures_df = ht_structures_df.replace('',np.nan)
+    need_to_fold = ht_structures_df[ht_structures_df['Avg pLDDT'].isna()]['UniProtID'].tolist()
+    with open('processed_data/fusionpdb/intermediates/uniprotids_not_in_afdb.txt','w') as f:
+        for uniprotid in need_to_fold:
+            f.write(f'{uniprotid}\n')
+
+    idmap = pd.read_csv(f'raw_data/fusionpdb/not_in_afdb_idmap.txt',sep='\t')
+    idmap = idmap[idmap['Entry'].isin(need_to_fold)].reset_index(drop=True)
+    idmap = idmap[['Entry','Sequence']].rename(columns={
+                                                        'Entry': 'ID'})
+    idmap['Length'] = idmap['Sequence'].apply(len)
+
+    log_update("Investigating heads and tails that were not in the AF2 database:")
+    log_update(f"\tMin length: {min(idmap['Length'])}")
+    log_update(f"\tMax length: {max(idmap['Length'])}")
+    idmap = idmap.sort_values(by='Length',ascending=True).reset_index(drop=True)
+    
+    # Q9NNW7
+    id='Q9NNW7'
+    if id in idmap['ID'].tolist():
+        ht_structures_df.loc[
+            ht_structures_df['UniProtID']=='Q9NNW7', 'Avg pLDDT'
+        ] = 91.68
+        ht_structures_df.loc[
+            ht_structures_df['UniProtID']=='Q9NNW7', 'Seq'
+        ] = idmap.loc[
+            idmap['ID']=='Q9NNW7', 'Sequence'
+        ].item()
+    
+    ## Q16881
+    id='Q16881'
+    if id in idmap['ID'].tolist():
+        ht_structures_df.loc[
+            ht_structures_df['UniProtID']==id, 'Avg pLDDT'
+        ] = 89.55
+        ht_structures_df.loc[
+            ht_structures_df['UniProtID']==id, 'Seq'
+        ] = idmap.loc[
+            idmap['ID']==id, 'Sequence'
+        ].item()
+    
+    # Q86V15
+    id='Q86V15'
+    if id in idmap['ID'].tolist():
+        ht_structures_df.loc[
+            ht_structures_df['UniProtID']==id, 'Avg pLDDT'
+        ] = 48.14
+        ht_structures_df.loc[
+            ht_structures_df['UniProtID']==id, 'Seq'
+        ] = idmap.loc[
+            idmap['ID']==id, 'Sequence'
+        ].item()
+            
+    return ht_structures_df
+
+def process_fusions_and_hts():
+    # Process the structures of fusion proteins downloaded from FusionPDB
+    level_2_3_structure_info_og = pd.read_csv('processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structure_links.csv')
+
+    # figure out which ones we have
+    folder = 'raw_data/fusionpdb/structures'
+    # get all the structure files in folder
+    files = os.listdir(folder)
+    log_update(f"total pdbs: {len(files)}")
+    log_update(f"examples: {files[:5]}")
+    
+    os.makedirs('processed_data/fusionpdb', exist_ok=True)  
+    
+    # process the full fusion pdbs
+    level_2_3_structure_info = process_fusionpdb_fusion_files(files, level_2_3_structure_info_og, folder, save_path='processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_structures_processed.csv')  
+    
+    # process the head and tail pdbs
+    level_2_3 = pd.read_csv(f'processed_data/fusionpdb/intermediates/giant_level2-3_fusion_protein_head_tail_info.csv') 
+    level_2_3['FusionGene'] = level_2_3['FusionGene'].str.replace('-','::') 
+    # Get the heads and tails, see how many unique proteins we're working with
+    heads = level_2_3['HGUniProtAcc'].tolist()
+    tails = level_2_3['TGUniProtAcc'].tolist()
+    ht = heads + tails
+    ht = set([x for x in ht if type(x)==str])
+    ht = set(','.join(ht).split(','))
+    log_update(f"Unique heads/tails: {len(ht)}")
+    
+    heads_tails_analyzed = process_fusionpdb_head_tail_files(list(ht), save_path='processed_data/fusionpdb/heads_tails_structural_data.csv') 
+    
+    # In the level_2_3 database, we only have the fusions with documented heads and tails.
+    level_2 = pd.read_csv(f'raw_data/fusionpdb/FusionPDB_level2_curated_09_05_2024.csv')
+    level_3 = pd.read_csv(f'raw_data/fusionpdb/FusionPDB_level3_curated_09_05_2024.csv')
+    joined_23 = pd.concat([level_2,level_3]).reset_index(drop=True)
+    joined_23['FusionGene'] = joined_23['FusionGene'].str.replace('-','::')     # use new notation with head::tail
+    log_update(f"\nnumber of duplicated fusion gene rows: {len(joined_23[joined_23['FusionGene'].duplicated()])}")
+    # make the dictionary
+    fo_gid_dict = dict(zip(joined_23['FusionGene'],joined_23['FusionGID']))
+    log_update(len(fo_gid_dict))
+
+    # let's clean giant level 2 and level 3
+    # first, drop anyting where Fold AA seq is nan. there is no fold.
+    level_2_3_structure_info_clean = level_2_3_structure_info.replace('',np.nan)    # make sure there are nans where there should be
+    level_2_3_structure_info_clean = level_2_3_structure_info_clean.dropna(subset=['Fold AA seq']).reset_index(drop=True)
+    log_update(f"length of processed structure file: {len(level_2_3_structure_info_clean)}")
+    level_2_3_structure_info_clean['pLDDT'] = level_2_3_structure_info_clean['Avg pLDDT'].round(2)
+    level_2_3_structure_info_clean = level_2_3_structure_info_clean.drop(columns=['Avg pLDDT'])
+    level_2_3_structure_info_clean['FusionGene'] = level_2_3_structure_info_clean['FusionGene'].str.replace('-','::')
+    level_2_3_structure_info_clean['FusionGID'] = level_2_3_structure_info_clean['FusionGene'].apply(lambda x: fo_gid_dict[x])
+    
+    # now let's use the FusionPDB database we processed as ground truth for sequence, rather than the webpage
+    log_update("Using FusionPDB as ground truth for sequences...")
+    raw_download = pd.read_csv('../../data/raw_data/FusionPDB.txt',sep='\t',header=None)
+    raw_download['FusionGene'] = raw_download[7]+ '::' + raw_download[11]
+    raw_download = raw_download.rename(columns={18:'Raw Download AA Seq'})
+    log_update(f"FusionPDB raw download size: {len(raw_download)}")
+
+    level_2_3_structure_info_clean_ids = set(level_2_3_structure_info_clean['FusionGene'].tolist())
+    level_2_3_structure_info_clean_seqs = set(level_2_3_structure_info_clean['Fold AA seq'].tolist())
+    raw_download_ids = set(raw_download['FusionGene'].tolist())
+    raw_download_seqs = set(raw_download['Raw Download AA Seq'].tolist())
+    log_update(f"Number of overlapping gene IDs: {len(level_2_3_structure_info_clean_ids.intersection(raw_download_ids))}")
+    log_update(f"Number of overlapping sequences: {len(level_2_3_structure_info_clean_seqs.intersection(raw_download_seqs))}")
+    # attempt a merge on Raw Download AA Seq with both. ofthe ogs
+    # Merging with the AlphaFold sequence
+    test_merge_1 = pd.merge(
+        level_2_3_structure_info_clean.rename(columns={'Fold AA seq': 'Raw Download AA Seq'}),
+        raw_download,
+        on=['FusionGene','Raw Download AA Seq'],
+        how='inner'
+    )
+    test_merge_1 = test_merge_1.drop(columns=['AA seq'])
+    test_merge_1['Seq Source'] = ['AlphaFold,Raw Download']*len(test_merge_1)
+    log_update(f"Merge on AlphaFold AA Seq and raw Download AA Seq. len={len(test_merge_1)}")
+    # Merging with the webpage sequence
+    test_merge_2 = pd.merge(
+        level_2_3_structure_info_clean.rename(columns={'AA seq': 'Raw Download AA Seq'}),
+        raw_download,
+        on=['FusionGene','Raw Download AA Seq'],
+        how='inner'
+    )
+    test_merge_2 = test_merge_2.drop(columns=['Fold AA seq'])
+    test_merge_2['Seq Source'] = ['Webpage,Raw Download']*len(test_merge_2)
+    log_update(f"Merge on Webpage AA Seq and Raw Download AA Seq. len={len(test_merge_2)}")
+
+    test_merge = pd.concat([test_merge_1,test_merge_2])
+    test_merge['Len(AA seq)'] = test_merge['Raw Download AA Seq'].apply(lambda x: len(x))
+    # drop duplicates
+    test_merge = test_merge.drop_duplicates().reset_index(drop=True) 
+    
+    # for anything that has a CIF, keep the CIF
+    log_update(f"len test_merge before keeping CIFs over identical PDBs: {len(test_merge)}")
+    test_merge = test_merge.sort_values(by='Structure Type',ascending=True).reset_index(drop=True).groupby(['Hgene', 'Hchr', 'Hbp', 'Hstrand', 'Tgene', 'Tchr',
+        'Tbp', 'Tstrand', 'Len(AA seq)', 'FusionGene',
+        'Level', 'Raw Download AA Seq', 'pLDDT', 'pLDDTs','FusionGID', 'Seq Source']).agg(
+            {
+                'Structure Link': 'first',
+                'Structure Type': 'first'
+            }
+        ).reset_index()
+    log_update(f"len after: {len(test_merge)}")
+    
+    # for anything with multiple seq sources, concatenate them
+    log_update(f"len test_merge before combining seq sources: {len(test_merge)}")
+    test_merge = test_merge.groupby(['Structure Link','Hgene', 'Hchr', 'Hbp', 'Hstrand', 'Tgene', 'Tchr',
+        'Tbp', 'Tstrand', 'Len(AA seq)', 'FusionGene','Structure Type',
+        'Level', 'Raw Download AA Seq', 'pLDDT', 'pLDDTs', 'FusionGID', ]).agg(
+            {
+                'Seq Source': lambda x: ','.join(x)
+            }
+        ).reset_index()
+    test_merge['Seq Source'] = test_merge['Seq Source'].apply(lambda x: ','.join(set(x.split(','))))
+    log_update(f"len after: {len(test_merge)}")
+    
+    # are there cases of multiple folds for the same sequence? miraculously, yes!
+    dup_seqs = test_merge[test_merge['Raw Download AA Seq'].duplicated()]['Raw Download AA Seq'].unique().tolist()
+
+    # for anything with multiple folds / seq, randomly choose the first one. 
+    log_update(f"len test_merge before randomly choosing first fold when one seq has multiple folds: {len(test_merge)}")
+    test_merge = test_merge.groupby(['Hgene', 'Hchr', 'Hbp', 'Hstrand', 'Tgene', 'Tchr',
+        'Tbp', 'Tstrand', 'Len(AA seq)', 'FusionGene',
+        'Level', 'Raw Download AA Seq', 'FusionGID', ]).agg(
+            {
+                'Structure Link': 'first',
+                'Structure Type': 'first',
+                'Seq Source': 'first',
+                'pLDDT': 'first',
+                'pLDDTs': 'first'
+            }
+        ).reset_index()
+    log_update(f"len after: {len(test_merge)}")
+    
+    # how many columns DO NOT have the right AlphaFold sequences?
+    source_str = test_merge['Seq Source'].value_counts().reset_index().rename(columns={'index': 'Seq Source','Seq Source': 'count'}).to_string(index=False)
+    source_str = "\t\t" + source_str.replace("\n","\n\t\t")
+    log_update(f"Distribution of sequence sources:\n{source_str}")
+    
+    # dropping anything where AF sequence is wrong. Don't want to use these.
+    test_merge = test_merge.loc[test_merge['Seq Source'].str.contains('AlphaFold')].reset_index(drop=True)
+    log_update(f"Dropped rows where AlphaFold sequence was incorrect. New DataFrame length: {len(test_merge)}")
+    # make sure there's only one FusionGID number for each sequence for each GID
+    assert len(test_merge[test_merge.duplicated(['FusionGID','Raw Download AA Seq'])])==0
+    
+    # round pLDDTs
+    test_merge['pLDDT'] = test_merge['pLDDT'].round(2)
+    
+    # Finally, select only the columns we want
+    test_merge_v2 = test_merge[
+    ['FusionGID', 'FusionGene', 'Raw Download AA Seq','Len(AA seq)', 'Hgene', 'Hchr', 'Hbp', 'Hstrand', 'Tgene', 'Tchr', 'Tbp', 'Tstrand',
+       'Level','Structure Link', 'Structure Type', 'pLDDT', 'pLDDTs', 'Seq Source']
+    ].rename(
+        columns={
+            'Raw Download AA Seq': 'Fusion_Seq',
+            'Seq Source': 'Fusion_Seq_Source',
+            'Structure Link': 'Fusion_Structure_Link',
+            'Structure Type': 'Fusion_Structure_Type',
+            'pLDDT': 'Fusion_pLDDT',
+            'pLDDTs': 'Fusion_AA_pLDDTs',
+            'Len(AA seq)': 'Fusion_Length'
+        }
+    )
+    log_update(f"Unique FusionGIDs: {len(test_merge_v2['FusionGID'].unique())}")
+    log_update(f"Number of structures: {len(test_merge_v2)}")
+    
+    # Note that test_merge_v2 will still have duplicates where it's same seq, different ID
+    log_update("\nChecking for duplicate sequences..")
+    log_update(f"\tThe structure-based fusion database of length {len(test_merge_v2)} has {len(test_merge_v2['Fusion_Seq'].unique())} unique fusion sequences.")
+    dup_seqs = test_merge_v2[test_merge_v2['Fusion_Seq'].duplicated()]['Fusion_Seq'].tolist()
+    dup_seqs_df = test_merge_v2.loc[test_merge_v2['Fusion_Seq'].isin(dup_seqs)].reset_index(drop=True)
+    dup_seqs_df['FusionGID'] = dup_seqs_df['FusionGID'].astype(str)
+    dup_seqs_df = dup_seqs_df.groupby('Fusion_Seq').agg({
+        'FusionGID': lambda x: ','.join(x),
+        'FusionGene': lambda x: ','.join(x)
+    })
+    dup_seqs_df_str = dup_seqs_df.to_string(index=False)
+    dup_seqs_df_str = "\t"+dup_seqs_df_str.replace("\n","\n\t")
+    log_update(f"\tShowing FUsionGIDs and FusionGenes for duplicated sequences below:\n{dup_seqs_df_str}")
+    
+    # round pLDDT column
+    heads_tails_analyzed['Avg pLDDT'] = heads_tails_analyzed['Avg pLDDT'].round(2)
+    # merge treating data as head data
+    level_2_3_v2 = pd.merge(
+        level_2_3,
+        heads_tails_analyzed.rename(columns={'UniProtID': 'HGUniProtAcc', 'Avg pLDDT': 'HG_pLDDT', 'All pLDDTs': 'HG_AA_pLDDTs', 'Seq': 'HG_Seq'}),
+        on='HGUniProtAcc',
+        how='left'
+    )
+    # merge treating data as tail data
+    level_2_3_v2 = pd.merge(
+        level_2_3_v2,
+        heads_tails_analyzed.rename(columns={'UniProtID': 'TGUniProtAcc', 'Avg pLDDT': 'TG_pLDDT', 'All pLDDTs': 'TG_AA_pLDDTs', 'Seq': 'TG_Seq'}),
+        on='TGUniProtAcc',
+        how='left'
+    )
+    
+    # giant_level2_3 with valid structures only, no duplicate sequences, and head and tail proteins' uniprot IDs, pLDDTs, and sequences
+    test_merge_v2.to_csv(f'processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv',index=False)
+    log_update("Saved file with all fusion structure pLDDTs to: processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv")
+    
+    # level_2_3 with the head and tail proteins' uniprot IDs, pLDDTs, and sequences
+    level_2_3_v2.to_csv(f'processed_data/fusionpdb/FusionPDB_level2-3_cleaned_FusionGID_info.csv',index=False)
+    log_update("Saved file with all fusion protein heads and tails, and their structure pLDDTs to: processed_data/fusionpdb/FusionPDB_level2-3_cleaned_FusionGID_info.csv")
+        
+def main():
+    with open_logfile("process_fusion_structures_log.txt"):
+        process_fusions_and_hts()
+        
+if __name__ == "__main__":
+    main()
+    

fuson_plm/benchmarking/caid/processed_data/CAID-2_Disorder_NOX_Processed.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a6947f68c9d85ddaf7888207dde0d75a55a6e542da7a67aa5f3eefcf62a1511
+size 324778

fuson_plm/benchmarking/caid/processed_data/IDP-CRF_Training_Dataset.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f575340f3c68db5b5dcf1734301c946fa2d3aacb6583bcfca53e0610604553f
+size 2966294

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/AlphaFold-disorder_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ec1ffc27bddbb7a456c66391845f9487477077ebbcb99537310a8ae224ead78
+size 977219

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/AlphaFold-rsa_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f63fa90e6a629a9a64f61e9409e33af2aa134e026197f91fbef3c75b1574d38
+size 977736

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/DISOPRED3-diso_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5334217cd5290148af34134a1417c13fc88464ea757de5ccdc89fde9110b26b4
+size 1772014

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/DeepIDP-2L_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:268038cb09ff3363dfa6322f547ff00cfb40b297e465e92fc281357d62df3215
+size 1772014

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/DisoPred_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:952bc241621271e642da19abacc245899801733d7fcd16db066e19ab1f716ef0
+size 1687118

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/Dispredict3_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6a7293ebd84a6430c399e950a4d56adbb208a9f0ca568c32a117d78c47553e7
+size 1772014

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/ESpritz-D_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9da4c18279b48542e6da32e4306665f3712d0eeb37d8028323adc1058f113ac2
+size 1772014

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/IDP-Fusion_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a15182f4cb34d78c68555904b9d1fb5a25a10a425ce9414647c2d2057ebe92a6
+size 1742594

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/IUPred3_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e84fbaeb072f3e10cd2c939f5022bdaa196061651463aadd7ad526f65d5b707a
+size 1772018

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/disomine_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a13376dcbbb0858ab091d2694a40a34bcbf99937a4faa9b70fcd61ee3a6984d
+size 1449978

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/flDPlr2_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66420600ab69289e2a1a1452b12936c9b289b7e922571eb6efb139b16eab4db7
+size 1736172

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/flDPnn2_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8de27d681d8c216a2a430ace8bfa130f106d4174fb066deec8f9e780e04f51c
+size 1736172

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/flDPnn_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abb9d0b9cceae039b17105f9baaf023dd6e0d0da7aa48bece988cdcc5f7928c3
+size 1772014

fuson_plm/benchmarking/caid/processed_data/caid2_competition_results/flDPtr_CAID-2_Disorder_NOX.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27385d5a1d9348fdd90f1cfb516cd36209f61d0fe113439be8f39113825b3a03
+size 1736172

fuson_plm/benchmarking/caid/processed_data/figures/fusion_disorder/plddt_sequence_EML4::ALK_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f3d9a65f1e8ba25cdad55f016bcde23f739c9526f29983c30579e7a49fcd409e
+size 404717

fuson_plm/benchmarking/caid/processed_data/figures/fusion_disorder/plddt_sequence_EWSR1::FLI1_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8043d8bc2652721158e7f0cd8d40a94951ed85b781b9af9a997d0c961c022d29
+size 77468

fuson_plm/benchmarking/caid/processed_data/figures/fusion_disorder/plddt_sequence_PAX3::FOXO1_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93546673955907592808d58145040b075bf6dd58d08eebc412ef53d23e060858
+size 30322

fuson_plm/benchmarking/caid/processed_data/figures/fusion_disorder/plddt_sequence_SS18::SSX1_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86380dc4bb6c213591550b659d4e4dbb11650b8d31c0a1c32f1ca0d02e90abd9
+size 61266

fuson_plm/benchmarking/caid/processed_data/figures/histograms/disorder_nox_histogram_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32c2dfea9a8f45cc1d4005c7aac2fd729f2d930a8d68b17e05dc92175e85bead
+size 3051

fuson_plm/benchmarking/caid/processed_data/figures/histograms/fusions_histogram_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3f0a468334748fe08f1a230ea0897d783bdaf52b4674a3469c4d4111cbcaf98
+size 34084

fuson_plm/benchmarking/caid/processed_data/figures/histograms/heads_histogram_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c3c119114f8f3a98e8f16bb9b531102e4fa61b7cfe2eea1e641439d7b4c9762
+size 19279

fuson_plm/benchmarking/caid/processed_data/figures/histograms/tails_histogram_source_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbce7586f8f1a39696f6562c482aa6541cc3c99b104a207278922a47ba0a8dab
+size 20101

fuson_plm/benchmarking/caid/processed_data/flDPnn_Training_Dataset.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9882e290f110d47282ae81451a2e073ff2f579a957baa828144ca604f91f191e
+size 447048

fuson_plm/benchmarking/caid/processed_data/flDPnn_Validation_Dataset.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6568731760f2cb9daf20f55fec0014e0ee3eeaaa32c929770a16275792c837d6
+size 61366

fuson_plm/benchmarking/caid/processed_data/fusionpdb/FusionPDB_level2-3_cleaned_FusionGID_info.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aae5d5e39c4c52a1d07f0c427177004692688d0d8f447fd6b4eb83b0adbad6f3
+size 24073232

fuson_plm/benchmarking/caid/processed_data/fusionpdb/FusionPDB_level2-3_cleaned_structure_info.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e93f109073115fddbc97d493674ad2acc515b3c37bc62db23461d2a3b469e0f
+size 13726330

fuson_plm/benchmarking/caid/processed_data/fusionpdb/fusion_heads_and_tails.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61e8a29ba88d431c54ac40def2a1551ce21c203c2f9bb2f49a87b14ba0b41493
+size 68031

fuson_plm/benchmarking/caid/processed_data/fusionpdb/heads_tails_structural_data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b6115ec413dccb68dfe3c60c239c625a42161be6b87be14f2b3c1566452587e
+size 15471991