Physics informed neural operator ode

.github/workflows/Tests.yml

@@ -37,6 +37,7 @@ jobs:
           - "Forward"
           - "DGM"
           - "NNODE"
+          - "PINOODE"
           - "NeuralAdapter"
           - "IntegroDiff"
     uses: "SciML/.github/.github/workflows/tests.yml@v1"

Project.toml

@@ -27,6 +27,7 @@ MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
+NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
@@ -79,12 +80,13 @@ LogDensityProblems = "2"
 Lux = "1.1.0"
 LuxCUDA = "0.3.3"
 LuxCore = "1.0.1"
-LuxLib = "1.3.2"
+LuxLib = "1.3"
 MCMCChains = "6"
 MLDataDevices = "1.2.0"
 MethodOfLines = "0.11.6"
 ModelingToolkit = "9.46"
 MonteCarloMeasurements = "1.1"
+NeuralOperators = "0.5"
 Optimisers = "0.3.3"
 Optimization = "4"
 OptimizationOptimJL = "0.4"

docs/Project.toml

@@ -16,6 +16,7 @@ MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
 NeuralPDE = "315f7962-48a3-4962-8226-d0f33b1235f0"
+NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
@@ -44,6 +45,7 @@ MethodOfLines = "0.11"
 ModelingToolkit = "9.7"
 MonteCarloMeasurements = "1"
 NeuralPDE = "5"
+NeuralOperators = "0.5"
 Optimization = "4"
 OptimizationOptimJL = "0.4"
 OptimizationOptimisers = "0.3"
@@ -53,4 +55,4 @@ Plots = "1.36"
 QuasiMonteCarlo = "0.3.2"
 Random = "1"
 Roots = "2.0"
-SpecialFunctions = "2.1"
+SpecialFunctions = "2.1"

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,5 @@
+# Physics-Informed Neural Operator (PINO) for ODEs
+
+```@docs
+PINOODE
+```

docs/src/tutorials/pino_ode.md

@@ -0,0 +1,99 @@
+# Physics Informed Neural Operator for ODEs
+
+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 ODEs
+
+In this section, we will define a parametric ODE and then learn it with a PINO using [`PINOODE`](@ref). The PINO will be trained to learn the mapping from the parameters of the ODE to its solution.
+
+```@example pino
+using Test
+using OptimizationOptimisers
+using Lux
+using Statistics, Random
+using NeuralOperators
+using NeuralPDE
+
+# Define the parametric ODE equation
+equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3]
+tspan = (0.0, 1.0)
+u0 = 1.0
+prob = ODEProblem(equation, u0, tspan)
+
+# Set the number of parameters for the ODE
+number_of_parameter = 3
+# Define the DeepONet architecture for the PINO
+deeponet = NeuralOperators.DeepONet(
+    Chain(
+        Dense(number_of_parameter => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10)),
+    Chain(Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast),
+        Dense(10 => 10, Lux.tanh_fast)))
+
+# Define the bounds for the parameters
+bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)]
+number_of_parameter_samples = 50
+# Define the training strategy
+strategy = StochasticTraining(20)
+# Define the optimizer
+opt = OptimizationOptimisers.Adam(0.03)
+alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy)
+# Solve the ODE problem using the PINOODE algorithm
+sol = solve(prob, alg, verbose = false, maxiters = 4000)
+```
+
+Now let's compare the prediction from the learned operator with the ground truth solution which is obtained by analytic solution of the parametric ODE.
+
+```@example pino
+using Plots
+
+function get_trainset(bounds, tspan, number_of_parameters, dt)
+    p_ = [range(start = b[1], length = number_of_parameters, stop = b[2]) for b in bounds]
+    p = vcat([collect(reshape(p_i, 1, size(p_i, 1))) for p_i in p_]...)
+    t_ = collect(tspan[1]:dt:tspan[2])
+    t = collect(reshape(t_, 1, size(t_, 1), 1))
+    (p, t)
+end
+
+# Compute the ground truth solution for each parameter
+ground_solution = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) + p[3] * t
+function ground_solution_f(p, t)
+    reduce(hcat,
+        [[ground_solution(u0, p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)])
+end
+
+# generate the solution with new parameters for test the model
+(p, t) = get_trainset(bounds, tspan, 50, 0.025)
+# compute the ground truth solution
+ground_solution_ = ground_solution_f(p, t)
+# predict the solution with the PINO model
+predict = sol.interp((p, t))
+
+# calculate the errors between the ground truth solution and the predicted solution
+errors = ground_solution_ - predict
+# calculate the mean error and the standard deviation of the errors
+mean_error = mean(errors)
+# calculate the standard deviation of the errors
+std_error = std(errors)
+
+p, t = get_trainset(bounds, tspan, 100, 0.01)
+ground_solution_ = ground_solution_f(p, t)
+predict = sol.interp((p, t))
+
+errors = ground_solution_ - predict
+mean_error = mean(errors)
+std_error = std(errors)
+
+# Plot the predicted solution and the ground truth solution as a filled contour plot
+# predict, represents the predicted solution for each parameter value and time
+plot(predict, linetype = :contourf)
+plot!(ground_solution_, 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[:, i], label = "Predicted")
+# 'ground' is the ground truth solution
+plot!(ground_solution_[:, i], label = "Ground truth")
+```

src/NeuralPDE.jl

@@ -16,6 +16,7 @@ using IntervalSets: infimum, supremum
 using LinearAlgebra: Diagonal
 using Lux: Lux, Chain, Dense, SkipConnection, StatefulLuxLayer
 using Lux: FromFluxAdaptor, recursive_eltype
+using NeuralOperators: DeepONet
 using LuxCore: LuxCore, AbstractLuxLayer, AbstractLuxWrapperLayer
 using MLDataDevices: CPUDevice, get_device
 using Optimisers: Optimisers, Adam
@@ -79,7 +80,7 @@ include("adaptive_losses.jl")
 
 include("ode_solve.jl")
 include("dae_solve.jl")
-
+include("pino_ode_solve.jl")
 include("transform_inf_integral.jl")
 include("discretize.jl")
 
@@ -90,6 +91,7 @@ include("PDE_BPINN.jl")
 
 include("dgm.jl")
 
+export PINOODE
 export NNODE, NNDAE
 export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde
 export PhysicsInformedNN, discretize