superneuroabm/tutorials/user_customized_stdp.py at main · ORNL/superneuroabm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
"""
user_customized_stdp -- SuperNeuroMAT-style bounded STDP with constant depression.

Tutorial-owned learning rule for the digits notebook. Registered at runtime with
``model.register_learning_rule(step_func=user_customized_stdp, step_func_path=...)``,
which must be called before ``setup()``. The function name must not collide with a
built-in rule (``exp_pair_wise_stdp_bounded`` etc.), or registration is rejected.

Learning rule
-------------
1. When the postsynaptic neuron spikes:
      - If the presynaptic neuron fired recently, potentiate using pre_trace.
      - Otherwise, depress by a_exp_post.

2. When the presynaptic neuron spikes but the postsynaptic neuron does not,
   depress the synapse by a_exp_post.

3. Clip the updated weight to [wmin, wmax].
"""

import cupy as cp
from cupyx import jit

from superneuroabm.step_functions.synapse.util import get_soma_spike
from sagesim.math_utils import clamp


@jit.rawkernel(device="cuda")
def user_customized_stdp(
    tick,
    agent_index,
    dt,
    I_bias,
    agent_ids,
    breeds,
    locations,
    synapse_params,
    learning_params,
    internal_states,
    learning_internal_states,
    synapse_history,
    input_spikes_tensor,
    output_spikes_tensor,
    internal_states_buffer,
    learning_internal_states_buffer,
):
    t_current = int(tick)

    # ------------------------------------------------------------
    # Synapse parameters
    # ------------------------------------------------------------
    weight = synapse_params[agent_index][0]
    synaptic_delay = synapse_params[agent_index][1]

    # ------------------------------------------------------------
    # Learning parameters
    # ------------------------------------------------------------
    tau_pre_stdp = learning_params[agent_index][1]
    tau_post_stdp = learning_params[agent_index][2]

    # Amount added to the presynaptic trace on a pre spike.
    a_exp_pre = learning_params[agent_index][3]

    # Constant depression magnitude.
    a_exp_post = learning_params[agent_index][4]

    stdp_history_length = learning_params[agent_index][5]
    wmin = learning_params[agent_index][6]
    wmax = learning_params[agent_index][7]

    # ------------------------------------------------------------
    # Learning internal states
    # ------------------------------------------------------------
    pre_trace = learning_internal_states[agent_index][0]
    post_trace = learning_internal_states[agent_index][1]
    dW = learning_internal_states[agent_index][2]

    # ------------------------------------------------------------
    # Pre- and postsynaptic neuron indices
    # ------------------------------------------------------------
    pre_soma_index = locations[agent_index][0]
    post_soma_index = locations[agent_index][1]

    # ------------------------------------------------------------
    # Read pre- and postsynaptic spikes
    # ------------------------------------------------------------
    pre_soma_spike = get_soma_spike(
        tick,
        agent_index,
        dt,
        I_bias,
        agent_ids,
        pre_soma_index,
        t_current,
        input_spikes_tensor,
        output_spikes_tensor,
    )

    post_soma_spike = get_soma_spike(
        tick,
        agent_index,
        dt,
        I_bias,
        agent_ids,
        post_soma_index,
        t_current,
        input_spikes_tensor,
        output_spikes_tensor,
    )

    # ------------------------------------------------------------
    # Update traces using current spikes
    # ------------------------------------------------------------
    pre_trace = (
        pre_trace * (1.0 - dt / tau_pre_stdp)
        + pre_soma_spike * a_exp_pre
    )

    post_trace = (
        post_trace * (1.0 - dt / tau_post_stdp)
        + post_soma_spike * a_exp_post
    )

    # ------------------------------------------------------------
    # Remove traces that have decayed below their active windows
    # ------------------------------------------------------------
    trace_threshold_pre = a_exp_pre / 5.0
    trace_threshold_post = a_exp_post / 5.0

    if pre_trace <= trace_threshold_pre:
        pre_trace = 0.0

    if post_trace <= trace_threshold_post:
        post_trace = 0.0

    # ------------------------------------------------------------
    # Weight update
    # ------------------------------------------------------------
    dW = 0.0

    if post_soma_spike > 0.0:
        # The postsynaptic neuron fired.
        if pre_trace > trace_threshold_pre:
            # The presynaptic neuron fired recently:
            # reinforce the synapse.
            dW = pre_trace
        else:
            # The postsynaptic neuron fired without recent activity
            # from this presynaptic neuron.
            dW = -a_exp_post

    elif pre_soma_spike > 0.0:
        # This presynaptic neuron fired, but this postsynaptic neuron
        # did not respond.
        dW = -a_exp_post

    # Clip the updated weight.
    weight = clamp(weight + dW, wmin, wmax)

    # ------------------------------------------------------------
    # Store updated values
    # ------------------------------------------------------------
    synapse_params[agent_index][0] = weight

    learning_internal_states[agent_index][0] = pre_trace
    learning_internal_states[agent_index][1] = post_trace
    learning_internal_states[agent_index][2] = dW

    # ------------------------------------------------------------
    # State tracking buffer
    # ------------------------------------------------------------
    buffer_idx = (
        t_current
        % len(learning_internal_states_buffer[agent_index])
    )

    learning_internal_states_buffer[agent_index][buffer_idx][0] = pre_trace
    learning_internal_states_buffer[agent_index][buffer_idx][1] = post_trace
    learning_internal_states_buffer[agent_index][buffer_idx][2] = dW