eigenangles/scripts/generate_poisson_correlation_matrix.py

from networkunit import models, tests, scores
import numpy as np
from quantities import ms, Hz
import argparse

class activity_model(models.stochastic_activity):
    params = {**models.stochastic_activity.default_params}

class corr_test(tests.correlation_matrix_test):
    score_type = scores.effect_size
    params = {'cluster_matrix':False,
              'remove_autocorr':False}

def generate_matrix(N, t_start, t_stop, binsize, rate, assembly_sizes,
                    correlations, bkgr_correlation, corr_method):

    params = {'size': N,
              't_start': t_start * ms,
              't_stop': t_stop * ms,
              'rate': rate * Hz,
              'statistic': 'poisson',
              'correlation_method': corr_method,
              'expected_bin_size': binsize * ms,
              'correlations': correlations,
              'assembly_sizes': assembly_sizes,
              'bkgr_correlation': bkgr_correlation,
              'max_pattern_length':100 * ms,
              'shuffle': False,
              'shuffle_seed': None}

    activity_model_inst = activity_model(**params)

    test = corr_test()
    return test.generate_prediction(activity_model_inst)


if __name__ == '__main__':
    CLI = argparse.ArgumentParser()
    CLI.add_argument("--N", nargs='?', type=int)
    CLI.add_argument("--t_start", nargs='?', type=float)
    CLI.add_argument("--t_stop", nargs='?', type=float)
    CLI.add_argument("--binsize", nargs='?', type=float)
    CLI.add_argument("--rate", nargs='?', type=float)
    CLI.add_argument("--assembly_sizes", nargs='?', type=lambda s: [int(i) for i in s.split(',')])
    CLI.add_argument("--correlations", nargs='?', type=lambda s: [float(i) for i in s.split(',')])
    CLI.add_argument("--bkgr_correlation", nargs='?', type=float)
    CLI.add_argument("--corr_method", nargs='?', type=str)
    CLI.add_argument("--output", nargs='?', type=str)
    args, unknown = CLI.parse_known_args()

    M = generate_matrix(N=args.N,
                        t_start=args.t_start,
                        t_stop=args.t_stop,
                        binsize=args.binsize,
                        rate=args.rate,
                        assembly_sizes=args.assembly_sizes,
                        correlations=args.correlations,
                        bkgr_correlation=args.bkgr_correlation,
                        corr_method=args.corr_method)

    np.save(args.output, M)