Pipeline_Functions.py

import os
import random
import numpy as np
import pandas as pd
import torch
import Data_Processing as dp
import PPI_Graph_Functions as pgf

def merge_pred_dicts(pred_dict_folder='./data/raw_networks/', save_folder='./data/merged_networks/', merge_func=np.nanmedian, min_reps=1, save_name_add='_merged.dat'):
    # Dynamic grouping of replicates from raw_networks folder
    files = [f for f in os.listdir(pred_dict_folder) if f.endswith('_weighted_pred_dict.dat')]
    address_dict = {}
    for file in files:
        if '_Sol_' in file:
            prefix = file.split('_Sol_')[0]
            if prefix not in address_dict:
                address_dict[prefix] = []
            address_dict[prefix].append(os.path.join(pred_dict_folder, file))

    # Sort to maintain replicate order
    for prefix in address_dict:
        address_dict[prefix].sort()

    for datatset, address_list in address_dict.items():
        save_address = os.path.join(save_folder, datatset + save_name_add)
        print("Merging into:", save_address)
        if os.path.exists(save_address):
            continue

        master_pred_dict = {}
        for address in address_list:
            pred_dict = dp.dat_to_dict(address)

            for pair, val in pred_dict.items():
                if pair in master_pred_dict:
                    master_pred_dict[pair].append(val)
                elif (pair[1],pair[0]) in master_pred_dict:
                    master_pred_dict[pair[1],pair[0]].append(val)
                else:
                    master_pred_dict[pair] = [val]

            del pred_dict

        merged_pred_dict = {}
        for pair, val_list in master_pred_dict.items():
            if len(val_list) >= min_reps:
                merged_pred_dict[pair] = merge_func(val_list)

        del master_pred_dict

        dp.network_dict_to_dat_file(merged_pred_dict, savename=save_address)


def build_ppigraphs(esm_dict, raw_folder='./data/merged_networks/', save_folder='./data/ppigraphs/', percentile=70, score_saver_percentile=99, gene_dict_percentile=90, min_interactors=20, max_degree=300):
    file_list = os.listdir(raw_folder)
    file_list = [os.path.join(raw_folder, file) for file in file_list if '.dat' in file]
    random.shuffle(file_list)

    for file in file_list:
        name = file.split('/')[-1].replace('.dat', '_ppigraph.dat')
        save_address = os.path.join(save_folder, name)
        if os.path.exists(save_address):
            continue

        print("Building graph for:", file)
        pred_dict = dp.dat_to_dict(file)

        ppi_graph_col = pgf.PPIGraphCollection(
            pred_dict=pred_dict, zs_cutoff=None, relv_zs_cutoff=None, percentile=percentile,
            score_saver_percentile=score_saver_percentile, gene_dict_percentile=gene_dict_percentile,
            min_interactors=min_interactors, max_degree=max_degree, esm2_dict=esm_dict,
            del_raw_pred_dict=True, name_dict=None, id_len=15, gen_attrs=False,
            gs_edges=None, gs_nodes=None, load_address=None, save_address=save_address,
            graph_load_mode=False, verbose=0, del_esm2_dict=False, prebuild_neighbors=False,
            num_workers=0, convert_to_pyg_graph=False, strict_nodes=False, rescue_prots=True)


def run_protea_predictions(esm_dict, model_address='./model/protea_weights.pth', raw_data_folder='./data/ppigraphs/', save_folder='./data/protea_predictions/', batch_size=2):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print('device:', device)

    model = pgf.Protea(device=device, pred_mode=True)
    model.load_state_dict(torch.load(model_address, map_location=device, weights_only=True))
    model = model.to(device)
    model.eval()

    file_list = [os.path.join(raw_data_folder, file) for file in os.listdir(raw_data_folder) if '.dat' in file]
    random.shuffle(file_list)

    for address_1 in file_list:
        file_1 = address_1.split('/')[-1].replace('.dat', '')

        ppi_graph_1 = pgf.PPIGraphCollection(esm2_dict=esm_dict, load_address=address_1)

        for address_2 in file_list:
            if address_1 == address_2:
                continue
            file_2 = address_2.split('/')[-1].replace('.dat', '')

            save_address = os.path.join(save_folder, file_1 + '_' + file_2 + '_homo_preds.csv')
            alt_save_address = os.path.join(save_folder, file_2 + '_' + file_1 + '_homo_preds.csv')

            if os.path.exists(save_address) or os.path.exists(alt_save_address):
                continue

            ppi_graph_2 = pgf.PPIGraphCollection(esm2_dict=esm_dict, load_address=address_2)
            pgf.protea_predictor(model, device, ppi_graph_1, ppi_graph_2, True,
                                 save_address, batch_size=batch_size)


