eigenangles/balanced_network/Snakefile

import config
import numpy as np
from pathlib import Path

wildcard_constraints:
    source_pop = 'E|I',
    target_pop = 'E|I',
    fraction = '[\d\.]+',
    f = '[\d\.]+',
    N = '\d+',
    epsilon = '[\d\.]+',
    eta = '[\d\.]+',
    mu = '[\d\.]+',
    seed = '\d+',
    seed_a = '\d+',
    seed_b = '\d+',
    sigma_ex = '[\d\.]+',
    sigma_in = '[\d\.]+',
    plot = '[\w\d\-]+',
    matrix = 'weights|correlations',
    measures = 'weights|correlations|ratecorr',
    protocol = '[\w\d\-\.]+',
    syndist = '[a-z\-]+',
    network_specs = '[\w\d\.]+',
    size = '[\d\.]+',
    length = '\d',
    connectors = '\d+',
    connectors_a = '\d+',
    connectors_b = '\d+',

rule rewire_and_redraw:
    input:
        expand('simulation_output/'\
             + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
             + 'sin{sigma_in}_{syndist}/seed_{seed}/{protocol}/{outfile}', #'{plot}.pdf',
            N = config.N,
            f = config.f,
            mu = config.mu,
            epsilon = config.epsilon,
            eta = config.eta,
            sigma_ex = config.sigma_ex,
            sigma_in = config.sigma_in,
            syndist = config.syndist,
            seed = config.seed,
            outfile = 'correlations.npy',
            # plot = ['correlations', 'weights',
            #         'spikes_50000-60000ms', 'spikes_59000-60000ms'],
            protocol = ['original']
                      + expand('shuffle_{frac}_{source}-{target}',
                               frac=config.shuffle_frac,
                               source=config.shuffle_source,
                               target=config.shuffle_target)
                      + expand('add_{source_frac}{source}-{target_frac}{target}',
                               source_frac=config.add_source_frac,
                               target_frac=config.add_target_frac,
                               source=config.add_source,
                               target=config.add_target)
                      + expand('cluster_{num}x{size}{target}_p{epsilon}',
                               num=config.cluster_number,
                               size=config.cluster_size,
                               target=config.cluster_pop,
                               epsilon=config.cluster_epsilon))

rule redraw_comparisons:
    input:
        expand('results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
             + 'sin{sigma_in}_{syndist}/redraw_{protocol}/{measure}_{seed_pair}.csv',
            N = config.N,
            f = config.f,
            mu = config.mu,
            epsilon = config.epsilon,
            eta = config.eta,
            sigma_ex = config.sigma_ex,
            sigma_in = config.sigma_in,
            syndist = config.syndist,
            seed_pair = config.seed_pairs[::int(len(config.seed)/2)],
            measure = ['correlations', 'weights', 'ratecorr'],
            protocol =
                       ['original']
                     # + expand('cluster_{num}x{size}{target}_p{epsilon}',
                     #           num=config.cluster_number,
                     #           size=config.cluster_size,
                     #           target=config.cluster_pop,
                     #           epsilon=config.cluster_epsilon)
                     # + expand('chain_{num}x{size}{target}_j{strength}_p{epsilon}',
                     #           num=config.chain_length,
                     #           size=config.chain_size,
                     #           target=config.chain_pop,
                     #           strength=config.chain_strength,
                     #           epsilon=config.chain_epsilon)
                     )

rule rewire_comparisons:
    input:
        expand('results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
             + 'sin{sigma_in}_{syndist}/{protocol}/{measure}_{seed}.csv',
            N = config.N,
            f = config.f,
            mu = config.mu,
            epsilon = config.epsilon,
            eta = config.eta,
            sigma_ex = config.sigma_ex,
            sigma_in = config.sigma_in,
            syndist = config.syndist,
            seed = config.seed,
            measure = ['correlations', 'weights', 'ratecorr'],
            protocol =
                        # expand('sheffle_{frac}_{source}-{target}',
                        #        frac=3500, #config.shuffle_frac,
                        #        source=config.shuffle_source,
                        #        target=config.shuffle_target)
                     # +
                      expand('add_{source_frac}{source}-{target_frac}{target}',
                               source_frac=config.add_source_frac,
                               target_frac=config.add_target_frac,
                               source=config.add_source,
                               target=config.add_target)
                     # + expand('cluster_{num}x{size}{target}_p{epsilon}',
                     #           num=config.cluster_number,
                     #           size=config.cluster_size,
                     #           target=config.cluster_pop,
                     #           epsilon=config.cluster_epsilon)
                       # expand('chain_{length}_{size}{cpop}_j{strength}_p{epsilon}_c{c_a}-{c_b}',
                       #          length=config.chain_length,
                       #          c_a=config.connectors[0],
                       #          c_b=config.connectors[1],
                       #          size=config.chain_size,
                       #          cpop=config.chain_pop,
                       #          strength=config.chain_strength,
                       #          epsilon=config.chain_epsilon)
            )

