fuson_plm/benchmarking/idr_prediction/split.py

from fuson_plm.utils.logging import open_logfile, log_update
from fuson_plm.utils.visualizing import set_font
from fuson_plm.benchmarking.idr_prediction.config import SPLIT
from fuson_plm.utils.splitting import split_clusters, check_split_validity
import os
import pandas as pd

def get_training_dfs(train, val, test, idr_db):
    """
    Remove unnecessary columns for efficient storing of train, validation, and test sets for benchmarking.
    Also, add the values using idr_db
    """
    log_update('\nMaking dataframes for IDR prediction benchmark...')
    
    # Delete cluster-related columns we don't need
    train = train.drop(columns=['representative seq_id','member seq_id', 'member length', 'representative seq']).rename(columns={'member seq':'Sequence'})
    val = val.drop(columns=['representative seq_id','member seq_id', 'member length', 'representative seq']).rename(columns={'member seq':'Sequence'})
    test = test.drop(columns=['representative seq_id','member seq_id', 'member length', 'representative seq']).rename(columns={'member seq':'Sequence'})
    
    # Add values and make one df for each one
    # idr_db values are in columns: asph,scaled_re,scaled_rg,scaling_exp
    value_cols = ['asph','scaled_re','scaled_rg','scaling_exp']
    return_dict = {}
    for col in value_cols:
        temp_train = pd.merge(train, idr_db[['Sequence',col]], on='Sequence',how='left').rename(columns={col:'Value'}).dropna(subset='Value')
        temp_val = pd.merge(val, idr_db[['Sequence',col]], on='Sequence',how='left').rename(columns={col:'Value'}).dropna(subset='Value')
        temp_test = pd.merge(test, idr_db[['Sequence',col]], on='Sequence',how='left').rename(columns={col:'Value'}).dropna(subset='Value')
        return_dict[col] = {
            'train': temp_train,
            'val': temp_val,
            'test': temp_test
        }
    
    return return_dict

def main():
    """
    """
    # Read all the input files
    LOG_PATH = "splitting_log.txt"
    IDR_DB_PATH = SPLIT.IDR_DB_PATH
    CLUSTER_OUTPUT_PATH = SPLIT.CLUSTER_OUTPUT_PATH
    RANDOM_STATE_1 = SPLIT.RANDOM_STATE_1
    TEST_SIZE_1 = SPLIT.TEST_SIZE_1
    RANDOM_STATE_2 = SPLIT.RANDOM_STATE_2
    TEST_SIZE_2 = SPLIT.TEST_SIZE_2
    
    # set font
    set_font()
    
    # Prepare the log file
    with open_logfile(LOG_PATH):
        log_update("Loaded data-splitting configurations from config.py")
        SPLIT.print_config(indent='\t')
            
        # Prepare directory to save results
        os.makedirs("splits",exist_ok=True)
            
        # Read the clusters and get a list of the representative IDs for splitting
        clusters = pd.read_csv(CLUSTER_OUTPUT_PATH)
        reps = clusters['representative seq_id'].unique().tolist()
        log_update(f"\nPreparing clusters...\n\tCollected {len(reps)} clusters for splitting")
        
        # Make the splits and extract the results
        splits = split_clusters(reps,
                                random_state_1 = RANDOM_STATE_1, test_size_1 = TEST_SIZE_1,
                                random_state_2= RANDOM_STATE_2, test_size_2 = TEST_SIZE_2) 
        X_train = splits['X_train']
        X_val = splits['X_val']
        X_test = splits['X_test']
        
        # Make slices of clusters dataframe for train, val, and test
        train_clusters = clusters.loc[clusters['representative seq_id'].isin(X_train)].reset_index(drop=True)
        val_clusters = clusters.loc[clusters['representative seq_id'].isin(X_val)].reset_index(drop=True)
        test_clusters = clusters.loc[clusters['representative seq_id'].isin(X_test)].reset_index(drop=True)
        
        # Check validity
        check_split_validity(train_clusters, val_clusters, test_clusters)
        
        # Print min and max sequence lengths
        min_train_seqlen = min(train_clusters['member seq'].str.len())
        max_train_seqlen = max(train_clusters['member seq'].str.len())
        min_val_seqlen = min(val_clusters['member seq'].str.len())
        max_val_seqlen = max(val_clusters['member seq'].str.len())
        min_test_seqlen = min(test_clusters['member seq'].str.len())
        max_test_seqlen = max(test_clusters['member seq'].str.len())
        log_update(f"\nLength breakdown summary...\n\tTrain: min seq length = {min_train_seqlen}, max seq length = {max_train_seqlen}")
        log_update(f"\nVal: min seq length = {min_val_seqlen}, max seq length = {max_val_seqlen}")
        log_update(f"\nTest: min seq length = {min_test_seqlen}, max seq length = {max_test_seqlen}")
        
        # cols = representative seq_id,member seq_id,representative seq,member seq
        train_clusters.to_csv("splits/train_cluster_split.csv",index=False)
        val_clusters.to_csv("splits/val_cluster_split.csv",index=False)
        test_clusters.to_csv("splits/test_cluster_split.csv",index=False)
        log_update('\nSaved cluster splits to splits/train_cluster_split.csv, splits/val_cluster_split.csv, splits/test_cluster_split.csv')
        cols=','.join(list(train_clusters.columns))
        log_update(f'\tColumns: {cols}')
        
        # Get final dataframes for training, and check their distributions
        idr_db = pd.read_csv(IDR_DB_PATH)
        train_dfs_dict = get_training_dfs(train_clusters, val_clusters, test_clusters, idr_db)
        os.makedirs('splits',exist_ok=True)
        train_test_values_dict = {}
        idr_property_name_dict = {'asph':'Asphericity','scaled_re':'End-to-End Distance (Re)','scaled_rg':'Radius of Gyration (Rg)','scaling_exp':'Scaling Exponent'}
        
        for idr_property, dfs in train_dfs_dict.items():
            os.makedirs(f"splits/{idr_property}", exist_ok=True)
            train_df = dfs['train']
            val_df = dfs['val']
            test_df = dfs['test']
            
            total_seqs = len(train_df)+len(val_df)+len(test_df)
            train_df.to_csv(f"splits/{idr_property}/train_df.csv",index=False)
            val_df.to_csv(f"splits/{idr_property}/val_df.csv",index=False)
            test_df.to_csv(f"splits/{idr_property}/test_df.csv",index=False)
            log_update(f"\nSaved {idr_property} training dataframes to splits/{idr_property}/train_df.csv, splits/{idr_property}/val_df.csv splits/test_df.csv")
            log_update(f"\tTrain sequences: {len(train_df)} ({100*len(train_df)/total_seqs:.2f}%)")
            log_update(f"\tVal sequences: {len(val_df)} ({100*len(val_df)/total_seqs:.2f}%)")
            log_update(f"\tTest sequences: {len(test_df)} ({100*len(test_df)/total_seqs:.2f}%)")
            log_update(f"\tTotal: {total_seqs}")
            
            # Make sure the lengths are right
            log_update(len(idr_db[idr_db[idr_property].notna()]))
            assert total_seqs == len(idr_db[idr_db[idr_property].notna()])
            
            train_test_values_dict[
                idr_property_name_dict[idr_property]
            ] = {
                'train': train_df['Value'].tolist(),
                'val': val_df['Value'].tolist(),
                'test': test_df['Value'].tolist()
            }

if __name__ == "__main__":
    main()