structured_inhibition_spatial_network_model/scripts/spatial_maps/spatial_network_layout.py

import random as rng
import numpy as np
from matplotlib.patches import Ellipse


class Interneuron:
    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


# TODO: move to plotting utils
def get_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)
    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_interneuron = []
    phase_entropy_per_interneuron = []
    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_interneuron.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_interneuron.append(phase_entropy)
    return phase_entropy_per_interneuron


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

# TODO: There must be some methods twice
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_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 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 = Interneuron(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 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)
    interneuron_a = long_axis
    interneuron_b = short_axis
    interneuron_n = number_of_interneurons
    x_n = int(np.sqrt(interneuron_n))
    y_n = int(np.sqrt(interneuron_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(interneuron_n):
        theta = rng.random() * 2. * np.pi
        pj_x = pj * np.cos(theta)
        pj_y = pj * np.sin(theta)
        inhibitory_axonal_clouds.append(Interneuron(((n % x_n) + 1) * x_step + pj_x,
                                                    (int(n / y_n) + 1) * y_step + pj_y,
                                                    interneuron_a, interneuron_b, theta))
    return inhibitory_axonal_clouds


def create_grid_of_excitatory_neurons(sheet_size, n_e_per_row, tuning_map):
    x = np.linspace(0, sheet_size, n_e_per_row)
    y = np.linspace(0, sheet_size, n_e_per_row)
    ex_neuron_positions = []
    ex_neuron_tuning = []
    for y_idx, yy in enumerate(y):
        for x_idx, xx in enumerate(x):
            noise_value = tuning_map(xx, yy)
            ex_neuron_positions.append([xx, yy])
            ex_neuron_tuning.append(noise_value)
    return ex_neuron_positions, ex_neuron_tuning