Physics informed neural operator ode #806
Conversation
|
@ChrisRackauckas I need help with packages version. Adding dependency of NeuralOperator.jl to project, fail CI. I've tried a little to line up suitable versions but not success. |
|
What's left here before review? |
@ChrisRackauckas all is done here, but it is necessary for the tests to pass, but this is related to Lux updates and not to the code in this PR |
|
These two tests don't pass and that doesn't seem lux related? https://github.com/SciML/NeuralPDE.jl/actions/runs/11405946017/job/31738592363?pr=806#step:6:1096 It just needs to match https://github.com/SciML/NeuralPDE.jl/actions/runs/11395770870/job/31708526179 and make sure the new tests pass, but the new tests you added don't pass. |
| function physics_loss( | ||
| phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where {C <: DeepONet, T} | ||
| p, t = x | ||
| f = prob.f | ||
| out = phi(x, θ) | ||
| if size(p, 1) == 1 | ||
| f_vec = reduce(hcat, | ||
| [reduce(vcat, [f(out[j, i], p[1, i], t[j]) for j in axes(t, 2)]) | ||
| for i in axes(p, 2)]) | ||
| else | ||
| f_vec = reduce(hcat, | ||
| [reduce(vcat, [f(out[j, i], p[:, i], t[j]) for j in axes(t, 2)]) | ||
| for i in axes(p, 2)]) | ||
| end | ||
| du = dfdx(phi, x, θ) | ||
| norm = prod(size(du)) | ||
| sum(abs2, du .- f_vec) / norm | ||
| end |
There was a problem hiding this comment.
why can't this use the ODE code?
There was a problem hiding this comment.
Here it has parameters and coordinates on input, unlike the usual ODE with just cord, also ODE does not process DeepONet and NeuralOperator in general. It can't be reduced to instance ODE task. If only on the contrary, ODE is a special case of PinoODE as a parametric ODE with one parameter
| function initial_condition_loss( | ||
| phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { | ||
| C <: DeepONet, T} | ||
| p, t = x | ||
| t0 = reshape([prob.tspan[1]], (1, 1, 1)) | ||
| x0 = (p, t0) | ||
| u = phi(x0, θ) | ||
| u0 = size(prob.u0, 1) == 1 ? fill(prob.u0, size(u)) : | ||
| reduce(vcat, [fill(u0, size(u)) for u0 in prob.u0]) | ||
| norm = prod(size(u0)) | ||
| sum(abs2, u .- u0) / norm | ||
| end |
There was a problem hiding this comment.
why can't this use the ODE code?
|
Seems a bit odd to take over |
it wasn't identified |
|
it is reason why CI fail LuxDL/LuxLib.jl#179 |
test/PINO_ode_tests.jl
Outdated
| @test ground_solution≈predict_sol rtol=0.05 | ||
| p, t = get_trainset(chain, bounds, 100, tspan, 0.01) | ||
| ground_solution = ground_analytic.(u0, p, t) | ||
| predict_sol = sol(reduce(vcat, (p, t))) |
There was a problem hiding this comment.
it doesn't match the documented interface though, so it's really odd. sol.original(...) would be allowed to do this, but what I was saying is that we shouldn't have an interface break here.
There was a problem hiding this comment.
How do you think it would be better ?
There was a problem hiding this comment.
sol(t) should do the normal interoplation w.r.t. p in the ODEProblem. You can overload sol.interp to do extra things and document that as well, but sol(t) is described as having very specific behavior which should be kept the same with all other ODE solvers.
src/pino_ode_solve.jl
Outdated
| SciMLBase.allowscomplex(::PINOODE) = true | ||
|
|
||
| function (sol::SciMLBase.AbstractODESolution)(t::AbstractArray) | ||
| p, _ = sol.t |
There was a problem hiding this comment.
I don't understand how .t would have p?
There was a problem hiding this comment.
how can we place t and parameters in ODEsolution separately without override t is how all input data(t and p), or is it better identify PINOODEsolution here?
There was a problem hiding this comment.
I would've guessed you'd do:
function (sol::SciMLBase.AbstractODESolution)(t::Union{Number,AbstractArray})
sol.interp(sol.prob.p, t)
end? If that works then this PR is complete.
There was a problem hiding this comment.
ok, but with Chain it need that size(t) == size(p). I guess, It need add the warning "t should be same size as p"
or
function (sol::SciMLBase.AbstractODESolution)(t::Union{Number,AbstractArray})
p = gen(sol.prob.p, size(t)) # generate p same size as t
sol.interp(p, t)
endThere was a problem hiding this comment.
I've implemented it with the warning
There was a problem hiding this comment.
alright, I think we're good to go.
|
🎉 🎉 🎉 🎉 🎉 🎉 |
|
🥳🥳🥳🥳🥳 |
Implementation Physics-informed neural operator method for solve parametric Ordinary Differential Equations (ODE) use DeepOnet.
#575
Checklist
https://arxiv.org/abs/2103.10974
https://arxiv.org/abs/2111.03794