structured_inhibition_spatial_network_model/scripts/spatial_network/hdi_optimization_via_syn_strength.py

import matplotlib.pyplot as plt
import noise
import numpy as np
from brian2.units import *
from scripts.spatial_maps.orientation_map import OrientationMap
from scripts.interneuron_placement import create_interneuron_sheet_entropy_max_orientation
import scripts.models as modellib
from scripts.ring_network.head_direction import get_head_direction_input, \
    ex_in_network
from tqdm import tqdm
from scripts.interneuron_placement import create_grid_of_excitatory_neurons, \
    create_interneuron_sheet_by_repulsive_force, get_excitatory_neurons_in_inhibitory_axonal_clouds

import multiprocessing
import itertools



trials_per_scale = 1

N_E = 900
N_I = 90

sheet_x = 450 * um
sheet_y = 450 * um

inhibitory_axon_long_axis = 100 * um
inhibitory_axon_short_axis = 25 * um

number_of_excitatory_neurons_per_row = int(np.sqrt(N_E))

'''
Neuron and synapse models
'''

excitatory_eqs = modellib.hodgkin_huxley_eqs_with_synaptic_conductance + modellib.eqs_ih
excitatory_params = modellib.hodgkin_huxley_params
excitatory_params.update(modellib.ih_params)
excitatory_params.update({"E_i": -80 * mV})
excitatory_params['ghbar'] = 0. * nS

lif_interneuron_eqs = """
dv/dt =1.0/tau* (-v + u_ext) :volt (unless refractory)
u_ext = u_ext_const : volt
"""

lif_interneuron_params = {
    "tau": 7 * ms,
    "v_threshold": -40 * mV,
    "v_reset": -60 * mV,
    "tau_refractory": 0.0 * ms,
    "u_ext_const": -50 * mV
}

lif_interneuron_options = {
    "threshold": "v>v_threshold",
    "reset": "v=v_reset",
    "refractory": "tau_refractory",
    'method': 'euler'
}

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

ie_synapse_model = modellib.exponential_synapse

ie_synapse_on_pre = modellib.exponential_synapse_on_pre
ie_synapse_param = modellib.exponential_synapse_params
ie_synapse_param["tau_syn"] = 2 * ms

'''
Tuning Maps
'''

# tuning_label = "Perlin"
tuning_label = "Orientation map"

# optimization_label = "Repulsive"
optimization_label = "Entropy Optimization"

ellipse_trial_sharpening_list = []
circle_trial_sharpening_list = []
no_conn_trial_sharpening_list = []

