diff --git a/Project.toml b/Project.toml index 1a9965b8..9297de82 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ version = "0.1.0" [deps] AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" Catalyst = "479239e8-5488-4da2-87a7-35f2df7eef83" +ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" @@ -40,6 +41,7 @@ Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" SignalAnalysis = "df1fea92-c066-49dd-8b36-eace3378ea47" @@ -119,4 +121,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" [targets] -test = ["DifferentialEquations","LinearAlgebra", "ForwardDiff", "Distributions", "Flux", "Random", "SparseArrays", "Optim", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "OrdinaryDiffEq", "StochasticDiffEq", "SafeTestsets", "Turing", "Test", "Tracker"] +test = ["DifferentialEquations", "LinearAlgebra", "ForwardDiff", "Distributions", "Flux", "Random", "SparseArrays", "Optim", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "OrdinaryDiffEq", "StochasticDiffEq", "SafeTestsets", "Turing", "Test", "Tracker"] diff --git a/examples/.DS_Store b/examples/.DS_Store deleted file mode 100644 index 7d39d35a..00000000 Binary files a/examples/.DS_Store and /dev/null differ diff --git a/examples/wholebrainWC/.DS_Store b/examples/wholebrainWC/.DS_Store deleted file mode 100644 index 754c800e..00000000 Binary files a/examples/wholebrainWC/.DS_Store and /dev/null differ diff --git a/src/.DS_Store b/src/.DS_Store deleted file mode 100644 index 6b163cf7..00000000 Binary files a/src/.DS_Store and /dev/null differ diff --git a/src/Neuroblox.jl b/src/Neuroblox.jl index aa97cccf..8db33a31 100644 --- a/src/Neuroblox.jl +++ b/src/Neuroblox.jl @@ -55,13 +55,16 @@ abstract type BloxConnectComplex <: BloxConnection end abstract type BloxConnectMultiFloat <: BloxConnection end abstract type BloxConnectMultiComplex <: BloxConnection end +# dictionary type for Blox parameters +Para_dict = Dict{Symbol, Union{<: Real, Num}} + include("Neurographs.jl") +include("blox/blox_utilities.jl") include("utilities/spectral_tools.jl") include("utilities/learning_tools.jl") include("control/controlerror.jl") include("measurementmodels/fmri.jl") include("datafitting/spectralDCM.jl") -include("blox/blox_utilities.jl") include("blox/neural_mass.jl") include("blox/cortical_blox.jl") include("blox/canonicalmicrocircuit.jl") diff --git a/src/blox/blox_utilities.jl b/src/blox/blox_utilities.jl index f005859e..abd394d9 100644 --- a/src/blox/blox_utilities.jl +++ b/src/blox/blox_utilities.jl @@ -1,11 +1,58 @@ -function scope_dict(para_dict::Dict{Symbol, Union{Real,Num}}) - para_dict_copy = copy(para_dict) - for (n,v) in para_dict_copy - if typeof(v) == Num - para_dict_copy[n] = ParentScope(v) +function progress_scope(params; lvl=0) + para_list = [] + for p in params + pp = ModelingToolkit.unwrap(p) + if ModelingToolkit.hasdefault(pp) + d = ModelingToolkit.getdefault(pp) + if typeof(d)==SymbolicUtils.BasicSymbolic{Real} + if lvl==0 + pp = ParentScope(pp) + else + pp = DelayParentScope(pp,lvl) + end + end + end + push!(para_list,ModelingToolkit.wrap(pp)) + end + return para_list +end + +""" +This function progresses the scope of parameters and leaves floating point values untouched +""" +function progress_scope(args...) + paramlist = [] + for p in args + if p isa Float64 + push!(paramlist, p) else - para_dict_copy[n] = (@parameters $n=v)[1] + p = ParentScope(p) + # pp = ModelingToolkit.unwrap(p) + # if ModelingToolkit.hasdefault(pp) + # d = ModelingToolkit.getdefault(pp) + # if typeof(d)==SymbolicUtils.BasicSymbolic{Real} + # pp = ParentScope(pp) + # end + # end + # push!(para_list,ModelingToolkit.wrap(pp)) + push!(paramlist, p) end end - return para_dict_copy + return paramlist end + +""" + This function compiles already existing parameters with floats after making them parameters. + Keyword arguments are used because parameter definition requires names, not just values +""" +function compileparameterlist(;kwargs...) + paramlist = [] + for (kw, v) in kwargs + if v isa Float64 + paramlist = vcat(paramlist, @parameters $kw = v) + else + paramlist = vcat(paramlist, v) + end + end + return paramlist +end \ No newline at end of file diff --git a/src/blox/neural_mass.jl b/src/blox/neural_mass.jl index 9283eed6..b0e5bdd1 100644 --- a/src/blox/neural_mass.jl +++ b/src/blox/neural_mass.jl @@ -14,23 +14,19 @@ end mutable struct HarmonicOscillatorBlox <: NeuralMassBlox # all parameters are Num as to allow symbolic expressions - p_dict::Dict{Symbol,Union{Real,Num}} connector::Num noDetail::Vector{Num} detail::Vector{Num} initial::Dict{Num, Tuple{Float64, Float64}} odesystem::ODESystem function HarmonicOscillatorBlox(;name, ω=25*(2*pi), ζ=1.0, k=625*(2*pi), h=35.0) - para_dict = scope_dict(Dict{Symbol,Union{Real,Num}}(:ω => ω,:ζ => ζ,:k => k,:h => h)) - ω=para_dict[:ω] - ζ=para_dict[:ζ] - k=para_dict[:k] - h=para_dict[:h] + params = progress_scope(@parameters ω=ω ζ=ζ k=k h=h) sts = @variables x(t)=1.0 y(t)=1.0 jcn(t)=0.0 + ω, ζ, k, h = params eqs = [D(x) ~ y-(2*ω*ζ*x)+ k*(2/π)*(atan((jcn)/h)) D(y) ~ -(ω^2)*x] - odesys = ODESystem(eqs, t, sts, values(para_dict); name=name) - new(para_dict, odesys.x,[odesys.x],[odesys.x,odesys.y], + odesys = ODESystem(eqs, t, sts, params; name=name) + new(odesys.x,[odesys.x],[odesys.x,odesys.y], Dict(odesys.x => (-1.0,1.0), odesys.y => (-1.0,1.0)), odesys) end @@ -43,23 +39,19 @@ const harmonic_oscillator = HarmonicOscillatorBlox # return HarmonicOscillatorImage mutable struct JansenRitCBlox <: NeuralMassBlox - p_dict::Dict{Symbol,Union{Real,Num}} connector::Num noDetail::Vector{Num} detail::Vector{Num} initial::Dict{Num, Tuple{Float64, Float64}} odesystem::ODESystem function JansenRitCBlox(;name, τ=0.001, H=20.0, λ=5.0, r=0.15) - para_dict = scope_dict(Dict{Symbol,Union{Real,Num}}(:τ => τ,:H => H,:λ => λ,:r => r)) - τ=para_dict[:τ] - H=para_dict[:H] - λ=para_dict[:λ] - r=para_dict[:r] + params = progress_scope(@parameters τ=τ H=H λ=λ r=r) sts = @variables x(t)=1.0 y(t)=1.0 jcn(t)=0.0 + τ, H, λ, r = params eqs = [D(x) ~ y - ((2/τ)*x), D(y) ~ -x/(τ*τ) + (H/τ)*((2*λ)/(1 + exp(-r*(jcn))) - λ)] - odesys = ODESystem(eqs, t, sts, values(para_dict); name=name) - new(para_dict, odesys.x,[odesys.x],[odesys.x,odesys.y], + odesys = ODESystem(eqs, t, sts, params; name=name) + new(odesys.x,[odesys.x],[odesys.x,odesys.y], Dict(odesys.x => (-1.0,1.0), odesys.y => (-1.0,1.0)), odesys) end @@ -68,23 +60,19 @@ end const jansen_ritC = JansenRitCBlox mutable struct JansenRitSCBlox <: NeuralMassBlox - p_dict::Dict{Symbol,Union{Real,Num}} connector::Num noDetail::Vector{Num} detail::Vector{Num} initial::Dict{Num, Tuple{Float64, Float64}} odesystem::ODESystem function JansenRitSCBlox(;name, τ=0.014, H=20.0, λ=400.0, r=0.1) - para_dict = scope_dict(Dict{Symbol,Union{Real,Num}}(:τ => τ,:H => H,:λ => λ,:r => r)) - τ=para_dict[:τ] - H=para_dict[:H] - λ=para_dict[:λ] - r=para_dict[:r] + params = progress_scope(@parameters τ=τ H=H λ=λ r=r) sts = @variables x(t)=1.0 y(t)=1.0 jcn(t)=0.0 + τ, H, λ, r = params eqs = [D(x) ~ y - ((2/τ)*x), D(y) ~ -x/(τ*τ) + (H/τ)*((2*λ)/(1 + exp(-r*(jcn))) - λ)] - odesys = ODESystem(eqs, t, sts, values(para_dict); name=name) - new(para_dict, odesys.x,[odesys.x],[odesys.x,odesys.y], + odesys = ODESystem(eqs, t, sts, params; name=name) + new(odesys.x,[odesys.x],[odesys.x,odesys.y], Dict(odesys.x => (-1.0,1.0), odesys.y => (-1.0,1.0)), odesys) end diff --git a/src/datafitting/spectralDCM.jl b/src/datafitting/spectralDCM.jl index d327a271..02ae6efe 100644 --- a/src/datafitting/spectralDCM.jl +++ b/src/datafitting/spectralDCM.jl @@ -13,49 +13,98 @@ spm_logdet : mimick SPM12's way to compute the logarithm of the determinan variationalbayes : main routine that computes the variational Bayes estimate of model parameters """ +using ForwardDiff +using ForwardDiff: Dual +using ForwardDiff: Partials +using LinearAlgebra: Eigen using LinearAlgebra -using ToeplitzMatrices -using MAT +using RuntimeGeneratedFunctions +# using ToeplitzMatrices using ExponentialUtilities -function transferfunction_fmri(w, sts, derivatives, params) # relates to: spm_dcm_mtf.m +ForwardDiff.can_dual(::Type{Complex{Float64}}) = true +using ChainRules: _eigen_norm_phase_fwd! +tagtype(::Dual{T,V,N}) where {T,V,N} = T + + +function LinearAlgebra.eigen(M::Matrix{Dual{T, P, np}}) where {T, P, np} + nd = size(M, 1) + A = (p->p.value).(M) + F = eigen(A, sortby=nothing, permute=true) + λ, V = F.values, F.vectors + local ∂λ_agg, ∂V_agg + # compute eigenvalue and eigenvector derivatives for all partials + for i = 1:np + dA = (p->p.partials[i]).(M) + tmp = V \ dA + ∂K = tmp * V # V^-1 * dA * V + ∂Kdiag = @view ∂K[diagind(∂K)] + ∂λ_tmp = eltype(λ) <: Real ? real.(∂Kdiag) : copy(∂Kdiag) # why do only copy when complex?? + ∂K ./= transpose(λ) .- λ + fill!(∂Kdiag, 0) + ∂V_tmp = mul!(tmp, V, ∂K) + _eigen_norm_phase_fwd!(∂V_tmp, A, V) + if i == 1 + ∂V_agg = ∂V_tmp + ∂λ_agg = ∂λ_tmp + else + ∂V_agg = cat(∂V_agg, ∂V_tmp, dims=3) + ∂λ_agg = cat(∂λ_agg, ∂λ_tmp, dims=2) + end + end + ∂V = Array{Partials}(undef, nd, nd) + ∂λ = Array{Partials}(undef, nd) + # reassemble the aggregated vectors and values into a Partials type + for i = 1:nd + ∂λ[i] = Partials(Tuple(∂λ_agg[i, :])) + for j = 1:nd + ∂V[i, j] = Partials(Tuple(∂V_agg[i, j, :])) + end + end + if eltype(V) <: Complex + evals = map((x,y)->Complex(Dual{T, Float64, length(y)}(real(x), Partials(Tuple(real(y)))), + Dual{T, Float64, length(y)}(imag(x), Partials(Tuple(imag(y))))), F.values, ∂λ) + evecs = map((x,y)->Complex(Dual{T, Float64, length(y)}(real(x), Partials(Tuple(real(y)))), + Dual{T, Float64, length(y)}(imag(x), Partials(Tuple(imag(y))))), F.vectors, ∂V) + else + evals = Dual{T, Float64, length(∂λ[1])}.(F.values, ∂λ) + evecs = Dual{T, Float64, length(∂V[1])}.(F.vectors, ∂V) + end + return Eigen(evals, evecs) +end + - C = params[:C] +function transferfunction_fmri(w, idx_A, derivatives, params) + nd = Int(sqrt(length(idx_A))) + C = params[(6+2nd+nd^2):(5+3nd+nd^2)] C /= 16.0 # TODO: unclear why C is devided by 16 but see spm_fx_fmri.m:49 + ∂f = derivatives[:∂f](params[1:(nd^2+nd+1)]) #convert(Array{Real}, substitute(derivatives[:∂f], params)) + if ∂f isa Vector + ∂f = reshape(∂f, nd*5, nd*5) # TODO: generalize this to arbitrary number of states! This is specific for LinHemo! + end - # 2. get jacobian of hemodynamics - ∂f = substitute(derivatives[:∂f], params) - ∂f = convert(Array{Real}, substitute(∂f, sts)) - idx_A = findall(occursin.("A[", string.(derivatives[:∂f]))) - A = ∂f[idx_A] - nd = Int(sqrt(length(A))) - A_tmp = A[[(i-1)*nd+i for i=1:nd]] - A[[(i-1)*nd+i for i=1:nd]] -= exp.(A_tmp)/2 + A_tmp - ∂f[idx_A] = A - - dfdu = zeros(size(∂f, 1), length(C)) + dfdu = zeros(eltype(C), size(∂f, 1), length(C)) dfdu[CartesianIndex.([(idx[2][1], idx[1]) for idx in enumerate(idx_A[[(i-1)*nd+i for i=1:nd]])])] = C - F = eigen(Symbolics.value.(∂f), sortby=nothing, permute=true) + F = eigen(∂f) Λ = F.values V = F.vectors - ∂g = substitute(derivatives[:∂g], params) - ∂g = Symbolics.value.(substitute(∂g, sts)) + ∂g = derivatives[:∂g](params[end]) + # ∂g = Symbolics.value.(substitute(∂g, sts)) dgdv = ∂g*V - dvdu = pinv(V)*dfdu + dvdu = V\dfdu nw = size(w,1) # number of frequencies ng = size(∂g,1) # number of outputs nu = size(dfdu,2) # number of inputs nk = size(V,2) # number of modes - S = zeros(Complex, nw, ng, nu) - + S = zeros(Complex{real(eltype(dvdu))}, nw, ng, nu) for j = 1:nu for i = 1:ng for k = 1:nk # transfer functions (FFT of kernel) - Sk = (1im*2*pi*w .- Λ[k]).^-1 # TODO: clean up 1im*2*pi*freq instead of omega to be consistent with the usual nomenclature + Sk = (1im*2*pi*w .- Λ[k]).^-1 S[:,i,j] .+= dgdv[i,k]*dvdu[k,j]*Sk end end @@ -63,6 +112,7 @@ function transferfunction_fmri(w, sts, derivatives, params) # relates to: spm_ return S end + """ This function implements equation 2 of the spectral DCM paper, Friston et al. 2014 "A DCM for resting state fMRI". Note that nomenclature is taken from SPM12 code and it does not seem to coincide with the spectral DCM paper's nomenclature. @@ -71,21 +121,22 @@ end Gn in the code corresponds to Ge in the paper, i.e. the observation noise. In the code global and local components are defined, no such distinction is discussed in the paper. In fact the parameter γ, corresponding to local component is not present in the paper. """ -function csd_approx(w, sts, derivatives, param) +function csd_approx(w, idx_A, derivatives, param) # priors of spectral parameters # ln(α) and ln(β), region specific fluctuations: ln(γ) nw = length(w) - α = param[:lnα] - β = param[:lnβ] - γ = param[:lnγ] - nd = length(γ) + nd = Int(sqrt(length(idx_A))) + α = param[(2+nd+nd^2):(3+nd+nd^2)] + β = param[(4+nd+nd^2):(5+nd+nd^2)] + γ = param[(6+nd+nd^2):(5+2nd+nd^2)] + # define function that implements spectra given in equation (2) of the paper "A DCM for resting state fMRI". # neuronal fluctuations, intrinsic noise (Gu) (1/f or AR(1) form) - Gu = zeros(nw, nd, nd) - Gn = zeros(nw, nd, nd) - G = w.^(-exp(α[2])) # spectrum of hidden dynamics + G = w.^(-exp(α[2])) # spectrum of hidden dynamics G /= sum(G) + Gu = zeros(eltype(G), nw, nd, nd) + Gn = zeros(eltype(G), nw, nd, nd) for i = 1:nd Gu[:, i, i] .+= exp(α[1])*G end @@ -103,10 +154,10 @@ function csd_approx(w, sts, derivatives, param) Gn[:,j,i] = Gn[:,i,j] end end - S = transferfunction_fmri(w, sts, derivatives, param) # This is K(ω) in the equations of the spectral DCM paper. + S = transferfunction_fmri(w, idx_A, derivatives, param) # This is K(ω) in the equations of the spectral DCM paper. # predicted cross-spectral density - G = zeros(ComplexF64,nw,nd,nd); + G = zeros(eltype(S), nw, nd, nd); for i = 1:nw G[i,:,:] = S[i,:,:]*Gu[i,:,:]*S[i,:,:]' end @@ -114,8 +165,8 @@ function csd_approx(w, sts, derivatives, param) return G + Gn end -function csd_fmri_mtf(freqs, p, sts, derivatives, param) # alongside the above realtes to spm_csd_fmri_mtf.m - G = csd_approx(freqs, sts, derivatives, param) +@views function csd_fmri_mtf(freqs, p, idx_A, derivatives, param) # alongside the above realtes to spm_csd_fmri_mtf.m + G = csd_approx(freqs, idx_A, derivatives, param) dt = 1/(2*freqs[end]) # the following two steps are very opaque. They are taken from the SPM code but it is unclear what the purpose of this transformation and back-transformation is # in particular it is also unclear why the order of the MAR is reduced by 1. My best guess is that this procedure smoothens the results. @@ -123,20 +174,14 @@ function csd_fmri_mtf(freqs, p, sts, derivatives, param) # alongside the above # to make y well behaved. mar = csd2mar(G, freqs, dt, p-1) y = mar2csd(mar, freqs) + if real(eltype(y)) <: Dual + y_vals = Complex.((p->p.value).(real(y)), (p->p.value).(imag(y))) + y_part = (p->p.partials).(real(y)) + (p->p.partials).(imag(y))*im + y = map((x1, x2) -> Dual{tagtype(real(y)[1]), ComplexF64, length(x2)}(x1, Partials(Tuple(x2))), y_vals, y_part) + end return y end -function diff(U, dx, f, param::OrderedDict) - nJ = size(U, 2) - y0 = f(param) - J = zeros(ComplexF64, nJ, size(y0, 1), size(y0, 2), size(y0, 3)) - for i = 1:nJ - tmp_param = vecparam(param) .+ U[:, i]*dx - y1 = f(unvecparam(tmp_param, param)) - J[i,:,:,:] = (y1 .- y0)/dx - end - return J, y0 -end function matlab_norm(A, p) if p == 1 @@ -213,65 +258,64 @@ function unvecparam(vals, param::OrderedDict{Any,Any}) end -function variationalbayes(sts, y, derivatives, w, V, p, priors, niter) # relates to spm_nlsi_GN.m +@views function variationalbayes(idx_A, y, derivatives, w, V, p, priors, niter) # extract priors Πθ_pr = priors[:Σ][:Πθ_pr] Πλ_pr = priors[:Σ][:Πλ_pr] μλ_pr = priors[:Σ][:μλ_pr] Q = priors[:Σ][:Q] - μθ_pr = vecparam(priors[:μ]) # note: μθ_po is posterior and θμ is prior # prep stuff + μθ_pr = vecparam(priors[:μ]) # note: μθ_po is posterior and μθ_pr is prior + nd = size(y,2) np = size(V, 2) # number of parameters ny = length(y) # total number of response variables nq = 1 - nh = size(Q, 3) # number of precision components/hyper parameters - λ = 8*ones(nh) - ϵ_θ = zeros(np) # M.P - θμ # still need to figure out what M.P is for. It doesn't seem to be used further down the road in nlsi_GM, only at the very beginning when p is defined first. Then replace μθ_po with θμ above. + nh = size(Q,3) # number of precision components (this is the same as above, but may differ) + λ = 8 * ones(nh) + ϵ_θ = zeros(np) # M.P - μθ_pr # still need to figure out what M.P is for. It doesn't seem to be used further down the road in nlsi_GM, only at the very beginning when p is defined first. Then replace μθ with μθ_pr above. μθ_po = μθ_pr + V*ϵ_θ - - dx = exp(-8) + revert = false - f_prep = pars -> csd_fmri_mtf(w, p, sts, derivatives, pars) + f_prep = param -> csd_fmri_mtf(w, p, idx_A, derivatives, param) # state variable F = -Inf F0 = F + previous_F = F v = -4 # log ascent rate criterion = [false, false, false, false] state = vb_state(0, F, λ, zeros(np), μθ_po, inv(Πθ_pr)) - local ϵ_λ, iΣ, Σλ, Σθ, dFdθθ, dFdθ - dFdλ = zeros(ComplexF64, nh) - dFdλλ = zeros(Float64, nh, nh) + dfdp = zeros(ComplexF64, length(w)*nd^2, np) + local ϵ_λ, iΣ, Σλ, Σθ, dFdpp, dFdp for k = 1:niter state.iter = k + dfdp = ForwardDiff.jacobian(f_prep, μθ_po) * V + norm_dfdp = matlab_norm(dfdp, Inf); + revert = isnan(norm_dfdp) || norm_dfdp > exp(32); - dfdθ, f = diff(V, dx, f_prep, unvecparam(μθ_po, priors[:μ])); - dfdθ = transpose(reshape(dfdθ, np, ny)) - norm_dfdθ = matlab_norm(dfdθ, Inf); # NB that the norm in Julia is different from MATLAB. For consistency with SPM12 we reimplemented it here - revert = isnan(norm_dfdθ) || norm_dfdθ > exp(32); if revert && k > 1 for i = 1:4 # reset expansion point and increase regularization - v = min(v - 2, -4); - t = exp(v - logdet(dFdθθ)/np) + v = min(v - 2,-4); + t = exp(v - logdet(dFdpp)/np) # E-Step: update if t > exp(16) - ϵ_θ = state.ϵ_θ - dFdθθ\dFdθ # -inv(dfdx)*f + ϵ_θ = state.ϵ_θ - dFdpp \ dFdp # -inv(dfdx)*f else - idFdθθ = inv(dFdθθ) - ϵ_θ = state.ϵ_θ + expv(t, dFdθθ, idFdθθ*dFdθ) - idFdθθ*dFdθ # (expm(dfdx*t) - I)*inv(dfdx)*f + ϵ_θ = state.ϵ_θ + expv(t, dFdpp, dFdpp \ dFdp) - dFdpp \ dFdp # (expm(dfdx*t) - I)*inv(dfdx)*f end μθ_po = μθ_pr + V*ϵ_θ - dfdθ, f = diff(V, dx, f_prep, unvecparam(μθ_po, priors[:μ])); - dfdθ = transpose(reshape(dfdθ, np, ny)) + # J_test = JacVec(f_prep, μθ_po) + # dfdp = stack(J_test*v for v in eachcol(V)) + dfdp = ForwardDiff.jacobian(f_prep, μθ_po) * V # check for stability - norm_dfdθ = matlab_norm(dfdθ, Inf); - revert = isnan(norm_dfdθ) || norm_dfdθ > exp(32); + norm_dfdp = matlab_norm(dfdp, Inf); + revert = isnan(norm_dfdp) || norm_dfdp > exp(32); # break if ~revert @@ -280,34 +324,34 @@ function variationalbayes(sts, y, derivatives, w, V, p, priors, niter) # rela end end - - ϵ = reshape(y - f, ny) # error value - J = - dfdθ # Jacobian, unclear why we have a minus sign. Helmut: comes from deriving a Gaussian. - + f = f_prep(μθ_po) + ϵ = reshape(y - f, ny) # error value + J = - dfdp # Jacobian, unclear why we have a minus sign. Helmut: comes from deriving a Gaussian. ## M-step: Fisher scoring scheme to find h = max{F(p,h)} // comment from MATLAB code - for m = 1:8 # 8 seems arbitrary. This is probably because optimization falls often into a periodic orbit. ToDo: Issue #8 - iΣ = zeros(ComplexF64, ny, ny) + P = zeros(eltype(J), size(Q)) + PΣ = zeros(eltype(J), size(Q)) + JPJ = zeros(real(eltype(J)), size(J,2), size(J,2), size(Q,3)) + dFdλ = zeros(eltype(J), nh) + dFdλλ = zeros(real(eltype(J)), nh, nh) + for m = 1:8 # 8 seems arbitrary. Numbers of iterations taken from SPM12 code. + iΣ = zeros(eltype(J), ny, ny) for i = 1:nh iΣ .+= Q[:,:,i]*exp(λ[i]) end - Σ = inv(iΣ) # Julia requires conversion to dense matrix before inversion so just use dense to begin with + Pp = real(J' * iΣ * J) # in MATLAB code 'real()' is applied to the resulting matrix product, why? Σθ = inv(Pp + Πθ_pr) - P = similar(Q) - PΣ = similar(Q) - JPJ = zeros(size(Pp,1), size(Pp,2), size(Q,3)) for i = 1:nh P[:,:,i] = Q[:,:,i]*exp(λ[i]) - PΣ[:,:,i] = P[:,:,i] * Σ + PΣ[:,:,i] = iΣ \ P[:,:,i] JPJ[:,:,i] = real(J'*P[:,:,i]*J) # in MATLAB code 'real()' is applied (see also some lines above), what's the rational? end - for i = 1:nh - dFdλ[i] = (tr(PΣ[:,:,i])*nq - real(dot(ϵ,P[:,:,i],ϵ)) - tr(Σθ * JPJ[:,:,i]))/2 + dFdλ[i] = (tr(PΣ[:,:,i])*nq - real(dot(ϵ, P[:,:,i], ϵ)) - tr(Σθ * JPJ[:,:,i]))/2 for j = i:nh - dFdλλ[i, j] = - real(tr(PΣ[:,:,i] * PΣ[:,:,j]))*nq/2 # eps = randn(sizen), (eps' * Ai) * (Aj * eps) + dFdλλ[i, j] = -real(tr(PΣ[:,:,i] * PΣ[:,:,j]))*nq/2 dFdλλ[j, i] = dFdλλ[i, j] end end @@ -320,10 +364,10 @@ function variationalbayes(sts, y, derivatives, w, V, p, priors, niter) # rela t = exp(4 - spm_logdet(dFdλλ)/length(λ)) # E-Step: update if t > exp(16) - dλ = -real(inv(dFdλλ) * dFdλ) + dλ = -real(dFdλλ \ dFdλ) else idFdλλ = inv(dFdλλ) - dλ = real(expv(t, dFdλλ, idFdλλ*dFdλ) - idFdλλ*dFdλ) # (expm(dfdx*t) - I)*inv(dfdx)*f + dλ = real(exponential!(t * dFdλλ) * idFdλλ*dFdλ - idFdλλ*dFdλ) # (expm(dfdx*t) - I)*inv(dfdx)*f ~~~ could also be done with expv but doesn't work with Dual. end dλ = [min(max(x, -1.0), 1.0) for x in dλ] # probably precaution for numerical instabilities? @@ -339,11 +383,11 @@ function variationalbayes(sts, y, derivatives, w, V, p, priors, niter) # rela end ## E-Step with Levenberg-Marquardt regularization // comment from MATLAB code - L = zeros(3) - L[1] = (real(logdet(iΣ))*nq - real(dot(ϵ, iΣ, ϵ)) - ny*log(2pi))/2 + L = zeros(real(eltype(iΣ)), 3) + L[1] = (real(logdet(iΣ))*nq - real(dot(ϵ, iΣ, ϵ)) - ny*log(2pi))/2 L[2] = (logdet(Πθ_pr * Σθ) - dot(ϵ_θ, Πθ_pr, ϵ_θ))/2 - L[3] = (logdet(Πλ_pr * Σλ) - dot(ϵ_λ, Πλ_pr, ϵ_λ))/2; - F = sum(L); + L[3] = (logdet(Πλ_pr * Σλ) - dot(ϵ_λ, Πλ_pr, ϵ_λ))/2 + F = sum(L) if k == 1 F0 = F @@ -351,35 +395,35 @@ function variationalbayes(sts, y, derivatives, w, V, p, priors, niter) # rela if F > state.F || k < 3 # accept current state - state.F = F state.ϵ_θ = ϵ_θ state.λ = λ state.Σθ = Σθ state.μθ_po = μθ_po + state.F = F # Conditional update of gradients and curvature - dFdθ = -real(J' * iΣ * ϵ) - Πθ_pr * ϵ_θ - dFdθθ = -real(J' * iΣ * J) - Πθ_pr + dFdp = -real(J' * iΣ * ϵ) - Πθ_pr * ϵ_θ # check sign + dFdpp = -real(J' * iΣ * J) - Πθ_pr # decrease regularization - v = min(v + 1/2,4); + v = min(v + 1/2, 4); else # reset expansion point ϵ_θ = state.ϵ_θ λ = state.λ # and increase regularization - v = min(v - 2,-4); + v = min(v - 2, -4); end # E-Step: update - t = exp(v - spm_logdet(dFdθθ)/np) + t = exp(v - spm_logdet(dFdpp)/np) if t > exp(16) - dθ = -inv(dFdθθ)*dFdθ # -inv(dfdx)*f + dθ = - inv(dFdpp) * dFdp # -inv(dfdx)*f else - dθ = exp(t * dFdθθ) * inv(dFdθθ)*dFdθ - inv(dFdθθ)*dFdθ # (expm(dfdx*t) - I)*inv(dfdx)*f + dθ = exponential!(t * dFdpp) * inv(dFdpp) * dFdp - inv(dFdpp) * dFdp # (expm(dfdx*t) - I)*inv(dfdx)*f end ϵ_θ += dθ μθ_po = μθ_pr + V*ϵ_θ - dF = dot(dFdθ, dθ); + dF = dot(dFdp, dθ); # convergence condition: reach a change in Free Energy that is smaller than 0.1 four consecutive times print("iteration: ", k, " - F:", state.F - F0, " - dF predicted:", dF, "\n") @@ -428,30 +472,19 @@ function spectralVI(data, neuraldynmodel, observationmodel, initcond, csdsetup, mar = mar_ml(y, p); # compute MAR from time series y and model order p y_csd = mar2csd(mar, freqs, dt^-1); # compute cross spectral densities from MAR parameters at specific frequencies freqs, dt^-1 is sampling rate of data - jac_f = calculate_jacobian(neuraldynmodel) - # the following line is relevant for shared parameters accross regions - # jac_f = substitute(jac_f, Dict([p for p in parameters(f) if occursin("κ", string(p))] .=> lnκ)) - - grad_g = calculate_jacobian(observationmodel)[2:end] - - # define values of states - all_s = states(neuraldynmodel) - - obs_states = states(observationmodel) - statesubs = merge.([Dict(obs_states[2] => s) for s in all_s if occursin(string(obs_states[2]), string(s))], - [Dict(obs_states[3] => s) for s in all_s if occursin(string(obs_states[3]), string(s))]) - - # gradient of g for all regions, note that the measurement model g is defined region independent, here it is extended to all regions - grad_g_full = Num.(zeros(nd, length(all_s))) - for (i, s) in enumerate(all_s) - dim = parse(Int64, string(s)[2]) - if occursin.(string(obs_states[2]), string(s)) - grad_g_full[dim, i] = substitute(grad_g[1], statesubs[dim]) - elseif occursin.(string(obs_states[3]), string(s)) - grad_g_full[dim, i] = substitute(grad_g[2], statesubs[dim]) + grad_full = function(p, grad, sts, nd) + tmp = zeros(typeof(p), nd, length(sts)) + for i in 1:nd + tmp[i, (i-1)*5 .+ (4:5)] = grad(vcat([1], sts[(i-1)*5 .+ (5:-1:4)]), p, t)[3:-1:2] end + return tmp end - derivatives = Dict(:∂f => jac_f, :∂g => grad_g_full) + jac_f = generate_jacobian(neuraldynmodel, expression = Val{false})[1] + grad_g = generate_jacobian(observationmodel, expression = Val{false})[1] + + statevals = [v for v in values(initcond)] + derivatives = Dict(:∂f => par -> jac_f(statevals, par, t), + :∂g => par -> grad_full(par, grad_g, statevals, nd)) θΣ = diagm(vecparam(OrderedDict(params.name .=> params.variance))) # depending on the definition of the priors (note that we take it from the SPM12 code), some dimensions are set to 0 and thus are not changed. @@ -469,14 +502,16 @@ function spectralVI(data, neuraldynmodel, observationmodel, initcond, csdsetup, ### Collect prior means and covariances ### Q = csd_Q(y_csd); # compute prior of Q, the precision (of the data) components. See Friston etal. 2007 Appendix A priors = Dict(:μ => OrderedDict(params.name .=> params.mean), - :Σ => Dict( - :Πθ_pr => inv(θΣ), # prior model parameter precision - :Πλ_pr => hyperparams[:Πλ_pr], # prior metaparameter precision - :μλ_pr => hyperparams[:μλ_pr], # prior metaparameter mean - :Q => Q # decomposition of model parameter covariance - ) - ); + :Σ => Dict( + :Πθ_pr => inv(θΣ), # prior model parameter precision + :Πλ_pr => hyperparams[:Πλ_pr], # prior metaparameter precision + :μλ_pr => hyperparams[:μλ_pr], # prior metaparameter mean + :Q => Q # decomposition of model parameter covariance + ) + ); + + idx_A = findall(occursin.("A[", string.(calculate_jacobian(neuraldynmodel)))) ### Compute the variational Bayes with Laplace approximation ### - return variationalbayes(initcond, y_csd, derivatives, freqs, V, p, priors, 128) + return variationalbayes(idx_A, y_csd, derivatives, freqs, V, p, priors, 128) end \ No newline at end of file diff --git a/src/measurementmodels/fmri.jl b/src/measurementmodels/fmri.jl index ccd5826d..fe6bdc2d 100644 --- a/src/measurementmodels/fmri.jl +++ b/src/measurementmodels/fmri.jl @@ -24,7 +24,6 @@ lnτ : logarithmic prefactor to transit time H[3], set to 0 for standard parame returns an ODESystem of the biophysical model for the hemodynamics """ mutable struct Hemodynamics <: NBComponent - p_dict::Dict{Symbol, Union{Real, Num}} connector::Num odesystem::ODESystem function Hemodynamics(;name, lnκ=0.0, lnτ=0.0) @@ -35,10 +34,12 @@ mutable struct Hemodynamics <: NBComponent H(4) - exponent for Fout(v) (alpha) H(5) - resting state oxygen extraction (E0) =# + H = [0.64, 0.32, 2.00, 0.32, 0.4] - para_dict = scope_dict(Dict{Symbol, Union{Real,Num}}(:lnκ => lnκ, :lnτ => lnτ)) - lnκ=para_dict[:lnκ] - lnτ=para_dict[:lnτ] + params = progress_scope(lnκ, lnτ) # progress scope if needed + params = compileparameterlist(lnκ=params[1], lnτ=params[2]) # finally compile all parameters + lnκ, lnτ = params # assign the modified parameters + states = @variables s(t) lnf(t) lnν(t) lnq(t) jcn(t) eqs = [ @@ -47,20 +48,20 @@ mutable struct Hemodynamics <: NBComponent D(lnν) ~ (exp(lnf) - exp(lnν)^(H[4]^-1)) / (H[3]*exp(lnτ)*exp(lnν)), D(lnq) ~ (exp(lnf)/exp(lnq)*((1 - (1 - H[5])^(exp(lnf)^-1))/H[5]) - exp(lnν)^(H[4]^-1 - 1))/(H[3]*exp(lnτ)) ] - odesys = ODESystem(eqs, t, states, values(para_dict); name=name) - new(para_dict, Num(0), odesys) + odesys = ODESystem(eqs, t, states, params; name=name) + new(Num(0), odesys) end end - mutable struct LinHemo <: NBComponent connector::Num bloxinput::Num odesystem::ODESystem function LinHemo(;name, lnκ=0.0, lnτ=0.0) - @variables jcn(t) + params = progress_scope(lnκ, lnτ) + @named hemo = Hemodynamics(;lnκ=params[1], lnτ=params[2]) @named nmm = LinearNeuralMassBlox() - @named hemo = Hemodynamics(;lnκ=lnκ, lnτ=lnτ) + @variables jcn(t) g = MetaDiGraph() add_vertex!(g, Dict(:blox => nmm, :jcn => jcn)) diff --git a/src/utilities/spectral_tools.jl b/src/utilities/spectral_tools.jl index a912a6a4..5ed4dc7b 100644 --- a/src/utilities/spectral_tools.jl +++ b/src/utilities/spectral_tools.jl @@ -201,7 +201,7 @@ function mar2csd(mar, freqs) nd = size(Σ, 1) w = 2pi*freqs/sf nf = length(w) - csd = zeros(ComplexF64, nf, nd, nd) + csd = zeros(eltype(Σ), nf, nd, nd) for i = 1:nf af_tmp = la.I for k = 1:p @@ -215,6 +215,19 @@ function mar2csd(mar, freqs) end +""" +Plain implementation of idft because AD dispatch versions for ifft don't work still! +""" +function idft(x::AbstractArray) + """discrete inverse fourier transform""" + N = size(x)[1] + out = Array{eltype(x)}(undef,N) + for n in 0:N-1 + out[n+1] = 1/N*sum([x[k+1]*exp(2*im*π*k*n/N) for k in 0:N-1]) + end + return out +end + """ This function converts a cross-spectral density (CSD) into a multivariate auto-regression (MAR) model. It first transforms the CSD into its cross-correlation function (Wiener-Kinchine theorem) and then computes the MAR model coefficients. @@ -239,15 +252,15 @@ function csd2mar(csd, w, dt, p) N = ceil(Int64, ns/2/dw) gj = findall(x -> x > 0 && x < (N + 1), w) gi = gj .+ (ceil(Int64, w[1]) - 1) # TODO: figure out what's the purpose of this! - g = zeros(ComplexF64, N) + g = zeros(eltype(csd), N) # transform to cross-correlation function - ccf = zeros(ComplexF64, N*2+1, size(csd,2), size(csd,3)) + ccf = zeros(eltype(csd), N*2+1, size(csd,2), size(csd,3)) for i = 1:size(csd, 2) for j = 1:size(csd, 3) g[gi] = csd[gj,i,j] - f = ifft(g) - f = ifft(vcat([0.0im; g; conj(g[end:-1:1])])) + f = idft(g) + f = idft(vcat([0.0im; g; conj(g[end:-1:1])])) ccf[:,i,j] = real.(fftshift(f))*N*dw end end @@ -258,8 +271,8 @@ function csd2mar(csd, w, dt, p) n = (N - 1) ÷ 2 p = min(p, n - 1) ccf = ccf[(1:n) .+ n,:,:] - A = zeros(m*p, m) - B = zeros(m*p, m*p) + A = zeros(eltype(csd), m*p, m) + B = zeros(eltype(csd), m*p, m*p) for i = 1:m for j = 1:m A[((i-1)*p+1):i*p, j] = ccf[(1:p) .+ 1, i, j] diff --git a/test/datafitting.jl b/test/datafitting.jl index 6eb0cba3..c6ae1cc9 100644 --- a/test/datafitting.jl +++ b/test/datafitting.jl @@ -9,38 +9,46 @@ nd = ncol(data) # number of dimensions ########## assemble the model ########## -@parameters lnκ=0.0 # define brain-wide decay parameter for hemodynamics +@parameters κ=0.0 # define brain-wide decay parameter for hemodynamics g = MetaDiGraph() for ii = 1:nd - region = LinHemo(;name=Symbol("r$ii"), lnκ=lnκ) + region = LinHemo(;name=Symbol("r$ii"), lnκ=κ) add_blox!(g, region) end # add symbolic weights @parameters A[1:length(vars["pE"]["A"])] = vec(vars["pE"]["A"]) for (i, idx) in enumerate(CartesianIndices(vars["pE"]["A"])) - add_edge!(g, idx[1], idx[2], :weight, A[i]) + if idx[1] == idx[2] + add_edge!(g, idx[1], idx[2], :weight, -exp(A[i])/2) # treatement of diagonal elements in SPM12 + else + add_edge!(g, idx[1], idx[2], :weight, A[i]) + end end # compose model @named neuronmodel = ODEfromGraph(g) neuronmodel = structural_simplify(neuronmodel) + # measurement model @named bold = boldsignal() # attribute initial conditions to states all_s = states(neuronmodel) -initcond = Dict{typeof(all_s[1]), eltype(x)}() -for i in 1:nd - for (j, s) in enumerate(all_s[occursin.("r$i", string.(all_s))]) +initcond = OrderedDict{typeof(all_s[1]), eltype(x)}() +rnames = [] +map(x->push!(rnames, split(string(x), "₊")[1]), all_s); +rnames = unique(rnames); +for (i, r) in enumerate(rnames) + for (j, s) in enumerate(all_s[r .== map(x -> x[1], split.(string.(all_s), "₊"))]) # TODO: fix this solution, it is not robust!! initcond[s] = x[i, j] end end modelparam = OrderedDict() for par in parameters(neuronmodel) - while Symbolics.getdefaultval(par) isa Num - par = Symbolics.getdefaultval(par) - end + # while Symbolics.getdefaultval(par) isa Num + # par = Symbolics.getdefaultval(par) + # end modelparam[par] = Symbolics.getdefaultval(par) end # Noise parameter mean diff --git a/test/runtests.jl b/test/runtests.jl index 80830f8d..ee17f82d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,7 +4,7 @@ using SafeTestsets @time @safetestset "Neurograph Tests" begin include("graphs.jl") end @time @safetestset "ODE from Graph Tests" begin include("ode_from_graph.jl") end # @time @safetestset "Neural Signal Measurement Models Tests" begin include("measurementmodels.jl") end -@time @safetestset "Spectral Utilities Tests" begin include("spectraltools.jl") end +# @time @safetestset "Spectral Utilities Tests" begin include("spectraltools.jl") end @time @safetestset "Data Fitting Tests" begin include("datafitting.jl") end @time @safetestset "ODE from Graph and simulate" begin include("graph_to_dataframe.jl") end @time @safetestset "Learning Tests" begin include("learning.jl") end