superneuroabm 1.0.0__py3-none-any.whl

superneuroabm/step_functions/soma/izh.py

@@ -0,0 +1,86 @@
+"""
+Izhikevich Neuron and weighted synapse step functions for spiking neural networks
+
+"""
+
+import cupy as cp
+from cupyx import jit
+
+
+@jit.rawkernel(device="cuda")
+def izh_soma_step_func(
+    tick,
+    agent_index,
+    globals,
+    agent_ids,
+    breeds,
+    locations,
+    neuron_params,  # k, vth, C, a, b,
+    learning_params,
+    internal_state,  # v, u
+    internal_learning_state,
+    synapse_history,  # Synapse delay
+    input_spikes_tensor,  # input spikes
+    output_spikes_tensor,
+    internal_states_buffer,
+    internal_learning_states_buffer,
+):
+    synapse_indices = locations[agent_index]  # network location is defined by neighbors
+
+    I_synapse = 0.0
+
+    # synapse_indices now contains pre-computed local indices (converted in SAGESim)
+    # No linear search needed!
+    for i in range(len(synapse_indices)):
+        synapse_index = int(synapse_indices[i])
+        if synapse_index >= 0 and not cp.isnan(synapse_indices[i]):
+            I_synapse += internal_state[synapse_index][0]
+
+    # Get the current time step value:
+    t_current = int(tick)
+    dt = globals[0]  # time step size
+    I_bias = globals[1]  # bias current
+
+    # NOTE: neuron_params would need to be as long as the max number of params in any spiking neuron model
+    k = neuron_params[agent_index][0]
+    vthr = neuron_params[agent_index][1]
+    C = neuron_params[agent_index][2]
+    a = neuron_params[agent_index][3]
+    b = neuron_params[agent_index][4]
+    vpeak = neuron_params[agent_index][5]
+    vrest = neuron_params[agent_index][6]
+    d = neuron_params[agent_index][7]
+    vreset = neuron_params[agent_index][8]
+    I_in = neuron_params[agent_index][9]
+
+    # From https://www.izhikevich.org/publications/spikes.htm
+    # v' = 0.04v^2 + 5v + 140 -u + I
+    # u' = a(bv - u)
+    # if v=30mV: v = c, u = u + d, spike
+
+    # dv = (k*(internal_state[my_idx]-vrest)*(internal_state[my_idx]-vthr)-u[my_idx]+I) / C
+    # internal_state: [0] - v, [1] - u
+    # NOTE: size of internal_state would need to be set as the maximum possible state variables of any spiking neuron
+
+    v = internal_state[agent_index][0]
+    u = internal_state[agent_index][1]
+
+    dv = (k * (v - vrest) * (v - vthr) - u + I_synapse + I_bias + I_in) / C
+    v = v + dt * dv * 1e3
+
+    u += dt * 1e3 * (a * (b * (v - vrest) - u))
+    # s = 1 * (v >= vthr)  # output spike
+    s = 1 * (v >= vpeak)  # output spike
+    u = u + d * s  # If spiked, update recovery variable
+    v = v * (1 - s) + vreset * s  # If spiked, reset membrane potential
+
+    internal_state[agent_index][0] = v
+    internal_state[agent_index][1] = u
+
+    output_spikes_tensor[agent_index][t_current] = s
+
+    # Safe buffer indexing: use modulo to prevent out-of-bounds access
+    # When tracking is disabled, buffer length is 1, so t_current % 1 = 0 always
+    buffer_idx = t_current % len(internal_states_buffer[agent_index])
+    internal_states_buffer[agent_index][buffer_idx][0] = v
+    internal_states_buffer[agent_index][buffer_idx][1] = u

superneuroabm/step_functions/soma/lif.py

