Bayesian inference for SDE models: a case study for an excitable stochastic-dynamical model

]up; add StaticArrays Distributions Bridge BridgeSDEInference#nextjournal

using Bridge

using StaticArrays

using BridgeSDEInference

using BridgeSDEInference: EulerMaruyamaBounded

using Random, LinearAlgebra, Distributions

const State = SArray{Tuple{2},T,1,2} where {T};

using Plots

gr(fmt=:png)

param = :regular # regular parametrization of the trajectories (no transformation of the coordinate processes)

x0 = State(-0.5, -0.6)

# time grid 

dt = 1/1000

T = 20.0

tt = 0.0:dt:T

# Exitatory behaviour (stable fixed point)

ε = 0.1 ; s =0.5 ; γ =1.5 ; β = 1.4 ; σ =0.3 ;

P = FitzhughDiffusion(param, ε, s, γ, β, σ);

Random.seed!(4)

X1, _ = simulate_segment(0.0, x0, P, tt);

skip = 10

p1_1 = plot(first.(X1.yy)[1:skip:end], last.(X1.yy)[1:skip:end], leg = false)

x1 = -1.2:0.01:1.2

x2 = [-1.2, 0.3]

x3 = [-0.9, 0.6]

plot!(x1, [x1 .- x1.^3 .+ s])

plot!(x2, γ.*x2 .+ β)

p1_2 = Plots.plot(X1.tt[1:skip:end],  first.(X1.yy)[1:skip:end], leg = false)

Plots.plot!(p1_2, X1.tt[1:skip:end],  last.(X1.yy)[1:skip:end])

p1 = plot(p1_1, p1_2, layout = (1,2))

# Oscillatory behaviour (unstable fixed point)

β = 0.6

P = FitzhughDiffusion(param, ε, s, γ, β, σ);

X2, _ = simulate_segment(0.0, x0, P, tt); 

p2_1 = plot(first.(X2.yy)[1:skip:end], last.(X2.yy)[1:skip:end], leg = false)

plot!(x1, [x1 .- x1.^3 .+ s])

plot!(x3, γ.*x3 .+ β)

p2_2 = Plots.plot(X2.tt[1:skip:end],  first.(X2.yy)[1:skip:end], leg = false)

Plots.plot!(p2_2, X2.tt[1:skip:end],  last.(X2.yy)[1:skip:end])

p2 = plot(p2_1, p2_2, layout = (1,2))

p3 = plot(p1, p2, layout = (2,1))

png(p3, "/results/out.png")

# starting point under :regular parametrisation

x0 = State(-0.5, -0.6)

ε = 0.1 ; s =-0.8 ; γ =1.5 ; β = 0.0 ; σ =0.3 ; 

P = FitzhughDiffusion(param, ε, s, γ, β, σ);

# time grid 

dt = 1/50000

T = 20.0

tt = 0.0:dt:T

Random.seed!(4)

X, _ = simulate_segment(0.0, x0, P, tt); 

# subsampling

num_obs = 100

skip = div(length(tt), num_obs)

Σ = [10^(-4)]

L = [1.0 0.0]

obs = (time = X.tt[1:skip:end],

  values = [Float64(rand(MvNormal(L*x, Σ))[1]) for x in X.yy[1:skip:end]]);

skip = 50

Plots.plot(X.tt[1:skip:end],  first.(X.yy)[1:skip:end], label = "X")

Plots.plot!(X.tt[1:skip:end],  last.(X.yy)[1:skip:end], label = "Y")

p = Plots.scatter!(obs.time, obs.values, markersize=1.5, label = "observations")

png(p, "/results/out.png")

# Take the real β, as it is fixed. 

θ_init = [ε, s, γ, β, σ].*(1 .+ (2*rand(5) .- 1).*[0.2, 0.2, 0.2, 0.0, 0.2]);

P_trgt = FitzhughDiffusion(param, θ_init...)

P_aux = [FitzhughDiffusionAux(param, θ_init..., t₀, u, T, v) for (t₀,T,u,v)

        in zip(obs.time[1:end-1], obs.time[2:end], obs.values[1:end-1], obs.values[2:end])]

# Container

model_setup = DiffusionSetup(P_trgt, P_aux, PartObs());

