simple.py

import numpy as np
from numpy.random import default_rng
import matplotlib
matplotlib.use('Qt5Agg')
# matplotlib.use('WebAgg')
import matplotlib.pyplot as plt


def decode_position(cpu4_reshaped, cpu4_mem_gain):
    """
    Decode position from sinusoid in to polar coordinates.
    Amplitude is distance, Angle is angle from nest outwards.
    Without offset angle gives the home vector.
    Input must have shape of (2, -1)
    """

    signal = np.sum(cpu4_reshaped, axis=0)
    fund_freq = np.fft.fft(signal)[1]
    angle = -np.angle(np.conj(fund_freq))
    distance = np.absolute(fund_freq) / cpu4_mem_gain

    return angle, distance


def pi2pi(angle, radians=True):
    """ Maps any angle to [-π, π) or [-180.0, 180.0) """
    if radians:
        half_c = np.pi
    else:
        half_c = 180.0
    return float((angle + half_c) % (2.0 * half_c) - half_c)


def zero2pi(angle, radians=True):
    """ Maps any angle to [0, 2π) or [0, 360.0) """
    if radians:
        full_c = 2*np.pi
    else:
        full_c = 360.0
    return float(angle % full_c)


def cart2pol(x, y):
    """ Converts cartesian coordinates to polar coordinates """
    rho = np.sqrt(x**2 + y**2)
    theta = np.arctan2(y, x)
    return rho, theta


def pol2cart(rho, theta):
    """ Converts polar coordinates to cartesian coordinates. """
    x = rho * np.cos(theta)
    y = rho * np.sin(theta)
    return x, y


def noisify(x, strength=0.01, strictly_positive=False, seed=None):
    arr = np.array(x).astype(float)

    generator = default_rng(seed=seed)

    if strictly_positive:
        return arr * generator.normal(loc=1, scale=strength, size=arr.shape)
    else:
        return arr + generator.normal(loc=0, scale=strength, size=arr.shape)


def sigmoid(x, slope=1.0, bias=0.5):
    arr = np.array(x).astype(float)
    return 1 / (1 + np.exp(-arr * slope - bias))


def normalise(x):
    """
    Rescales the values in the array x to be in the [0, 1] range
    """
    arr = np.array(x).astype(float)
    arr -= np.nanmin(arr)
    arr /= np.nanmax(arr)
    return np.nan_to_num(arr, nan=0.0, posinf=1.0)


def rescale(x, low=-1, high=1):
    """
    Rescales the values in the array x to be in the [low, high] range
    """
    spread = high - low
    arr = normalise(x)
    return arr * spread + low


def show_weights(arr, cmap='magma'):
    plt.imshow(arr, cmap=cmap)
    plt.show()


##