def analyze_intra_interactome_regulation(homo_folder='./data/protea_predictions/', exp_list=None, relv_cond_list=None, output_folder='./data/intra_interactome_regulation/'):
    if relv_cond_list is None:
        relv_cond_list = [0]

    if exp_list is None:
        print('the relevant experiments need to be listed for this to work!')
        return

    # Create folder if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)

    files = [file for file in os.listdir(homo_folder) if '.csv' in file]

    master_homo_dict = dp.build_homo_comparision_dict(files, homo_folder, exp_list, {},
                                    name_split='_merged_ppigraph_', include_singles=True)

    mock_relative_master_homo_dict = dp.build_relative_homo_comparison_dict(master_homo_dict, relv_cond_list, inverted_relv_cond_list=False)
    infection_relative_master_homo_dict = dp.build_relative_homo_comparison_dict(master_homo_dict, relv_cond_list, inverted_relv_cond_list=True)
    timepoint_paired_to_mock_dict = dp.build_timepoint_infection_vs_mock_dict(mock_relative_master_homo_dict, infection_relative_master_homo_dict)
    diag_dist_dict = dp.build_diag_dist_dict(timepoint_paired_to_mock_dict, dist_comb_func=np.nanmean, dist_aggr_func=None)

    median_diag_dist_dict = {}
    for exp, prot_dict in diag_dist_dict.items():
        median_diag_dist_dict[exp] = {}
        for prot, hpi_dict in prot_dict.items():
            median_diag_dist_dict[exp][prot] = {}
            for hpi, val_list in hpi_dict.items():
                median_diag_dist_dict[exp][prot][hpi] = np.nanmedian(val_list)

    # Save median_diag_dist_dict using save_object
    save_path = os.path.join(output_folder, 'median_diag_dist_dict')
    dp.save_object(median_diag_dist_dict, save_path)

    # For each exp, create a table
    for exp, prot_dict in median_diag_dist_dict.items():
        # Get all unique hpis across all proteins for this exp
        all_hpis = set()
        for prot, hpi_dict in prot_dict.items():
            all_hpis.update(hpi_dict.keys())
        
        # Sort hpis as numbers from lowest to highest
        sorted_hpis = sorted(list(all_hpis), key=lambda x: float(x))

        # Build table rows
        rows = []
        for prot, hpi_dict in prot_dict.items():
            row = {'Protein': prot}
            for hpi in sorted_hpis:
                row[hpi] = hpi_dict.get(hpi, np.nan)
            rows.append(row)

        if not rows:
            continue

        df = pd.DataFrame(rows)
        # Reorder columns to have Protein first, followed by sorted hpis
        df = df[['Protein'] + sorted_hpis]
        
        csv_save_path = os.path.join(output_folder, f"{exp}_intra_infection_regulation.csv")
        df.to_csv(csv_save_path, index=False)
        print(f"Saved regulation table to: {csv_save_path}")