def get_fwhm_for_corr_len_and_seed(corr_len, seed, tuning_center):
    print(corr_len, tuning_center, seed)
    if tuning_label == "Perlin":  # TODO: How to handle scale in Perlin
        tuning_map = lambda x, y: noise.pnoise2(x / 100.0, y / 100.0, octaves=2) * np.pi
    elif tuning_label == "Orientation map":
        map = OrientationMap(number_of_excitatory_neurons_per_row + 1, number_of_excitatory_neurons_per_row + 1,
                             corr_len, sheet_x / um, sheet_y / um, seed)
        # map.improve(10)
        try:
            map.load_orientation_map()
        except:
            print(
                'No map yet with {}x{} pixels and {} pixel correllation length and {} seed'.format(map.x_dim, map.y_dim,
                                                                                                   map.corr_len,
                                                                                                   map.rnd_seed))
            return -1,-1,-1
        tuning_map = lambda x, y: map.tuning(x, y)

    ex_positions, ex_tunings = create_grid_of_excitatory_neurons(sheet_x / um, sheet_y / um,
                                                                 number_of_excitatory_neurons_per_row, tuning_map)

    inhibitory_radial_axis = np.sqrt(inhibitory_axon_long_axis * inhibitory_axon_short_axis)

    if optimization_label == "Repulsive":
        inhibitory_axonal_clouds = create_interneuron_sheet_by_repulsive_force(N_I, inhibitory_axon_long_axis / um,
                                                            inhibitory_axon_short_axis / um, sheet_x / um,
                                                            sheet_y / um, random_seed=2, n_iterations=1000)

        inhibitory_axonal_circles = create_interneuron_sheet_by_repulsive_force(N_I, inhibitory_radial_axis / um,
                                                             inhibitory_radial_axis / um, sheet_x / um,
                                                             sheet_y / um, random_seed=2, n_iterations=1000)
    elif optimization_label == "Entropy Optimization":
        inhibitory_axonal_clouds, ellipse_single_trial_entropy = create_interneuron_sheet_entropy_max_orientation(
            ex_positions, ex_tunings, N_I, inhibitory_axon_long_axis / um,
            inhibitory_axon_short_axis / um, sheet_x / um,
            sheet_y / um, trial_orientations=30)
        inhibitory_axonal_circles, circle_single_trial_entropy = create_interneuron_sheet_entropy_max_orientation(
            ex_positions, ex_tunings, N_I, inhibitory_radial_axis / um,
            inhibitory_radial_axis / um, sheet_x / um,
            sheet_y / um, trial_orientations=1)

    '''
    Connectvities
    '''

    # Spatial network with ellipsoid axons
    ie_connections = get_excitatory_neurons_in_inhibitory_axonal_clouds(ex_positions, inhibitory_axonal_clouds)

    inhibitory_synapse_strength = 30 * nS
    in_ex_weights = np.zeros((N_I, N_E)) * nS
    for interneuron_idx, connected_excitatory_idxs in enumerate(ie_connections):
        in_ex_weights[interneuron_idx, connected_excitatory_idxs] = inhibitory_synapse_strength

    excitatory_synapse_strength = 1 * mV
    ex_in_weights = np.where(in_ex_weights > 0 * nS, excitatory_synapse_strength, 0 * mV).T * volt

    # Spatial network with circular axons
    ie_connections_circle = get_excitatory_neurons_in_inhibitory_axonal_clouds(ex_positions,
                                                                               inhibitory_axonal_circles)

    in_ex_weights_circle = np.zeros((N_I, N_E)) * nS
    for interneuron_idx, connected_excitatory_idxs in enumerate(ie_connections_circle):
        in_ex_weights_circle[interneuron_idx, connected_excitatory_idxs] = inhibitory_synapse_strength

    excitatory_synapse_strength = 1 * mV
    ex_in_weights_circle = np.where(in_ex_weights_circle > 0 * nS, excitatory_synapse_strength, 0 * mV).T * volt

    # No synapses
    no_conn_ie = np.zeros((N_I, N_E)) * nS
    no_conn_ei = np.zeros((N_E, N_I)) * mV

    '''
    Prepare nets
    '''
    nets = []

    connectivity_label = ["No synapse", "Ellipsoid", "Circle"]
    connectivities = [(no_conn_ei, no_conn_ie), (ex_in_weights, in_ex_weights),
                      (ex_in_weights_circle, in_ex_weights_circle)]
    for ei_weights, ie_weights in connectivities:
        net = ex_in_network(N_E, N_I, excitatory_eqs, excitatory_params, lif_interneuron_eqs,
                            lif_interneuron_params,
                            lif_interneuron_options, ei_synapse_model, ei_synapse_on_pre,
                            ei_synapse_param,
                            ei_weights, ie_synapse_model, ie_synapse_on_pre,
                            ie_synapse_param, ie_weights, random_seed=2)
        nets.append(net)

    '''
    Head direction input
    '''

    ex_input_baseline = 0.0 * nA

    peak_phase = tuning_center
    input_sharpness = 1
    direction_input = get_head_direction_input(peak_phase, input_sharpness)
    max_head_direction_input_amplitude = 0.5 * nA
    input_to_excitatory_population = ex_input_baseline + max_head_direction_input_amplitude * direction_input(
        np.array(ex_tunings))

    '''
    Run simulation
    '''
    skip = 100 * ms
    length = 1100 * ms
    duration = skip + length

    for net in nets:
        excitatory_neurons = net["excitatory_neurons"]
        excitatory_neurons.I = input_to_excitatory_population
        net.run(duration)

    '''
    Get spatial map of rates
    '''

    def get_rates(spike_monitor, min_time=0 * ms):
        spike_trains = spike_monitor.spike_trains()

        isis = [np.ediff1d(np.extract(spike_times / ms > skip / ms, spike_times / ms)) * ms for spike_times in
                spike_trains.values()]

        rates = np.array([1.0 / np.mean(isi / ms) if isi.shape[0] != 0 else 0 for isi in isis]) * khertz
        return rates

    excitatory_rates = [get_rates(net["excitatory_spike_monitor"]) for net in nets]

    '''
    Get rate distribution over angles
    '''

    return ex_tunings, excitatory_rates