class SynapticWeights:

    def __init__(self, noise=0.0, seed=None):

        # Constants
        self.N_TL2 = 16
        self.N_CL1 = 16
        self.N_TB1 = 8
        self.N_TN1 = 2
        self.N_TN2 = 2
        self.N_CPU4 = 16
        self.N_CPU1a = 14
        self.N_CPU1b = 2
        self.N_CPU1 = self.N_CPU1a + self.N_CPU1b

        self._gen_weights()

        self._seed = seed
        self.noise = noise

    def _gen_weights(self):
        """
        Generate synaptic connections matrices based on the Central Complex anatomy.
        """
        self._binary_matrices = {}

        def _gen_TB1_TB1():
            """
            TB1-TB1 connections are a ring attractor network
            """
            x_values = np.linspace(0, 2 * np.pi, self.N_TB1, endpoint=False)  # Generate N_TB1 values between 0 and 2π
            sinus = np.sin(x_values)

            # Each row is shifted by 1 column
            weights = np.zeros([self.N_TB1, self.N_TB1])
            for i in range(self.N_TB1):
                weights[i, :] = np.roll(sinus, i)

            # Roll everything left so that the diagonal is 0
            weights = np.roll(weights, self.N_TB1 // 4, axis=1)

            return normalise(weights)

        self._binary_matrices['tb1_tb1'] = _gen_TB1_TB1()

        # All other connections are only binary matrices
        self._binary_matrices['cl1_tb1'] = np.tile(np.eye(self.N_TB1, dtype=float), 2)
        self._binary_matrices['tb1_cpu1a'] = np.tile(np.eye(self.N_TB1, dtype=float), (2, 1))[1:self.N_CPU1a + 1, :]
        self._binary_matrices['tb1_cpu1b'] = np.array([
            [0, 0, 0, 0, 0, 0, 0, 1],
            [1, 0, 0, 0, 0, 0, 0, 0]
        ], dtype=float)
        self._binary_matrices['tb1_cpu4'] = np.tile(np.eye(self.N_TB1, dtype=float), (2, 1))
        self._binary_matrices['tn_cpu4'] = np.array([
            [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
        ], dtype=float).T
        self._binary_matrices['cpu4_cpu1a'] = np.array([
            [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
            [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        ], dtype=float)
        self._binary_matrices['cpu4_cpu1b'] = np.array([
            [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],  # 8
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],  # 9
        ], dtype=float)
        self._binary_matrices['cpu1a_motor'] = np.array([
            [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]], dtype=float)
        self._binary_matrices['cpu1b_motor'] = np.array([
            [0, 1],
            [1, 0]
        ], dtype=float)
        self._binary_matrices['pontin_cpu1a'] = np.array([
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],  # 2
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
            [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
            [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],  # 15
        ], dtype=float)
        self._binary_matrices['pontin_cpu1b'] = np.array([
            [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],  # 8
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],  # 9
        ], dtype=float)
        self._binary_matrices['cpu4_pontin'] = np.eye(self.N_CPU4, dtype=float)

    @property
    def noise(self):
        return float(self._noise)

    @noise.setter
    def noise(self, val):
        self._gen_weights()
        if val != 0:
            for k, m in self._binary_matrices.items():
                self._binary_matrices[k] = noisify(m, strength=val, strictly_positive=True, seed=self._seed)
        self._noise = val

    @property
    def noisy(self):
        return self._noise != 0

    def __getitem__(self, name):
        return self._binary_matrices[name.lower().replace('-', '_')]

    def __getattr__(self, name):
        if name.lower() in self._binary_matrices:
            return self._binary_matrices[name.lower()]
        else:
            raise AttributeError(f"No such attribute: {name}")

    def __repr__(self):
        return f"SynapticWeights(noise={self.noise}, seed={self._seed})"


##

# TUNED PARAMETERS:
tl2_slope = 6.8
tl2_bias = 3.0

cl1_slope = 3.0
cl1_bias = -0.5

tb1_slope = 5.0
tb1_bias = 0.0

cpu4_slope = 5.0
cpu4_bias = 2.5

cpu1_slope = 5.0
cpu1_bias = 2.5

motor_slope = 1.0
motor_bias = 3.0

pontin_slope = 5.0
pontin_bias = 2.5

##

class CentralComplex:

    def __init__(self,
                 syn,
                 cpu4_mem_gain=0.005
                 ):

        self.syn = syn

        # The cell properties (for sigmoid function)
        self.tl2_slope = tl2_slope
        self.tl2_bias = tl2_bias
        self.cl1_bias = cl1_bias
        self.cl1_slope = cl1_slope
        self.tb1_slope = tb1_slope
        self.tb1_bias = tb1_bias
        self.cpu4_slope = cpu4_slope
        self.cpu4_bias = cpu4_bias
        # self.cpu1_slope = cpu1_slope      # TODO - WHY
        # self.cpu1_bias = cpu1_bias
        self.cpu1_slope = 7.5
        self.cpu1_bias = -1.0
        self.motor_slope = motor_slope
        self.motor_bias = motor_bias
        self.pontin_slope = pontin_slope
        self.pontin_bias = pontin_bias

        #

        self.cpu4_mem_gain = cpu4_mem_gain
        self.smoothed_flow = 0
        self.cpu4_mem_gain *= 0.5       # TODO - WHY

        self.neuron_types = []

    def tl2_output(self, theta):

        preferred_angles = np.tile(np.linspace(0, 2*np.pi, self.syn.N_TL2//2, endpoint=False), 2)
        # preferred_angles = np.linspace(0, 2*np.pi, self.syn.N_TL2, endpoint=False)
        angular_diff = theta - preferred_angles

        out = (np.cos(angular_diff) + 1) * 0.5

        return out

    def cl1_output(self, tl2):

        out = sigmoid(-tl2)

        return normalise(out)

    def tb1_output(self, cl1, tb1):

        # Proportion of input from CL1 vs TB1
        prop_cl1 = 2/3
        prop_tb1 = 1.0 - prop_cl1

        input_from_cl1 = np.dot(self.syn.CL1_TB1, cl1)
        input_from_tb1 = - np.dot(self.syn.TB1_TB1, tb1)    # TB1 are mutually inhibitory, so negative sign here

        out = sigmoid(prop_cl1 * input_from_cl1 + prop_tb1 * input_from_tb1)

        return normalise(out)

    def get_flow(self, heading, velocity, filter_steps=0):
        """Calculate optic flow depending on preference angles. [L, R]"""

        rf_angle = np.pi / 4.0

        left = [np.sin(heading - rf_angle), np.cos(heading - rf_angle)]
        right = [np.sin(heading + rf_angle), np.cos(heading + rf_angle)]

        flow = np.dot(np.array([left, right]), velocity)

        # If we are low-pass filtering speed signals (fading memory)
        if filter_steps > 0:
            self.smoothed_flow = (1.0 / filter_steps * flow + (1.0 - 1.0 / filter_steps) * self.smoothed_flow)
            flow = self.smoothed_flow
        return flow

    # def tns_output(self, travel_direction, speed):
    #     """
    #         L1                L2
    #     -45 (-π/4)        +45 (π/4)
    #
    #         L3                L4
    #     -135 (-3π/4)      +135 (3π/4)
    #     """
    #
    #     rf_angle = np.pi / 4.0
    #
    #     L1 = np.abs(speed) * np.cos(travel_direction - (- rf_angle))
    #     L2 = np.abs(speed) * np.cos(travel_direction - (+ rf_angle))
    #     L3 = np.abs(speed) * np.cos(travel_direction - (- 3 * rf_angle))
    #     L4 = np.abs(speed) * np.cos(travel_direction - (+ 3 * rf_angle))
    #
    #     return np.array([L1, L2, L3, L4])
    #
    # a= np.array([[1., 0., 1., 0.],
    #              [1., 0., 1., 0.],
    #              [1., 0., 1., 0.],
    #              [1., 0., 1., 0.],
    #              [1., 0., 1., 0.],
    #              [1., 0., 1., 0.],
    #              [1., 0., 1., 0.],
    #              [1., 0., 1., 0.],
    #              [0., 1., 0., 1.],
    #              [0., 1., 0., 1.],
    #              [0., 1., 0., 1.],
    #              [0., 1., 0., 1.],
    #              [0., 1., 0., 1.],
    #              [0., 1., 0., 1.],
    #              [0., 1., 0., 1.],
    #              [0., 1., 0., 1.]])


    def tn1_output(self, flow):
        output = (1.0 - flow) / 2.0
        # if self.noise > 0.0:
        #     output += np.random.normal(scale=self.noise, size=flow.shape)
        return np.clip(output, 0.0, 1.0)

    def tn2_output(self, flow):
        output = flow
        # if self.noise > 0.0:
        #     output += np.random.normal(scale=self.noise, size=flow.shape)
        return np.clip(output, 0.0, 1.0)


    def cpu4_update(self, cpu4_mem, tb1, tn1, tn2):
        """Memory neurons update.
        cpu4[0-7] store optic flow peaking at left 45 deg
        cpu[8-15] store optic flow peaking at right 45 deg."""
        mem_update = np.dot(self.syn.TN_CPU4, tn2)
        mem_update -= np.dot(self.syn.TB1_CPU4, tb1)
        mem_update = np.clip(mem_update, 0, 1)
        mem_update *= self.cpu4_mem_gain
        cpu4_mem += mem_update
        cpu4_mem -= 0.125 * self.cpu4_mem_gain

        return np.clip(cpu4_mem, 0.0, 1.0)

    def cpu4_update_holo(self, cpu4_mem, tb1, tn1, tn2):
        cpu4_mem_reshaped = cpu4_mem.reshape(2, -1)
        mem_update = (0.5 - tn1.reshape(2, 1)) * (1.0 - tb1)
        mem_update -= 0.5 * (0.5 - tn1.reshape(2, 1))

        # Constant purely to visualise same as rate-based model
        cpu4_mem_reshaped += self.cpu4_mem_gain * mem_update
        return np.clip(cpu4_mem_reshaped.reshape(-1), 0.0, 1.0)

    def cpu4_output(self, cpu4_mem):
        """The output from memory neuron, based on current calcium levels."""
        return sigmoid(cpu4_mem, self.cpu4_slope, self.cpu4_bias)

    def pontin_output(self, cpu4):
        inputs = np.dot(self.syn.CPU4_pontin, cpu4)
        return sigmoid(inputs, self.pontin_slope, self.pontin_bias)

    def cpu1a_output(self, tb1, cpu4):
        """The memory and direction used together to get population code for
        heading."""
        inputs = 0.5 * np.dot(self.syn.CPU4_CPU1a, cpu4)

        pontin = 0.5 * self.pontin_output(cpu4)
        inputs -= np.dot(self.syn.pontin_CPU1a, pontin)
        inputs -= np.dot(self.syn.TB1_CPU1a, tb1)

        return sigmoid(inputs, self.cpu1_slope, self.cpu1_bias)

    def cpu1b_output(self, tb1, cpu4):
        """The memory and direction used together to get population code for
        heading."""
        inputs = 0.5 * np.dot(self.syn.CPU4_CPU1b, cpu4)

        pontin = 0.5 * self.pontin_output(cpu4)
        inputs -= np.dot(self.syn.pontin_CPU1b, pontin)
        inputs -= np.dot(self.syn.TB1_CPU1b, tb1)

        return sigmoid(inputs, self.cpu1_slope, self.cpu1_bias)

    def cpu1_output(self, tb1, cpu4):
        cpu1a = self.cpu1a_output(tb1, cpu4)
        cpu1b = self.cpu1b_output(tb1, cpu4)
        return np.hstack([cpu1b[-1], cpu1a, cpu1b[0]])

    def motor_output(self, cpu1):
        """outputs a scalar where sign determines left or right turn."""
        cpu1a = cpu1[1:-1]
        cpu1b = np.array([cpu1[-1], cpu1[0]])
        motor = np.dot(self.syn.CPU1a_motor, cpu1a)
        motor += np.dot(self.syn.CPU1b_motor, cpu1b)
        output = (motor[0] - motor[1]) * 0.25  # To kill the noise a bit!
        return output

    def decode_cpu4(self, cpu4):
        """Shifts both CPU4 by +1 and -1 column to cancel 45 degree flow
        preference. When summed single sinusoid should point home."""
        cpu4_reshaped = cpu4.reshape(2, -1)
        cpu4_shifted = np.vstack([np.roll(cpu4_reshaped[0], 1),
                                  np.roll(cpu4_reshaped[1], -1)])
        return decode_position(cpu4_shifted, self.cpu4_mem_gain * 2.0)

    def decode_cpu4_holo(self, cpu4):
        """Shifts both CPU4 by +1 and -1 column to cancel 45 degree flow
        preference. When summed single sinusoid should point home."""
        cpu4_reshaped = cpu4.reshape(2, -1)
        cpu4_shifted = np.vstack([np.roll(cpu4_reshaped[0], 1),
                                  np.roll(cpu4_reshaped[1], -1)])
        return decode_position(cpu4_shifted, self.cpu4_mem_gain)

##