# Observation scheme

L = @SMatrix [1. 0.]     

Σ = @SMatrix [10^(-4)]      

set_observations!(model_setup, [L for _ in P_aux], [Σ for _ in P_aux],

                  obs.values, obs.time)

# Imputation grid

dt = 1/200

set_imputation_grid!(model_setup, dt)

# Prio distribution on (X_0, Y_0) 

set_x0_prior!(model_setup, GsnStartingPt(x0, @SMatrix [.1 0; 0 .1]), x0);

initialise!(eltype(x0), model_setup, Vern7(), false, NoChangePt(100))

# Further setting 

set_auxiliary!(model_setup; skip_for_save=2, adaptive_prop=NoAdaptation());

mcmc_setup = MCMCSetup(

      Imputation(NoBlocking(), 0.975, Vern7()),

      ParamUpdate(MetropolisHastingsUpdt(), 1, θ_init,

                  UniformRandomWalk(0.5, true), ImproperPosPrior(),

                  UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P_aux, 1))

                  ),

      ParamUpdate(MetropolisHastingsUpdt(), 2, θ_init,

                  UniformRandomWalk(0.5, false), ImproperPrior(),

                  UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P_aux, 2))

                  ),

      ParamUpdate(MetropolisHastingsUpdt(), 3, θ_init,

                  UniformRandomWalk(0.5, true), ImproperPosPrior(),

                  UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P_aux, 3))

                  ),

      ParamUpdate(MetropolisHastingsUpdt(), 5, θ_init,

                  UniformRandomWalk(0.5, true), ImproperPosPrior(),

                  UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P_aux, 5))

                  ))

schedule = MCMCSchedule(10^4, [[1, 2, 3, 4, 5]],

                        (save=10^3, verbose=10^3, warm_up=100,

                         readjust=(x->x%100==0), fuse=(x->false)));

Random.seed!(4)

GC.gc()

out = mcmc(mcmc_setup, schedule, model_setup);

θ1_names = ["eps" "s" "gamma" "sigma"]

chain2 = out[2].θ_chain[1:10:end]

p = plot(hcat( [i[1] for i in chain2 ], 

  		[i[2] for i in chain2], [i[3] for i in chain2], 

  		[i[5] for i in chain2]), label = θ1_names,layout = (4,1))

png(p, "/results/out.png")

ss = 5

p1 = scatter(obs.time, obs.values, color = :green, markersize=1.5)

Plots.plot!(p1, X.tt[1:ss:end],  first.(X.yy)[1:ss:end], alpha=0.5, lw = 1.8, color=:green, label = "Xtrue")

Plots.plot!(p1, X.tt[1:ss:end],  last.(X.yy)[1:ss:end], alpha=0.5, lw = 1.3, color=:green, label = "Ytrue")

for j in 1:length(out[1].paths)

		Plots.plot!(p1, out[1].time[1:ss:end],  [x[1] for x in out[1].paths[j]][1:ss:end], 

    	color = :blue, alpha = 0.1, leg = false)

  	Plots.plot!(p1, out[1].time[1:ss:end],  [x[2] for x in out[1].paths[j]][1:ss:end], 

    	color = :red, alpha = 0.1, leg = false )

end

png(p1, "/results/out.png")

# Save median and 90% credible sets of the empirical posterior distribution of the parameters as a dataframe

using CSV

using DataFrames

burnin = 10^4

θ_names = ("eps","s","gamma","beta","sigma")

θ = [ε, s, γ, β, σ]

df = DataFrame(par = String[], true_val = Float64[], median = Float64[], Q1 = Float64[], Q99 = Float64[])

for j in 1:5

  push!(df, [θ_names[j], θ[j], median([i[j] for i in out[2].θ_chain[burnin:end]]),

      quantile([i[j] for i in out[2].θ_chain[burnin:end]], 0.01),

      quantile([i[j] for i in out[2].θ_chain[burnin:end]], 0.99)])

end

CSV.write("results/posterior.csv", df);

df[!,2:5] = round.(df[!,2:5]; digits = 3)

df

using StatsPlots

p = []

for j in [1, 2,3,5]

  push!(p, StatsPlots.histogram([i[j] for i in out[2].θ_chain[burnin:end]], label = θ_names[j]))

