eigenangles/brunel_network/scripts/firing_rate_correlation.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 load_rate_vectors(path):
    try:
        spiketrains = load_spiketrains(path)
        spiketrains_exc = filter_spiketrains(spiketrains, neuron_type='excitatory')
        spiketrains_inh = filter_spiketrains(spiketrains, neuron_type='inhibitory')
        return [calc_rate_vector(st) for st in [spiketrains,
                                                spiketrains_exc, spiketrains_inh]]

    except FileNotFoundError:
        return [np.array([]), np.array([]), np.array([])]

def filter_spiketrains(spiketrains, **kwargs):
    for key, value in kwargs.items():
        spiketrains = [st for st in spiketrains if st.annotations[key] == value]
    return spiketrains

def load_spiketrains(path):
    io = neo.io.get_io(path)
    return io.read_block().segments[0].spiketrains

def calc_rate_vector(spiketrains):
    return np.array([mean_firing_rate(st) for st in spiketrains])

def calc_vector_correlation(vector_a, vector_b):
    if len(vector_a) and len(vector_b):
        return np.corrcoef(vector_a, vector_b)[0,1]
    else:
        return np.nan

if __name__ == '__main__':
    CLI = argparse.ArgumentParser()
    CLI.add_argument("--input", nargs='?', type=lambda s: s.split(' '))
    CLI.add_argument("--output", nargs='?', type=Path)
    args, unknown = CLI.parse_known_args()

    if len(args.input) != 2:
        raise ValueError(f'Expected two input files, got {len(args.input)}!')

    rates_a, rates_exc_a, rates_inh_a = load_rate_vectors(args.input[0])
    rates_b, rates_exc_b, rates_inh_b = load_rate_vectors(args.input[1])

    ratecorr = calc_vector_correlation(rates_a, rates_b)
    ratecorr_exc = calc_vector_correlation(rates_exc_a, rates_exc_b)
    ratecorr_inh = calc_vector_correlation(rates_inh_a, rates_inh_b)

    params = dict([(k.strip('-'),v) for k,v in zip(unknown[:-1:2],unknown[1::2])])
    params.update(rate_correlation=ratecorr,
                  rate_correlation_exc=ratecorr_exc,
                  rate_correlation_inh=ratecorr_inh)

    df = pd.Series(params).to_frame().T

    df.to_csv(args.output)