ODESystem usage in NeuralPDE/NNODE

New to Julia
1

Hi everyone,
I have to infer some parameters (~ 10^4, forget for now about the possiblity of doing this) of interactions strength for a ~ 90 variables network of ODEs whose interactions links I know beforehand. I wanted to do this through NeuralPDE, but before I wanted to be able to reproduce the existing examples. My constraint is that I have to be able to define the variables and parameters as vectors since first, they are many, and second they may change in the future.
I am having problem implementing a NNODE parameter inference for a system of ODEs through an ODEsystem instance by using vectors for parameters, variables and equations. The code is the following:

using ModelingToolkit, NeuralPDE, OrdinaryDiffEq, Lux, Random, OptimizationOptimJL, LineSearches, Plots

@independent_variables t
@parameters p[1:4]
@variables (u(t))[1:2]
D = Differential(t)

#Define equations
eqs = Vector{Equation}(undef,2)
eqs[1] = D(u[1]) ~ p[1] * u[1] - p[2] * u[2] * u[1];
eqs[2] = D(u[2]) ~ -p[3] * u[2] + p[4] * u[2] * u[1];

#Define time span, initial conditions and ODESystem
tspan = (0.0, 5.0)
u0 = [5.0, 5.0]
@named sys = ODESystem(eqs,t)

#Define parameters
true_p = [1.5, 1.0, 3.0, 1.0]

#Assign initial conditions and parameters
init_conds = Vector{Pair{Num, Int64}}(undef,2);
init_conds[1] = u[1] => u0[1];
init_conds[2] = u[2] => u0[2];
init_params = Vector{Pair{Num, Float64}}(undef,4);
init_params[1] = p[1] => true_p[1];
init_params[2] = p[2] => true_p[2];
init_params[3] = p[3] => true_p[3];
init_params[4] = p[4] => true_p[4];;

#Create ODEProblem
sys = structural_simplify(sys);
prob = ODEProblem(sys,init_conds,tspan,init_params)

#Create true data to be fitted to
prob_data = remake(prob, p = true_p)
sol_data = solve(prob_data, Tsit5(), saveat = 0.01)
t_ = sol_data.t
u_ = reduce(hcat, sol_data.u)

#Create NN
rng = Random.default_rng()
Random.seed!(rng, 0)
n = 15
chain = Chain(Dense(1, n, Ļƒ), Dense(n, n, Ļƒ), Dense(n, n, Ļƒ), Dense(n, 2))
ps, st = Lux.setup(rng, chain) |> f64

#Define loss function
additional_loss(phi, Īø) = sum(abs2, phi(t_, Īø) .- u_) / size(u_, 2)

#Define optimizer and alg function
opt = LBFGS(linesearch = BackTracking())
alg = NNODE(chain, opt, ps; strategy = WeightedIntervalTraining([0.7, 0.2, 0.1], 500),
    param_estim = true, additional_loss)

#Minimize
sol = solve(prob, alg, verbose = true, abstol = 1e-8, maxiters = 5000, saveat = t_)

Parameters, variables and equations are introduced as vectors because thatā€™s how I have to initialize them for the larger system in the future. The error I get is the following:

AssertionError: The NNODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).

I encountered this error also when trying to implement the NeuralPDE optimization throgh a function to pass to ODEProblem when this function modified the differential in place. In order to avoid this, I just initialize an empty array in the function and filled it with the corresponding gradients for the dynamics, but then NNODE complains that I am changing variables in place and as such cannot do gradient calculations.
How to avoid an in-place implementation of an ODEsystem? And isnā€™t there a direct, simple way to just initialize a system through indexing in a vector for parameters, variables and equations? For now I still havenā€™t found a convenient and relatively easy way of doing this, and I donā€™t see how someone could do otherwise for very large systems.

Thank you very much.

2

Definitely NeuralPDE and physics-informed neural networks are the wrong way to do this. SciMLSensitivity.jl and using direct inference will be a few orders of magnitude faster.

Do instead prob = ODEProblem{false}(sys,init_conds,tspan,init_params) to enforce the out-of-place function build.

3

Hi Chris,
Thanks a lot. I will have a look at SciLMSensitivity.
Is this because NeuralPDE scales badly with system size, or are there any more specific reasons for this?

4

Physics-informed neural networks scale pretty poorly and are not numerically robust to many behaviors such as stiffness, even with all of the little tweaks. Differentiable simulation techniques, like adjoint differentiation of ODEs, is just a few orders magnitude faster and more robust in any real scenario. And I donā€™t see that changing any time soon, itā€™s just a lot faster to use high order convergent schemes than to replace that with a simple neural network. The only things PINNs can really claim is that they are more parallelizable, so if the stock price of NVIDIA drops a bit then maybe you can grab enough GPUs to match the performance of the other technique.

There are still some other uses of PINNs, such as for building surrogates of solutions, but solving inverse problems is not one of them that is supported by any real benchmark (I know there are some papers that say it works, but they donā€™t have a baseline that actually tests against fast adjoint method implementations of other inverse problem solvers, and generally when I reproduce those their PINN is a few orders of magnitude slower than the most standard ā€œdifferentiate the ODEā€, soā€¦)

1 Like