structured_inhibition_spatial_network_model/scripts/interneuron_placement.py

import random as rng

import matplotlib.pyplot as plt
import noise
import numpy as np
from brian2.units import *
from matplotlib.patches import Ellipse
from tqdm import tqdm

from scripts.spatial_maps.orientation_maps.orientation_map import OrientationMap


class Pickle:

    def __init__(self, x, y, a, b, phi):
        self.x = x
        self.y = y
        self.a = a  # Semimajor axis
        self.b = b  # Semiminor axis
        self.phi = phi
        self.c = 2.0 * a  # Sum distance
        self.foc_len = np.sqrt(a ** 2 - b ** 2)

        self.area = np.pi * a * b

    def get_p1(self):
        x1 = self.x - self.foc_len * np.cos(self.phi)
        y1 = self.y - self.foc_len * np.sin(self.phi)
        return x1, y1

    def get_p2(self):
        x2 = self.x + self.foc_len * np.cos(self.phi)
        y2 = self.y + self.foc_len * np.sin(self.phi)
        return x2, y2

    def get_ellipse(self):
        ell = Ellipse((self.x, self.y), 2.0 * self.a, 2.0 * self.b,
                      angle=self.phi * 180. / np.pi, linewidth=1, fill=False, zorder=2)
        return ell

    def get_area(self):
        return np.pi * self.a * self.b


def plot_neural_sheet(ex_positions, ex_tunings, axonal_clouds=None):
    X, Y = get_correct_position_mesh(ex_positions)

    n_ex = int(np.sqrt(len(ex_positions)))
    head_dir_preference = np.array(ex_tunings).reshape((n_ex, n_ex))

    fig = plt.figure(figsize=(4, 4))
    ax = fig.add_subplot(111)

    c = ax.pcolor(X, Y, head_dir_preference, vmin=-np.pi, vmax=np.pi, cmap="twilight")
    fig.colorbar(c, ax=ax, label="Orientation")
    if axonal_clouds is not None:
        for i, p in enumerate(axonal_clouds):
            ell = p.get_ellipse()
            # print(ell)
            ax.add_artist(ell)


def get_position_mesh(ex_positions):
    n_ex = int(np.sqrt(len(ex_positions)))
    X = np.array([x for x, _ in ex_positions]).reshape((n_ex, n_ex))
    Y = np.array([y for _, y in ex_positions]).reshape((n_ex, n_ex))
    return X, Y

def get_correct_position_mesh(ex_positions):
    n_ex = int(np.sqrt(len(ex_positions)))

    x_wrong = np.array([x for x, _ in ex_positions]).reshape((n_ex, n_ex))
    y_wrong = np.array([y for _, y in ex_positions]).reshape((n_ex, n_ex))

    x = x_wrong[0, :]
    y = y_wrong[:, 0]

    xmin = np.min(x)
    xmax = np.max(x)
    dx = (xmax - xmin) / (x.shape[0] - 1)

    ymin = np.min(y)
    ymax = np.max(y)
    dy = (ymax - ymin) / (y.shape[0] - 1)
    # print(np.arange(xmin, xmax + 2 * dx, dx) - dx / 2.)
    # X, Y = np.meshgrid(np.arange(xmin, xmax + 2 * dx, dx) - dx / 2., np.arange(ymin, ymax + 2 * dy, dy) - dy / 2.)
    # print(np.linspace(xmin, xmax + dx, n_ex + 1, endpoint=True) - dx / 2)
    X, Y = np.meshgrid(np.linspace(xmin, xmax + dx, n_ex + 1, endpoint=True) - dx / 2, np.linspace(ymin, ymax + dy, n_ex + 1, endpoint=True) - dy / 2)


    return X, Y