# rule chain_comparisons:
#     input:
#         expand('results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
#              + 'sin{sigma_in}_{syndist}/{protocol}/{measure}_{seed}.csv',
#             N = config.N,
#             f = config.f,
#             mu = config.mu,
#             epsilon = config.epsilon,
#             eta = config.eta,
#             sigma_ex = config.sigma_ex,
#             sigma_in = config.sigma_in,
#             syndist = config.syndist,
#             seed = config.seed,
#             measure = ['correlations', 'weights', 'ratecorr'],
#             protocol = expand('shuffle_{frac}_{source}-{target}',
#                                frac=config.shuffle_frac,
#                                source=config.shuffle_source,
#                                target=config.shuffle_target)
#         )


# REWIRE EXPERIMENTS

def id_span(N, f, pop):
    N, f = int(N), float(f)
    N_exc = int(N*f)
    if pop == 'E':
        return (0, N_exc)
    elif pop =='I':
        return (N_exc, N)
    else:
        raise ValueError

rule shuffle_weights:
    input:
        script = 'scripts/shuffle_weights.py',
        weights = '{root}/N{N}_f{f}_{specs}/original/weights.npy'
    output:
        weights = '{root}/N{N}_f{f}_{specs}/'\
                + 'shuffle_{fraction}_{source_pop}-{target_pop}/weights.npy'
    params:
        source_span = lambda w: id_span(w.N, w.f, w.source_pop),
        target_span = lambda w: id_span(w.N, w.f, w.target_pop),
    shell:
        """
        python {input.script} --input "{input.weights}" \
                              --output "{output.weights}" \
                              --source_span {params.source_span} \
                              --target_span {params.target_span} \
                              --fraction {wildcards.fraction}
        """

use rule shuffle_weights as sheffle_weights with:
    output:
        weights = '{root}/N{N}_f{f}_{specs}/'\
                + 'sheffle_{fraction}_{source_pop}-{target_pop}/weights.npy'


rule add_weights:
    input:
        script = 'scripts/add_weights.py',
        weights = '{root}/N{N}_f{f}_{specs}/original/weights.npy'
    output:
        weights = '{root}/N{N}_f{f}_{specs}/'\
                + 'add_{source_frac}{source_pop}-{target_frac}{target_pop}'\
                + '/weights.npy'
    params:
        source_span = lambda w: id_span(w.N, w.f, w.source_pop),
        target_span = lambda w: id_span(w.N, w.f, w.target_pop),
        weight_mean = lambda w: config.J_ex if w.source_pop == 'E' else config.J_in,
        weight_std =  lambda w: config.sigma_ex if w.source_pop == 'E' else config.sigma_in,
        syndist = config.syndist
    shell:
        """
        python {input.script} --input "{input.weights}" \
                              --output "{output}" \
                              --source_span {params.source_span} \
                              --target_span {params.target_span} \
                              --source_fraction {wildcards.source_frac} \
                              --target_fraction {wildcards.target_frac} \
                              --weight_mean {params.weight_mean} \
                              --weight_std {params.weight_std} \
                              --syndist {params.syndist}
        """

rule cluster_weights:
    input:
        script = 'scripts/cluster_weights.py',
        weights = '{root}/N{N}_f{f}_{specs}/original/weights.npy'
    output:
        weights = '{root}/N{N}_f{f}_{specs}/'\
                + 'cluster_{num}x{size}{cpop}_p{epsilon}/weights.npy'
    params:
        pop_span = lambda w: id_span(w.N, w.f, w.cpop),
        weight_mean = lambda w: config.J_ex if w.cpop == 'E' else config.J_in,
        weight_std =  lambda w: config.sigma_ex if w.cpop == 'E' else config.sigma_in,
        syndist = config.syndist
    shell:
        """
        python {input.script} --input "{input.weights}" \
                              --output "{output}" \
                              --cluster_number {wildcards.num} \
                              --cluster_fraction {wildcards.size} \
                              --cluster_prob {wildcards.epsilon} \
                              --pop_span {params.pop_span} \
                              --weight_mean {params.weight_mean} \
                              --weight_std {params.weight_std} \
                              --syndist {params.syndist}
        """

