I gave this a try in NeuralPDE.jl, but like all my other attempts, it didn’t work well at all - can you spot any issues?
using ModelingToolkit
using DomainSets
using NeuralPDE
using Flux
using Optimization
using OptimizationOptimisers
using OptimizationOptimJL
using Plots
@parameters t
@variables S(..) I(..) R(..)
It = Integral(t in DomainSets.ClosedInterval(0,t))
eqs = [S(t) ~ S(0.) - It(0.5*S(t)*I(t)),
I(t) ~ I(0.) + It(0.5*S(t)*I(t)) - It(0.25*I(t)),
R(t) ~ R(0.) + It(0.25*I(t))];
bcs = [S(0) ~ 0.99, I(0) ~ 0.01, R(0) ~ 0.0];
domains = [t ∈ (0.0, 40.0)];
@named pde_system = PDESystem(eqs, bcs, domains, [t], [S(t), I(t), R(t)]);
numhid = 4
chain = [Chain(Dense(1, numhid, Flux.σ),
Dense(numhid, numhid, Flux.σ),
Dense(numhid, 1, abs)) for i in 1:3] |> f64;
strategy_ = GridTraining(0.1);
discretization = PhysicsInformedNN(chain, strategy_);
prob = NeuralPDE.discretize(pde_system, discretization);
callback = function (p,l)
println("Current loss is: $l")
return false
end;
res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.1); callback = callback, maxiters=20);
prob = remake(prob, u0 = res.minimizer)
res = Optimization.solve(prob, OptimizationOptimJL.BFGS(); callback = callback, maxiters=20);
ts = [infimum(d.domain):0.1:supremum(d.domain) for d in domains][1]
phi = discretization.phi
np = length(res.u) ÷ 3
acum = [0;accumulate(+, [np,np,np])]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
Spred = hcat([phi[1]([t],res.u[sep[1]]) for t in ts]...)'
Ipred = hcat([phi[2]([t],res.u[sep[2]]) for t in ts]...)'
Rpred = hcat([phi[3]([t],res.u[sep[3]]) for t in ts]...)'
plot(ts, Spred, label="S")
plot!(ts, Ipred, label="I")
plot!(ts, Rpred, label="R")