def get_entropy_of_axonal_coverage(ex_neuron_tuning, ie_connections):
    ex_phase_vals_per_pickle = []
    phase_entropy_per_pickle = []
    n_hist_bins = 10
    for connected_idxs in ie_connections:
        ex_phase_vals_of_p = [ex_neuron_tuning[idx] for idx in connected_idxs]
        phase_entropy = 0.0

        ex_phase_vals_per_pickle.append(ex_phase_vals_of_p)
        ex_phase_hist, _ = np.histogram(ex_phase_vals_of_p, n_hist_bins)

        n_ex_of_p = float(len(ex_phase_vals_of_p))
        ex_phase_hist = 1. / n_ex_of_p * ex_phase_hist.astype(float)

        for n_phi in ex_phase_hist:
            if n_phi != 0:
                phase_entropy -= n_phi * np.log(n_phi)

        phase_entropy_per_pickle.append(phase_entropy)
    return phase_entropy_per_pickle


def get_excitatory_neurons_in_inhibitory_axonal_clouds(ex_neuron_positions, inhibitory_axonal_clouds):
    ie_connections = []
    for p in inhibitory_axonal_clouds:
        ie_conn_p = []

        x_p_1, y_p_1 = p.get_p1()
        x_p_2, y_p_2 = p.get_p2()

        for ex_idx, ex_pos in enumerate(ex_neuron_positions):
            x_ex, y_ex = ex_pos
            d_1 = np.sqrt((x_ex - x_p_1) ** 2 + (y_ex - y_p_1) ** 2)
            d_2 = np.sqrt((x_ex - x_p_2) ** 2 + (y_ex - y_p_2) ** 2)
            if d_1 + d_2 < p.c:
                ie_conn_p.append(ex_idx)

        ie_connections.append(ie_conn_p)
    return ie_connections


def get_excitatory_phases_in_inhibitory_axon(ex_neuron_positions, ex_neuron_tunings, inhibitory_neuron):
    ex_phase_vals = []

    x_p_1, y_p_1 = inhibitory_neuron.get_p1()
    x_p_2, y_p_2 = inhibitory_neuron.get_p2()

    for ex_idx, ex_pos in enumerate(ex_neuron_positions):
        x_ex, y_ex = ex_pos
        d_1 = np.sqrt((x_ex - x_p_1) ** 2 + (y_ex - y_p_1) ** 2)
        d_2 = np.sqrt((x_ex - x_p_2) ** 2 + (y_ex - y_p_2) ** 2)
        if d_1 + d_2 < inhibitory_neuron.c:
            ex_phase_vals.append(ex_neuron_tunings[ex_idx])

    return ex_phase_vals


def get_interneuron_entropy(ex_phase_vals, entropy_bins):
    phase_entropy = 0.

    ex_phase_hist, _ = np.histogram(ex_phase_vals, entropy_bins)
    ex_phase_prob = ex_phase_hist / len(ex_phase_vals)
    for prob in ex_phase_prob:
        if prob > 0.0:
            phase_entropy += -prob * np.log(prob)
    return phase_entropy


