docs: refactor/cleanup ode tutorials

docs/pages.jl

@@ -1,6 +1,6 @@
 pages = ["index.md",
     "ODE PINN Tutorials" => Any["Introduction to NeuralPDE for ODEs" => "tutorials/ode.md",
-                                "Bayesian PINNs for Coupled ODEs - Lotka-Volterra" => "examples/Lotka_Volterra_BPINNs.md"
+                                "Bayesian PINNs for Coupled ODEs" => "tutorials/Lotka_Volterra_BPINNs.md"
                                 #"examples/nnrode_example.md", # currently incorrect
                                 ],
     "PDE PINN Tutorials" => Any["Introduction to NeuralPDE for PDEs" => "tutorials/pdesystem.md",

docs/src/manual/ode.md

@@ -3,3 +3,9 @@
 ```@docs
 NNODE
 ```
+
+# Bayesian inference with PINNs
+
+```@docs
+BNNODE
+```

docs/src/tutorials/Lotka_Volterra_BPINNs.md

@@ -0,0 +1,114 @@
+# Bayesian Physics informed Neural Network ODEs Solvers
+
+Bayesian inference for PINNs provides an approach to ODE solution finding and parameter estimation with quantified uncertainty.
+
+## 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:
+
+```math
+\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 ideal parameters $\alpha = 1.5$, $\beta = 1$, $\gamma = 3$, and $\delta = 1$ and initial conditions $x(0) = y(0) = 1$.
+
+We then solve the equations and estimate the parameters of the model with priors for $\alpha$, $\beta$, $\gamma$ and $\delta$ as  `Normal(1,2)`, `Normal(2,2)`, `Normal(2,2)` and `Normal(0,2)` using a neural network.
+
+```@example bpinn
+using NeuralPDE, Lux, Plots, OrdinaryDiffEq, Distributions, Random
+
+function lotka_volterra(u, p, t)
+    # Model parameters.
+    α, β, γ, δ = p
+    # Current state.
+    x, y = u
+
+    # Evaluate differential equations.
+    dx = (α - β * y) * x # prey
+    dy = (δ * x - γ) * y # predator
+
+    return [dx, dy]
+end
+
+# initial-value problem.
+u0 = [1.0, 1.0]
+p = [1.5, 1.0, 3.0, 1.0]
+tspan = (0.0, 4.0)
+prob = ODEProblem(lotka_volterra, u0, tspan, p)
+```
+
+With the [`saveat`](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/) argument, we can specify that the solution is stored only at `saveat` time units.
+
+```@example bpinn
+# Solve using OrdinaryDiffEq.jl solver
+dt = 0.01
+solution = solve(prob, Tsit5(); saveat = dt)
+```
+
+We generate noisy observations to use for the parameter estimation task in this tutorial. To make the example more realistic we add random normally distributed noise to the simulation.
+
+```@example bpinn
+# Dataset creation for parameter estimation (30% noise)
+time = solution.t
+u = hcat(solution.u...)
+x = u[1, :] + (u[1, :]) .* (0.3 .* randn(length(u[1, :])))
+y = u[2, :] + (u[2, :]) .* (0.3 .* randn(length(u[2, :])))
+dataset = [x, y, time]
+
+# Plotting the data which will be used
+plot(time, x, label = "noisy x")
+plot!(time, y, label = "noisy y")
+plot!(solution, labels = ["x" "y"])
+```
+
+Lets define a PINN.
+
+```@example bpinn
+# Neural Networks must have 2 outputs as u -> [dx,dy] in function lotka_volterra()
+chain = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh),
+    Lux.Dense(6, 2))
+```
+
+The dataset we generated can be passed for doing parameter estimation using provided priors in `param` keyword argument for [`BNNODE`](@ref).
+
+```@example bpinn
+alg = BNNODE(chain;
+    dataset = dataset,
+    draw_samples = 1000,
+    l2std = [0.1, 0.1],
+    phystd = [0.1, 0.1],
+    priorsNNw = (0.0, 3.0),
+    param = [
+        Normal(1, 2),
+        Normal(2, 2),
+        Normal(2, 2),
+        Normal(0, 2)], progress = false)
+
+sol_pestim = solve(prob, alg; saveat = dt)
+
+nothing #hide
+```
+
+The solution for the ODE is retured as a nested vector `sol_flux_pestim.ensemblesol`. Here, [$x$ , $y$] would be returned.
+
+```@example bpinn
+# plotting solution for x,y for chain
+plot(time, sol_pestim.ensemblesol[1], label = "estimated x")
+plot!(time, sol_pestim.ensemblesol[2], label = "estimated y")
+
+# comparing it with the original solution
+plot!(solution, labels = ["true x" "true y"])
+```
+
+We can see the estimated ODE parameters by -
+
+```@example bpinn
+sol_pestim.estimated_ode_params
+```
+
+We can see it is close to the true values of the parameters.

docs/src/tutorials/ode.md

@@ -6,69 +6,70 @@
     equations with DifferentialEquations.jl tutorial](https://docs.sciml.ai/DiffEqDocs/stable/tutorials/ode_example/)
     before reading this tutorial.
 
-This tutorial is an introduction to using physics-informed neural networks (PINNs)
-for solving ordinary differential equations (ODEs). In contrast to the later
-parts of this documentation which use the symbolic interface, here we will focus on
-the simplified `NNODE` which uses the `ODEProblem` specification for the ODE.
+This tutorial is an introduction to using physics-informed neural networks (PINNs) for solving ordinary differential equations (ODEs). In contrast to the later parts of this documentation which use the symbolic interface, here we will focus on the simplified [`NNODE`](@ref) which uses the `ODEProblem` specification for the ODE.
+
 Mathematically, the `ODEProblem` defines a problem:
 
 ```math
 u' = f(u,p,t)
 ```
 
-for ``t \in (t_0,t_f)`` with an initial condition ``u(t_0) = u_0``. With physics-informed
-neural networks, we choose a neural network architecture `NN` to represent the solution `u`
-and seek parameters `p` such that `NN' = f(NN,p,t)` for all points in the domain.
-When this is satisfied sufficiently closely, then `NN` is thus a solution to the differential
-equation.
+for ``t \in (t_0,t_f)`` with an initial condition ``u(t_0) = u_0``. With physics-informed neural networks, we choose a neural network architecture `NN` to represent the solution `u` and seek parameters `p` such that `NN' = f(NN,p,t)` for all points in the domain. When this is satisfied sufficiently closely, then `NN` is thus a solution to the differential equation.
 
 ## Solving an ODE with NNODE
 
