# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # STDP Basics, 'offline'
#
# Demonstrates calculate dw from the stdp learning rule, applied to spike processes and firing records (as can be obtained from the model.step_time end timestep callback).

import numpy as np
from spikeflow import firing_to_spike_process, firings_to_spike_processes
from spikeflow import spike_process_delta_times
from spikeflow import STDPParams, stdp_offline_dw, stdp_offline_dw_process, stdp_offline_dw_processes, stdp_offline_dw_firings

# # Make some data: create a couple firing records and spike processes

# +
f1 = np.array([True, False, False, True, True, False, False, True])
f2 = np.array([False, False, True, False, False, True, True, False])
print('firing 1:', f1.shape, f1)
print('firing 2:', f2.shape, f2)

s1 = firing_to_spike_process(f1)
s2 = firing_to_spike_process(f2)
print('spike process 1:', s1.shape, s1)
print('spike process 2:', s2.shape, s2)
# -

# # Apply STDP learning rule to calculate change in weight in various ways

stdp_params = STDPParams(APlus=1.0, AMinus=1.0, TauPlus=10.0, TauMinus=10.0)

# ## ... calculate from time deltas, to check the math

# +
delta_times = spike_process_delta_times(s1, s2)
print('Delta times shape', delta_times.shape)
print(delta_times)

dw = stdp_offline_dw(delta_times, stdp_params)
print('dw  shape', dw.shape)
print(dw)
# -

# ## ... directly from spike processes

dw_p = stdp_offline_dw_process(s1, s2, stdp_params)
print('SUM dw from time deltas:', sum(dw))
print('SUM dw from spike processes:', dw_p)

# ## ... but like a 1x1 weight matrix

dW = stdp_offline_dw_processes(np.ones((1,1), dtype=np.float32), [s1], [s2], stdp_params)
print('dW  shape', dW.shape)
print(dW)

# ## ... or like a bigger, sparse weight matrix

W = np.array([[1, 0, 0], [0, 1, 1]], dtype=np.float32)
dW = stdp_offline_dw_processes(W, [s1, s1], [s2, s2, s1], stdp_params)
print('W  shape', W.shape)
print(W)
print('dW  shape', dW.shape)
print(dW)

# ## ... same but from firing matrices
#
# (as can be extracted from model.step_time's end timestep callback)

# +
f_in = np.vstack([f1, f1]).T
f_out = np.vstack([f2, f2, f1]).T
print('f_in  shape', f_in.shape)
print('f_out  shape', f_out.shape)

dWf = stdp_offline_dw_firings(W, f_in, f_out, stdp_params)
print('dWf  shape', dWf.shape)
print(dWf)

StarcoderdataJupyterScriptsDedupFiltered