@@ -0,0 +1,98 @@
+"""
+LIF Neuron and weighted synapse step functions for spiking neural networks
+
+"""
+
+from cupyx import jit
+import cupy as cp
+
+
+@jit.rawkernel(device="cuda")
+def lif_soma_step_func(  # NOTE: update the name to soma_step_func from neuron_step_func
+    tick,
+    agent_index,
+    globals,
+    agent_ids,
+    breeds,
+    locations,
+    neuron_params,  # k, vth, C, a, b,
+    learning_params,
+    internal_state,  # v, u
+    internal_learning_state,
+    synapse_history,  # Synapse delay
+    input_spikes_tensor,  # input spikes
+    output_spikes_tensor,
+    internal_states_buffer,
+    internal_learning_states_buffer,
+):
+    synapse_indices = locations[agent_index]  # Now contains local indices instead of IDs
+
+    I_synapse = 0.0
+
+    # synapse_indices now contains pre-computed local indices (converted in SAGESim)
+    # No linear search needed!
+    for i in range(len(synapse_indices)):
+        synapse_index = int(synapse_indices[i])
+        if synapse_index >= 0 and not cp.isnan(synapse_indices[i]):
+            I_synapse += internal_state[synapse_index][0]
+
+    # Get the current time step value:
+    t_current = int(tick)  # Check if tcount is needed or if we ca use this directly.
+    dt = globals[0]  # time step size
+    I_bias = globals[1]  # bias current
+
+    # NOTE: neuron_params would need to as long as the max number of params in any spiking neuron model
+    # Neuron Parameter
+    C = neuron_params[agent_index][0]  # membrane capacitance
+    R = neuron_params[agent_index][1]  # Leak resistance
+    vthr = neuron_params[agent_index][2]  # spike threshold
+    tref = neuron_params[agent_index][3]  # refractory period
+    vrest = neuron_params[agent_index][4]  # resting potential
+    vreset = neuron_params[agent_index][5]  # reset potential
+    tref_allows_integration = neuron_params[agent_index][
+        6
+    ]  # whether to allow integration during refractory period
+    I_in = neuron_params[agent_index][7]  # input current
+    scaling_factor = neuron_params[agent_index][
+        8
+    ]  # scaling factor for synaptic current
+    # vreset = neuron_params[agent_index][8]
+    # I_in = neuron_params[agent_index][9]
+
+    # NOTE: size of internal_state would need to be set as the maximum possible state variables of any spiking neuron
+    # Internal state variables
+    v = internal_state[agent_index][0]  # membrane potential
+    tcount = internal_state[agent_index][
+        1
+    ]  # time count from the start of the simulation
+    tlast = internal_state[agent_index][2]  # last spike time
+
+    # Calculate the membrane potential update
+    dv = (vrest - v) / (R * C) + (I_synapse * scaling_factor + I_bias + I_in) / C
+
+    v += (
+        (dv * dt)
+        if ((dt * tcount) > (tlast + tref)) or tref_allows_integration
+        else 0.0
+    )
+
+    #if tlast > 0 else 1 # output spike only happens if the membrane potential exceeds the threshold and the neuron is not in refractory period.
+    s = 1.0 * ((v >= vthr) and (( dt * tcount > tlast + tref) if tlast > 0 else True))
+
+
+    tlast = tlast * (1 - s) + dt * tcount * s
+    v = v * (1 - s) + vreset * s  # If spiked, reset membrane potential
+
+    internal_state[agent_index][0] = v
+    internal_state[agent_index][1] += 1
+    internal_state[agent_index][2] = tlast
+
+    output_spikes_tensor[agent_index][t_current] = s
+
+    # Safe buffer indexing: use modulo to prevent out-of-bounds access
+    # When tracking is disabled, buffer length is 1, so t_current % 1 = 0 always
+    buffer_idx = t_current % len(internal_states_buffer[agent_index])
+    internal_states_buffer[agent_index][buffer_idx][0] = v
+    internal_states_buffer[agent_index][buffer_idx][1] = internal_state[agent_index][1] + 1
+    internal_states_buffer[agent_index][buffer_idx][2] = tlast
+ 

