eigenangles/scripts/utils.py

import pandas as pd
from matplotlib.offsetbox import AnchoredOffsetbox, TextArea, HPacker, VPacker


int_types = ['N', 'binsize', 'bin_num', 'seed_a', 'seed_b', 'seed']
float_types = ['t_start', 't_stop', 'rate', 'bkgr_correlation',
                'score', 'pvalue', 'epsilon', 'sigma_ex', 'sigma_in', 'simtime',
                'f', 'mu']
dtypes = dict.fromkeys(int_types, 'Int64') | dict.fromkeys(float_types, float)

colormap = {'weights': "#870156",
            'correlations': "#039857",
            'ratio': "#D29402",
            'rate_correlation': "#2c53af",
            'rate_correlation_exc': "#2c53af",
            'exc': "#2c53af",
            'rate_correlation_inh': "#e8685b",
            'inh': "#e8685b",
           }

def load_df(path):
    df = pd.read_csv(path, dtype=dtypes)
    df.drop(df.columns[df.columns.str.contains('unnamed', case=False)],
            axis=1, inplace=True)
    return df


def compact_column(df, compact_column='matrix', expand_column=['score', 'pvalue']):
    compact_values = df[compact_column].unique()
    if len(compact_values) <= 1:
        return df
    elif len(compact_values) > 2:
        raise ValueError('Can not compact on column with more than 2 values')

    df_left, df_right = [df[df[compact_column] == compact_value].copy()
                         for compact_value in compact_values]
    df_left.drop(compact_column, axis=1, inplace=True)
    df_right.drop(compact_column, axis=1, inplace=True)

    out = expand_column if type(expand_column) == list else [expand_column]
    out += [compact_column]
    on = [col for col in df_left.columns.to_list() if col not in out]

    df_compact = df_left.merge(df_right, how='left', on=on,
                               suffixes=(f'_{compact_values[0]}',
                                         f'_{compact_values[1]}'))
    return df_compact


def stack_columns(df, prefixes=['pvalue'], suffixes=['weights', 'correlations'],
                  new_column='matrix'):
    dfs = []
    for suffix in suffixes:
        data = df.copy()
        data[new_column] = suffix
        for prefix in prefixes:
            data[prefix] = data[f'{prefix}_{suffix}']
            for suf in suffixes:
                data.drop(f'{prefix}_{suf}', axis=1, inplace=True)
        dfs += [data]
    return pd.concat(dfs, axis=0, ignore_index=True)


def multicolor_ylabel(ax, list_of_strings, list_of_colors, anchorpad=0,
                      bbox_to_anchor=(-0.05, 0.5), **kw):
    ax.set_ylabel('')
    textprops=dict(ha='left', va='bottom', rotation=90, **kw)
    boxes = [TextArea(text, textprops=textprops | dict(color=color))
             for text, color in zip(list_of_strings[::-1], list_of_colors[::-1])]
    ybox = VPacker(children=boxes, align="center", pad=0, sep=5)
    anchored_ybox = AnchoredOffsetbox(loc='center', child=ybox, pad=anchorpad,
                                      frameon=False, bbox_to_anchor=bbox_to_anchor,
                                      bbox_transform=ax.transAxes, borderpad=0.)
    ax.add_artist(anchored_ybox)
    return None