def create_interneuron_sheet_by_noise(ex_positions, ex_tunings, number_of_interneurons, long_axis, short_axis, max_x,
                                      max_y, random_seed=None,
                                      n_iterations=100):
    '''
        Start with uniform distribution
        '''
    if random_seed is not None:
        rng.seed(random_seed)
    pickle_a = long_axis
    pickle_b = short_axis
    pickle_n = number_of_interneurons
    x_n = int(np.sqrt(pickle_n))
    y_n = int(np.sqrt(pickle_n))
    x_step = int(max_x / (x_n + 1))
    y_step = int(max_y / (y_n + 1))
    inhibitory_axonal_clouds = []
    for n in range(pickle_n):
        inhibitory_axonal_clouds.append(Pickle(((n % x_n) + 1) * x_step, \
                                               (int(n / y_n) + 1) * y_step, \
                                               pickle_a, pickle_b, rng.random() * 2. * np.pi))

    '''
    Determine initial tunings and entropy values
    '''
    entropy_bins = np.linspace(-np.pi, np.pi, 30)

    ex_tunings_of_interneurons = []
    entropies_of_interneurons = []

    for i_neuron in inhibitory_axonal_clouds:
        ex_phase_vals = get_excitatory_phases_in_inhibitory_axon(ex_positions, ex_tunings, i_neuron)
        ex_tunings_of_interneurons.append(ex_phase_vals)
        entropies_of_interneurons.append(get_interneuron_entropy(ex_phase_vals, entropy_bins))

    '''
    (Randomly) perturb interneuron and see if entropy increases, if so update position and axonal tunings
    '''
    # print('Mean Entropy before rotation ', np.mean(entropies_of_interneurons))

    pert_phi = np.pi / 4.0
    pert_phi_decrease = pert_phi / (n_iterations + 1)
    for it in tqdm(range(n_iterations), desc="Optimizing axonal cloud distribution"):
        rotation_perturbation = pert_phi - it * pert_phi_decrease

        for i_neuron_id, i_neuron in enumerate(inhibitory_axonal_clouds):
            perturbed_interneuron = Pickle(i_neuron.x, i_neuron.y, long_axis, short_axis,
                                           i_neuron.phi + rotation_perturbation)
            perturbed_phase_vals = get_excitatory_phases_in_inhibitory_axon(ex_positions, ex_tunings,
                                                                            perturbed_interneuron)
            perturbed_entropy = get_interneuron_entropy(perturbed_phase_vals, entropy_bins)

            if perturbed_entropy > entropies_of_interneurons[i_neuron_id]:
                ex_tunings_of_interneurons.append(ex_phase_vals)
                entropies_of_interneurons[i_neuron_id] = perturbed_entropy
                i_neuron.phi = perturbed_interneuron.phi
            else:
                perturbed_interneuron = Pickle(i_neuron.x, i_neuron.y, long_axis, short_axis,
                                               i_neuron.phi - rotation_perturbation)
                perturbed_phase_vals = get_excitatory_phases_in_inhibitory_axon(ex_positions, ex_tunings,
                                                                                perturbed_interneuron)
                perturbed_entropy = get_interneuron_entropy(perturbed_phase_vals, entropy_bins)

                if perturbed_entropy > entropies_of_interneurons[i_neuron_id]:
                    ex_tunings_of_interneurons.append(ex_phase_vals)
                    entropies_of_interneurons[i_neuron_id] = perturbed_entropy
                    i_neuron.phi = perturbed_interneuron.phi
    # print('Mean Entropy after rotation', np.mean(entropies_of_interneurons))

    return inhibitory_axonal_clouds, np.mean(entropies_of_interneurons)


def create_interneuron_sheet_entropy_max_orientation(ex_positions, ex_tunings, number_of_interneurons, long_axis,
                                                     short_axis, max_x, max_y, trial_orientations=30,
                                                     number_of_bins=30, placement_jitter=0.0, seed=0):
    inhibitory_axonal_clouds = initialize_axonal_clouds_on_grid_with_random_orientation(number_of_interneurons,
                                                                                        long_axis, short_axis, max_x,
                                                                                        max_y, placement_jitter, seed)
    interneuron_positions = [(inh.x, inh.y) for inh in inhibitory_axonal_clouds]

    '''
    Determine initial tunings and entropy values
    '''
    entropies_of_interneurons = [0.]
    if trial_orientations > 0:
        entropies_of_interneurons, optimal_orientations = get_optimal_orientations_for_maximum_entropy(ex_positions,
                                                                                                   ex_tunings,
                                                                                                   interneuron_positions,
                                                                                                   long_axis,
                                                                                                   short_axis,
                                                                                                   number_of_bins,
                                                                                                   trial_orientations)

        for opt_orientation, axon in zip(optimal_orientations, inhibitory_axonal_clouds):
            axon.phi = opt_orientation

    return inhibitory_axonal_clouds, np.mean(entropies_of_interneurons)

