Multivariate stochastic differential equations with `Bridge`

This IJulia script gives a tour for my package Bridge with focus on multivariate stochastic differential equations. I use Makie.jl for the visualisations.

1.
Installation

First we need to install Bridge and a couple of other packages:

# using Pkg
Pkg.pkg"add Bridge#master StaticArrays Colors GeometryTypes Distributions"

33.5s

The installation of Makie is a bit tricky and is detailed in the README.md file of Makie.

2.
Setting the scene

The next few lines load the needed packages and some scripts.

using Bridge, StaticArrays, Makie, Bridge.Models, Colors, GeometryTypes
using Random, LinearAlgebra
popdisplay() # pop makie display
nothing

46.4s

include(joinpath(dirname(pathof(Bridge)), "../extra/makie.jl"));

1.3s

Some definitions.

Random.seed!(5)
sphere = Sphere(Point3f0(0,0,0), 1.0f0)
circle = Sphere(Point2f0(0,0), 1.0f0)
perspective = @SArray Float32[0.433 0.901 -0.0 1.952; -0.237 0.114 0.965 -20.43; 0.869 -0.418 0.263 -90.271; 0.0 0.0 0.0 1.0];

1.6s

3.
Time

Bridge mostly works with fix time grid methods. To get started, define a grid of time points tt say in the interval [0, 5] on which you want to simulate the process.

T = 5.00
n = 10001 # total length
dt = T/(n - 1)
tt = 0.0:dt:T
;

0.1s

4.
Space

Bridge interacts nicely with StaticArrays. We use SVector{3,Float64} for points in 3d space. In Bridge.Models the alias ℝ{3} == SVector{3,Float64} is defined. Because I often use MCMC methods and have to sample thousands of solutions, I try to make sure the functions are fast and have minimal overhead. Using SVectors helps alot.

ℝ{3}

0.6s

5.
3D Wiener process or Brownian motion

Bridge.jl is a statistical toolbox for diffusion processes and stochastic differential equations. The simplest diffusion process is a Brownian motion. The distribution and concept of a Brownian motion is represented by the object Wiener{T}() where T is the value type. As long as randn(T) is defined, Wiener{T}() can be sampled.

Wiener{Float64}()
Wiener{Complex{Float64}}()

0.7s

But now for 3d Brownian motion...

Wiener{ℝ{3}}()

0.6s

Use sample to exactly sample a 3d Wiener process on at the time points tt.

W = sample(tt, Wiener{ℝ{3}}())

1.0s

The function sample returns a SamplePath X. SamplePath is the time series object of Bridge.jl, basically a struct with a vector of time points X.tt and a vector of locations X.yy.
The script extra/makie.jl defines a recipe for plotting SamplePaths with Makie.

W = sample(tt[1:10:end], Wiener{ℝ{3}}())
scatter(W, markersize = 0.09, marker = circle, color = viridis(1+n÷10))

0.8s

lines(W, color = viridis(1+n÷10))

0.8s

W = sample(tt[1:10:end], Wiener())
scatter(W, markersize = 0.03, marker = circle, color = viridis(1+n÷10))

0.8s

lines(W, color = viridis(1+n÷10))

0.7s

# Figure 1: Brownian motion path

scene = lines(W.yy, linewidth = 0.5, color = viridis(length(W.yy)))
scatter!(scene, [W.yy[1]], markersize = 0.09, marker = circle, color = viridis(1)) # starting point

0.8s

Figure 1. Brownian motion in 3d. Colors indicate progress of time.

6.
Lorenz system of ordinary differential equations

Bridge.jl is mostly concerned with stochastic differential equations, but we can also solve ordinary differiential equations

\frac{d}{dt} x(t) = F(t, x(t)).

As a stochastic differential equation can be seen as ordinary differential equation with noise, let's start with an ordinary one and add noise in a second step.
The Lorenz system is famous and nice looking 3d system of ordinary differential equations.

F(t, x, s = 10.0, ρ = 28.0, β = 8/3) = ℝ{3}(s*(x[2] - x[1]), x[1]*(ρ - x[3]) - x[2], x[1]*x[2] - β*x[3])
x0 = ℝ{3}(1.508870, -1.531271, 25.46091)
;

0.2s