corr_len_range = range(0, 451, 15)
corr_len_range_len = len(corr_len_range)
tuning_range_len = 12
tuning_range = np.linspace(-np.pi,np.pi,tuning_range_len,endpoint=False)
seed_range = range(10)
seed_range_len = len(seed_range)
pool_arguments = itertools.product(corr_len_range, seed_range, tuning_range)

use_saved_array = True
if not use_saved_array:

    pool = multiprocessing.Pool()
    data = pool.starmap(get_fwhm_for_corr_len_and_seed,[*pool_arguments])
    print(type(data))
    # print(data)
    # data_array = np.reshape(np.array(data),(corr_len_range_len,seed_range_len,tuning_range_len,3))

    np.save('../../simulations/2020_02_27_head_direction_index_over_noise_scale/data.npy', np.array(data))
else:
    data = np.load('../../simulations/2020_02_27_head_direction_index_over_noise_scale/data.npy', allow_pickle=True)

no_conn_trial_hdi_array = np.array((corr_len_range_len,seed_range_len,tuning_range_len))
ellipse_trial_hdi_array = np.array((corr_len_range_len,seed_range_len,tuning_range_len))
circle_trial_hdi_array = np.array((corr_len_range_len,seed_range_len,tuning_range_len))
# pool_arguments = itertools.product(corr_len_range, seed_range, tuning_range)

# ex_tunings, excitatory_rates = data[0]
# plt.plot(ex_tunings, excitatory_rates[0] / hertz)
# plt.show()
# for id, corr_len, tuning_center, seed in enumerate(pool_arguments):
#     ex_tunings, excitatory_rates = data[id]
#     print(id, corr_len, tuning_center, seed)
circle_mean_hdi_overall = []
no_conn_mean_hdi_overall = []
ellipse_mean_hdi_overall = []

tuning_vectors_test = []

for cl_id, corr_len in enumerate(tqdm(corr_len_range)):
    circle_mean_hdi_per_corr_len = []
    no_conn_mean_hdi_per_corr_len = []
    ellipse_mean_hdi_per_corr_len = []
    for s_id, seed in enumerate(seed_range):
        circle_tuning_vector_list = np.zeros((N_E, 2))
        ellipse_tuning_vector_list = np.zeros((N_E, 2))
        no_conn_tuning_vector_list = np.zeros((N_E, 2))

        circle_rates_sum = np.zeros(N_E)
        ellipse_rates_sum = np.zeros(N_E)
        no_conn_rates_sum = np.zeros(N_E)
        for t_id, tuning_center in enumerate(tuning_range):
            total_id = t_id + tuning_range_len * s_id + tuning_range_len*seed_range_len * cl_id
            # print(cl_id, s_id, t_id, total_id)
            ex_tunings, excitatory_rates = data[total_id]
            circle_tuning_vector_list_at_tuning = np.array([np.array([np.cos(tuning_center), np.sin(tuning_center)]) \
                                                            * rate for rate in excitatory_rates[2]])
            circle_tuning_vector_list = circle_tuning_vector_list + circle_tuning_vector_list_at_tuning
            circle_rates_sum += excitatory_rates[2]

            ellipse_tuning_vector_list_at_tuning = np.array([np.array([np.cos(tuning_center), np.sin(tuning_center)]) \
                                                            * rate for rate in excitatory_rates[1]])
            ellipse_tuning_vector_list = ellipse_tuning_vector_list + ellipse_tuning_vector_list_at_tuning
            ellipse_rates_sum += excitatory_rates[1]

            no_conn_tuning_vector_list_at_tuning = np.array([np.array([np.cos(tuning_center), np.sin(tuning_center)]) \
                                                            * rate for rate in excitatory_rates[0]])
            no_conn_tuning_vector_list = no_conn_tuning_vector_list + no_conn_tuning_vector_list_at_tuning
            no_conn_rates_sum += excitatory_rates[0]

            if cl_id == 0 and s_id == 0:
                print('rates: \n', excitatory_rates[0][0])
                print('vectors: \n', no_conn_tuning_vector_list_at_tuning[0])
                print('vec norm: \n', np.linalg.norm(no_conn_tuning_vector_list_at_tuning[0]))
                tuning_vectors_test.append(no_conn_tuning_vector_list_at_tuning[0])

        circle_hdi_list = [np.linalg.norm(vec) / rate_sum for vec, rate_sum in zip(circle_tuning_vector_list,circle_rates_sum)]
        circle_hdi_mean_over_pop_list = np.sum(circle_hdi_list) / len(circle_hdi_list)
        circle_mean_hdi_per_corr_len.append(circle_hdi_mean_over_pop_list)

        ellipse_hdi_list = [np.linalg.norm(vec) / rate_sum for vec, rate_sum in
                           zip(ellipse_tuning_vector_list, ellipse_rates_sum)]
        ellipse_hdi_mean_over_pop_list = np.sum(ellipse_hdi_list) / len(ellipse_hdi_list)
        ellipse_mean_hdi_per_corr_len.append(ellipse_hdi_mean_over_pop_list)

        no_conn_hdi_list = [np.linalg.norm(vec) / rate_sum for vec, rate_sum in
                           zip(no_conn_tuning_vector_list, no_conn_rates_sum)]
        no_conn_hdi_mean_over_pop_list = np.sum(no_conn_hdi_list) / len(no_conn_hdi_list)
        no_conn_mean_hdi_per_corr_len.append(no_conn_hdi_mean_over_pop_list)
        # print(circle_tuning_vector_mean)
    circle_mean_hdi_overall.append(circle_mean_hdi_per_corr_len)
    no_conn_mean_hdi_overall.append(no_conn_mean_hdi_per_corr_len)
    ellipse_mean_hdi_overall.append(ellipse_mean_hdi_per_corr_len)