def analyze_cross_interactome_regulation(homo_folder='./data/protea_predictions/', exp_list=None, relv_comparison_dict=None, alignment_func=None, output_folder='./data/cross_interactome_regulation/'):
    if exp_list is None:
        print('the relevant experiments need to be listed for this to work!')
        return

    if relv_comparison_dict is None:
        print('relv_comparison_dict must be provided!')
        return

    if alignment_func is None:
        # Default alignment function: matches identical timepoints
        def default_alignment_func(exp_1, exp_2, cond_1, cond_2):
            if cond_1 == cond_2:
                return cond_1
            return None
        alignment_func = default_alignment_func

    # Create folder if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)

    files = [file for file in os.listdir(homo_folder) if '.csv' in file]

    master_homo_dict = dp.build_homo_comparision_dict(files, homo_folder, exp_list, relv_comparison_dict,
                                    name_split='_merged_ppigraph_', include_singles=False)

    relative_master_homo_dict = dp.build_relative_homo_comparison_dict(master_homo_dict, relv_cond_list=[], inverted_relv_cond_list=True)

    synced_relative_master_homo_dict = {}
    for exp_pair, prot_dict in relative_master_homo_dict.items():
        synced_relative_master_homo_dict[exp_pair] = {}
        for prot, cond_dict in prot_dict.items():
            synced_relative_master_homo_dict[exp_pair][prot] = {}
            for cond_pair, val in cond_dict.items():
                aligned_hpi = alignment_func(exp_pair[0], exp_pair[1], cond_pair[0], cond_pair[1])
                if aligned_hpi is not None:
                    synced_relative_master_homo_dict[exp_pair][prot][aligned_hpi] = val

    # Build homo_dict for slope calculations
    homo_dict = {}
    for exp_pair, prot_dict in synced_relative_master_homo_dict.items():
        homo_dict[exp_pair] = {}
        for prot, cond_dict in prot_dict.items():
            homo_dict[exp_pair][prot] = {}
            for cond, val in cond_dict.items():
                homo_dict[exp_pair][prot][cond] = val

    homo_slope_dict = {}
    for exp_pair, prot_dict in homo_dict.items():
        homo_slope_dict[exp_pair] = {}
        for prot, hpi_dict in prot_dict.items():
            if len(hpi_dict) == 0:
                continue

            sorted_hpi_list = [(hpi, np.nanmean(val_list)) for hpi, val_list in hpi_dict.items()]
            sorted_hpi_list.sort()

            sorted_hpi_list_vals = np.array([val for (hpi, val) in sorted_hpi_list])
            if len(sorted_hpi_list_vals) == 0 or sorted_hpi_list_vals[0] == 0:
                continue

            sorted_hpi_list_vals = sorted_hpi_list_vals / sorted_hpi_list_vals[0]

            slope_list = []
            for i in range(len(sorted_hpi_list_vals) - 1):
                val_1 = sorted_hpi_list_vals[i]
                val_2 = sorted_hpi_list_vals[i+1]
                slope = val_2 - val_1
                slope_list.append(slope)
            if len(slope_list) == 0:
                continue

            homo_slope_dict[exp_pair][prot] = np.nansum(slope_list)

    # Save homo_slope_dict using save_object
    save_path = os.path.join(output_folder, 'homo_slope_dict')
    dp.save_object(homo_slope_dict, save_path)

    saved_pairs = set()
    # For each exp_pair, create a table with median prediction values
    for exp_pair, prot_dict in synced_relative_master_homo_dict.items():
        exp_1, exp_2 = exp_pair
        
        canonical_pair = tuple(sorted(exp_pair))
        if canonical_pair in saved_pairs:
            continue
        saved_pairs.add(canonical_pair)
        
        # Get all unique hpis across all proteins for this exp_pair
        all_hpis = set()
        for prot, hpi_dict in prot_dict.items():
            all_hpis.update(hpi_dict.keys())
        
        # Sort hpis numerically
        sorted_hpis = sorted(list(all_hpis), key=lambda x: float(x))

        rows = []
        for prot, hpi_dict in prot_dict.items():
            row = {'Protein': prot}
            for hpi in sorted_hpis:
                val_list = hpi_dict.get(hpi, [])
                if isinstance(val_list, (list, tuple, np.ndarray)):
                    row[hpi] = np.nanmedian(val_list) if len(val_list) > 0 else np.nan
                else:
                    row[hpi] = val_list if pd.notnull(val_list) else np.nan
            
            slope = np.nan
            if exp_pair in homo_slope_dict:
                slope = homo_slope_dict[exp_pair].get(prot, np.nan)
            elif (exp_pair[1], exp_pair[0]) in homo_slope_dict:
                slope = homo_slope_dict[(exp_pair[1], exp_pair[0])].get(prot, np.nan)
            row['Slope'] = slope
            rows.append(row)

        if not rows:
            continue

        df = pd.DataFrame(rows)
        df = df[['Protein'] + sorted_hpis + ['Slope']]

        csv_save_path = os.path.join(output_folder, f"{exp_1}_{exp_2}_cross_infection_regulation.csv")
        df.to_csv(csv_save_path, index=False)
        print(f"Saved cross regulation table to: {csv_save_path}")