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)

u: ComponentVector{Float64}(x = (layer_1 = (weight = [0.750964081208812; 0.3283996834700114; … ; 1.6464641561826752; 0.3205421046605763;;], bias = [-0.13471797140844882; -0.008152543143825719; … ; -0.6000764800893694; -0.12028593582201608;;]), layer_2 = (weight = [0.13187651620105306 -0.2427275740213474 … -0.28374685483143786 -1.0826927684854837; -0.16889055282033183 0.1455950863630056 … 0.05978114801071167 0.6087374870950415; … ; -0.7835451641676421 -1.3675741391105054 … 0.10662737075323676 -0.2578471625229923; 0.21340189986848757 0.818121473499046 … 0.4266102292226471 0.17164258464917537], bias = [0.17832793832700805; 0.19895060869684866; … ; 0.11595276063280631; 0.15760593420975927;;]), layer_3 = (weight = [1.189846315179575 -0.2368620841471681 … 0.6418700871874653 0.5307558137588445; 0.06594523902282935 -0.5684667626995322 … -0.14064065252496374 0.8010436302167909; … ; -0.46428574922393734 -1.2856308740941578 … 1.0968668650296203 -0.7754068327861993; -0.45874041582988756 -0.2654955572538791 … -0.4292608655290291 0.48133542886249014], bias = [-0.3997677439643047; 0.19657418080407998; … ; -0.6455024338307105; 0.06270866030945645;;]), layer_4 = (weight = [-1.330778788963445 -0.06196480094884854 … -1.6323998296114353 0.610110763518234], bias = [1.0057353322552407;;])), y = (layer_1 = (weight = [-1.587432603716084; -0.4082014375829665; … ; -2.4604706135123853; 1.8771758034711972;;], bias = [0.07521957884224177; 0.42217434315563873; … ; -0.5925782975914129; 0.21280258595546034;;]), layer_2 = (weight = [0.002412410684109055 0.9184007107625649 … 0.4540075503410871 -0.41050459447996696; -0.5598665859378051 -0.14553376165357376 … -0.111086519476513 -0.05731057552394625; … ; 0.4856843809469001 0.7509231946768434 … 0.8355070099019613 0.10208674846723824; -0.2750423744578333 0.5677560788167143 … 0.639936236999094 0.5386669950872259], bias = [-0.2740423075636897; -0.33766823655016115; … ; -0.04540740399625323; 0.47344973620969144;;]), layer_3 = (weight = [0.260349478364943 0.22051113212644946 … 0.34625508355646245 0.3898360240307486; 0.5909399107364334 -0.17058360279246124 … -0.5853616474106695 -0.45488646727528037; … ; -0.447794637423782 0.06627274778733167 … 0.7462739009509882 -0.05658554134114637; 0.8357405564708686 -0.6055371901820774 … 0.7809470346421488 0.6747473612555885], bias = [-0.15095359621467136; 0.025969604937815107; … ; 0.06047784008581245; 0.4436259058437015;;]), layer_4 = (weight = [-1.181239286538214 1.3784981615599226 … -0.715048732042173 -0.961520866101497], bias = [0.280371889999883;;])))

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