I was surprised that I couldn't find an efficient implementation of STDP anywhere. So here it is: a fast PyTorch implementation for offline STDP.
python3 -m venv env
source env/bin/activate
pip3 install -r requirements.txt
python calc_STDP.py
The relevant code snippet is this:
The function calc_dW_from_spike_trains implements the idea of implementing STDP via spike-tracking I got from the Neuromatch Academy Bonus Tutorial: Spike-timing dependent plasticity (STDP) (Content creators: Qinglong Gu, Songtin Li, John Murray, Richard Naud, Arvind Kumar).
This implementation is offline/batched, meaning that STDP is not calculated during every time step (online), but across a recorded time series of n_time_steps steps (e.g. a stimulus presentation).
Comes with two STDP kernels for free: exp_kernel and box_kernel.