Data_Processing.py

import pickle
import pandas as pd
import numpy as np


def save_object(obj, save_name):
    f = open(save_name + ".pkl", "wb")
    pickle.dump(obj, f, -1)
    f.close()


def load_object(file_name):
    f = open(file_name + ".pkl", "rb")
    obj = pickle.load(f)
    f.close()
    return obj


def load_esm_token_dict(address, weighted=True, merge=True):
    raw_token_dict = load_object(address)
    prot_name_dict = {}

    # Get all protein names and assoicated token names since some proteins are split across multiple tokens
    for name in raw_token_dict.keys():
        if name == 'FOO':
            continue
        name_list = name.split('*')
        if len(name_list) == 1:
            prot_name_dict[name] = [name]
        else:
            prot = name_list[0]
            if prot not in prot_name_dict:
                prot_name_dict[prot] = []
            prot_name_dict[prot].append(name)

    token_dict = {}
    for prot, name_list in prot_name_dict.items():
        if merge:
            if len(name_list) == 1:
                token_dict[prot] = raw_token_dict[name_list[0]]
            else:
                sub_token_dict = {}
                for name in name_list:
                    sub_token_dict[name] = raw_token_dict[name]

                merged_token = merge_esm2_tokens(sub_token_dict, weighted=weighted)
                token_dict[prot] = merged_token

        else:
            if len(name_list) == 1:
                token_dict[prot] = [raw_token_dict[name_list[0]]]
            else:
                token_dict[prot] = []
                for name in name_list:
                    token_dict[prot].append(raw_token_dict[name])

    return token_dict


def create_gene_interaction_dict(pair_dict):
    gene_dict = {}
    for pair, value in pair_dict.items():
        gene1 = pair[0]
        gene2 = pair[1]
        if gene1 not in gene_dict:
            gene_dict[gene1] = {'genes': [], 'values': [], 'z_scores':[]}
        if gene2 not in gene_dict:
            gene_dict[gene2] = {'genes': [], 'values': [], 'z_scores':[]}

        gene_dict[gene1]['genes'].append(gene2)
        gene_dict[gene2]['genes'].append(gene1)
        gene_dict[gene1]['values'].append(value)
        gene_dict[gene2]['values'].append(value)

    for gene, dict_ in gene_dict.items():
        values = dict_['values']
        mean = np.nanmean(values)
        std = np.nanstd(values)
        if std == 0:
            std = 1

        for value in values:
            z_score = (value - mean) / std
            gene_dict[gene]['z_scores'].append(z_score)

    return gene_dict


def network_dict_to_dat_file(network_dict, savename='./data/test', gene_id_dict=None, n=2, z_score=False):
    if '.dat' in savename:
        f = open(savename, "w")
    else:
        f = open(savename + '.dat', "w")
    for key, value in network_dict.items():
        line = ''
        skip_flag = False
        for i in range(n):
            gene = key[i]
            if gene == None:
                skip_flag = True
                break
            if gene_id_dict != None:
                if gene in gene_id_dict:
                    gene = gene_id_dict[gene]
            line = line + '\t' + gene
        if skip_flag:
            continue

        line = line +'\t'+ str(value) + '\n'
        f.write(line)

    f.close()


def dat_to_dict(address,n=2):
    out_dict = {}
    with open(address, 'rt') as f:
        for line in f:
            try:
                gene_list = []
                if line[0] == '\t':
                    line=line[1:]
                value = float(line.split('\t')[n].replace('\n', '').replace('[', '').replace(']', ''))
                for i in range(n):
                    gene_list.append(line.split('\t')[i])
                out_dict[tuple(gene_list)] = float(value)
            except:
                continue

    return out_dict


