JuliaでIzhikevichモデルの実装をする。JuliaによるSNNの実装はSpikingNeuralNetworks.jlが既にある。この記事の実装ではJuliaの書き方を勉強すべく大幅に参考にした。以下はJulia 1.4.0で実行。ブログでは何故かJuliaのSyntax highlightがされてない。

モデルの定義

まず必要なパッケージを読み込む。Parameters.jlはstructのデフォルト値の設定のために必要。structの前に@with_kwというmacroを書くことでstructをデコレートできる。ただ、この記事を書いた後に@bicycle1885氏に以下のように教えていただいたので、Base.@kwdefを用いることにした。感謝…！ 

using Plots
using Base: @kwdef
using Parameters: @unpack # or using UnPack

変更しない定数を保持するstructのIZParameterと、変数を保持するmutable structのIZを作成する。

@kwdef struct IZParameter{FT}
    a::FT = 0.02
    b::FT = 0.2
    c::FT = -65
    d::FT = 8.0
    vrest::FT = -60.0
    vpeak::FT = 30.0
end

@kwdef mutable struct IZ{FT}
    param::IZParameter = IZParameter{FT}()
    N::Int32
    v::Vector{FT} = fill(param.vrest, N)
    u::Vector{FT} = param.b * v 
    fire::Vector{Bool} = zeros(Bool, N)
end

型をFTで指定していますが、これはFloatTypeの略。インスタンス作成時にneurons = IZ{Float32}(N=N)などとして型を指定する。元々Float32と指定していたが、この記事を書いた後に黒木玄さん(@genkuroki)にFloat32と型を限定するのは良くないと指摘していただき、次のようなnotebookを作成していただいた：https://nbviewer.jupyter.org/gist/genkuroki/41824601d6905cdab159e0c88e925ff8

変数の更新式 (forループを使う場合)

次に変数を更新する関数updateIZ!を書く。関数名に!を付けるのは、引数を関数の中で変更する場合は関数名に!記号を付けるというJuliaの習慣らしい。

function updateIZ!(variable::IZ, param::IZParameter, I::Vector, dt)
    @unpack N, v, u, fire = variable
    @unpack a, b, c, d, vrest, vpeak = param
    @inbounds for i = 1:N
        v[i] += dt * (0.04f0v[i]^2 + 5f0v[i] + 140f0 - u[i] + I[i])
        u[i] += dt * (a * (b * v[i] - u[i]))
    end
    @inbounds for i = 1:N
        fire[i] = v[i] >= vpeak
        v[i] = ifelse(fire[i], c, v[i])
        u[i] += ifelse(fire[i], d, 0)
    end
end

updateIZ! (generic function with 1 method)

forループを用いて1つのニューロンごとに膜電位vや回復変数uを更新する。SNN本ではPythonを使ったので後述するベクトル表記を使用したが、Juliaなら問題なし。macro @inboundsを使うと配列の境界確認が無くなり高速化されるらしい (Bounds checking · The Julia Language)。

シミュレーションの実行

いくつかの定数を設定してシミュレーションを実行する。恐らくもっと早い書き方があると思うのですが… (forループを関数内で書いた方が速いという話は聞いているが以下ではべた書きしている)。

T = 400 # ms
dt = 0.01f0 # ms
nt = Int32(T/dt) # number of timesteps
N = 1 # ニューロンの数

# 入力刺激
t = Array{Float32}(1:nt)*dt
I = repeat(10f0 * ((t .> 50) - (t .> 350)), 1, N)  # injection current

# 記録用
varr = zeros(Float32, nt, N)
uarr = zeros(Float32, nt, N)

# modelの定義
neurons = IZ{Float32}(N=N)

# simulation
@time for i = 1:nt
    updateIZ!(neurons, neurons.param, I[i, :], dt)
    varr[i, :] = neurons.v
    uarr[i, :] = neurons.u
end

  0.293661 seconds (768.67 k allocations: 24.452 MiB, 4.13% gc time)

ニューロンの膜電位v、回復変数u、刺激電流Iの描画をする。

p1 = plot(t, varr[:, 1])
p2 = plot(t, uarr[:, 1])
p3 = plot(t, I[:, 1])
plot(p1, p2, p3, 
    xlabel = ["" "" "Times (ms)"], 
    ylabel= ["Membrane\n potential (mV)" "Recovery\n current (pA)" "Injection\n current (pA)"],
    layout = grid(3, 1, heights=[0.5, 0.25, 0.25]), legend = false)
