Physics informed neural operator ode

.github/workflows/CI.yml

@@ -26,6 +26,7 @@ jobs:
           - Logging
           - Forward
           - DGM
+          - ODEPINO
           - NNODE
           - NeuralAdapter
           - IntegroDiff

Project.toml

@@ -97,4 +97,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "CUDA", "SafeTestsets", "OptimizationOptimJL", "Pkg", "OrdinaryDiffEq", "LineSearches", "LuxCUDA", "Flux", "MethodOfLines"]
+test = ["Aqua", "Test", "CUDA", "SafeTestsets", "OptimizationOptimJL", "Pkg", "OrdinaryDiffEq", "LineSearches", "LuxCUDA", "Flux", "MethodOfLines"]

docs/pages.jl

@@ -3,6 +3,7 @@ pages = ["index.md",
         "Bayesian PINNs for Coupled ODEs" => "tutorials/Lotka_Volterra_BPINNs.md",
         "PINNs DAEs" => "tutorials/dae.md",
         "Parameter Estimation with PINNs for ODEs" => "tutorials/ode_parameter_estimation.md",
+        "Physics informed Neural Operator ODEs" => "tutorials/pino_ode.md",
         "Deep Galerkin Method" => "tutorials/dgm.md"        #"examples/nnrode_example.md", # currently incorrect
     ],
     "PDE PINN Tutorials" => Any[
@@ -31,6 +32,7 @@ pages = ["index.md",
         "manual/training_strategies.md",
         "manual/adaptive_losses.md",
         "manual/logging.md",
-        "manual/neural_adapters.md"],
+        "manual/neural_adapters.md",
+        "manual/pino_ode.md"],
     "Developer Documentation" => Any["developer/debugging.md"]
 ]

docs/src/manual/pino_ode.md

@@ -0,0 +1,11 @@
+# Physics-Informed Neural operator for solve ODEs
+
+```@docs
+PINOODE
+```
+
+```@docs
+DeepONet
+```
+
+

docs/src/tutorials/pino_ode.md

@@ -0,0 +1,66 @@
+# Physics Informed Neural Operator for ODEs Solvers
+
+This tutorial provides an example of how to use the Physics Informed Neural Operator (PINO) for solving a family of parametric ordinary differential equations (ODEs).
+
+## Operator Learning  for a family of parametric ODE.
+
+In this section, we will define a parametric ODE and solve it using a PINO. The PINO will be trained to learn the mapping from the parameters of the ODE to its solution.
+
+```@example pino
+using Test
+using OrdinaryDiffEq, OptimizationOptimisers
+using Lux
+using Statistics, Random
+using NeuralPDE
+
+equation = (u, p, t) -> cos(p * t)
+tspan = (0.0f0, 1.0f0)
+u0 = 1.0f0
+prob = ODEProblem(equation, u0, tspan)
+
+# Define the architecture of the neural network that will be used as the PINO.
+branch = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10))
+trunk = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast))
+deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing)
+
+bounds = (p = [0.1f0, pi],)
+db = (bounds.p[2] - bounds.p[1]) / 50
+dt = (tspan[2] - tspan[1]) / 40
+strategy = NeuralPDE.GridTraining([db, dt])
+opt = OptimizationOptimisers.Adam(0.03)
+alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy)
+sol = solve(prob, alg, verbose = false, maxiters = 2000)
+predict = sol.u
+```
+
+Now let's compare the prediction from the learned operator with the ground truth solution which is obtained by analytic solution the parametric ODE. Where 
+Compare prediction with ground truth.
+
+```@example pino
+using Plots
+# Compute the ground truth solution for each parameter
+ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p)
+p_ = bounds.p[1]:strategy.dx[1]:bounds.p[2]
+p = reshape(p_, 1, size(p_)[1], 1)
+ground_solution = ground_analytic.(u0, p, sol.t.trunk)
+
+# Plot the predicted solution and the ground truth solution as a filled contour plot
+# sol.u[1, :, :], represents the predicted solution for each parameter value and time
+plot(predict[1, :, :], linetype = :contourf)
+plot!(ground_solution[1, :, :], linetype = :contourf)
+```
+
+```@example pino
+# 'i' is the index of the parameter 'p' in the dataset 
+i = 20
+# 'predict' is the predicted solution from the PINO model
+plot(predict[1, i, :], label = "Predicted")
+# 'ground' is the ground truth solution
+plot!(ground_solution[1, i, :], label = "Ground truth")
+```

src/NeuralPDE.jl

@@ -33,6 +33,7 @@ using UnPack: @unpack
 import ChainRulesCore, Lux, ComponentArrays
 using Lux: FromFluxAdaptor
 using ChainRulesCore: @non_differentiable
+#using NeuralOperators
 
 RuntimeGeneratedFunctions.init(@__MODULE__)
 