Note that  F(t,x)F(t, x)F(t,x) returns a 3d vector, we have written the Lorenz system as vector valued differential equation.
  sss,  ρ\rhoρ and  β\betaβ are parameters governing the system. With following parameters chosen by Lorenz the system shows chaotic behaviour:

s0 = 10
ρ0 = 28.0
β0 = 8/3
θ0 = ℝ{3}(s0, ρ0, β0)
;

0.2s

Compute a solution with solve. The argument BS3() tells solve to use an order 3 Bogacki–Shampine method.

U, err = solve(BS3(), tt, x0, F)
round(err, digits=5)

0.8s

F2(t,x,_) = F(t,x)
solve!(BS3(), F2, U, x0, nothing)

0.7s

# Figure 2: Solution of the Lorenz system

scene = lines( U.yy, linewidth = 0.8, color = viridis(n))
scatter!(scene, [U.yy[1]], markersize=0.4, marker = circle, color = viridis(1))
center!(scene)
#set_perspective!(scene, perspective)
rotate_cam!(scene, 0.4, 0.2, 0.)
scene

1.4s

Figure 2. A solution of the deterministic Lorenz system.

7.
Stochastic Lorenz system

A corresponding stochastic differential equation has the following form

\frac{d}{dt} X(t) = F(t, X(t)) + \sigma(t,X(t)) W(t).

For the example, we choose  σ=5I\sigma = 5Iσ=5I.

σ = (t,x)->5I
X = solve(EulerMaruyama(), x0, W, (F, σ))
;

0.2s

As the driving Brownian motion path provides a set of time points W.tt, the argument tt is dropped. solve has also an in-place version solve!.

@time solve!(EulerMaruyama(), X, x0, W, (F, σ));

1.5s

Note the solver is quite efficient.

# Figure 3: Sample path

scene = lines(X.yy, linewidth = 0.5, color = viridis(length(X.yy)))
scatter!(scene, [X.yy[1]], markersize=0.09, marker = circle, color = viridis(1))
center!(scene)
#set_perspective!(scene, perspective)
rotate_cam!(scene, 0.4, 0.2, 0.)
scene

0.8s

Figure 3. Sample of the solution of the stochastic Lorenz system.

8.
Parameter inference for the stochastic Lorenz system

The likelihood for the parameter  θ=(s,ρ,β)\theta = (s, \rho, \beta)θ=(s,ρ,β) is given by Girsanov's theorem. The stochastic Lorenz system is defined in Bridge.Model and takes a parameter triple θ.

function loglikelihood(θ, θref, X) 
    P = Lorenz(θ, SDiagonal(5.0, 5.0, 5.0))
    Pref = Lorenz(θref, SDiagonal(5.0, 5.0, 5.0)) # Reference measure
    girsanov(X, P, Pref)
end

0.8s

Choose a reference measure. We only estimate ρ and β.

S = 9.0:0.05:11.0
Ρ = 26:0.05:30
Β = 2.0:0.02:4.0
θref = s0, Ρ[end÷2], Β[end÷2] 
@show θref;

1.2s

llsurface = [loglikelihood((s0, ρ, β), θref, X) for ρ in Ρ, β in Β];

1.7s

# Figure 4: Likelihood surface

Scene(resolution = (200, 200))
llsurfaces = (llsurface .- mean(llsurface))/std(llsurface)
llsurfaces0 = llsurfaces[first(searchsorted(Ρ,ρ0)), first(searchsorted(Β,β0))]
scene = surface( Ρ, Β, llsurfaces, colormap = :Spectral)

l = Point3f0(ρ0, β0, 0.0)
u = Point3f0(ρ0, β0, 1.2*llsurfaces0)
lines!(Point3f0[l,(u+2l)/3, (2u+l)/3, u], linewidth=3.5, color=:darkred)

i,j = Tuple(argmax(llsurfaces))
scatter!([Point3f0(Ρ[i],Β[j], maximum(llsurfaces))], markersize=0.1, marker = circle, color = :white)
#axis3d!(scene, Ρ[1]:1.0:Ρ[end], Β[1]:1.0:Β[end], minimum(llsurfaces):1.0:maximum(llsurfaces))
center!(scene)
#set_perspective!(scene, Float32[-0.7788 0.6272 -0.0 20.1757; -0.3824 -0.4748 0.7926 13.1915; 0.4972 0.6173 0.6097 -23.9617; 0.0 0.0 0.0 1.0])

