utils.py

# utils.py

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable


# ---- Connectivity & Math Functions ----

def log_epsilon(x, epsilon=1e-10):
    """ Safely computes log10 of an array, handling values close to zero. """
    return np.log10(np.maximum(x, epsilon))


def get_w_pre_post(P, p_pre, p_post, epsilon=1e-20, diagonal_zero=True):
    """ Calculates the weight matrix based on pre/post-synaptic probabilities. """
    outer = np.outer(p_post, p_pre)
    # To avoid division by zero
    outer[outer < epsilon] = epsilon

    x = P / outer
    w = log_epsilon(x, epsilon)

    if diagonal_zero:
        w[np.diag_indices_from(w)] = 0
    return w


def get_beta(p, epsilon=1e-10):
    """ Calculates the bias term 'beta' from probabilities. """
    probability = np.copy(p)
    probability[p < epsilon] = epsilon
    return np.log10(probability)


def softmax(input_vector, G=1.0, minicolumns=2):
    """ Calculates the softmax function within each hypercolumn. """
    x = np.copy(input_vector)
    x_size = x.size
    x = np.reshape(x, (x_size // minicolumns, minicolumns))
    x = G * x

    e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
    dist = e_x / e_x.sum(axis=1, keepdims=True)
    return dist.flatten()


def strict_max(x, minicolumns):
    """ Implements a strict winner-take-all within each hypercolumn. """
    x_reshaped = np.reshape(x, (x.size // minicolumns, minicolumns))
    z = np.zeros_like(x_reshaped)
    max_indices = np.argmax(x_reshaped, axis=1)
    z[np.arange(len(max_indices)), max_indices] = 1
    return z.flatten()


# ---- Data Transformation ----

def build_ortogonal_patterns(hypercolumns, minicolumns):
    """ Creates a dictionary of one-hot encoded orthogonal patterns. """
    n_units = hypercolumns * minicolumns
    patterns_dic = {}
    for i in range(minicolumns):
        pattern = np.zeros(n_units)
        for j in range(hypercolumns):
            pattern[j * minicolumns + i] = 1
        patterns_dic[i] = pattern
    return patterns_dic


# ---- Analysis Functions ----

def calculate_angle_from_history(manager):
    """ Calculates the cosine similarity between network states and stored patterns. """
    history = manager.history
    patterns_dic = manager.patterns_dic
    o = history['o']

    if o.shape[0] == 0:
        raise ValueError("History of unit activities 'o' is empty.")

    stored_pattern_indexes = sorted(manager.stored_patterns_indexes)
    num_patterns = max(stored_pattern_indexes) + 1
    angles = np.zeros((o.shape[0], num_patterns))

    for i, state in enumerate(o):
        norm_state = np.linalg.norm(state)
        if norm_state == 0: continue

        for pattern_index in stored_pattern_indexes:
            pattern = patterns_dic[pattern_index]
            dot_product = np.dot(state, pattern)
            norm_pattern = np.linalg.norm(pattern)
            angles[i, pattern_index] = dot_product / (norm_state * norm_pattern)

    return angles


def calculate_winning_pattern_from_distances(distances):
    """ Determines the winning pattern at each time step. """
    return np.argmax(distances, axis=1)


def calculate_patterns_timings(winning_patterns, dt, remove=0):
    """ Calculates the duration for which each pattern was active. """
    if len(winning_patterns) == 0:
        return []

    change_indices = np.where(np.diff(winning_patterns) != 0)[0]
    # Add the last index to capture the final pattern's duration
    indexes = np.append(change_indices, len(winning_patterns) - 1)

    patterns = winning_patterns[indexes]
    timings = []
    last_idx = -1
    for pattern, idx in zip(patterns, indexes):
        duration = (idx - last_idx) * dt
        if duration >= remove:
            timings.append((pattern, duration, (last_idx + 1) * dt))
        last_idx = idx

    return timings


def calculate_recall_success(manager, T_recall, T_cue, I_cue, target_sequence):
    """ Checks if the recalled sequence matches the target sequence. """
    manager.run_network_recall(T_recall=T_recall, T_cue=T_cue, I_cue=I_cue)

    angles = calculate_angle_from_history(manager)
    if angles.shape[0] == 0: return False

    winning_patterns = calculate_winning_pattern_from_distances(angles)
    timings = calculate_patterns_timings(winning_patterns, manager.dt, remove=0.010)
    recalled_sequence = [p[0] for p in timings]

    # Check if recalled sequence starts with the target sequence
    if len(recalled_sequence) < len(target_sequence):
        return False
    return recalled_sequence[:len(target_sequence)] == target_sequence


def calculate_recall_time_quantities(manager, T_recall, T_cue, n_trials, sequences):
    """ Runs multiple recall trials and computes performance metrics. """
    success_count = 0
    target_sequence = sequences[0]
    I_cue = target_sequence[0]

    for _ in range(n_trials):
        if calculate_recall_success(manager, T_recall, T_cue, I_cue, target_sequence):
            success_count += 1

    success_rate = (success_count / n_trials) * 100.0

    # Analyze timings from the last trial
    timings = calculate_patterns_timings(
        calculate_winning_pattern_from_distances(calculate_angle_from_history(manager)), manager.dt, remove=0.010)

    durations = [t[1] for t in timings[:len(target_sequence)]]
    total_time = sum(durations) if durations else 0
    mean_time = np.mean(durations) if len(durations) > 2 else 0
    std_time = np.std(durations) if len(durations) > 2 else 0

    return total_time, mean_time, std_time, success_rate, timings


# ---- Plotting Functions ----

def plot_weight_matrix(nn, ampa=True, ax=None):
    """ Plots the network's weight matrix. """
    if ax is None:
        fig, ax = plt.subplots(figsize=(12, 10))

    w = nn.w_ampa if ampa else nn.w
    title = 'AMPA Connectivity' if ampa else 'NMDA Connectivity'

    w_plot = w[:nn.minicolumns, :nn.minicolumns]
    v_max = np.max(np.abs(w_plot))
    v_min = -v_max

    cmap = matplotlib.cm.RdBu_r
    im = ax.imshow(w_plot, cmap=cmap, interpolation='none', vmin=v_min, vmax=v_max)
    ax.set_title(title)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.1)
    ax.figure.colorbar(im, ax=ax, cax=cax)


def plot_network_activity_angle(manager):
    """ Plots the winning pattern over time based on angle similarity. """
    fig, ax = plt.subplots(figsize=(16, 8))

    n_patterns = manager.nn.minicolumns
    T_total = manager.T_total

    angles = calculate_angle_from_history(manager)
    winning = calculate_winning_pattern_from_distances(angles) + 1  # Use +1 for color mapping

    # Filter out low-similarity states for visualization
    angles[angles < 0.1] = 0
    filtered_angles = angles * np.arange(1, angles.shape[1] + 1)

    cmap = matplotlib.cm.Paired
    cmap.set_under('white')
    extent = [0, n_patterns, T_total, 0]

    im = ax.imshow(filtered_angles, aspect='auto', interpolation='none', cmap=cmap, vmin=0.9, vmax=n_patterns,
                   extent=extent)
    ax.set_title('Sequence of Winning Patterns During Recall')
    ax.set_xlabel('Pattern Index')
    ax.set_ylabel('Time (s)')

    ticks = np.arange(1, n_patterns + 1)
    ax.set_xticks(np.arange(0.5, n_patterns + 0.5))
    ax.set_xticklabels(np.arange(0, n_patterns))

    # Colorbar
    fig.subplots_adjust(right=0.8)
    cbar_ax = fig.add_axes([0.82, 0.15, 0.03, 0.7])
    fig.colorbar(im, cax=cbar_ax, ticks=ticks, spacing='proportional')