-Let's solve the simple ODE:
+Let's solve a simple ODE:
 
 ```math
 u' = \cos(2\pi t)
 ```
 
-for ``t \in (0,1)`` and ``u_0 = 0`` with `NNODE`. First, we define the `ODEProblem` as we would
-with any other DifferentialEquations.jl solver. This looks like:
+for ``t \in (0,1)`` and ``u_0 = 0`` with [`NNODE`](@ref). First, we define an `ODEProblem` as we would for defining an ODE using DifferentialEquations.jl interface. This looks like:
 
 ```@example nnode1
-using NeuralPDE, Flux, OptimizationOptimisers
+using NeuralPDE
 
-linear(u, p, t) = cos(2pi * t)
-tspan = (0.0f0, 1.0f0)
-u0 = 0.0f0
+linear(u, p, t) = cos(t * 2 * pi)
+tspan = (0.0, 1.0)
+u0 = 0.0
 prob = ODEProblem(linear, u0, tspan)
 ```
 
-Now, to define the `NNODE` solver, we must choose a neural network architecture. To do this, we
-will use the [Flux.jl library](https://fluxml.ai/) to define a multilayer perceptron (MLP)
-with one hidden layer of 5 nodes and a sigmoid activation function. This looks like:
+Now, to define the [`NNODE`](@ref) solver, we must choose a neural network architecture. To do this, we will use the [Lux.jl](https://lux.csail.mit.edu/) to define a multilayer perceptron (MLP) with one hidden layer of 5 nodes and a sigmoid activation function. This looks like:
 
 ```@example nnode1
-chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1))
+using Lux, Random
+
+rng = Random.default_rng()
+Random.seed!(rng, 0)
+chain = Chain(Dense(1, 5, σ), Dense(5, 1))
+ps, st = Lux.setup(rng, chain) |> Lux.f64
 ```
 
-Now we must choose an optimizer to define the `NNODE` solver. A common choice is `ADAM`, with
-a tunable rate , which we will set to `0.1`. In general, this rate parameter should be
-decreased if the solver's loss tends to be unsteady (sometimes rise “too much”), but should be
-as large as possible for efficiency. Thus, the definition of the `NNODE` solver is as follows:
+Now we must choose an optimizer to define the [`NNODE`](@ref) solver. A common choice is `Adam`, with a tunable rate, which we will set to `0.1`. In general, this rate parameter should be decreased if the solver's loss tends to be unsteady (sometimes rise “too much”), but should be as large as possible for efficiency. We use `Adam` from [OptimizationOptimisers](https://docs.sciml.ai/Optimization/stable/optimization_packages/optimisers/). Thus, the definition of the [`NNODE`](@ref) solver is as follows:
 
 ```@example nnode1
-opt = OptimizationOptimisers.Adam(0.1)
-alg = NeuralPDE.NNODE(chain, opt)
+using OptimizationOptimisers
+
+opt = Adam(0.1)
+alg = NNODE(chain, opt, init_params = ps)
 ```
 
-Once these pieces are together, we call `solve` just like with any other `ODEProblem` solver.
-Let's turn on `verbose` so we can see the loss over time during the training process:
+Once these pieces are together, we call `solve` just like with any other `ODEProblem`. Let's turn on `verbose` so we can see the loss over time during the training process:
 
 ```@example nnode1
-sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, abstol = 1.0f-6,
-            maxiters = 200)
+sol = solve(prob, alg, verbose = true, maxiters = 2000, saveat = 0.01)
+```
+
+Now lets compare the predictions from the learned network with the ground truth which we can obtain by numerically solving the ODE.
+
+```@example nnode1
+using OrdinaryDiffEq, Plots
+
+ground_truth = solve(prob, Tsit5(), saveat = 0.01)
+
+plot(ground_truth, label = "ground truth")
+plot!(sol.t, sol.u, label = "pred")
 ```
 
-And that's it: the neural network solution was computed by training the neural network and
-returned in the standard DifferentialEquations.jl `ODESolution` format. For more information
-on handling the solution, consult
-[the DifferentialEquations.jl solution handling section](https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/).
+And that's it: the neural network solution was computed by training the neural network and returned in the standard DifferentialEquations.jl `ODESolution` format. For more information on handling the solution, consult [here](https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/).