@@ -47,6 +48,8 @@ include("adaptive_losses.jl")
 include("ode_solve.jl")
 # include("rode_solve.jl")
 include("dae_solve.jl")
+include("neural_operators.jl")
+include("pino_ode_solve.jl")
 include("transform_inf_integral.jl")
 include("discretize.jl")
 include("neural_adapter.jl")
@@ -55,7 +58,8 @@ include("BPINN_ode.jl")
 include("PDE_BPINN.jl")
 include("dgm.jl")
 
-export NNODE, NNDAE,
+
+export NNODE, NNDAE, PINOODE, DeepONet, SomeStrategy #TODO remove SomeStrategy
        PhysicsInformedNN, discretize,
        GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining,
        WeightedIntervalTraining,

src/neural_operators.jl

@@ -0,0 +1,101 @@
+abstract type NeuralOperator <: Lux.AbstractExplicitLayer end
+
+"""
+DeepONet(branch,trunk)
+"""
+
+"""
+    DeepONet(branch,trunk,linear=nothing)
+
+`DeepONet` is differential neural operator focused for solving physic-informed parametric ODEs.
+
+DeepONet uses two neural networks, referred to as the "branch" and "trunk", to approximate
+the solution of a differential equation. The branch network takes the spatial variables as
+input and the trunk network takes the temporal variables as input. The final output is
+the dot product of the outputs of the branch and trunk networks.
+
+DeepONet is composed of two separate neural networks referred to as the "branch" and "trunk",
+respectively. The branch net takes on input represents a function  evaluated at a collection
+of fixed locations in some boundsand returns a features embedding. The trunk net takes the
+continuous coordinates as inputs, and outputs a features embedding. The final output of the
+DeepONet, the outputs of the branch and trunk networks are merged together via a dot product.
+
+## Positional Arguments
+*  `branch`: A branch neural network.
+*  `trunk`: A trunk neural network.
+
+## Keyword Arguments
+* `linear`: A linear layer to apply to the output of the branch and trunk networks.
+
+## Example
+
+```julia
+branch = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10))
+trunk = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast),
+    Lux.Dense(10, 10, Lux.tanh_fast))
+linear = Lux.Chain(Lux.Dense(10, 1))
+
+deeponet = DeepONet(branch, trunk; linear= linear)
+
+a = rand(1, 50, 40)
+b = rand(1, 1, 40)
+x = (branch = a, trunk = b)
+θ, st = Lux.setup(Random.default_rng(), deeponet)
+y, st = deeponet(x, θ, st)
+```
+
+## References
+* Lu Lu, Pengzhan Jin, George Em Karniadakis "DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators"
+* Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets"
+"""
+struct DeepONet{L <: Union{Nothing, Lux.AbstractExplicitLayer }} <: NeuralOperator
+    branch::Lux.AbstractExplicitLayer
+    trunk::Lux.AbstractExplicitLayer
+    linear::L
+end
+
+function DeepONet(branch, trunk; linear=nothing)
+    DeepONet(branch, trunk, linear)
+end
+
+function Lux.setup(rng::AbstractRNG, l::DeepONet)
+    branch, trunk, linear = l.branch, l.trunk, l.linear
+    θ_branch, st_branch = Lux.setup(rng, branch)
+    θ_trunk, st_trunk = Lux.setup(rng, trunk)
+    θ = (branch = θ_branch, trunk = θ_trunk)
+    st = (branch = st_branch, trunk = st_trunk)
+    if linear !== nothing
+        θ_liner, st_liner = Lux.setup(rng, linear)
+        θ = (θ..., liner = θ_liner)
+        st = (st..., liner = st_liner)
+    end
+    θ, st
+end
+
+Lux.initialstates(::AbstractRNG, ::DeepONet) = NamedTuple()
+
+@inline function (f::DeepONet)(x::NamedTuple, θ, st::NamedTuple)
+    x_branch, x_trunk = x.branch, x.trunk
+    branch, trunk = f.branch, f.trunk
+    st_branch, st_trunk = st.branch, st.trunk
+    θ_branch, θ_trunk = θ.branch, θ.trunk
+    out_b, st_b = branch(x_branch, θ_branch, st_branch)
+    out_t, st_t = trunk(x_trunk, θ_trunk, st_trunk)
+    if f.linear !== nothing
+        linear = f.linear
+        θ_liner, st_liner = θ.liner, st.liner
+        # out = sum(out_b .* out_t, dims = 1)
+        out_ = out_b .* out_t
+        out, st_liner = linear(out_, θ_liner, st_liner)
+        out = sum(out, dims = 1)
+        return out, (branch = st_b, trunk = st_t, liner = st_liner)
+    else
+        out = sum(out_b .* out_t, dims = 1)
+        return out, (branch = st_b, trunk = st_t)
+    end
+end