5.8s

Figure 4. (Log-) likelihood surface. A line marks the true parameter value, a circle the maximum likelihood estimate

Shift+Enter to run

9.
Markov chain Monte Carlo

In my work I am interested in Bayesian methods for inference for stochastic differential equations. To compute the posterior distribution of the parameters on naturally employes Markov Chain Monte Carlo (MCMC) methods.
Julia is a very good match for MCMC computations: They are sequential and cannot be vectorized. In programming languages with slow loops this is a problem and probabilistic programming libraries are used. For Julia, those too exists, but we may also just stay with Julia.

Shift+Enter to run

# MCMC sampler

logπ(s, ρ, β) = 1.0
function mcmc(X, logπ, θstart, θref; iterations = 10000)
    θ = θstart
    Θ = [θ]
    ll = -Inf
    lπ = logπ(θ...)
    for i in 1:iterations
        θᵒ = θ + 0.1*randn(ℝ{3})
        lπᵒ = logπ(θᵒ...)
        llᵒ = loglikelihood(θᵒ, θref, X)
        if rand() < exp(llᵒ - ll + lπᵒ - lπ)
            θ, lπ, ll = θᵒ, lπᵒ, llᵒ 
        end
        push!(Θ, θ)
    end
    Θ
end

# MCMC experiment

θref = S[end÷2], Ρ[end÷2], Β[end÷2] 
@time Θ = mcmc(X, logπ, ℝ{3}(9.,30.,2.0), θref; iterations = 10000)
@show mean(Θ)
@show θ0
;

1.1s

# Figure 5: Traceplot

scene = Scene(resolution = (900, 900))
Θs = [Point3f0(Θ[i]+0.01randn(ℝ{3})) for i in 1:1:1000] # subsample
scatter!(scene, Θs, markersize=0.02, marker = circle, color=RGBA(0.0, 0.0, 0.5, 0.3) )
lines!(scene, Θs, linewidth=0.5, color=RGBA(0.0, 0.0, 0.5, 0.3) )
#lines(Θ[end-10000:10:end], linewidth=0.2, color=:black)

for k in 1:3
    p = ℝ{3}(ntuple(n->n != k, 3))
    (lines!(scene, [Θs[i].*p .+ ℝ{3}(S[1],Ρ[1],Β[1]).*(1 .- p) for i in 1:length(Θs)], linewidth=0.4, color=RGB(0.6,0.7,0.8) ))
end

scatter!(scene, [ℝ{3}(s0, ρ0, β0)], markersize=0.1, marker = '+', color = :darkred)
Ps = [ℝ{3}(ntuple(n->n!=i,3)) for i in 1:3]
#axis3d!(8.0:1.0:12.0, Ρ[1]:1.0:Ρ[end], Β[1]:1.0:Β[end])

scatter!(scene, [ℝ{3}(s0, ρ0, β0).*p  .+ ℝ{3}(S[1],Ρ[1],Β[1]).*(1 .- p) for p in Ps], 
    markersize=0.08, marker = '+', color = RGB(0.8,0.5,0.5))



rotate_cam!(scene, -0.2, -0.2, 0.)
scene

4.0s

Figure 5. Samples of the MCMC chain for the posterior distribution (black) and true value (red). Projections on the  sss- ρ\rhoρ-plane, the  ρ\rhoρ- β\betaβ-plane and the  β\betaβ- sss-plane in gray, gray-red.

Shift+Enter to run

Multivariate stochastic differential equations with Bridge

1. Installation

2. Setting the scene

3. Time

4. Space

5. 3D Wiener process or Brownian motion

6. Lorenz system of ordinary differential equations

7. Stochastic Lorenz system

8. Parameter inference for the stochastic Lorenz system

9. Markov chain Monte Carlo

Multivariate stochastic differential equations with `Bridge`

1.
Installation

2.
Setting the scene

3.
Time

4.
Space

5.
3D Wiener process or Brownian motion

6.
Lorenz system of ordinary differential equations

7.
Stochastic Lorenz system

8.
Parameter inference for the stochastic Lorenz system

9.
Markov chain Monte Carlo