end

png(plot(p[1], p[2], p[3], p[4], layout = (2,2)), "/results/out.png")

# Point estimates

using CSV

thetas = round.(CSV.read(
posterior.csv)[!, :median]; digits = 2)

# Needed for sliders as well

thetalims = [(0.0,1.0), (-5.0,5.0), (0.0,10.0), (0.0,10.0), (0.0,1.0)]

truethetas = [0.1, -0.8, 1.5, 0.0, 0.3];

#thetas = truethetas # should true parameters or point estimates marked in the app?

# Histogram plots to be used as slider background

using Plots

using Plots.PlotMeasures

pl = Vector{Any}(undef, 5)

for j in [1, 2, 3, 4, 5] 

  	pl[j] = Plots.scatter([thetas[j]], [2.5],

        showaxis=false, grid=false, legend=false, border=false, size=(230,100),

        markersize=2.5, color=:black, xlim=thetalims[j], ylim=(0,100), left_margin=-9mm, margin=-5mm, background=RGB(0.95,0.95,0.95), background_color_outside=:white)

  	if @isdefined out

    qs = quantile([i[j] for i in out[2].θ_chain[burnin:end]], 0.0:0.1:1.0)

        Plots.plot!(pl[j], Bridge.piecewise(SamplePath(qs, 0.1*(thetalims[j][2]-thetalims[j][1]) ./diff(qs)))...,

            linewidth=0, color=nothing, fillrange = 0, fillalpha = 0.5,

      	fillcolor = :red)

  	end

		savefig(pl[j], "results/slider$j.png");

end

] add Hyperscript Markdown JSServe WGLMakie#master AbstractPlotting Colors;

using JSServe, WGLMakie, AbstractPlotting

using JSServe: JSServe.DOM, @js_str, onjs

using Colors

using Random

using WGLMakie: scatter, scatter!

# fallback values if statistical analysis is not run

if !@isdefined thetas

   thetas = [0.1, -0.8, 1.5, 0.0, 0.3]

end

if !@isdefined thetalims

   thetalims = [(0.0,1.0), (-5.0,5.0), (0.0,10.0), (0.0,10.0), (0.0,1.0)]

end

WGLMakie.activate!()

using Hyperscript, Markdown

using JSServe, Observables

using JSServe: Application, Session, evaljs, linkjs, div, Asset

using JSServe: @js_str, onjs, Button, TextField, Slider, JSString, Dependency, with_session

using JSServe.DOM

# load histogram images to use as slider background 

sliderbg = [

  JSServe.Asset(
slider1.png), 

  JSServe.Asset(
slider2.png), 

  JSServe.Asset(
slider3.png), 

  JSServe.Asset(
slider4.png), 

  JSServe.Asset(
slider5.png)

]

styles = map(1:5) do i

  return css("#slider$(i)",

    paddingTop = "10px",

    height = "50px",

    backgroundSize = "115px 50px",

    backgroundRepeat = "no-repeat",

    backgroundPosition = "center center",

    backgroundImage = sliderbg[i])

end