rule hub_weights:
    input:
        script = 'scripts/hub_weights.py',
        weights = '{root}/N{N}_f{f}_{specs}/original/weights.npy'
    output:
        weights = '{root}/N{N}_f{f}_{specs}/' \
                + 'hub_{size}{hpop}_j{strength}_p{epsilon}/weights.npy'
    params:
        pop_span = lambda w: id_span(w.N, w.f, w.hpop),
    shell:
        """
        python {input.script} --input "{input.weights}" \
                              --output "{output.weights}" \
                              --hub_size {wildcards.size} \
                              --hub_prob {wildcards.epsilon} \
                              --hub_strength {wildcards.strength} \
                              --pop_span {params.pop_span}
        """

def chain_sections(w):
    min_size = int(w.length)
    sections = list(range(min_size, int(w.length)+min_size))[::-1]
    norm = sum(sections)
    return [i/norm for i in sections]

rule chain_weights:
    input:
        script = 'scripts/chain_weights.py',
        weights = '{root}/N{N}_f{f}_{network_specs}/seed_{seed}/original/weights.npy'
    output:
        weights = '{root}/N{N}_f{f}_{network_specs}/seed_{seed}/' \
                + 'chain_{length}_{size}{cpop}' \
                + '_j{strength}_p{epsilon}_c{connectors}/weights.npy'
    params:
        pop_span = lambda w: id_span(w.N, w.f, w.cpop),
        chain_sections = chain_sections
    shell:
        """
        python {input.script} --input "{input.weights}" \
                              --output "{output.weights}" \
                              --chain_size {wildcards.size} \
                              --chain_prob {wildcards.epsilon} \
                              --chain_strength {wildcards.strength} \
                              --pop_span {params.pop_span} \
                              --chain_sections "{params.chain_sections}" \
                              --connectors {wildcards.connectors} \
                              --seed {wildcards.seed}
        """


# COMPARE NETWORK ACTIVITY AND CONNECTIVITY

rule compare_redrawn_networks:
    input:
        script = "../scripts/eigenangle_test.py",
        matrix_a = 'simulation_output/'\
                 + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
                 + '_sin{sigma_in}_{syndist}/seed_{seed_a}/{protocol}/{matrix}.npy',
        matrix_b = 'simulation_output/'\
                 + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
                 + '_sin{sigma_in}_{syndist}/seed_{seed_b}/{protocol}/{matrix}.npy'
    output:
        temp('results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
           + '_sin{sigma_in}_{syndist}/redraw_{protocol}/'\
           + '{matrix}_{seed_a}-{seed_b}.json')
    params:
        bin_num = config.bin_num,
        is_connectivity = lambda w: True if ('weights' in w.matrix) else False,
        shuffle_neuron_ids = lambda w, output: 'redraw' in str(output)
    shell:
        """
        python {input.script} --matrix_a {input.matrix_a} \
                              --matrix_b {input.matrix_b} \
                              --output {output} \
                              --N {wildcards.N} \
                              --bin_num {params.bin_num} \
                              --mu {wildcards.mu} \
                              --f {wildcards.f} \
                              --epsilon {wildcards.epsilon} \
                              --sigma_ex {wildcards.sigma_ex} \
                              --sigma_in {wildcards.sigma_in} \
                              --is_connectivity {params.is_connectivity} \
                              --shuffle_neuron_ids {params.shuffle_neuron_ids}
        """

use rule compare_redrawn_networks as compare_rewired_network with:
    input:
        script = "../scripts/eigenangle_test.py",
        matrix_a = 'simulation_output/'\
                 + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
                 + '_sin{sigma_in}_{syndist}/seed_{seed}/original/{matrix}.npy',
        matrix_b = 'simulation_output/'\
                 + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
                 + '_sin{sigma_in}_{syndist}/seed_{seed}/{protocol}/{matrix}.npy',
    output:
        temp('results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
           + '_sin{sigma_in}_{syndist}/{protocol}/{matrix}_{seed}.json')

use rule compare_redrawn_networks as compare_chain_networks with:
    input:
        script = "../scripts/eigenangle_test.py",
        matrix_a = 'simulation_output/'\
                 + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
                 + '_sin{sigma_in}_{syndist}/seed_{seed}/chain_{specs}_'\
                 + 'c{connectors_a}/{matrix}.npy',
        matrix_b = 'simulation_output/'\
                 + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
                 + '_sin{sigma_in}_{syndist}/seed_{seed}/chain_{specs}_'\
                 + 'c{connectors_b}/{matrix}.npy',
    output:
        'results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}'\
           + '_sin{sigma_in}_{syndist}/chain_{specs}_'\
           + 'c{connectors_a}-{connectors_b}/{matrix}_{seed}.json'