def merge_esm2_tokens(token_dict, weighted=True):
    if weighted:
        weighted_token_list = []
        weight_list = []
        for sub_seq, token in token_dict.items():
            weight = float(sub_seq.split('*')[-1])
            weighted_token = token * weight

            weighted_token_list.append(weighted_token)
            weight_list.append(weight)

        weighted_token_list = np.array(weighted_token_list)
        weight_list = np.array(weight_list)
        merged_token = np.sum(weighted_token_list, axis=0) / np.sum(weight_list)
    else:
        token_list = []
        for sub_seq, token in token_dict.items():
            token_list.append(token)

        merged_token = np.nanmean(np.array(token_list), axis=0)

    return merged_token


def point_to_line_distance(x, y, a, b, c, absolute=False):
    """
    Formula for an upwards diagonal line is a = 1, b = -1, c = 0
    """

    if absolute:
        d = abs(a * x + b * y + c) / (a ** 2 + b ** 2) ** (1 / 2)
    else:
        d = -(a * x + b * y + c) / (a ** 2 + b ** 2) ** (1 / 2)

    return d


def load_interhomo_file(address,homo=True):
    table = pd.read_csv(address)
    interhomo_dict = {}
    if homo:
        prot_list = list(table['Protein 1'])
        val_list = list(table['Score'])
        for i in range(len(prot_list)):
            interhomo_dict[prot_list[i]] = val_list[i]
    else:
        prot_1_list = list(table['Protein 1'])
        prot_2_list = list(table['Protein 2'])
        val_list = list(table['Score'])
        for i in range(len(prot_1_list)):
            interhomo_dict[prot_1_list[i],prot_2_list[i]] = val_list[i]

    return interhomo_dict


def build_homo_comparision_dict(files, files_folder, exp_list, relv_comparison_dict,
                                name_split='_merged_ppigraph_', include_singles=True, single_exp=False):
    master_homo_dict = {}

    for file in files:
        name_1 = file.split(name_split)[0]
        if single_exp:
            exp_1 = name_1
        else:
            exp_1 = name_1.split('_')[0]

        if single_exp:
            cond_1 = 'single'
        else:
            if len(name_1.split('_')) == 2:
                cond_1 = int(name_1.split('_')[1])
            else:
                cond_1 = int(name_1.split('_')[2])
        name_2 = file.split(name_split)[1]
        if single_exp:
            exp_2 = name_2
        else:
            exp_2 = name_2.split('_')[0]

        if single_exp:
            cond_2 = 'single'
        else:
            if len(name_2.split('_')) == 2:
                cond_2 = int(name_2.split('_')[1])
            else:
                cond_2 = int(name_2.split('_')[2])

        flag = True
        for exp in exp_list:
            if exp in name_1:
                exp_1 = exp
                flag = False
                break
        if flag:
            continue

        flag = True
        for exp in exp_list:
            if exp in name_2:
                exp_2 = exp
                flag = False
                break

        if flag:
            continue

        flag = True
        if exp_1 == exp_2:
            flag = False
        elif exp_1 != exp_2:
            for relv_comps in relv_comparison_dict.values():
                if exp_1 in relv_comps and exp_2 in relv_comps:
                    flag = False
                    break
        if flag:
            continue

        if exp_1 == exp_2:
            if include_singles:
                if exp_1 not in master_homo_dict:
                    master_homo_dict[exp_1] = {}
            else:
                continue

            master_homo_dict[exp_1][cond_1, cond_2] = load_interhomo_file(files_folder + file, homo=True)
        else:
            if (exp_1, exp_2) not in master_homo_dict and (exp_2, exp_1) not in master_homo_dict:
                master_homo_dict[exp_1, exp_2] = {}
            if (exp_1, exp_2) in master_homo_dict:
                master_homo_dict[exp_1, exp_2][cond_1, cond_2] = load_interhomo_file(files_folder + file, homo=True)
            else:
                master_homo_dict[exp_2, exp_1][cond_2, cond_1] = load_interhomo_file(files_folder + file, homo=True)

    return master_homo_dict