superneuroabm/step_functions/soma/lif_soma_adaptive_thr.py

@@ -0,0 +1,111 @@
+"""
+LIF Neuron and weighted synapse step functions for spiking neural networks
+
+"""
+
+from cupyx import jit
+import cupy as cp
+
+
+@jit.rawkernel(device="cuda")
+def lif_soma_adaptive_thr_step_func(  # NOTE: update the name to soma_step_func from neuron_step_func
+    tick,
+    agent_index,
+    globals,
+    agent_ids,
+    breeds,
+    locations,
+    neuron_params,  # k, vth, C, a, b,
+    learning_params,
+    internal_state,  # v, u
+    internal_learning_state,
+    synapse_history,  # Synapse delay
+    input_spikes_tensor,  # input spikes
+    output_spikes_tensor,
+    internal_states_buffer,
+    internal_learning_states_buffer,
+):
+    synapse_indices = locations[agent_index]  # Now contains local indices instead of IDs
+
+    I_synapse = 0.0
+
+    # synapse_indices now contains pre-computed local indices (converted in SAGESim)
+    # No linear search needed!
+    for i in range(len(synapse_indices)):
+        synapse_index = int(synapse_indices[i])
+        if synapse_index >= 0 and not cp.isnan(synapse_indices[i]):
+            I_synapse += internal_state[synapse_index][0]
+
+    # Get the current time step value:
+    t_current = int(tick)  # Check if tcount is needed or if we ca use this directly.
+    dt = globals[0]  # time step size
+    I_bias = globals[1]  # bias current
+
+    # NOTE: neuron_params would need to as long as the max number of params in any spiking neuron model
+    # Neuron Parameter
+    C = neuron_params[agent_index][0]  # membrane capacitance
+    R = neuron_params[agent_index][1]  # Leak resistance
+    vthr_initial = neuron_params[agent_index][2]  # initial spike threshold
+    tref = neuron_params[agent_index][3]  # refractory period
+    vrest = neuron_params[agent_index][4]  # resting potential
+    vreset = neuron_params[agent_index][5]  # reset potential
+    tref_allows_integration = neuron_params[agent_index][6]  # whether to allow integration during refractory period
+    I_in = neuron_params[agent_index][7]  # input current
+    scaling_factor = neuron_params[agent_index][8]  # scaling factor for synaptic current
+    # Adaptive threshold parameters
+    delta_thr      = neuron_params[agent_index][9]     # threshold increment
+    tau_decay_thr  = neuron_params[agent_index][10]    # threshold decay constant
+
+
+
+    # vreset = neuron_params[agent_index][8]
+    # I_in = neuron_params[agent_index][9]
+
+    # NOTE: size of internal_state would need to be set as the maximum possible state variables of any spiking neuron
+    # Internal state variables
+    v = internal_state[agent_index][0]  # membrane potential
+    tcount = internal_state[agent_index][
+        1
+    ]  # time count from the start of the simulation
+    tlast = internal_state[agent_index][2]  # last spike time
+    vthr = internal_state[agent_index][3]  # spike threshold (updated value)
+    # Calculate the membrane potential update
+    dv = (vrest - v) / (R * C) + (I_synapse * scaling_factor + I_bias + I_in) / C
+
+    v += (
+        (dv * dt)
+        if ((dt * tcount) > (tlast + tref)) or tref_allows_integration
+        else 0.0
+    )
+
+    #if tlast > 0 else 1 # output spike only happens if the membrane potential exceeds the threshold and the neuron is not in refractory period.
+    s = 1.0 * ((v >= vthr) and (( dt * tcount > tlast + tref) if tlast > 0 else True))
+    # -----------------------------
+    # ADAPTIVE THRESHOLD UPDATE
+    # -----------------------------
+    if s == 1.0:
+        # Increase threshold after spike
+        vthr += delta_thr
+    else:
+        # Exponential decay: vthr = vrest_thr + (vthr - vrest_thr)*exp(-dt/tau)
+        vthr = vthr_initial + (vthr - vthr_initial) * cp.exp(-dt / tau_decay_thr)
+
+
+    tlast = tlast * (1 - s) + dt * tcount * s
+    v = v * (1 - s) + vreset * s  # If spiked, reset membrane potential
+
+    internal_state[agent_index][0] = v
+    internal_state[agent_index][1] += 1
+    internal_state[agent_index][2] = tlast
+    # Write back updated threshold state
+    internal_state[agent_index][3] =vthr
+
+    output_spikes_tensor[agent_index][t_current] = s
+
+    # Safe buffer indexing: use modulo to prevent out-of-bounds access
+    # When tracking is disabled, buffer length is 1, so t_current % 1 = 0 always
+    buffer_idx = t_current % len(internal_states_buffer[agent_index])
+    internal_states_buffer[agent_index][buffer_idx][0] = v
+    internal_states_buffer[agent_index][buffer_idx][1] = internal_state[agent_index][1] + 1
+    internal_states_buffer[agent_index][buffer_idx][2] = tlast
+    internal_states_buffer[agent_index][buffer_idx][3] = vthr