def create_interneuron_sheet_entropy_min_orientation(ex_positions, ex_tunings, number_of_interneurons, long_axis,
                                                     short_axis, max_x, max_y, trial_orientations=30,
                                                     number_of_bins=30, placement_jitter=0.0):
    inhibitory_axonal_clouds = initialize_axonal_clouds_on_grid_with_random_orientation(number_of_interneurons,
                                                                                        long_axis, short_axis, max_x,
                                                                                        max_y, placement_jitter)
    interneuron_positions = [(inh.x, inh.y) for inh in inhibitory_axonal_clouds]

    '''
    Determine initial tunings and entropy values
    '''
    entropies_of_interneurons, optimal_orientations = get_optimal_orientations_for_minimum_entropy(ex_positions,
                                                                                                   ex_tunings,
                                                                                                   interneuron_positions,
                                                                                                   long_axis,
                                                                                                   short_axis,
                                                                                                   number_of_bins,
                                                                                                   trial_orientations)

    for opt_orientation, axon in zip(optimal_orientations, inhibitory_axonal_clouds):
        axon.phi = opt_orientation

    return inhibitory_axonal_clouds, np.mean(entropies_of_interneurons)


def get_optimal_orientations_for_maximum_entropy(ex_positions, ex_tunings, interneuron_positions, long_axis, short_axis,
                                                 number_of_bins, trial_orientations):
    entropy_bins = np.linspace(-np.pi, np.pi, number_of_bins)
    entropies_of_interneurons = []
    optimal_orientations = []
    # for i_neuron in tqdm(inhibitory_axonal_clouds, desc="Maximizing entropy by rotation"):
    for x, y in interneuron_positions:
        orientations = np.linspace(-np.pi, np.pi, trial_orientations)
        entropy_for_orientation = []
        for orientation in orientations:
            perturbed_interneuron = Pickle(x, y, long_axis, short_axis, orientation)
            ex_phase_vals = get_excitatory_phases_in_inhibitory_axon(ex_positions, ex_tunings, perturbed_interneuron)
            entropy_for_orientation.append(get_interneuron_entropy(ex_phase_vals, entropy_bins))
        optimal_orientations.append(orientations[np.argmax(entropy_for_orientation)])
        entropies_of_interneurons.append(np.max(entropy_for_orientation))
    return entropies_of_interneurons, optimal_orientations

def get_optimal_orientations_for_minimum_entropy(ex_positions, ex_tunings, interneuron_positions, long_axis, short_axis,
                                                 number_of_bins, trial_orientations):
    entropy_bins = np.linspace(-np.pi, np.pi, number_of_bins)
    entropies_of_interneurons = []
    optimal_orientations = []
    # for i_neuron in tqdm(inhibitory_axonal_clouds, desc="Maximizing entropy by rotation"):
    for x, y in interneuron_positions:
        orientations = np.linspace(-np.pi, np.pi, trial_orientations)
        entropy_for_orientation = []
        for orientation in orientations:
            perturbed_interneuron = Pickle(x, y, long_axis, short_axis, orientation)
            ex_phase_vals = get_excitatory_phases_in_inhibitory_axon(ex_positions, ex_tunings, perturbed_interneuron)
            entropy_for_orientation.append(get_interneuron_entropy(ex_phase_vals, entropy_bins))
        optimal_orientations.append(orientations[np.argmin(entropy_for_orientation)])
        entropies_of_interneurons.append(np.min(entropy_for_orientation))
    return entropies_of_interneurons, optimal_orientations


def initialize_axonal_clouds_on_grid_with_random_orientation(number_of_interneurons, long_axis, short_axis, sheet_x,
                                                             sheet_y, placement_jitter=0.0, seed=0):
    '''
    Start with uniform distribution
    '''
    rng.seed(seed)
    pickle_a = long_axis
    pickle_b = short_axis
    pickle_n = number_of_interneurons
    x_n = int(np.sqrt(pickle_n))
    y_n = int(np.sqrt(pickle_n))
    x_step = int(sheet_x / (x_n + 1))
    y_step = int(sheet_y / (y_n + 1))
    inhibitory_axonal_clouds = []
    pj = placement_jitter
    for n in range(pickle_n):
        theta = rng.random() * 2. * np.pi
        pj_x = pj * np.cos(theta)
        pj_y = pj * np.sin(theta)
        inhibitory_axonal_clouds.append(Pickle(((n % x_n) + 1) * x_step + pj_x, \
                                               (int(n / y_n) + 1) * y_step + pj_y, \
                                               pickle_a, pickle_b, theta))
    return inhibitory_axonal_clouds


