Getting started with ODEs

This tutorial provides a step-by-step guide to solve the Lotka-Volterra system of ordinary differential equations (ODEs).

using ModelingToolkit
using Sophon, IntervalSets
using Optimization, OptimizationOptimJL, Zygote
using Plots

# Defining parameters and variables
@parameters t
@variables x(..), y(..)

# Define the differential operator
Dₜ = Differential(t)
p = [1.5, 1.0, 3.0, 1.0]

# Setting up the system of equations
eqs = [Dₜ(x(t)) ~ p[1] * x(t) - p[2] * x(t) * y(t),
      Dₜ(y(t)) ~ -p[3] * y(t) + p[4] * x(t) * y(t)]

# Defining the domain
domain = [t ∈ 0 .. 3.0]

# Defining the initial conditions
bcs = [x(0.0) ~ 1.0, y(0.0) ~ 1.0]

@named lotka_volterra = PDESystem(eqs, bcs, domain, [t], [x(t), y(t)])

\[ \begin{align} \frac{\mathrm{d} x\left( t \right)}{\mathrm{d}t} =& 1.5 x\left( t \right) - x\left( t \right) y\left( t \right) \\ \frac{\mathrm{d} y\left( t \right)}{\mathrm{d}t} =& - 3 y\left( t \right) + x\left( t \right) y\left( t \right) \end{align} \]

In this part of the tutorial, we will employ BetaRandomSampler to generate training data using a Beta distribution. This introduces a soft causality into the training data, enhancing the effectiveness of the learning process.

# Constructing the physics-informed neural network (PINN)
pinn = PINN(x = FullyConnected(1, 1, sin; hidden_dims=8, num_layers=3),
            y = FullyConnected(1, 1, sin; hidden_dims=8, num_layers=3))

# Setting up the sampler, training strategy and problem
sampler = BetaRandomSampler(200, 1)
strategy = NonAdaptiveTraining(1,50)
prob = Sophon.discretize(lotka_volterra, pinn, sampler, strategy)

# Solving the problem using BFGS optimization
@showprogress res = Optimization.solve(prob, BFGS(); maxiters=1000)

retcode: Failure
u: ComponentVector{Float64}(x = (layer_1 = (weight = [0.8351408518995401; -0.9541922192290322; … ; 2.7144574385538967; 1.599641393751547;;], bias = [0.41714439496495564; 0.21532225289909612; … ; -0.03867321326296458; -0.03419951799035754;;]), layer_2 = (weight = [0.578570642934694 -0.33676409132599894 … -0.37686775565877173 0.9392820247653858; 0.2710260446990231 0.9599222831637794 … -0.182445523306519 0.9316089356576772; … ; -0.0007637531748919139 -0.4672226764870716 … -0.6888185736010156 0.27365517819465096; -0.1700038987635988 -0.1300852554251798 … -0.4482385935178413 -0.17431392596006412], bias = [0.06815933303637696; -0.19409577244962714; … ; -0.14223120354663263; -0.38133080630547617;;]), layer_3 = (weight = [-0.6155437165101663 0.12064704434277118 … 0.29174082196125367 0.7438672282363753; -0.02806767454893661 0.9597141856362779 … -0.6704204252287786 0.7732952758549674; … ; -0.4620189852370677 0.03541809166694957 … -0.4407777502998581 -0.18934529495045027; -0.29888800238692603 1.041838572941174 … -1.009149165411456 -0.29568304413255153], bias = [-0.00935774706427194; -0.5003511384731265; … ; 0.16901016053006393; -0.23199856014418704;;]), layer_4 = (weight = [0.4659352428001471 0.6210609184544491 … -1.619909739682488 0.7422718569144456], bias = [1.2875378709306984;;])), y = (layer_1 = (weight = [1.2650188513097138; -2.1299628027959185; … ; -1.6018386244910905; -0.40728247342389434;;], bias = [0.4882137648993693; -0.8516322036072205; … ; 0.020344683040969886; -0.36091063342950724;;]), layer_2 = (weight = [-0.47981757033316824 -0.5360182357807084 … -0.2723352016910647 -0.14270072412206494; -0.2776162651179807 -0.7344311041067417 … 0.5121312693881402 -0.10655393862145073; … ; -0.5881084799064208 0.10147546108258558 … 0.3513848535179934 0.2338398214087507; 0.21997988725486953 -0.14023948324786842 … 0.09678352349869036 0.5562322993397529], bias = [0.13957519379544067; 0.3410128644781268; … ; -0.3437079101775504; 0.2739446525431784;;]), layer_3 = (weight = [-0.20491648064869716 0.3039714591339555 … -0.06056753004361014 -0.522842650328342; 0.5599635658148778 0.3384346421911367 … -0.6114330221001418 -0.08778841821485256; … ; 0.5237695202692946 0.8616835430379399 … 0.10745905260710749 -0.22383508749916928; 0.1589481898967056 0.11716927839054965 … -0.06826239924316242 0.3700041709454553], bias = [0.08786519077465263; 0.001821474695719219; … ; -0.32537799697573627; 0.36607950491584784;;]), layer_4 = (weight = [0.3466050770141191 -0.7149183393831824 … -0.7046120123515138 0.4931731828626906], bias = [0.7034745017740908;;])))

Next, we'll compare our results with a reference solution to verify our computations.

using OrdinaryDiffEq

function f(u, p, t)
    return [p[1] * u[1] - p[2] * u[1] * u[2], -p[3] * u[2] + p[4] * u[1] * u[2]]
end

p = [1.5, 1.0, 3.0, 1.0]
u0 = [1.0, 1.0]
prob_oop = ODEProblem{false}(f, u0, (0.0, 3.0), p)
true_sol = solve(prob_oop, Tsit5(), saveat=0.01)

phi = pinn.phi
ts = [true_sol.t...;;]
x_pred = phi.x(ts, res.u.x)
y_pred = phi.y(ts, res.u.y)

plot(vec(ts), vec(x_pred), label="x_pred")
plot!(vec(ts), vec(y_pred), label="y_pred")
plot!(true_sol)

While the initial results are encouraging, we can further refine our model. By remaking the sampler, we can gradually transition to the uniform distribution for improved results.

# Adjusting the sampler to uniform distribution and re-solving the problem
for α in [0.6, 0.8, 1.0] # when α = 1.0, it is equivalent to uniform sampling
    sampler = remake(sampler; α=α)
    data = Sophon.sample(lotka_volterra, sampler)
    prob = remake(prob; p=data, u0=res.u)
    res = Optimization.solve(prob, BFGS(); maxiters=1000)
end

# Generating new predictions and calculating the absolute error
x_pred = phi.x(ts, res.u.x)
y_pred = phi.y(ts, res.u.y)
maximum(sum(abs2, vcat(x_pred, y_pred) .- stack(true_sol.u); dims=1)) # print the absolute error