superneuroabm/step_functions/synapse/single_exp.py

@@ -0,0 +1,71 @@
+"""
+Single exponential synapse step functions for spiking neural networks
+
+"""
+
+import cupy as cp
+import numpy as np
+from cupyx import jit
+
+from superneuroabm.step_functions.synapse.util import get_soma_spike
+from sagesim.utils import get_neighbor_data_from_tensor
+
+
+@jit.rawkernel(device="cuda")
+def synapse_single_exp_step_func(
+    tick,
+    agent_index,
+    globals,
+    agent_ids,
+    breeds,
+    locations,
+    synapse_params,  # scale, time constant (tau_rise and tau_fall)
+    learning_params,
+    internal_state,  #
+    internal_learning_state,
+    synapse_history,  # delay
+    input_spikes_tensor,  # input spikes
+    output_spikes_tensor,
+    internal_states_buffer,
+    internal_learning_states_buffer,
+):
+    t_current = int(tick)
+
+    dt = globals[0]  # time step size
+
+    weight = synapse_params[agent_index][0]
+    synaptic_delay = synapse_params[agent_index][1]
+    scale = synapse_params[agent_index][2]
+    tau_fall = synapse_params[agent_index][3]
+    tau_rise = synapse_params[agent_index][4]
+
+    # locations[agent_index] = [pre_soma_index, post_soma_index]
+    # SAGESim has already converted agent IDs to local indices
+    pre_soma_index, post_soma_index = (
+        locations[agent_index][0],
+        locations[agent_index][1],
+    )
+
+    spike = get_soma_spike(
+        tick,
+        agent_index,
+        globals,
+        agent_ids,
+        pre_soma_index,
+        t_current,
+        input_spikes_tensor,
+        output_spikes_tensor,
+    )
+
+    I_synapse = internal_state[agent_index][0]
+
+    I_synapse = I_synapse * (1 - dt / tau_fall) + spike * scale * weight
+
+    internal_state[agent_index][0] = I_synapse
+
+    # Safe buffer indexing: use modulo to prevent out-of-bounds access
+    # When tracking is disabled, buffer length is 1, so t_current % 1 = 0 always
+    buffer_idx = t_current % len(internal_states_buffer[agent_index])
+    internal_states_buffer[agent_index][buffer_idx][0] = I_synapse
+    internal_states_buffer[agent_index][buffer_idx][1] = spike
+    internal_states_buffer[agent_index][buffer_idx][2] = t_current

superneuroabm/step_functions/synapse/stdp/Low_resolution_synapse.py