0 100 200 300 400 -75 -50 -25 0 25 Membrane potential (mV) 0 100 200 300 400 -10 -10 -8 -6 -4 -2 0 Recovery current (pA) 0 100 200 300 400 0.0 2.5 5.0 7.5 10.0 Times (ms) Injection current (pA)

変数の更新式 (ベクトルを使う場合)

更新式をベクトルで表現する。

function update_vecIZ!(variable::IZ, param::IZParameter, I::Vector, dt)
    @unpack N, v, u, fire = variable
    @unpack a, b, c, d, vrest, vpeak = param

    v[:] += dt * (0.04f0v.^2 + 5f0v .+ 140f0 - u + I)
    u[:] += dt * (a * (b * v - u))
    
    fire = v .>= vpeak
    u[:] += d * fire
    v[:] = v .* (1 .- fire) + c * fire
end

update_vecIZ! (generic function with 1 method)

なお、上の関数においてv[:] += ...ではなくv += ...としていたために変数の更新が行われず嵌った。Mutating function in Julia (function that modifies its arguments) - Stack Overflow によると

function right1_fill_with_twos!(v::Vector{Int64})
    v[:]=[2 for ii in 1:length(v)]
end

は機能して

function wrong1_fill_with_twos!(v::Vector{Int64})
    v=[2 for ii in 1:length(v)]
end

はダメらしい。

同様にシミュレーションを実行。Pythonと異なり、forループを用いた更新の方が速い。

# 記録用
varr2 = zeros(Float32, nt, N)
uarr2 = zeros(Float32, nt, N)

neurons = IZ{Float32}(N=N)

@time for i = 1:nt
    update_vecIZ!(neurons, neurons.param, I[i, :], dt)
    varr2[i, :] = neurons.v
    uarr2[i, :] = neurons.u
end

  1.162311 seconds (3.32 M allocations: 184.048 MiB, 5.89% gc time)

結果を表示する。

p1 = plot(t, varr2[:, 1])
p2 = plot(t, uarr2[:, 1])
p3 = plot(t, I[:, 1])
plot(p1, p2, p3, 
    xlabel = ["" "" "Times (ms)"], 
    ylabel= ["Membrane\n potential (mV)" "Recovery\n current (pA)" "Injection\n current (pA)"],
    layout = grid(3, 1, heights=[0.5, 0.25, 0.25]), legend = false)
0 100 200 300 400 -75 -50 -25 0 25 Membrane potential (mV) 0 100 200 300 400 -10 -10 -8 -6 -4 -2 0 Recovery current (pA) 0 100 200 300 400 0.0 2.5 5.0 7.5 10.0 Times (ms) Injection current (pA)

frequency-current (F-I) curve

入力を一定電流とし、その強さを1からmaxcurrentまで変更したときの発火率の変化を見る。入力電流はrangeで作成する。linspaceは廃止されたらしい(cf. Julia 0.6 から 1.x への移植)。

T = 5000 # ms
dt = 0.01f0 # ms
nt = Int32(T/dt) # number of timesteps

N = 200 # ニューロンの数

# 入力刺激
maxcurrent = 10
t = Array{Float32}(1:nt)*dt
I = Array{Float32}(range(1,maxcurrent,length=N)) # injection current

# 記録用
spikearr = zeros(Int32, nt, N)

# modelの定義
neurons = IZ{Float32}(N=N)

# simulation
@time for i = 1:nt
    updateIZ!(neurons, neurons.param, I[:], dt)
    spikearr[i, :] = neurons.fire
end

  1.931489 seconds (3.57 M allocations: 484.492 MiB, 8.79% gc time)

発火率を計算する。

fr = sum(spikearr, dims=1) / (T*1e-3)

1×200 Array{Float64,2}:
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  22.0  22.0  22.2  22.2  22.4  22.4

結果を描画する。

plot(I, fr[:], 
    xlabel="Injection current (pA)", 
    ylabel="Firing rate (Hz)",
    legend=false)
2 4 6 8 10 0 5 10 15 20 Injection current (pA) Firing rate (Hz)

記事を書いた動機

『Juliaで学ぶ計算論的神経科学』的な本を2～5年後ぐらいに何らかの形式で出したいと思っています。とりあえずSNN本(ゼロから作るSpiking Neural Networks)のJulia版をJupyter Notebook形式で公開して既存の内容についてコメントを貰いたいと思っています。Juliaに強い方、この記事のコードでもっと良い書き方があればお教えいただきたく存じます。

関連資料

Twitterで検索したらDifferentialEquations.jlを用いた実装を公開されている方がいた。