# 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"

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
include(joinpath(dirname(pathof(Bridge)), "../extra/makie.jl"));

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];

## 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
;

## 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}

## 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}}()

But now for 3d Brownian motion...

Wiener{ℝ{3}}()

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

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

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))
lines(W, color = viridis(1+n÷10))
W = sample(tt[1:10:end], Wiener())
scatter(W, markersize = 0.03, marker = circle, color = viridis(1+n÷10))
lines(W, color = viridis(1+n÷10))
# 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

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

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)
;

Note that $F(t, x)$ returns a 3d vector, we have written the Lorenz system as vector valued differential equation.

$s$, $\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)
;

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)
F2(t,x,_) = F(t,x)
solve!(BS3(), F2, U, x0, nothing)
# 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

Figure 2. A solution of the deterministic Lorenz system.

## 7. Stochastic Lorenz system

A corresponding stochastic differential equation has the following form

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

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

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, σ));

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

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

## 8. Parameter inference for the stochastic Lorenz system

The likelihood for the parameter $\theta = (s, \rho, \beta)$ 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

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;
llsurface = [loglikelihood((s0, ρ, β), θref, X) for ρ in Ρ, β in Β];
# 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])

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

## 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.

# 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
;
# 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

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