Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Bayes SDE notebook #446

Draft
wants to merge 7 commits into
base: main
Choose a base branch
from
Draft
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions _quarto.yml
Original file line number Diff line number Diff line change
@@ -80,6 +80,8 @@ website:
href: tutorials/09-variational-inference/index.qmd
- text: "Bayesian Differential Equations"
href: tutorials/10-bayesian-differential-equations/index.qmd
- text: "Bayesian Stochastic Differential Equations"
href: tutorials/10-bayesian-stochastic-differential-equations/index.qmd
- text: "Probabilistic PCA"
href: tutorials/11-probabilistic-pca/index.qmd
- tutorials/13-seasonal-time-series/index.qmd
2,809 changes: 2,809 additions & 0 deletions tutorials/10-bayesian-stochastic-differential-equations/Manifest.toml

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
[deps]
DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"
Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
179 changes: 179 additions & 0 deletions tutorials/10-bayesian-stochastic-differential-equations/index.qmd
Original file line number Diff line number Diff line change
@@ -0,0 +1,179 @@
---
title: Bayesian Estimation of Differential Equations
engine: julia
---

```{julia}
#| echo: false
#| output: false
using Pkg;
Pkg.instantiate();
```

::: {.callout-note collapse="true"}

## This Part is from [Bayesian Differential Equations Notebook](../10-bayesian-differential-equations/)

```{julia}
using Turing
using DifferentialEquations
using DifferentialEquations.EnsembleAnalysis
# Load StatsPlots for visualizations and diagnostics.
using StatsPlots
using LinearAlgebra
# Set a seed for reproducibility.
using Random
Random.seed!(14);
```

## The Lotka-Volterra Model

The Lotka–Volterra equations, also known as the predator–prey equations, are a pair of first-order nonlinear differential equations.
These differential equations are frequently used to describe the dynamics of biological systems in which two species interact, one as a predator and the other as prey.
The populations change through time according to the pair of equations

$$
\begin{aligned}
\frac{\mathrm{d}x}{\mathrm{d}t} &= (\alpha - \beta y(t))x(t), \\
\frac{\mathrm{d}y}{\mathrm{d}t} &= (\delta x(t) - \gamma)y(t)
\end{aligned}
$$

where $x(t)$ and $y(t)$ denote the populations of prey and predator at time $t$, respectively, and $\alpha, \beta, \gamma, \delta$ are positive parameters.

We implement the Lotka-Volterra model and simulate it with parameters $\alpha = 1.5$, $\beta = 1$, $\gamma = 3$, and $\delta = 1$ and initial conditions $x(0) = y(0) = 1$.

```{julia}
# Define Lotka-Volterra model.
function lotka_volterra(du, u, p, t)
# Model parameters.
α, β, γ, δ = p
# Current state.
x, y = u
# Evaluate differential equations.
du[1] = (α - β * y) * x # prey
du[2] = (δ * x - γ) * y # predator
return nothing
end
# Define initial-value problem.
u0 = [1.0, 1.0]
p = [1.5, 1.0, 3.0, 1.0]
tspan = (0.0, 10.0)
prob = ODEProblem(lotka_volterra, u0, tspan, p)
# Plot simulation.
plot(solve(prob, Tsit5()))
```