function dom_handler(session, request)

  	# fetch initial parameter (initial slider settings)

    eps = thetas[1];

    s = thetas[2];

    gamma = thetas[3];

    beta = thetas[4];

    si = thetas[5];

    rebirth = 0.001; # how often a particle is "reborn" at random position

  	sl = 101 # slider sub-divisions

    # slider and field for sigma

    slider5 = JSServe.Slider(range(thetalims[5]..., length=sl), si)

    nrs5 = JSServe.NumberInput(si)

    linkjs(session, slider5.value, nrs5.value)

    # slider and field for beta

    slider4 = JSServe.Slider(range(thetalims[4]..., length=sl), beta)

    nrs4 = JSServe.NumberInput(beta)

    linkjs(session, slider4.value, nrs4.value)

    # slider and field for gamma

    slider3 = JSServe.Slider(range(thetalims[3]..., length=sl), gamma)

    nrs3 = JSServe.NumberInput(gamma)

    linkjs(session, slider3.value, nrs3.value)

    # slider and field for s

    slider2 = JSServe.Slider(range(thetalims[2]..., length=sl), s)

    nrs2 = JSServe.NumberInput(s)

    linkjs(session, slider2.value, nrs2.value)

    # slider and field for eps

    slider1 = JSServe.Slider(range(thetalims[1]..., length=sl), eps)

    nrs1 = JSServe.NumberInput(eps)

    linkjs(session, slider1.value, nrs1.value)

    # slider and field for rebirth

    slider6 = JSServe.Slider(0.0:0.0001:0.005, rebirth)

    nrs6 = JSServe.NumberInput(rebirth)

    linkjs(session, slider6.value, nrs6.value)

    # init

    R = (1.5, 3.0) # plot area

    R1, R2 = R

    limits = FRect(-R[1], -R[2], 2R[1], 2R[2])

    n = 400 # no of particles

    K = 150 # display K past positions of particle fading out

    dt = 0.0005 # time step

    sqrtdt = sqrt(dt)

		particlecss = DOM.style(

    	css("input[type='number']", width = "135px"), 

 			css("input[type='range']", width="135px")

		)

    ms1 = 0.02 # markersize particles

    ms2 = 0.02 # markersize isokline

  	# plot particles, initially at random positions

    scene = WGLMakie.scatter(repeat(R1*(2rand(n) .- 1), outer=K), repeat(R2*(2rand(n) .- 1),outer=K), color = fill((:white,0f0), n*K),

        backgroundcolor = RGB{Float32}(0.04, 0.11, 0.22), markersize = ms1,

        glowwidth = 0.005, glowcolor = :white,

        resolution=(500,500), limits = limits,

        )

  	# style plot

    axis = scene[Axis]

    axis[:grid, :linewidth] =  (0.3, 0.3)

    axis[:grid, :linecolor] = (RGBA{Float32}(0.5, 0.7, 1.0, 0.3),RGBA{Float32}(0.5, 0.7, 1.0, 0.3))

    axis[:names][:textsize] = (0.0,0.0)

    axis[:ticks, :textcolor] = (RGBA{Float32}(0.5, 0.7, 1.0, 0.5),RGBA{Float32}(0.5, 0.7, 1.0, 0.5))

    splot = scene[end]

  	# plot isoklines

    WGLMakie.scatter!(scene, -R1:0.01:R1, (-R1:0.01:R1) .- (-R1:0.01:R1).^3 .+ s, color = RGBA{Float32}(0.5, 0.7, 1.0, 0.8), markersize=ms2)

    kplot1 = scene[end]

    WGLMakie.scatter!(scene, -R1:0.01:R1, gamma*(-R1:0.01:R1) .+ beta , color = RGBA{Float32}(0.5, 0.7, 1.0, 0.8), markersize=ms2)

    kplot2 = scene[end]

  	# set up threejs scene

    three, canvas = WGLMakie.three_display(session, scene)

    js_scene = WGLMakie.to_jsscene(three, scene)

    mesh = js_scene.getObjectByName(string(objectid(splot)))

    mesh1 = js_scene.getObjectByName(string(objectid(kplot1)))

    mesh2 = js_scene.getObjectByName(string(objectid(kplot2)))

    # init javascript

    evaljs(session,  js"""

        iter = 1; // iteration number

    		// fetch parameters

        eps = $(eps);

        s = $(s);

        gamma = $(gamma);

        beta = $(beta);

        si = $(si);

        R1 = $(R1);

        R2 = $(R2);

    	  rebirth = $(rebirth);

    		// update functions for isoklines

        updateklinebeta = function (value){

            beta = value;

            var mesh = $(mesh2);

            var positions = mesh.geometry.attributes.offset.array;

            for ( var i = 0, l = positions.length; i < l; i += 2 ) {

                        positions[i+1] = beta + positions[i]*gamma;

                }

            mesh.geometry.attributes.offset.needsUpdate = true;

            //mesh.geometry.attributes.color.needsUpdate = true;

        }

        updateklinegamma = function (value){

            gamma = value;

            var mesh = $(mesh2);

            var positions = mesh.geometry.attributes.offset.array;

            for ( var i = 0, l = positions.length; i < l; i += 2 ) {

                        positions[i+1] = beta + positions[i]*gamma;

                }

            mesh.geometry.attributes.offset.needsUpdate = true;

            //mesh.geometry.attributes.color.needsUpdate = true;

        }

        updateklines = function (value){

            s = value;

            var mesh = $(mesh1);

            var positions = mesh.geometry.attributes.offset.array;

            for ( var i = 0, l = positions.length; i < l; i += 2 ) {

                        positions[i+1] = positions[i] - positions[i]*positions[i]*positions[i] + s;

                }

            mesh.geometry.attributes.offset.needsUpdate = true;

            //mesh.geometry.attributes.color.needsUpdate = true;

        }

    		// move particles every x milliseconds

        setInterval(

            function (){

                function randn_bm() {

                    var u = 0, v = 0;

                    while(u === 0) u = Math.random(); //Converting [0,1) to (0,1)

                    while(v === 0) v = Math.random();

                    return Math.sqrt( -2.0 * Math.log( u ) ) * Math.cos( 2.0 * Math.PI * v );

                }

                var mu = 0.2;

                var mesh = $(mesh);

                var K = $(K);

                var n = $(n);

                var dt = $(dt);

                iter++

                var sqrtdt = $(sqrtdt);

                k = iter%K;

                var positions = mesh.geometry.attributes.offset.array;

                var color = mesh.geometry.attributes.color.array;

                for ( var i = 0; i < n; i++ ) {

                    inew = k*2*n + 2*i;

                    iold = ((K + k - 1)%K)*2*n + 2*i;

                    positions[inew] = positions[iold] + dt/eps*((1 - positions[iold]*positions[iold])*positions[iold] - positions[iold+1] + s); // x

                    positions[inew+1] = positions[iold+1] + dt*(-positions[iold+1] + gamma*positions[iold] + beta) + si*sqrtdt*randn_bm();

                    color[k*4*n + 4*i] = 1.0;

                    color[k*4*n + 4*i + 1] = 1.0;

                    color[k*4*n + 4*i + 2] = 1.0;

                    color[k*4*n + 4*i + 3] = 1.0;

                    if (Math.random() < rebirth)

                    {

                        positions[inew] = (2*Math.random()-1)*R1;

                        positions[inew+1] = (2*Math.random()-1)*R2;

                    }

                }

                for ( var k = 0; k < K; k++ ) {

                    for ( var i = 0; i < n; i++ ) {

                        color[k*4*n + 4*i + 3] = 0.98*color[k*4*n + 4*i + 3];

                    }

                }

                mesh.geometry.attributes.color.needsUpdate = true;

                mesh.geometry.attributes.offset.needsUpdate = true;

            }

        , 15); // milliseconds

    """)

  	# react on slider movements

    onjs(session, slider1.value, js"""function (value){

        eps = value;

    }""")

    onjs(session, slider2.value, js"""function (value){

        updateklines(value);

    }""")

      onjs(session, slider3.value, js"""function (value){

        updateklinegamma(value);

    }""")

    onjs(session, slider4.value, js"""function (value){

        updateklinebeta(value);

    }""")

    onjs(session, slider5.value, js"""function (value){

        si = value;

    }""")

  	onjs(session, slider6.value, js"""function (value){

        rebirth = value;

    }""")

  	# arrange canvas and sliders as html elements

    dom = DOM.div(DOM.style(styles...), particlecss, DOM.p(canvas), DOM.p("Parameters"), DOM.table(  

    DOM.tr(DOM.td("ε ($eps)"), DOM.td("s ($s)"), DOM.td("γ ($gamma)"), DOM.td("β ($beta)"), DOM.td("σ ($si)")), 

    DOM.tr(

        DOM.td(DOM.div(slider1, id="slider1"), DOM.div(nrs1)),

        DOM.td(DOM.div(slider2, id="slider2"), DOM.div(nrs2)),

        DOM.td(DOM.div(slider3, id="slider3"), DOM.div(nrs3)),

        DOM.td(DOM.div(slider4, id="slider4"), DOM.div(nrs4)),

        DOM.td(DOM.div(slider5, id="slider5"), DOM.div(nrs5))

       ),

    DOM.tr(

        DOM.td("rebirth", DOM.div(slider6, id="slider6"), DOM.div(nrs6)),

  	)))

    println("running...")

    dom

end

# attach handler to current session

JSServe.with_session() do session, request

    dom_handler(session, request)

end