structured_inhibition_spatial_network_model/scripts/competition_gain/run_micro_network.py

import json

import numpy
from brian2.units import *
from pypet import Environment, cartesian_product
from pypet.brian2 import Brian2MonitorResult

from scripts import models
from scripts.ring_network.head_direction import ex_in_network
from scripts.spatial_network.head_direction_index_over_noise_scale import calculate_rates
from scripts.spatial_network.run_entropy_maximisation_orientation_map import DATA_FOLDER, LOG_FOLDER

NO_CONNECTIONS = 0
RECIPROCAL = 1
UNI_DIRECTIONAL = 2

TRAJ_NAME = "competition_in_microcircuit"


def run_micro_networks(traj):
    N_E = 10
    number_of_inhibitory_neurons_per_population = 2
    architecture = traj.network.architecture
    excitatory_synapse_strength = traj.synapses.ex_in.strength * mV
    inhibitory_synapse_strength = traj.synapses.in_ex.strength * nS
    baseline = traj.input.baseline * nA
    input_to_pop_1 = traj.input.to_1 * nA
    input_to_pop_2 = traj.input.to_2 * nA
    duration = 500 * ms

    excitatory_eqs = models.hodgkin_huxley_eqs_with_synaptic_conductance + '''
        ih = 0*amp: amp
        '''  # add noise?

    excitatory_params = models.hodgkin_huxley_params
    excitatory_params.update({
        "E_i": -100 * mV
    })

    inhibitory_eqs = models.lif_eqs
    inhibitory_options = models.lif_options
    inhibitory_params = models.lif_params

    ei_synapse_model = models.delta_synapse_model
    ei_synapse_on_pre = models.delta_synapse_on_pre
    ei_synapse_param = models.delta_synapse_param

    ie_synapse_model = models.exponential_synapse
    ie_synapse_on_pre = models.exponential_synapse_on_pre
    ie_synapse_param = models.exponential_synapse_params

    number_of_excitatory_neurons_per_population = int(N_E / 2.0)
    if architecture == NO_CONNECTIONS:
        N_I = number_of_inhibitory_neurons_per_population
        ex_in_weights = numpy.zeros((N_E, N_I)) * volt
        in_ex_weights = numpy.zeros((N_I, N_E)) * siemens
    elif architecture == RECIPROCAL:
        N_I = number_of_inhibitory_neurons_per_population
        ex_in_weights = numpy.ones((N_E, N_I)) * excitatory_synapse_strength
        in_ex_weights = numpy.ones((N_I, N_E)) * inhibitory_synapse_strength
    elif architecture == UNI_DIRECTIONAL:
        N_I = 2 * number_of_inhibitory_neurons_per_population
        # each population has its own interneurons which then project to the other population
        ex_in_weights = numpy.zeros((N_E, N_I))
        ex_in_weights[:number_of_excitatory_neurons_per_population, :number_of_inhibitory_neurons_per_population] = 1
        ex_in_weights[number_of_excitatory_neurons_per_population:, number_of_inhibitory_neurons_per_population:] = 1
        ex_in_weights = ex_in_weights * excitatory_synapse_strength

        in_ex_weights = numpy.zeros((N_I, N_E))
        in_ex_weights[:number_of_inhibitory_neurons_per_population, number_of_excitatory_neurons_per_population:] = 1
        in_ex_weights[number_of_inhibitory_neurons_per_population:, :number_of_excitatory_neurons_per_population] = 1
        in_ex_weights = in_ex_weights * inhibitory_synapse_strength

    else:
        raise RuntimeError("Specify a valid network architecture.")

    net = ex_in_network(N_E, N_I, excitatory_eqs, excitatory_params, inhibitory_eqs,
                        inhibitory_params, inhibitory_options, ei_synapse_model, ei_synapse_on_pre,
                        ei_synapse_param,
                        ex_in_weights, ie_synapse_model, ie_synapse_on_pre,
                        ie_synapse_param, in_ex_weights, random_seed=4, dt=0.1 * ms)

    net["excitatory_neurons"].I[:int(N_E / 2.0)] = baseline + input_to_pop_1
    net["excitatory_neurons"].I[int(N_E / 2.0):] = baseline + input_to_pop_2
    net.run(duration=duration)

    traj.f_add_result(Brian2MonitorResult, 'spikes.e', net["excitatory_spike_monitor"],
                      comment='The spiketimes of the excitatory population')
    traj.f_add_result(Brian2MonitorResult, 'spikes.i', net["inhibitory_spike_monitor"],
                      comment='The spiketimes of the inhibitory population')
    excitatory_firing_rates = calculate_rates(net["excitatory_spike_monitor"].spike_trains().values())
    ex1 = numpy.mean(excitatory_firing_rates[:int(N_E / 2.0)])
    ex2 = numpy.mean(excitatory_firing_rates[int(N_E / 2.0):])
    return ex1, ex2


def add_parameters_to_trajectory(traj, parameter_dict, name_list=[]):
    if "value" in parameter_dict.keys():
        parameter_name = ".".join(name_list)
        traj.f_add_parameter(parameter_name, parameter_dict["value"],
                             comment=parameter_dict["comment"])
    else:
        for key in parameter_dict.keys():
            add_parameters_to_trajectory(traj, parameter_dict[key], name_list + [key])


def micronetwork_postproc(traj, result_list):
    runs = [run for run, rates in result_list]
    ex_1 = [rates[0] / hertz for run, rates in result_list]
    ex_2 = [rates[1] / hertz for run, rates in result_list]
    architectures = traj.par.network.f_get("architecture").f_get_range()
    input_1 = traj.par.input.f_get("to_1").f_get_range()
    # input_2 = traj.par.input.f_get("to_2").f_get_range()

    traj.f_add_result('summary', {"run": runs,
                                  "architecture": architectures,
                                  "f_1": ex_1,
                                  "f_2": ex_2,
                                  "I_1": input_1
                                  })


if __name__ == "__main__":
    env = Environment(trajectory=TRAJ_NAME, comment="Compare competition in microcircuits with "
                                                    "relay "
                                                    "interneurons or hub "
                                                    "interneurons", multiproc=True,
                      filename=DATA_FOLDER,
                      overwrite_file=True, ncores=3, log_folder=LOG_FOLDER)

    traj = env.trajectory
    with open("micro_network_parameter_space.json") as f:
        parameter_dict = json.load(f)
    add_parameters_to_trajectory(traj, parameter_dict)

    exploration_dict = {
        "network.architecture": [0, 1, 2],
        "input.to_1": [float(x) for x in numpy.arange(0.0, 0.7, 0.2)]
    }

    env.add_postprocessing(micronetwork_postproc)
    traj.f_explore(cartesian_product(exploration_dict))
    env.run(run_micro_networks)
    env.disable_logging()