def create_interneuron_sheet_by_repulsive_force(number_of_interneurons, long_axis, short_axis, max_x, max_y,
                                                random_seed=None,
                                                n_iterations=100):
    # '''
    # Start with random pickle distribution
    # '''
    # if random_seed is not None:
    #     rng.seed(random_seed)
    # pickle_a = long_axis
    # pickle_b = short_axis
    # pickle_n = number_of_interneurons
    # x_rng_min = pickle_a
    # x_rng_max = max_x - pickle_a
    # y_rng_min = pickle_a
    # y_rng_max = max_y - pickle_a
    # inhibitory_axonal_clouds = []
    # for n in range(pickle_n):
    #     inhibitory_axonal_clouds.append(Pickle(rng.uniform(x_rng_min, x_rng_max), \
    #                                            rng.uniform(y_rng_min, y_rng_max), \
    #                                            pickle_a, pickle_b, rng.random() * 2. * np.pi))

    '''
        Start with uniform distribution
        '''
    if random_seed is not None:
        rng.seed(random_seed)
    pickle_a = long_axis
    pickle_b = short_axis
    pickle_n = number_of_interneurons
    x_n = int(np.sqrt(pickle_n))
    y_n = int(np.sqrt(pickle_n))
    x_step = int(max_x / (x_n + 1))
    y_step = int(max_y / (y_n + 1))
    inhibitory_axonal_clouds = []
    for n in range(pickle_n):
        inhibitory_axonal_clouds.append(AxonalCloud(((n % x_n) + 1) * x_step, \
                                                    (int(n / y_n) + 1) * y_step, \
                                                    pickle_a, pickle_b, rng.random() * 2. * np.pi))

    # plotting()
    '''
    Distribute pickles by repulsive "force"
    '''

    f_scale = 0.2

    def point_repulsion(x_i, y_i, x_j, y_j):
        d_x = x_i - x_j
        d_y = y_i - y_j

        # Distance is torus-like as to approach more homogenous distribution

        if d_x > max_x / 2.:
            d_x = -max_x + d_x
        elif d_x < -max_x / 2.:
            d_x = max_x + d_x
        if d_y > max_y / 2.:
            d_y = -max_y + d_y
        elif d_y < -max_y / 2.:
            d_y = max_y + d_y
        d = np.sqrt(d_x ** 2 + d_y ** 2)
        f_i = f_scale * 1. / (d + 1) ** 2
        f_x = f_i * d_x / d
        f_y = f_i * d_y / d
        return f_x, f_y

    def pickle_repulsion(p_i, p_j):
        phi = p_i.phi
        p_i_x1, p_i_y1 = p_i.get_p1()
        p_i_x2, p_i_y2 = p_i.get_p2()

        p_j_x1, p_j_y1 = p_j.get_p1()
        p_j_x2, p_j_y2 = p_j.get_p2()

        f_p1p1_x, f_p1p1_y = point_repulsion(p_i_x1, p_i_y1, p_j_x1, p_j_y1)  # p1-p1
        f_p1p2_x, f_p1p2_y = point_repulsion(p_i_x1, p_i_y1, p_j_x2, p_j_y2)  # p1-p2

        f_p2p1_x, f_p2p1_y = point_repulsion(p_i_x2, p_i_y2, p_j_x1, p_j_y1)  # p2-p1
        f_p2p2_x, f_p2p2_y = point_repulsion(p_i_x2, p_i_y2, p_j_x2, p_j_y2)  # p2-p2

        f_p1_x = f_p1p1_x + f_p1p2_x
        f_p1_y = f_p1p1_y + f_p1p2_y

        f_p2_x = f_p2p1_x + f_p2p2_x
        f_p2_y = f_p2p1_y + f_p2p2_y

        m_p1 = (f_p1_x * np.sin(phi) - f_p1_y * np.cos(phi))
        m_p2 = (-f_p2_x * np.sin(phi) + f_p2_y * np.cos(phi))

        f_x_sum = f_p1_x + f_p2_x
        f_y_sum = f_p1_y + f_p2_y
        m_sum = m_p1 + m_p2

        return f_x_sum, f_y_sum, m_sum

    for i in tqdm(range(n_iterations), desc="Optimizing axonal cloud distribution"):
        delta_xyphi = []
        for i, p_i in enumerate(inhibitory_axonal_clouds):
            delta_x = 0.
            delta_y = 0.
            delta_phi = 0.
            for p_j in inhibitory_axonal_clouds:
                if p_i is not p_j:
                    del_x, del_y, del_m = pickle_repulsion(p_i, p_j)
                    delta_x += del_x
                    delta_y += del_y
                    delta_phi += del_m
            delta_xyphi.append([delta_x, delta_y, delta_phi])
        # plotting(delta_xy)
        for i, p_i in enumerate(inhibitory_axonal_clouds):
            x_new = p_i.x + delta_xyphi[i][0]
            y_new = p_i.y + delta_xyphi[i][1]
            p_i.phi = p_i.phi + delta_xyphi[i][2]

            # Torus-like distance so pickles can cross borders and behave like a subset of a bigger distribution

            if x_new < max_x and x_new > 0.0:
                p_i.x = x_new
            elif x_new < 0.:
                p_i.x = max_x + x_new
            elif x_new > max_x:
                p_i.x = x_new - max_x

            if y_new < max_y and y_new > 0.0:
                p_i.y = y_new
            elif y_new < 0.:
                p_i.y = max_y + y_new
            elif y_new > max_y:
                p_i.y = y_new - max_y
        # plotting()
    # plotting()
    return inhibitory_axonal_clouds


