eigenangles/brunel_network/scripts/process_result_dataframe.py

import argparse
import pandas as pd
import numpy as np
from elephant.statistics import mean_firing_rate
import neo
from pathlib import Path
import os
import sys
sys.path.append(str(Path.cwd().parents[0] / 'scripts'))
from utils import load_df, compact_column


def add_logp_column(df, col_name_str='pvalue', log_prefix='log_'):
    col_names = [col for col in df.columns if col_name_str in col]
    for col_name in col_names:
        pvalue = df[col_name].to_numpy().astype(float)
        pvalue[~np.isfinite(pvalue)] = np.nan
        df[col_name] = pvalue
        log_pvalue = np.array([np.log10(p) if p else np.nan for p in pvalue])
        log_pvalue[~np.isfinite(log_pvalue)] = np.nan
        df[f'{log_prefix}{col_name}'] = log_pvalue
    return df

def calc_ratio(row, numerator, denominator, inf_value=np.nan):
    if row[denominator] == 0:
        return inf_value
    else:
        return row[numerator] / row[denominator]

def add_pvalue_ratio(df, ratio_keys=['correlations', 'weights'],
                         ratio_columns=['pvalue', 'score']):

    for col in ratio_columns:
        df[f'{col}_ratio'] = df.apply(
                lambda row: calc_ratio(row,
                                       f'{col}_{ratio_keys[0]}',
                                       f'{col}_{ratio_keys[1]}'), axis=1)
    return df

def add_rate_correlation_column(df, ratecorr_df, on=['protocol', 'seeds']):
    return df.merge(ratecorr_df, how='outer', on=None)

def sort_protocol_values(protocol):
    if 'ffle' in protocol:
        _ , fraction, populations = protocol.split('_')
        pop_order = np.array(['E-E', 'E-I', 'I-E', 'I-I'])
        value = np.where(populations == pop_order)[0][0]
        return float(fraction) + value/100
    elif 'add' in protocol:
        source, target = protocol.split('_')[-1].split('-')
        source, target = source.strip('E'), target.strip('E')
        return float(source)*100 + float(target)
    else:
        return 0

def order_df(df):
    key_function = lambda protocols: [sort_protocol_values(p) for p in protocols]
    df = df.sort_values('protocol', key=key_function)
    return df

if __name__ == '__main__':
    CLI = argparse.ArgumentParser()
    CLI.add_argument("--output", nargs='?', type=Path)
    CLI.add_argument("--comparison_df", nargs='?', type=Path)
    CLI.add_argument("--ratecorr_df", nargs='?', type=Path, default=None)
    args, unknown = CLI.parse_known_args()

    df = load_df(args.comparison_df)

    df = compact_column(df)
    df = add_pvalue_ratio(df)
    df = add_logp_column(df)

    if args.ratecorr_df is not None:
        df = add_rate_correlation_column(df, load_df(args.ratecorr_df))

    df = order_df(df)

    df.to_csv(args.output)