def build_relative_homo_comparison_dict(master_homo_dict, relv_cond_list, inverted_relv_cond_list=False,
                                        singles_as_pairs=False):
    relative_master_homo_dict = {}

    if len(relv_cond_list) == 1:
        single_cond_flag = True
    else:
        single_cond_flag = False

    for exp, cond_dict in master_homo_dict.items():
        if singles_as_pairs:
            if type(exp) != tuple:
                exp = (exp,exp)
        relative_master_homo_dict[exp] = {}
        for pair, prot_dict in cond_dict.items():
            cond_1 = pair[0]
            cond_2 = pair[1]
            if inverted_relv_cond_list:
                if cond_1 in relv_cond_list:
                    continue
                elif cond_2 in relv_cond_list:
                    continue
            else:
                if cond_1 in relv_cond_list:
                    if single_cond_flag:
                        cond = cond_2
                elif cond_2 in relv_cond_list:
                    if single_cond_flag:
                        cond = cond_1
                else:
                    continue

            for prot, val in prot_dict.items():
                if prot not in relative_master_homo_dict[exp]:
                    relative_master_homo_dict[exp][prot] = {}
                if single_cond_flag and not inverted_relv_cond_list:
                    relative_master_homo_dict[exp][prot][cond] = val
                else:
                    relative_master_homo_dict[exp][prot][pair] = val

    return relative_master_homo_dict


def build_timepoint_infection_vs_mock_dict(mock_dict, infection_dict):
    timepoint_paired_to_mock_master_homo_dict = {}

    for exp, prot_dict in infection_dict.items():
        timepoint_paired_to_mock_master_homo_dict[exp] = {}

        mock_prot_dict = mock_dict[exp]
        for prot, hpi_dict in prot_dict.items():
            timepoint_paired_to_mock_master_homo_dict[exp][prot] = {}
            if prot not in mock_prot_dict:
                continue
            mock_hpi_dict = mock_prot_dict[prot]
            for hpi_pair, val in hpi_dict.items():
                hpi_1 = hpi_pair[0]
                hpi_2 = hpi_pair[1]
                if hpi_1 in mock_hpi_dict:
                    mock_1_val = mock_hpi_dict[hpi_1]
                else:
                    mock_1_val = None

                if hpi_2 in mock_hpi_dict:
                    mock_2_val = mock_hpi_dict[hpi_2]
                else:
                    mock_2_val = None

                if mock_1_val == None and mock_2_val == None:
                    continue
                elif mock_1_val == None:
                    mock_1_val = mock_2_val
                elif mock_2_val == None:
                    mock_2_val = mock_1_val

                timepoint_paired_to_mock_master_homo_dict[exp][prot][hpi_pair] = (mock_1_val, mock_2_val, val)

    return timepoint_paired_to_mock_master_homo_dict


def build_diag_dist_dict(timepoint_paired_dict, dist_comb_func=np.nanmean, dist_aggr_func=np.nansum):
    diag_dist_dict = {}

    for exp, prot_dict in timepoint_paired_dict.items():
        diag_dist_dict[exp] = {}
        for prot, cond_pair_dict in prot_dict.items():
            if dist_aggr_func != None:
                dist_list = []
            else:
                dist_dict = {}
            for cond_pair, vals in cond_pair_dict.items():
                mock_1_val, mock_2_val, val = vals

                x, y = mock_1_val, val
                dist_1 = point_to_line_distance(x, y, a=1, b=-1, c=0, absolute=False)

                x, y = mock_2_val, val
                dist_2 = point_to_line_distance(x, y, a=1, b=-1, c=0, absolute=False)

                dist = dist_comb_func([dist_1, dist_2])
                if dist_aggr_func != None:
                    dist_list.append(dist)
                else:
                    for cond in cond_pair:
                        if cond not in dist_dict:
                            dist_dict[cond] = []
                        dist_dict[cond].append(dist)


            if dist_aggr_func != None:
                if len(dist_list) == 0:
                    continue
                diag_dist_dict[exp][prot] = dist_aggr_func(dist_list)
            else:
                if len(dist_dict) == 0:
                    continue
                diag_dist_dict[exp][prot] = dist_dict

    return diag_dist_dict