def create_grid_of_excitatory_neurons(max_x, max_y, number_of_excitatory_neurons_per_row, tuning_map):
    grid_x = number_of_excitatory_neurons_per_row
    grid_y = number_of_excitatory_neurons_per_row
    x = np.linspace(0, max_x, grid_x)
    y = np.linspace(0, max_y, grid_y)
    ex_neuron_positions = []
    ex_neuron_tuning = []
    head_dir_preference = np.zeros((grid_x, grid_y))
    for y_idx, yy in enumerate(y):
        for x_idx, xx in enumerate(x):
            p_noise = tuning_map(xx, yy)
            head_dir_preference[x_idx, y_idx] = p_noise
            ex_neuron_positions.append([xx, yy])
            ex_neuron_tuning.append(p_noise)
    return ex_neuron_positions, ex_neuron_tuning


def test_run():
    '''
    Neural sheet
    '''

    # To go from the 3d data obtained in Peng 2017 to a 2d sheet is not straightforward
    # Assuming a depth of 100um, the density becomes a 1000 interneurons/mm2 and thus about 3000-4000 exneurons/mm2
    # A long 200um short 50um interneuron axon would cover about 90 partners
    # On the other hand, the typical contacts per excitatory neuron would be A_total/(A_I*N_I) = 30

    N_E = 400
    N_I = 40

    sheet_x = 300 * um
    sheet_y = 300 * um

    inhibitory_axon_long_axis = 100 * um
    inhibitory_axon_short_axis = 25 * um

    number_of_excitatory_neurons_per_row = int(np.sqrt(N_E))

    '''
    Tuning Maps
    '''

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

    if tuning_label == "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(20, 20, 7, sheet_x / um, sheet_y / um)
        map.improve(10)
        tuning_map = lambda x, y: map.orientation_map_float(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_axonal_clouds = create_interneuron_sheet(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_clouds = create_interneuron_sheet_by_noise(ex_positions, ex_tunings, N_I,
                                                                 inhibitory_axon_long_axis / um,
                                                                 inhibitory_axon_short_axis / um, sheet_x / um,
                                                                 sheet_y / um, random_seed=3, n_iterations=10)

    plot_neural_sheet(ex_positions, ex_tunings, inhibitory_axonal_clouds)


if __name__ == "__main__":
    test_run()