We generate noisy observations to use for the parameter estimation tasks in this tutorial.
With the [`saveat` argument](https://docs.sciml.ai/latest/basics/common_solver_opts/) we specify that the solution is stored only at `0.1` time units.
To make the example more realistic we add random normally distributed noise to the simulation.

```{julia}
sol = solve(prob, Tsit5(); saveat=0.1)
odedata = Array(sol) + 0.8 * randn(size(Array(sol)))
# Plot simulation and noisy observations.
plot(sol; alpha=0.3)
scatter!(sol.t, odedata'; color=[1 2], label="")
```

Alternatively, we can use real-world data from Hudson’s Bay Company records (an Stan implementation with slightly different priors can be found here: https://mc-stan.org/users/documentation/case-studies/lotka-volterra-predator-prey.html).

:::

## Inference of a Stochastic Differential Equation

A [Stochastic Differential Equation (SDE)](https://diffeq.sciml.ai/stable/tutorials/sde_example/) is a differential equation that has a stochastic (noise) term in the expression of the derivatives.
Here we fit a stochastic version of the Lokta-Volterra system.

We use a quasi-likelihood approach in which all trajectories of a solution are compared instead of a reduction such as mean, this increases the robustness of fitting and makes the likelihood more identifiable.
We use SOSRI to solve the equation.

```{julia}
u0 = [1.0, 1.0]
tspan = (0.0, 10.0)
function multiplicative_noise!(du, u, p, t)
x, y = u
du[1] = p[5] * x
return du[2] = p[6] * y
end
p = [1.5, 1.0, 3.0, 1.0, 0.1, 0.1]
function lotka_volterra!(du, u, p, t)
x, y = u
α, β, γ, δ = p
du[1] = dx = α * x - β * x * y
return du[2] = dy = δ * x * y - γ * y
end
prob_sde = SDEProblem(lotka_volterra!, multiplicative_noise!, u0, tspan, p)
ensembleprob = EnsembleProblem(prob_sde)
data = solve(ensembleprob, SOSRI(); saveat=0.1, trajectories=1000)
plot(EnsembleSummary(data))
# We generate new noisy observations based on the stochastic model for the parameter estimation tasks in this tutorial.
# We create our observations by adding random normally distributed noise to the mean of the ensemble simulation.
sdedata = reduce(hcat, timeseries_steps_mean(data).u) + 0.8 * randn(size(reduce(hcat, timeseries_steps_mean(data).u)))
```

```{julia}
@model function fitlv_sde(data, prob)
# Prior distributions.
σ ~ InverseGamma(2, 3)
α ~ truncated(Normal(1.3, 0.5); lower=0.5, upper=2.5)
β ~ truncated(Normal(1.2, 0.25); lower=0.5, upper=2)
γ ~ truncated(Normal(3.2, 0.25); lower=2.2, upper=4)
δ ~ truncated(Normal(1.2, 0.25); lower=0.5, upper=2)
ϕ1 ~ truncated(Normal(0.12, 0.3); lower=0.05, upper=0.25)
ϕ2 ~ truncated(Normal(0.12, 0.3); lower=0.05, upper=0.25)
# Simulate stochastic Lotka-Volterra model.
p = [α, β, γ, δ, ϕ1, ϕ2]
remake(prob, p = p)
ensembleprob = EnsembleProblem(prob)
predicted = solve(ensembleprob, SOSRI(); saveat=0.1, trajectories = 1000)
# Early exit if simulation could not be computed successfully.
for i in 1:length(predicted)
if !SciMLBase.successful_retcode(predicted[i])
Turing.@addlogprob! -Inf
return nothing
end
end
# Observations.
# We compute the likelihood for each trajectory of our simulation in order to better approximate the overall likelihood of our choice of parameters
for j in 1:length(predicted)
for i in 1:length(predicted[j])
data[:, i] ~ MvNormal(predicted[j][i], σ^2 * I)
end
end
return nothing
end;
```

The probabilistic nature of the SDE solution makes the likelihood function noisy which poses a challenge for NUTS since the gradient is changing with every calculation.
Therefore we use NUTS with a low target acceptance rate of `0.25` and specify a set of initial parameters.
SGHMC might be a more suitable algorithm to be used here.

```{julia}
model_sde = fitlv_sde(sdedata, prob_sde)
chain_sde = sample(
model_sde,
NUTS(0.25),
5000;
initial_params=[1.5, 1.3, 1.2, 2.7, 1.2, 0.12, 0.12],
progress=false,
)
plot(chain_sde)
```