@@ -0,0 +1,117 @@
+import numpy as np
+from cupyx import jit
+
+from superneuroabm.step_functions.synapse.util import get_pre_soma_spike
+
+@jit.rawkernel(device="cuda")
+def synapse_single_exp_step_func(
+    tick,
+    agent_index,
+    globals,
+    agent_ids,
+    breeds,
+    locations,
+    synapse_params,          # Layout modified to add learning sigmas
+    learning_params,         # Adding sigmas for learning noise
+    internal_state,          # I_synapse
+    internal_learning_state, # Added w_eff, flags, rng state
+    synapse_history,         
+    input_spikes_tensor,
+    output_spikes_tensor,
+    internal_states_buffer,
+    internal_learning_states_buffer,
+):
+    # --- tiny device helpers (inline for rawkernel styles) -------------------
+    def _xorshift32(state_u32):
+        # returns (new_state, uint32)
+        x = state_u32 & 0xFFFFFFFF
+        x ^= (x << 13) & 0xFFFFFFFF
+        x ^= (x >> 17) & 0xFFFFFFFF
+        x ^= (x << 5)  & 0xFFFFFFFF
+        return x, x
+
+    def _u01_from_u32(u32):
+        # map to (0,1); avoid exact 0
+        return (u32 + 1.0) * (1.0 / 4294967297.0)
+
+    def _randn(state_u32):
+        # Box–Muller using two uniforms
+        s, u1_i = _xorshift32(state_u32)
+        s, u2_i = _xorshift32(s)
+        u1 = _u01_from_u32(u1_i)
+        u2 = _u01_from_u32(u2_i)
+        # sqrt(-2 ln u1) * cos(2 pi u2)
+        r = math.sqrt(-2.0 * math.log(u1))
+        z = r * math.cos(6.283185307179586 * u2)
+        return s, z  # new_state, standard normal
+    # -------------------------------------------------------------------------
+
+    t_current = int(tick)
+    dt = globals[0]
+
+    # ---- read params ---------------------------------------------------------
+    w_nom          = synapse_params[agent_index][0]
+    synaptic_delay = synapse_params[agent_index][1]
+    scale          = synapse_params[agent_index][2]
+    tau_fall       = synapse_params[agent_index][3]
+    tau_rise       = synapse_params[agent_index][4]
+
+    # per-synapse standard deviations for weight noise
+    sigma_prog     = synapse_params[agent_index][5] # standard deviation for initial programming
+    sigma_stdp     = synapse_params[agent_index][6] # standard deviation for STDP updates
+
+    # ---- persistent learning state ------------------------------------------
+    w_eff          = internal_learning_state[agent_index][0]
+    is_programmed  = int(internal_learning_state[agent_index][1]) # bool as int
+    rng_state_u32  = int(internal_learning_state[agent_index][2])
+    needs_reprog   = int(internal_learning_state[agent_index][3])
+
+    # seed rng at first use (deterministic per synapse)
+    if rng_state_u32 == 0:
+        # simple seed from agent_index (offset avoids zero)
+        rng_state_u32 = (1664525 * (agent_index + 1) + 1013904223) & 0xFFFFFFFF
+        if rng_state_u32 == 0:
+            rng_state_u32 = 123456789
+
+    # initial programming (once)
+    if is_programmed == 0:
+        rng_state_u32, z = _randn(rng_state_u32)
+        # choose additive or multiplicative model;
+        # (a) additive: w_eff = w_nom + sigma_prog * z
+        # (b) multiplicative: w_eff = w_nom * (1.0 + sigma_prog * z)
+        w_eff = w_nom + sigma_prog * z   # Additive, Katie's STDP is multiplicative in the orignal code
+        is_programmed = 1
+        needs_reprog  = 0
+
+    # reprogram after STDP changed nominal weight
+    if needs_reprog == 1:
+        rng_state_u32, z = _randn(rng_state_u32)
+        # apply STDP programming noise
+        w_eff = w_nom + sigma_stdp * z   # or multiplicative form
+        needs_reprog = 0
+
+    # store back persistent learning state
+    internal_learning_state[agent_index][0] = w_eff
+    internal_learning_state[agent_index][1] = is_programmed
+    internal_learning_state[agent_index][2] = rng_state_u32
+    internal_learning_state[agent_index][3] = needs_reprog
+
+    # ---- spikes and current update (unchanged except weight source) ----------
+    location_data = locations[agent_index]
+    pre_soma_id = -1 if cp.isnan(location_data[1]) else location_data[0]
+    post_soma_id = location_data[0] if cp.isnan(location_data[1]) else location_data[1]
+
+    spike = get_pre_soma_spike(
+        tick, agent_index, globals, agent_ids,
+        pre_soma_id, t_current, input_spikes_tensor, output_spikes_tensor
+    )
+
+    I_synapse = internal_state[agent_index][0]
+    # use w_eff (constant between programming events)
+    I_synapse = I_synapse * (1 - dt / tau_fall) + spike * scale * w_eff
+
+    internal_state[agent_index][0] = I_synapse
+    internal_states_buffer[agent_index][t_current][0] = I_synapse
+    internal_states_buffer[agent_index][t_current][1] = spike
+    internal_states_buffer[agent_index][t_current][2] = t_current
+    internal_states_buffer[agent_index][t_current][3] = pre_soma_id