#
# ruleorder:  calc_firing_rate_correlation > score_to_dataframe #> compare_chain_networks > compare_rewired_network

def get_spikes(wildcards):
    network_specs = wildcards.network_specs
    protocol = wildcards.protocol
    seeds = wildcards.seeds
    path = lambda prcl, s: f'simulation_output/{network_specs}/' \
                         + f'seed_{s}/{prcl}/spikes.pkl'
    if 'redraw' in protocol:
        protocol = protocol.strip('redraw_')
        return [path(protocol, seed) for seed in seeds.split('-')]
    elif 'chain' in protocol:
        protocol, connectors = protocol.split('_c')
        connectors = connectors.split('-')
        return [path(f'{protocol}_c{c}', seeds) for c in connectors]
    else:
        return [path(p, seeds) for p in ['original', protocol]]

rule calc_firing_rate_correlation:
    input:
        script = 'scripts/firing_rate_correlation.py',
        spikes = get_spikes
    params:
        sim_folder = 'simulation_output'
    output:
        'results/{network_specs}/{protocol}/ratecorr_{seeds}.csv'
    shell:
        """
        python {input.script} --output "{output}" \
                              --input "{input.spikes}" \
                              --protocol {wildcards.protocol} \
                              --seeds {wildcards.seeds}
        """

# BUILD AND SIMULATE NETWORK

rule build_network:
    input:
        script = 'scripts/build_network.py',
        config = 'config.py'
    params:
        out_config = lambda w, output: Path(output.weights).parents[1] / 'config.yml'
    output:
        weights = 'simulation_output/'\
                + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
                + 'sin{sigma_in}_{syndist}/seed_{seed}/original/weights.npy'
    shell:
        """
        python {input.script} --weights_path {output.weights} \
                              --out_config {params.out_config} \
                              --network_config {input.config} \
                              --N {wildcards.N} \
                              --f {wildcards.f} \
                              --mu {wildcards.mu} \
                              --sigma_ex {wildcards.sigma_ex} \
                              --sigma_in {wildcards.sigma_in} \
                              --epsilon {wildcards.epsilon} \
                              --seed {wildcards.seed} \
                              --syndist {wildcards.syndist}
        """

rule simulate_network:
    input:
        script = 'scripts/simulate_network.py',
        weights = 'simulation_output/'\
                + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
                + 'sin{sigma_in}_{syndist}/seed_{seed}/{protocol}/weights.npy',
    output:
        spikes_ex = temp('simulation_output/'\
                  + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
                  + 'sin{sigma_in}_{syndist}/seed_{seed}/{protocol}/spikes_ex.gdf'),
        spikes_in = temp('simulation_output/'\
                  + 'N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}_'\
                  + 'sin{sigma_in}_{syndist}/seed_{seed}/{protocol}/spikes_in.gdf'),
    params:
        simtime = config.simtime,
        config_path = 'config.py'
    shell:
        """
        python {input.script} --spikes_ex_path {output.spikes_ex} \
                              --spikes_in_path {output.spikes_in} \
                              --weights_path {input.weights} \
                              --network_config {params.config_path} \
                              --N {wildcards.N} \
                              --f {wildcards.f} \
                              --mu {wildcards.mu} \
                              --sigma_ex {wildcards.sigma_ex} \
                              --sigma_in {wildcards.sigma_in} \
                              --epsilon {wildcards.epsilon} \
                              --simtime {params.simtime} \
                              --seed {wildcards.seed} \
                              --eta {wildcards.eta}
        """

# TRANSFORM RULES

rule nest_to_neo:
    input:
        script = 'scripts/nest_to_neo.py',
        spikes_ex = 'simulation_output/N{N}_f{f}_{specs}/spikes_ex.gdf',
        spikes_in = 'simulation_output/N{N}_f{f}_{specs}/spikes_in.gdf',
    output:
        'simulation_output/N{N}_f{f}_{specs}/spikes.pkl'
    params:
        t_start = 0,
        t_stop = config.simtime
    shell:
        """
        python {input.script} --spikes_ex "{input.spikes_ex}" \
                              --spikes_in "{input.spikes_in}" \
                              --output {output} \
                              --t_stop {params.t_stop} \
                              --t_start {params.t_start} \
                              --N {wildcards.N} \
                              --f {wildcards.f}
        """