plt.figure()
plt.scatter(np.array(tuning_vectors_test)[:,0],np.array(tuning_vectors_test)[:,1])
plt.show()

# print(data.shape)
#
# for i in range(corr_len_range_len):
#     no_conn_trial_sharpening_list.append(data[i,:,0])
#     ellipse_trial_sharpening_list.append(data[i,:,1])
#     circle_trial_sharpening_list.append(data[i,:,2])
# print(circle_trial_sharpening_list)


ellipse_sharpening_mean = np.array([np.mean(i) for i in ellipse_mean_hdi_overall])
circle_sharpening_mean = np.array([np.mean(i) for i in circle_mean_hdi_overall])
no_conn_sharpening_mean = np.array([np.mean(i) for i in no_conn_mean_hdi_overall])

ellipse_sharpening_std_dev = np.array([np.std(i) for i in ellipse_mean_hdi_overall])
circle_sharpening_std_dev = np.array([np.std(i) for i in circle_mean_hdi_overall])
no_conn_sharpening_std_dev = np.array([np.std(i) for i in no_conn_mean_hdi_overall])
# print(ellipse_trial_sharpening_list)
# print(ellipse_entropy_std_dev)

plt.figure()
plt.plot(corr_len_range,circle_sharpening_mean, label='Circle', marker='o',color='C1')
plt.fill_between(corr_len_range,circle_sharpening_mean-circle_sharpening_std_dev,circle_sharpening_mean+circle_sharpening_std_dev,color='C1',alpha=0.4)
plt.plot(corr_len_range,ellipse_sharpening_mean, label='Ellipse', marker='o',color='C2')
plt.fill_between(corr_len_range,ellipse_sharpening_mean-ellipse_sharpening_std_dev,ellipse_sharpening_mean+ellipse_sharpening_std_dev,color='C2',alpha=0.4)
plt.plot(corr_len_range,no_conn_sharpening_mean, label='No Conn.', marker='o',color='C3')
plt.fill_between(corr_len_range,no_conn_sharpening_mean-no_conn_sharpening_std_dev,no_conn_sharpening_mean+no_conn_sharpening_std_dev,color='C3',alpha=0.4)
plt.xlabel('Correlation length')
plt.ylabel('Head Direction Index')
plt.legend()
plt.show()