superneuroabm/step_functions/synapse/stdp/Three-bit_exp_pair_wise.py

@@ -0,0 +1,130 @@
+"""
+Exponential STDP (Spike-Timing Dependent Plasticity) step function for spiking neural networks
+
+"""
+
+import cupy as cp
+from cupyx import jit
+
+from superneuroabm.step_functions.synapse.util import get_soma_spike
+
+
+@jit.rawkernel(device="cuda")
+def exp_pair_wise_stdp_quantized(
+    tick,
+    agent_index,
+    globals,
+    agent_ids,
+    breeds,
+    locations,
+    connectivity,
+    synapse_params,  # scale, time constant (tau_rise and tau_fall)
+    learning_params,
+    internal_state,  #
+    internal_learning_state,  # learning state variables
+    synapse_history,  # delay
+    input_spikes_tensor,  # input spikes
+    output_spikes_tensor,
+    internal_states_buffer,
+    internal_learning_states_buffer,
+):
+    t_current = int(tick)
+
+    dt = globals[0]  # time step size
+
+    # Get the synapse parameters:
+    weight = synapse_params[agent_index][0]
+    synaptic_delay = synapse_params[agent_index][1]
+
+    # Get the learning parameters:
+    # stdpType = 0 # Parsed in the learning rule selector
+    tau_pre_stdp = learning_params[agent_index][1]
+    tau_post_stdp = learning_params[agent_index][2]
+    a_exp_pre = learning_params[agent_index][3]
+    a_exp_post = learning_params[agent_index][4]
+    stdp_history_length = learning_params[agent_index][5]
+    # Wmax, Wmin
+
+    pre_trace = internal_learning_state[agent_index][0]
+    post_trace = internal_learning_state[agent_index][1]
+    dW = internal_learning_state[agent_index][2]
+
+    # Get pre and post soma IDs from connectivity (contains agent IDs, not converted by SAGESim)
+    pre_soma_id = connectivity[agent_index][0]
+    post_soma_id = connectivity[agent_index][1]
+
+    # Get the pre-soma spike
+    pre_soma_spike = get_soma_spike(
+        tick,
+        agent_index,
+        globals,
+        agent_ids,
+        pre_soma_id,
+        t_current,
+        input_spikes_tensor,
+        output_spikes_tensor,
+    )
+
+    post_soma_spike = get_soma_spike(
+        tick,
+        agent_index,
+        globals,
+        agent_ids,
+        post_soma_id,
+        t_current,
+        input_spikes_tensor,
+        output_spikes_tensor,
+    )
+
+    pre_trace = pre_trace * (1 - dt / tau_pre_stdp) + pre_soma_spike * a_exp_pre
+    post_trace = post_trace * (1 - dt / tau_post_stdp) + post_soma_spike * a_exp_post
+    dW = pre_trace * post_soma_spike - post_trace * pre_soma_spike
+
+    weight += dW  # Update the weight
+
+    # === 3-bit quantization ===
+    wmin = 0#learning_params[agent_index][6]  # assuming stored in learning_params
+    wmax = 14#learning_params[agent_index][7]
+    num_levels = 8  # 3 bits -> 8 quantization levels#learning_params[agent_index][8]
+    delta = (wmax - wmin) / (num_levels - 1)
+    weight = cp.clip(weight, wmin, wmax)
+    quantized_weight = cp.round((weight - wmin) / delta) * delta + wmin
+    weight = quantized_weight
+    # ==========================
+
+    synapse_params[agent_index][0] = weight  # Update quantized weight
+
+    internal_learning_state[agent_index][0] = pre_trace
+    internal_learning_state[agent_index][1] = post_trace
+    internal_learning_state[agent_index][2] = dW
+
+    # Safe buffer indexing: use modulo to prevent out-of-bounds access
+    # When tracking is disabled, buffer length is 1, so t_current % 1 = 0 always
+    buffer_idx = t_current % len(internal_learning_states_buffer[agent_index])
+    internal_learning_states_buffer[agent_index][buffer_idx][0] = pre_trace
+    internal_learning_states_buffer[agent_index][buffer_idx][1] = post_trace
+    internal_learning_states_buffer[agent_index][buffer_idx][2] = dW
+
+    # spike_pre_[t_current] = pre_soma_spike #spike_pre_ is an array of size (stdp_history_length, number of input neurons), pre_soma_spike is (number of input neurons,)
+    # spike_post_[:, t_current] = post_soma_spike#spike_post_ is an array of size (number of output neurons,stdp_history_length), post_soma_spike is (number of output neurons,)
+    # trace_pre_[t_current] = pre_trace #Corresponding traces an array of size (stdp_history_length,number of input neurons), pre_trace is (number of input neurons,)
+    # trace_post_[:, t_current] = post_trace #Corresponding traces is an array of size (number of output neurons,stdp_history_length)
+
+    # if t_current == stdp_history_length:
+    #     dW = cp.dot(spike_post_, trace_pre_)#(1,stdp_history_length) dot (stdp_history_length,1) we might need additional learning rate and multiplicative STDP*(wmax - W)*
+    #     dW -=cp.dot(trace_post_, spike_pre_)#(1,stdp_history_length) dot (stdp_history_length,1),  add learning rat*W for multiplicative STDP
+    #     clipped_dW = cp.clip(dW / stdp_history_length, dw_max, dw_min)  # Clip the weight change if needed
+    #     weight = cp.clip(weight+clipped_dW,wmin, wmax)  # Update the weight
+    #     #reset the traces and spikes buffers
+    #     spike_pre_ = cp.zeros((stdp_history_length, number_of_input_neurons), dtype=cp.float32)
+    #     spike_post_ = cp.zeros((number_of_output_neurons, stdp_history_length), dtype=cp.float32)
+    #     trace_pre_ = cp.zeros((stdp_history_length, number_of_input_neurons), dtype=cp.float32)
+    #     trace_post_ = cp.zeros((number_of_output_neurons, stdp_history_length), dtype=cp.float32)
+
+    internal_state[agent_index][2] = pre_trace
+    internal_state[agent_index][3] = post_trace
+
+    # Safe buffer indexing for internal_states_buffer (reuse buffer_idx from above)
+    state_buffer_idx = t_current % len(internal_states_buffer[agent_index])
+    internal_states_buffer[agent_index][state_buffer_idx][2] = post_soma_spike
+    internal_states_buffer[agent_index][state_buffer_idx][3] = post_trace