rule create_correlation_matrix_from_spikes:
    input:
        script = 'scripts/correlation_matrix_from_spikes.py',
        spikes = '{spikes_dir}/spikes.pkl',
    output:
        correlation = '{spikes_dir}/correlations.npy'
    params:
        t_start = config.cut_inital_time,
        t_stop = config.simtime,
        bin_size = config.bin_size
    shell:
        """
        python {input.script} --spikes "{input.spikes}" \
                              --output {output} \
                              --t_stop {params.t_stop} \
                              --t_start {params.t_start} \
                              --bin_size {params.bin_size}
        """

rule score_to_dataframe:
    input:
        script = '../scripts/score_to_dataframe.py',
        data = 'results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}' \
             + '_sin{sigma_in}_{syndist}/{protocol}/{matrix}_{seeds}.json'
    output:
        'results/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}_sex{sigma_ex}' \
        + '_sin{sigma_in}_{syndist}/{protocol}/{matrix}_{seeds}.csv'
    params:
        bin_num = config.bin_num,
        simtime = config.simtime
    shell:
        """
        python {input.script} --output "{output}" \
                              --input "{input.data}" \
                              --params {wildcards} \
                              --bin_num {params.bin_num} \
                              --simtime {params.simtime}
        """

rule merge_comparison_results:
    input:
        script = '../scripts/merge_dataframes.py',
    params:
        exclude = ['merged_ratecorr_results', 'merged_comparison_results',
                   'rewiring_results', 'ratecorr']
    output:
        temp('{dir}/merged_comparison_results.csv')
    shell:
        """
        python {input.script} --output "{output}" \
                              --exclude "{params.exclude}"
        """

use rule merge_comparison_results as merge_ratecorr_results with:
    params:
        exclude = ['merged_ratecorr_results', 'merged_comparison_results',
                   'rewiring_results', 'weights', 'correlations']
    output:
        temp('{dir}/merged_ratecorr_results.csv')

rule process_result_dataframe:
    input:
        script = 'scripts/process_result_dataframe.py',
        comparison_df = '{dir}/merged_comparison_results.csv',
        ratecorr_df = '{dir}/merged_ratecorr_results.csv'
    output:
        '{dir}/rewiring_results.csv'
    shell:
        """
        python {input.script}  --output "{output}" \
                               --comparison_df "{input.comparison_df}" \
                               --ratecorr_df "{input.ratecorr_df}"
        """

# PLOT SIMULATION OUTPUT

rule plot_spiketrains:
    input:
        script = '../scripts/plot_spiketrains.py',
        spikes = 'simulation_output/{network_specs}/seed_{seed}/{protocol}/spikes.pkl'
    output:
        'images/rasterplot/{network_specs}_seed{seed}_{protocol}_spikes{t_start}-{t_stop}ms.pdf'
    shell:
        """
        python {input.script} --spikes "{input.spikes}" \
                              --output "{output}" \
                              --t_start {wildcards.t_start} \
                              --t_stop {wildcards.t_stop}
        """

rule plot_matrix:
    input:
        script = '../scripts/plot_matrix.py',
        matrix = 'simulation_output/{network_specs}/seed_{seed}/{protocol}/{matrix}.npy'
    output:
        'images/{matrix}/{network_specs}_seed{seed}_{protocol}.pdf'
    shell:
        """
        python {input.script} --input "{input.matrix}" \
                              --output "{output}"
        """

rule plot_eigenspectrum:
    input:
        script = 'scripts/plot_eigenspectrum.py',
        files = expand('simulation_output/N{{N}}_f{{f}}_mu{{mu}}_p{{epsilon}}' \
                     + '_eta{{eta}}_sex{{sigma_ex}}_sin{{sigma_in}}_{{syndist}}' \
                     + '/seed_{seed}/{{protocol}}/weights.npy',
                       seed = range(1,9))
    output:
        'images/eigenspectrum/N{N}_f{f}_mu{mu}_p{epsilon}_eta{eta}' \
      + '_sex{sigma_ex}_sin{sigma_in}_{syndist}_{protocol}.pdf'
    shell:
        """
        python {input.script} --input "{input.files}" \
                              --output "{output}" \
                              --N {wildcards.N} \
                              --f {wildcards.f} \
                              --mu {wildcards.mu} \
                              --sigma_ex {wildcards.sigma_ex} \
                              --sigma_in {wildcards.sigma_in} \
                              --epsilon {wildcards.epsilon} \
        """

rule plot_pvalues:
    input:
        script = 'scripts/plot_pvalue_{plot_name}.py',
        dataframe = 'results/{network_specs}/rewiring_results.csv'
    output:
        'images/pvalue_{plot_name}/{network_specs}-{protocol}.pdf'
    shell:
        """
        python {input.script} --input "{input.dataframe}" \
                              --output "{output}" \
                              --protocol {wildcards.protocol}
        """