From 9988925358676196d2f61a5f23e223ce701a4027 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 12 Feb 2024 21:02:15 +0400 Subject: [PATCH 001/153] pino_ode_draft --- src/NeuralPDE.jl | 3 +- src/pino_ode_solve.jl | 180 +++++++++++++++++++++++++++++++++++++++++ test/PINO_ode_tests.jl | 40 +++++++++ 3 files changed, 222 insertions(+), 1 deletion(-) create mode 100644 src/pino_ode_solve.jl create mode 100644 test/PINO_ode_tests.jl diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 0f217b07d4..7d6e703e3b 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -45,6 +45,7 @@ include("adaptive_losses.jl") include("ode_solve.jl") # include("rode_solve.jl") include("dae_solve.jl") +include("pino_ode_solve.jl") include("transform_inf_integral.jl") include("discretize.jl") include("neural_adapter.jl") @@ -52,7 +53,7 @@ include("advancedHMC_MCMC.jl") include("BPINN_ode.jl") include("PDE_BPINN.jl") -export NNODE, NNDAE, +export NNODE, NNDAE, PINOODE PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl new file mode 100644 index 0000000000..bd008de7ad --- /dev/null +++ b/src/pino_ode_solve.jl @@ -0,0 +1,180 @@ +struct PINOODE{C, O, T, P, B, K} <: DiffEqBase.AbstractODEAlgorithm + chain::C + opt::O + training_mapping::T + init_params::P + batch::B + kwargs::K +end + +function PINOODE(chain, + opt, + training_mapping, + init_params = nothing; + batch = nothing, + kwargs...) + !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) + PINOODE(chain, opt, training_mapping, init_params, batch, kwargs) +end + +""" + PINOPhi(chain::Lux.AbstractExplicitLayer, t, u0, st) +""" +mutable struct PINOPhi{C, T, U, S} + chain::C + t0::T + u0::U + st::S + function PINOPhi(chain::Lux.AbstractExplicitLayer, t::Number, u0, st) + new{typeof(chain), typeof(t), typeof(u0), typeof(st)}(chain, t, u0, st) + end +end + +function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, t, u0, init_params) + θ, st = Lux.setup(Random.default_rng(), chain) + if init_params === nothing + init_params = ComponentArrays.ComponentArray(θ) + else + init_params = ComponentArrays.ComponentArray(init_params) + end + PINOPhi(chain, t, u0, st), init_params +end + +function (f::PINOPhi{C, T, U})(t::AbstractMatrix, + θ) where {C <: Lux.AbstractExplicitLayer, T, U} + # Batch via data as row vectors + # y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ.depvar)), t'), + # θ.depvar, + # f.st) + y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), + θ, + f.st) + ChainRulesCore.@ignore_derivatives f.st = st + # f.u0 .+ (t' .- f.t0) .* y + y +end + +function inner_data_loss(phi, ts, θ, a, ground_u, i) + u_size = size(ground_u[i]) + a_arr = fill(a[i], u_size) + input_data = reduce(vcat, [ts, a_arr]) + y = phi(input_data, θ) + y - ground_u[i] +end +function data_loss(phi,tspan, θ, training_mapping) + a, ground_u = training_mapping + t0, t_end = tspan + size_1 = size(ground_u[1]) + size_2 = size(ground_u)[1] + ts = Float32.(reshape(collect(range(t0, stop = t_end, length = size_1[2])), size_1)) + loss = [inner_data_loss(phi, ts, θ, a, ground_u,i) + for i in 1:size_2] + # reduce(vcat, loss) + sum(abs2, reduce(vcat, loss)) / (size_1[2] * size_2) +end + +function inner_loss(phi, tspan, θ, training_mapping) + loss = data_loss(phi, tspan, θ, training_mapping ) #+ physics_loss() + loss +end + +function generate_loss(phi, tspan, training_mapping) #training_mapping::Tuple + function loss(θ, _) + # inner_loss(phi, f, tspan, θ, p,training_mapping ) + sum(abs2, inner_loss(phi, tspan, θ, training_mapping)) + end + return loss +end + +function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, + alg::PINOODE, + args...; + dt = nothing, + abstol = 1.0f-6, + reltol = 1.0f-3, + verbose = false, + saveat = nothing, + maxiters = nothing) + u0 = prob.u0 + tspan = prob.tspan + f = prob.f + p = prob.p + t0 = tspan[1] + # param_estim = alg.param_estim + + #hidden layer + chain = alg.chain + opt = alg.opt + + #train points generation + init_params = alg.init_params + + # mapping between functional space of some vararible 'a' of equation (for example initial + # codition {u(t0 x)} or parameter) and solution of equation u(t) + training_mapping = alg.training_mapping + + !(chain isa Lux.AbstractExplicitLayer) && + error("Only Lux.AbstractExplicitLayer neural networks are supported") + + phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) + + # init_params = ComponentArrays.ComponentArray(init_params) + # init_params = ComponentArrays.ComponentArray(;depvar = ComponentArrays.ComponentArray(init_params)) + @show phi(rand(2,1), init_params) + @show phi(rand(2, 10), init_params) + + isinplace(prob) && + throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) + + try + # TODO + phi(rand(2, 1), init_params) + catch err + if isa(err, DimensionMismatch) + throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) + else + throw(err) + end + end + + # batch = if alg.batch === nothing + # true + # else + # alg.batch + # end + + inner_f = generate_loss(phi, tspan, training_mapping) + + # Creates OptimizationFunction Object from total_loss + total_loss(θ, _) = inner_f(θ, phi) + + + # Optimization Algo for Training Strategies + opt_algo = Optimization.AutoZygote() + + # Creates OptimizationFunction Object from total_loss + optf = OptimizationFunction(total_loss, opt_algo) + + # iteration = 0 + # callback = function (p, l) + # iteration += 1 + # verbose && println("Current loss is: $l, Iteration: $iteration") + # l < abstol + # end + + callback = function (p, l) + println("Current loss is: $l") + return false + end + + # init_params = ComponentArrays.ComponentArray(θ) + + # Zygote.gradient(θ -> total_loss(θ, 1), init_params) + + optprob = OptimizationProblem(optf, init_params) + res = solve(optprob, opt; callback = callback, maxiters = maxiters) # alg.kwargs...) + + # PINOsolution(fullsolution) + # res + (total_loss,optprob,opt,callback, maxiters) +end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl new file mode 100644 index 0000000000..ebc4d0bf10 --- /dev/null +++ b/test/PINO_ode_tests.jl @@ -0,0 +1,40 @@ +using Test +using OrdinaryDiffEq,OptimizationOptimisers +using Flux, Lux +using Statistics, Random +using NeuralOperators +# using NeuralPDE + +@testset "Example 1" begin + linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + linear = (u, p, t) -> cos(p * t) + tspan = (0.0, 2.0) + u0 = 0.0 + a = Float32.([i for i in 0.1:0.05:pi]) + + u_output_ = [] + for i in 1:length(a) + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a[i]) + sol1 = solve(prob, Tsit5(); saveat = 0.02) + push!(u_output_, Float32.(sol1.u')) + end + #TODO u0 -> [u0,a]? + prob = ODEProblem(linear, u0, tspan) + chain = Lux.Chain(Lux.Dense(2, 5, Lux.σ), Lux.Dense(5, 1)) + # opt = OptimizationOptimisers.Adam(0.1) + opt = OptimizationOptimisers.Adam(0.03) + training_mapping = (a, u_output_) + alg = NeuralPDE.PINOODE(chain, opt, training_mapping) + + sol = solve(prob, + alg, verbose = true, + maxiters = 2000, abstol = 1.0f-7) + + total_loss + total_loss, optprob, opt, callback, maxiters = sol + phi(rand(2, 10), init_params) + + total_loss(init_params, 1) + + using Plots +end From 5fc72b9708acb66e65e3ec7e23bfb2f9784631de Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 13 Feb 2024 17:27:14 +0400 Subject: [PATCH 002/153] pino ode prot --- Project.toml | 1 + src/pino_ode_solve.jl | 150 +++++++++++++++++++++-------------------- test/PINO_ode_tests.jl | 54 ++++++++++----- 3 files changed, 115 insertions(+), 90 deletions(-) diff --git a/Project.toml b/Project.toml index a8e43f9b08..90790ff847 100644 --- a/Project.toml +++ b/Project.toml @@ -24,6 +24,7 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" +NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1" diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index bd008de7ad..a5f4fc6a8e 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,9 +1,9 @@ -struct PINOODE{C, O, T, P, B, K} <: DiffEqBase.AbstractODEAlgorithm +struct PINOODE{C, O, T, P, K} <: DiffEqBase.AbstractODEAlgorithm chain::C opt::O training_mapping::T init_params::P - batch::B + minibatch::Int kwargs::K end @@ -11,77 +11,98 @@ function PINOODE(chain, opt, training_mapping, init_params = nothing; - batch = nothing, + minibatch = 0, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, training_mapping, init_params, batch, kwargs) + PINOODE(chain, opt, training_mapping, init_params, minibatch, kwargs) end +# mutable struct Phi{C, S} +# f::C +# st::S +# function Phi(chain::Lux.AbstractExplicitLayer) +# st = Lux.initialstates(Random.default_rng(), chain) +# new{typeof(chain), typeof(st)}(chain, st) +# end +# end + +# function (f::Phi{<:Lux.AbstractExplicitLayer})(x::Number, θ) +# y, st = f.f(adapt(parameterless_type(ComponentArrays.getdata(θ)), [x]), θ, f.st) + +# ChainRulesCore.@ignore_derivatives f.st = st +# y +# end + +# function (f::Phi{<:Lux.AbstractExplicitLayer})(x::AbstractArray, θ) +# y, st = f.f(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) +# ChainRulesCore.@ignore_derivatives f.st = st +# y +# end + +# function (f::Phi{<:Optimisers.Restructure})(x, θ) +# f.f(θ)(adapt(parameterless_type(θ), x)) +# end + +# function (f::PINOPhi{C})(t, θ) where {C <: Optimisers.Restructure} +# f.f(θ)(t) +# end +# Zygote.gradient(θ -> sum(abs2, data_loss(phi, tspan, θ, training_mapping)), init_params) + +# function inner_loss(phi::PINOPhi{C}, tspan, θ, training_mapping::Tuple) where {C} +# loss = data_loss(phi, tspan, θ, training_mapping ) #+ physics_loss() +# loss +# end + """ - PINOPhi(chain::Lux.AbstractExplicitLayer, t, u0, st) + PINOPhi(chain::Lux.AbstractExplicitLayer, t, st) """ -mutable struct PINOPhi{C, T, U, S} +mutable struct PINOPhi{C, S} chain::C - t0::T - u0::U st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, t::Number, u0, st) - new{typeof(chain), typeof(t), typeof(u0), typeof(st)}(chain, t, u0, st) + function PINOPhi(chain::Lux.AbstractExplicitLayer, st) + new{typeof(chain), typeof(st)}(chain, st) end end -function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, t, u0, init_params) +function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) if init_params === nothing init_params = ComponentArrays.ComponentArray(θ) else init_params = ComponentArrays.ComponentArray(init_params) end - PINOPhi(chain, t, u0, st), init_params + PINOPhi(chain, st), init_params +end + +function (f::PINOPhi{C})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer} + y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), [t]), θ, f.st) + ChainRulesCore.@ignore_derivatives f.st = st + # f.u0 .+ (t .- f.t0) .* y + y end -function (f::PINOPhi{C, T, U})(t::AbstractMatrix, - θ) where {C <: Lux.AbstractExplicitLayer, T, U} +function (f::PINOPhi{C})(t::AbstractArray, θ) where {C <: Lux.AbstractExplicitLayer} # Batch via data as row vectors - # y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ.depvar)), t'), - # θ.depvar, - # f.st) - y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), - θ, - f.st) + y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st # f.u0 .+ (t' .- f.t0) .* y y end -function inner_data_loss(phi, ts, θ, a, ground_u, i) - u_size = size(ground_u[i]) - a_arr = fill(a[i], u_size) - input_data = reduce(vcat, [ts, a_arr]) - y = phi(input_data, θ) - y - ground_u[i] -end -function data_loss(phi,tspan, θ, training_mapping) - a, ground_u = training_mapping - t0, t_end = tspan - size_1 = size(ground_u[1]) - size_2 = size(ground_u)[1] - ts = Float32.(reshape(collect(range(t0, stop = t_end, length = size_1[2])), size_1)) - loss = [inner_data_loss(phi, ts, θ, a, ground_u,i) - for i in 1:size_2] - # reduce(vcat, loss) - sum(abs2, reduce(vcat, loss)) / (size_1[2] * size_2) +function inner_data_loss(phi::PINOPhi{C}, θ, in_, out_) where {C} + phi(in_, θ) - out_ end -function inner_loss(phi, tspan, θ, training_mapping) - loss = data_loss(phi, tspan, θ, training_mapping ) #+ physics_loss() - loss +function data_loss(phi::PINOPhi{C}, tspan, θ, training_mapping) where {C} + input_set, output_set = training_mapping + data_set_size = size(input_set)[1] * size(input_set[1])[2] + loss = reduce(vcat,[inner_data_loss(phi, θ, in_, out_) for (in_, out_) in zip(input_set, output_set)]) + loss / data_set_size end -function generate_loss(phi, tspan, training_mapping) #training_mapping::Tuple +function generate_loss(phi::PINOPhi{C}, tspan, training_mapping::Tuple) where {C} function loss(θ, _) - # inner_loss(phi, f, tspan, θ, p,training_mapping ) - sum(abs2, inner_loss(phi, tspan, θ, training_mapping)) + sum(abs2, data_loss(phi, tspan, θ, training_mapping)) end return loss end @@ -89,46 +110,38 @@ end function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, alg::PINOODE, args...; - dt = nothing, + # dt = nothing, abstol = 1.0f-6, reltol = 1.0f-3, verbose = false, saveat = nothing, maxiters = nothing) - u0 = prob.u0 + # u0 = prob.u0 ? TODO tspan = prob.tspan f = prob.f p = prob.p - t0 = tspan[1] # param_estim = alg.param_estim - #hidden layer chain = alg.chain opt = alg.opt - - #train points generation init_params = alg.init_params # mapping between functional space of some vararible 'a' of equation (for example initial - # codition {u(t0 x)} or parameter) and solution of equation u(t) + # condition {u(t0 x)} or parameter) join and solution of equation u(t) training_mapping = alg.training_mapping !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) + phi, init_params = generate_pino_phi_θ(chain, init_params) - # init_params = ComponentArrays.ComponentArray(init_params) - # init_params = ComponentArrays.ComponentArray(;depvar = ComponentArrays.ComponentArray(init_params)) - @show phi(rand(2,1), init_params) - @show phi(rand(2, 10), init_params) + init_params = ComponentArrays.ComponentArray(init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - # TODO - phi(rand(2, 1), init_params) + phi(first(training_mapping[1]), init_params) catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) @@ -148,33 +161,22 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, # Creates OptimizationFunction Object from total_loss total_loss(θ, _) = inner_f(θ, phi) - # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() # Creates OptimizationFunction Object from total_loss optf = OptimizationFunction(total_loss, opt_algo) - # iteration = 0 - # callback = function (p, l) - # iteration += 1 - # verbose && println("Current loss is: $l, Iteration: $iteration") - # l < abstol - # end - + iteration = 0 callback = function (p, l) - println("Current loss is: $l") - return false + iteration += 1 + verbose && println("Current loss is: $l, Iteration: $iteration") + l < abstol end - # init_params = ComponentArrays.ComponentArray(θ) - - # Zygote.gradient(θ -> total_loss(θ, 1), init_params) - optprob = OptimizationProblem(optf, init_params) - res = solve(optprob, opt; callback = callback, maxiters = maxiters) # alg.kwargs...) + res = solve(optprob, opt; callback, maxiters, alg.kwargs...) # PINOsolution(fullsolution) - # res - (total_loss,optprob,opt,callback, maxiters) + (res, phi) end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ebc4d0bf10..502592f6f9 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -3,38 +3,60 @@ using OrdinaryDiffEq,OptimizationOptimisers using Flux, Lux using Statistics, Random using NeuralOperators -# using NeuralPDE +using NeuralPDE @testset "Example 1" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0, 2.0) u0 = 0.0 - a = Float32.([i for i in 0.1:0.05:pi]) + + #generate data set + t0, t_end = tspan + instances_size = 100 + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) + batch_size = 50 + as = [i for i in range(0.1, stop = pi/2, length = batch_size)] + + patamaters_set = [] + for a_i in as + a_arr = fill(a_i, instances_size)' + s = reduce(vcat, [ts, a_arr]) + push!(patamaters_set, s) + end u_output_ = [] - for i in 1:length(a) - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a[i]) - sol1 = solve(prob, Tsit5(); saveat = 0.02) - push!(u_output_, Float32.(sol1.u')) + for a_i in as + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) + sol1 = solve(prob, Tsit5(); saveat = 0.0204) + push!(u_output_, sol1(range_).u') end + + """ + Set of training data: + * input data: mesh of 't' paired with set of parameters 'a': + * output data: set of corresponding parameter 'a' solutions u(t){a} + """ + training_mapping = (patamaters_set, u_output_) + #TODO u0 -> [u0,a]? prob = ODEProblem(linear, u0, tspan) - chain = Lux.Chain(Lux.Dense(2, 5, Lux.σ), Lux.Dense(5, 1)) - # opt = OptimizationOptimisers.Adam(0.1) + chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) opt = OptimizationOptimisers.Adam(0.03) - training_mapping = (a, u_output_) + alg = NeuralPDE.PINOODE(chain, opt, training_mapping) - sol = solve(prob, + res, phi = solve(prob, alg, verbose = true, - maxiters = 2000, abstol = 1.0f-7) + maxiters = 4000, abstol = 1.0f-10) - total_loss - total_loss, optprob, opt, callback, maxiters = sol - phi(rand(2, 10), init_params) - total_loss(init_params, 1) + predict = reduce(hcat, [phi(patamaters_set[i], res.u)' for i in 1:batch_size]) + grpound = reduce(hcat, [u_output_[i]' for i in 1:batch_size]) + @test grpound≈predict atol=3 - using Plots + # i = 30 + # plot(predict[:, i]) + # plot!(grpound[:, i]) end From 25e476443a98e9d65513d77b317ee6cadd419b32 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 13 Feb 2024 18:32:25 +0400 Subject: [PATCH 003/153] add physics loss function --- src/pino_ode_solve.jl | 72 +++++++++++++++++-------------------------- 1 file changed, 28 insertions(+), 44 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index a5f4fc6a8e..7eda9114b1 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -17,42 +17,6 @@ function PINOODE(chain, PINOODE(chain, opt, training_mapping, init_params, minibatch, kwargs) end -# mutable struct Phi{C, S} -# f::C -# st::S -# function Phi(chain::Lux.AbstractExplicitLayer) -# st = Lux.initialstates(Random.default_rng(), chain) -# new{typeof(chain), typeof(st)}(chain, st) -# end -# end - -# function (f::Phi{<:Lux.AbstractExplicitLayer})(x::Number, θ) -# y, st = f.f(adapt(parameterless_type(ComponentArrays.getdata(θ)), [x]), θ, f.st) - -# ChainRulesCore.@ignore_derivatives f.st = st -# y -# end - -# function (f::Phi{<:Lux.AbstractExplicitLayer})(x::AbstractArray, θ) -# y, st = f.f(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) -# ChainRulesCore.@ignore_derivatives f.st = st -# y -# end - -# function (f::Phi{<:Optimisers.Restructure})(x, θ) -# f.f(θ)(adapt(parameterless_type(θ), x)) -# end - -# function (f::PINOPhi{C})(t, θ) where {C <: Optimisers.Restructure} -# f.f(θ)(t) -# end -# Zygote.gradient(θ -> sum(abs2, data_loss(phi, tspan, θ, training_mapping)), init_params) - -# function inner_loss(phi::PINOPhi{C}, tspan, θ, training_mapping::Tuple) where {C} -# loss = data_loss(phi, tspan, θ, training_mapping ) #+ physics_loss() -# loss -# end - """ PINOPhi(chain::Lux.AbstractExplicitLayer, t, st) """ @@ -89,20 +53,40 @@ function (f::PINOPhi{C})(t::AbstractArray, θ) where {C <: Lux.AbstractExplicitL y end -function inner_data_loss(phi::PINOPhi{C}, θ, in_, out_) where {C} - phi(in_, θ) - out_ +function dfdx(phi::PINOPhi, t::AbstractArray, θ) + ε = [sqrt(eps(eltype(t))), zero(eltype(t))] + (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) +end + +function inner_physics_loss(phi::PINOPhi{C}, f, θ, in_) where {C} + out_ = phi(in_, θ) + dudt = dfdx(phi, in_, θ) + ts = in_[[1], :] + ps = in_[[2], :] + fs = f.(out_, ps, ts) + dudt - fs +end + +function physics_loss(phi::PINOPhi{C}, f, θ, training_mapping::Tuple{Vector, Vector}) where {C} + input_set, _ = training_mapping + data_set_size = size(input_set)[1] * size(input_set[1])[2] + loss = reduce(vcat, + [inner_physics_loss(phi, f, θ, in_) for in_ in input_set]) + sum(abs2, loss) / data_set_size end -function data_loss(phi::PINOPhi{C}, tspan, θ, training_mapping) where {C} +function data_loss(phi::PINOPhi{C}, θ, training_mapping::Tuple{Vector, Vector}) where {C} input_set, output_set = training_mapping data_set_size = size(input_set)[1] * size(input_set[1])[2] - loss = reduce(vcat,[inner_data_loss(phi, θ, in_, out_) for (in_, out_) in zip(input_set, output_set)]) - loss / data_set_size + inner_data_loss(phi::PINOPhi{C}, θ, in_, out_) where {C} = phi(in_, θ) - out_ + loss = reduce(vcat, + [inner_data_loss(phi, θ, in_, out_) for (in_, out_) in zip(input_set, output_set)]) + sum(abs2,loss) / data_set_size end -function generate_loss(phi::PINOPhi{C}, tspan, training_mapping::Tuple) where {C} +function generate_loss(phi::PINOPhi{C}, f, training_mapping::Tuple{Vector, Vector}) where {C} function loss(θ, _) - sum(abs2, data_loss(phi, tspan, θ, training_mapping)) + data_loss(phi, θ, training_mapping) + physics_loss(phi, f, θ, training_mapping) end return loss end @@ -156,7 +140,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, # alg.batch # end - inner_f = generate_loss(phi, tspan, training_mapping) + inner_f = generate_loss(phi, f, training_mapping) # Creates OptimizationFunction Object from total_loss total_loss(θ, _) = inner_f(θ, phi) From 5888e1247f13c3829f7d49116c55ab1c6f9bafc4 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 13 Feb 2024 18:58:48 +0400 Subject: [PATCH 004/153] add doc draft --- docs/pages.jl | 1 + docs/src/tutorials/pino_ode.md | 3 +++ src/pino_ode_solve.jl | 32 ++++++++++++++++++++++++++++++-- test/runtests.jl | 5 +++++ 4 files changed, 39 insertions(+), 2 deletions(-) create mode 100644 docs/src/tutorials/pino_ode.md diff --git a/docs/pages.jl b/docs/pages.jl index c1016b4c23..4d167fe265 100644 --- a/docs/pages.jl +++ b/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 Opeator ODEs" => "tutorials/pino_ode.md" #"examples/nnrode_example.md", # currently incorrect ], "PDE PINN Tutorials" => Any["Introduction to NeuralPDE for PDEs" => "tutorials/pdesystem.md", diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md new file mode 100644 index 0000000000..3df6586105 --- /dev/null +++ b/docs/src/tutorials/pino_ode.md @@ -0,0 +1,3 @@ +# Physics informed Neural Opeator ODEs Solvers + +## some example TODO \ No newline at end of file diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 7eda9114b1..3eef13a8d1 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,3 +1,32 @@ +""" + PINOODE(chain, + OptimizationOptimisers.Adam(0.1), + training_mapping + init_params = nothing; + kwargs...) + +## Positional Arguments + +* `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.AbstractExplicitLayer`. +* `opt`: The optimizer to train the neural network. +* `training_mapping`: TODO +* `init_params`: The initial parameter of the neural network. By default, this is `nothing` + which thus uses the random initialization provided by the neural network library. + +## Keyword Arguments +* `minibatch`: TODO + +## Examples + +TODO +```julia + +``` + + +## References +Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" +""" struct PINOODE{C, O, T, P, K} <: DiffEqBase.AbstractODEAlgorithm chain::C opt::O @@ -19,6 +48,7 @@ end """ PINOPhi(chain::Lux.AbstractExplicitLayer, t, st) + TODO """ mutable struct PINOPhi{C, S} chain::C @@ -41,7 +71,6 @@ end function (f::PINOPhi{C})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), [t]), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - # f.u0 .+ (t .- f.t0) .* y y end @@ -49,7 +78,6 @@ function (f::PINOPhi{C})(t::AbstractArray, θ) where {C <: Lux.AbstractExplicitL # Batch via data as row vectors y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - # f.u0 .+ (t' .- f.t0) .* y y end diff --git a/test/runtests.jl b/test/runtests.jl index e21f554fb4..09d3c56fd4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -32,6 +32,11 @@ end @time @safetestset "NNDAE" begin include("NNDAE_tests.jl") end end + if GROUP == "All" || GROUP == "ODEPINO" + @time @safetestset "pino ode" begin include("PINO_ode_tests.jl") + end + end + if GROUP == "All" || GROUP == "NNPDE2" @time @safetestset "Additional Loss" begin include("additional_loss_tests.jl") end @time @safetestset "Direction Function Approximation" begin include("direct_function_tests.jl") end From b6ae2b4e500e7dad7c26606cbc17c89503a9a849 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 15 Feb 2024 20:07:02 +0400 Subject: [PATCH 005/153] add test mapping inital condition u0 -> solution u(t){u0} --- test/PINO_ode_tests.jl | 65 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 60 insertions(+), 5 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 502592f6f9..aabe50b208 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -22,15 +22,16 @@ using NeuralPDE patamaters_set = [] for a_i in as a_arr = fill(a_i, instances_size)' - s = reduce(vcat, [ts, a_arr]) - push!(patamaters_set, s) + t_and_p = Float32.(reshape(reduce(vcat, [ts, a_arr]), 2, instances_size,1)) + push!(patamaters_set, t_and_p) end u_output_ = [] for a_i in as prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) sol1 = solve(prob, Tsit5(); saveat = 0.0204) - push!(u_output_, sol1(range_).u') + reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) + push!(u_output_, reshape_sol) end """ @@ -49,14 +50,68 @@ using NeuralPDE res, phi = solve(prob, alg, verbose = true, - maxiters = 4000, abstol = 1.0f-10) + maxiters = 1000, abstol = 1.0f-10) + predict = reduce(vcat, [phi(patamaters_set[i], res.u) for i in 1:batch_size]) + grpound = reduce(vcat, [u_output_[i] for i in 1:batch_size]) + @test grpound≈predict atol=3 + + # i = 10 + # plot(predict[:, i]) + # plot!(grpound[:, i]) +end + +@testset "Example 2" begin + linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + linear = (u, p, t) -> cos(p * t) + tspan = (0.0, 2.0) + # u0 = 0.0 + p = 1.0 + #generate data set + t0, t_end = tspan + instances_size = 100 + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) + batch_size = 50 + u0s = [i for i in range(0.0, stop = pi / 2, length = batch_size)] + + initial_condition_set = [] + for u0_i in u0s + u0_i_arr = fill(u0_i, instances_size)' + t_and_u0 = reshape(reduce(vcat, [ts, u0_i_arr]), 2, instances_size, 1) + push!(initial_condition_set, t_and_u0) + end + + u_output_ = [] + for u0_i in u0s + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) + sol1 = solve(prob, Tsit5(); saveat = 0.0204) + reshape_sol = reshape(sol1(range_).u', 1, instances_size, 1) + push!(u_output_, reshape_sol) + end + + """ + Set of training data: + * input data: mesh of 't' paired with set of initial conditions 'a': + * output data: set of corresponding parameter 'a' solutions u(t){a} + """ + training_mapping = (initial_condition_set, u_output_) + + u0 = 0 + prob = ODEProblem(linear, u0, tspan, p) + chain = Lux.Chain(Lux.Dense(2, 10, Lux.σ), Lux.Dense(10, 10, Lux.σ), Lux.Dense(10, 1)) + opt = OptimizationOptimisers.Adam(0.03) + alg = NeuralPDE.PINOODE(chain, opt, training_mapping) + + res, phi = solve(prob, + alg, verbose = true, + maxiters = 2000, abstol = 1.0f-10) predict = reduce(hcat, [phi(patamaters_set[i], res.u)' for i in 1:batch_size]) grpound = reduce(hcat, [u_output_[i]' for i in 1:batch_size]) @test grpound≈predict atol=3 - # i = 30 + # i = 2 # plot(predict[:, i]) # plot!(grpound[:, i]) end From ad0902a1ec42e81cd1baab31b55a40327a22a8c3 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 16 Feb 2024 19:34:34 +0400 Subject: [PATCH 006/153] support train family ode by initial conditions --- src/pino_ode_solve.jl | 119 +++++++++++++++++++++++++++-------------- test/PINO_ode_tests.jl | 77 ++++++++++++++------------ 2 files changed, 124 insertions(+), 72 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 3eef13a8d1..28d5bcfab9 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,7 +1,7 @@ """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), - training_mapping + train_set init_params = nothing; kwargs...) @@ -9,7 +9,7 @@ * `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.AbstractExplicitLayer`. * `opt`: The optimizer to train the neural network. -* `training_mapping`: TODO +* `train_set`: * `init_params`: The initial parameter of the neural network. By default, this is `nothing` which thus uses the random initialization provided by the neural network library. @@ -23,14 +23,24 @@ TODO ``` - ## References Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct PINOODE{C, O, T, P, K} <: DiffEqBase.AbstractODEAlgorithm + +struct TRAINSET{} + input_data::Any + output_data::Any + u0::Bool +end + +function TRAINSET(input_data, output_data; u0 = false) + TRAINSET(input_data, output_data, u0) +end + +struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm chain::C opt::O - training_mapping::T + train_set::TRAINSET init_params::P minibatch::Int kwargs::K @@ -38,46 +48,53 @@ end function PINOODE(chain, opt, - training_mapping, + train_set, init_params = nothing; minibatch = 0, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, training_mapping, init_params, minibatch, kwargs) + PINOODE(chain, opt, train_set, init_params, minibatch, kwargs) end """ PINOPhi(chain::Lux.AbstractExplicitLayer, t, st) TODO """ -mutable struct PINOPhi{C, S} +mutable struct PINOPhi{C, T, U, S} chain::C + t0::T + u0::U st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, st) - new{typeof(chain), typeof(st)}(chain, st) + function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) + new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) end end -function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) +function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, + t0, + u0, + init_params) θ, st = Lux.setup(Random.default_rng(), chain) if init_params === nothing init_params = ComponentArrays.ComponentArray(θ) else init_params = ComponentArrays.ComponentArray(init_params) end - PINOPhi(chain, st), init_params + PINOPhi(chain, t0, u0, st), init_params end -function (f::PINOPhi{C})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer} +function (f::PINOPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer, T, U} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), [t]), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - y + first(y) end -function (f::PINOPhi{C})(t::AbstractArray, θ) where {C <: Lux.AbstractExplicitLayer} +function (f::PINOPhi{C, T, U})(t::AbstractArray, + θ) where {C <: Lux.AbstractExplicitLayer, T, U} # Batch via data as row vectors y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st + # f.u0 .+ (t[[1], :, :] .- f.t0) .* y y end @@ -86,35 +103,63 @@ function dfdx(phi::PINOPhi, t::AbstractArray, θ) (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end -function inner_physics_loss(phi::PINOPhi{C}, f, θ, in_) where {C} +function inner_physics_loss(phi::PINOPhi{C, T, U}, f, θ, in_::AbstractArray) where {C, T, U} + ts = in_[[1], :, :] #TODO remove dependence on dimension + ps = in_[[2], :, :] out_ = phi(in_, θ) dudt = dfdx(phi, in_, θ) - ts = in_[[1], :] - ps = in_[[2], :] fs = f.(out_, ps, ts) dudt - fs end -function physics_loss(phi::PINOPhi{C}, f, θ, training_mapping::Tuple{Vector, Vector}) where {C} - input_set, _ = training_mapping +function inner_physics_loss(phi::PINOPhi{C, T, U}, + f, + θ, + in_::AbstractArray, + p::AbstractArray) where {C, T, U} + ts = in_[[1], :, :] + out_ = phi(in_, θ) + dudt = dfdx(phi, in_, θ) + fs = f.(out_, p, ts) + dudt - fs +end + +function physics_loss(phi::PINOPhi{C, T, U}, f, θ, train_set::TRAINSET, p) where {C, T, U} + input_set = train_set.input_data data_set_size = size(input_set)[1] * size(input_set[1])[2] - loss = reduce(vcat, - [inner_physics_loss(phi, f, θ, in_) for in_ in input_set]) + if train_set.u0 == false + loss = reduce(vcat, + [inner_physics_loss(phi, f, θ, in_) for in_ in input_set]) + else #train_set.u!==nothing + p = fill(p, 1, size(input_set[1])[2], 1) + loss = reduce(vcat, + [inner_physics_loss(phi, f, θ, in_, p) for in_ in input_set]) + end sum(abs2, loss) / data_set_size end -function data_loss(phi::PINOPhi{C}, θ, training_mapping::Tuple{Vector, Vector}) where {C} - input_set, output_set = training_mapping +function inner_data_loss(phi::PINOPhi{C, T, U}, + θ, + in_::AbstractArray, + out_::AbstractArray) where {C, T, U} + phi(in_, θ) - out_ +end + +function data_loss(phi::PINOPhi{C, T, U}, + θ, + train_set::TRAINSET) where {C, T, U} + input_set, output_set = train_set.input_data, train_set.output_data data_set_size = size(input_set)[1] * size(input_set[1])[2] - inner_data_loss(phi::PINOPhi{C}, θ, in_, out_) where {C} = phi(in_, θ) - out_ loss = reduce(vcat, [inner_data_loss(phi, θ, in_, out_) for (in_, out_) in zip(input_set, output_set)]) - sum(abs2,loss) / data_set_size + sum(abs2, loss) / data_set_size end -function generate_loss(phi::PINOPhi{C}, f, training_mapping::Tuple{Vector, Vector}) where {C} +function generate_loss(phi::PINOPhi{C, T, U}, + f, + train_set::TRAINSET, p) where {C, T, U} function loss(θ, _) - data_loss(phi, θ, training_mapping) + physics_loss(phi, f, θ, training_mapping) + data_loss(phi, θ, train_set) + physics_loss(phi, f, θ, train_set, p) end return loss end @@ -128,10 +173,11 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, verbose = false, saveat = nothing, maxiters = nothing) - # u0 = prob.u0 ? TODO tspan = prob.tspan + t0 = tspan[1] f = prob.f p = prob.p + u0 = prob.u0 # param_estim = alg.param_estim chain = alg.chain @@ -139,13 +185,13 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, init_params = alg.init_params # mapping between functional space of some vararible 'a' of equation (for example initial - # condition {u(t0 x)} or parameter) join and solution of equation u(t) - training_mapping = alg.training_mapping + # condition {u(t0 x)} or parameter p) join and solution of equation u(t) + train_set = alg.train_set !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - phi, init_params = generate_pino_phi_θ(chain, init_params) + phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) init_params = ComponentArrays.ComponentArray(init_params) @@ -153,7 +199,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - phi(first(training_mapping[1]), init_params) + phi(first(train_set.input_data), init_params) #TODO first(train_set.input_data) catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) @@ -162,13 +208,8 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end end - # batch = if alg.batch === nothing - # true - # else - # alg.batch - # end - inner_f = generate_loss(phi, f, training_mapping) + inner_f = generate_loss(phi, f, train_set, p) # Creates OptimizationFunction Object from total_loss total_loss(θ, _) = inner_f(θ, phi) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index aabe50b208..10d5cf3395 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,5 +1,5 @@ using Test -using OrdinaryDiffEq,OptimizationOptimisers +using OrdinaryDiffEq, OptimizationOptimisers using Flux, Lux using Statistics, Random using NeuralOperators @@ -9,20 +9,19 @@ using NeuralPDE linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0, 2.0) - u0 = 0.0 - + u0 = 2.0 #generate data set t0, t_end = tspan instances_size = 100 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) batch_size = 50 - as = [i for i in range(0.1, stop = pi/2, length = batch_size)] - + as = [i for i in range(0.1, stop = pi / 2, length = batch_size)] + # p = ps TODO patamaters_set = [] for a_i in as a_arr = fill(a_i, instances_size)' - t_and_p = Float32.(reshape(reduce(vcat, [ts, a_arr]), 2, instances_size,1)) + t_and_p = Float32.(reshape(reduce(vcat, [ts, a_arr]), 2, instances_size, 1)) push!(patamaters_set, t_and_p) end @@ -39,34 +38,30 @@ using NeuralPDE * input data: mesh of 't' paired with set of parameters 'a': * output data: set of corresponding parameter 'a' solutions u(t){a} """ - training_mapping = (patamaters_set, u_output_) - #TODO u0 -> [u0,a]? + train_set = NeuralPDE.TRAINSET(patamaters_set, u_output_) + prob = ODEProblem(linear, u0, tspan) chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, training_mapping) + alg = NeuralPDE.PINOODE(chain, opt, train_set) res, phi = solve(prob, alg, verbose = true, - maxiters = 1000, abstol = 1.0f-10) - - predict = reduce(vcat, [phi(patamaters_set[i], res.u) for i in 1:batch_size]) - grpound = reduce(vcat, [u_output_[i] for i in 1:batch_size]) - @test grpound≈predict atol=3 + maxiters = 2000, abstol = 1.0f-10) - # i = 10 - # plot(predict[:, i]) - # plot!(grpound[:, i]) + predict = reduce(vcat, [phi(train_set.input_data[i], res.u) for i in 1:batch_size]) + ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) + @test ground≈predict atol=0.5 end @testset "Example 2" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0, 2.0) - # u0 = 0.0 - p = 1.0 + p = pi + u0 = 2 #generate data set t0, t_end = tspan instances_size = 100 @@ -74,10 +69,11 @@ end ts = reshape(collect(range_), 1, instances_size) batch_size = 50 u0s = [i for i in range(0.0, stop = pi / 2, length = batch_size)] - initial_condition_set = [] + u0_arr = [] for u0_i in u0s - u0_i_arr = fill(u0_i, instances_size)' + u0_i_arr = reshape(fill(u0_i, instances_size)', 1, instances_size, 1) + push!(u0_arr, u0_i_arr) t_and_u0 = reshape(reduce(vcat, [ts, u0_i_arr]), 2, instances_size, 1) push!(initial_condition_set, t_and_u0) end @@ -95,23 +91,38 @@ end * input data: mesh of 't' paired with set of initial conditions 'a': * output data: set of corresponding parameter 'a' solutions u(t){a} """ - training_mapping = (initial_condition_set, u_output_) - - u0 = 0 + train_set = NeuralPDE.TRAINSET(initial_condition_set, u_output_; u0 = true) prob = ODEProblem(linear, u0, tspan, p) - chain = Lux.Chain(Lux.Dense(2, 10, Lux.σ), Lux.Dense(10, 10, Lux.σ), Lux.Dense(10, 1)) + chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, training_mapping) - + alg = NeuralPDE.PINOODE(chain, opt, train_set) res, phi = solve(prob, alg, verbose = true, maxiters = 2000, abstol = 1.0f-10) - predict = reduce(hcat, [phi(patamaters_set[i], res.u)' for i in 1:batch_size]) - grpound = reduce(hcat, [u_output_[i]' for i in 1:batch_size]) - @test grpound≈predict atol=3 + predict = reduce(vcat, [phi(train_set.input_data[i], res.u) for i in 1:batch_size]) + ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) + @test ground≈predict atol=1 +end + +@testset "lotka volterra" begin + function lotka_volterra(u, p, t) + # Model parameters. + α, β, γ, δ = p + # Current state. + x, y = u + + # Evaluate differential equations. + dx = (α - β * y) * x # prey + dy = (δ * x - γ) * y # predator + + return [dx, dy] + end + u0 = [1.0, 1.0] + p = [1.5, 1.0, 3.0, 1.0] + tspan = (0.0, 4.0) + prob = ODEProblem(lotka_volterra, u0, tspan, p) - # i = 2 - # plot(predict[:, i]) - # plot!(grpound[:, i]) + dt = 0.01 + solution = solve(prob, Tsit5(); saveat = dt) end From c7b10f8e69e54cb60bb67cf673a538458f3a83da Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 28 Feb 2024 19:15:22 +0400 Subject: [PATCH 007/153] generate prob set --- src/pino_ode_solve.jl | 126 ++++++++++++++++++++++------------------- test/PINO_ode_tests.jl | 105 ++++++++++++++++++++++------------ 2 files changed, 137 insertions(+), 94 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 28d5bcfab9..ccd46cd23a 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -26,15 +26,15 @@ TODO ## References Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ - -struct TRAINSET{} - input_data::Any - output_data::Any - u0::Bool +#TODO +struct TRAINSET{} #T + input_data::Vector{ODEProblem} + output_data::Vector{Array} + isu0::Bool end -function TRAINSET(input_data, output_data; u0 = false) - TRAINSET(input_data, output_data, u0) +function TRAINSET(input_data, output_data; isu0 = false) + TRAINSET(input_data, output_data, isu0) end struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm @@ -60,19 +60,17 @@ end PINOPhi(chain::Lux.AbstractExplicitLayer, t, st) TODO """ -mutable struct PINOPhi{C, T, U, S} +mutable struct PINOPhi{C, T, S} chain::C t0::T - u0::U st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) - new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) + function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, st) + new{typeof(chain), typeof(t0), typeof(st)}(chain, t0, st) end end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, t0, - u0, init_params) θ, st = Lux.setup(Random.default_rng(), chain) if init_params === nothing @@ -80,7 +78,7 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, else init_params = ComponentArrays.ComponentArray(init_params) end - PINOPhi(chain, t0, u0, st), init_params + PINOPhi(chain, t0, st), init_params end function (f::PINOPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer, T, U} @@ -103,63 +101,74 @@ function dfdx(phi::PINOPhi, t::AbstractArray, θ) (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end -function inner_physics_loss(phi::PINOPhi{C, T, U}, f, θ, in_::AbstractArray) where {C, T, U} - ts = in_[[1], :, :] #TODO remove dependence on dimension - ps = in_[[2], :, :] - out_ = phi(in_, θ) - dudt = dfdx(phi, in_, θ) - fs = f.(out_, ps, ts) - dudt - fs -end - function inner_physics_loss(phi::PINOPhi{C, T, U}, - f, θ, - in_::AbstractArray, - p::AbstractArray) where {C, T, U} - ts = in_[[1], :, :] + ts::AbstractArray, + prob::ODEProblem, + isu0::Bool) where {C, T, U} + u0 = prob.u0 + p = prob.p + f = prob.f + if isu0 == true + in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2])]) + else + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + end out_ = phi(in_, θ) dudt = dfdx(phi, in_, θ) fs = f.(out_, p, ts) dudt - fs end -function physics_loss(phi::PINOPhi{C, T, U}, f, θ, train_set::TRAINSET, p) where {C, T, U} - input_set = train_set.input_data - data_set_size = size(input_set)[1] * size(input_set[1])[2] - if train_set.u0 == false - loss = reduce(vcat, - [inner_physics_loss(phi, f, θ, in_) for in_ in input_set]) - else #train_set.u!==nothing - p = fill(p, 1, size(input_set[1])[2], 1) - loss = reduce(vcat, - [inner_physics_loss(phi, f, θ, in_, p) for in_ in input_set]) - end - sum(abs2, loss) / data_set_size +function physics_loss(phi::PINOPhi{C, T, U}, + θ, + ts::AbstractArray, + train_set::TRAINSET ) where {C, T, U} + prob_set, output_data = train_set.input_data, train_set.output_data + norm = size(output_data)[1] * size(output_data[1])[2] + loss = reduce(vcat, + [inner_physics_loss(phi, θ, ts, prob, train_set.isu0) for prob in prob_set]) + sum(abs2, loss) / norm end function inner_data_loss(phi::PINOPhi{C, T, U}, θ, - in_::AbstractArray, - out_::AbstractArray) where {C, T, U} + ts::AbstractArray, + prob::ODEProblem, + out_::AbstractArray, + isu0::Bool) where {C, T, U} + u0 = prob.u0 + p = prob.p + f = prob.f + if isu0 == true + in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2])]) + else + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + end phi(in_, θ) - out_ end function data_loss(phi::PINOPhi{C, T, U}, θ, - train_set::TRAINSET) where {C, T, U} - input_set, output_set = train_set.input_data, train_set.output_data - data_set_size = size(input_set)[1] * size(input_set[1])[2] + ts::AbstractArray, + train_set::TRAINSET + ) where {C, T, U} + prob_set, output_data = train_set.input_data, train_set.output_data + norm = size(output_data)[1] * size(output_data[1])[2] loss = reduce(vcat, - [inner_data_loss(phi, θ, in_, out_) for (in_, out_) in zip(input_set, output_set)]) - sum(abs2, loss) / data_set_size + [inner_data_loss(phi, θ, ts, prob, out_, train_set.isu0) + for (prob, out_) in zip(prob_set, output_data)]) + sum(abs2, loss) / norm end -function generate_loss(phi::PINOPhi{C, T, U}, - f, - train_set::TRAINSET, p) where {C, T, U} +function generate_loss(phi::PINOPhi{C, T, U}, train_set::TRAINSET, tspan) where {C, T, U} + t0 = tspan[1] + t_end = tspan[2] + instances_size = size(train_set.output_data[1])[2] + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) function loss(θ, _) - data_loss(phi, θ, train_set) + physics_loss(phi, f, θ, train_set, p) + data_loss(phi, θ, ts, train_set) + physics_loss(phi, θ, ts, train_set) end return loss end @@ -175,9 +184,9 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, maxiters = nothing) tspan = prob.tspan t0 = tspan[1] - f = prob.f - p = prob.p - u0 = prob.u0 + # f = prob.f + # p = prob.p + # u0 = prob.u0 # param_estim = alg.param_estim chain = alg.chain @@ -191,7 +200,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) + phi, init_params = generate_pino_phi_θ(chain, t0, init_params) init_params = ComponentArrays.ComponentArray(init_params) @@ -199,7 +208,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - phi(first(train_set.input_data), init_params) #TODO first(train_set.input_data) + phi(rand(chain.layers.layer_1.in_dims, 10), init_params) #TODO input data catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) @@ -208,11 +217,10 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end end - - inner_f = generate_loss(phi, f, train_set, p) - - # Creates OptimizationFunction Object from total_loss - total_loss(θ, _) = inner_f(θ, phi) + # dt + total_loss = generate_loss(phi, train_set, tspan) + # total_loss_(init_params, nothing) + # Zygote.gradient(p -> total_loss_(p, nothing), init_params) # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 10d5cf3395..ff2d4387d4 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -17,19 +17,21 @@ using NeuralPDE ts = reshape(collect(range_), 1, instances_size) batch_size = 50 as = [i for i in range(0.1, stop = pi / 2, length = batch_size)] - # p = ps TODO - patamaters_set = [] - for a_i in as - a_arr = fill(a_i, instances_size)' - t_and_p = Float32.(reshape(reduce(vcat, [ts, a_arr]), 2, instances_size, 1)) - push!(patamaters_set, t_and_p) - end - - u_output_ = [] + # # p = ps TODO + # patamaters_set = [] + # for a_i in as + # a_arr = fill(a_i, instances_size)' + # t_and_p = Float32.(reshape(reduce(vcat, [ts, a_arr]), 2, instances_size, 1)) + # push!(patamaters_set, t_and_p) + # end + + u_output_ = Array{Float32, 3}[] + prob_set = [] for a_i in as prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) sol1 = solve(prob, Tsit5(); saveat = 0.0204) reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) + push!(prob_set, prob) push!(u_output_, reshape_sol) end @@ -38,20 +40,21 @@ using NeuralPDE * input data: mesh of 't' paired with set of parameters 'a': * output data: set of corresponding parameter 'a' solutions u(t){a} """ - - train_set = NeuralPDE.TRAINSET(patamaters_set, u_output_) - + train_set = TRAINSET(prob_set, u_output_) + #TODO u0 ? prob = ODEProblem(linear, u0, tspan) chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, train_set) + alg = PINOODE(chain, opt, train_set) res, phi = solve(prob, alg, verbose = true, maxiters = 2000, abstol = 1.0f-10) - predict = reduce(vcat, [phi(train_set.input_data[i], res.u) for i in 1:batch_size]) + predict = reduce(vcat, + [phi(reduce(vcat, [ts, fill(train_set.input_data[i].p, 1, size(ts)[2])]), res.u) + for i in 1:batch_size]) ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) @test ground≈predict atol=0.5 end @@ -69,20 +72,22 @@ end ts = reshape(collect(range_), 1, instances_size) batch_size = 50 u0s = [i for i in range(0.0, stop = pi / 2, length = batch_size)] - initial_condition_set = [] - u0_arr = [] - for u0_i in u0s - u0_i_arr = reshape(fill(u0_i, instances_size)', 1, instances_size, 1) - push!(u0_arr, u0_i_arr) - t_and_u0 = reshape(reduce(vcat, [ts, u0_i_arr]), 2, instances_size, 1) - push!(initial_condition_set, t_and_u0) - end - - u_output_ = [] + # initial_condition_set = [] + # u0_arr = [] + # for u0_i in u0s + # u0_i_arr = reshape(fill(u0_i, instances_size)', 1, instances_size, 1) + # push!(u0_arr, u0_i_arr) + # t_and_u0 = reshape(reduce(vcat, [ts, u0_i_arr]), 2, instances_size, 1) + # push!(initial_condition_set, t_and_u0) + # end + + u_output_ = Array{Float32, 3}[] + prob_set = [] for u0_i in u0s prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) sol1 = solve(prob, Tsit5(); saveat = 0.0204) reshape_sol = reshape(sol1(range_).u', 1, instances_size, 1) + push!(prob_set, prob) push!(u_output_, reshape_sol) end @@ -90,19 +95,22 @@ end Set of training data: * input data: mesh of 't' paired with set of initial conditions 'a': * output data: set of corresponding parameter 'a' solutions u(t){a} - """ - train_set = NeuralPDE.TRAINSET(initial_condition_set, u_output_; u0 = true) + """ + train_set = TRAINSET(prob_set, u_output_; isu0 = true); + #TODO u0 ? prob = ODEProblem(linear, u0, tspan, p) - chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) + chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)); opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, train_set) + alg = PINOODE(chain, opt, train_set); res, phi = solve(prob, alg, verbose = true, maxiters = 2000, abstol = 1.0f-10) - predict = reduce(vcat, [phi(train_set.input_data[i], res.u) for i in 1:batch_size]) + predict = reduce(vcat, + [phi(reduce(vcat, [ts, fill(train_set.input_data[i].u0, 1, size(ts)[2])]), res.u) + for i in 1:batch_size]) ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) - @test ground≈predict atol=1 + @test ground≈predict atol=0.5 end @testset "lotka volterra" begin @@ -111,18 +119,45 @@ end α, β, γ, δ = p # Current state. x, y = u - # Evaluate differential equations. dx = (α - β * y) * x # prey dy = (δ * x - γ) * y # predator - return [dx, dy] end u0 = [1.0, 1.0] p = [1.5, 1.0, 3.0, 1.0] tspan = (0.0, 4.0) - prob = ODEProblem(lotka_volterra, u0, tspan, p) - dt = 0.01 - solution = solve(prob, Tsit5(); saveat = dt) + + instances_size = 100 + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) + batch_size = 50 + ps = [p + randn(4) for _ in 1:batch_size] + + u_output_ = [] + prob_set = [] + for p_i in ps + prob = ODEProblem(lotka_volterra, u0, tspan, p_i) + solution = solve(prob, Tsit5(); saveat = dt) + reshape_sol = Float32.(reduce(hcat,solution(range_).u)) + push!(prob_set, prob) + push!(u_output_, reshape_sol) + end + + train_set = TRAINSET(prob_set, u_output_) + #TODO u0 ? + prob = ODEProblem(linear, u0, tspan, p) + chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)) + opt = OptimizationOptimisers.Adam(0.03) + alg = PINOODE(chain, opt, train_set) + res, phi = solve(prob, + alg, verbose = true, + maxiters = 2000, abstol = 1.0f-10) + + predict = reduce(vcat, + [phi(reduce(vcat, [ts, fill(train_set.input_data[i].p, 1, size(ts)[2])]), res.u) + for i in 1:batch_size]) + ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) + @test ground≈predict atol=2 end From 911ec4e809c610367a06a5672754efa111151b48 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 29 Feb 2024 11:32:59 +0400 Subject: [PATCH 008/153] ode system --- src/pino_ode_solve.jl | 30 +++++++++++++++++++++++------- test/PINO_ode_tests.jl | 22 +++++++++++++--------- 2 files changed, 36 insertions(+), 16 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index ccd46cd23a..bfbb43baf0 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -97,7 +97,7 @@ function (f::PINOPhi{C, T, U})(t::AbstractArray, end function dfdx(phi::PINOPhi, t::AbstractArray, θ) - ε = [sqrt(eps(eltype(t))), zero(eltype(t))] + ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end @@ -112,11 +112,24 @@ function inner_physics_loss(phi::PINOPhi{C, T, U}, if isu0 == true in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2])]) else - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + if p isa Number + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + elseif p isa Vector + in_ = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2]))]) + else + error("p should be a number or a vector") + end end out_ = phi(in_, θ) dudt = dfdx(phi, in_, θ) - fs = f.(out_, p, ts) + if p isa Number + fs = f.(out_, p, ts) + elseif p isa Vector + fs = reduce(hcat, [f(out_[:, i], p, ts[i]) for i in 1:size(out_, 2)]) + else + error("p should be a number or a vector") + end + dudt - fs end @@ -143,7 +156,13 @@ function inner_data_loss(phi::PINOPhi{C, T, U}, if isu0 == true in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2])]) else - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + if p isa Number + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + elseif p isa Vector + in_ = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2]))]) + else + error("p should be a number or a vector") + end end phi(in_, θ) - out_ end @@ -217,10 +236,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end end - # dt total_loss = generate_loss(phi, train_set, tspan) - # total_loss_(init_params, nothing) - # Zygote.gradient(p -> total_loss_(p, nothing), init_params) # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ff2d4387d4..6e7fbc7621 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -40,13 +40,13 @@ using NeuralPDE * input data: mesh of 't' paired with set of parameters 'a': * output data: set of corresponding parameter 'a' solutions u(t){a} """ - train_set = TRAINSET(prob_set, u_output_) + train_set = TRAINSET(prob_set, u_output_); #TODO u0 ? prob = ODEProblem(linear, u0, tspan) - chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) + chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)); opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(chain, opt, train_set) + alg = PINOODE(chain, opt, train_set); res, phi = solve(prob, alg, verbose = true, @@ -133,31 +133,35 @@ end range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) batch_size = 50 - ps = [p + randn(4) for _ in 1:batch_size] + ps = [p .+ i*[0.015, 0.01, 0.03, 0.01] for i in 1:batch_size] u_output_ = [] prob_set = [] + plot_set =[] for p_i in ps prob = ODEProblem(lotka_volterra, u0, tspan, p_i) solution = solve(prob, Tsit5(); saveat = dt) reshape_sol = Float32.(reduce(hcat,solution(range_).u)) + push!(plot_set,solution) push!(prob_set, prob) push!(u_output_, reshape_sol) end - train_set = TRAINSET(prob_set, u_output_) + train_set = TRAINSET(prob_set, u_output_); #TODO u0 ? prob = ODEProblem(linear, u0, tspan, p) - chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)) + chain = Lux.Chain(Lux.Dense(5, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)) opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(chain, opt, train_set) + alg = PINOODE(chain, opt, train_set); res, phi = solve(prob, alg, verbose = true, maxiters = 2000, abstol = 1.0f-10) predict = reduce(vcat, - [phi(reduce(vcat, [ts, fill(train_set.input_data[i].p, 1, size(ts)[2])]), res.u) + [phi(reduce(vcat, + [ts, reduce(hcat, fill(train_set.input_data[i].p, 1, size(ts)[2]))]), + res.u) for i in 1:batch_size]) ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) - @test ground≈predict atol=2 + @test ground≈predict atol=5 end From a4ff4135ceaa3274190e406f49a3a52a07dd215e Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 1 Mar 2024 18:30:19 +0400 Subject: [PATCH 009/153] support Fourier Neural Operator --- src/NeuralPDE.jl | 1 + src/pino_ode_solve.jl | 45 +++++++++++-------- test/PINO_ode_tests.jl | 99 ++++++++++++++++++++---------------------- 3 files changed, 74 insertions(+), 71 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 7d6e703e3b..b175da44f0 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -31,6 +31,7 @@ using SciMLBase: @add_kwonly, parameterless_type using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using ChainRulesCore: @non_differentiable +using NeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index bfbb43baf0..416964bbf0 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -53,6 +53,7 @@ function PINOODE(chain, minibatch = 0, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) + #TODO transform convert complex numbers to zero PINOODE(chain, opt, train_set, init_params, minibatch, kwargs) end @@ -92,12 +93,12 @@ function (f::PINOPhi{C, T, U})(t::AbstractArray, # Batch via data as row vectors y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - # f.u0 .+ (t[[1], :, :] .- f.t0) .* y y end function dfdx(phi::PINOPhi, t::AbstractArray, θ) - ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] + ε = [sqrt(eps(eltype(t))), zero(eltype(t))] + # ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end @@ -110,35 +111,38 @@ function inner_physics_loss(phi::PINOPhi{C, T, U}, p = prob.p f = prob.f if isu0 == true - in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2])]) + in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) + #TODO for all case p and u0 + # u0 isa Vector + # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) else if p isa Number - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) elseif p isa Vector - in_ = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2]))]) + #TODO nno for Vector + inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) + in_ = reshape(inner, size(inner)..., 1) else error("p should be a number or a vector") end end out_ = phi(in_, θ) - dudt = dfdx(phi, in_, θ) if p isa Number - fs = f.(out_, p, ts) + fs = f.f.(out_, p, ts) elseif p isa Vector - fs = reduce(hcat, [f(out_[:, i], p, ts[i]) for i in 1:size(out_, 2)]) + fs = reduce(hcat, [f.f(out_[:, i], p, ts[i]) for i in 1:size(out_, 2)]) else error("p should be a number or a vector") end - - dudt - fs + NeuralOperators.l₂loss(dfdx(phi, in_, θ), fs) end function physics_loss(phi::PINOPhi{C, T, U}, θ, ts::AbstractArray, - train_set::TRAINSET ) where {C, T, U} + train_set::TRAINSET) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data - norm = size(output_data)[1] * size(output_data[1])[2] + norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] loss = reduce(vcat, [inner_physics_loss(phi, θ, ts, prob, train_set.isu0) for prob in prob_set]) sum(abs2, loss) / norm @@ -154,26 +158,29 @@ function inner_data_loss(phi::PINOPhi{C, T, U}, p = prob.p f = prob.f if isu0 == true - in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2])]) + in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) + #TODO for all case p and u0 + # u0 isa Vector + # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) else if p isa Number in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) elseif p isa Vector - in_ = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2]))]) + inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) + in_ = reshape(inner, size(inner)..., 1) else error("p should be a number or a vector") end end - phi(in_, θ) - out_ + NeuralOperators.l₂loss(phi(in_, θ), out_) end function data_loss(phi::PINOPhi{C, T, U}, θ, ts::AbstractArray, - train_set::TRAINSET - ) where {C, T, U} + train_set::TRAINSET) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data - norm = size(output_data)[1] * size(output_data[1])[2] + norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] loss = reduce(vcat, [inner_data_loss(phi, θ, ts, prob, out_, train_set.isu0) for (prob, out_) in zip(prob_set, output_data)]) @@ -227,7 +234,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - phi(rand(chain.layers.layer_1.in_dims, 10), init_params) #TODO input data + # phi(rand(5, 100, 1), init_params) #TODO input data catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 6e7fbc7621..10f0ffde70 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -9,21 +9,14 @@ using NeuralPDE linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0, 2.0) - u0 = 2.0 + u0 = 0.0 #generate data set t0, t_end = tspan - instances_size = 100 + instances_size = 40 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) - batch_size = 50 - as = [i for i in range(0.1, stop = pi / 2, length = batch_size)] - # # p = ps TODO - # patamaters_set = [] - # for a_i in as - # a_arr = fill(a_i, instances_size)' - # t_and_p = Float32.(reshape(reduce(vcat, [ts, a_arr]), 2, instances_size, 1)) - # push!(patamaters_set, t_and_p) - # end + batch_size = 40 + as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] u_output_ = Array{Float32, 3}[] prob_set = [] @@ -37,23 +30,27 @@ using NeuralPDE """ Set of training data: - * input data: mesh of 't' paired with set of parameters 'a': - * output data: set of corresponding parameter 'a' solutions u(t){a} + * input data: set of parameters 'a': + * output data: set of solutions u(t){a} corresponding parameter 'a' """ - train_set = TRAINSET(prob_set, u_output_); + train_set = NeuralPDE.TRAINSET(prob_set, u_output_); #TODO u0 ? - prob = ODEProblem(linear, u0, tspan) + prob = ODEProblem(linear, u0, tspan, 0) chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)); + flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), + σ = gelu); opt = OptimizationOptimisers.Adam(0.03) - - alg = PINOODE(chain, opt, train_set); + alg = NeuralPDE.PINOODE(chain, opt, train_set); res, phi = solve(prob, alg, verbose = true, - maxiters = 2000, abstol = 1.0f-10) + maxiters = 1000, abstol = 1.0f-10) predict = reduce(vcat, - [phi(reduce(vcat, [ts, fill(train_set.input_data[i].p, 1, size(ts)[2])]), res.u) + [phi( + reshape(reduce(vcat, [ts, fill(train_set.input_data[i].p, 1, size(ts)[2])]), + 2, instances_size, 1), + res.u) for i in 1:batch_size]) ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) @test ground≈predict atol=0.5 @@ -72,45 +69,40 @@ end ts = reshape(collect(range_), 1, instances_size) batch_size = 50 u0s = [i for i in range(0.0, stop = pi / 2, length = batch_size)] - # initial_condition_set = [] - # u0_arr = [] - # for u0_i in u0s - # u0_i_arr = reshape(fill(u0_i, instances_size)', 1, instances_size, 1) - # push!(u0_arr, u0_i_arr) - # t_and_u0 = reshape(reduce(vcat, [ts, u0_i_arr]), 2, instances_size, 1) - # push!(initial_condition_set, t_and_u0) - # end - - u_output_ = Array{Float32, 3}[] prob_set = [] + u_output_ = Array{Float32, 3}[] for u0_i in u0s prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) - sol1 = solve(prob, Tsit5(); saveat = 0.0204) - reshape_sol = reshape(sol1(range_).u', 1, instances_size, 1) + sol = reshape([prob.f.analytic(u0_i, p, t) for t in ts], 1, instances_size, 1) + push!(u_output_, sol) push!(prob_set, prob) - push!(u_output_, reshape_sol) end """ Set of training data: - * input data: mesh of 't' paired with set of initial conditions 'a': - * output data: set of corresponding parameter 'a' solutions u(t){a} + * input data: set of initial conditions 'a': + * output data: set of solutions u(t){a} corresponding initial conditions 'a' """ - train_set = TRAINSET(prob_set, u_output_; isu0 = true); + train_set = NeuralPDE.TRAINSET(prob_set, u_output_; isu0 = true); #TODO u0 ? - prob = ODEProblem(linear, u0, tspan, p) + prob = ODEProblem(linear, 0, tspan, p) chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)); + flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), + σ = gelu); opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(chain, opt, train_set); + alg = NeuralPDE.PINOODE(chain, opt, train_set); res, phi = solve(prob, alg, verbose = true, - maxiters = 2000, abstol = 1.0f-10) + maxiters = 1000, abstol = 1.0f-10) predict = reduce(vcat, - [phi(reduce(vcat, [ts, fill(train_set.input_data[i].u0, 1, size(ts)[2])]), res.u) + [phi( + reshape(reduce(vcat, [ts, fill(train_set.input_data[i].u0, 1, size(ts)[2])]), + 2, instances_size, 1), + res.u) for i in 1:batch_size]) ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) - @test ground≈predict atol=0.5 + @test ground≈predict atol=1. end @testset "lotka volterra" begin @@ -129,38 +121,41 @@ end tspan = (0.0, 4.0) dt = 0.01 - instances_size = 100 + instances_size = 50 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) batch_size = 50 - ps = [p .+ i*[0.015, 0.01, 0.03, 0.01] for i in 1:batch_size] + ps = [p .+ i * [0.015, 0.01, 0.03, 0.01] for i in 1:batch_size] u_output_ = [] prob_set = [] - plot_set =[] + plot_set = [] for p_i in ps prob = ODEProblem(lotka_volterra, u0, tspan, p_i) solution = solve(prob, Tsit5(); saveat = dt) - reshape_sol = Float32.(reduce(hcat,solution(range_).u)) - push!(plot_set,solution) + reshape_sol = Float32.(reduce(hcat, solution(range_).u)) + push!(plot_set, solution) push!(prob_set, prob) push!(u_output_, reshape_sol) end - train_set = TRAINSET(prob_set, u_output_); + train_set = NeuralPDE.TRAINSET(prob_set, u_output_) #TODO u0 ? prob = ODEProblem(linear, u0, tspan, p) - chain = Lux.Chain(Lux.Dense(5, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)) + chain = Lux.Chain(Lux.Dense(5, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)); + flat_no = FourierNeuralOperator(ch = (5, 16, 16, 16, 16, 16, 32, 2), modes = (16,), + σ = gelu); opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(chain, opt, train_set); + alg = NeuralPDE.PINOODE(chain, opt, train_set) res, phi = solve(prob, alg, verbose = true, - maxiters = 2000, abstol = 1.0f-10) + maxiters = 1500, abstol = 1.0f-10) predict = reduce(vcat, - [phi(reduce(vcat, - [ts, reduce(hcat, fill(train_set.input_data[i].p, 1, size(ts)[2]))]), - res.u) + [phi( + reduce(vcat, + [ts, reduce(hcat, fill(train_set.input_data[i].p, 1, size(ts)[2]))]), + res.u) for i in 1:batch_size]) ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) @test ground≈predict atol=5 From cbf01aff29760366144b004a99cfe124a12d31c7 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 5 Mar 2024 20:18:44 +0400 Subject: [PATCH 010/153] update --- src/pino_ode_solve.jl | 200 +++++++++++++++++++++++++++++++---------- test/PINO_ode_tests.jl | 60 +++++++++---- 2 files changed, 199 insertions(+), 61 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 416964bbf0..e383bfec7e 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -29,7 +29,7 @@ Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Eq #TODO struct TRAINSET{} #T input_data::Vector{ODEProblem} - output_data::Vector{Array} + output_data::Array isu0::Bool end @@ -61,17 +61,19 @@ end PINOPhi(chain::Lux.AbstractExplicitLayer, t, st) TODO """ -mutable struct PINOPhi{C, T, S} +mutable struct PINOPhi{C, T, U, S} chain::C t0::T + u0::U st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, st) - new{typeof(chain), typeof(t0), typeof(st)}(chain, t0, st) + function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) + new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0,u0, st) end end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, t0, + u0, init_params) θ, st = Lux.setup(Random.default_rng(), chain) if init_params === nothing @@ -79,25 +81,44 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, else init_params = ComponentArrays.ComponentArray(init_params) end - PINOPhi(chain, t0, st), init_params + PINOPhi(chain, t0, u0, st), init_params end -function (f::PINOPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer, T, U} - y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), [t]), θ, f.st) - ChainRulesCore.@ignore_derivatives f.st = st - first(y) -end +# function (f::PINOPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer, T, U} +# y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), [t]), θ, f.st) +# ChainRulesCore.@ignore_derivatives f.st = st +# first(y) +# end function (f::PINOPhi{C, T, U})(t::AbstractArray, θ) where {C <: Lux.AbstractExplicitLayer, T, U} # Batch via data as row vectors - y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) + # y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) + y, st = f.chain(t, θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - y + # y + f.u0 .+ (t[[1],:,:] .- f.t0) .* y end +# feature_dims = 2:(ndims(t) - 1) +# loss = sum( t, dims = feature_dims) + # loss = sum(.√(sum(abs2, 𝐲̂ - 𝐲, dims = feature_dims))) + # y_norm = sum(.√(sum(abs2, 𝐲, dims = feature_dims))) + + # return loss / y_norm +# function dfdx(phi::PINOPhi, t::AbstractArray, θ) +# ε = sqrt(eps(eltype(t))) +# εs = [ε, zero(eltype(t))] +# # ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] +# # ChainRulesCore.@ignore_derivatives tl = t .+ ε +# tl = t .+ ε +# tr = t +# (phi(tl, θ) - phi(tr, θ)) ./ ε +# end + function dfdx(phi::PINOPhi, t::AbstractArray, θ) ε = [sqrt(eps(eltype(t))), zero(eltype(t))] + #TODO ε is size of t ? # ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end @@ -106,27 +127,30 @@ function inner_physics_loss(phi::PINOPhi{C, T, U}, θ, ts::AbstractArray, prob::ODEProblem, - isu0::Bool) where {C, T, U} + isu0::Bool, + in_) where {C, T, U} u0 = prob.u0 p = prob.p f = prob.f - if isu0 == true - in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) - #TODO for all case p and u0 - # u0 isa Vector - # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) - else - if p isa Number - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) - elseif p isa Vector - #TODO nno for Vector - inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) - in_ = reshape(inner, size(inner)..., 1) - else - error("p should be a number or a vector") - end - end + # if isu0 == true + # #TODO data should be generate before train ? + # in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) + # #TODO for all case p and u0 + # # u0 isa Vector + # # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) + # else + # if p isa Number + # in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) + # elseif p isa Vector + # #TODO nno for Vector + # inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) + # in_ = reshape(inner, size(inner)..., 1) + # else + # error("p should be a number or a vector") + # end + # end out_ = phi(in_, θ) + # fs = f.f.(out_, p, ts) if p isa Number fs = f.f.(out_, p, ts) elseif p isa Vector @@ -137,15 +161,23 @@ function inner_physics_loss(phi::PINOPhi{C, T, U}, NeuralOperators.l₂loss(dfdx(phi, in_, θ), fs) end + function physics_loss(phi::PINOPhi{C, T, U}, θ, ts::AbstractArray, - train_set::TRAINSET) where {C, T, U} + train_set::TRAINSET, + input_data_set) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data - norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] - loss = reduce(vcat, - [inner_physics_loss(phi, θ, ts, prob, train_set.isu0) for prob in prob_set]) - sum(abs2, loss) / norm + f = prob_set[1].f + # norm = prod(size(output_data)) + # norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] + # loss = reduce(vcat, + # [inner_physics_loss(phi, θ, ts, prob, train_set.isu0, in_) + # for (in_, prob) in zip(inputdata, prob_set)]) + # sum(abs2, loss) / norm + ps = [prob.p for prob in prob_set]' + fs = f.f.(output_data, ps, ts) + loss = NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) end function inner_data_loss(phi::PINOPhi{C, T, U}, @@ -153,7 +185,8 @@ function inner_data_loss(phi::PINOPhi{C, T, U}, ts::AbstractArray, prob::ODEProblem, out_::AbstractArray, - isu0::Bool) where {C, T, U} + isu0::Bool, + in_) where {C, T, U} u0 = prob.u0 p = prob.p f = prob.f @@ -164,7 +197,7 @@ function inner_data_loss(phi::PINOPhi{C, T, U}, # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) else if p isa Number - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2])]) + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2],1)]) elseif p isa Vector inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) in_ = reshape(inner, size(inner)..., 1) @@ -178,23 +211,98 @@ end function data_loss(phi::PINOPhi{C, T, U}, θ, ts::AbstractArray, - train_set::TRAINSET) where {C, T, U} + train_set::TRAINSET, + input_data_set) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data - norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] - loss = reduce(vcat, - [inner_data_loss(phi, θ, ts, prob, out_, train_set.isu0) - for (prob, out_) in zip(prob_set, output_data)]) - sum(abs2, loss) / norm + # norm = prod(size(output_data)) + # norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] + # loss = reduce(vcat, + # [inner_data_loss(phi, θ, ts, prob, out_, train_set.isu0, in_) + # for (prob, out_, in_) in zip(prob_set, output_data, input_data_set)]) + # sum(abs2, loss) / norm + loss = NeuralOperators.l₂loss(phi(input_data_set, θ), output_data) +end + +function generate_data(ts, prob_set, isu0) + input_data_set = [] + input_data_set_right = [] + for prob in prob_set + u0 = prob.u0 + p = prob.p + f = prob.f + ε = sqrt(eps(eltype(ts))) + tsr = ts .+ ε + if isu0 == true + #TODO data should be generate before train ? + in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) + + #TODO for all case p and u0 + # u0 isa Vector + # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) + else + if p isa Number + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) + in_r = reduce(vcat, [tsr, fill(p, 1, size(ts)[2], 1)]) + + elseif p isa Vector + #TODO nno for Vector + inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) + in_ = reshape(inner, size(inner)..., 1) + else + error("p should be a number or a vector") + end + end + push!(input_data_set, in_) + push!(input_data_set_right, in_r) + end + input_data_set, input_data_set_right +end + +function generate_data_matrix(ts, prob_set, isu0) + + batch_size = size(prob_set)[1] + instances_size = size(ts)[2] + dims = 2 + input_data_set = Array{Float32, 3}(undef, dims, instances_size, batch_size) + for (i,prob) in enumerate(prob_set) + u0 = prob.u0 + p = prob.p + f = prob.f + if isu0 == true + #TODO data should be generate before train ? + in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) + + #TODO for all case p and u0 + # u0 isa Vector + # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) + else + if p isa Number + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) + elseif p isa Vector + #TODO nno for Vector + inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) + in_ = reshape(inner, size(inner)..., 1) + else + error("p should be a number or a vector") + end + end + input_data_set[:, :, i] = in_ + end + input_data_set end function generate_loss(phi::PINOPhi{C, T, U}, train_set::TRAINSET, tspan) where {C, T, U} t0 = tspan[1] t_end = tspan[2] - instances_size = size(train_set.output_data[1])[2] + instances_size = size(train_set.output_data)[2] + # instances_size = size(train_set.output_data[1])[2] range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) + prob_set, output_data = train_set.input_data, train_set.output_data + input_data_set = generate_data_matrix(ts, prob_set, train_set.isu0) function loss(θ, _) - data_loss(phi, θ, ts, train_set) + physics_loss(phi, θ, ts, train_set) + data_loss(phi, θ, ts, train_set, input_data_set) + + physics_loss(phi, θ, ts, train_set, input_data_set) end return loss end @@ -212,7 +320,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, t0 = tspan[1] # f = prob.f # p = prob.p - # u0 = prob.u0 + u0 = prob.u0 # param_estim = alg.param_estim chain = alg.chain @@ -226,7 +334,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - phi, init_params = generate_pino_phi_θ(chain, t0, init_params) + phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) init_params = ComponentArrays.ComponentArray(init_params) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 10f0ffde70..5ac9f100f2 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,6 +1,6 @@ using Test using OrdinaryDiffEq, OptimizationOptimisers -using Flux, Lux +using Lux using Statistics, Random using NeuralOperators using NeuralPDE @@ -8,25 +8,34 @@ using NeuralPDE @testset "Example 1" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) - tspan = (0.0, 2.0) - u0 = 0.0 + tspan = (0.0f0, 2.0f0) + u0 = 0.0f0 #generate data set t0, t_end = tspan - instances_size = 40 + instances_size = 50 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) - batch_size = 40 + batch_size = 50 as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] - u_output_ = Array{Float32, 3}[] + u_output_ = Array{Float32, 3}(undef, 1, instances_size, batch_size) prob_set = [] - for a_i in as + for (i, a_i) in enumerate(as) prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) sol1 = solve(prob, Tsit5(); saveat = 0.0204) reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) push!(prob_set, prob) - push!(u_output_, reshape_sol) + u_output_[:, :, i] = reshape_sol end + # u_output_ = Array{Float32, 3}[] + # prob_set = [] + # for a_i in as + # prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) + # sol1 = solve(prob, Tsit5(); saveat = 0.0204) + # reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) + # push!(prob_set, prob) + # push!(u_output_, reshape_sol) + # end """ Set of training data: @@ -36,15 +45,24 @@ using NeuralPDE train_set = NeuralPDE.TRAINSET(prob_set, u_output_); #TODO u0 ? prob = ODEProblem(linear, u0, tspan, 0) - chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)); + chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), - σ = gelu); + σ = gelu) + # flat_no(rand(2, 100, 1)) + # Random.default_rng() + # luxm = Lux.transform(flat_no) + # θ, st = Lux.setup(Random.default_rng(), luxm) + # luxm(rand(Float32, 2, 40, 1), θ, st)[1] + # pk(c, θ) = luxm(rand(2, 40, 1), θ, st)[1] + # Zygote.gradient(θ -> sum(abs2, pk(rand(2, 100, 1), θ)), θ) + # NeuralOperators.l₂loss(pk(rand(2, 100, 1), θ), rand(1,100,1)) + opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, train_set); + alg = NeuralPDE.PINOODE(chain, opt, train_set) res, phi = solve(prob, alg, verbose = true, - maxiters = 1000, abstol = 1.0f-10) + maxiters = 400, abstol = 1.0f-10) predict = reduce(vcat, [phi( @@ -53,10 +71,22 @@ using NeuralPDE res.u) for i in 1:batch_size]) ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) - @test ground≈predict atol=0.5 + @test ground≈predict atol=1 +end + + +function plot_() + # Animate + anim = @animate for (i) in 1:batch_size + plot(predict[i, :], label = "Predicted") + plot!(ground[i, :], label = "Ground truth") + end + gif(anim, "pino.gif", fps = 10) end -@testset "Example 2" begin +plot_() + +"Example 2" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0, 2.0) @@ -90,7 +120,7 @@ end flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu); opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, train_set); + alg = NeuralPDE.PINOODE(chain, opt, train_set) res, phi = solve(prob, alg, verbose = true, maxiters = 1000, abstol = 1.0f-10) From db50090cba9c6b3743e8026574d12c07ab99193f Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 6 Mar 2024 18:32:33 +0400 Subject: [PATCH 011/153] clear up code --- src/pino_ode_solve.jl | 186 +++++++---------------------------------- test/PINO_ode_tests.jl | 67 ++++++++------- 2 files changed, 61 insertions(+), 192 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index e383bfec7e..0d34c6a6e5 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -14,11 +14,10 @@ which thus uses the random initialization provided by the neural network library. ## Keyword Arguments -* `minibatch`: TODO +* `minibatch`: ## Examples -TODO ```julia ``` @@ -26,8 +25,7 @@ TODO ## References Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -#TODO -struct TRAINSET{} #T +struct TRAINSET{} #TODO #T <: Number input_data::Vector{ODEProblem} output_data::Array isu0::Bool @@ -53,7 +51,6 @@ function PINOODE(chain, minibatch = 0, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - #TODO transform convert complex numbers to zero PINOODE(chain, opt, train_set, init_params, minibatch, kwargs) end @@ -67,7 +64,7 @@ mutable struct PINOPhi{C, T, U, S} u0::U st::S function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) - new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0,u0, st) + new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) end end @@ -84,128 +81,45 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, PINOPhi(chain, t0, u0, st), init_params end -# function (f::PINOPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer, T, U} -# y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), [t]), θ, f.st) -# ChainRulesCore.@ignore_derivatives f.st = st -# first(y) -# end - function (f::PINOPhi{C, T, U})(t::AbstractArray, θ) where {C <: Lux.AbstractExplicitLayer, T, U} # Batch via data as row vectors - # y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) - y, st = f.chain(t, θ, f.st) + y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - # y - f.u0 .+ (t[[1],:,:] .- f.t0) .* y + f.u0 .+ (t[[1], :, :] .- f.t0) .* y end -# feature_dims = 2:(ndims(t) - 1) -# loss = sum( t, dims = feature_dims) - # loss = sum(.√(sum(abs2, 𝐲̂ - 𝐲, dims = feature_dims))) - # y_norm = sum(.√(sum(abs2, 𝐲, dims = feature_dims))) - - # return loss / y_norm -# function dfdx(phi::PINOPhi, t::AbstractArray, θ) -# ε = sqrt(eps(eltype(t))) -# εs = [ε, zero(eltype(t))] -# # ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] -# # ChainRulesCore.@ignore_derivatives tl = t .+ ε -# tl = t .+ ε -# tr = t -# (phi(tl, θ) - phi(tr, θ)) ./ ε -# end +function dfdx_rand_matrix(phi::PINOPhi, t::AbstractArray, θ) + ε_ = sqrt(eps(eltype(t))) + d = Normal{eltype(t)}(0.0f0, ε_) + size_ = size(t) .- (1, 0, 0) + eps_ = ε_ .+ rand(d, size_) .* ε_ + zeros_ = zeros(eltype(t), size_) + ε = cat(eps_, zeros_, dims = 1) + (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) +end function dfdx(phi::PINOPhi, t::AbstractArray, θ) ε = [sqrt(eps(eltype(t))), zero(eltype(t))] - #TODO ε is size of t ? # ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end -function inner_physics_loss(phi::PINOPhi{C, T, U}, - θ, - ts::AbstractArray, - prob::ODEProblem, - isu0::Bool, - in_) where {C, T, U} - u0 = prob.u0 - p = prob.p - f = prob.f - # if isu0 == true - # #TODO data should be generate before train ? - # in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) - # #TODO for all case p and u0 - # # u0 isa Vector - # # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) - # else - # if p isa Number - # in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) - # elseif p isa Vector - # #TODO nno for Vector - # inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) - # in_ = reshape(inner, size(inner)..., 1) - # else - # error("p should be a number or a vector") - # end - # end - out_ = phi(in_, θ) - # fs = f.f.(out_, p, ts) - if p isa Number - fs = f.f.(out_, p, ts) - elseif p isa Vector - fs = reduce(hcat, [f.f(out_[:, i], p, ts[i]) for i in 1:size(out_, 2)]) - else - error("p should be a number or a vector") - end - NeuralOperators.l₂loss(dfdx(phi, in_, θ), fs) -end - - function physics_loss(phi::PINOPhi{C, T, U}, θ, ts::AbstractArray, train_set::TRAINSET, input_data_set) where {C, T, U} - prob_set, output_data = train_set.input_data, train_set.output_data - f = prob_set[1].f - # norm = prod(size(output_data)) - # norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] - # loss = reduce(vcat, - # [inner_physics_loss(phi, θ, ts, prob, train_set.isu0, in_) - # for (in_, prob) in zip(inputdata, prob_set)]) - # sum(abs2, loss) / norm - ps = [prob.p for prob in prob_set]' - fs = f.f.(output_data, ps, ts) - loss = NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) -end - -function inner_data_loss(phi::PINOPhi{C, T, U}, - θ, - ts::AbstractArray, - prob::ODEProblem, - out_::AbstractArray, - isu0::Bool, - in_) where {C, T, U} - u0 = prob.u0 - p = prob.p - f = prob.f - if isu0 == true - in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) - #TODO for all case p and u0 - # u0 isa Vector - # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) + prob_set, output_data = train_set.input_data, train_set.output_data #TODO + f = prob_set[1].f #TODO one f for all + out_ = phi(input_data_set, θ) + if train_set.isu0 === false + ps = [prob.p for prob in prob_set] #TODO do it within generator for data else - if p isa Number - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2],1)]) - elseif p isa Vector - inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) - in_ = reshape(inner, size(inner)..., 1) - else - error("p should be a number or a vector") - end + error("WIP") end - NeuralOperators.l₂loss(phi(in_, θ), out_) + fs = cat([f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3) + NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) end function data_loss(phi::PINOPhi{C, T, U}, @@ -214,64 +128,20 @@ function data_loss(phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data - # norm = prod(size(output_data)) - # norm = size(output_data)[1] * size(output_data[1])[2] * size(output_data[1])[1] - # loss = reduce(vcat, - # [inner_data_loss(phi, θ, ts, prob, out_, train_set.isu0, in_) - # for (prob, out_, in_) in zip(prob_set, output_data, input_data_set)]) - # sum(abs2, loss) / norm - loss = NeuralOperators.l₂loss(phi(input_data_set, θ), output_data) + NeuralOperators.l₂loss(phi(input_data_set, θ), output_data) end function generate_data(ts, prob_set, isu0) - input_data_set = [] - input_data_set_right = [] - for prob in prob_set - u0 = prob.u0 - p = prob.p - f = prob.f - ε = sqrt(eps(eltype(ts))) - tsr = ts .+ ε - if isu0 == true - #TODO data should be generate before train ? - in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) - - #TODO for all case p and u0 - # u0 isa Vector - # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) - else - if p isa Number - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) - in_r = reduce(vcat, [tsr, fill(p, 1, size(ts)[2], 1)]) - - elseif p isa Vector - #TODO nno for Vector - inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) - in_ = reshape(inner, size(inner)..., 1) - else - error("p should be a number or a vector") - end - end - push!(input_data_set, in_) - push!(input_data_set_right, in_r) - end - input_data_set, input_data_set_right -end - -function generate_data_matrix(ts, prob_set, isu0) - batch_size = size(prob_set)[1] instances_size = size(ts)[2] dims = 2 input_data_set = Array{Float32, 3}(undef, dims, instances_size, batch_size) - for (i,prob) in enumerate(prob_set) + for (i, prob) in enumerate(prob_set) u0 = prob.u0 p = prob.p f = prob.f if isu0 == true - #TODO data should be generate before train ? in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) - #TODO for all case p and u0 # u0 isa Vector # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) @@ -295,11 +165,11 @@ function generate_loss(phi::PINOPhi{C, T, U}, train_set::TRAINSET, tspan) where t0 = tspan[1] t_end = tspan[2] instances_size = size(train_set.output_data)[2] - # instances_size = size(train_set.output_data[1])[2] range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) - prob_set, output_data = train_set.input_data, train_set.output_data - input_data_set = generate_data_matrix(ts, prob_set, train_set.isu0) + + prob_set, output_data = train_set.input_data, train_set.output_data #TODO one format data + input_data_set = generate_data(ts, prob_set, train_set.isu0) function loss(θ, _) data_loss(phi, θ, ts, train_set, input_data_set) + physics_loss(phi, θ, ts, train_set, input_data_set) @@ -328,7 +198,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, init_params = alg.init_params # mapping between functional space of some vararible 'a' of equation (for example initial - # condition {u(t0 x)} or parameter p) join and solution of equation u(t) + # condition {u(t0 x)} or parameter p) and solution of equation u(t) train_set = alg.train_set !(chain isa Lux.AbstractExplicitLayer) && diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 5ac9f100f2..cb1de3ecd9 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -5,7 +5,7 @@ using Statistics, Random using NeuralOperators using NeuralPDE -@testset "Example 1" begin +@testset "Example p" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) @@ -18,7 +18,7 @@ using NeuralPDE batch_size = 50 as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] - u_output_ = Array{Float32, 3}(undef, 1, instances_size, batch_size) + u_output_ = zeros(Float32, 1, instances_size, batch_size) prob_set = [] for (i, a_i) in enumerate(as) prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) @@ -27,15 +27,6 @@ using NeuralPDE push!(prob_set, prob) u_output_[:, :, i] = reshape_sol end - # u_output_ = Array{Float32, 3}[] - # prob_set = [] - # for a_i in as - # prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) - # sol1 = solve(prob, Tsit5(); saveat = 0.0204) - # reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - # push!(prob_set, prob) - # push!(u_output_, reshape_sol) - # end """ Set of training data: @@ -45,48 +36,46 @@ using NeuralPDE train_set = NeuralPDE.TRAINSET(prob_set, u_output_); #TODO u0 ? prob = ODEProblem(linear, u0, tspan, 0) - chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) + chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 1)) flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) - # flat_no(rand(2, 100, 1)) - # Random.default_rng() - # luxm = Lux.transform(flat_no) - # θ, st = Lux.setup(Random.default_rng(), luxm) - # luxm(rand(Float32, 2, 40, 1), θ, st)[1] - # pk(c, θ) = luxm(rand(2, 40, 1), θ, st)[1] - # Zygote.gradient(θ -> sum(abs2, pk(rand(2, 100, 1), θ)), θ) - # NeuralOperators.l₂loss(pk(rand(2, 100, 1), θ), rand(1,100,1)) - + η₀ = 1.0f-2 opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, train_set) + alg = NeuralPDE.PINOODE(flat_no, opt, train_set) res, phi = solve(prob, alg, verbose = true, - maxiters = 400, abstol = 1.0f-10) + maxiters = 500) + + input_data_set_2 = Array{Float32, 3}(undef, 2, instances_size, batch_size) + for (i, prob) in enumerate(prob_set) + in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) + input_data_set_2[:, :, i] = in_ + end + predict_2 = phi(input_data_set_2, res.u) + predict = phi(input_data_set, res.u) + ground = output_data - predict = reduce(vcat, - [phi( - reshape(reduce(vcat, [ts, fill(train_set.input_data[i].p, 1, size(ts)[2])]), - 2, instances_size, 1), - res.u) - for i in 1:batch_size]) - ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) @test ground≈predict atol=1 end - function plot_() # Animate anim = @animate for (i) in 1:batch_size - plot(predict[i, :], label = "Predicted") - plot!(ground[i, :], label = "Ground truth") + plot(predict[1, :, i], label = "Predicted") + plot!(ground[1, :,i], label = "Ground truth") end gif(anim, "pino.gif", fps = 10) end plot_() -"Example 2" begin +"Example u0" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0, 2.0) @@ -108,6 +97,16 @@ plot_() push!(prob_set, prob) end + u_output_ = zeros(Float32, 1, instances_size, batch_size) + prob_set = [] + for (i, u0_i) in enumerate(u0s) + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) + sol1 = solve(prob, Tsit5(); saveat = 0.0204) + reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) + push!(prob_set, prob) + u_output_[:, :, i] = reshape_sol + end + """ Set of training data: * input data: set of initial conditions 'a': From 60cda5315defe26fe2a469dd6d7858b21e11e829 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 7 Mar 2024 19:46:55 +0400 Subject: [PATCH 012/153] update tests --- src/pino_ode_solve.jl | 65 ++++++++++----- test/PINO_ode_tests.jl | 174 +++++++++++++++++++---------------------- 2 files changed, 126 insertions(+), 113 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 0d34c6a6e5..010da00b60 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -40,6 +40,7 @@ struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm opt::O train_set::TRAINSET init_params::P + #TODO remove minibatch for a while minibatch::Int kwargs::K end @@ -50,6 +51,7 @@ function PINOODE(chain, init_params = nothing; minibatch = 0, kwargs...) + #TODO fnn trasform !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) PINOODE(chain, opt, train_set, init_params, minibatch, kwargs) end @@ -100,8 +102,8 @@ function dfdx_rand_matrix(phi::PINOPhi, t::AbstractArray, θ) end function dfdx(phi::PINOPhi, t::AbstractArray, θ) - ε = [sqrt(eps(eltype(t))), zero(eltype(t))] - # ε = [sqrt(eps(eltype(t))), zeros(eltype(t), phi.chain.layers.layer_1.in_dims - 1)...] + # ε = [sqrt(eps(eltype(t))), zero(eltype(t))] + ε = [sqrt(eps(eltype(t))), zeros(eltype(t), size(t)[1] - 1)...] (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end @@ -112,19 +114,31 @@ function physics_loss(phi::PINOPhi{C, T, U}, input_data_set) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data #TODO f = prob_set[1].f #TODO one f for all + p = prob_set[1].p out_ = phi(input_data_set, θ) - if train_set.isu0 === false - ps = [prob.p for prob in prob_set] #TODO do it within generator for data + if train_set.isu0 == true + p = prob_set[1].p + fs = f.f.(out_, p, ts) else - error("WIP") + ps = [prob.p for prob in prob_set] + if p isa Number + fs = cat( + [f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3) + else + #TODO generalize + fs = cat( + [reduce( + hcat, [f.f(out_[:, j, [i]], p, ts) for j in axes(out_[:, :, [i]], 2)]) + for (i, p) in enumerate(ps)]..., + dims = 3) + end end - fs = cat([f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3) NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) end function data_loss(phi::PINOPhi{C, T, U}, θ, - ts::AbstractArray, + ts::AbstractArray, #remove unessasry train_set::TRAINSET, input_data_set) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data @@ -134,7 +148,7 @@ end function generate_data(ts, prob_set, isu0) batch_size = size(prob_set)[1] instances_size = size(ts)[2] - dims = 2 + dims = isu0 ? length(prob_set[1].u0) + 1 : length(prob_set[1].p) + 1 input_data_set = Array{Float32, 3}(undef, dims, instances_size, batch_size) for (i, prob) in enumerate(prob_set) u0 = prob.u0 @@ -161,17 +175,10 @@ function generate_data(ts, prob_set, isu0) input_data_set end -function generate_loss(phi::PINOPhi{C, T, U}, train_set::TRAINSET, tspan) where {C, T, U} - t0 = tspan[1] - t_end = tspan[2] - instances_size = size(train_set.output_data)[2] - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - - prob_set, output_data = train_set.input_data, train_set.output_data #TODO one format data - input_data_set = generate_data(ts, prob_set, train_set.isu0) +function generate_loss( + phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts) where {C, T, U} function loss(θ, _) - data_loss(phi, θ, ts, train_set, input_data_set) + + data_loss(phi, θ, ts, train_set, input_data_set) + #TODO train_set + input_data_set physics_loss(phi, θ, ts, train_set, input_data_set) end return loss @@ -204,7 +211,23 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) + #TODO format data + t0 = tspan[1] + t_end = tspan[2] + instances_size = size(train_set.output_data)[2] + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) + prob_set, output_data = train_set.input_data, train_set.output_data + isu0 = train_set.isu0 + input_data_set = generate_data(ts, prob_set, isu0) + + if isu0 + u0 = input_data_set[[2], :, :] + phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) + else + u0 = prob.u0 + phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) + end init_params = ComponentArrays.ComponentArray(init_params) @@ -212,7 +235,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - # phi(rand(5, 100, 1), init_params) #TODO input data + phi(input_data_set, init_params) catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) @@ -221,7 +244,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end end - total_loss = generate_loss(phi, train_set, tspan) + total_loss = generate_loss(phi, train_set, input_data_set, ts) # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index cb1de3ecd9..83d06f3689 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -33,70 +33,46 @@ using NeuralPDE * input data: set of parameters 'a': * output data: set of solutions u(t){a} corresponding parameter 'a' """ - train_set = NeuralPDE.TRAINSET(prob_set, u_output_); +train_set = TRAINSET(prob_set, u_output_); #TODO u0 ? prob = ODEProblem(linear, u0, tspan, 0) chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 1)) + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 1)) flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) η₀ = 1.0f-2 opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(flat_no, opt, train_set) + alg = PINOODE(flat_no, opt, train_set) + res, phi = solve(prob, alg, verbose = true, maxiters = 200) - res, phi = solve(prob, - alg, verbose = true, - maxiters = 500) - - input_data_set_2 = Array{Float32, 3}(undef, 2, instances_size, batch_size) + input_data_set = Array{Float32, 3}(undef, 2, instances_size, batch_size) for (i, prob) in enumerate(prob_set) - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) - input_data_set_2[:, :, i] = in_ + in_ = reduce(vcat, [ts, fill(prob.p, 1, size(ts)[2], 1)]) + input_data_set[:, :, i] = in_ end - predict_2 = phi(input_data_set_2, res.u) predict = phi(input_data_set, res.u) - ground = output_data - + ground = u_output_ @test ground≈predict atol=1 end -function plot_() - # Animate - anim = @animate for (i) in 1:batch_size - plot(predict[1, :, i], label = "Predicted") - plot!(ground[1, :,i], label = "Ground truth") - end - gif(anim, "pino.gif", fps = 10) -end - -plot_() - -"Example u0" begin +"Example u0" +begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) - tspan = (0.0, 2.0) - p = pi - u0 = 2 + tspan = (0.0f0, 2.0f0) + p = Float32(pi) + u0 = 2.0f0 #generate data set t0, t_end = tspan - instances_size = 100 + instances_size = 50 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) batch_size = 50 - u0s = [i for i in range(0.0, stop = pi / 2, length = batch_size)] - prob_set = [] - u_output_ = Array{Float32, 3}[] - for u0_i in u0s - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) - sol = reshape([prob.f.analytic(u0_i, p, t) for t in ts], 1, instances_size, 1) - push!(u_output_, sol) - push!(prob_set, prob) - end - + u0s = [i for i in range(0.0f0, stop = pi / 2.0f0, length = batch_size)] u_output_ = zeros(Float32, 1, instances_size, batch_size) prob_set = [] for (i, u0_i) in enumerate(u0s) @@ -109,29 +85,29 @@ plot_() """ Set of training data: - * input data: set of initial conditions 'a': - * output data: set of solutions u(t){a} corresponding initial conditions 'a' + * input data: set of initial conditions 'u0': + * output data: set of solutions u(t){u0} corresponding initial conditions 'u0' """ - train_set = NeuralPDE.TRAINSET(prob_set, u_output_; isu0 = true); + train_set = TRAINSET(prob_set, u_output_; isu0 = true) #TODO u0 ? - prob = ODEProblem(linear, 0, tspan, p) - chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)); - flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), - σ = gelu); - opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, train_set) + prob = ODEProblem(linear, 0.0f0, tspan, p) + # chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) + fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), + σ = gelu) + opt = OptimizationOptimisers.Adam(0.001) + alg = PINOODE(fno, opt, train_set) res, phi = solve(prob, alg, verbose = true, - maxiters = 1000, abstol = 1.0f-10) - - predict = reduce(vcat, - [phi( - reshape(reduce(vcat, [ts, fill(train_set.input_data[i].u0, 1, size(ts)[2])]), - 2, instances_size, 1), - res.u) - for i in 1:batch_size]) - ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) - @test ground≈predict atol=1. + maxiters = 200) + + input_data_set = Array{Float32, 3}(undef, 2, instances_size, batch_size) + for (i, prob) in enumerate(prob_set) + in_ = reduce(vcat, [ts, fill(prob.u0, 1, size(ts)[2], 1)]) + input_data_set[:, :, i] = in_ + end + predict = phi(input_data_set, res.u) + ground = u_output_ + @test ground≈predict atol=1.0 end @testset "lotka volterra" begin @@ -145,47 +121,61 @@ end dy = (δ * x - γ) * y # predator return [dx, dy] end - u0 = [1.0, 1.0] - p = [1.5, 1.0, 3.0, 1.0] - tspan = (0.0, 4.0) - dt = 0.01 - instances_size = 50 + u0 = [1.0f0, 1.0f0] + p = Float32[1.5, 1.0, 3.0, 1.0] + tspan = (0.0f0, 4.0f0) + dt = 0.01f0 + + instances_size = 100 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) batch_size = 50 - ps = [p .+ i * [0.015, 0.01, 0.03, 0.01] for i in 1:batch_size] - - u_output_ = [] + ps = [p .+ i * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size] + u_output_ = zeros(Float32, 2, instances_size, batch_size) prob_set = [] - plot_set = [] - for p_i in ps + for (i, p_i) in enumerate(ps) prob = ODEProblem(lotka_volterra, u0, tspan, p_i) solution = solve(prob, Tsit5(); saveat = dt) - reshape_sol = Float32.(reduce(hcat, solution(range_).u)) - push!(plot_set, solution) + reshape_sol_ = reduce(hcat, solution(range_).u) + reshape_sol = Float32.(reshape(reshape_sol_, 2, instances_size, 1)) push!(prob_set, prob) - push!(u_output_, reshape_sol) + u_output_[:, :, i] = reshape_sol end - train_set = NeuralPDE.TRAINSET(prob_set, u_output_) + train_set = TRAINSET(prob_set, u_output_); #TODO u0 ? - prob = ODEProblem(linear, u0, tspan, p) - chain = Lux.Chain(Lux.Dense(5, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)); + prob = ODEProblem(lotka_volterra_matrix, u0, tspan, p) + chain = Lux.Chain(Lux.Dense(5, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)) flat_no = FourierNeuralOperator(ch = (5, 16, 16, 16, 16, 16, 32, 2), modes = (16,), - σ = gelu); - opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(chain, opt, train_set) - res, phi = solve(prob, - alg, verbose = true, - maxiters = 1500, abstol = 1.0f-10) - - predict = reduce(vcat, - [phi( - reduce(vcat, - [ts, reduce(hcat, fill(train_set.input_data[i].p, 1, size(ts)[2]))]), - res.u) - for i in 1:batch_size]) - ground = reduce(vcat, [train_set.output_data[i] for i in 1:batch_size]) + σ = gelu) + opt = OptimizationOptimisers.Adam(0.01) + alg = PINOODE(flat_no, opt, train_set); + res, phi = solve(prob, alg, verbose = true, maxiters = 300) + + input_data_set = Array{Float32, 3}(undef, 5, instances_size, batch_size) + for (i, prob) in enumerate(prob_set) + inner = reduce(vcat, [ts, reduce(hcat, fill(prob.p, 1, size(ts)[2], 1))]) + in_ = reshape(inner, size(inner)..., 1) + input_data_set[:, :, i] = in_ + end + + predict = phi(input_data_set, res.u) + ground = u_output_ + @test ground≈predict atol=5 end + +using Plots +function plot_() + # Animate + anim = @animate for (i) in 1:batch_size + plot(predict[1, :, i], label = "Predicted") + plot!(predict[2, :, i], label = "Predicted") + plot!(ground[1, :, i], label = "Ground truth") + plot!(ground[2, :, i], label = "Ground truth") + end + gif(anim, "pino.gif", fps = 10) +end + +plot_() From 5781cb5f2a373bbb0ceed53c1431ae0e6b994828 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 8 Mar 2024 17:02:36 +0400 Subject: [PATCH 013/153] gpu --- .buildkite/pipeline.yml | 1 + .github/workflows/CI.yml | 1 + src/pino_ode_solve.jl | 15 ++++---- test/PINO_ode_tests.jl | 16 +-------- test/PINO_ode_tests_gpu.jl | 70 ++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 5 +++ 6 files changed, 86 insertions(+), 22 deletions(-) create mode 100644 test/PINO_ode_tests_gpu.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index ff287e15bd..b894d6041c 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -25,6 +25,7 @@ steps: - "NNODE" - "NeuralAdapter" - "IntegroDiff" + - "PINO_GPU" env: BUILDKITE_PLUGIN_JULIA_VERSION: "{{matrix.version}}" GROUP: "{{matrix.group}}" diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 59f046cc9f..18aa1ac2aa 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -27,6 +27,7 @@ jobs: - Forward - NeuralAdapter - DGM + - ODEPINO version: - "1" steps: diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 010da00b60..d4570cd81e 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -40,8 +40,6 @@ struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm opt::O train_set::TRAINSET init_params::P - #TODO remove minibatch for a while - minibatch::Int kwargs::K end @@ -49,11 +47,10 @@ function PINOODE(chain, opt, train_set, init_params = nothing; - minibatch = 0, kwargs...) #TODO fnn trasform !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, train_set, init_params, minibatch, kwargs) + PINOODE(chain, opt, train_set, init_params, kwargs) end """ @@ -88,7 +85,8 @@ function (f::PINOPhi{C, T, U})(t::AbstractArray, # Batch via data as row vectors y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - f.u0 .+ (t[[1], :, :] .- f.t0) .* y + ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[[1], :, :]) + f.u0 .+ (ts .- f.t0) .* y end function dfdx_rand_matrix(phi::PINOPhi, t::AbstractArray, θ) @@ -116,6 +114,7 @@ function physics_loss(phi::PINOPhi{C, T, U}, f = prob_set[1].f #TODO one f for all p = prob_set[1].p out_ = phi(input_data_set, θ) + ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), ts) if train_set.isu0 == true p = prob_set[1].p fs = f.f.(out_, p, ts) @@ -133,6 +132,7 @@ function physics_loss(phi::PINOPhi{C, T, U}, dims = 3) end end + # fs = adapt(parameterless_type(ComponentArrays.getdata(θ)), fs) NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) end @@ -142,6 +142,7 @@ function data_loss(phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data + output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data) NeuralOperators.l₂loss(phi(input_data_set, θ), output_data) end @@ -222,7 +223,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, input_data_set = generate_data(ts, prob_set, isu0) if isu0 - u0 = input_data_set[[2], :, :] + u0 = input_data_set[2:end, :, :] phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) else u0 = prob.u0 @@ -238,7 +239,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, phi(input_data_set, init_params) catch err if isa(err, DimensionMismatch) - throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) + throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) #TODO change message else throw(err) end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 83d06f3689..07f0691460 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -151,7 +151,7 @@ end σ = gelu) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(flat_no, opt, train_set); - res, phi = solve(prob, alg, verbose = true, maxiters = 300) + res, phi = solve(prob, alg, verbose = true, maxiters = 200) input_data_set = Array{Float32, 3}(undef, 5, instances_size, batch_size) for (i, prob) in enumerate(prob_set) @@ -165,17 +165,3 @@ end @test ground≈predict atol=5 end - -using Plots -function plot_() - # Animate - anim = @animate for (i) in 1:batch_size - plot(predict[1, :, i], label = "Predicted") - plot!(predict[2, :, i], label = "Predicted") - plot!(ground[1, :, i], label = "Ground truth") - plot!(ground[2, :, i], label = "Ground truth") - end - gif(anim, "pino.gif", fps = 10) -end - -plot_() diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl new file mode 100644 index 0000000000..47b6bf75d1 --- /dev/null +++ b/test/PINO_ode_tests_gpu.jl @@ -0,0 +1,70 @@ +using Test +using OrdinaryDiffEq +using Lux +using ComponentArrays +using NeuralOperators +using OptimizationOptimisers +using Random +using LuxCUDA +using NeuralPDE + +CUDA.allowscalar(false) +const gpud = gpu_device() + +@testset "Example p" begin + linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + linear = (u, p, t) -> cos(p * t) + tspan = (0.0f0, 2.0f0) + u0 = 0.0f0 + #generate data set + t0, t_end = tspan + instances_size = 50 + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) + batch_size = 50 + as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] + + u_output_ = zeros(Float32, 1, instances_size, batch_size) + prob_set = [] + for (i, a_i) in enumerate(as) + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) + sol1 = solve(prob, Tsit5(); saveat = 0.0204) + reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) + push!(prob_set, prob) + u_output_[:, :, i] = reshape_sol + end + + """ + Set of training data: + * input data: set of parameters 'a': + * output data: set of solutions u(t){a} corresponding parameter 'a' + """ + train_set = NeuralPDE.TRAINSET(prob_set, u_output_) + + #TODO u0 ? + prob = ODEProblem(linear, u0, tspan, 0) + inner = 20 + chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, 1)) + ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud + + flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), + σ = gelu) + + opt = OptimizationOptimisers.Adam(0.03) + alg = NeuralPDE.PINOODE(flat_no, opt, train_set, ps) + res, phi = solve(prob, alg, verbose = true, maxiters = 200) + + input_data_set = Array{Float32, 3}(undef, 2, instances_size, batch_size) + for (i, prob) in enumerate(prob_set) + in_ = reduce(vcat, [ts, fill(prob.p, 1, size(ts)[2], 1)]) + input_data_set[:, :, i] = in_ + end + predict = phi(input_data_set, res.u) |> cpu + ground = u_output_ + @test ground≈predict atol=1 +end diff --git a/test/runtests.jl b/test/runtests.jl index aec012166a..52b6e875fe 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -69,6 +69,11 @@ end if !is_APPVEYOR && GROUP == "GPU" @safetestset "NNPDE_gpu_Lux" begin include("NNPDE_tests_gpu_Lux.jl") end end + if !is_APPVEYOR && GROUP == "PINO_GPU" + @safetestset "PINo ode gpu" begin include("PINO_ode_tests_gpu.jl.jl") + end + end + if GROUP == "All" || GROUP == "DGM" @time @safetestset "Deep Galerkin solver" begin include("dgm_test.jl") end From b562ebea7158172f79b3d3ac509052193c07a8d8 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 8 Mar 2024 17:06:50 +0400 Subject: [PATCH 014/153] fix spelling --- docs/pages.jl | 2 +- docs/src/tutorials/pino_ode.md | 2 +- src/pino_ode_solve.jl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/pages.jl b/docs/pages.jl index 06b1eaf80c..e2586bfa63 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -3,7 +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 Opeator ODEs" => "tutorials/pino_ode.md", + "Physics informed Neural Operator ODEs" => "tutorials/pino_ode.md", "Deep Galerkin Method" => "tutorials/dgm.md" #"examples/nnrode_example.md", # currently incorrect ], diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 3df6586105..d7ca3102b0 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -1,3 +1,3 @@ -# Physics informed Neural Opeator ODEs Solvers +# Physics informed Neural Operator ODEs Solvers ## some example TODO \ No newline at end of file diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index d4570cd81e..d67c8e49a2 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -48,7 +48,7 @@ function PINOODE(chain, train_set, init_params = nothing; kwargs...) - #TODO fnn trasform + #TODO fnn transform !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) PINOODE(chain, opt, train_set, init_params, kwargs) end From eb2e9c1caa0d0ea30170ea92cbfe9e599b13a966 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 15 Mar 2024 18:01:31 +0400 Subject: [PATCH 015/153] add PINOsolution --- src/pino_ode_solve.jl | 52 +++++++++++++++++++----------- test/PINO_ode_tests.jl | 65 +++++++++++--------------------------- test/PINO_ode_tests_gpu.jl | 10 ++---- 3 files changed, 54 insertions(+), 73 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index d67c8e49a2..f04ff4aec9 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -25,7 +25,7 @@ ## References Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct TRAINSET{} #TODO #T <: Number +struct TRAINSET{} input_data::Vector{ODEProblem} output_data::Array isu0::Bool @@ -39,22 +39,26 @@ struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm chain::C opt::O train_set::TRAINSET + is_data_loss::Bool + is_physics_loss::Bool init_params::P kwargs::K end function PINOODE(chain, opt, - train_set, - init_params = nothing; + train_set; + is_data_loss = true, + is_physics_loss = true, + init_params = nothing, kwargs...) #TODO fnn transform !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, train_set, init_params, kwargs) + PINOODE(chain, opt, train_set, is_data_loss, is_physics_loss, init_params, kwargs) end """ - PINOPhi(chain::Lux.AbstractExplicitLayer, t, st) + PINOPhi(chain::Lux.AbstractExplicitLayer, t0,u0, st) TODO """ mutable struct PINOPhi{C, T, U, S} @@ -66,6 +70,11 @@ mutable struct PINOPhi{C, T, U, S} new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) end end +struct PINOsolution{} + predict::Array + res::SciMLBase.OptimizationSolution + phi::PINOPhi +end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, t0, @@ -100,7 +109,6 @@ function dfdx_rand_matrix(phi::PINOPhi, t::AbstractArray, θ) end function dfdx(phi::PINOPhi, t::AbstractArray, θ) - # ε = [sqrt(eps(eltype(t))), zero(eltype(t))] ε = [sqrt(eps(eltype(t))), zeros(eltype(t), size(t)[1] - 1)...] (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end @@ -124,7 +132,6 @@ function physics_loss(phi::PINOPhi{C, T, U}, fs = cat( [f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3) else - #TODO generalize fs = cat( [reduce( hcat, [f.f(out_[:, j, [i]], p, ts) for j in axes(out_[:, :, [i]], 2)]) @@ -132,13 +139,12 @@ function physics_loss(phi::PINOPhi{C, T, U}, dims = 3) end end - # fs = adapt(parameterless_type(ComponentArrays.getdata(θ)), fs) NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) end function data_loss(phi::PINOPhi{C, T, U}, θ, - ts::AbstractArray, #remove unessasry + ts::AbstractArray, #TODO remove unessasry train_set::TRAINSET, input_data_set) where {C, T, U} prob_set, output_data = train_set.input_data, train_set.output_data @@ -158,13 +164,11 @@ function generate_data(ts, prob_set, isu0) if isu0 == true in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) #TODO for all case p and u0 - # u0 isa Vector # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) else if p isa Number in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) elseif p isa Vector - #TODO nno for Vector inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) in_ = reshape(inner, size(inner)..., 1) else @@ -177,10 +181,20 @@ function generate_data(ts, prob_set, isu0) end function generate_loss( - phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts) where {C, T, U} + phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts, + is_data_loss, is_physics_loss) where { + C, T, U} function loss(θ, _) - data_loss(phi, θ, ts, train_set, input_data_set) + #TODO train_set + input_data_set - physics_loss(phi, θ, ts, train_set, input_data_set) + if is_data_loss + data_loss(phi, θ, ts, train_set, input_data_set) + elseif is_physics_loss + physics_loss(phi, θ, ts, train_set, input_data_set) + elseif is_data_loss && is_physics_loss + data_loss(phi, θ, ts, train_set, input_data_set) + + physics_loss(phi, θ, ts, train_set, input_data_set) + else + error("data loss or physics loss should be true") + end end return loss end @@ -204,6 +218,8 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, chain = alg.chain opt = alg.opt init_params = alg.init_params + is_data_loss = alg.is_data_loss + is_physics_loss = alg.is_physics_loss # mapping between functional space of some vararible 'a' of equation (for example initial # condition {u(t0 x)} or parameter p) and solution of equation u(t) @@ -212,7 +228,6 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - #TODO format data t0 = tspan[1] t_end = tspan[2] instances_size = size(train_set.output_data)[2] @@ -245,7 +260,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end end - total_loss = generate_loss(phi, train_set, input_data_set, ts) + total_loss = generate_loss(phi, train_set, input_data_set, ts, is_data_loss, is_physics_loss) # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() @@ -262,7 +277,6 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - - # PINOsolution(fullsolution) - (res, phi) + predict = phi(input_data_set, res.u) + PINOsolution(predict, res, phi) end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 07f0691460..c0d15eb39a 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -33,28 +33,21 @@ using NeuralPDE * input data: set of parameters 'a': * output data: set of solutions u(t){a} corresponding parameter 'a' """ -train_set = TRAINSET(prob_set, u_output_); + train_set = TRAINSET(prob_set, u_output_); #TODO u0 ? prob = ODEProblem(linear, u0, tspan, 0) chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 1)) - flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), - σ = gelu) - η₀ = 1.0f-2 + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 1)) + # flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), + # σ = gelu) opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(flat_no, opt, train_set) - res, phi = solve(prob, alg, verbose = true, maxiters = 200) - - input_data_set = Array{Float32, 3}(undef, 2, instances_size, batch_size) - for (i, prob) in enumerate(prob_set) - in_ = reduce(vcat, [ts, fill(prob.p, 1, size(ts)[2], 1)]) - input_data_set[:, :, i] = in_ - end - predict = phi(input_data_set, res.u) + alg = PINOODE(chain, opt, train_set; is_data_loss = true, is_physics_loss = true) + pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) + predict = pino_solution.predict ground = u_output_ @test ground≈predict atol=1 end @@ -91,21 +84,11 @@ begin train_set = TRAINSET(prob_set, u_output_; isu0 = true) #TODO u0 ? prob = ODEProblem(linear, 0.0f0, tspan, p) - # chain = Lux.Chain(Lux.Dense(2, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 1)) - fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), - σ = gelu) + fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) opt = OptimizationOptimisers.Adam(0.001) alg = PINOODE(fno, opt, train_set) - res, phi = solve(prob, - alg, verbose = true, - maxiters = 200) - - input_data_set = Array{Float32, 3}(undef, 2, instances_size, batch_size) - for (i, prob) in enumerate(prob_set) - in_ = reduce(vcat, [ts, fill(prob.u0, 1, size(ts)[2], 1)]) - input_data_set[:, :, i] = in_ - end - predict = phi(input_data_set, res.u) + pino_solution = solve(prob, alg, verbose = true, maxiters = 200) + predict = pino_solution.predict ground = u_output_ @test ground≈predict atol=1.0 end @@ -143,25 +126,15 @@ end u_output_[:, :, i] = reshape_sol end - train_set = TRAINSET(prob_set, u_output_); + train_set = NeuralPDE.TRAINSET(prob_set, u_output_); #TODO u0 ? - prob = ODEProblem(lotka_volterra_matrix, u0, tspan, p) - chain = Lux.Chain(Lux.Dense(5, 20, Lux.σ), Lux.Dense(20, 20, Lux.σ), Lux.Dense(20, 2)) + prob = ODEProblem(lotka_volterra, u0, tspan, p) flat_no = FourierNeuralOperator(ch = (5, 16, 16, 16, 16, 16, 32, 2), modes = (16,), σ = gelu) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(flat_no, opt, train_set); - res, phi = solve(prob, alg, verbose = true, maxiters = 200) - - input_data_set = Array{Float32, 3}(undef, 5, instances_size, batch_size) - for (i, prob) in enumerate(prob_set) - inner = reduce(vcat, [ts, reduce(hcat, fill(prob.p, 1, size(ts)[2], 1))]) - in_ = reshape(inner, size(inner)..., 1) - input_data_set[:, :, i] = in_ - end - - predict = phi(input_data_set, res.u) + alg = NeuralPDE.PINOODE(flat_no, opt, train_set, is_data_loss = true, is_physics_loss = false) + pino_solution = solve(prob, alg, verbose = true, maxiters = 500) + predict = pino_solution.predict ground = u_output_ - @test ground≈predict atol=5 end diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 47b6bf75d1..222e7fb3a0 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -57,14 +57,8 @@ const gpud = gpu_device() opt = OptimizationOptimisers.Adam(0.03) alg = NeuralPDE.PINOODE(flat_no, opt, train_set, ps) - res, phi = solve(prob, alg, verbose = true, maxiters = 200) - - input_data_set = Array{Float32, 3}(undef, 2, instances_size, batch_size) - for (i, prob) in enumerate(prob_set) - in_ = reduce(vcat, [ts, fill(prob.p, 1, size(ts)[2], 1)]) - input_data_set[:, :, i] = in_ - end - predict = phi(input_data_set, res.u) |> cpu + pino_solution = solve(prob, alg, verbose = true, maxiters = 200) + predict = pino_solution.predict |> cpu ground = u_output_ @test ground≈predict atol=1 end From ca2b80248533e55b400fcdfe48c7ca2f436385e4 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 15 Mar 2024 18:07:05 +0400 Subject: [PATCH 016/153] add NeuralOperator deps --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 11628f97f2..ea60fc823e 100644 --- a/Project.toml +++ b/Project.toml @@ -65,6 +65,7 @@ LuxCUDA = "0.3" MCMCChains = "6" ModelingToolkit = "8" MonteCarloMeasurements = "1" +NeuralOperators = "0.4.8" Optim = "1.7.8" Optimization = "3" OptimizationOptimJL = "0.2" From 5b8d226565223ae931228b95eff6ea8608579027 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 18 Mar 2024 18:14:15 +0400 Subject: [PATCH 017/153] update Project.toml --- Project.toml | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/Project.toml b/Project.toml index 8af87d292f..f087415093 100644 --- a/Project.toml +++ b/Project.toml @@ -40,11 +40,11 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] -Adapt = "4" +Adapt = "3, 4" AdvancedHMC = "0.6.1" Aqua = "0.8" ArrayInterface = "7.7" -CUDA = "5.2" +CUDA = "4.4.1, 5.2" ChainRulesCore = "1.18" ComponentArrays = "0.15.8" Cubature = "1.5" @@ -63,25 +63,25 @@ LogDensityProblems = "2" Lux = "0.5.14" LuxCUDA = "0.3.2" MCMCChains = "6" -NeuralOperators = "0.4.8" MethodOfLines = "0.10.7" -ModelingToolkit = "8.75" +ModelingToolkit = "8.73, 8.75" MonteCarloMeasurements = "1.1" +NeuralOperators = "0.4.8" Optim = "1.7.8" -Optimization = "3.22" +Optimization = "3.19.3, 3.22" OptimizationOptimJL = "0.2.1" -OptimizationOptimisers = "0.2.1" -OrdinaryDiffEq = "6.70" +OptimizationOptimisers = "0.1.6, 0.2.1" +OrdinaryDiffEq = "6.66.0" Pkg = "1" QuasiMonteCarlo = "0.3.2" Random = "1" Reexport = "1.2" RuntimeGeneratedFunctions = "0.5.11" SafeTestsets = "0.1" -SciMLBase = "2.24" +SciMLBase = "2,10, 2.24" Statistics = "1.10" SymbolicUtils = "1.4" -Symbolics = "5.17" +Symbolics = "5,11, 5.17" Test = "1" UnPack = "1" Zygote = "0.6.68" From c711b0de29c830758b390d7d95f6ea93293e6ca9 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 18 Mar 2024 18:29:48 +0400 Subject: [PATCH 018/153] update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index f087415093..270b43c52f 100644 --- a/Project.toml +++ b/Project.toml @@ -61,7 +61,7 @@ LineSearches = "7.2" LinearAlgebra = "1" LogDensityProblems = "2" Lux = "0.5.14" -LuxCUDA = "0.3.2" +LuxCUDA = "0.3" MCMCChains = "6" MethodOfLines = "0.10.7" ModelingToolkit = "8.73, 8.75" From 2ad84acb6d249e524d1b261efe7bcf50a6c31528 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 18 Mar 2024 18:52:24 +0400 Subject: [PATCH 019/153] update Project.toml --- Project.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 270b43c52f..8a650ea680 100644 --- a/Project.toml +++ b/Project.toml @@ -64,13 +64,13 @@ Lux = "0.5.14" LuxCUDA = "0.3" MCMCChains = "6" MethodOfLines = "0.10.7" -ModelingToolkit = "8.73, 8.75" +ModelingToolkit = "8.73" MonteCarloMeasurements = "1.1" NeuralOperators = "0.4.8" Optim = "1.7.8" -Optimization = "3.19.3, 3.22" +Optimization = "3.19.3" OptimizationOptimJL = "0.2.1" -OptimizationOptimisers = "0.1.6, 0.2.1" +OptimizationOptimisers = "0.1.6" OrdinaryDiffEq = "6.66.0" Pkg = "1" QuasiMonteCarlo = "0.3.2" @@ -78,10 +78,10 @@ Random = "1" Reexport = "1.2" RuntimeGeneratedFunctions = "0.5.11" SafeTestsets = "0.1" -SciMLBase = "2,10, 2.24" +SciMLBase = "2.10" Statistics = "1.10" SymbolicUtils = "1.4" -Symbolics = "5,11, 5.17" +Symbolics = "5.11" Test = "1" UnPack = "1" Zygote = "0.6.68" From a71c4697f3d0cf283054f0d57fc16e590aa25239 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 18 Mar 2024 18:57:45 +0400 Subject: [PATCH 020/153] update --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 8a650ea680..f62b5ea0ad 100644 --- a/Project.toml +++ b/Project.toml @@ -63,7 +63,7 @@ LogDensityProblems = "2" Lux = "0.5.14" LuxCUDA = "0.3" MCMCChains = "6" -MethodOfLines = "0.10.7" +MethodOfLines = "0.10" ModelingToolkit = "8.73" MonteCarloMeasurements = "1.1" NeuralOperators = "0.4.8" From fb730030b10ed6c8e309fb2bfe887982f0e776c7 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 18 Mar 2024 19:21:37 +0400 Subject: [PATCH 021/153] update --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index f62b5ea0ad..faf455eae3 100644 --- a/Project.toml +++ b/Project.toml @@ -69,7 +69,7 @@ MonteCarloMeasurements = "1.1" NeuralOperators = "0.4.8" Optim = "1.7.8" Optimization = "3.19.3" -OptimizationOptimJL = "0.2.1" +OptimizationOptimJL = "0.1.14" OptimizationOptimisers = "0.1.6" OrdinaryDiffEq = "6.66.0" Pkg = "1" From 6c524ab28888e7c2312393aa99f213b2a89bcef2 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 18 Mar 2024 19:54:24 +0400 Subject: [PATCH 022/153] update --- Project.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index faf455eae3..c1e9839806 100644 --- a/Project.toml +++ b/Project.toml @@ -44,7 +44,7 @@ Adapt = "3, 4" AdvancedHMC = "0.6.1" Aqua = "0.8" ArrayInterface = "7.7" -CUDA = "4.4.1, 5.2" +CUDA = "4, 5" ChainRulesCore = "1.18" ComponentArrays = "0.15.8" Cubature = "1.5" @@ -69,7 +69,7 @@ MonteCarloMeasurements = "1.1" NeuralOperators = "0.4.8" Optim = "1.7.8" Optimization = "3.19.3" -OptimizationOptimJL = "0.1.14" +OptimizationOptimJL = "0.1" OptimizationOptimisers = "0.1.6" OrdinaryDiffEq = "6.66.0" Pkg = "1" @@ -79,7 +79,7 @@ Reexport = "1.2" RuntimeGeneratedFunctions = "0.5.11" SafeTestsets = "0.1" SciMLBase = "2.10" -Statistics = "1.10" +Statistics = "1" SymbolicUtils = "1.4" Symbolics = "5.11" Test = "1" From 2addf6c5786ecd5cd5978ca597385ff2fd8bbba6 Mon Sep 17 00:00:00 2001 From: KirilZubov Date: Mon, 18 Mar 2024 12:05:01 -0400 Subject: [PATCH 023/153] update gpu test --- src/pino_ode_solve.jl | 5 ++-- test/PINO_ode_tests.jl | 46 ---------------------------- test/PINO_ode_tests_gpu.jl | 61 +++++++++++++++++++++++++++++++++----- 3 files changed, 57 insertions(+), 55 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index f04ff4aec9..9f1f3421b2 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -94,8 +94,9 @@ function (f::PINOPhi{C, T, U})(t::AbstractArray, # Batch via data as row vectors y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[[1], :, :]) - f.u0 .+ (ts .- f.t0) .* y + ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[1:size(y)[1], :, :]) + f_ = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0) + f_ .+ (ts .- f.t0) .* y end function dfdx_rand_matrix(phi::PINOPhi, t::AbstractArray, θ) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index c0d15eb39a..f3df3efa9e 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -92,49 +92,3 @@ begin ground = u_output_ @test ground≈predict atol=1.0 end - -@testset "lotka volterra" begin - function lotka_volterra(u, p, t) - # Model parameters. - α, β, γ, δ = p - # Current state. - x, y = u - # Evaluate differential equations. - dx = (α - β * y) * x # prey - dy = (δ * x - γ) * y # predator - return [dx, dy] - end - - u0 = [1.0f0, 1.0f0] - p = Float32[1.5, 1.0, 3.0, 1.0] - tspan = (0.0f0, 4.0f0) - dt = 0.01f0 - - instances_size = 100 - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - batch_size = 50 - ps = [p .+ i * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size] - u_output_ = zeros(Float32, 2, instances_size, batch_size) - prob_set = [] - for (i, p_i) in enumerate(ps) - prob = ODEProblem(lotka_volterra, u0, tspan, p_i) - solution = solve(prob, Tsit5(); saveat = dt) - reshape_sol_ = reduce(hcat, solution(range_).u) - reshape_sol = Float32.(reshape(reshape_sol_, 2, instances_size, 1)) - push!(prob_set, prob) - u_output_[:, :, i] = reshape_sol - end - - train_set = NeuralPDE.TRAINSET(prob_set, u_output_); - #TODO u0 ? - prob = ODEProblem(lotka_volterra, u0, tspan, p) - flat_no = FourierNeuralOperator(ch = (5, 16, 16, 16, 16, 16, 32, 2), modes = (16,), - σ = gelu) - opt = OptimizationOptimisers.Adam(0.01) - alg = NeuralPDE.PINOODE(flat_no, opt, train_set, is_data_loss = true, is_physics_loss = false) - pino_solution = solve(prob, alg, verbose = true, maxiters = 500) - predict = pino_solution.predict - ground = u_output_ - @test ground≈predict atol=5 -end diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 222e7fb3a0..9664114ddb 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -33,17 +33,18 @@ const gpud = gpu_device() push!(prob_set, prob) u_output_[:, :, i] = reshape_sol end + u_output_ = u_output_ |> gpud """ Set of training data: * input data: set of parameters 'a': * output data: set of solutions u(t){a} corresponding parameter 'a' """ - train_set = NeuralPDE.TRAINSET(prob_set, u_output_) + train_set = TRAINSET(prob_set, u_output_) #TODO u0 ? prob = ODEProblem(linear, u0, tspan, 0) - inner = 20 + inner = 50 chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ), Lux.Dense(inner, inner, Lux.σ), Lux.Dense(inner, inner, Lux.σ), @@ -51,14 +52,60 @@ const gpud = gpu_device() Lux.Dense(inner, inner, Lux.σ), Lux.Dense(inner, 1)) ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud + opt = OptimizationOptimisers.Adam(0.03) + alg = PINOODE(chain, opt, train_set; init_params = ps) + pino_solution = solve(prob, alg, verbose = true, maxiters = 1000) + predict = pino_solution.predict |> cpu + ground = u_output_ |> cpu + @test ground≈predict atol=1 +end + +@testset "lotka volterra" begin + function lotka_volterra(u, p, t) + # Model parameters. + α, β, γ, δ = p + # Current state. + x, y = u + # Evaluate differential equations. + dx = (α - β * y) * x # prey + dy = (δ * x - γ) * y # predator + return [dx, dy] + end + + u0 = [1.0f0, 1.0f0] + p = Float32[1.5, 1.0, 3.0, 1.0] + tspan = (0.0f0, 4.0f0) + dt = 0.01f0 + + instances_size = 100 + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) + batch_size = 50 + ps = [p .+ i * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size] + u_output_ = zeros(Float32, 2, instances_size, batch_size) + prob_set = [] + for (i, p_i) in enumerate(ps) + prob = ODEProblem(lotka_volterra, u0, tspan, p_i) + solution = solve(prob, Tsit5(); saveat = dt) + reshape_sol_ = reduce(hcat, solution(range_).u) + reshape_sol = Float32.(reshape(reshape_sol_, 2, instances_size, 1)) + push!(prob_set, prob) + u_output_[:, :, i] = reshape_sol + end - flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), + train_set = TRAINSET(prob_set, u_output_) + #TODO u0 ? + prob = ODEProblem(lotka_volterra, u0, tspan, p) + flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,), σ = gelu) + flat_no = Lux.transform(flat_no) + ps = Lux.setup(Random.default_rng(), flat_no)[1] |> ComponentArray |> gpud - opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(flat_no, opt, train_set, ps) - pino_solution = solve(prob, alg, verbose = true, maxiters = 200) + opt = OptimizationOptimisers.Adam(0.001) + alg = PINOODE( + flat_no, opt, train_set; init_params = ps, is_data_loss = true, is_physics_loss = true) + pino_solution = solve(prob, alg, verbose = true, maxiters = 1000) predict = pino_solution.predict |> cpu ground = u_output_ - @test ground≈predict atol=1 + @test ground≈predict atol=2 end From c058f5e86cfd33e1ff192dbf1b141d1b879cec82 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 18 Mar 2024 20:10:08 +0400 Subject: [PATCH 024/153] updste version Statistics --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index c1e9839806..4331cf2d64 100644 --- a/Project.toml +++ b/Project.toml @@ -79,7 +79,7 @@ Reexport = "1.2" RuntimeGeneratedFunctions = "0.5.11" SafeTestsets = "0.1" SciMLBase = "2.10" -Statistics = "1" +Statistics = "1.10.0" SymbolicUtils = "1.4" Symbolics = "5.11" Test = "1" From 48c762b0b60f4727dde8ad8b4888c74d6a3d7559 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 16:43:09 +0400 Subject: [PATCH 025/153] add docs --- docs/src/tutorials/pino_ode.md | 63 +++++++++++++++++++++++++- src/pino_ode_solve.jl | 82 ++++++++++++++++------------------ test/PINO_ode_tests.jl | 11 +++-- test/PINO_ode_tests_gpu.jl | 6 +-- 4 files changed, 108 insertions(+), 54 deletions(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index d7ca3102b0..6459a91c5d 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -1,3 +1,64 @@ # Physics informed Neural Operator ODEs Solvers -## some example TODO \ No newline at end of file +This tutorial is an introduction to using physics-informed neural operator (PINOs) for solving family of parametric ordinary diferential equations (ODEs). + + +## Solving a family of parametric ODE. + +```@example pino +using Test +using OrdinaryDiffEq, OptimizationOptimisers +using Lux +using Statistics, Random +using NeuralOperators +using NeuralPDE + +linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) +linear = (u, p, t) -> cos(p * t) +tspan = (0.0f0, 2.0f0) +u0 = 0.0f0 +``` + +Generate a dataset for learning a given family of ODEs where the parameter 'a' is varied. The dataset is generated by solving the ODE for different values of 'a' and storing the solutions. The dataset is then used to train the PINO model: +* input data: set of parameters 'a', +* output data: set of solutions u(t){a} corresponding parameter 'a'. + +```@example pino +t0, t_end = tspan +instances_size = 50 +range_ = range(t0, stop = t_end, length = instances_size) +ts = reshape(collect(range_), 1, instances_size) +batch_size = 50 +as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] + +u_output_ = zeros(Float32, 1, instances_size, batch_size) +prob_set = [] +for (i, a_i) in enumerate(as) + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) + sol1 = solve(prob, Tsit5(); saveat = 0.0204) + reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) + push!(prob_set, prob) + u_output_[:, :, i] = reshape_sol +end +train_set = TRAINSET(prob_set, u_output_); +``` + +Here it used the PINO method to train the given family of parametric ODEs. + +```@example pino +prob = ODEProblem(linear, u0, tspan, 0) +flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), + σ = gelu) +opt = OptimizationOptimisers.Adam(0.03) +alg = PINOODE(flat_no, opt, train_set; is_data_loss = true, is_physics_loss = true) +pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) +predict = pino_solution.predict +ground = u_output_ +``` + +Now let's compare the predictions from the learned operator with the ground truth solution which is obtained early by numerically solving the parametric ODE. Where 'i' is the index of the parameter 'a' in the dataset. + +```@example pino +plot(predict[1, :, i], label = "Predicted") +plot!(ground[1, :, i], label = "Ground truth") +``` \ No newline at end of file diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 9f1f3421b2..8ed273888d 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,40 +1,32 @@ """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), - train_set - init_params = nothing; + train_set, + is_data_loss =true, + is_physics_loss =true, + init_params, kwargs...) -## Positional Arguments +The method is that combine training data and physics constraints +to learn the solution operator of a given family of parametric Ordinary Differential Equations (ODE). -* `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.AbstractExplicitLayer`. +## Positional Arguments +* `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`. + `Flux.Chain` will be converted to `Lux` using `Lux.transform`. * `opt`: The optimizer to train the neural network. -* `train_set`: -* `init_params`: The initial parameter of the neural network. By default, this is `nothing` - which thus uses the random initialization provided by the neural network library. +* `train_set`: Contains 'input data' - sr of parameters 'a' and output data - set of solutions + u(t){a} corresponding initial conditions 'u0'. ## Keyword Arguments -* `minibatch`: - -## Examples - -```julia - -``` +* `is_data_loss` Includes or off a loss function for training on the data set. +* `is_physics_loss`: Includes or off loss function training on physics-informed approach. +* `init_params`: The initial parameter of the neural network. By default, this is `nothing` + which thus uses the random initialization provided by the neural network library. +* `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. ## References Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct TRAINSET{} - input_data::Vector{ODEProblem} - output_data::Array - isu0::Bool -end - -function TRAINSET(input_data, output_data; isu0 = false) - TRAINSET(input_data, output_data, isu0) -end - struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm chain::C opt::O @@ -57,10 +49,16 @@ function PINOODE(chain, PINOODE(chain, opt, train_set, is_data_loss, is_physics_loss, init_params, kwargs) end -""" - PINOPhi(chain::Lux.AbstractExplicitLayer, t0,u0, st) - TODO -""" +struct TRAINSET{} + input_data::Vector{ODEProblem} + output_data::Array + isu0::Bool +end + +function TRAINSET(input_data, output_data; isu0 = false) + TRAINSET(input_data, output_data, isu0) +end + mutable struct PINOPhi{C, T, U, S} chain::C t0::T @@ -91,7 +89,6 @@ end function (f::PINOPhi{C, T, U})(t::AbstractArray, θ) where {C <: Lux.AbstractExplicitLayer, T, U} - # Batch via data as row vectors y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[1:size(y)[1], :, :]) @@ -119,9 +116,8 @@ function physics_loss(phi::PINOPhi{C, T, U}, ts::AbstractArray, train_set::TRAINSET, input_data_set) where {C, T, U} - prob_set, output_data = train_set.input_data, train_set.output_data #TODO - f = prob_set[1].f #TODO one f for all - p = prob_set[1].p + prob_set, _ = train_set.input_data, train_set.output_data + f = prob_set[1].f out_ = phi(input_data_set, θ) ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), ts) if train_set.isu0 == true @@ -132,12 +128,14 @@ function physics_loss(phi::PINOPhi{C, T, U}, if p isa Number fs = cat( [f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3) - else + elseif p isa Vector fs = cat( [reduce( hcat, [f.f(out_[:, j, [i]], p, ts) for j in axes(out_[:, :, [i]], 2)]) for (i, p) in enumerate(ps)]..., dims = 3) + else + error("p should be a number or a vector") end end NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) @@ -145,10 +143,9 @@ end function data_loss(phi::PINOPhi{C, T, U}, θ, - ts::AbstractArray, #TODO remove unessasry train_set::TRAINSET, input_data_set) where {C, T, U} - prob_set, output_data = train_set.input_data, train_set.output_data + _, output_data = train_set.input_data, train_set.output_data output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data) NeuralOperators.l₂loss(phi(input_data_set, θ), output_data) end @@ -164,8 +161,6 @@ function generate_data(ts, prob_set, isu0) f = prob.f if isu0 == true in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) - #TODO for all case p and u0 - # in_ = reduce(vcat, [ts, reduce(hcat, fill(u0, 1, size(ts)[2], 1))]) else if p isa Number in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) @@ -187,11 +182,11 @@ function generate_loss( C, T, U} function loss(θ, _) if is_data_loss - data_loss(phi, θ, ts, train_set, input_data_set) + data_loss(phi, θ, train_set, input_data_set) elseif is_physics_loss physics_loss(phi, θ, ts, train_set, input_data_set) elseif is_data_loss && is_physics_loss - data_loss(phi, θ, ts, train_set, input_data_set) + + data_loss(phi, θ, train_set, input_data_set) + physics_loss(phi, θ, ts, train_set, input_data_set) else error("data loss or physics loss should be true") @@ -211,9 +206,9 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, maxiters = nothing) tspan = prob.tspan t0 = tspan[1] + u0 = prob.u0 # f = prob.f # p = prob.p - u0 = prob.u0 # param_estim = alg.param_estim chain = alg.chain @@ -234,7 +229,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, instances_size = size(train_set.output_data)[2] range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) - prob_set, output_data = train_set.input_data, train_set.output_data + prob_set, _ = train_set.input_data, train_set.output_data isu0 = train_set.isu0 input_data_set = generate_data(ts, prob_set, isu0) @@ -255,13 +250,14 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, phi(input_data_set, init_params) catch err if isa(err, DimensionMismatch) - throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) #TODO change message + throw(DimensionMismatch("Dimensions of input data and chain should match")) else throw(err) end end - total_loss = generate_loss(phi, train_set, input_data_set, ts, is_data_loss, is_physics_loss) + total_loss = generate_loss( + phi, train_set, input_data_set, ts, is_data_loss, is_physics_loss) # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index f3df3efa9e..93732dce28 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -30,11 +30,10 @@ using NeuralPDE """ Set of training data: - * input data: set of parameters 'a': - * output data: set of solutions u(t){a} corresponding parameter 'a' + * input data: set of parameters 'a' + * output data: set of solutions u(t){a} corresponding parameter 'a'. """ train_set = TRAINSET(prob_set, u_output_); - #TODO u0 ? prob = ODEProblem(linear, u0, tspan, 0) chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), @@ -78,11 +77,11 @@ begin """ Set of training data: - * input data: set of initial conditions 'u0': - * output data: set of solutions u(t){u0} corresponding initial conditions 'u0' + * input data: set of initial conditions 'u0' + * output data: set of solutions u(t){u0} corresponding initial conditions 'u0'. """ train_set = TRAINSET(prob_set, u_output_; isu0 = true) - #TODO u0 ? + #TODO we argument u0 but dont actualy use u0 because we use only set of u0 for generate train set from prob_set prob = ODEProblem(linear, 0.0f0, tspan, p) fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) opt = OptimizationOptimisers.Adam(0.001) diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 9664114ddb..1ae44c2adb 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -37,12 +37,11 @@ const gpud = gpu_device() """ Set of training data: - * input data: set of parameters 'a': - * output data: set of solutions u(t){a} corresponding parameter 'a' + * input data: set of parameters 'a', + * output data: set of solutions u(t){a} corresponding parameter 'a'. """ train_set = TRAINSET(prob_set, u_output_) - #TODO u0 ? prob = ODEProblem(linear, u0, tspan, 0) inner = 50 chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ), @@ -94,7 +93,6 @@ end end train_set = TRAINSET(prob_set, u_output_) - #TODO u0 ? prob = ODEProblem(lotka_volterra, u0, tspan, p) flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,), σ = gelu) From 671cedc7c71a8b4275da6e77165db261ddcd9ecc Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 16:45:19 +0400 Subject: [PATCH 026/153] fix --- src/pino_ode_solve.jl | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 8ed273888d..abdd6f852f 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,3 +1,13 @@ +struct TRAINSET{} + input_data::Vector{ODEProblem} + output_data::Array + isu0::Bool +end + +function TRAINSET(input_data, output_data; isu0 = false) + TRAINSET(input_data, output_data, isu0) +end + """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), @@ -49,16 +59,6 @@ function PINOODE(chain, PINOODE(chain, opt, train_set, is_data_loss, is_physics_loss, init_params, kwargs) end -struct TRAINSET{} - input_data::Vector{ODEProblem} - output_data::Array - isu0::Bool -end - -function TRAINSET(input_data, output_data; isu0 = false) - TRAINSET(input_data, output_data, isu0) -end - mutable struct PINOPhi{C, T, U, S} chain::C t0::T From fe8a819cb92f7bcd8d39997880ce6b3d9c4f6158 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 16:57:47 +0400 Subject: [PATCH 027/153] fix --- Project.toml | 2 +- docs/src/tutorials/pino_ode.md | 1 + test/PINO_ode_tests.jl | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 4331cf2d64..1d34bdfa6c 100644 --- a/Project.toml +++ b/Project.toml @@ -70,7 +70,7 @@ NeuralOperators = "0.4.8" Optim = "1.7.8" Optimization = "3.19.3" OptimizationOptimJL = "0.1" -OptimizationOptimisers = "0.1.6" +OptimizationOptimisers = "0.1" OrdinaryDiffEq = "6.66.0" Pkg = "1" QuasiMonteCarlo = "0.3.2" diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 6459a91c5d..721b63c344 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -59,6 +59,7 @@ ground = u_output_ Now let's compare the predictions from the learned operator with the ground truth solution which is obtained early by numerically solving the parametric ODE. Where 'i' is the index of the parameter 'a' in the dataset. ```@example pino +i=1 plot(predict[1, :, i], label = "Predicted") plot!(ground[1, :, i], label = "Ground truth") ``` \ No newline at end of file diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 93732dce28..0d35b5c166 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -45,7 +45,7 @@ using NeuralPDE # σ = gelu) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(chain, opt, train_set; is_data_loss = true, is_physics_loss = true) - pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) + pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ @test ground≈predict atol=1 From 2c87492bbda101685f597a070d81ca716038980d Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 17:21:22 +0400 Subject: [PATCH 028/153] fix TRAINSET --- src/NeuralPDE.jl | 2 +- test/PINO_ode_tests_gpu.jl | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 18bbdf36c0..30e59486e7 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -55,7 +55,7 @@ include("BPINN_ode.jl") include("PDE_BPINN.jl") include("dgm.jl") -export NNODE, NNDAE, PINOODE +export NNODE, NNDAE, PINOODE, TRAINSET PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 1ae44c2adb..ef14ad5001 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -53,7 +53,7 @@ const gpud = gpu_device() ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(chain, opt, train_set; init_params = ps) - pino_solution = solve(prob, alg, verbose = true, maxiters = 1000) + pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) predict = pino_solution.predict |> cpu ground = u_output_ |> cpu @test ground≈predict atol=1 @@ -102,7 +102,7 @@ end opt = OptimizationOptimisers.Adam(0.001) alg = PINOODE( flat_no, opt, train_set; init_params = ps, is_data_loss = true, is_physics_loss = true) - pino_solution = solve(prob, alg, verbose = true, maxiters = 1000) + pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) predict = pino_solution.predict |> cpu ground = u_output_ @test ground≈predict atol=2 From 3905b1433af1e193618f1e9d40ad473612e654ef Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 17:26:43 +0400 Subject: [PATCH 029/153] fix typo --- src/NeuralPDE.jl | 29 +++++++++++++++-------------- test/PINO_ode_tests.jl | 2 +- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 30e59486e7..68e7073ab4 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -24,14 +24,15 @@ using Symbolics: wrap, unwrap, arguments, operation using SymbolicUtils using AdvancedHMC, LogDensityProblems, LinearAlgebra, Functors, MCMCChains using MonteCarloMeasurements: Particles -using ModelingToolkit: value, nameof, toexpr, build_expr, expand_derivatives, Interval, infimum, supremum +using ModelingToolkit: value, nameof, toexpr, build_expr, expand_derivatives, Interval, + infimum, supremum import DomainSets -using DomainSets: Domain, ClosedInterval, AbstractInterval, leftendpoint, rightendpoint, ProductDomain +using DomainSets: Domain, ClosedInterval, AbstractInterval, leftendpoint, rightendpoint, + ProductDomain using SciMLBase: @add_kwonly, parameterless_type using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using ChainRulesCore: @non_differentiable -using NeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) @@ -56,16 +57,16 @@ include("PDE_BPINN.jl") include("dgm.jl") export NNODE, NNDAE, PINOODE, TRAINSET - PhysicsInformedNN, discretize, - GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, - WeightedIntervalTraining, - build_loss_function, get_loss_function, - generate_training_sets, get_variables, get_argument, get_bounds, - get_numeric_integral, symbolic_discretize, - AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, - MiniMaxAdaptiveLoss, LogOptions, - ahmc_bayesian_pinn_ode, BNNODE, ahmc_bayesian_pinn_pde, vector_to_parameters, - BPINNsolution, BayesianPINN, - DeepGalerkin + PhysicsInformedNN, discretize, + GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, + WeightedIntervalTraining, + build_loss_function, get_loss_function, + generate_training_sets, get_variables, get_argument, get_bounds, + get_numeric_integral, symbolic_discretize, + AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, + MiniMaxAdaptiveLoss, LogOptions, + ahmc_bayesian_pinn_ode, BNNODE, ahmc_bayesian_pinn_pde, vector_to_parameters, + BPINNsolution, BayesianPINN, + DeepGalerkin end # module diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 0d35b5c166..e733b16a0c 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -81,7 +81,7 @@ begin * output data: set of solutions u(t){u0} corresponding initial conditions 'u0'. """ train_set = TRAINSET(prob_set, u_output_; isu0 = true) - #TODO we argument u0 but dont actualy use u0 because we use only set of u0 for generate train set from prob_set + #TODO we argument u0 but dont actually use u0 because we use only set of u0 for generate train set from prob_set prob = ODEProblem(linear, 0.0f0, tspan, p) fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) opt = OptimizationOptimisers.Adam(0.001) From 652fb72a84e2598c662490878c1c6bd080470277 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 18:35:30 +0400 Subject: [PATCH 030/153] add dep NeuralOperators --- src/NeuralPDE.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 68e7073ab4..ab25ec93ba 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -33,6 +33,7 @@ using SciMLBase: @add_kwonly, parameterless_type using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using ChainRulesCore: @non_differentiable +using NeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) From f74a46f530b8e6ed9aeea42ac2df13b0c1bc9a7c Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 18:55:57 +0400 Subject: [PATCH 031/153] fix test --- test/PINO_ode_tests.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index e733b16a0c..a3ebfea88b 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -51,8 +51,7 @@ using NeuralPDE @test ground≈predict atol=1 end -"Example u0" -begin +@testset "Example u0" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) From c163e2852436d67350d21ad8875ec08f172efc2a Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 19:02:22 +0400 Subject: [PATCH 032/153] fix gpu test --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 52b6e875fe..babd817752 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -70,7 +70,7 @@ end @safetestset "NNPDE_gpu_Lux" begin include("NNPDE_tests_gpu_Lux.jl") end end if !is_APPVEYOR && GROUP == "PINO_GPU" - @safetestset "PINo ode gpu" begin include("PINO_ode_tests_gpu.jl.jl") + @safetestset "PINO ode gpu" begin include("PINO_ode_tests_gpu.jl") end end From b52c4c2a3bd1aa556fcf1bbce7e84133579cfba1 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 19 Mar 2024 19:17:27 +0400 Subject: [PATCH 033/153] fix tests --- test/PINO_ode_tests.jl | 2 +- test/PINO_ode_tests_gpu.jl | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index a3ebfea88b..321b093c81 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -85,7 +85,7 @@ end fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) opt = OptimizationOptimisers.Adam(0.001) alg = PINOODE(fno, opt, train_set) - pino_solution = solve(prob, alg, verbose = true, maxiters = 200) + pino_solution = solve(prob, alg, verbose = false, maxiters = 200) predict = pino_solution.predict ground = u_output_ @test ground≈predict atol=1.0 diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index ef14ad5001..6358d484b3 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -53,7 +53,7 @@ const gpud = gpu_device() ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(chain, opt, train_set; init_params = ps) - pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) + pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict |> cpu ground = u_output_ |> cpu @test ground≈predict atol=1 @@ -75,6 +75,7 @@ end p = Float32[1.5, 1.0, 3.0, 1.0] tspan = (0.0f0, 4.0f0) dt = 0.01f0 + t0, t_end = tspan instances_size = 100 range_ = range(t0, stop = t_end, length = instances_size) From b8292650334cf5526a28f06af15b980b501520ed Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 20 Mar 2024 16:48:54 +0400 Subject: [PATCH 034/153] remove dep NeuralOperator --- Project.toml | 28 +++++++++++++--------------- src/NeuralPDE.jl | 2 +- src/pino_ode_solve.jl | 11 +++++++++-- test/PINO_ode_tests.jl | 12 +++++++++--- test/PINO_ode_tests_gpu.jl | 17 ++++++++++++----- 5 files changed, 44 insertions(+), 26 deletions(-) diff --git a/Project.toml b/Project.toml index 1d34bdfa6c..206234887c 100644 --- a/Project.toml +++ b/Project.toml @@ -24,7 +24,6 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" -NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1" @@ -40,11 +39,11 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] -Adapt = "3, 4" +Adapt = "4" AdvancedHMC = "0.6.1" Aqua = "0.8" ArrayInterface = "7.7" -CUDA = "4, 5" +CUDA = "5.2" ChainRulesCore = "1.18" ComponentArrays = "0.15.8" Cubature = "1.5" @@ -61,27 +60,26 @@ LineSearches = "7.2" LinearAlgebra = "1" LogDensityProblems = "2" Lux = "0.5.14" -LuxCUDA = "0.3" +LuxCUDA = "0.3.2" MCMCChains = "6" -MethodOfLines = "0.10" -ModelingToolkit = "8.73" +MethodOfLines = "0.10.7" +ModelingToolkit = "8.75" MonteCarloMeasurements = "1.1" -NeuralOperators = "0.4.8" Optim = "1.7.8" -Optimization = "3.19.3" -OptimizationOptimJL = "0.1" -OptimizationOptimisers = "0.1" -OrdinaryDiffEq = "6.66.0" +Optimization = "3.22" +OptimizationOptimJL = "0.2.1" +OptimizationOptimisers = "0.2.1" +OrdinaryDiffEq = "6.70" Pkg = "1" QuasiMonteCarlo = "0.3.2" Random = "1" Reexport = "1.2" RuntimeGeneratedFunctions = "0.5.11" SafeTestsets = "0.1" -SciMLBase = "2.10" -Statistics = "1.10.0" +SciMLBase = "2.24" +Statistics = "1.10" SymbolicUtils = "1.4" -Symbolics = "5.11" +Symbolics = "5.17" Test = "1" UnPack = "1" Zygote = "0.6.68" @@ -101,4 +99,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4" [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"] \ No newline at end of file diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index ab25ec93ba..dbb37b15e2 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -33,7 +33,7 @@ using SciMLBase: @add_kwonly, parameterless_type using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using ChainRulesCore: @non_differentiable -using NeuralOperators +#using NeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index abdd6f852f..13452ba5dc 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -111,6 +111,13 @@ function dfdx(phi::PINOPhi, t::AbstractArray, θ) (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) end +function l₂loss(𝐲̂, 𝐲) + feature_dims = 2:(ndims(𝐲) - 1) + loss = sum(.√(sum(abs2, 𝐲̂ - 𝐲, dims = feature_dims))) + y_norm = sum(.√(sum(abs2, 𝐲, dims = feature_dims))) + return loss / y_norm +end + function physics_loss(phi::PINOPhi{C, T, U}, θ, ts::AbstractArray, @@ -138,7 +145,7 @@ function physics_loss(phi::PINOPhi{C, T, U}, error("p should be a number or a vector") end end - NeuralOperators.l₂loss(dfdx(phi, input_data_set, θ), fs) + l₂loss(dfdx(phi, input_data_set, θ), fs) end function data_loss(phi::PINOPhi{C, T, U}, @@ -147,7 +154,7 @@ function data_loss(phi::PINOPhi{C, T, U}, input_data_set) where {C, T, U} _, output_data = train_set.input_data, train_set.output_data output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data) - NeuralOperators.l₂loss(phi(input_data_set, θ), output_data) + l₂loss(phi(input_data_set, θ), output_data) end function generate_data(ts, prob_set, isu0) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 321b093c81..e361af17e5 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -2,7 +2,7 @@ using Test using OrdinaryDiffEq, OptimizationOptimisers using Lux using Statistics, Random -using NeuralOperators +#using NeuralOperators using NeuralPDE @testset "Example p" begin @@ -82,9 +82,15 @@ end train_set = TRAINSET(prob_set, u_output_; isu0 = true) #TODO we argument u0 but dont actually use u0 because we use only set of u0 for generate train set from prob_set prob = ODEProblem(linear, 0.0f0, tspan, p) - fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) + # fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) + chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 1)) opt = OptimizationOptimisers.Adam(0.001) - alg = PINOODE(fno, opt, train_set) + alg = PINOODE(chain, opt, train_set) pino_solution = solve(prob, alg, verbose = false, maxiters = 200) predict = pino_solution.predict ground = u_output_ diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 6358d484b3..1f9ac664a6 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -2,7 +2,7 @@ using Test using OrdinaryDiffEq using Lux using ComponentArrays -using NeuralOperators +#using NeuralOperators using OptimizationOptimisers using Random using LuxCUDA @@ -95,10 +95,17 @@ end train_set = TRAINSET(prob_set, u_output_) prob = ODEProblem(lotka_volterra, u0, tspan, p) - flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,), - σ = gelu) - flat_no = Lux.transform(flat_no) - ps = Lux.setup(Random.default_rng(), flat_no)[1] |> ComponentArray |> gpud + # flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,), + # σ = gelu) + # flat_no = Lux.transform(flat_no) + inner = 50 + chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, inner, Lux.σ), + Lux.Dense(inner, 1)) + ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.001) alg = PINOODE( From 9870858284706431cbff72f2b54147e9d4beebb6 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 20 Mar 2024 16:54:31 +0400 Subject: [PATCH 035/153] update tests --- test/PINO_ode_tests.jl | 11 ++++++----- test/PINO_ode_tests_gpu.jl | 4 ++-- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index e361af17e5..4dacf28f51 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -33,7 +33,7 @@ using NeuralPDE * input data: set of parameters 'a' * output data: set of solutions u(t){a} corresponding parameter 'a'. """ - train_set = TRAINSET(prob_set, u_output_); + train_set = NeuralPDE.TRAINSET(prob_set, u_output_); prob = ODEProblem(linear, u0, tspan, 0) chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), @@ -44,7 +44,8 @@ using NeuralPDE # flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), # σ = gelu) opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(chain, opt, train_set; is_data_loss = true, is_physics_loss = true) + alg = NeuralPDE.PINOODE( + chain, opt, train_set; is_data_loss = true, is_physics_loss = true) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ @@ -79,7 +80,7 @@ end * input data: set of initial conditions 'u0' * output data: set of solutions u(t){u0} corresponding initial conditions 'u0'. """ - train_set = TRAINSET(prob_set, u_output_; isu0 = true) + train_set = NeuralPDE.TRAINSET(prob_set, u_output_; isu0 = true) #TODO we argument u0 but dont actually use u0 because we use only set of u0 for generate train set from prob_set prob = ODEProblem(linear, 0.0f0, tspan, p) # fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) @@ -90,8 +91,8 @@ end Lux.Dense(16, 32, Lux.σ), Lux.Dense(32, 1)) opt = OptimizationOptimisers.Adam(0.001) - alg = PINOODE(chain, opt, train_set) - pino_solution = solve(prob, alg, verbose = false, maxiters = 200) + alg = NeuralPDE.PINOODE(chain, opt, train_set) + pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ @test ground≈predict atol=1.0 diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 1f9ac664a6..22de9b9555 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -110,8 +110,8 @@ end opt = OptimizationOptimisers.Adam(0.001) alg = PINOODE( flat_no, opt, train_set; init_params = ps, is_data_loss = true, is_physics_loss = true) - pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) + pino_solution = solve(prob, alg, verbose = false, maxiters = 4000) predict = pino_solution.predict |> cpu ground = u_output_ - @test ground≈predict atol=2 + @test ground≈predict atol=5 end From c5e734806bb1c68eea364f25296bacc44b3278bb Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 20 Mar 2024 16:56:20 +0400 Subject: [PATCH 036/153] fix --- test/PINO_ode_tests.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 4dacf28f51..4df60b2bb8 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -33,7 +33,7 @@ using NeuralPDE * input data: set of parameters 'a' * output data: set of solutions u(t){a} corresponding parameter 'a'. """ - train_set = NeuralPDE.TRAINSET(prob_set, u_output_); + train_set = TRAINSET(prob_set, u_output_); prob = ODEProblem(linear, u0, tspan, 0) chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), @@ -44,7 +44,7 @@ using NeuralPDE # flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), # σ = gelu) opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE( + alg = PINOODE( chain, opt, train_set; is_data_loss = true, is_physics_loss = true) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict @@ -80,7 +80,7 @@ end * input data: set of initial conditions 'u0' * output data: set of solutions u(t){u0} corresponding initial conditions 'u0'. """ - train_set = NeuralPDE.TRAINSET(prob_set, u_output_; isu0 = true) + train_set = TRAINSET(prob_set, u_output_; isu0 = true) #TODO we argument u0 but dont actually use u0 because we use only set of u0 for generate train set from prob_set prob = ODEProblem(linear, 0.0f0, tspan, p) # fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) @@ -91,7 +91,7 @@ end Lux.Dense(16, 32, Lux.σ), Lux.Dense(32, 1)) opt = OptimizationOptimisers.Adam(0.001) - alg = NeuralPDE.PINOODE(chain, opt, train_set) + alg = PINOODE(chain, opt, train_set) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ From cefb901280b963e4d7a0f2c10aea9c3a6bd1257a Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 20 Mar 2024 16:58:00 +0400 Subject: [PATCH 037/153] fix --- test/PINO_ode_tests_gpu.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 22de9b9555..147f46dbe3 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -109,7 +109,7 @@ end opt = OptimizationOptimisers.Adam(0.001) alg = PINOODE( - flat_no, opt, train_set; init_params = ps, is_data_loss = true, is_physics_loss = true) + chain, opt, train_set; init_params = ps, is_data_loss = true, is_physics_loss = true) pino_solution = solve(prob, alg, verbose = false, maxiters = 4000) predict = pino_solution.predict |> cpu ground = u_output_ From 10169cbcd0cf2b0dd6ca972af40b353452424a58 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 20 Mar 2024 17:00:46 +0400 Subject: [PATCH 038/153] fix --- test/PINO_ode_tests_gpu.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 147f46dbe3..24cf7226b5 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -99,12 +99,12 @@ end # σ = gelu) # flat_no = Lux.transform(flat_no) inner = 50 - chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ), + chain = Lux.Chain(Lux.Dense(5, inner, Lux.σ), Lux.Dense(inner, inner, Lux.σ), Lux.Dense(inner, inner, Lux.σ), Lux.Dense(inner, inner, Lux.σ), Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, 1)) + Lux.Dense(inner, 2)) ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.001) From 2f2be69e5d1e22a63e30a8c96f939d1e1e675653 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 21 Mar 2024 18:45:40 +0400 Subject: [PATCH 039/153] fine tuning --- src/NeuralPDE.jl | 2 +- src/pino_ode_solve.jl | 133 +++++++++++++++++++++++++------------ test/PINO_ode_tests.jl | 28 ++++++-- test/PINO_ode_tests_gpu.jl | 6 +- 4 files changed, 119 insertions(+), 50 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index dbb37b15e2..b589af2d03 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -57,7 +57,7 @@ include("BPINN_ode.jl") include("PDE_BPINN.jl") include("dgm.jl") -export NNODE, NNDAE, PINOODE, TRAINSET +export NNODE, NNDAE, PINOODE, TRAINSET, EquationSolving, OperatorLearning PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 13452ba5dc..c1e00a2d4c 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -8,6 +8,35 @@ function TRAINSET(input_data, output_data; isu0 = false) TRAINSET(input_data, output_data, isu0) end +mutable struct PINOPhi{C, T, U, S} + chain::C + t0::T + u0::U + st::S + function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) + new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) + end +end + +struct PINOsolution{} + predict::Array + res::SciMLBase.OptimizationSolution + phi::PINOPhi + input_data_set::Array +end + +abstract type PINOPhases end +struct OperatorLearning <: PINOPhases + is_data_loss::Bool + is_physics_loss::Bool +end +function OperatorLearning(; is_data_loss = true, is_physics_loss = true) + OperatorLearning(is_data_loss, is_physics_loss) +end +struct EquationSolving <: PINOPhases + pino_solution::PINOsolution +end + """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), @@ -15,6 +44,7 @@ end is_data_loss =true, is_physics_loss =true, init_params, + #TODO update docstring kwargs...) The method is that combine training data and physics constraints @@ -41,37 +71,20 @@ struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm chain::C opt::O train_set::TRAINSET - is_data_loss::Bool - is_physics_loss::Bool + pino_phase::PINOPhases init_params::P kwargs::K end function PINOODE(chain, opt, - train_set; - is_data_loss = true, - is_physics_loss = true, + train_set, + pino_phase; init_params = nothing, kwargs...) - #TODO fnn transform + #TODO fnn transform check !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, train_set, is_data_loss, is_physics_loss, init_params, kwargs) -end - -mutable struct PINOPhi{C, T, U, S} - chain::C - t0::T - u0::U - st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) - new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) - end -end -struct PINOsolution{} - predict::Array - res::SciMLBase.OptimizationSolution - phi::PINOPhi + PINOODE(chain, opt, train_set, pino_phase, init_params, kwargs) end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, @@ -125,10 +138,10 @@ function physics_loss(phi::PINOPhi{C, T, U}, input_data_set) where {C, T, U} prob_set, _ = train_set.input_data, train_set.output_data f = prob_set[1].f + p = prob_set[1].p out_ = phi(input_data_set, θ) ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), ts) if train_set.isu0 == true - p = prob_set[1].p fs = f.f.(out_, p, ts) else ps = [prob.p for prob in prob_set] @@ -157,7 +170,7 @@ function data_loss(phi::PINOPhi{C, T, U}, l₂loss(phi(input_data_set, θ), output_data) end -function generate_data(ts, prob_set, isu0) +function generate_data(ts, prob_set::Vector{ODEProblem}, isu0) batch_size = size(prob_set)[1] instances_size = size(ts)[2] dims = isu0 ? length(prob_set[1].u0) + 1 : length(prob_set[1].p) + 1 @@ -165,7 +178,7 @@ function generate_data(ts, prob_set, isu0) for (i, prob) in enumerate(prob_set) u0 = prob.u0 p = prob.p - f = prob.f + # f = prob.f if isu0 == true in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) else @@ -185,8 +198,9 @@ end function generate_loss( phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts, - is_data_loss, is_physics_loss) where { + pino_phase::OperatorLearning) where { C, T, U} + is_data_loss, is_physics_loss = pino_phase.is_data_loss, pino_phase.is_physics_loss function loss(θ, _) if is_data_loss data_loss(phi, θ, train_set, input_data_set) @@ -202,6 +216,31 @@ function generate_loss( return loss end +function finetune_loss(phi::PINOPhi{C, T, U}, + θ, + train_set::TRAINSET, + input_data_set, + pino_phase::EquationSolving) where {C, T, U} + _, output_data = train_set.input_data, train_set.output_data + output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data) + pino_solution = pino_phase.pino_solution + learned_operator = pino_solution.phi + predict = learned_operator(input_data_set, pino_solution.res.u) + l₂loss(phi(input_data_set, θ), predict) +end + +function generate_loss( + phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts, + pino_phase::EquationSolving) where { + C, T, U} + a = 1 / 100 + function loss(θ, _) + physics_loss(phi, θ, ts, train_set, input_data_set) + + a * finetune_loss(phi, θ, train_set, input_data_set, pino_phase) + end + return loss +end + function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, alg::PINOODE, args...; @@ -212,18 +251,16 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, saveat = nothing, maxiters = nothing) tspan = prob.tspan - t0 = tspan[1] + t0, t_end = tspan[1], tspan[2] u0 = prob.u0 + p = prob.p # f = prob.f - # p = prob.p # param_estim = alg.param_estim chain = alg.chain opt = alg.opt init_params = alg.init_params - is_data_loss = alg.is_data_loss - is_physics_loss = alg.is_physics_loss - + pino_phase = alg.pino_phase # mapping between functional space of some vararible 'a' of equation (for example initial # condition {u(t0 x)} or parameter p) and solution of equation u(t) train_set = alg.train_set @@ -231,23 +268,26 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - t0 = tspan[1] - t_end = tspan[2] instances_size = size(train_set.output_data)[2] range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) - prob_set, _ = train_set.input_data, train_set.output_data + prob_set, output_set = train_set.input_data, train_set.output_data isu0 = train_set.isu0 input_data_set = generate_data(ts, prob_set, isu0) + # input_data_set = if pino_phase == EquationSolving + # generate_data(ts, [prob], isu0) + # elseif pino_phase == OperatorLearning + # generate_data(ts, prob_set, isu0) + # else + # error("pino_phase should be EquationSolving or OperatorLearning") + # end - if isu0 + if isu0 #TODO remove the block u0 = input_data_set[2:end, :, :] - phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) else u0 = prob.u0 - phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) end - + phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) init_params = ComponentArrays.ComponentArray(init_params) isinplace(prob) && @@ -263,8 +303,19 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end end - total_loss = generate_loss( - phi, train_set, input_data_set, ts, is_data_loss, is_physics_loss) + if pino_phase isa EquationSolving + #TODO bad code rewrite,the parameter must uniquely match the index + #TODO doenst need TRAINSET for EquationSolving + find(as, a) = findfirst(x -> isapprox(x.p, a.p), as) + index = find(prob_set, prob) + input_data_set = input_data_set[:, :, [index]] + train_set = TRAINSET(prob_set[index:index], output_set[:, :, [index]], isu0) + total_loss = generate_loss(phi, train_set, input_data_set, ts, pino_phase) + elseif pino_phase isa OperatorLearning + total_loss = generate_loss(phi, train_set, input_data_set, ts, pino_phase) + else + error("pino_phase should be EquationSolving or OperatorLearning") + end # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() @@ -282,5 +333,5 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) predict = phi(input_data_set, res.u) - PINOsolution(predict, res, phi) + PINOsolution(predict, res, phi, input_data_set) end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 4df60b2bb8..9c53e528a7 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -2,7 +2,7 @@ using Test using OrdinaryDiffEq, OptimizationOptimisers using Lux using Statistics, Random -#using NeuralOperators +# using NeuralOperators using NeuralPDE @testset "Example p" begin @@ -33,8 +33,9 @@ using NeuralPDE * input data: set of parameters 'a' * output data: set of solutions u(t){a} corresponding parameter 'a'. """ - train_set = TRAINSET(prob_set, u_output_); - prob = ODEProblem(linear, u0, tspan, 0) + train_set = TRAINSET(prob_set, u_output_) + p = pi / 2 + prob = ODEProblem(linear, u0, tspan, p) chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), @@ -44,14 +45,28 @@ using NeuralPDE # flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), # σ = gelu) opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE( - chain, opt, train_set; is_data_loss = true, is_physics_loss = true) + # pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true) + pino_phase = OperatorLearning( + is_data_loss = true, is_physics_loss = true) + alg = PINOODE(chain, opt, train_set, pino_phase) + # pino_solution = learn() pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ @test ground≈predict atol=1 + + pino_phase = EquationSolving(pino_solution) + alg = PINOODE(chain, opt, train_set, pino_phase) + pino_solution = solve(prob, alg, verbose = true, maxiters = 2000) + + find(as, a) = findfirst(x -> isapprox(x.p, a.p), as) + index = find(prob_set, prob) + predict = pino_solution.predict + ground = u_output_[:,:, [index]] + @test ground≈predict atol=0.1 end + @testset "Example u0" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) @@ -91,7 +106,8 @@ end Lux.Dense(16, 32, Lux.σ), Lux.Dense(32, 1)) opt = OptimizationOptimisers.Adam(0.001) - alg = PINOODE(chain, opt, train_set) + pino_phase = OperatorLearning() + alg = PINOODE(chain, opt, train_set, pino_phase) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 24cf7226b5..378ce9ddd4 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -52,7 +52,8 @@ const gpud = gpu_device() Lux.Dense(inner, 1)) ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(chain, opt, train_set; init_params = ps) + pino_phase = OperatorLearning() + alg = PINOODE(chain, opt, train_set, pino_phase; init_params = ps) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict |> cpu ground = u_output_ |> cpu @@ -108,8 +109,9 @@ end ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.001) + pino_phase = OperatorLearning() alg = PINOODE( - chain, opt, train_set; init_params = ps, is_data_loss = true, is_physics_loss = true) + chain, opt, train_set, pino_phase; init_params = ps) pino_solution = solve(prob, alg, verbose = false, maxiters = 4000) predict = pino_solution.predict |> cpu ground = u_output_ From c38a78faa57ef6670372de8504850c7253caba36 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 22 Mar 2024 17:46:13 +0400 Subject: [PATCH 040/153] fine tunnning update --- docs/src/manual/pino_ode.md | 21 ++++ docs/src/tutorials/pino_ode.md | 68 ++++++++--- src/pino_ode_solve.jl | 215 +++++++++++++++++++-------------- test/PINO_ode_tests.jl | 102 ++++++++++------ test/PINO_ode_tests_gpu.jl | 64 +++++++--- 5 files changed, 308 insertions(+), 162 deletions(-) create mode 100644 docs/src/manual/pino_ode.md diff --git a/docs/src/manual/pino_ode.md b/docs/src/manual/pino_ode.md new file mode 100644 index 0000000000..379f84c46b --- /dev/null +++ b/docs/src/manual/pino_ode.md @@ -0,0 +1,21 @@ +# Physics-Informed Neural operator for solve ODEs + +```@docs +PINOODE +``` + +```@docs +TRAINSET +``` + +```@docs +PINOsolution +``` + +```@docs +OperatorLearning +``` + +```@docs +EquationSolving +``` diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 721b63c344..2e0ab0ed43 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -2,21 +2,24 @@ This tutorial is an introduction to using physics-informed neural operator (PINOs) for solving family of parametric ordinary diferential equations (ODEs). +#TODO two phase -## Solving a family of parametric ODE. +## Operator Learning for a family of parametric ODE. ```@example pino using Test using OrdinaryDiffEq, OptimizationOptimisers using Lux using Statistics, Random -using NeuralOperators +# using NeuralOperators using NeuralPDE linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) u0 = 0.0f0 +p = pi / 2f0 +prob = ODEProblem(linear, u0, tspan, p) ``` Generate a dataset for learning a given family of ODEs where the parameter 'a' is varied. The dataset is generated by solving the ODE for different values of 'a' and storing the solutions. The dataset is then used to train the PINO model: @@ -34,24 +37,33 @@ as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] u_output_ = zeros(Float32, 1, instances_size, batch_size) prob_set = [] for (i, a_i) in enumerate(as) - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) - sol1 = solve(prob, Tsit5(); saveat = 0.0204) + prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) + sol1 = solve(prob_, Tsit5(); saveat = 0.0204) reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob) + push!(prob_set, prob_) u_output_[:, :, i] = reshape_sol end -train_set = TRAINSET(prob_set, u_output_); +train_set = TRAINSET(prob_set, u_output_) ``` -Here it used the PINO method to train the given family of parametric ODEs. +Here it used the PINO method to learning operator of the given family of parametric ODEs. ```@example pino -prob = ODEProblem(linear, u0, tspan, 0) -flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), - σ = gelu) -opt = OptimizationOptimisers.Adam(0.03) -alg = PINOODE(flat_no, opt, train_set; is_data_loss = true, is_physics_loss = true) -pino_solution = solve(prob, alg, verbose = false, maxiters = 1000) +chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 32, Lux.σ), + Lux.Dense(32, 1)) +# flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), +# σ = gelu) + +opt = OptimizationOptimisers.Adam(0.01) +pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true) + +alg = PINOODE(chain, opt, pino_phase) +pino_solution = solve( + prob, alg, verbose = false, maxiters = 3000) predict = pino_solution.predict ground = u_output_ ``` @@ -59,7 +71,35 @@ ground = u_output_ Now let's compare the predictions from the learned operator with the ground truth solution which is obtained early by numerically solving the parametric ODE. Where 'i' is the index of the parameter 'a' in the dataset. ```@example pino -i=1 +using Plots +i=45 plot(predict[1, :, i], label = "Predicted") plot!(ground[1, :, i], label = "Ground truth") +``` + +Now to move on the stage of solving a certain equation using a trained operator and physics + +## Solve ODE using learned operator family of parametric ODE for fine tuning. +```@example pino +dt = (t_end - t0) / instances_size +pino_phase = EquationSolving(dt, pino_solution) +chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 1)) +alg = PINOODE(chain, opt, pino_phase) +fine_tune_solution = solve( prob, alg, verbose = false, maxiters = 2000) + +fine_tune_predict = fine_tune_solution.predict +operator_predict = pino_solution.phi( + fine_tune_solution.input_data_set, pino_solution.res.u) +ground_fine_tune = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :]) +``` + +Compare prediction with ground truth. + +```@example pino +plot(operator_predict[1, :, 1], label = "operator_predict") +plot!(fine_tune_predict[1, :, 1], label = "fine_tune_predict") +plot!(ground_fine_tune[1, :, 1], label = "Ground truth") ``` \ No newline at end of file diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index c1e00a2d4c..15fcd4eb64 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,11 +1,12 @@ +""" +TRAINSET(input_data, output_data; isu0 = false) +## Positional Arguments +* input_data: variables set 'a' of equation (for example initial condition {u(t0 x)} or parameter p +* output_data: set of solutions u(t){a} corresponding parameter 'a'. +""" struct TRAINSET{} input_data::Vector{ODEProblem} output_data::Array - isu0::Bool -end - -function TRAINSET(input_data, output_data; isu0 = false) - TRAINSET(input_data, output_data, isu0) end mutable struct PINOPhi{C, T, U, S} @@ -18,6 +19,14 @@ mutable struct PINOPhi{C, T, U, S} end end +""" +PINOsolution(predict, res, phi, input_data_set) +## Positional Arguments +* predict: The predicted solution. +* res: The optimization solution. +* phi: The solution operator. +* input_data_set: The input data set. +""" struct PINOsolution{} predict::Array res::SciMLBase.OptimizationSolution @@ -26,25 +35,48 @@ struct PINOsolution{} end abstract type PINOPhases end + +""" +OperatorLearning(train_set; is_data_loss = true, is_physics_loss = true) +## Positional Arguments +* train_set: Contains 'input data' - set of parameters 'a' and output data - set of solutions +* u(t){a} corresponding initial conditions 'u0'. + +## Keyword Arguments +* is_data_loss: Includes or off a loss function for training on the data set. +* is_physics_loss: Includes or off loss function training on physics-informed approach. +""" struct OperatorLearning <: PINOPhases + train_set::TRAINSET is_data_loss::Bool is_physics_loss::Bool end -function OperatorLearning(; is_data_loss = true, is_physics_loss = true) - OperatorLearning(is_data_loss, is_physics_loss) +function OperatorLearning(train_set; is_data_loss = true, is_physics_loss = true) + OperatorLearning(train_set, is_data_loss, is_physics_loss) end + +""" +EquationSolving(dt, pino_solution) +## Positional Arguments +* dt: The time step. +* pino_solution: Contains the solution of the operator learning phase. +""" struct EquationSolving <: PINOPhases + dt::Number pino_solution::PINOsolution + is_finetune_loss::Bool + is_physics_loss::Bool +end + +function EquationSolving(dt, pino_solution; is_finetune_loss = true, is_physics_loss = true) + EquationSolving(dt, pino_solution, is_finetune_loss, is_physics_loss) end """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), - train_set, - is_data_loss =true, - is_physics_loss =true, + pino_phase; init_params, - #TODO update docstring kwargs...) The method is that combine training data and physics constraints @@ -54,14 +86,12 @@ to learn the solution operator of a given family of parametric Ordinary Differen * `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`. `Flux.Chain` will be converted to `Lux` using `Lux.transform`. * `opt`: The optimizer to train the neural network. -* `train_set`: Contains 'input data' - sr of parameters 'a' and output data - set of solutions - u(t){a} corresponding initial conditions 'u0'. +* `pino_phase`: The phase of the PINN algorithm, either `OperatorLearning` or `EquationSolving`. ## Keyword Arguments -* `is_data_loss` Includes or off a loss function for training on the data set. -* `is_physics_loss`: Includes or off loss function training on physics-informed approach. * `init_params`: The initial parameter of the neural network. By default, this is `nothing` which thus uses the random initialization provided by the neural network library. +* isu0: If true, the input data set contains initial conditions 'u0'. * `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. ## References @@ -70,21 +100,21 @@ Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Eq struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm chain::C opt::O - train_set::TRAINSET pino_phase::PINOPhases init_params::P + isu0::Bool kwargs::K end function PINOODE(chain, opt, - train_set, pino_phase; init_params = nothing, + isu0 = false, kwargs...) #TODO fnn transform check !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, train_set, pino_phase, init_params, kwargs) + PINOODE(chain, opt, pino_phase, init_params, isu0, kwargs) end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, @@ -132,23 +162,25 @@ function l₂loss(𝐲̂, 𝐲) end function physics_loss(phi::PINOPhi{C, T, U}, + prob::ODEProblem, θ, ts::AbstractArray, - train_set::TRAINSET, - input_data_set) where {C, T, U} - prob_set, _ = train_set.input_data, train_set.output_data - f = prob_set[1].f - p = prob_set[1].p + prob_set::Array, + input_data_set, + isu0:: Bool) where {C, T, U} + f = prob.f out_ = phi(input_data_set, θ) ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), ts) - if train_set.isu0 == true + if isu0 == true + p = prob.p fs = f.f.(out_, p, ts) else ps = [prob.p for prob in prob_set] - if p isa Number + first_p = ps[1] + if first_p isa Number fs = cat( [f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3) - elseif p isa Vector + elseif first_p isa Vector fs = cat( [reduce( hcat, [f.f(out_[:, j, [i]], p, ts) for j in axes(out_[:, :, [i]], 2)]) @@ -163,14 +195,13 @@ end function data_loss(phi::PINOPhi{C, T, U}, θ, - train_set::TRAINSET, + output_data::Array, input_data_set) where {C, T, U} - _, output_data = train_set.input_data, train_set.output_data output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data) l₂loss(phi(input_data_set, θ), output_data) end -function generate_data(ts, prob_set::Vector{ODEProblem}, isu0) +function generate_data(ts, prob_set, isu0) batch_size = size(prob_set)[1] instances_size = size(ts)[2] dims = isu0 ? length(prob_set[1].u0) + 1 : length(prob_set[1].p) + 1 @@ -196,33 +227,10 @@ function generate_data(ts, prob_set::Vector{ODEProblem}, isu0) input_data_set end -function generate_loss( - phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts, - pino_phase::OperatorLearning) where { - C, T, U} - is_data_loss, is_physics_loss = pino_phase.is_data_loss, pino_phase.is_physics_loss - function loss(θ, _) - if is_data_loss - data_loss(phi, θ, train_set, input_data_set) - elseif is_physics_loss - physics_loss(phi, θ, ts, train_set, input_data_set) - elseif is_data_loss && is_physics_loss - data_loss(phi, θ, train_set, input_data_set) + - physics_loss(phi, θ, ts, train_set, input_data_set) - else - error("data loss or physics loss should be true") - end - end - return loss -end - function finetune_loss(phi::PINOPhi{C, T, U}, θ, - train_set::TRAINSET, input_data_set, pino_phase::EquationSolving) where {C, T, U} - _, output_data = train_set.input_data, train_set.output_data - output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data) pino_solution = pino_phase.pino_solution learned_operator = pino_solution.phi predict = learned_operator(input_data_set, pino_solution.res.u) @@ -230,13 +238,44 @@ function finetune_loss(phi::PINOPhi{C, T, U}, end function generate_loss( - phi::PINOPhi{C, T, U}, train_set::TRAINSET, input_data_set, ts, - pino_phase::EquationSolving) where { + phi::PINOPhi{C, T, U}, prob::ODEProblem, prob_set::Array, input_data_set, ts, + pino_phase::EquationSolving, isu0::Bool) where { C, T, U} a = 1 / 100 + is_finetune_loss, is_physics_loss = pino_phase.is_finetune_loss, + pino_phase.is_physics_loss function loss(θ, _) - physics_loss(phi, θ, ts, train_set, input_data_set) + - a * finetune_loss(phi, θ, train_set, input_data_set, pino_phase) + if is_finetune_loss + finetune_loss(phi, θ, input_data_set, pino_phase) + elseif is_physics_loss + physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) + elseif is_finetune_loss && is_physics_loss + physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) + + a * finetune_loss(phi, θ, input_data_set, pino_phase) + else + error("finetune loss or physics loss should be true") + end + end + return loss +end + +function generate_loss( + phi::PINOPhi{C, T, U},prob::ODEProblem, train_set::TRAINSET, input_data_set, ts, + pino_phase::OperatorLearning, isu0::Bool) where { + C, T, U} + is_data_loss, is_physics_loss = pino_phase.is_data_loss, pino_phase.is_physics_loss + prob_set, output_data = train_set.input_data, train_set.output_data + function loss(θ, _) + if is_data_loss + data_loss(phi, θ, output_data, input_data_set) + elseif is_physics_loss + physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) + elseif is_data_loss && is_physics_loss + data_loss(phi, θ, output_data, input_data_set) + + physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) + else + error("data loss or physics loss should be true") + end end return loss end @@ -244,45 +283,40 @@ end function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, alg::PINOODE, args...; - # dt = nothing, abstol = 1.0f-6, reltol = 1.0f-3, verbose = false, saveat = nothing, maxiters = nothing) - tspan = prob.tspan + @unpack tspan, u0, p, f = prob t0, t_end = tspan[1], tspan[2] - u0 = prob.u0 - p = prob.p - # f = prob.f - # param_estim = alg.param_estim - - chain = alg.chain - opt = alg.opt - init_params = alg.init_params - pino_phase = alg.pino_phase - # mapping between functional space of some vararible 'a' of equation (for example initial - # condition {u(t0 x)} or parameter p) and solution of equation u(t) - train_set = alg.train_set + @unpack chain, opt, pino_phase, init_params, isu0 = alg !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - instances_size = size(train_set.output_data)[2] - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - prob_set, output_set = train_set.input_data, train_set.output_data - isu0 = train_set.isu0 - input_data_set = generate_data(ts, prob_set, isu0) - # input_data_set = if pino_phase == EquationSolving - # generate_data(ts, [prob], isu0) - # elseif pino_phase == OperatorLearning - # generate_data(ts, prob_set, isu0) - # else - # error("pino_phase should be EquationSolving or OperatorLearning") - # end - - if isu0 #TODO remove the block + if pino_phase isa EquationSolving + @unpack dt, pino_solution = pino_phase + range_ = collect(t0:dt:t_end) + ts = reshape(collect(range_), 1, size(range_)[1]) + input_data_set = generate_data(ts, [prob], isu0) + if isu0 + pino_solution.phi.u0 = input_data_set[2:end, :, :] + end + elseif pino_phase isa OperatorLearning + # mapping between functional space of some vararible 'a' of equation (for example initial + # condition {u(t0 x)} or parameter p) and solution of equation u(t) + train_set = pino_phase.train_set + instances_size = size(train_set.output_data)[2] + range_ = range(t0, stop = t_end, length = instances_size) + ts = reshape(collect(range_), 1, instances_size) + prob_set, _ = train_set.input_data, train_set.output_data + input_data_set = generate_data(ts, prob_set, isu0) + else + error("pino_phase should be EquationSolving or OperatorLearning") + end + + if isu0 u0 = input_data_set[2:end, :, :] else u0 = prob.u0 @@ -304,15 +338,10 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end if pino_phase isa EquationSolving - #TODO bad code rewrite,the parameter must uniquely match the index - #TODO doenst need TRAINSET for EquationSolving - find(as, a) = findfirst(x -> isapprox(x.p, a.p), as) - index = find(prob_set, prob) - input_data_set = input_data_set[:, :, [index]] - train_set = TRAINSET(prob_set[index:index], output_set[:, :, [index]], isu0) - total_loss = generate_loss(phi, train_set, input_data_set, ts, pino_phase) + prob_set = [prob] + total_loss = generate_loss(phi, prob, prob_set, input_data_set, ts, pino_phase,isu0) elseif pino_phase isa OperatorLearning - total_loss = generate_loss(phi, train_set, input_data_set, ts, pino_phase) + total_loss = generate_loss(phi, prob, train_set, input_data_set, ts, pino_phase, isu0) else error("pino_phase should be EquationSolving or OperatorLearning") end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 9c53e528a7..628e68c998 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -10,7 +10,9 @@ using NeuralPDE linear = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) u0 = 0.0f0 - #generate data set + p = pi / 2 + prob = ODEProblem(linear, u0, tspan, p) + #generate data t0, t_end = tspan instances_size = 50 range_ = range(t0, stop = t_end, length = instances_size) @@ -21,58 +23,65 @@ using NeuralPDE u_output_ = zeros(Float32, 1, instances_size, batch_size) prob_set = [] for (i, a_i) in enumerate(as) - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) - sol1 = solve(prob, Tsit5(); saveat = 0.0204) + prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) + sol1 = solve(prob_, Tsit5(); saveat = 0.0204) reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob) + push!(prob_set, prob_) u_output_[:, :, i] = reshape_sol end - """ - Set of training data: - * input data: set of parameters 'a' + Training data: + * input data: set of parameters 'a', * output data: set of solutions u(t){a} corresponding parameter 'a'. - """ + """ train_set = TRAINSET(prob_set, u_output_) - p = pi / 2 - prob = ODEProblem(linear, u0, tspan, p) + + # operator learning phase + # init neural network chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1)) # flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), # σ = gelu) - opt = OptimizationOptimisers.Adam(0.03) - # pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true) - pino_phase = OperatorLearning( - is_data_loss = true, is_physics_loss = true) - alg = PINOODE(chain, opt, train_set, pino_phase) - # pino_solution = learn() - pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) + + opt = OptimizationOptimisers.Adam(0.01) + pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true) + + alg = PINOODE(chain, opt, pino_phase) + pino_solution = solve( + prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ @test ground≈predict atol=1 + # equation solving phase + dt = (t_end - t0) / instances_size + pino_phase = EquationSolving(dt, pino_solution) + chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 1)) + alg = PINOODE(chain, opt, pino_phase) + fine_tune_solution = solve( + prob, alg, verbose = false, maxiters = 2000) - pino_phase = EquationSolving(pino_solution) - alg = PINOODE(chain, opt, train_set, pino_phase) - pino_solution = solve(prob, alg, verbose = true, maxiters = 2000) - - find(as, a) = findfirst(x -> isapprox(x.p, a.p), as) - index = find(prob_set, prob) - predict = pino_solution.predict - ground = u_output_[:,:, [index]] - @test ground≈predict atol=0.1 + fine_tune_predict = fine_tune_solution.predict + operator_predict = pino_solution.phi( + fine_tune_solution.input_data_set, pino_solution.res.u) + ground = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :]) + @test ground≈fine_tune_predict atol=0.1 + @test operator_predict≈fine_tune_predict rtol=0.1 end - @testset "Example u0" begin linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) linear = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) p = Float32(pi) u0 = 2.0f0 + prob = ODEProblem(linear, u0, tspan, p) #generate data set t0, t_end = tspan instances_size = 50 @@ -83,33 +92,48 @@ end u_output_ = zeros(Float32, 1, instances_size, batch_size) prob_set = [] for (i, u0_i) in enumerate(u0s) - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) - sol1 = solve(prob, Tsit5(); saveat = 0.0204) + prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) + sol1 = solve(prob_, Tsit5(); saveat = 0.0204) reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob) + push!(prob_set, prob_) u_output_[:, :, i] = reshape_sol end - + # operator learning phase """ Set of training data: * input data: set of initial conditions 'u0' * output data: set of solutions u(t){u0} corresponding initial conditions 'u0'. """ - train_set = TRAINSET(prob_set, u_output_; isu0 = true) - #TODO we argument u0 but dont actually use u0 because we use only set of u0 for generate train set from prob_set - prob = ODEProblem(linear, 0.0f0, tspan, p) + train_set = TRAINSET(prob_set, u_output_) # fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1)) opt = OptimizationOptimisers.Adam(0.001) - pino_phase = OperatorLearning() - alg = PINOODE(chain, opt, train_set, pino_phase) + pino_phase = OperatorLearning(train_set) + alg = PINOODE(chain, opt, pino_phase; isu0 = true) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict ground = u_output_ - @test ground≈predict atol=1.0 + @test ground≈predict atol=1. + + # equation solving phase + dt = (t_end -t0) / instances_size + pino_phase = EquationSolving(dt, pino_solution) + chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 1)) + alg = PINOODE(chain, opt, pino_phase; isu0 = true) + fine_tune_solution = solve(prob, alg, verbose = false, maxiters = 2000) + + fine_tune_predict = fine_tune_solution.predict + operator_predict = pino_solution.phi( + fine_tune_solution.input_data_set, pino_solution.res.u) + ground_fine_tune = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :]) + @test ground_fine_tune≈fine_tune_predict atol=1 + @test operator_predict≈fine_tune_predict rtol=0.1 end diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl index 378ce9ddd4..919e6012ed 100644 --- a/test/PINO_ode_tests_gpu.jl +++ b/test/PINO_ode_tests_gpu.jl @@ -16,21 +16,23 @@ const gpud = gpu_device() linear = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) u0 = 0.0f0 + p = pi / 2.0f0 + prob = ODEProblem(linear, u0, tspan, p) #generate data set t0, t_end = tspan instances_size = 50 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) batch_size = 50 - as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] + as = [Float32(i) for i in range(0.1, stop = pi / 2.0f0, length = batch_size)] u_output_ = zeros(Float32, 1, instances_size, batch_size) prob_set = [] for (i, a_i) in enumerate(as) - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) - sol1 = solve(prob, Tsit5(); saveat = 0.0204) + prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) + sol1 = solve(prob_, Tsit5(); saveat = 0.0204) reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob) + push!(prob_set, prob_) u_output_[:, :, i] = reshape_sol end u_output_ = u_output_ |> gpud @@ -42,7 +44,6 @@ const gpud = gpu_device() """ train_set = TRAINSET(prob_set, u_output_) - prob = ODEProblem(linear, u0, tspan, 0) inner = 50 chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ), Lux.Dense(inner, inner, Lux.σ), @@ -52,12 +53,26 @@ const gpud = gpu_device() Lux.Dense(inner, 1)) ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.03) - pino_phase = OperatorLearning() - alg = PINOODE(chain, opt, train_set, pino_phase; init_params = ps) + pino_phase = OperatorLearning(train_set) + alg = PINOODE(chain, opt, pino_phase; init_params = ps) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict |> cpu ground = u_output_ |> cpu @test ground≈predict atol=1 + + dt = (t_end - t0) / instances_size + pino_phase = EquationSolving(dt, pino_solution) + alg = PINOODE(chain, opt, pino_phase; init_params = ps) + fine_tune_solution = solve( + prob, alg, verbose = false, maxiters = 2000) + + fine_tune_predict = fine_tune_solution.predict |> cpu + operator_predict = pino_solution.phi( + fine_tune_solution.input_data_set, pino_solution.res.u) |> cpu + input_data_set_ = fine_tune_solution.input_data_set[[1], :, :] |> cpu + ground_fine_tune = linear_analytic.(u0, p, input_data_set_) + @test ground_fine_tune≈fine_tune_predict atol=0.5 + @test operator_predict≈fine_tune_predict rtol=0.1 end @testset "lotka volterra" begin @@ -76,26 +91,26 @@ end p = Float32[1.5, 1.0, 3.0, 1.0] tspan = (0.0f0, 4.0f0) dt = 0.01f0 + prob = ODEProblem(lotka_volterra, u0, tspan, p) t0, t_end = tspan - instances_size = 100 range_ = range(t0, stop = t_end, length = instances_size) ts = reshape(collect(range_), 1, instances_size) batch_size = 50 - ps = [p .+ i * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size] + ps = [p .+ (i - 1) * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size] u_output_ = zeros(Float32, 2, instances_size, batch_size) prob_set = [] for (i, p_i) in enumerate(ps) - prob = ODEProblem(lotka_volterra, u0, tspan, p_i) - solution = solve(prob, Tsit5(); saveat = dt) + prob_ = ODEProblem(lotka_volterra, u0, tspan, p_i) + solution = solve(prob_, Tsit5(); saveat = dt) reshape_sol_ = reduce(hcat, solution(range_).u) reshape_sol = Float32.(reshape(reshape_sol_, 2, instances_size, 1)) - push!(prob_set, prob) + push!(prob_set, prob_) u_output_[:, :, i] = reshape_sol end train_set = TRAINSET(prob_set, u_output_) - prob = ODEProblem(lotka_volterra, u0, tspan, p) + # flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,), # σ = gelu) # flat_no = Lux.transform(flat_no) @@ -109,11 +124,28 @@ end ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud opt = OptimizationOptimisers.Adam(0.001) - pino_phase = OperatorLearning() - alg = PINOODE( - chain, opt, train_set, pino_phase; init_params = ps) + pino_phase = OperatorLearning(train_set) + alg = PINOODE(chain, opt, pino_phase; init_params = ps) pino_solution = solve(prob, alg, verbose = false, maxiters = 4000) predict = pino_solution.predict |> cpu ground = u_output_ @test ground≈predict atol=5 + + dt = (t_end - t0) / instances_size + pino_phase = EquationSolving(dt, pino_solution; is_finetune_loss = true,is_physics_loss = true) + chain = Lux.Chain(Lux.Dense(5, 16, Lux.σ), + Lux.Dense(16, 16, Lux.σ), + Lux.Dense(16, 32, Lux.σ), + Lux.Dense(32, 2)) + ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud + alg = PINOODE(chain, opt, pino_phase; init_params = ps) + fine_tune_solution = solve(prob, alg, verbose = false, maxiters = 2000) + + fine_tune_predict = fine_tune_solution.predict |> cpu + operator_predict = pino_solution.phi( + fine_tune_solution.input_data_set, pino_solution.res.u) |> cpu + input_data_set_ = fine_tune_solution.input_data_set[[1], :, :] |> cpu + ground_fine_tune = u_output_[:, :, [1]] + @test ground_fine_tune ≈ fine_tune_predict[:, 1:100, :] atol = 3 + @test operator_predict≈fine_tune_predict rtol=0.1 end From 1e38676f785ad0d8ad86f708f1d9686e1a9f8ee7 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 22 Mar 2024 18:01:12 +0400 Subject: [PATCH 041/153] fix --- test/PINO_ode_tests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 628e68c998..1a75690862 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -71,7 +71,7 @@ using NeuralPDE operator_predict = pino_solution.phi( fine_tune_solution.input_data_set, pino_solution.res.u) ground = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :]) - @test ground≈fine_tune_predict atol=0.1 + @test ground≈fine_tune_predict atol=1. @test operator_predict≈fine_tune_predict rtol=0.1 end From 2cc1d1f299498b108cc68581d59bd471c26a5815 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 25 Apr 2024 19:00:15 +0400 Subject: [PATCH 042/153] implement DeepONet, refactor pinoode --- src/NeuralPDE.jl | 3 +- src/neural_operators.jl | 47 ++++++ src/pino_ode_solve.jl | 364 ++++++++++++---------------------------- test/PINO_ode_tests.jl | 199 +++++++++------------- 4 files changed, 237 insertions(+), 376 deletions(-) create mode 100644 src/neural_operators.jl diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index bd04c036ab..d9f401151d 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -48,6 +48,7 @@ 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") @@ -58,7 +59,7 @@ include("PDE_BPINN.jl") include("dgm.jl") -export NNODE, NNDAE, PINOODE, TRAINSET, EquationSolving, OperatorLearning +export NNODE, NNDAE, PINOODE, DeepONet PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, diff --git a/src/neural_operators.jl b/src/neural_operators.jl new file mode 100644 index 0000000000..7544675832 --- /dev/null +++ b/src/neural_operators.jl @@ -0,0 +1,47 @@ +#TODO: Add docstrings +""" +DeepONet(branch,trunk) +""" +struct DeepONet{} <: Lux.AbstractExplicitLayer + branch::Lux.AbstractExplicitLayer + trunk::Lux.AbstractExplicitLayer +end + +function Lux.setup(rng::AbstractRNG, l::DeepONet) + branch, trunk = l.branch, l.trunk + θ_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) + θ, st +end + +# function Lux.initialparameters(rng::AbstractRNG, e::DeepONet) +# code +# end + +Lux.initialstates(::AbstractRNG, ::DeepONet) = NamedTuple() + +""" +example: + +branch = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 1)) +trunk = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 1)) +a = rand(1, 100, 10) +t = rand(1, 1, 10) +x = (branch = a, trunk = t) + +deeponet = DeepONet(branch, trunk) +θ, st = Lux.setup(Random.default_rng(), deeponet) +y = deeponet(x, θ, st) +""" +@inline function (f::DeepONet)(x::NamedTuple, θ, st::NamedTuple) + parameters, cord = 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(parameters, θ_branch, st_branch) + out_t, st_t = trunk(cord, θ_trunk, st_trunk) + out = out_b' * out_t + return out, (branch = st_b, trunk = st_t) +end diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 15fcd4eb64..597131ee0d 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,77 +1,3 @@ -""" -TRAINSET(input_data, output_data; isu0 = false) -## Positional Arguments -* input_data: variables set 'a' of equation (for example initial condition {u(t0 x)} or parameter p -* output_data: set of solutions u(t){a} corresponding parameter 'a'. -""" -struct TRAINSET{} - input_data::Vector{ODEProblem} - output_data::Array -end - -mutable struct PINOPhi{C, T, U, S} - chain::C - t0::T - u0::U - st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) - new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) - end -end - -""" -PINOsolution(predict, res, phi, input_data_set) -## Positional Arguments -* predict: The predicted solution. -* res: The optimization solution. -* phi: The solution operator. -* input_data_set: The input data set. -""" -struct PINOsolution{} - predict::Array - res::SciMLBase.OptimizationSolution - phi::PINOPhi - input_data_set::Array -end - -abstract type PINOPhases end - -""" -OperatorLearning(train_set; is_data_loss = true, is_physics_loss = true) -## Positional Arguments -* train_set: Contains 'input data' - set of parameters 'a' and output data - set of solutions -* u(t){a} corresponding initial conditions 'u0'. - -## Keyword Arguments -* is_data_loss: Includes or off a loss function for training on the data set. -* is_physics_loss: Includes or off loss function training on physics-informed approach. -""" -struct OperatorLearning <: PINOPhases - train_set::TRAINSET - is_data_loss::Bool - is_physics_loss::Bool -end -function OperatorLearning(train_set; is_data_loss = true, is_physics_loss = true) - OperatorLearning(train_set, is_data_loss, is_physics_loss) -end - -""" -EquationSolving(dt, pino_solution) -## Positional Arguments -* dt: The time step. -* pino_solution: Contains the solution of the operator learning phase. -""" -struct EquationSolving <: PINOPhases - dt::Number - pino_solution::PINOsolution - is_finetune_loss::Bool - is_physics_loss::Bool -end - -function EquationSolving(dt, pino_solution; is_finetune_loss = true, is_physics_loss = true) - EquationSolving(dt, pino_solution, is_finetune_loss, is_physics_loss) -end - """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), @@ -95,26 +21,38 @@ to learn the solution operator of a given family of parametric Ordinary Differen * `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. ## References -Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" +* Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets" +* Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct PINOODE{C, O, P, K} <: DiffEqBase.AbstractODEAlgorithm +struct PINOODE{C, O, B, I, S, K} <: SciMLBase.AbstractODEAlgorithm chain::C opt::O - pino_phase::PINOPhases - init_params::P + bounds::B + init_params::I isu0::Bool + strategy::S kwargs::K end function PINOODE(chain, opt, - pino_phase; + bounds; init_params = nothing, - isu0 = false, + isu0 = false, #TODOD remove + strategy = nothing, kwargs...) - #TODO fnn transform check !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, pino_phase, init_params, isu0, kwargs) + PINOODE(chain, opt, bounds, init_params, isu0, strategy, kwargs) +end + +mutable struct PINOPhi{C, T, U, S} + chain::C + t0::T + u0::U + st::S + function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) + new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) + end end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, @@ -130,159 +68,76 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, PINOPhi(chain, t0, u0, st), init_params end -function (f::PINOPhi{C, T, U})(t::AbstractArray, - θ) where {C <: Lux.AbstractExplicitLayer, T, U} - y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) - ChainRulesCore.@ignore_derivatives f.st = st - ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[1:size(y)[1], :, :]) - f_ = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0) - f_ .+ (ts .- f.t0) .* y -end - -function dfdx_rand_matrix(phi::PINOPhi, t::AbstractArray, θ) - ε_ = sqrt(eps(eltype(t))) - d = Normal{eltype(t)}(0.0f0, ε_) - size_ = size(t) .- (1, 0, 0) - eps_ = ε_ .+ rand(d, size_) .* ε_ - zeros_ = zeros(eltype(t), size_) - ε = cat(eps_, zeros_, dims = 1) - (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) -end - -function dfdx(phi::PINOPhi, t::AbstractArray, θ) - ε = [sqrt(eps(eltype(t))), zeros(eltype(t), size(t)[1] - 1)...] - (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) -end - -function l₂loss(𝐲̂, 𝐲) - feature_dims = 2:(ndims(𝐲) - 1) - loss = sum(.√(sum(abs2, 𝐲̂ - 𝐲, dims = feature_dims))) - y_norm = sum(.√(sum(abs2, 𝐲, dims = feature_dims))) - return loss / y_norm -end +#TODO update +# function (f::PINOPhi{C, T, U})(t::AbstractArray, +# θ) where {C <: Lux.AbstractExplicitLayer, T, U} +# y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) +# ChainRulesCore.@ignore_derivatives f.st = st +# ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[1:size(y)[1], :, :]) +# u_0 = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0) +# u_0 .+ (ts .- f.t0) .* y +# end -function physics_loss(phi::PINOPhi{C, T, U}, - prob::ODEProblem, - θ, - ts::AbstractArray, - prob_set::Array, - input_data_set, - isu0:: Bool) where {C, T, U} +#TODO C <: DeepONet +function (f::PINOPhi{C, T, U})(x::NamedTuple, θ) where {C, T, U} + y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) + ChainRulesCore.@ignore_derivatives f.st = st + a, t = x.branch, x.trunk + ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t) + u0_ = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0) + u0_ .+ (ts .- f.t0) .* y +end + +# function dfdx(phi::PINOPhi, t::AbstractArray, θ) +# ε = [sqrt(eps(eltype(t))), zeros(eltype(t), size(t)[1] - 1)...] +# (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) +# end + +#TODO C <: DeepONet +function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ) where {C, T, U} + t = x.trunk + ε = [sqrt(eps(eltype(t)))] + phi_trunk(x, θ) = phi.chain.trunk(x, θ.trunk, phi.st.trunk)[1] + du_trunk_ = (phi_trunk(t .+ ε, θ) .- phi_trunk(t, θ)) ./ sqrt(eps(eltype(t))) + u_branch = phi.chain.branch(x.branch, θ.branch, phi.st.branch)[1] + u_branch' .* du_trunk_ +end + +# function l₂loss(𝐲̂, 𝐲) +# feature_dims = 2:(ndims(𝐲) - 1) +# loss = sum(.√(sum(abs2, 𝐲̂ - 𝐲, dims = feature_dims))) +# y_norm = sum(.√(sum(abs2, 𝐲, dims = feature_dims))) +# return loss / y_norm +# end + +function physics_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U} f = prob.f - out_ = phi(input_data_set, θ) - ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), ts) - if isu0 == true - p = prob.p - fs = f.f.(out_, p, ts) - else - ps = [prob.p for prob in prob_set] - first_p = ps[1] - if first_p isa Number - fs = cat( - [f.f.(out_[:, :, [i]], p, ts) for (i, p) in enumerate(ps)]..., dims = 3) - elseif first_p isa Vector - fs = cat( - [reduce( - hcat, [f.f(out_[:, j, [i]], p, ts) for j in axes(out_[:, :, [i]], 2)]) - for (i, p) in enumerate(ps)]..., - dims = 3) - else - error("p should be a number or a vector") - end - end - l₂loss(dfdx(phi, input_data_set, θ), fs) -end - -function data_loss(phi::PINOPhi{C, T, U}, - θ, - output_data::Array, - input_data_set) where {C, T, U} - output_data = adapt(parameterless_type(ComponentArrays.getdata(θ)), output_data) - l₂loss(phi(input_data_set, θ), output_data) + ps , ts = x.branch, x.trunk + norm = size(x.branch)[2] * size(x.trunk)[2] + sum(abs2, dfdx(phi, x, θ) - f.(phi(x, θ), ps, ts)) / norm end -function generate_data(ts, prob_set, isu0) - batch_size = size(prob_set)[1] - instances_size = size(ts)[2] - dims = isu0 ? length(prob_set[1].u0) + 1 : length(prob_set[1].p) + 1 - input_data_set = Array{Float32, 3}(undef, dims, instances_size, batch_size) - for (i, prob) in enumerate(prob_set) - u0 = prob.u0 - p = prob.p - # f = prob.f - if isu0 == true - in_ = reduce(vcat, [ts, fill(u0, 1, size(ts)[2], 1)]) - else - if p isa Number - in_ = reduce(vcat, [ts, fill(p, 1, size(ts)[2], 1)]) - elseif p isa Vector - inner = reduce(vcat, [ts, reduce(hcat, fill(p, 1, size(ts)[2], 1))]) - in_ = reshape(inner, size(inner)..., 1) - else - error("p should be a number or a vector") - end - end - input_data_set[:, :, i] = in_ - end - input_data_set -end - -function finetune_loss(phi::PINOPhi{C, T, U}, - θ, - input_data_set, - pino_phase::EquationSolving) where {C, T, U} - pino_solution = pino_phase.pino_solution - learned_operator = pino_solution.phi - predict = learned_operator(input_data_set, pino_solution.res.u) - l₂loss(phi(input_data_set, θ), predict) -end - -function generate_loss( - phi::PINOPhi{C, T, U}, prob::ODEProblem, prob_set::Array, input_data_set, ts, - pino_phase::EquationSolving, isu0::Bool) where { - C, T, U} - a = 1 / 100 - is_finetune_loss, is_physics_loss = pino_phase.is_finetune_loss, - pino_phase.is_physics_loss - function loss(θ, _) - if is_finetune_loss - finetune_loss(phi, θ, input_data_set, pino_phase) - elseif is_physics_loss - physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) - elseif is_finetune_loss && is_physics_loss - physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) + - a * finetune_loss(phi, θ, input_data_set, pino_phase) - else - error("finetune loss or physics loss should be true") - end - end - return loss +function get_trainset(bounds, tspan, strategy) + #TODO dt -> instances_size + instances_size = 100 + p = range(bounds.p[1], stop = bounds.p[2], length = instances_size) + t = range(tspan[1], stop = tspan[2], length = instances_size) + x = (branch = collect(p)', trunk = collect(t)') + x end -function generate_loss( - phi::PINOPhi{C, T, U},prob::ODEProblem, train_set::TRAINSET, input_data_set, ts, - pino_phase::OperatorLearning, isu0::Bool) where { - C, T, U} - is_data_loss, is_physics_loss = pino_phase.is_data_loss, pino_phase.is_physics_loss - prob_set, output_data = train_set.input_data, train_set.output_data +#TODO GridTraining +function generate_loss(strategy, prob::ODEProblem, phi, bounds, tspan) + x = get_trainset(bounds, tspan, strategy) function loss(θ, _) - if is_data_loss - data_loss(phi, θ, output_data, input_data_set) - elseif is_physics_loss - physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) - elseif is_data_loss && is_physics_loss - data_loss(phi, θ, output_data, input_data_set) + - physics_loss(phi, prob, θ, ts, prob_set, input_data_set, isu0) - else - error("data loss or physics loss should be true") - end + physics_loss(phi, prob, x, θ) end - return loss end -function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, +function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, alg::PINOODE, args...; + dt = nothing, abstol = 1.0f-6, reltol = 1.0f-3, verbose = false, @@ -290,37 +145,15 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, maxiters = nothing) @unpack tspan, u0, p, f = prob t0, t_end = tspan[1], tspan[2] - @unpack chain, opt, pino_phase, init_params, isu0 = alg + @unpack chain, opt, bounds, init_params, isu0 = alg !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - if pino_phase isa EquationSolving - @unpack dt, pino_solution = pino_phase - range_ = collect(t0:dt:t_end) - ts = reshape(collect(range_), 1, size(range_)[1]) - input_data_set = generate_data(ts, [prob], isu0) - if isu0 - pino_solution.phi.u0 = input_data_set[2:end, :, :] - end - elseif pino_phase isa OperatorLearning - # mapping between functional space of some vararible 'a' of equation (for example initial - # condition {u(t0 x)} or parameter p) and solution of equation u(t) - train_set = pino_phase.train_set - instances_size = size(train_set.output_data)[2] - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - prob_set, _ = train_set.input_data, train_set.output_data - input_data_set = generate_data(ts, prob_set, isu0) - else - error("pino_phase should be EquationSolving or OperatorLearning") + if !any(in(keys(bounds)), (:u0, :p)) + error("bounds should contain u0 or p only") end - if isu0 - u0 = input_data_set[2:end, :, :] - else - u0 = prob.u0 - end phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) init_params = ComponentArrays.ComponentArray(init_params) @@ -328,7 +161,8 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - phi(input_data_set, init_params) + x = (branch = rand(length(bounds), 10), trunk = rand(1, 10)) + phi(x, init_params) catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of input data and chain should match")) @@ -337,13 +171,17 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end end - if pino_phase isa EquationSolving - prob_set = [prob] - total_loss = generate_loss(phi, prob, prob_set, input_data_set, ts, pino_phase,isu0) - elseif pino_phase isa OperatorLearning - total_loss = generate_loss(phi, prob, train_set, input_data_set, ts, pino_phase, isu0) - else - error("pino_phase should be EquationSolving or OperatorLearning") + strategy = nothing + + inner_f = generate_loss(strategy, prob, phi, bounds, tspan) + + function total_loss(θ, _) + inner_f(θ, nothing) + #TODO add loss + # L2_loss = inner_f(θ, nothing) + # if !(additional_loss isa Nothing) + # L2_loss = L2_loss + additional_loss(phi, θ) + # end end # Optimization Algo for Training Strategies @@ -361,6 +199,16 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - predict = phi(input_data_set, res.u) - PINOsolution(predict, res, phi, input_data_set) + res, phi + + #TODO build_solution + # if saveat isa Number + # ts = tspan[1]:saveat:tspan[2] + # end + + # if u0 isa Number + # u = [first(phi(t, res.u)) for t in ts] + # end + # sol = SciMLBase.build_solution(prob, alg, ts, u; + # sol end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 1a75690862..50efbe5b43 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -2,138 +2,103 @@ using Test using OrdinaryDiffEq, OptimizationOptimisers using Lux using Statistics, Random -# using NeuralOperators using NeuralPDE @testset "Example p" begin - linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - linear = (u, p, t) -> cos(p * t) + equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) u0 = 0.0f0 - p = pi / 2 - prob = ODEProblem(linear, u0, tspan, p) - #generate data - t0, t_end = tspan - instances_size = 50 - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - batch_size = 50 - as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] - - u_output_ = zeros(Float32, 1, instances_size, batch_size) - prob_set = [] - for (i, a_i) in enumerate(as) - prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) - sol1 = solve(prob_, Tsit5(); saveat = 0.0204) - reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob_) - u_output_[:, :, i] = reshape_sol - end - """ - Training data: - * input data: set of parameters 'a', - * output data: set of solutions u(t){a} corresponding parameter 'a'. - """ - train_set = TRAINSET(prob_set, u_output_) - - # operator learning phase - # init neural network - chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 32, Lux.σ), - Lux.Dense(32, 1)) - # flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), - # σ = gelu) + prob = ODEProblem(equation, u0, tspan) + # prob = PINOODEProblem(equation, tspan)? + + # init neural operator + branch = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1)) + trunk = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1)) + deeponet = NeuralPDE.DeepONet(branch, trunk) + + θ, st = Lux.setup(Random.default_rng(), deeponet) + a = rand(1, 10) + t = rand(1, 10) + x = (branch = a, trunk = t) + y, st = deeponet(x, θ, st) + + bounds = (p = [0.1, pi / 2],) + #instances_size = 100 TODO remove dt -> instances_size + opt = OptimizationOptimisers.Adam(0.1) + alg = NeuralPDE.PINOODE(deeponet, opt, bounds) + sol, phi = solve(prob, alg, dt = 0.1, verbose = true, maxiters = 2000) + + instances_size = 100 + p = range(bounds.p[1], stop = bounds.p[2], length = instances_size) + t = range(tspan[1], stop = tspan[2], length = instances_size) + x = (branch = collect(p)', trunk = collect(t)') + predict = phi(x, sol.u) + + ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) + ground_solution = ground_func.(u0, x.branch', x.trunk) + @test ground_solution ≈ predict atol=1 +end +@testset "Example with data" begin + equation = (u, p, t) -> cos(p * t) + tspan = (0.0f0, 2.0f0) + u0 = 0.0f0 + prob = ODEProblem(linear, u0, tspan) + + # init neural operator + deeponet = DeepONet(branch, trunk) opt = OptimizationOptimisers.Adam(0.01) - pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true) + bounds = (p = [0, pi / 2]) + function data_loss() + #code + end + alg = NeuralPDE.PINOODE(chain, opt, bounds; add_loss = data_loss) + sol = solve(prob, alg, verbose = false, maxiters = 2000) + predict = sol.predict - alg = PINOODE(chain, opt, pino_phase) - pino_solution = solve( - prob, alg, verbose = false, maxiters = 2000) - predict = pino_solution.predict - ground = u_output_ + ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) + ground = ground_func(..) @test ground≈predict atol=1 - # equation solving phase - dt = (t_end - t0) / instances_size - pino_phase = EquationSolving(dt, pino_solution) - chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 1)) - alg = PINOODE(chain, opt, pino_phase) - fine_tune_solution = solve( - prob, alg, verbose = false, maxiters = 2000) - - fine_tune_predict = fine_tune_solution.predict - operator_predict = pino_solution.phi( - fine_tune_solution.input_data_set, pino_solution.res.u) - ground = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :]) - @test ground≈fine_tune_predict atol=1. - @test operator_predict≈fine_tune_predict rtol=0.1 end @testset "Example u0" begin - linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - linear = (u, p, t) -> cos(p * t) + equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 2.0f0) p = Float32(pi) - u0 = 2.0f0 - prob = ODEProblem(linear, u0, tspan, p) - #generate data set - t0, t_end = tspan - instances_size = 50 - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - batch_size = 50 - u0s = [i for i in range(0.0f0, stop = pi / 2.0f0, length = batch_size)] - u_output_ = zeros(Float32, 1, instances_size, batch_size) - prob_set = [] - for (i, u0_i) in enumerate(u0s) - prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0_i, tspan, p) - sol1 = solve(prob_, Tsit5(); saveat = 0.0204) - reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob_) - u_output_[:, :, i] = reshape_sol - end - # operator learning phase - """ - Set of training data: - * input data: set of initial conditions 'u0' - * output data: set of solutions u(t){u0} corresponding initial conditions 'u0'. - """ - train_set = TRAINSET(prob_set, u_output_) - # fno = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), σ = gelu) - chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 32, Lux.σ), - Lux.Dense(32, 1)) + # u0 = 2.0f0 + prob = ODEProblem(linear, 0, tspan, p) + prob = PINOODEProblem(linear, tspan, p) + + neuraloperator = DeepONet(branch, trunk) + opt = OptimizationOptimisers.Adam(0.001) - pino_phase = OperatorLearning(train_set) - alg = PINOODE(chain, opt, pino_phase; isu0 = true) + bounds = (u0 = [0, 2],) + alg = PINOODE(chain, opt, bounds) pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) predict = pino_solution.predict - ground = u_output_ - @test ground≈predict atol=1. - - # equation solving phase - dt = (t_end -t0) / instances_size - pino_phase = EquationSolving(dt, pino_solution) - chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 1)) - alg = PINOODE(chain, opt, pino_phase; isu0 = true) - fine_tune_solution = solve(prob, alg, verbose = false, maxiters = 2000) - - fine_tune_predict = fine_tune_solution.predict - operator_predict = pino_solution.phi( - fine_tune_solution.input_data_set, pino_solution.res.u) - ground_fine_tune = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :]) - @test ground_fine_tune≈fine_tune_predict atol=1 - @test operator_predict≈fine_tune_predict rtol=0.1 + + ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) + ground = ground_func(...) + @test ground≈predict atol=1.0 +end + +@testset "Example u0 and p" begin + equation = (u, p, t) -> cos(p * t) + tspan = (0.0f0, 2.0f0) + p = Float32(pi) + # u0 = 2.0f0 + prob = ODEProblem(linear, 0, tspan, p) + prob = PINOODEProblem(linear, tspan, p) + + neuraloperator = DeepONet(branch, trunk) + + opt = OptimizationOptimisers.Adam(0.001) + bounds = (u0 = [0, 2],) + alg = PINOODE(chain, opt, bounds) + pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) + predict = pino_solution.predict + + ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) + ground = ground_func(...) + @test ground≈predict atol=1.0 end From f873541c19f0e2069574a8609615ad97d171740e Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 2 May 2024 14:35:37 +0400 Subject: [PATCH 043/153] pure PINO with DeepOnet --- docs/pages.jl | 3 +- docs/src/manual/pino_ode.md | 12 +-- docs/src/tutorials/pino_ode.md | 127 +++++++++-------------- src/NeuralPDE.jl | 2 +- src/neural_operators.jl | 105 ++++++++++++++----- src/pino_ode_solve.jl | 182 ++++++++++++++++++--------------- test/PINO_ode_tests.jl | 156 +++++++++++++--------------- test/PINO_ode_tests_gpu.jl | 151 --------------------------- test/runtests.jl | 5 - 9 files changed, 298 insertions(+), 445 deletions(-) delete mode 100644 test/PINO_ode_tests_gpu.jl diff --git a/docs/pages.jl b/docs/pages.jl index 020d330407..54b706f28c 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -32,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"] ] diff --git a/docs/src/manual/pino_ode.md b/docs/src/manual/pino_ode.md index 379f84c46b..72b5385924 100644 --- a/docs/src/manual/pino_ode.md +++ b/docs/src/manual/pino_ode.md @@ -5,17 +5,7 @@ PINOODE ``` ```@docs -TRAINSET +DeepONet ``` -```@docs -PINOsolution -``` -```@docs -OperatorLearning -``` - -```@docs -EquationSolving -``` diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 2e0ab0ed43..2d6d053852 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -1,8 +1,6 @@ -# Physics informed Neural Operator ODEs Solvers +# Physics Informed Neural Operator for ODEs Solvers -This tutorial is an introduction to using physics-informed neural operator (PINOs) for solving family of parametric ordinary diferential equations (ODEs). - -#TODO two phase +This tutorial provides an example of how using Physics Informed Neural Operator (PINO) for solving family of parametric ordinary differential equations (ODEs). ## Operator Learning for a family of parametric ODE. @@ -11,95 +9,62 @@ using Test using OrdinaryDiffEq, OptimizationOptimisers using Lux using Statistics, Random -# using NeuralOperators using NeuralPDE -linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) -linear = (u, p, t) -> cos(p * t) -tspan = (0.0f0, 2.0f0) -u0 = 0.0f0 -p = pi / 2f0 -prob = ODEProblem(linear, u0, tspan, p) +equation = (u, p, t) -> cos(p * t) +tspan = (0.0f0, 1.0f0) +u0 = 1.0f0 +prob = ODEProblem(equation, u0, tspan) + +# initilize DeepONet operator +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],) +#TODO add truct +strategy = (branch_size = 50, trunk_size = 40) +# strategy = (branch_size = 50, dt = 0.1)? +opt = OptimizationOptimisers.Adam(0.03) +alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) + +sol = solve(prob, alg, verbose = true, maxiters = 2000) ``` -Generate a dataset for learning a given family of ODEs where the parameter 'a' is varied. The dataset is generated by solving the ODE for different values of 'a' and storing the solutions. The dataset is then used to train the PINO model: -* input data: set of parameters 'a', -* output data: set of solutions u(t){a} corresponding parameter 'a'. +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 -t0, t_end = tspan -instances_size = 50 -range_ = range(t0, stop = t_end, length = instances_size) -ts = reshape(collect(range_), 1, instances_size) -batch_size = 50 -as = [Float32(i) for i in range(0.1, stop = pi / 2, length = batch_size)] - -u_output_ = zeros(Float32, 1, instances_size, batch_size) -prob_set = [] -for (i, a_i) in enumerate(as) - prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) - sol1 = solve(prob_, Tsit5(); saveat = 0.0204) - reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob_) - u_output_[:, :, i] = reshape_sol -end -train_set = TRAINSET(prob_set, u_output_) +using Plots +# Compute the ground truth solution for each parameter value and time in the solution +# The '.' operator is used to apply the functd ion element-wise +ground_analytic = (u0, p, t) -> begin u0 + sin(p * t) / (p) +p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) +p = reshape(p_, 1, branch_size, 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(sol.u[1, :, :], linetype = :contourf) +plot!(ground_solution[1, :, :], linetype = :contourf) ``` -Here it used the PINO method to learning operator of the given family of parametric ODEs. ```@example pino -chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 32, Lux.σ), - Lux.Dense(32, 1)) -# flat_no = FourierNeuralOperator(ch = (2, 16, 16, 16, 16, 16, 32, 1), modes = (16,), -# σ = gelu) - -opt = OptimizationOptimisers.Adam(0.01) -pino_phase = OperatorLearning(train_set, is_data_loss = true, is_physics_loss = true) - -alg = PINOODE(chain, opt, pino_phase) -pino_solution = solve( - prob, alg, verbose = false, maxiters = 3000) -predict = pino_solution.predict -ground = u_output_ -``` +using Plots -Now let's compare the predictions from the learned operator with the ground truth solution which is obtained early by numerically solving the parametric ODE. Where 'i' is the index of the parameter 'a' in the dataset. +# 'i' is the index of the parameter 'a' in the dataset +i = 45 -```@example pino -using Plots -i=45 +# 'predict' is the predicted solution from the PINO model +# 'ground' is the ground truth solution plot(predict[1, :, i], label = "Predicted") plot!(ground[1, :, i], label = "Ground truth") ``` - -Now to move on the stage of solving a certain equation using a trained operator and physics - -## Solve ODE using learned operator family of parametric ODE for fine tuning. -```@example pino -dt = (t_end - t0) / instances_size -pino_phase = EquationSolving(dt, pino_solution) -chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 1)) -alg = PINOODE(chain, opt, pino_phase) -fine_tune_solution = solve( prob, alg, verbose = false, maxiters = 2000) - -fine_tune_predict = fine_tune_solution.predict -operator_predict = pino_solution.phi( - fine_tune_solution.input_data_set, pino_solution.res.u) -ground_fine_tune = linear_analytic.(u0, p, fine_tune_solution.input_data_set[[1], :, :]) -``` - -Compare prediction with ground truth. - -```@example pino -plot(operator_predict[1, :, 1], label = "operator_predict") -plot!(fine_tune_predict[1, :, 1], label = "fine_tune_predict") -plot!(ground_fine_tune[1, :, 1], label = "Ground truth") -``` \ No newline at end of file diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index d9f401151d..30ac9b32d3 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -59,7 +59,7 @@ include("PDE_BPINN.jl") include("dgm.jl") -export NNODE, NNDAE, PINOODE, DeepONet +export NNODE, NNDAE, PINOODE, DeepONet, SomeStrategy #TODO remove SomeStrategy PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, diff --git a/src/neural_operators.jl b/src/neural_operators.jl index 7544675832..ddc87dcfce 100644 --- a/src/neural_operators.jl +++ b/src/neural_operators.jl @@ -1,47 +1,102 @@ -#TODO: Add docstrings +abstract type NeuralOperator <: Lux.AbstractExplicitLayer end + """ DeepONet(branch,trunk) """ -struct DeepONet{} <: Lux.AbstractExplicitLayer + +""" + 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 = l.branch, l.trunk + 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 -# function Lux.initialparameters(rng::AbstractRNG, e::DeepONet) -# code -# end - Lux.initialstates(::AbstractRNG, ::DeepONet) = NamedTuple() -""" -example: - -branch = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 1)) -trunk = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 1)) -a = rand(1, 100, 10) -t = rand(1, 1, 10) -x = (branch = a, trunk = t) - -deeponet = DeepONet(branch, trunk) -θ, st = Lux.setup(Random.default_rng(), deeponet) -y = deeponet(x, θ, st) -""" @inline function (f::DeepONet)(x::NamedTuple, θ, st::NamedTuple) - parameters, cord = x.branch, x.trunk + 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(parameters, θ_branch, st_branch) - out_t, st_t = trunk(cord, θ_trunk, st_trunk) - out = out_b' * out_t - return out, (branch = st_b, trunk = st_t) + 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 diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 597131ee0d..09d4fdefc4 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,36 +1,42 @@ """ - PINOODE(chain, + PINOODE + +PINOODE(chain, OptimizationOptimisers.Adam(0.1), - pino_phase; - init_params, + bounds; + init_params = nothing, + strategy = nothing kwargs...) -The method is that combine training data and physics constraints -to learn the solution operator of a given family of parametric Ordinary Differential Equations (ODE). +Algorithm for solving paramentric ordinary differential equations using a physics-informed +neural operator, which is used as a solver for a parametrized `ODEProblem`. ## Positional Arguments * `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`. `Flux.Chain` will be converted to `Lux` using `Lux.transform`. * `opt`: The optimizer to train the neural network. -* `pino_phase`: The phase of the PINN algorithm, either `OperatorLearning` or `EquationSolving`. +* `bounds`: A dictionary containing the bounds for the parameters of the neural network +in which will be train the prediction of parametric ODE. ## Keyword Arguments * `init_params`: The initial parameter of the neural network. By default, this is `nothing` which thus uses the random initialization provided by the neural network library. -* isu0: If true, the input data set contains initial conditions 'u0'. +* `strategy`: The strategy for training the neural network. +* `additional_loss`: additional function to the main one. For example, add training on data. * `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. ## References * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets" * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct PINOODE{C, O, B, I, S, K} <: SciMLBase.AbstractODEAlgorithm +struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy}, AL <: Union{Nothing, Function},K} <: + SciMLBase.AbstractODEAlgorithm chain::C opt::O bounds::B init_params::I - isu0::Bool strategy::S + additional_loss::AL #TODO kwargs::K end @@ -38,17 +44,26 @@ function PINOODE(chain, opt, bounds; init_params = nothing, - isu0 = false, #TODOD remove strategy = nothing, + additional_loss = nothing, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, bounds, init_params, isu0, strategy, kwargs) + PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs) +end + +#TODO cool name for strategy +struct SomeStrategy{} <: AbstractTrainingStrategy + branch_size::Int + trunk_size::Int +end +function SomeStrategy(; branch_size = 10, trunk_size=10) + SomeStrategy(branch_size, trunk_size) end mutable struct PINOPhi{C, T, U, S} chain::C t0::T - u0::U + u0::U#TODO remove u0, t0? st::S function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) @@ -68,100 +83,99 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, PINOPhi(chain, t0, u0, st), init_params end -#TODO update -# function (f::PINOPhi{C, T, U})(t::AbstractArray, -# θ) where {C <: Lux.AbstractExplicitLayer, T, U} -# y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t), θ, f.st) -# ChainRulesCore.@ignore_derivatives f.st = st -# ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t[1:size(y)[1], :, :]) -# u_0 = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0) -# u_0 .+ (ts .- f.t0) .* y -# end - -#TODO C <: DeepONet function (f::PINOPhi{C, T, U})(x::NamedTuple, θ) where {C, T, U} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st - a, t = x.branch, x.trunk - ts = adapt(parameterless_type(ComponentArrays.getdata(θ)), t) - u0_ = adapt(parameterless_type(ComponentArrays.getdata(θ)), f.u0) - u0_ .+ (ts .- f.t0) .* y + y end -# function dfdx(phi::PINOPhi, t::AbstractArray, θ) -# ε = [sqrt(eps(eltype(t))), zeros(eltype(t), size(t)[1] - 1)...] -# (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) -# end - -#TODO C <: DeepONet -function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ) where {C, T, U} - t = x.trunk - ε = [sqrt(eps(eltype(t)))] - phi_trunk(x, θ) = phi.chain.trunk(x, θ.trunk, phi.st.trunk)[1] - du_trunk_ = (phi_trunk(t .+ ε, θ) .- phi_trunk(t, θ)) ./ sqrt(eps(eltype(t))) - u_branch = phi.chain.branch(x.branch, θ.branch, phi.st.branch)[1] - u_branch' .* du_trunk_ +function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ, branch_left) where {C, T, U} + x_trunk = x.trunk + x_left = (branch = branch_left, trunk = x_trunk .+ sqrt(eps(eltype(x_trunk)))) + x_right = (branch = x.branch, trunk = x_trunk) + (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(x_trunk))) end -# function l₂loss(𝐲̂, 𝐲) -# feature_dims = 2:(ndims(𝐲) - 1) -# loss = sum(.√(sum(abs2, 𝐲̂ - 𝐲, dims = feature_dims))) -# y_norm = sum(.√(sum(abs2, 𝐲, dims = feature_dims))) -# return loss / y_norm -# end - function physics_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U} + x_ , p = x + norm = prod(size(x_.branch)) + branch_left = prob.f.(nothing, p, x_.trunk .+ sqrt(eps(eltype(x_.trunk)))) + sum(abs2, vec(dfdx(phi, x_, θ, branch_left)) .- vec(x_.branch)) / norm +end + +function operator_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U} + x_, u0_ = x + norm = prod(size(u0_)) + sum(abs2, vec(phi(x_, θ)) .- vec(u0_)) / norm +end + +#TODO inside loss couse call f(),return (p, t) +function get_trainset(branch_size, trunk_size, bounds, tspan, prob) + p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) + p = reshape(p_, 1, branch_size,1) + t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) + t = reshape(t_, 1, 1, trunk_size) f = prob.f - ps , ts = x.branch, x.trunk - norm = size(x.branch)[2] * size(x.trunk)[2] - sum(abs2, dfdx(phi, x, θ) - f.(phi(x, θ), ps, ts)) / norm + #TODO phi + f_ = f.(nothing, p, t) + x = (branch = f_, trunk = t) + x , p end -function get_trainset(bounds, tspan, strategy) - #TODO dt -> instances_size - instances_size = 100 - p = range(bounds.p[1], stop = bounds.p[2], length = instances_size) - t = range(tspan[1], stop = tspan[2], length = instances_size) - x = (branch = collect(p)', trunk = collect(t)') - x +#TODO return (t0, u0)? +function get_trainset_operator(branch_size, trunk_size, bounds, tspan, prob) + p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) + p = reshape(p_, 1, branch_size, 1) + t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) + t = reshape(t_, 1, 1, trunk_size) + f = prob.f + #TODO phi + f_ = f.(nothing, p, t[1]) + x_ = (branch = f_, trunk = t) + + t0 = x_.trunk[:, :, [1]] + # f_2 = f.(nothing, p, t0) + f_2 = f_[:, :, [1]] + x = (branch = f_2, trunk =t0) + u0_ = fill(prob.u0, size(f_2)) + x, u0_ end -#TODO GridTraining function generate_loss(strategy, prob::ODEProblem, phi, bounds, tspan) - x = get_trainset(bounds, tspan, strategy) + branch_size, trunk_size = strategy.branch_size, strategy.trunk_size + #TODO move to loss + x = get_trainset(branch_size, trunk_size, bounds, tspan, prob) + x_op = get_trainset_operator(branch_size, trunk_size, bounds, tspan, prob) function loss(θ, _) - physics_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) + operator_loss(phi, prob, x_op, θ) end end function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, alg::PINOODE, args...; - dt = nothing, - abstol = 1.0f-6, + abstol = 1.0f-8, reltol = 1.0f-3, verbose = false, saveat = nothing, maxiters = nothing) @unpack tspan, u0, p, f = prob - t0, t_end = tspan[1], tspan[2] - @unpack chain, opt, bounds, init_params, isu0 = alg + @unpack chain, opt, bounds, init_params, strategy = alg !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") if !any(in(keys(bounds)), (:u0, :p)) - error("bounds should contain u0 or p only") + error("bounds should contain u0 and p only") end - phi, init_params = generate_pino_phi_θ(chain, t0, u0, init_params) - init_params = ComponentArrays.ComponentArray(init_params) + phi, init_params = generate_pino_phi_θ(chain, 0, u0, init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - x = (branch = rand(length(bounds), 10), trunk = rand(1, 10)) + x = (branch = rand(1, 10,10), trunk = rand(1, 1, 10)) phi(x, init_params) catch err if isa(err, DimensionMismatch) @@ -171,13 +185,12 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end end - strategy = nothing - inner_f = generate_loss(strategy, prob, phi, bounds, tspan) + #TODO impl add loss function total_loss(θ, _) inner_f(θ, nothing) - #TODO add loss + # L2_loss = inner_f(θ, nothing) # if !(additional_loss isa Nothing) # L2_loss = L2_loss + additional_loss(phi, θ) @@ -199,16 +212,19 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - res, phi - - #TODO build_solution - # if saveat isa Number - # ts = tspan[1]:saveat:tspan[2] - # end - - # if u0 isa Number - # u = [first(phi(t, res.u)) for t in ts] - # end - # sol = SciMLBase.build_solution(prob, alg, ts, u; - # sol + + branch_size, trunk_size = strategy.branch_size, strategy.trunk_size + x, p = get_trainset(branch_size, trunk_size, bounds, tspan, prob) + u = phi(x, res.u) + + sol = SciMLBase.build_solution(prob, alg, x, u; + k = res, dense = true, + calculate_error = false, + retcode = ReturnCode.Success, + original = res, + resid = res.objective) + SciMLBase.has_analytic(prob.f) && + SciMLBase.calculate_solution_errors!(sol; timeseries_errors = true, + dense_errors = false) + sol end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 50efbe5b43..5049c121d7 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -4,101 +4,83 @@ using Lux using Statistics, Random using NeuralPDE -@testset "Example p" begin +# dG(u(t,p),t) = u(t,p) +@testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 2.0f0) - u0 = 0.0f0 - prob = ODEProblem(equation, u0, tspan) - # prob = PINOODEProblem(equation, tspan)? + tspan = (0.0f0, 1.0f0) + u0 = 1.0f0 + prob = ODEProblem(equation, u0, tspan) + + 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) + a = rand(1, 50, 40) + b = rand(1, 1, 40) + x = (branch = a, trunk = b) + θ, st = Lux.setup(Random.default_rng(), deeponet) + c = deeponet(x, θ, st)[1] - # init neural operator - branch = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1)) - trunk = Lux.Chain(Lux.Dense(1, 32, Lux.σ), Lux.Dense(32, 32, Lux.σ), Lux.Dense(32, 1)) - deeponet = NeuralPDE.DeepONet(branch, trunk) + bounds = (p = [0.1f0, pi],) - θ, st = Lux.setup(Random.default_rng(), deeponet) - a = rand(1, 10) - t = rand(1, 10) - x = (branch = a, trunk = t) - y, st = deeponet(x, θ, st) - - bounds = (p = [0.1, pi / 2],) - #instances_size = 100 TODO remove dt -> instances_size - opt = OptimizationOptimisers.Adam(0.1) - alg = NeuralPDE.PINOODE(deeponet, opt, bounds) - sol, phi = solve(prob, alg, dt = 0.1, verbose = true, maxiters = 2000) - - instances_size = 100 - p = range(bounds.p[1], stop = bounds.p[2], length = instances_size) - t = range(tspan[1], stop = tspan[2], length = instances_size) - x = (branch = collect(p)', trunk = collect(t)') - predict = phi(x, sol.u) - - ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) - ground_solution = ground_func.(u0, x.branch', x.trunk) - @test ground_solution ≈ predict atol=1 -end + strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) -@testset "Example with data" begin - equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 2.0f0) - u0 = 0.0f0 - prob = ODEProblem(linear, u0, tspan) + opt = OptimizationOptimisers.Adam(0.03) + alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) - # init neural operator - deeponet = DeepONet(branch, trunk) - opt = OptimizationOptimisers.Adam(0.01) - bounds = (p = [0, pi / 2]) - function data_loss() - #code - end - alg = NeuralPDE.PINOODE(chain, opt, bounds; add_loss = data_loss) sol = solve(prob, alg, verbose = false, maxiters = 2000) - predict = sol.predict - ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) - ground = ground_func(..) - @test ground≈predict atol=1 -end + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) + p = reshape(p_, 1, branch_size, 1) + ground_solution = ground_analytic.(u0, p, sol.t.trunk) -@testset "Example u0" begin - equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 2.0f0) - p = Float32(pi) - # u0 = 2.0f0 - prob = ODEProblem(linear, 0, tspan, p) - prob = PINOODEProblem(linear, tspan, p) - - neuraloperator = DeepONet(branch, trunk) - - opt = OptimizationOptimisers.Adam(0.001) - bounds = (u0 = [0, 2],) - alg = PINOODE(chain, opt, bounds) - pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) - predict = pino_solution.predict - - ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) - ground = ground_func(...) - @test ground≈predict atol=1.0 + @test ground_solution≈sol.u rtol=0.01 end -@testset "Example u0 and p" begin - equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 2.0f0) - p = Float32(pi) - # u0 = 2.0f0 - prob = ODEProblem(linear, 0, tspan, p) - prob = PINOODEProblem(linear, tspan, p) - - neuraloperator = DeepONet(branch, trunk) - - opt = OptimizationOptimisers.Adam(0.001) - bounds = (u0 = [0, 2],) - alg = PINOODE(chain, opt, bounds) - pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) - predict = pino_solution.predict - - ground_func = (u0, p, t) -> u0 + sin(p * t) / (p) - ground = ground_func(...) - @test ground≈predict atol=1.0 +@testset "Example du = p*t^2 " begin + equation = (u, p, t) -> p * t^2 + tspan = (0.0f0, 1.0f0) + u0 = 0.f0 + prob = ODEProblem(equation, u0, tspan) + + # init neural operator + branch = Lux.Chain( + Lux.Dense(1, 10, Lux.tanh_fast), + Lux.Dense(10, 10, Lux.tanh_fast), + Lux.Dense(10, 10, Lux.tanh_fast)) + 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 = NeuralPDE.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) + c = deeponet(x, θ, st)[1] + + bounds = (p = [0.0f0, 2.f0],) + strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) + opt = OptimizationOptimisers.Adam(0.03) + alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) + + sol = solve(prob, alg, verbose = false, maxiters = 2000) + ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 + p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) + p = reshape(p_, 1, branch_size, 1) + ground_solution = ground_analytic.(u0, p, sol.t.trunk) + + @test ground_solution≈sol.u rtol=0.01 end + +# @testset "Example with data" begin +# end diff --git a/test/PINO_ode_tests_gpu.jl b/test/PINO_ode_tests_gpu.jl deleted file mode 100644 index 919e6012ed..0000000000 --- a/test/PINO_ode_tests_gpu.jl +++ /dev/null @@ -1,151 +0,0 @@ -using Test -using OrdinaryDiffEq -using Lux -using ComponentArrays -#using NeuralOperators -using OptimizationOptimisers -using Random -using LuxCUDA -using NeuralPDE - -CUDA.allowscalar(false) -const gpud = gpu_device() - -@testset "Example p" begin - linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - linear = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 2.0f0) - u0 = 0.0f0 - p = pi / 2.0f0 - prob = ODEProblem(linear, u0, tspan, p) - #generate data set - t0, t_end = tspan - instances_size = 50 - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - batch_size = 50 - as = [Float32(i) for i in range(0.1, stop = pi / 2.0f0, length = batch_size)] - - u_output_ = zeros(Float32, 1, instances_size, batch_size) - prob_set = [] - for (i, a_i) in enumerate(as) - prob_ = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, a_i) - sol1 = solve(prob_, Tsit5(); saveat = 0.0204) - reshape_sol = Float32.(reshape(sol1(range_).u', 1, instances_size, 1)) - push!(prob_set, prob_) - u_output_[:, :, i] = reshape_sol - end - u_output_ = u_output_ |> gpud - - """ - Set of training data: - * input data: set of parameters 'a', - * output data: set of solutions u(t){a} corresponding parameter 'a'. - """ - train_set = TRAINSET(prob_set, u_output_) - - inner = 50 - chain = Lux.Chain(Lux.Dense(2, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, 1)) - ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud - opt = OptimizationOptimisers.Adam(0.03) - pino_phase = OperatorLearning(train_set) - alg = PINOODE(chain, opt, pino_phase; init_params = ps) - pino_solution = solve(prob, alg, verbose = false, maxiters = 2000) - predict = pino_solution.predict |> cpu - ground = u_output_ |> cpu - @test ground≈predict atol=1 - - dt = (t_end - t0) / instances_size - pino_phase = EquationSolving(dt, pino_solution) - alg = PINOODE(chain, opt, pino_phase; init_params = ps) - fine_tune_solution = solve( - prob, alg, verbose = false, maxiters = 2000) - - fine_tune_predict = fine_tune_solution.predict |> cpu - operator_predict = pino_solution.phi( - fine_tune_solution.input_data_set, pino_solution.res.u) |> cpu - input_data_set_ = fine_tune_solution.input_data_set[[1], :, :] |> cpu - ground_fine_tune = linear_analytic.(u0, p, input_data_set_) - @test ground_fine_tune≈fine_tune_predict atol=0.5 - @test operator_predict≈fine_tune_predict rtol=0.1 -end - -@testset "lotka volterra" begin - function lotka_volterra(u, p, t) - # Model parameters. - α, β, γ, δ = p - # Current state. - x, y = u - # Evaluate differential equations. - dx = (α - β * y) * x # prey - dy = (δ * x - γ) * y # predator - return [dx, dy] - end - - u0 = [1.0f0, 1.0f0] - p = Float32[1.5, 1.0, 3.0, 1.0] - tspan = (0.0f0, 4.0f0) - dt = 0.01f0 - prob = ODEProblem(lotka_volterra, u0, tspan, p) - t0, t_end = tspan - instances_size = 100 - range_ = range(t0, stop = t_end, length = instances_size) - ts = reshape(collect(range_), 1, instances_size) - batch_size = 50 - ps = [p .+ (i - 1) * Float32[0.000, 0.0, 0.001, 0.01] for i in 1:batch_size] - u_output_ = zeros(Float32, 2, instances_size, batch_size) - prob_set = [] - for (i, p_i) in enumerate(ps) - prob_ = ODEProblem(lotka_volterra, u0, tspan, p_i) - solution = solve(prob_, Tsit5(); saveat = dt) - reshape_sol_ = reduce(hcat, solution(range_).u) - reshape_sol = Float32.(reshape(reshape_sol_, 2, instances_size, 1)) - push!(prob_set, prob_) - u_output_[:, :, i] = reshape_sol - end - - train_set = TRAINSET(prob_set, u_output_) - - # flat_no = FourierNeuralOperator(ch = (5, 64, 64, 64, 64, 64, 128, 2), modes = (16,), - # σ = gelu) - # flat_no = Lux.transform(flat_no) - inner = 50 - chain = Lux.Chain(Lux.Dense(5, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, inner, Lux.σ), - Lux.Dense(inner, 2)) - ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud - - opt = OptimizationOptimisers.Adam(0.001) - pino_phase = OperatorLearning(train_set) - alg = PINOODE(chain, opt, pino_phase; init_params = ps) - pino_solution = solve(prob, alg, verbose = false, maxiters = 4000) - predict = pino_solution.predict |> cpu - ground = u_output_ - @test ground≈predict atol=5 - - dt = (t_end - t0) / instances_size - pino_phase = EquationSolving(dt, pino_solution; is_finetune_loss = true,is_physics_loss = true) - chain = Lux.Chain(Lux.Dense(5, 16, Lux.σ), - Lux.Dense(16, 16, Lux.σ), - Lux.Dense(16, 32, Lux.σ), - Lux.Dense(32, 2)) - ps = Lux.setup(Random.default_rng(), chain)[1] |> ComponentArray |> gpud - alg = PINOODE(chain, opt, pino_phase; init_params = ps) - fine_tune_solution = solve(prob, alg, verbose = false, maxiters = 2000) - - fine_tune_predict = fine_tune_solution.predict |> cpu - operator_predict = pino_solution.phi( - fine_tune_solution.input_data_set, pino_solution.res.u) |> cpu - input_data_set_ = fine_tune_solution.input_data_set[[1], :, :] |> cpu - ground_fine_tune = u_output_[:, :, [1]] - @test ground_fine_tune ≈ fine_tune_predict[:, 1:100, :] atol = 3 - @test operator_predict≈fine_tune_predict rtol=0.1 -end diff --git a/test/runtests.jl b/test/runtests.jl index ac8a4b6e36..0990c63bb7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -99,11 +99,6 @@ end include("NNPDE_tests_gpu_Lux.jl") end end - if !is_APPVEYOR && GROUP == "PINO_GPU" - @safetestset "PINO ode gpu" begin include("PINO_ode_tests_gpu.jl") - end - end - if GROUP == "All" || GROUP == "DGM" @time @safetestset "Deep Galerkin solver" begin From 0b4de350631f68feb1e729fb2ca2576559cfe5b3 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 6 May 2024 19:49:12 +0400 Subject: [PATCH 044/153] support case right side eq depends on 'u' --- docs/src/tutorials/pino_ode.md | 9 ++-- src/pino_ode_solve.jl | 93 ++++++++++++++++++---------------- test/PINO_ode_tests.jl | 48 +++++++++++++++--- 3 files changed, 94 insertions(+), 56 deletions(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 2d6d053852..b3de246b0e 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -16,7 +16,6 @@ tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) -# initilize DeepONet operator branch = Lux.Chain( Lux.Dense(1, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast), @@ -29,13 +28,13 @@ trunk = Lux.Chain( deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing) bounds = (p = [0.1f0, pi],) -#TODO add truct -strategy = (branch_size = 50, trunk_size = 40) -# strategy = (branch_size = 50, dt = 0.1)? + +strategy = SomeStrategy(branch_size = 50, trunk_size = 40) + opt = OptimizationOptimisers.Adam(0.03) alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) -sol = solve(prob, alg, verbose = true, maxiters = 2000) +sol = solve(prob, alg, verbose = false, maxiters = 2000) ``` 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 diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 09d4fdefc4..412ee07876 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -51,12 +51,13 @@ function PINOODE(chain, PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs) end -#TODO cool name for strategy +#TODO change to GridStrategy struct SomeStrategy{} <: AbstractTrainingStrategy branch_size::Int trunk_size::Int end -function SomeStrategy(; branch_size = 10, trunk_size=10) + +function SomeStrategy(;branch_size = 10, trunk_size=10) SomeStrategy(branch_size, trunk_size) end @@ -89,65 +90,65 @@ function (f::PINOPhi{C, T, U})(x::NamedTuple, θ) where {C, T, U} y end -function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ, branch_left) where {C, T, U} - x_trunk = x.trunk - x_left = (branch = branch_left, trunk = x_trunk .+ sqrt(eps(eltype(x_trunk)))) - x_right = (branch = x.branch, trunk = x_trunk) - (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(x_trunk))) +# function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ, branch_left) where {C, T, U} +# x_trunk = x.trunk +# x_left = (branch = branch_left, trunk = x_trunk .+ sqrt(eps(eltype(x_trunk)))) +# x_right = (branch = x.branch, trunk = x_trunk) +# (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(x_trunk))) +# end + +#TODO C <: DeepONet +function dfdx(phi::PINOPhi{C, T, U}, x::Tuple, θ, prob::ODEProblem) where {C, T, U} + p, t = x + f = prob.f + branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) + trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t + x_left = (branch = branch_left, trunk = trunk_left) + x_right = (branch = branch_right, trunk = trunk_right) + (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(t))) end function physics_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U} - x_ , p = x - norm = prod(size(x_.branch)) - branch_left = prob.f.(nothing, p, x_.trunk .+ sqrt(eps(eltype(x_.trunk)))) - sum(abs2, vec(dfdx(phi, x_, θ, branch_left)) .- vec(x_.branch)) / norm + p, t = x + f = prob.f + #TODO If du = f(u,p,t), where f = g(p,t)*u so it will wrong, f(0, p, t) = g(p,t)*0 = 0 + #work correct only with function like du = f(p,t) + g(u) + du = vec(dfdx(phi, x, θ, prob)) + + tuple = (branch = f.(0, p, t), trunk = t) + out = phi(tuple, θ) + f_ = vec(f.(out, p, t)) + norm = prod(size(out)) + sum(abs2, du .- f_) / norm end function operator_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U} - x_, u0_ = x + p, t = x + f = prob.f + t0 = t[:, :, [1]] + f_0 = f.(0, p, t0) + tuple = (branch = f_0, trunk = t0) + out = phi(tuple, θ) + u = vec(out) + u0_ = fill(prob.u0, size(out)) + u0 = vec(u0_) norm = prod(size(u0_)) - sum(abs2, vec(phi(x_, θ)) .- vec(u0_)) / norm + sum(abs2, u .- u0) / norm end -#TODO inside loss couse call f(),return (p, t) function get_trainset(branch_size, trunk_size, bounds, tspan, prob) p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) p = reshape(p_, 1, branch_size,1) t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) t = reshape(t_, 1, 1, trunk_size) - f = prob.f - #TODO phi - f_ = f.(nothing, p, t) - x = (branch = f_, trunk = t) - x , p -end - -#TODO return (t0, u0)? -function get_trainset_operator(branch_size, trunk_size, bounds, tspan, prob) - p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) - p = reshape(p_, 1, branch_size, 1) - t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) - t = reshape(t_, 1, 1, trunk_size) - f = prob.f - #TODO phi - f_ = f.(nothing, p, t[1]) - x_ = (branch = f_, trunk = t) - - t0 = x_.trunk[:, :, [1]] - # f_2 = f.(nothing, p, t0) - f_2 = f_[:, :, [1]] - x = (branch = f_2, trunk =t0) - u0_ = fill(prob.u0, size(f_2)) - x, u0_ + (p,t) end function generate_loss(strategy, prob::ODEProblem, phi, bounds, tspan) branch_size, trunk_size = strategy.branch_size, strategy.trunk_size - #TODO move to loss x = get_trainset(branch_size, trunk_size, bounds, tspan, prob) - x_op = get_trainset_operator(branch_size, trunk_size, bounds, tspan, prob) function loss(θ, _) - physics_loss(phi, prob, x, θ) + operator_loss(phi, prob, x_op, θ) + physics_loss(phi, prob, x, θ) + operator_loss(phi, prob, x, θ) end end @@ -159,14 +160,15 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, verbose = false, saveat = nothing, maxiters = nothing) - @unpack tspan, u0, p, f = prob + @unpack tspan, u0, f = prob @unpack chain, opt, bounds, init_params, strategy = alg !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - if !any(in(keys(bounds)), (:u0, :p)) - error("bounds should contain u0 and p only") + #TODO support for u0 + if !any(in(keys(bounds)), (:p,)) + error("bounds should contain p only") end phi, init_params = generate_pino_phi_θ(chain, 0, u0, init_params) @@ -214,7 +216,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, res = solve(optprob, opt; callback, maxiters, alg.kwargs...) branch_size, trunk_size = strategy.branch_size, strategy.trunk_size - x, p = get_trainset(branch_size, trunk_size, bounds, tspan, prob) + p, t = get_trainset(branch_size, trunk_size, bounds, tspan, prob) + x = (branch = f.(0, p,t), trunk = t) u = phi(x, res.u) sol = SciMLBase.build_solution(prob, alg, x, u; diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 5049c121d7..cadf5d51ba 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -5,6 +5,7 @@ using Statistics, Random using NeuralPDE # dG(u(t,p),t) = u(t,p) +# dG(u(t, p), t) = f(G,u(t, p)) @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -30,20 +31,58 @@ using NeuralPDE bounds = (p = [0.1f0, pi],) strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) + #TODO + # branch_size = 50, trunk_size = 40 + # strategy = NeuralPDE.GridStretagy([1 / branch_size, 1 / trunk_size]) opt = OptimizationOptimisers.Adam(0.03) alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + sol = solve(prob, alg, verbose = true, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) - p = reshape(p_, 1, branch_size, 1) + p = reshape(p_, 1, strategy.branch_size, 1) ground_solution = ground_analytic.(u0, p, sol.t.trunk) @test ground_solution≈sol.u rtol=0.01 end +@testset "Example du = cos(p * t) + u" begin + eq(u, p, t) = cos(p * t) + u + tspan = (0.0f0, 1.0f0) + u0 = 1.0f0 + prob = ODEProblem(eq, u0, tspan) + 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) + + bounds = (p = [0.1f0, 2],) + + strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) + + opt = OptimizationOptimisers.Adam(0.01) + alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) + + sol = solve(prob, alg, verbose = true, maxiters = 3000) + + #if u0 == 1 + ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2+2)*exp(t) ) / (p^2 + 1) + + p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) + p = reshape(p_, 1, strategy.branch_size, 1) + ground_solution = ground_analytic_.(u0, p, sol.t.trunk) + + @test ground_solution≈sol.u rtol=0.01 +end + @testset "Example du = p*t^2 " begin equation = (u, p, t) -> p * t^2 tspan = (0.0f0, 1.0f0) @@ -76,11 +115,8 @@ end sol = solve(prob, alg, verbose = false, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) - p = reshape(p_, 1, branch_size, 1) + p = reshape(p_, 1, strategy.branch_size, 1) ground_solution = ground_analytic.(u0, p, sol.t.trunk) @test ground_solution≈sol.u rtol=0.01 end - -# @testset "Example with data" begin -# end From 68e38d5de9b37bc28d777f570787b62dcae7abc0 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 6 May 2024 21:10:44 +0400 Subject: [PATCH 045/153] example with data loss --- src/pino_ode_solve.jl | 72 +++++++++++++++++------------------------- test/PINO_ode_tests.jl | 62 ++++++++++++++++++++++++------------ 2 files changed, 70 insertions(+), 64 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 412ee07876..3225d8aed6 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -29,14 +29,15 @@ in which will be train the prediction of parametric ODE. * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets" * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy}, AL <: Union{Nothing, Function},K} <: +struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy}, + AL <: Union{Nothing, Function}, K} <: SciMLBase.AbstractODEAlgorithm chain::C opt::O bounds::B init_params::I strategy::S - additional_loss::AL #TODO + additional_loss::AL kwargs::K end @@ -57,58 +58,46 @@ struct SomeStrategy{} <: AbstractTrainingStrategy trunk_size::Int end -function SomeStrategy(;branch_size = 10, trunk_size=10) +function SomeStrategy(; branch_size = 10, trunk_size = 10) SomeStrategy(branch_size, trunk_size) end -mutable struct PINOPhi{C, T, U, S} +mutable struct PINOPhi{C, S} chain::C - t0::T - u0::U#TODO remove u0, t0? st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, t0, u0, st) - new{typeof(chain), typeof(t0), typeof(u0), typeof(st)}(chain, t0, u0, st) + function PINOPhi(chain::Lux.AbstractExplicitLayer, st) + new{typeof(chain), typeof(st)}(chain, st) end end -function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, - t0, - u0, - init_params) +function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) if init_params === nothing init_params = ComponentArrays.ComponentArray(θ) else init_params = ComponentArrays.ComponentArray(init_params) end - PINOPhi(chain, t0, u0, st), init_params + PINOPhi(chain, st), init_params end -function (f::PINOPhi{C, T, U})(x::NamedTuple, θ) where {C, T, U} +function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st y end -# function dfdx(phi::PINOPhi{C, T, U}, x::NamedTuple, θ, branch_left) where {C, T, U} -# x_trunk = x.trunk -# x_left = (branch = branch_left, trunk = x_trunk .+ sqrt(eps(eltype(x_trunk)))) -# x_right = (branch = x.branch, trunk = x_trunk) -# (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(x_trunk))) -# end - #TODO C <: DeepONet -function dfdx(phi::PINOPhi{C, T, U}, x::Tuple, θ, prob::ODEProblem) where {C, T, U} +function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C, T} p, t = x f = prob.f - branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) + branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t x_left = (branch = branch_left, trunk = trunk_left) x_right = (branch = branch_right, trunk = trunk_right) (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(t))) end -function physics_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U} +function physics_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} p, t = x f = prob.f #TODO If du = f(u,p,t), where f = g(p,t)*u so it will wrong, f(0, p, t) = g(p,t)*0 = 0 @@ -122,7 +111,7 @@ function physics_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, sum(abs2, du .- f_) / norm end -function operator_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, T, U} +function operator_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} p, t = x f = prob.f t0 = t[:, :, [1]] @@ -136,19 +125,19 @@ function operator_loss(phi::PINOPhi{C, T, U}, prob::ODEProblem, x, θ) where {C, sum(abs2, u .- u0) / norm end -function get_trainset(branch_size, trunk_size, bounds, tspan, prob) +function get_trainset(branch_size, trunk_size, bounds, tspan) p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) - p = reshape(p_, 1, branch_size,1) + p = reshape(p_, 1, branch_size, 1) t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) t = reshape(t_, 1, 1, trunk_size) - (p,t) + (p, t) end function generate_loss(strategy, prob::ODEProblem, phi, bounds, tspan) branch_size, trunk_size = strategy.branch_size, strategy.trunk_size - x = get_trainset(branch_size, trunk_size, bounds, tspan, prob) + x = get_trainset(branch_size, trunk_size, bounds, tspan) function loss(θ, _) - physics_loss(phi, prob, x, θ) + operator_loss(phi, prob, x, θ) + operator_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end @@ -161,7 +150,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, saveat = nothing, maxiters = nothing) @unpack tspan, u0, f = prob - @unpack chain, opt, bounds, init_params, strategy = alg + @unpack chain, opt, bounds, init_params, strategy, additional_loss = alg !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") @@ -170,14 +159,13 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, if !any(in(keys(bounds)), (:p,)) error("bounds should contain p only") end - - phi, init_params = generate_pino_phi_θ(chain, 0, u0, init_params) + phi, init_params = generate_pino_phi_θ(chain, init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - x = (branch = rand(1, 10,10), trunk = rand(1, 1, 10)) + x = (branch = rand(1, 10, 10), trunk = rand(1, 1, 10)) phi(x, init_params) catch err if isa(err, DimensionMismatch) @@ -189,14 +177,12 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, inner_f = generate_loss(strategy, prob, phi, bounds, tspan) - #TODO impl add loss function total_loss(θ, _) - inner_f(θ, nothing) - - # L2_loss = inner_f(θ, nothing) - # if !(additional_loss isa Nothing) - # L2_loss = L2_loss + additional_loss(phi, θ) - # end + L2_loss = inner_f(θ, nothing) + if !(additional_loss isa Nothing) + L2_loss = L2_loss + additional_loss(phi, θ) + end + L2_loss end # Optimization Algo for Training Strategies @@ -216,8 +202,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, res = solve(optprob, opt; callback, maxiters, alg.kwargs...) branch_size, trunk_size = strategy.branch_size, strategy.trunk_size - p, t = get_trainset(branch_size, trunk_size, bounds, tspan, prob) - x = (branch = f.(0, p,t), trunk = t) + p, t = get_trainset(branch_size, trunk_size, bounds, tspan) + x = (branch = f.(0, p, t), trunk = t) u = phi(x, res.u) sol = SciMLBase.build_solution(prob, alg, x, u; diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index cadf5d51ba..0141d67cdc 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -10,7 +10,7 @@ using NeuralPDE equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0f0 - prob = ODEProblem(equation, u0, tspan) + prob = ODEProblem(equation, u0, tspan) branch = Lux.Chain( Lux.Dense(1, 10, Lux.tanh_fast), @@ -30,7 +30,7 @@ using NeuralPDE bounds = (p = [0.1f0, pi],) - strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) + strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) #TODO # branch_size = 50, trunk_size = 40 # strategy = NeuralPDE.GridStretagy([1 / branch_size, 1 / trunk_size]) @@ -38,9 +38,9 @@ using NeuralPDE opt = OptimizationOptimisers.Adam(0.03) alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) p = reshape(p_, 1, strategy.branch_size, 1) ground_solution = ground_analytic.(u0, p, sol.t.trunk) @@ -66,15 +66,16 @@ end bounds = (p = [0.1f0, 2],) - strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) + strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) opt = OptimizationOptimisers.Adam(0.01) alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) #if u0 == 1 - ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2+2)*exp(t) ) / (p^2 + 1) + ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / + (p^2 + 1) p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) p = reshape(p_, 1, strategy.branch_size, 1) @@ -83,13 +84,12 @@ end @test ground_solution≈sol.u rtol=0.01 end -@testset "Example du = p*t^2 " begin +@testset "Example with data du = p*t^2" begin equation = (u, p, t) -> p * t^2 tspan = (0.0f0, 1.0f0) - u0 = 0.f0 + u0 = 0.0f0 prob = ODEProblem(equation, u0, tspan) - # init neural operator branch = Lux.Chain( Lux.Dense(1, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast), @@ -99,24 +99,44 @@ end Lux.Dense(10, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast)) linear = Lux.Chain(Lux.Dense(10, 1)) - deeponet = NeuralPDE.DeepONet(branch, trunk; linear= linear) + deeponet = NeuralPDE.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) - c = deeponet(x, θ, st)[1] + bounds = (p = [0.0f0, 10.0f0],) + + strategy = NeuralPDE.SomeStrategy(branch_size = 60, trunk_size = 50) - bounds = (p = [0.0f0, 2.f0],) - strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 + function get_trainset(branch_size, trunk_size, bounds, tspan) + p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) + p = reshape(p_, 1, branch_size, 1) + t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) + t = reshape(t_, 1, 1, trunk_size) + (p, t) + end + branch_size, trunk_size = strategy.branch_size, strategy.trunk_size + p,t = get_trainset(branch_size, trunk_size, bounds, tspan) + function get_data() + sol = ground_analytic.(u0, p, t) + x = equation.(sol, p, t) + tuple = (branch = x, trunk = t) + sol, tuple + end + data, tuple = get_data() + function additional_loss(phi, θ) + u = phi(tuple, θ) + norm = prod(size(u)) + sum(abs2, u .- data) / norm + end + alg = NeuralPDE.PINOODE( + deeponet, opt, bounds; strategy = strategy, additional_loss = additional_loss) + sol = solve(prob, alg, verbose = false, maxiters = 2000) + p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) p = reshape(p_, 1, strategy.branch_size, 1) ground_solution = ground_analytic.(u0, p, sol.t.trunk) - @test ground_solution≈sol.u rtol=0.01 + @test ground_solution≈sol.u rtol=0.003 end From 345040be87df013534e008062d205b9de02dd84b Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 8 May 2024 16:06:21 +0400 Subject: [PATCH 046/153] support GridTraining --- src/pino_ode_solve.jl | 54 ++++++++++++++++++++---------------------- test/PINO_ode_tests.jl | 43 +++++++++++++++++---------------- 2 files changed, 48 insertions(+), 49 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 3225d8aed6..68b54b2da8 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -51,17 +51,6 @@ function PINOODE(chain, !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs) end - -#TODO change to GridStrategy -struct SomeStrategy{} <: AbstractTrainingStrategy - branch_size::Int - trunk_size::Int -end - -function SomeStrategy(; branch_size = 10, trunk_size = 10) - SomeStrategy(branch_size, trunk_size) -end - mutable struct PINOPhi{C, S} chain::C st::S @@ -80,14 +69,13 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) PINOPhi(chain, st), init_params end -function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C, T} +function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C <: NeuralOperator, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st y end -#TODO C <: DeepONet -function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C, T} +function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T} p, t = x f = prob.f branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) @@ -97,7 +85,8 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C, T} (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(t))) end -function physics_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} +function physics_loss( + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} p, t = x f = prob.f #TODO If du = f(u,p,t), where f = g(p,t)*u so it will wrong, f(0, p, t) = g(p,t)*0 = 0 @@ -111,7 +100,8 @@ function physics_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} sum(abs2, du .- f_) / norm end -function operator_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} +function operator_loss( + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} p, t = x f = prob.f t0 = t[:, :, [1]] @@ -125,17 +115,17 @@ function operator_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} sum(abs2, u .- u0) / norm end -function get_trainset(branch_size, trunk_size, bounds, tspan) - p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) - p = reshape(p_, 1, branch_size, 1) - t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) - t = reshape(t_, 1, 1, trunk_size) +function get_trainset(strategy::GridTraining, bounds, tspan) + db, dt = strategy.dx + p_ = bounds.p[1]:db:bounds.p[2] + p = reshape(p_, 1, size(p_)[1], 1) + t_ = collect(tspan[1]:dt:tspan[2]) + t = reshape(t_, 1, 1, size(t_)[1]) (p, t) end -function generate_loss(strategy, prob::ODEProblem, phi, bounds, tspan) - branch_size, trunk_size = strategy.branch_size, strategy.trunk_size - x = get_trainset(branch_size, trunk_size, bounds, tspan) +function generate_loss(strategy::GridTraining, prob::ODEProblem, phi, bounds, tspan) + x = get_trainset(strategy, bounds, tspan) function loss(θ, _) operator_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end @@ -155,11 +145,11 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - #TODO support for u0 + #TODO implement for u0 if !any(in(keys(bounds)), (:p,)) error("bounds should contain p only") end - phi, init_params = generate_pino_phi_θ(chain, init_params) + phi, init_params = generate_pino_phi_θ(chain, init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) @@ -175,6 +165,15 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end end + if strategy isa GridTraining + if length(strategy.dx) !== 2 + throw(ArgumentError("The discretization should have two elements dx= [db,dt], + steps for branch and trunk bounds")) + end + else + throw(ArgumentError("Only GridTraining strategy is supported")) + end + inner_f = generate_loss(strategy, prob, phi, bounds, tspan) function total_loss(θ, _) @@ -201,8 +200,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - branch_size, trunk_size = strategy.branch_size, strategy.trunk_size - p, t = get_trainset(branch_size, trunk_size, bounds, tspan) + p, t = get_trainset(strategy, bounds, tspan) x = (branch = f.(0, p, t), trunk = t) u = phi(x, res.u) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 0141d67cdc..324d9459d5 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -29,20 +29,15 @@ using NeuralPDE c = deeponet(x, θ, st)[1] bounds = (p = [0.1f0, pi],) - - strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) - #TODO - # branch_size = 50, trunk_size = 40 - # strategy = NeuralPDE.GridStretagy([1 / branch_size, 1 / trunk_size]) - + 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) - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) - p = reshape(p_, 1, strategy.branch_size, 1) + 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) @test ground_solution≈sol.u rtol=0.01 @@ -66,7 +61,9 @@ end bounds = (p = [0.1f0, 2],) - strategy = NeuralPDE.SomeStrategy(branch_size = 50, trunk_size = 40) + db = (bounds.p[2] - bounds.p[1]) / 50 + dt = (tspan[2] - tspan[1]) / 40 + strategy = NeuralPDE.GridTraining([db, dt]) opt = OptimizationOptimisers.Adam(0.01) alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) @@ -77,8 +74,8 @@ end ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / (p^2 + 1) - p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) - p = reshape(p_, 1, strategy.branch_size, 1) + 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) @test ground_solution≈sol.u rtol=0.01 @@ -103,7 +100,9 @@ end bounds = (p = [0.0f0, 10.0f0],) - strategy = NeuralPDE.SomeStrategy(branch_size = 60, trunk_size = 50) + db = (bounds.p[2] - bounds.p[1]) / 50 + dt = (tspan[2] - tspan[1]) / 40 + strategy = NeuralPDE.GridTraining([db, dt]) opt = OptimizationOptimisers.Adam(0.03) @@ -116,27 +115,29 @@ end t = reshape(t_, 1, 1, trunk_size) (p, t) end - branch_size, trunk_size = strategy.branch_size, strategy.trunk_size - p,t = get_trainset(branch_size, trunk_size, bounds, tspan) function get_data() sol = ground_analytic.(u0, p, t) x = equation.(sol, p, t) tuple = (branch = x, trunk = t) sol, tuple end + + branch_size, trunk_size = 50, 40 + p,t = get_trainset(branch_size, trunk_size, bounds, tspan) data, tuple = get_data() - function additional_loss(phi, θ) + + function additional_loss_(phi, θ) u = phi(tuple, θ) norm = prod(size(u)) sum(abs2, u .- data) / norm end alg = NeuralPDE.PINOODE( - deeponet, opt, bounds; strategy = strategy, additional_loss = additional_loss) + deeponet, opt, bounds; strategy = strategy, additional_loss = additional_loss_) sol = solve(prob, alg, verbose = false, maxiters = 2000) - p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) - p = reshape(p_, 1, strategy.branch_size, 1) + 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) - @test ground_solution≈sol.u rtol=0.003 + @test ground_solution≈sol.u rtol=0.005 end From 60005c830e1a6a058d9f74bc7b336a09f8baddb5 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 8 May 2024 16:58:43 +0400 Subject: [PATCH 047/153] update doc --- docs/src/tutorials/pino_ode.md | 39 ++++++++++++++++------------------ src/neural_operators.jl | 1 - test/PINO_ode_tests.jl | 1 - 3 files changed, 18 insertions(+), 23 deletions(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index b3de246b0e..fee09f74b3 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -1,9 +1,11 @@ # Physics Informed Neural Operator for ODEs Solvers -This tutorial provides an example of how using Physics Informed Neural Operator (PINO) for solving family of parametric ordinary differential equations (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 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 @@ -14,8 +16,9 @@ using NeuralPDE equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0f0 -prob = ODEProblem(equation, u0, tspan) +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), @@ -24,17 +27,16 @@ 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],) - -strategy = SomeStrategy(branch_size = 50, trunk_size = 40) - +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 @@ -42,28 +44,23 @@ Compare prediction with ground truth. ```@example pino using Plots -# Compute the ground truth solution for each parameter value and time in the solution -# The '.' operator is used to apply the functd ion element-wise -ground_analytic = (u0, p, t) -> begin u0 + sin(p * t) / (p) -p_ = range(bounds.p[1], stop = bounds.p[2], length = strategy.branch_size) -p = reshape(p_, 1, branch_size, 1) +# 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(sol.u[1, :, :], linetype = :contourf) +plot(predict[1, :, :], linetype = :contourf) plot!(ground_solution[1, :, :], linetype = :contourf) ``` - ```@example pino -using Plots - -# 'i' is the index of the parameter 'a' in the dataset -i = 45 - +# '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(predict[1, :, i], label = "Predicted") -plot!(ground[1, :, i], label = "Ground truth") +plot!(ground_solution[1, i, :], label = "Ground truth") ``` diff --git a/src/neural_operators.jl b/src/neural_operators.jl index ddc87dcfce..ed3c16db1d 100644 --- a/src/neural_operators.jl +++ b/src/neural_operators.jl @@ -53,7 +53,6 @@ y, st = deeponet(x, θ, st) * 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 diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 324d9459d5..ad4871c044 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -4,7 +4,6 @@ using Lux using Statistics, Random using NeuralPDE -# dG(u(t,p),t) = u(t,p) # dG(u(t, p), t) = f(G,u(t, p)) @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) From fbef5c79ae3d3bde64ce32e4ae6028bc11a6a1b0 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 12 Jun 2024 17:02:49 +0400 Subject: [PATCH 048/153] PINOPhi immutable --- src/pino_ode_solve.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 68b54b2da8..59bfd026d0 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -51,7 +51,7 @@ function PINOODE(chain, !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs) end -mutable struct PINOPhi{C, S} +struct PINOPhi{C, S} chain::C st::S function PINOPhi(chain::Lux.AbstractExplicitLayer, st) From 0fb7d2379b6c4f1f28559fb0a59dd5a96760c9d1 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 12 Jun 2024 17:04:00 +0400 Subject: [PATCH 049/153] Update src/pino_ode_solve.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/pino_ode_solve.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 59bfd026d0..ae1d12a8c7 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -51,6 +51,7 @@ function PINOODE(chain, !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs) end + struct PINOPhi{C, S} chain::C st::S From b6639ea5213737c3b4590aca09361f1ccc4b1682 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 12 Jun 2024 17:04:32 +0400 Subject: [PATCH 050/153] Update src/pino_ode_solve.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/pino_ode_solve.jl | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index ae1d12a8c7..5e39ecbf65 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,31 +1,32 @@ """ - PINOODE - -PINOODE(chain, - OptimizationOptimisers.Adam(0.1), - bounds; - init_params = nothing, - strategy = nothing - kwargs...) + PINOODE(chain, + OptimizationOptimisers.Adam(0.1), + bounds; + init_params = nothing, + strategy = nothing + kwargs...) Algorithm for solving paramentric ordinary differential equations using a physics-informed neural operator, which is used as a solver for a parametrized `ODEProblem`. ## Positional Arguments + * `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`. - `Flux.Chain` will be converted to `Lux` using `Lux.transform`. + `Flux.Chain` will be converted to `Lux` using `adapt(FromFluxAdaptor(false, false), chain)` * `opt`: The optimizer to train the neural network. * `bounds`: A dictionary containing the bounds for the parameters of the neural network in which will be train the prediction of parametric ODE. ## Keyword Arguments -* `init_params`: The initial parameter of the neural network. By default, this is `nothing` - which thus uses the random initialization provided by the neural network library. + +* `init_params`: The initial parameter of the neural network. By default, this is `nothing`, + which thus uses the random initialization provided by the neural network library. * `strategy`: The strategy for training the neural network. * `additional_loss`: additional function to the main one. For example, add training on data. * `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. ## References + * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets" * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ From cc1fadaab372f07078f1d35fa3cba974dfbee046 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 12 Jun 2024 17:07:46 +0400 Subject: [PATCH 051/153] remove SomeStrategy --- src/NeuralPDE.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 30ac9b32d3..d9f401151d 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -59,7 +59,7 @@ include("PDE_BPINN.jl") include("dgm.jl") -export NNODE, NNDAE, PINOODE, DeepONet, SomeStrategy #TODO remove SomeStrategy +export NNODE, NNDAE, PINOODE, DeepONet PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, From bb0863b7f0e8e06dcfad3f584733c1fab852f30e Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 12 Jun 2024 18:18:13 +0400 Subject: [PATCH 052/153] rename operator_loss to inital_condition_loss --- src/pino_ode_solve.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 5e39ecbf65..9769571bf7 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -102,7 +102,7 @@ function physics_loss( sum(abs2, du .- f_) / norm end -function operator_loss( +function inital_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} p, t = x f = prob.f @@ -129,7 +129,7 @@ end function generate_loss(strategy::GridTraining, prob::ODEProblem, phi, bounds, tspan) x = get_trainset(strategy, bounds, tspan) function loss(θ, _) - operator_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) + inital_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end From ea7c638081429fa81e527702bc2f629f8c62ed0a Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 14 Jun 2024 15:32:05 +0400 Subject: [PATCH 053/153] mutable PINOPhi --- src/pino_ode_solve.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 9769571bf7..a530a1d57c 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -53,7 +53,7 @@ function PINOODE(chain, PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs) end -struct PINOPhi{C, S} +mutable struct PINOPhi{C, S} chain::C st::S function PINOPhi(chain::Lux.AbstractExplicitLayer, st) @@ -148,7 +148,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, error("Only Lux.AbstractExplicitLayer neural networks are supported") #TODO implement for u0 - if !any(in(keys(bounds)), (:p,)) + if !any(in(keys(bounds)), (:p, :u0)) error("bounds should contain p only") end phi, init_params = generate_pino_phi_θ(chain, init_params) From 09f4891e022afb83e38ff30e1872435095816d89 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 14 Jun 2024 16:12:45 +0400 Subject: [PATCH 054/153] support u0 is param --- src/pino_ode_solve.jl | 18 ++++++--- test/PINO_ode_tests.jl | 87 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 98 insertions(+), 7 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index a530a1d57c..4a5b432cf2 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -103,7 +103,7 @@ function physics_loss( end function inital_condition_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ, bounds) where {C <: DeepONet, T} p, t = x f = prob.f t0 = t[:, :, [1]] @@ -111,15 +111,23 @@ function inital_condition_loss( tuple = (branch = f_0, trunk = t0) out = phi(tuple, θ) u = vec(out) - u0_ = fill(prob.u0, size(out)) - u0 = vec(u0_) + #TODO + if any(in(keys(bounds)), (:u0,)) + u0_ = p + u0 = p + else + u0_ = fill(prob.u0, size(out)) + u0 = vec(u0_) + end norm = prod(size(u0_)) sum(abs2, u .- u0) / norm end function get_trainset(strategy::GridTraining, bounds, tspan) db, dt = strategy.dx - p_ = bounds.p[1]:db:bounds.p[2] + v = values(bounds)[1] + #TODO for all v + p_ = v[1]:db:v[2] p = reshape(p_, 1, size(p_)[1], 1) t_ = collect(tspan[1]:dt:tspan[2]) t = reshape(t_, 1, 1, size(t_)[1]) @@ -129,7 +137,7 @@ end function generate_loss(strategy::GridTraining, prob::ODEProblem, phi, bounds, tspan) x = get_trainset(strategy, bounds, tspan) function loss(θ, _) - inital_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) + inital_condition_loss(phi, prob, x, θ, bounds) + physics_loss(phi, prob, x, θ) end end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ad4871c044..122e3ecbcb 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -43,10 +43,10 @@ using NeuralPDE end @testset "Example du = cos(p * t) + u" begin - eq(u, p, t) = cos(p * t) + u + eq_(u, p, t) = cos(p * t) + u tspan = (0.0f0, 1.0f0) u0 = 1.0f0 - prob = ODEProblem(eq, u0, tspan) + prob = ODEProblem(eq_, u0, tspan) branch = Lux.Chain( Lux.Dense(1, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast), @@ -140,3 +140,86 @@ end @test ground_solution≈sol.u rtol=0.005 end + +#u0 +@testset "Example du = cos(p * t)" begin + equation = (u, p, t) -> cos(p * t) + tspan = (0.0f0, 1.0f0) + u0 = 1.0f0 + prob = ODEProblem(equation, u0, tspan) + + 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 = (u0 = [0.5f0, 2.f0],) + db = (bounds.u0[2] - bounds.u0[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 = true, maxiters = 2000) + + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + p_ = bounds.u0[1]:strategy.dx[1]:bounds.u0[2] + p = reshape(p_, 1, size(p_)[1], 1) + ground_solution = ground_analytic.(u0, p, sol.t.trunk) + + @test ground_solution≈sol.u rtol=0.01 +end + +plot(sol.u[1, :, :], linetype = :contourf) +plot(ground_solution[1, :, :], linetype = :contourf) + +function plot_() + # Animate + anim = @animate for (i) in 1:51 + plot(ground_solution[1, i, :], label = "Ground") + # plot(equation_[1, i, :], label = "equation") + plot!(sol.u[1, i, :], label = "Predicted") + end + gif(anim, "pino.gif", fps = 15) +end + +plot_() + +#vector outputs and multiple parameters +@testset "Example du = cos(p * t)" begin + equation = (u, p, t) -> cos(p1 * t) + p2 + tspan = (0.0f0, 1.0f0) + u0 = 1.0f0 + prob = ODEProblem(equation, u0, tspan) + + 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 = (p1 = [0.1f0, pi], p2 = [0.1f0, 2.f0], u0 = [0.0f0, 2.0f0]) + db = (bounds.u0[2] - bounds.u0[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) + 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) + + @test ground_solution≈sol.u rtol=0.01 +end From 188ceecc2cbe10916fa63e16deebee3004f6bdb8 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 17 Jun 2024 19:33:07 +0400 Subject: [PATCH 055/153] update --- src/NeuralPDE.jl | 1 - src/pino_ode_solve.jl | 56 ++++++++++++++++++++++++++++-------------- test/PINO_ode_tests.jl | 51 +++++++------------------------------- 3 files changed, 46 insertions(+), 62 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index d9f401151d..5b3a6e1ca0 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -33,7 +33,6 @@ using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor using ChainRulesCore: @non_differentiable -#using NeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 4a5b432cf2..ddcd669f04 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -80,7 +80,12 @@ end function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T} p, t = x f = prob.f - branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) + # if any(in(keys(bounds)), (:u0,)) + # branch_left, branch_right = f.(0, zeros(size(p)), t .+ sqrt(eps(eltype(p)))), + # f.(0, zeros(size(p)), t) + # else + branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) + # end trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t x_left = (branch = branch_left, trunk = trunk_left) x_right = (branch = branch_right, trunk = trunk_right) @@ -91,11 +96,19 @@ function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} p, t = x f = prob.f - #TODO If du = f(u,p,t), where f = g(p,t)*u so it will wrong, f(0, p, t) = g(p,t)*0 = 0 - #work correct only with function like du = f(p,t) + g(u) + # TODO If du = f(u,p,t), where f = g(p,t)*u so it will wrong, f(0, p, t) = g(p,t)*0 = 0 + # work correct only with function like du = f(p,t) + g(u) + # tuple = (branch = p, trunk = t) + # out = phi(tuple, θ) + # f.(out, p, t) ? du = vec(dfdx(phi, x, θ, prob)) - - tuple = (branch = f.(0, p, t), trunk = t) + # if any(in(keys(bounds)), (:u0,)) + # f_ = f.(0, zeros(size(p)), t) + # else + f_ = f.(0, p, t) + # end + #TODO if DeepONet else error + tuple = (branch = f_, trunk = t) out = phi(tuple, θ) f_ = vec(f.(out, p, t)) norm = prod(size(out)) @@ -103,23 +116,26 @@ function physics_loss( end function inital_condition_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x, θ, bounds) where {C <: DeepONet, T} + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} p, t = x f = prob.f t0 = t[:, :, [1]] - f_0 = f.(0, p, t0) - tuple = (branch = f_0, trunk = t0) - out = phi(tuple, θ) - u = vec(out) #TODO - if any(in(keys(bounds)), (:u0,)) - u0_ = p - u0 = p - else + # if any(in(keys(bounds)), (:u0,)) + # u0_ = collect(p) + # u0 = vec(u0_) + # #TODO f_0 = f.(0, p, t0) and p call as u0 + # f_0 = f.(0, zeros(size(u0_)), t0) + # else u0_ = fill(prob.u0, size(out)) u0 = vec(u0_) - end - norm = prod(size(u0_)) + f_0 = f.(0, p, t0) + # end + tuple = (branch = f_0, trunk = t0) + out = phi(tuple, θ) + u = vec(out) + + norm = prod(size(u0)) sum(abs2, u .- u0) / norm end @@ -137,7 +153,7 @@ end function generate_loss(strategy::GridTraining, prob::ODEProblem, phi, bounds, tspan) x = get_trainset(strategy, bounds, tspan) function loss(θ, _) - inital_condition_loss(phi, prob, x, θ, bounds) + physics_loss(phi, prob, x, θ) + inital_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end @@ -155,8 +171,10 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - #TODO implement for u0 - if !any(in(keys(bounds)), (:p, :u0)) + # if !any(in(keys(bounds)), (:p, :u0)) + # error("bounds should contain p only") + # end + if !any(in(keys(bounds)), (:p,)) error("bounds should contain p only") end phi, init_params = generate_pino_phi_θ(chain, init_params) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 122e3ecbcb..bf00d076a1 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -20,7 +20,7 @@ using NeuralPDE Lux.Dense(10, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast)) - deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing) + deeponet = DeepONet(branch, trunk; linear = nothing) a = rand(1, 50, 40) b = rand(1, 1, 40) x = (branch = a, trunk = b) @@ -30,10 +30,10 @@ using NeuralPDE bounds = (p = [0.1f0, pi],) db = (bounds.p[2] - bounds.p[1]) / 50 dt = (tspan[2] - tspan[1]) / 40 - strategy = NeuralPDE.GridTraining([db, dt]) + strategy = GridTraining([db, dt]) opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + alg = PINOODE(deeponet, opt, bounds; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) 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) @@ -141,46 +141,12 @@ end @test ground_solution≈sol.u rtol=0.005 end -#u0 -@testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 1.0f0) - u0 = 1.0f0 - prob = ODEProblem(equation, u0, tspan) - - 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 = (u0 = [0.5f0, 2.f0],) - db = (bounds.u0[2] - bounds.u0[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 = true, maxiters = 2000) - - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - p_ = bounds.u0[1]:strategy.dx[1]:bounds.u0[2] - p = reshape(p_, 1, size(p_)[1], 1) - ground_solution = ground_analytic.(u0, p, sol.t.trunk) - - @test ground_solution≈sol.u rtol=0.01 -end - plot(sol.u[1, :, :], linetype = :contourf) -plot(ground_solution[1, :, :], linetype = :contourf) +plot!(ground_solution[1, :, :], linetype = :contourf) function plot_() # Animate - anim = @animate for (i) in 1:51 + anim = @animate for (i) in 1:41 plot(ground_solution[1, i, :], label = "Ground") # plot(equation_[1, i, :], label = "equation") plot!(sol.u[1, i, :], label = "Predicted") @@ -192,7 +158,8 @@ plot_() #vector outputs and multiple parameters @testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> cos(p1 * t) + p2 + # equation = (u, p, t) -> cos(p1 * t) + p2 + equation = (u, p, t) -> cos(p[1] * t) + p[2] tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) @@ -208,7 +175,7 @@ plot_() deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing) - bounds = (p1 = [0.1f0, pi], p2 = [0.1f0, 2.f0], u0 = [0.0f0, 2.0f0]) + bounds = (p1 = [0.1f0, pi], p2 = [0.1f0, 2.0f0]) db = (bounds.u0[2] - bounds.u0[1]) / 50 dt = (tspan[2] - tspan[1]) / 40 strategy = NeuralPDE.GridTraining([db, dt]) From eb005c6908e4f265788deec1c124e4fee95f6fca Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 18 Jun 2024 19:10:24 +0400 Subject: [PATCH 056/153] update multiple parameters task --- src/pino_ode_solve.jl | 102 ++++++++++++++++++++++++++--------------- test/PINO_ode_tests.jl | 57 +++++++++++++---------- 2 files changed, 98 insertions(+), 61 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index ddcd669f04..42f0da4b57 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -80,18 +80,26 @@ end function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T} p, t = x f = prob.f - # if any(in(keys(bounds)), (:u0,)) - # branch_left, branch_right = f.(0, zeros(size(p)), t .+ sqrt(eps(eltype(p)))), - # f.(0, zeros(size(p)), t) - # else - branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) - # end + branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t x_left = (branch = branch_left, trunk = trunk_left) x_right = (branch = branch_right, trunk = trunk_right) (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(t))) end +# function physics_loss( +# phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} +# p, t = x +# f = prob.f +# du = vec(dfdx(phi, x, θ, prob)) +# f_ = f.(0, p, t) +# tuple = (branch = f_, trunk = t) +# out = phi(tuple, θ) +# f_ = vec(f.(out, p, t)) +# norm = prod(size(out)) +# sum(abs2, du .- f_) / norm +# end + function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} p, t = x @@ -101,50 +109,69 @@ function physics_loss( # tuple = (branch = p, trunk = t) # out = phi(tuple, θ) # f.(out, p, t) ? + #TODO if DeepONet else err du = vec(dfdx(phi, x, θ, prob)) - # if any(in(keys(bounds)), (:u0,)) - # f_ = f.(0, zeros(size(p)), t) - # else - f_ = f.(0, p, t) - # end - #TODO if DeepONet else error - tuple = (branch = f_, trunk = t) + fs_ = hcat([hcat([f(0, p[:, i, :], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]...) + fs = reshape(fs_, (1, size(fs_)...)) + tuple = (branch = p, trunk = t) out = phi(tuple, θ) - f_ = vec(f.(out, p, t)) + + tuple = (branch = fs, trunk = t) #TODO -> tuple = (branch = p, trunk = t) + out = phi(tuple, θ) + # f_ = vec(f.(out, p, t)) norm = prod(size(out)) sum(abs2, du .- f_) / norm end -function inital_condition_loss( +# function initial_condition_loss( +# phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} +# p, t = x +# f = prob.f +# t0 = t[:, :, [1]] +# f_0 = f.(0, p, t0) +# tuple = (branch = f_0, trunk = t0) +# out = phi(tuple, θ) +# u = vec(out) +# u0_ = fill(prob.u0, size(out)) +# u0 = vec(u0_) + +# norm = prod(size(u0)) +# sum(abs2, u .- u0) / norm +# end + +function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} p, t = x f = prob.f - t0 = t[:, :, [1]] - #TODO - # if any(in(keys(bounds)), (:u0,)) - # u0_ = collect(p) - # u0 = vec(u0_) - # #TODO f_0 = f.(0, p, t0) and p call as u0 - # f_0 = f.(0, zeros(size(u0_)), t0) - # else - u0_ = fill(prob.u0, size(out)) - u0 = vec(u0_) - f_0 = f.(0, p, t0) - # end - tuple = (branch = f_0, trunk = t0) + t0 = t[:, :, [5]] + fs_0 = hcat([f(0, p[:, i, :], t[j]) for i in axes(p, 2)]...) + fs0 = reshape(fs_0, (1, size(fs_0)...)) + tuple = (branch = fs0, trunk = t0) out = phi(tuple, θ) u = vec(out) + u0_ = fill(prob.u0, size(out)) + u0 = vec(u0_) norm = prod(size(u0)) sum(abs2, u .- u0) / norm end +# function get_trainset(strategy::GridTraining, bounds, tspan) +# db, dt = strategy.dx +# v = values(bounds)[1] +# #TODO for all v +# p_ = v[1]:db:v[2] +# p = reshape(p_, 1, size(p_)[1], 1) +# t_ = collect(tspan[1]:dt:tspan[2]) +# t = reshape(t_, 1, 1, size(t_)[1]) +# (p, t) +# end + function get_trainset(strategy::GridTraining, bounds, tspan) - db, dt = strategy.dx - v = values(bounds)[1] - #TODO for all v - p_ = v[1]:db:v[2] - p = reshape(p_, 1, size(p_)[1], 1) + dt = strategy.dx + size_of_p = 50 + p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] + p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...) t_ = collect(tspan[1]:dt:tspan[2]) t = reshape(t_, 1, 1, size(t_)[1]) (p, t) @@ -174,9 +201,10 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, # if !any(in(keys(bounds)), (:p, :u0)) # error("bounds should contain p only") # end - if !any(in(keys(bounds)), (:p,)) - error("bounds should contain p only") - end + #TODO new p + # if !any(in(keys(bounds)), (:p,)) + # error("bounds should contain p only") + # end phi, init_params = generate_pino_phi_θ(chain, init_params) isinplace(prob) && @@ -194,7 +222,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end if strategy isa GridTraining - if length(strategy.dx) !== 2 + if length(strategy.dx) !== 2 #TODO ? throw(ArgumentError("The discretization should have two elements dx= [db,dt], steps for branch and trunk bounds")) end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index bf00d076a1..6f4be1ff94 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -28,6 +28,7 @@ using NeuralPDE c = deeponet(x, θ, st)[1] bounds = (p = [0.1f0, pi],) + # bounds = [0.1f0, pi] db = (bounds.p[2] - bounds.p[1]) / 50 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining([db, dt]) @@ -141,25 +142,14 @@ end @test ground_solution≈sol.u rtol=0.005 end -plot(sol.u[1, :, :], linetype = :contourf) -plot!(ground_solution[1, :, :], linetype = :contourf) - -function plot_() - # Animate - anim = @animate for (i) in 1:41 - plot(ground_solution[1, i, :], label = "Ground") - # plot(equation_[1, i, :], label = "equation") - plot!(sol.u[1, i, :], label = "Predicted") - end - gif(anim, "pino.gif", fps = 15) -end - -plot_() - #vector outputs and multiple parameters @testset "Example du = cos(p * t)" begin - # equation = (u, p, t) -> cos(p1 * t) + p2 - equation = (u, p, t) -> cos(p[1] * t) + p[2] + function equation1(u, p, t) + p1, p2 = p[1], p[2] + cos(p1 * t) + p2 + end + + equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) @@ -173,20 +163,39 @@ plot_() Lux.Dense(10, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast)) - deeponet = NeuralPDE.DeepONet(branch, trunk; linear = nothing) - - bounds = (p1 = [0.1f0, pi], p2 = [0.1f0, 2.0f0]) - db = (bounds.u0[2] - bounds.u0[1]) / 50 + deeponet = DeepONet(branch, trunk; linear = nothing) + # p1 = [0.1f0, pi]; p2 = [0.1f0, 2.0f0] + # bounds = (p = [p1, p2],) + #TODO add size_of_p = 50 + bounds = [[0.1f0, pi], [0.1f0, 2.0f0]] + # db = 0.025f0 dt = (tspan[2] - tspan[1]) / 40 - strategy = NeuralPDE.GridTraining([db, dt]) + strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.03) - alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) + alg = PINOODE(deeponet, opt, bounds; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 2000) - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + 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) @test ground_solution≈sol.u rtol=0.01 end + + + +plot(sol.u[1, :, :], linetype = :contourf) +plot!(ground_solution[1, :, :], linetype = :contourf) + +function plot_() + # Animate + anim = @animate for (i) in 1:41 + plot(ground_solution[1, i, :], label = "Ground") + # plot(equation_[1, i, :], label = "equation") + plot!(sol.u[1, i, :], label = "Predicted") + end + gif(anim, "pino.gif", fps = 15) +end + +plot_() From d754e300c7effe1dac1245d449627c69ea38ecaa Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 19 Jun 2024 19:01:46 +0400 Subject: [PATCH 057/153] add ParametricFunction --- src/NeuralPDE.jl | 1 + src/pino_ode_solve.jl | 139 +++++++++++++++++++++-------------------- test/PINO_ode_tests.jl | 87 ++++++++++++++------------ 3 files changed, 119 insertions(+), 108 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 5b3a6e1ca0..c97530df3d 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -33,6 +33,7 @@ using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor using ChainRulesCore: @non_differentiable +# using LuxNeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 42f0da4b57..ef092b36e9 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,3 +1,8 @@ +struct ParametricFunction{} + function_ ::Union{Nothing, Function} + bounds::Any +end + """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), @@ -30,12 +35,12 @@ in which will be train the prediction of parametric ODE. * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets" * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy}, +struct PINOODE{C, O, I, S <: Union{Nothing, AbstractTrainingStrategy}, AL <: Union{Nothing, Function}, K} <: SciMLBase.AbstractODEAlgorithm chain::C opt::O - bounds::B + parametric_function::ParametricFunction init_params::I strategy::S additional_loss::AL @@ -44,13 +49,13 @@ end function PINOODE(chain, opt, - bounds; + parametric_function; init_params = nothing, strategy = nothing, additional_loss = nothing, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, bounds, init_params, strategy, additional_loss, kwargs) + PINOODE(chain, opt, parametric_function, init_params, strategy, additional_loss, kwargs) end mutable struct PINOPhi{C, S} @@ -78,12 +83,11 @@ function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C <: NeuralOperator, T} end function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T} - p, t = x - f = prob.f - branch_left, branch_right = f.(0, p, t .+ sqrt(eps(eltype(p)))), f.(0, p, t) + pfs, p, t = x + # branch_left, branch_right = pfs, pfs trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t - x_left = (branch = branch_left, trunk = trunk_left) - x_right = (branch = branch_right, trunk = trunk_right) + x_left = (branch = pfs, trunk = trunk_left) + x_right = (branch = pfs, trunk = trunk_right) (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(t))) end @@ -99,30 +103,6 @@ end # norm = prod(size(out)) # sum(abs2, du .- f_) / norm # end - -function physics_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} - p, t = x - f = prob.f - # TODO If du = f(u,p,t), where f = g(p,t)*u so it will wrong, f(0, p, t) = g(p,t)*0 = 0 - # work correct only with function like du = f(p,t) + g(u) - # tuple = (branch = p, trunk = t) - # out = phi(tuple, θ) - # f.(out, p, t) ? - #TODO if DeepONet else err - du = vec(dfdx(phi, x, θ, prob)) - fs_ = hcat([hcat([f(0, p[:, i, :], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]...) - fs = reshape(fs_, (1, size(fs_)...)) - tuple = (branch = p, trunk = t) - out = phi(tuple, θ) - - tuple = (branch = fs, trunk = t) #TODO -> tuple = (branch = p, trunk = t) - out = phi(tuple, θ) - # f_ = vec(f.(out, p, t)) - norm = prod(size(out)) - sum(abs2, du .- f_) / norm -end - # function initial_condition_loss( # phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} # p, t = x @@ -134,24 +114,35 @@ end # u = vec(out) # u0_ = fill(prob.u0, size(out)) # u0 = vec(u0_) - # norm = prod(size(u0)) # sum(abs2, u .- u0) / norm # end -function initial_condition_loss( +function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} - p, t = x + pfs, p, t = x f = prob.f - t0 = t[:, :, [5]] - fs_0 = hcat([f(0, p[:, i, :], t[j]) for i in axes(p, 2)]...) - fs0 = reshape(fs_0, (1, size(fs_0)...)) - tuple = (branch = fs0, trunk = t0) + du = vec(dfdx(phi, x, θ, prob)) + tuple = (branch = pfs, trunk = t) out = phi(tuple, θ) - u = vec(out) - u0_ = fill(prob.u0, size(out)) - u0 = vec(u0_) + # if size(p)[1] == 1 + fs = f.(out, p, t) + f_ = vec(fs) + # else + # f_ = reduce(vcat,[reduce(vcat, [f(out[i], p[i], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]) + # end + norm = prod(size(out)) + sum(abs2, du .- f_) / norm +end +function initial_condition_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} + pfs, p, t = x + t0 = t[:, :, [1]] + pfs0 = pfs[:, :, [1]] + tuple = (branch = pfs0, trunk = t0) + out = phi(tuple, θ) + u = vec(out) + u0 = vec(fill(prob.u0, size(out))) norm = prod(size(u0)) sum(abs2, u .- u0) / norm end @@ -167,20 +158,31 @@ end # (p, t) # end -function get_trainset(strategy::GridTraining, bounds, tspan) +function get_trainset( + strategy::GridTraining, parametric_function::ParametricFunction, tspan) + @unpack function_, bounds = parametric_function dt = strategy.dx + #TODO size_of_p = 50 - p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] - p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...) + if bounds isa Tuple + p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) + p = collect(reshape(p_, 1, size(p_)[1], 1)) + else + p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] + p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...) + end + t_ = collect(tspan[1]:dt:tspan[2]) t = reshape(t_, 1, 1, size(t_)[1]) - (p, t) + pfs = function_.(p,t) + (pfs, p, t) end -function generate_loss(strategy::GridTraining, prob::ODEProblem, phi, bounds, tspan) - x = get_trainset(strategy, bounds, tspan) +function generate_loss( + strategy::GridTraining, prob::ODEProblem, phi, parametric_function::ParametricFunction, tspan) + x = get_trainset(strategy, parametric_function, tspan) function loss(θ, _) - inital_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) + initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end @@ -193,18 +195,15 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, saveat = nothing, maxiters = nothing) @unpack tspan, u0, f = prob - @unpack chain, opt, bounds, init_params, strategy, additional_loss = alg + @unpack chain, opt, parametric_function, init_params, strategy, additional_loss = alg + + if !isa(chain, DeepONet) + error("Only DeepONet neural networks are supported") + end !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") - # if !any(in(keys(bounds)), (:p, :u0)) - # error("bounds should contain p only") - # end - #TODO new p - # if !any(in(keys(bounds)), (:p,)) - # error("bounds should contain p only") - # end phi, init_params = generate_pino_phi_θ(chain, init_params) isinplace(prob) && @@ -221,16 +220,14 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end end - if strategy isa GridTraining - if length(strategy.dx) !== 2 #TODO ? - throw(ArgumentError("The discretization should have two elements dx= [db,dt], - steps for branch and trunk bounds")) - end - else + if strategy === nothing + dt = (tspan[2] - tspan[1]) / 50 + strategy = GridTraining(dt) + elseif !isa(strategy, GridTraining) throw(ArgumentError("Only GridTraining strategy is supported")) end - inner_f = generate_loss(strategy, prob, phi, bounds, tspan) + inner_f = generate_loss(strategy, prob, phi, parametric_function, tspan) function total_loss(θ, _) L2_loss = inner_f(θ, nothing) @@ -240,6 +237,10 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, L2_loss end + # TODO delete + # total_loss(θ, 0) + # Zygote.gradient(θ -> total_loss(θ, 0), θ) + # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() @@ -256,11 +257,11 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - p, t = get_trainset(strategy, bounds, tspan) - x = (branch = f.(0, p, t), trunk = t) - u = phi(x, res.u) + pfs, p, t = get_trainset(strategy, parametric_function, tspan) + tuple = (branch = pfs, trunk = t) + u = phi(tuple, res.u) - sol = SciMLBase.build_solution(prob, alg, x, u; + sol = SciMLBase.build_solution(prob, alg, tuple, u; k = res, dense = true, calculate_error = false, retcode = ReturnCode.Success, diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 6f4be1ff94..6677ba99b6 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -25,24 +25,45 @@ using NeuralPDE b = rand(1, 1, 40) x = (branch = a, trunk = b) θ, st = Lux.setup(Random.default_rng(), deeponet) - c = deeponet(x, θ, st)[1] - bounds = (p = [0.1f0, pi],) - # bounds = [0.1f0, pi] - db = (bounds.p[2] - bounds.p[1]) / 50 + c = deeponet(x, θ, st)[1] + function_(p, t) = cos(p * t) + bounds = (0.1f0, pi) + parametric_function = ParametricFunction(function_, bounds) dt = (tspan[2] - tspan[1]) / 40 - strategy = GridTraining([db, dt]) - opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + strategy = GridTraining(dt) + opt = OptimizationOptimisers.Adam(0.01) + alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy) + sol = solve(prob, alg, verbose = false, maxiters = 5000) + sol.original + # TODO intrepretation output another mesh + # x = (branch = p, trunk = t) + # phi(sol.original.u) + # sol. 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) + size_of_p = 50 + p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) + p = collect(reshape(p_, 1, size(p_)[1], 1)) ground_solution = ground_analytic.(u0, p, sol.t.trunk) - @test ground_solution≈sol.u rtol=0.01 + @test ground_solution≈sol.u rtol=0.05 end +plot(sol.u[1, :, :], linetype = :contourf) +plot!(ground_solution[1, :, :], linetype = :contourf) + +function plot_() + # Animate + anim = @animate for (i) in 1:41 + plot(ground_solution[1, i, :], label = "Ground") + # plot(equation_[1, i, :], label = "equation") + plot!(sol.u[1, i, :], label = "Predicted") + end + gif(anim, "pino.gif", fps = 15) +end + +plot_() + @testset "Example du = cos(p * t) + u" begin eq_(u, p, t) = cos(p * t) + u tspan = (0.0f0, 1.0f0) @@ -57,30 +78,35 @@ end Lux.Dense(10, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast)) - deeponet = NeuralPDE.DeepONet(branch, trunk) + deeponet = DeepONet(branch, trunk) - bounds = (p = [0.1f0, 2],) + parametric_function = Parametric_Function(nothing, bounds) - db = (bounds.p[2] - bounds.p[1]) / 50 dt = (tspan[2] - tspan[1]) / 40 - strategy = NeuralPDE.GridTraining([db, dt]) + strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.01) - alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) + alg = PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 10000) #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / (p^2 + 1) - 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) + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + size_of_p = 50 + p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] + p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...) + ground_solution = ground_analytic.(u0, p, sol.t.trunk) - @test ground_solution≈sol.u rtol=0.01 + @test ground_solution≈sol.u rtol=0.05 end + + + + @testset "Example with data du = p*t^2" begin equation = (u, p, t) -> p * t^2 tspan = (0.0f0, 1.0f0) @@ -100,7 +126,7 @@ end bounds = (p = [0.0f0, 10.0f0],) - db = (bounds.p[2] - bounds.p[1]) / 50 + # db = (bounds.p[2] - bounds.p[1]) / 50 dt = (tspan[2] - tspan[1]) / 40 strategy = NeuralPDE.GridTraining([db, dt]) @@ -182,20 +208,3 @@ end @test ground_solution≈sol.u rtol=0.01 end - - - -plot(sol.u[1, :, :], linetype = :contourf) -plot!(ground_solution[1, :, :], linetype = :contourf) - -function plot_() - # Animate - anim = @animate for (i) in 1:41 - plot(ground_solution[1, i, :], label = "Ground") - # plot(equation_[1, i, :], label = "equation") - plot!(sol.u[1, i, :], label = "Predicted") - end - gif(anim, "pino.gif", fps = 15) -end - -plot_() From 214b178c9b33c156f9ff83f23b08aebf5d2721a9 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 20 Jun 2024 17:28:24 +0400 Subject: [PATCH 058/153] vector outputs and multiple parameters --- src/pino_ode_solve.jl | 40 +++++++------- test/PINO_ode_tests.jl | 120 +++++++++++++++++++++-------------------- 2 files changed, 83 insertions(+), 77 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index ef092b36e9..6e45ef02e0 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,5 +1,5 @@ struct ParametricFunction{} - function_ ::Union{Nothing, Function} + function_::Union{Nothing, Function} bounds::Any end @@ -83,12 +83,14 @@ function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C <: NeuralOperator, T} end function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T} + # @unpack function_, bounds = parametric_function + # branch_left, branch_right = function_.(p, t), function_.(p, t .+ sqrt(eps(eltype(t)))) pfs, p, t = x - # branch_left, branch_right = pfs, pfs + branch_left, branch_right = p, p trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t - x_left = (branch = pfs, trunk = trunk_left) - x_right = (branch = pfs, trunk = trunk_right) - (phi(x_left, θ) .- phi(x_right, θ)) / sqrt(eps(eltype(t))) + x_left = (branch = branch_left, trunk = trunk_left) + x_right = (branch = branch_right, trunk = trunk_right) + (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end # function physics_loss( @@ -122,24 +124,25 @@ function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} pfs, p, t = x f = prob.f - du = vec(dfdx(phi, x, θ, prob)) - tuple = (branch = pfs, trunk = t) + tuple = (branch = p, trunk = t) out = phi(tuple, θ) - # if size(p)[1] == 1 + if size(p)[1] == 1 fs = f.(out, p, t) - f_ = vec(fs) - # else - # f_ = reduce(vcat,[reduce(vcat, [f(out[i], p[i], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]) - # end - norm = prod(size(out)) - sum(abs2, du .- f_) / norm + f_vec= vec(fs) + # out_ = vec(out) + else + f_vec = reduce(vcat,[[f(out[i], p[:, i, 1], t[j]) for i in axes(p, 2)] for j in axes(t, 3)]) + end + du = vec(dfdx(phi, x, θ, prob)) + norm = prod(size(du)) + sum(abs2, du .- f_vec) / norm end function initial_condition_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} pfs, p, t = x t0 = t[:, :, [1]] - pfs0 = pfs[:, :, [1]] - tuple = (branch = pfs0, trunk = t0) + # pfs0 = pfs[:, :, [1]] + tuple = (branch = p, trunk = t0) out = phi(tuple, θ) u = vec(out) u0 = vec(fill(prob.u0, size(out))) @@ -210,7 +213,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - x = (branch = rand(1, 10, 10), trunk = rand(1, 1, 10)) + in_dim = chain.branch.layers.layer_1.in_dims + x = (branch = rand(in_dim, 10, 10), trunk = rand(1, 1, 10)) phi(x, init_params) catch err if isa(err, DimensionMismatch) @@ -258,7 +262,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, res = solve(optprob, opt; callback, maxiters, alg.kwargs...) pfs, p, t = get_trainset(strategy, parametric_function, tspan) - tuple = (branch = pfs, trunk = t) + tuple = (branch = p, trunk = t) u = phi(tuple, res.u) sol = SciMLBase.build_solution(prob, alg, tuple, u; diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 6677ba99b6..a0039fc361 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -7,7 +7,7 @@ using NeuralPDE # dG(u(t, p), t) = f(G,u(t, p)) @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 1.0f0) + tspan = (0.0f0, 2.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) @@ -27,15 +27,17 @@ using NeuralPDE θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet(x, θ, st)[1] - function_(p, t) = cos(p * t) - bounds = (0.1f0, pi) + function_(p, t) = cos(p*t) + bounds = (pi, 2pi) parametric_function = ParametricFunction(function_, bounds) dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 5000) - sol.original + sol = solve(prob, alg, verbose = true, maxiters = 3000) + + phi(tuple, sol.original.u) + sol.original.objective # TODO intrepretation output another mesh # x = (branch = p, trunk = t) # phi(sol.original.u) @@ -46,24 +48,9 @@ using NeuralPDE p = collect(reshape(p_, 1, size(p_)[1], 1)) ground_solution = ground_analytic.(u0, p, sol.t.trunk) - @test ground_solution≈sol.u rtol=0.05 -end - -plot(sol.u[1, :, :], linetype = :contourf) -plot!(ground_solution[1, :, :], linetype = :contourf) - -function plot_() - # Animate - anim = @animate for (i) in 1:41 - plot(ground_solution[1, i, :], label = "Ground") - # plot(equation_[1, i, :], label = "equation") - plot!(sol.u[1, i, :], label = "Predicted") - end - gif(anim, "pino.gif", fps = 15) + @test ground_solution≈sol.u rtol=0.01 end -plot_() - @testset "Example du = cos(p * t) + u" begin eq_(u, p, t) = cos(p * t) + u tspan = (0.0f0, 1.0f0) @@ -79,34 +66,29 @@ plot_() Lux.Dense(10, 10, Lux.tanh_fast)) deeponet = DeepONet(branch, trunk) - - parametric_function = Parametric_Function(nothing, bounds) - + function_(p, t) = cos(p * t) + bounds = (0.1f0, 2.f0) + parametric_function = ParametricFunction(function_, bounds) + sol.original.objective dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, bounds; strategy = strategy) - - sol = solve(prob, alg, verbose = false, maxiters = 10000) + alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy) + sol = solve(prob, alg, verbose = false, maxiters = 5000) + sol.original.objective #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / (p^2 + 1) - - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) size_of_p = 50 - p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] - p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...) - ground_solution = ground_analytic.(u0, p, sol.t.trunk) + p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) + p = collect(reshape(p_, 1, size(p_)[1], 1)) + ground_solution = ground_analytic_.(u0, p, sol.t.trunk) - @test ground_solution≈sol.u rtol=0.05 + @test ground_solution≈sol.u rtol=0.01 end - - - - @testset "Example with data du = p*t^2" begin equation = (u, p, t) -> p * t^2 tspan = (0.0f0, 1.0f0) @@ -122,20 +104,22 @@ end Lux.Dense(10, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast)) linear = Lux.Chain(Lux.Dense(10, 1)) - deeponet = NeuralPDE.DeepONet(branch, trunk; linear = linear) + deeponet = DeepONet(branch, trunk; linear = linear) - bounds = (p = [0.0f0, 10.0f0],) + function_(p, t) = cos(p * t) + bounds = (0.0f0, 10.0f0) + parametric_function = ParametricFunction(function_, bounds) # db = (bounds.p[2] - bounds.p[1]) / 50 dt = (tspan[2] - tspan[1]) / 40 - strategy = NeuralPDE.GridTraining([db, dt]) + strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.03) #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 function get_trainset(branch_size, trunk_size, bounds, tspan) - p_ = range(bounds.p[1], stop = bounds.p[2], length = branch_size) + p_ = range(bounds[1], stop = bounds[1], length = branch_size) p = reshape(p_, 1, branch_size, 1) t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) t = reshape(t_, 1, 1, trunk_size) @@ -157,11 +141,12 @@ end norm = prod(size(u)) sum(abs2, u .- data) / norm end - alg = NeuralPDE.PINOODE( - deeponet, opt, bounds; strategy = strategy, additional_loss = additional_loss_) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + alg = PINOODE( + deeponet, opt, parametric_function; strategy = strategy, additional_loss = additional_loss_) + sol = solve(prob, alg, verbose = true, maxiters = 2000) - p_ = bounds.p[1]:strategy.dx[1]:bounds.p[2] + size_of_p = 50 + p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) p = reshape(p_, 1, size(p_)[1], 1) ground_solution = ground_analytic.(u0, p, sol.t.trunk) @@ -175,13 +160,13 @@ end cos(p1 * t) + p2 end - equation = (u, p, t) -> cos(p * t) + equation = (u, p, t) -> p[1]*cos(p[2] * t) tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) branch = Lux.Chain( - Lux.Dense(1, 10, Lux.tanh_fast), + Lux.Dense(2, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast), Lux.Dense(10, 10)) trunk = Lux.Chain( @@ -189,22 +174,39 @@ end Lux.Dense(10, 10, Lux.tanh_fast), Lux.Dense(10, 10, Lux.tanh_fast)) - deeponet = DeepONet(branch, trunk; linear = nothing) - # p1 = [0.1f0, pi]; p2 = [0.1f0, 2.0f0] - # bounds = (p = [p1, p2],) + deeponet = DeepONet(branch, trunk) + #TODO add size_of_p = 50 - bounds = [[0.1f0, pi], [0.1f0, 2.0f0]] - # db = 0.025f0 + function_(p, t) = cos(p * t) + bounds = [(0.1f0, pi), (1.0f0, 2.0f0)] + parametric_function = ParametricFunction(function_, bounds) dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(deeponet, opt, bounds; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) - - 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) + alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 3000) + ga = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) + p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] + p = vcat([collect(reshape(p_i, 1, size(p_i)[1], 1)) for p_i in p_]...) + ground_solution_ = f_vec = reduce( + hcat, [reduce( + vcat, [ga(u0, p[:, i, 1], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]) + ground_solution = reshape(ground_solution_, 1, size(ground_solution_)...) @test ground_solution≈sol.u rtol=0.01 end + +# plot(sol.u[1, :, :], linetype = :contourf) +# plot!(ground_solution[1, :, :], linetype = :contourf) + +# function plot_() +# # Animate +# anim = @animate for (i) in 1:41 +# plot(ground_solution[1, i, :], label = "Ground") +# # plot(equation_[1, i, :], label = "equation") +# plot!(sol.u[1, i, :], label = "Predicted") +# end +# gif(anim, "pino.gif", fps = 10) +# end + +# plot_() From 7d81063845fef23fc1e87de928eb82049aa51404 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 21 Jun 2024 16:32:07 +0400 Subject: [PATCH 059/153] clear code, rm ParametricFunction --- src/pino_ode_solve.jl | 97 ++++++++++-------------------------------- test/PINO_ode_tests.jl | 75 +++++++++++--------------------- 2 files changed, 48 insertions(+), 124 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 6e45ef02e0..527535f213 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,8 +1,3 @@ -struct ParametricFunction{} - function_::Union{Nothing, Function} - bounds::Any -end - """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), @@ -35,12 +30,13 @@ in which will be train the prediction of parametric ODE. * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets" * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct PINOODE{C, O, I, S <: Union{Nothing, AbstractTrainingStrategy}, +struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy}, AL <: Union{Nothing, Function}, K} <: SciMLBase.AbstractODEAlgorithm chain::C opt::O - parametric_function::ParametricFunction + bounds::B + number_of_parameters::Int init_params::I strategy::S additional_loss::AL @@ -49,13 +45,14 @@ end function PINOODE(chain, opt, - parametric_function; + bounds, + number_of_parameters; init_params = nothing, strategy = nothing, additional_loss = nothing, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, parametric_function, init_params, strategy, additional_loss, kwargs) + PINOODE(chain, opt, bounds,number_of_parameters, init_params, strategy, additional_loss, kwargs) end mutable struct PINOPhi{C, S} @@ -82,10 +79,8 @@ function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C <: NeuralOperator, T} y end -function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: DeepONet, T} - # @unpack function_, bounds = parametric_function - # branch_left, branch_right = function_.(p, t), function_.(p, t .+ sqrt(eps(eltype(t)))) - pfs, p, t = x +function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} + p, t = x branch_left, branch_right = p, p trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t x_left = (branch = branch_left, trunk = trunk_left) @@ -93,53 +88,25 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ, prob::ODEProblem) where {C <: De (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end -# function physics_loss( -# phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} -# p, t = x -# f = prob.f -# du = vec(dfdx(phi, x, θ, prob)) -# f_ = f.(0, p, t) -# tuple = (branch = f_, trunk = t) -# out = phi(tuple, θ) -# f_ = vec(f.(out, p, t)) -# norm = prod(size(out)) -# sum(abs2, du .- f_) / norm -# end -# function initial_condition_loss( -# phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} -# p, t = x -# f = prob.f -# t0 = t[:, :, [1]] -# f_0 = f.(0, p, t0) -# tuple = (branch = f_0, trunk = t0) -# out = phi(tuple, θ) -# u = vec(out) -# u0_ = fill(prob.u0, size(out)) -# u0 = vec(u0_) -# norm = prod(size(u0)) -# sum(abs2, u .- u0) / norm -# end - function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} - pfs, p, t = x + p, t = x f = prob.f tuple = (branch = p, trunk = t) out = phi(tuple, θ) if size(p)[1] == 1 fs = f.(out, p, t) f_vec= vec(fs) - # out_ = vec(out) else f_vec = reduce(vcat,[[f(out[i], p[:, i, 1], t[j]) for i in axes(p, 2)] for j in axes(t, 3)]) end - du = vec(dfdx(phi, x, θ, prob)) + du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm end function initial_condition_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} - pfs, p, t = x + p, t = x t0 = t[:, :, [1]] # pfs0 = pfs[:, :, [1]] tuple = (branch = p, trunk = t0) @@ -150,40 +117,26 @@ function initial_condition_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) whe sum(abs2, u .- u0) / norm end -# function get_trainset(strategy::GridTraining, bounds, tspan) -# db, dt = strategy.dx -# v = values(bounds)[1] -# #TODO for all v -# p_ = v[1]:db:v[2] -# p = reshape(p_, 1, size(p_)[1], 1) -# t_ = collect(tspan[1]:dt:tspan[2]) -# t = reshape(t_, 1, 1, size(t_)[1]) -# (p, t) -# end - -function get_trainset( - strategy::GridTraining, parametric_function::ParametricFunction, tspan) - @unpack function_, bounds = parametric_function +function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspan) dt = strategy.dx - #TODO - size_of_p = 50 - if bounds isa Tuple - p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) + if size(bounds)[1] == 1 + bound = bounds[1] + p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) p = collect(reshape(p_, 1, size(p_)[1], 1)) else - p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] + 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], 1)) for p_i in p_]...) end t_ = collect(tspan[1]:dt:tspan[2]) t = reshape(t_, 1, 1, size(t_)[1]) - pfs = function_.(p,t) - (pfs, p, t) + (p, t) end function generate_loss( - strategy::GridTraining, prob::ODEProblem, phi, parametric_function::ParametricFunction, tspan) - x = get_trainset(strategy, parametric_function, tspan) + strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) + x = get_trainset(strategy, bounds, number_of_parameters, tspan) function loss(θ, _) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end @@ -198,7 +151,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, saveat = nothing, maxiters = nothing) @unpack tspan, u0, f = prob - @unpack chain, opt, parametric_function, init_params, strategy, additional_loss = alg + @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss=alg if !isa(chain, DeepONet) error("Only DeepONet neural networks are supported") @@ -231,7 +184,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, throw(ArgumentError("Only GridTraining strategy is supported")) end - inner_f = generate_loss(strategy, prob, phi, parametric_function, tspan) + inner_f = generate_loss(strategy, prob, phi, bounds, number_of_parameters, tspan) function total_loss(θ, _) L2_loss = inner_f(θ, nothing) @@ -241,10 +194,6 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, L2_loss end - # TODO delete - # total_loss(θ, 0) - # Zygote.gradient(θ -> total_loss(θ, 0), θ) - # Optimization Algo for Training Strategies opt_algo = Optimization.AutoZygote() @@ -261,7 +210,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - pfs, p, t = get_trainset(strategy, parametric_function, tspan) + p, t = get_trainset(strategy, bounds, number_of_parameters, tspan) tuple = (branch = p, trunk = t) u = phi(tuple, res.u) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index a0039fc361..8226386acc 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,5 +1,5 @@ using Test -using OrdinaryDiffEq, OptimizationOptimisers +using OptimizationOptimisers using Lux using Statistics, Random using NeuralPDE @@ -27,24 +27,22 @@ using NeuralPDE θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet(x, θ, st)[1] - function_(p, t) = cos(p*t) - bounds = (pi, 2pi) - parametric_function = ParametricFunction(function_, bounds) + bounds = [(pi, 2pi)] + number_of_parameters = 50 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 3000) - phi(tuple, sol.original.u) sol.original.objective # TODO intrepretation output another mesh # x = (branch = p, trunk = t) # phi(sol.original.u) # sol. ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - size_of_p = 50 - p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) + #TDOD another number_of_parameters + p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = collect(reshape(p_, 1, size(p_)[1], 1)) ground_solution = ground_analytic.(u0, p, sol.t.trunk) @@ -66,23 +64,21 @@ end Lux.Dense(10, 10, Lux.tanh_fast)) deeponet = DeepONet(branch, trunk) - function_(p, t) = cos(p * t) - bounds = (0.1f0, 2.f0) - parametric_function = ParametricFunction(function_, bounds) - sol.original.objective + bounds = [(0.1f0, 2.f0)] + number_of_parameters = 40 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 5000) + sol = solve(prob, alg, verbose = false, maxiters = 3000) sol.original.objective #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / (p^2 + 1) - size_of_p = 50 - p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) + + p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = collect(reshape(p_, 1, size(p_)[1], 1)) ground_solution = ground_analytic_.(u0, p, sol.t.trunk) @@ -106,11 +102,8 @@ end linear = Lux.Chain(Lux.Dense(10, 1)) deeponet = DeepONet(branch, trunk; linear = linear) - function_(p, t) = cos(p * t) - bounds = (0.0f0, 10.0f0) - parametric_function = ParametricFunction(function_, bounds) - - # db = (bounds.p[2] - bounds.p[1]) / 50 + bounds = [(0.0f0, 10.0f0)] + number_of_parameters = 60 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) @@ -133,24 +126,23 @@ end end branch_size, trunk_size = 50, 40 - p,t = get_trainset(branch_size, trunk_size, bounds, tspan) - data, tuple = get_data() + p,t = get_trainset(branch_size, trunk_size, bounds[1], tspan) + data, tuple_ = get_data() function additional_loss_(phi, θ) - u = phi(tuple, θ) + u = phi(tuple_, θ) norm = prod(size(u)) sum(abs2, u .- data) / norm end alg = PINOODE( - deeponet, opt, parametric_function; strategy = strategy, additional_loss = additional_loss_) + deeponet, opt, bounds, number_of_parameters; strategy = strategy, + additional_loss = additional_loss_) sol = solve(prob, alg, verbose = true, maxiters = 2000) - - size_of_p = 50 - p_ = range(start = bounds[1], length = size_of_p, stop = bounds[2]) + p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = reshape(p_, 1, size(p_)[1], 1) ground_solution = ground_analytic.(u0, p, sol.t.trunk) - @test ground_solution≈sol.u rtol=0.005 + @test ground_solution≈sol.u rtol=0.01 end #vector outputs and multiple parameters @@ -175,38 +167,21 @@ end Lux.Dense(10, 10, Lux.tanh_fast)) deeponet = DeepONet(branch, trunk) - - #TODO add size_of_p = 50 - function_(p, t) = cos(p * t) bounds = [(0.1f0, pi), (1.0f0, 2.0f0)] - parametric_function = ParametricFunction(function_, bounds) + number_of_parameters = 50 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(deeponet, opt, parametric_function; strategy = strategy) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 3000) ga = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) - p_ = [range(start = b[1], length = size_of_p, stop = b[2]) for b in bounds] + 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], 1)) for p_i in p_]...) + t = sol.t.trunk ground_solution_ = f_vec = reduce( hcat, [reduce( vcat, [ga(u0, p[:, i, 1], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]) ground_solution = reshape(ground_solution_, 1, size(ground_solution_)...) @test ground_solution≈sol.u rtol=0.01 end - -# plot(sol.u[1, :, :], linetype = :contourf) -# plot!(ground_solution[1, :, :], linetype = :contourf) - -# function plot_() -# # Animate -# anim = @animate for (i) in 1:41 -# plot(ground_solution[1, i, :], label = "Ground") -# # plot(equation_[1, i, :], label = "equation") -# plot!(sol.u[1, i, :], label = "Predicted") -# end -# gif(anim, "pino.gif", fps = 10) -# end - -# plot_() From f5e2b067f8a86ab7397fa8e8cd8fe1e1d62252ff Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 21 Jun 2024 18:09:46 +0400 Subject: [PATCH 060/153] begin migrate LuxNeuralOperators and add QuasiRandomTraining --- Project.toml | 1 + docs/src/tutorials/pino_ode.md | 16 +++++++++------- src/pino_ode_solve.jl | 9 +++++++++ test/PINO_ode_tests.jl | 11 +++++++++++ 4 files changed, 30 insertions(+), 7 deletions(-) diff --git a/Project.toml b/Project.toml index da8598554a..c7e3634d99 100644 --- a/Project.toml +++ b/Project.toml @@ -20,6 +20,7 @@ Integrals = "de52edbc-65ea-441a-8357-d3a637375a31" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index fee09f74b3..97a4576d95 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -8,7 +8,7 @@ In this section, we will define a parametric ODE and solve it using a PINO. The ```@example pino using Test -using OrdinaryDiffEq, OptimizationOptimisers +using OptimizationOptimisers using Lux using Statistics, Random using NeuralPDE @@ -27,14 +27,16 @@ 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) +deeponet = DeepONet(branch, trunk; linear = nothing) -bounds = (p = [0.1f0, pi],) -db = (bounds.p[2] - bounds.p[1]) / 50 +bounds = [(0.1f0, pi)] +number_of_parameters = 50 +db = (bounds.[1][2] - bounds.[1][1]) / 50 dt = (tspan[2] - tspan[1]) / 40 -strategy = NeuralPDE.GridTraining([db, dt]) +strategy = GridTraining(dt) +strategy = QuasiRandomTraining(points) opt = OptimizationOptimisers.Adam(0.03) -alg = NeuralPDE.PINOODE(deeponet, opt, bounds; strategy = strategy) +alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 2000) predict = sol.u ``` @@ -46,7 +48,7 @@ Compare prediction with ground truth. 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_ = bounds[1][1]:strategy.dx:bounds[1][2] p = reshape(p_, 1, size(p_)[1], 1) ground_solution = ground_analytic.(u0, p, sol.t.trunk) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 527535f213..9f1a5b5ebc 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -79,6 +79,7 @@ function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C <: NeuralOperator, T} y end +#TODO migrate to LuxNeuralOperators.DeepONet function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} p, t = x branch_left, branch_right = p, p @@ -142,6 +143,14 @@ function generate_loss( end end +function generate_loss( + strategy::QuasiRandomTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) + #TODO + function loss(θ, _) + initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) + end +end + function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, alg::PINOODE, args...; diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 8226386acc..de1cebdd65 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -3,6 +3,16 @@ using OptimizationOptimisers using Lux using Statistics, Random using NeuralPDE +using LuxNeuralOperators + +u = rand(1, 5) +y = rand(1, 10, 5) +don_ = LuxNeuralOperators.DeepONet(Chain(Dense(1 => 32), Dense(32 => 32), Dense(32 => 16)), + Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16))) +ps, st = Lux.setup(Random.default_rng(), don_) + +don_((u, y), ps, st) +@inferred don_((u, y), ps, st) # dG(u(t, p), t) = f(G,u(t, p)) @testset "Example du = cos(p * t)" begin @@ -31,6 +41,7 @@ using NeuralPDE number_of_parameters = 50 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) + strategy = QuasiRandomTraining(points) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 3000) From 30a5134e20862a7159f87e55010bbea73d0ed061 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 25 Jun 2024 20:05:37 +0400 Subject: [PATCH 061/153] migrate to LuxNeuralOperators.DeepOnet --- src/NeuralPDE.jl | 3 +- src/neural_operators.jl | 101 --------------------------- src/pino_ode_solve.jl | 68 ++++++++++-------- test/PINO_ode_tests.jl | 151 +++++++++++++++++++--------------------- 4 files changed, 110 insertions(+), 213 deletions(-) delete mode 100644 src/neural_operators.jl diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index c97530df3d..2a2b1fa374 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -33,7 +33,7 @@ using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor using ChainRulesCore: @non_differentiable -# using LuxNeuralOperators +using LuxNeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) @@ -48,7 +48,6 @@ 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") diff --git a/src/neural_operators.jl b/src/neural_operators.jl deleted file mode 100644 index ed3c16db1d..0000000000 --- a/src/neural_operators.jl +++ /dev/null @@ -1,101 +0,0 @@ -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 diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 9f1a5b5ebc..12bf1fee29 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -52,7 +52,8 @@ function PINOODE(chain, additional_loss = nothing, kwargs...) !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, bounds,number_of_parameters, init_params, strategy, additional_loss, kwargs) + PINOODE(chain, opt, bounds, number_of_parameters, + init_params, strategy, additional_loss, kwargs) end mutable struct PINOPhi{C, S} @@ -73,45 +74,48 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) PINOPhi(chain, st), init_params end -function (f::PINOPhi{C, T})(x::NamedTuple, θ) where {C <: NeuralOperator, T} +function (f::PINOPhi{C, T})(x, θ) where {C, T} #C <: NeuralOperator y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st y end #TODO migrate to LuxNeuralOperators.DeepONet -function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} +function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C, T} #C <: DeepONet p, t = x branch_left, branch_right = p, p trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t - x_left = (branch = branch_left, trunk = trunk_left) - x_right = (branch = branch_right, trunk = trunk_right) + x_left = (branch_left, trunk_left) + x_right = (branch_right, trunk_right) (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end function physics_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} + phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where {C, T} #C <: DeepONet p, t = x f = prob.f - tuple = (branch = p, trunk = t) - out = phi(tuple, θ) + # x = (p, t) + out = phi(x, θ) if size(p)[1] == 1 - fs = f.(out, p, t) - f_vec= vec(fs) + fs = f.(out, p, vec(t)) + # fs = f.(0, p, vec(t)) + f_vec = vec(fs) else - f_vec = reduce(vcat,[[f(out[i], p[:, i, 1], t[j]) for i in axes(p, 2)] for j in axes(t, 3)]) + f_vec = reduce( + vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) end du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm end -function initial_condition_loss(phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C <: DeepONet, T} +function initial_condition_loss( + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} #C <: DeepONet #TODO migrate to LuxNeuralOperators.DeepONet p, t = x - t0 = t[:, :, [1]] + t0 = t[:, [1], :] # pfs0 = pfs[:, :, [1]] - tuple = (branch = p, trunk = t0) - out = phi(tuple, θ) + x0 = (p, t0) + out = phi(x0, θ) u = vec(out) u0 = vec(fill(prob.u0, size(out))) norm = prod(size(u0)) @@ -120,18 +124,18 @@ end function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspan) dt = strategy.dx - if size(bounds)[1] == 1 + if size(bounds)[1] == 1 bound = bounds[1] p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) - p = collect(reshape(p_, 1, size(p_)[1], 1)) + p = collect(reshape(p_, 1, size(p_)[1])) else 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], 1)) for p_i in p_]...) + p = vcat([collect(reshape(p_i, 1, size(p_i)[1])) for p_i in p_]...) end t_ = collect(tspan[1]:dt:tspan[2]) - t = reshape(t_, 1, 1, size(t_)[1]) + t = reshape(t_, 1, size(t_)[1], 1) (p, t) end @@ -160,11 +164,12 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, saveat = nothing, maxiters = nothing) @unpack tspan, u0, f = prob - @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss=alg + @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg - if !isa(chain, DeepONet) - error("Only DeepONet neural networks are supported") - end + #TODO migrate to LuxNeuralOperators.DeepONet + # if !isa(chain, DeepONet) + # error("Only DeepONet neural networks are supported") + # end !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") @@ -175,8 +180,11 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - in_dim = chain.branch.layers.layer_1.in_dims - x = (branch = rand(in_dim, 10, 10), trunk = rand(1, 1, 10)) + in_dim = chain.layers.branch.layers.layer_1.in_dims + # x = (branch = rand(in_dim, 10, 10), trunk = rand(1, 1, 10)) + u = rand(in_dim, number_of_parameters) + v = rand(1, 10, 1) + x = (u, v) phi(x, init_params) catch err if isa(err, DimensionMismatch) @@ -189,8 +197,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, if strategy === nothing dt = (tspan[2] - tspan[1]) / 50 strategy = GridTraining(dt) - elseif !isa(strategy, GridTraining) - throw(ArgumentError("Only GridTraining strategy is supported")) + elseif !(strategy isa GridTraining || strategy isa QuasiRandomTraining) + throw(ArgumentError("Only GridTraining and QuasiRandomTraining strategy is supported")) end inner_f = generate_loss(strategy, prob, phi, bounds, number_of_parameters, tspan) @@ -220,10 +228,10 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, res = solve(optprob, opt; callback, maxiters, alg.kwargs...) p, t = get_trainset(strategy, bounds, number_of_parameters, tspan) - tuple = (branch = p, trunk = t) - u = phi(tuple, res.u) + x = (p, t) + u = phi(x, res.u) - sol = SciMLBase.build_solution(prob, alg, tuple, u; + sol = SciMLBase.build_solution(prob, alg, x, u; k = res, dense = true, calculate_error = false, retcode = ReturnCode.Success, diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index de1cebdd65..2ac0a64f03 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -2,47 +2,39 @@ using Test using OptimizationOptimisers using Lux using Statistics, Random -using NeuralPDE using LuxNeuralOperators - -u = rand(1, 5) -y = rand(1, 10, 5) -don_ = LuxNeuralOperators.DeepONet(Chain(Dense(1 => 32), Dense(32 => 32), Dense(32 => 16)), - Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16))) -ps, st = Lux.setup(Random.default_rng(), don_) - -don_((u, y), ps, st) -@inferred don_((u, y), ps, st) +using NeuralPDE # dG(u(t, p), t) = f(G,u(t, p)) @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 2.0f0) + tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) - - 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 = DeepONet(branch, trunk; linear = nothing) - a = rand(1, 50, 40) - b = rand(1, 1, 40) - x = (branch = a, trunk = b) + deeponet = LuxNeuralOperators.DeepONet( + Chain( + Dense(1 => 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))) + u = rand(1, 50) + v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) + c = deeponet((u, v), θ, st)[1] + + branch = deeponet.layers.branch + θ, st = Lux.setup(Random.default_rng(), branch) + b = branch(u, θ, st)[1] + + trunk = deeponet.layers.trunk + θ, st = Lux.setup(Random.default_rng(), trunk) + t = trunk(v, θ, st)[1] - c = deeponet(x, θ, st)[1] bounds = [(pi, 2pi)] number_of_parameters = 50 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) - strategy = QuasiRandomTraining(points) - opt = OptimizationOptimisers.Adam(0.01) + # strategy = QuasiRandomTraining(50) + opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 3000) @@ -54,9 +46,10 @@ don_((u, y), ps, st) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) #TDOD another number_of_parameters p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = collect(reshape(p_, 1, size(p_)[1], 1)) - ground_solution = ground_analytic.(u0, p, sol.t.trunk) + p = collect(reshape(p_, 1, size(p_)[1])) + ground_solution = ground_analytic.(u0, p, vec(sol.t[2])) + @test ground_solution≈sol.u rtol=0.1 @test ground_solution≈sol.u rtol=0.01 end @@ -65,17 +58,12 @@ end tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(eq_, u0, tspan) - 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 = DeepONet(branch, trunk) - bounds = [(0.1f0, 2.f0)] + deeponet = LuxNeuralOperators.DeepONet( + Chain( + Dense(1 => 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))) + bounds = [(0.1f0, 2.0f0)] number_of_parameters = 40 dt = (tspan[2] - tspan[1]) / 40 strategy = GridTraining(dt) @@ -83,15 +71,15 @@ end opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 3000) + sol = solve(prob, alg, verbose = true, maxiters = 3000) sol.original.objective #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / (p^2 + 1) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = collect(reshape(p_, 1, size(p_)[1], 1)) - ground_solution = ground_analytic_.(u0, p, sol.t.trunk) + p = collect(reshape(p_, 1, size(p_)[1])) + ground_solution = ground_analytic_.(u0, p, vec(sol.t[2])) @test ground_solution≈sol.u rtol=0.01 end @@ -102,16 +90,11 @@ end u0 = 0.0f0 prob = ODEProblem(equation, u0, tspan) - branch = Lux.Chain( - Lux.Dense(1, 10, Lux.tanh_fast), - Lux.Dense(10, 10, Lux.tanh_fast), - Lux.Dense(10, 10, Lux.tanh_fast)) - 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) + deeponet = LuxNeuralOperators.DeepONet( + Chain( + Dense(1 => 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))) bounds = [(0.0f0, 10.0f0)] number_of_parameters = 60 @@ -122,22 +105,29 @@ end #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 + function get_trainset(branch_size, trunk_size, bounds, tspan) - p_ = range(bounds[1], stop = bounds[1], length = branch_size) - p = reshape(p_, 1, branch_size, 1) + p_ = range(bounds[1][1], stop = bounds[1][2], length = branch_size) + p = reshape(p_, 1, branch_size) t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) - t = reshape(t_, 1, 1, trunk_size) + t = reshape(t_, 1, trunk_size, 1) (p, t) end + function get_data() - sol = ground_analytic.(u0, p, t) - x = equation.(sol, p, t) - tuple = (branch = x, trunk = t) - sol, tuple + sol = ground_analytic.(u0, p, vec(t)) + #x = equation.(sol, p, vec(t)) + tuple_ = (p, t) + sol, tuple_ end + u = rand(1, 50) + v = rand(1, 40, 1) + θ, st = Lux.setup(Random.default_rng(), deeponet) + c = deeponet((u, v), θ, st)[1] + branch_size, trunk_size = 50, 40 - p,t = get_trainset(branch_size, trunk_size, bounds[1], tspan) + p, t = get_trainset(branch_size, trunk_size, bounds, tspan) data, tuple_ = get_data() function additional_loss_(phi, θ) @@ -145,13 +135,14 @@ end norm = prod(size(u)) sum(abs2, u .- data) / norm end + alg = PINOODE( deeponet, opt, bounds, number_of_parameters; strategy = strategy, additional_loss = additional_loss_) sol = solve(prob, alg, verbose = true, maxiters = 2000) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = reshape(p_, 1, size(p_)[1], 1) - ground_solution = ground_analytic.(u0, p, sol.t.trunk) + p = reshape(p_, 1, size(p_)[1]) + ground_solution = ground_analytic.(u0, p, vec(sol.t[2])) @test ground_solution≈sol.u rtol=0.01 end @@ -163,21 +154,22 @@ end cos(p1 * t) + p2 end - equation = (u, p, t) -> p[1]*cos(p[2] * t) + equation = (u, p, t) -> p[1] * cos(p[2] * t) tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) - branch = Lux.Chain( - Lux.Dense(2, 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 = LuxNeuralOperators.DeepONet( + Chain( + Dense(2 => 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))) + + u = rand(2, 50) + v = rand(1, 40, 1) + θ, st = Lux.setup(Random.default_rng(), deeponet) + c = deeponet((u, v), θ, st)[1] - deeponet = DeepONet(branch, trunk) bounds = [(0.1f0, pi), (1.0f0, 2.0f0)] number_of_parameters = 50 dt = (tspan[2] - tspan[1]) / 40 @@ -188,11 +180,10 @@ end ga = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) 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], 1)) for p_i in p_]...) - t = sol.t.trunk - ground_solution_ = f_vec = reduce( - hcat, [reduce( - vcat, [ga(u0, p[:, i, 1], t[j]) for i in axes(p, 2)]) for j in axes(t, 3)]) - ground_solution = reshape(ground_solution_, 1, size(ground_solution_)...) + p = vcat([collect(reshape(p_i, 1, size(p_i)[1])) for p_i in p_]...) + t = sol.t[2] + ground_solution = reduce(hcat, + [[ga(u0, p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + @test ground_solution≈sol.u rtol=0.01 end From 2818f34305c4ae112ff34e29ac4b6e7a072118dc Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 26 Jun 2024 17:22:38 +0400 Subject: [PATCH 062/153] add StochasticTraining --- src/pino_ode_solve.jl | 56 ++++++++++++++++++++++++++++-------------- test/PINO_ode_tests.jl | 26 +++++++++----------- 2 files changed, 50 insertions(+), 32 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 12bf1fee29..1d34c72c52 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,3 +1,4 @@ +#TODO rewrite doc strings """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), @@ -74,14 +75,14 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) PINOPhi(chain, st), init_params end -function (f::PINOPhi{C, T})(x, θ) where {C, T} #C <: NeuralOperator +function (f::PINOPhi{C, T})( + x, θ) where {C <: Lux.AbstractExplicitLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st y end -#TODO migrate to LuxNeuralOperators.DeepONet -function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C, T} #C <: DeepONet +function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: CompactLuxLayer{:DeepONet,}, T} p, t = x branch_left, branch_right = p, p trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t @@ -91,14 +92,13 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C, T} #C <: DeepONet end function physics_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where {C, T} #C <: DeepONet + phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { + C <: CompactLuxLayer{:DeepONet,}, T} p, t = x f = prob.f - # x = (p, t) out = phi(x, θ) if size(p)[1] == 1 fs = f.(out, p, vec(t)) - # fs = f.(0, p, vec(t)) f_vec = vec(fs) else f_vec = reduce( @@ -110,10 +110,10 @@ function physics_loss( end function initial_condition_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} #C <: DeepONet #TODO migrate to LuxNeuralOperators.DeepONet + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { + C <: CompactLuxLayer{:DeepONet,}, T} p, t = x - t0 = t[:, [1], :] - # pfs0 = pfs[:, :, [1]] + t0 = reshape([prob.tspan[1]], (1, 1, 1)) # t[:, [1], :] x0 = (p, t0) out = phi(x0, θ) u = vec(out) @@ -147,14 +147,35 @@ function generate_loss( end end +function get_trainset(strategy::StochasticTraining, bounds, number_of_parameters, tspan) + if size(bounds)[1] == 1 + bound = bounds[1] + p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] + else + p = reduce(vcat, + [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] + for bound in bounds]) + end + t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points,1) .+ tspan[1] + (p, t) +end + function generate_loss( - strategy::QuasiRandomTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) - #TODO + strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) function loss(θ, _) + x = get_trainset(strategy, bounds, number_of_parameters, tspan) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end +struct PINOODEInterpolation{T <: PINOPhi, T2} + phi::T + θ::T2 +end + +#TODO +# (f::NNODEInterpolation)(t, ...) = f.phi(t, f.θ) + function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, alg::PINOODE, args...; @@ -166,10 +187,9 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, @unpack tspan, u0, f = prob @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg - #TODO migrate to LuxNeuralOperators.DeepONet - # if !isa(chain, DeepONet) - # error("Only DeepONet neural networks are supported") - # end + if !isa(chain, CompactLuxLayer{:DeepONet,}) + error("Only DeepONet neural networks are supported with PINO ODE") + end !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") @@ -181,7 +201,6 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, try in_dim = chain.layers.branch.layers.layer_1.in_dims - # x = (branch = rand(in_dim, 10, 10), trunk = rand(1, 1, 10)) u = rand(in_dim, number_of_parameters) v = rand(1, 10, 1) x = (u, v) @@ -197,8 +216,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, if strategy === nothing dt = (tspan[2] - tspan[1]) / 50 strategy = GridTraining(dt) - elseif !(strategy isa GridTraining || strategy isa QuasiRandomTraining) - throw(ArgumentError("Only GridTraining and QuasiRandomTraining strategy is supported")) + elseif !(strategy isa GridTraining || strategy isa StochasticTraining) + throw(ArgumentError("Only GridTraining and StochasticTraining strategy is supported")) end inner_f = generate_loss(strategy, prob, phi, bounds, number_of_parameters, tspan) @@ -233,6 +252,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, sol = SciMLBase.build_solution(prob, alg, x, u; k = res, dense = true, + interp = PINOODEInterpolation(phi, res.u), calculate_error = false, retcode = ReturnCode.Success, original = res, diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 2ac0a64f03..d40ea42876 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -31,26 +31,24 @@ using NeuralPDE bounds = [(pi, 2pi)] number_of_parameters = 50 - dt = (tspan[2] - tspan[1]) / 40 - strategy = GridTraining(dt) - # strategy = QuasiRandomTraining(50) - opt = OptimizationOptimisers.Adam(0.03) + # dt = (tspan[2] - tspan[1]) / 40 + # strategy = GridTraining(dt) + strategy = StochasticTraining(40) + opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) - + sol = solve(prob, alg, verbose = true, maxiters = 5000) sol.original.objective - # TODO intrepretation output another mesh - # x = (branch = p, trunk = t) - # phi(sol.original.u) - # sol. + # TODO intrepretation output with few mesh ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - #TDOD another number_of_parameters p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = collect(reshape(p_, 1, size(p_)[1])) - ground_solution = ground_analytic.(u0, p, vec(sol.t[2])) + t_ = collect(tspan[1]:dt:tspan[2]) + t = collect(reshape(t_, 1, size(t_)[1], 1)) + ground_solution = ground_analytic.(u0, p, t_) + predict_sol = sol.interp.phi((p,t), sol.interp.θ) - @test ground_solution≈sol.u rtol=0.1 - @test ground_solution≈sol.u rtol=0.01 + @test ground_solution≈predict_sol rtol=0.1 + @test ground_solution≈predict_sol rtol=0.01 end @testset "Example du = cos(p * t) + u" begin From 88956933e698ac61640660aa1627d4a1229d0f6d Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 27 Jun 2024 16:49:30 +0400 Subject: [PATCH 063/153] add interpolation --- src/pino_ode_solve.jl | 18 ++++++++++-------- test/PINO_ode_tests.jl | 37 ++++++++++++++++++++++--------------- 2 files changed, 32 insertions(+), 23 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 1d34c72c52..753d3575a8 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -97,7 +97,7 @@ function physics_loss( p, t = x f = prob.f out = phi(x, θ) - if size(p)[1] == 1 + if size(p,1) == 1 fs = f.(out, p, vec(t)) f_vec = vec(fs) else @@ -124,18 +124,18 @@ end function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspan) dt = strategy.dx - if size(bounds)[1] == 1 + if size(bounds,1) == 1 bound = bounds[1] p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) - p = collect(reshape(p_, 1, size(p_)[1])) + p = collect(reshape(p_, 1, size(p_,1))) else 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_]...) + p = vcat([collect(reshape(p_i, 1, size(p_i,1))) for p_i in p_]...) end t_ = collect(tspan[1]:dt:tspan[2]) - t = reshape(t_, 1, size(t_)[1], 1) + t = reshape(t_, 1, size(t_,1), 1) (p, t) end @@ -148,7 +148,7 @@ function generate_loss( end function get_trainset(strategy::StochasticTraining, bounds, number_of_parameters, tspan) - if size(bounds)[1] == 1 + if size(bounds,1) == 1 bound = bounds[1] p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] else @@ -173,8 +173,10 @@ struct PINOODEInterpolation{T <: PINOPhi, T2} θ::T2 end -#TODO -# (f::NNODEInterpolation)(t, ...) = f.phi(t, f.θ) +(f::PINOODEInterpolation)(x) = f.phi(x, f.θ) + +SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation" +SciMLBase.allowscomplex(::PINOODE) = true function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, alg::PINOODE, diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index d40ea42876..ae90643fb0 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -31,23 +31,30 @@ using NeuralPDE bounds = [(pi, 2pi)] number_of_parameters = 50 - # dt = (tspan[2] - tspan[1]) / 40 - # strategy = GridTraining(dt) strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 5000) - sol.original.objective - # TODO intrepretation output with few mesh + sol = solve(prob, alg, verbose = false, maxiters = 2000) + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = collect(reshape(p_, 1, size(p_)[1])) - t_ = collect(tspan[1]:dt:tspan[2]) - t = collect(reshape(t_, 1, size(t_)[1], 1)) - ground_solution = ground_analytic.(u0, p, t_) - predict_sol = sol.interp.phi((p,t), sol.interp.θ) + function get_trainset(bounds, tspan , number_of_parameters, dt) + p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) + p = collect(reshape(p_, 1, size(p_, 1))) + t_ = collect(tspan[1]:dt:tspan[2]) + t = collect(reshape(t_, 1, size(t_, 1), 1)) + (p,t) + end + p,t = get_trainset(bounds, tspan, number_of_parameters, dt) + + ground_solution = ground_analytic.(u0, p, vec(t)) + predict_sol = sol.interp((p, t)) + + @test ground_solution≈predict_sol rtol=0.01 + + p, t = get_trainset(bounds, tspan, 100, 0.01) + ground_solution = ground_analytic.(u0, p, vec(t)) + predict_sol = sol.interp((p, t)) - @test ground_solution≈predict_sol rtol=0.1 @test ground_solution≈predict_sol rtol=0.01 end @@ -76,7 +83,7 @@ end (p^2 + 1) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = collect(reshape(p_, 1, size(p_)[1])) + p = collect(reshape(p_, 1, size(p_,1))) ground_solution = ground_analytic_.(u0, p, vec(sol.t[2])) @test ground_solution≈sol.u rtol=0.01 @@ -163,7 +170,7 @@ end Chain(Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast))) - u = rand(2, 50) + u = rand(2, 50, 1) v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] @@ -178,7 +185,7 @@ end ga = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) 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_]...) + p = vcat([collect(reshape(p_i, 1, size(p_i,1))) for p_i in p_]...) t = sol.t[2] ground_solution = reduce(hcat, [[ga(u0, p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) From 06ce5173c1d56c0b2441c688ef87d399967d8255 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 27 Jun 2024 17:44:12 +0400 Subject: [PATCH 064/153] update doc --- docs/src/manual/pino_ode.md | 6 --- docs/src/tutorials/pino_ode.md | 85 ++++++++++++++++++++++------------ src/pino_ode_solve.jl | 2 +- test/PINO_ode_tests.jl | 47 ++++++++++++------- 4 files changed, 86 insertions(+), 54 deletions(-) diff --git a/docs/src/manual/pino_ode.md b/docs/src/manual/pino_ode.md index 72b5385924..6177341dc7 100644 --- a/docs/src/manual/pino_ode.md +++ b/docs/src/manual/pino_ode.md @@ -3,9 +3,3 @@ ```@docs PINOODE ``` - -```@docs -DeepONet -``` - - diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 97a4576d95..0a7229d621 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -11,34 +11,32 @@ using Test using OptimizationOptimisers using Lux using Statistics, Random +using LuxNeuralOperators using NeuralPDE -equation = (u, p, t) -> cos(p * t) +equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] 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 = DeepONet(branch, trunk; linear = nothing) - -bounds = [(0.1f0, pi)] +number_of_parameter = 3 +deeponet = LuxNeuralOperators.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))) + +u = rand(3, 50) +v = rand(1, 40, 1) +θ, st = Lux.setup(Random.default_rng(), deeponet) +c = deeponet((u, v), θ, st)[1] + +bounds = [(1.0f0, pi), (1.0f0, 2.0f0), (2.0f0, 3.0f0)] number_of_parameters = 50 -db = (bounds.[1][2] - bounds.[1][1]) / 50 -dt = (tspan[2] - tspan[1]) / 40 -strategy = GridTraining(dt) -strategy = QuasiRandomTraining(points) +strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) -sol = solve(prob, alg, verbose = false, maxiters = 2000) -predict = sol.u +sol = solve(prob, alg, verbose = true, maxiters = 3000) ``` 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 @@ -46,23 +44,52 @@ Compare prediction with ground truth. ```@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_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) -p_ = bounds[1][1]:strategy.dx:bounds[1][2] -p = reshape(p_, 1, size(p_)[1], 1) -ground_solution = ground_analytic.(u0, p, sol.t.trunk) +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 + +(p,t) = get_trainset(bounds, tspan, 50, 0.025f0) +ground_solution_ = ground_solution_f(p,t) +predict = sol.interp((p,t)) + +# Calculate the mean error and the standard deviation of the errors +errors = ground_solution_ - predict +mean_error = mean(errors) +std_error = std(errors) + +# generate the solution with new parameters for test the model +p,t = get_trainset(bounds, tspan, 100, 0.01f0) +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 -# sol.u[1, :, :], represents the predicted solution for each parameter value and time -plot(predict[1, :, :], linetype = :contourf) -plot!(ground_solution[1, :, :], linetype = :contourf) +# 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 +i = 5 # 'predict' is the predicted solution from the PINO model -plot(predict[1, i, :], label = "Predicted") +plot(predict[:, i], label = "Predicted") # 'ground' is the ground truth solution -plot!(ground_solution[1, i, :], label = "Ground truth") +plot!(ground_solution_[:, i], label = "Ground truth") ``` + + diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 753d3575a8..70e96d5739 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,4 +1,3 @@ -#TODO rewrite doc strings """ PINOODE(chain, OptimizationOptimisers.Adam(0.1), @@ -17,6 +16,7 @@ neural operator, which is used as a solver for a parametrized `ODEProblem`. * `opt`: The optimizer to train the neural network. * `bounds`: A dictionary containing the bounds for the parameters of the neural network in which will be train the prediction of parametric ODE. +* `number_of_parameters`: The number of points of train set in parameters boundaries. ## Keyword Arguments diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ae90643fb0..cebda9f681 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -154,41 +154,52 @@ end #vector outputs and multiple parameters @testset "Example du = cos(p * t)" begin - function equation1(u, p, t) - p1, p2 = p[1], p[2] - cos(p1 * t) + p2 - end - - equation = (u, p, t) -> p[1] * cos(p[2] * t) + equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) + number_of_parameter = 3 deeponet = LuxNeuralOperators.DeepONet( Chain( - Dense(2 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10)), + 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))) - u = rand(2, 50, 1) + u = rand(3, 50) v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] - bounds = [(0.1f0, pi), (1.0f0, 2.0f0)] + bounds = [(1.0f0, pi), (1.0f0, 2.0f0), (2.0f0, 3.0f0)] number_of_parameters = 50 - dt = (tspan[2] - tspan[1]) / 40 - strategy = GridTraining(dt) + strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 3000) - ga = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) - 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 = sol.t[2] - ground_solution = reduce(hcat, - [[ga(u0, p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + 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 - @test ground_solution≈sol.u rtol=0.01 + 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 + + (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 + + p, t = get_trainset(bounds, tspan, 100, 0.01f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 end From 7953468e9e5ff96373d8c897a0749b8e4da6b45a Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 1 Jul 2024 17:40:49 +0400 Subject: [PATCH 065/153] output vector --- src/pino_ode_solve.jl | 49 ++++++++++++++++++++-------- test/PINO_ode_tests.jl | 72 ++++++++++++++++++++++++++++++++++++++---- 2 files changed, 101 insertions(+), 20 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 70e96d5739..69b5fb7ce5 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -52,7 +52,9 @@ function PINOODE(chain, strategy = nothing, additional_loss = nothing, kwargs...) - !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) + !(chain isa Lux.AbstractExplicitLayer || + any(x -> isa(x, Lux.AbstractExplicitLayer), chain)) && + (chain = Lux.transform(chain)) PINOODE(chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss, kwargs) end @@ -161,7 +163,15 @@ function get_trainset(strategy::StochasticTraining, bounds, number_of_parameters end function generate_loss( - strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) + strategy::StochasticTraining, prob::ODEProblem, phi::Vector, bounds, number_of_parameters, tspan) + function loss(θ, _) + x = get_trainset(strategy, bounds, number_of_parameters, tspan) + initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) + end +end + +function generate_loss( + strategy::StochasticTraining, prob::ODEProblem, phi::PINOPhi, bounds, number_of_parameters, tspan) function loss(θ, _) x = get_trainset(strategy, bounds, number_of_parameters, tspan) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) @@ -175,7 +185,7 @@ end (f::PINOODEInterpolation)(x) = f.phi(x, f.θ) -SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation" +SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural operator interpolation" SciMLBase.allowscomplex(::PINOODE) = true function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, @@ -187,26 +197,37 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, saveat = nothing, maxiters = nothing) @unpack tspan, u0, f = prob - @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg + @unpack chain, opt, bounds, number_of_parameters, init_params_, strategy, additional_loss = alg - if !isa(chain, CompactLuxLayer{:DeepONet,}) + if !isa(chain, CompactLuxLayer{:DeepONet,}) && !isa(chain, Vector) #!any(x -> isa(x, CompactLuxLayer{:DeepONet,}), chain) error("Only DeepONet neural networks are supported with PINO ODE") end - !(chain isa Lux.AbstractExplicitLayer) && - error("Only Lux.AbstractExplicitLayer neural networks are supported") - - phi, init_params = generate_pino_phi_θ(chain, init_params) + # !(chain isa Lux.AbstractExplicitLayer) && + # error("Only Lux.AbstractExplicitLayer neural networks are supported") + if !isa(chain, Vector) + phi, init_params = generate_pino_phi_θ(chain[1], init_params_) + else + phi, init_params = [], [] + for c in chain + p, ip = generate_pino_phi_θ(c, init_params_) + push!(phi, p) + push!(init_params, ip) + end + end isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - in_dim = chain.layers.branch.layers.layer_1.in_dims - u = rand(in_dim, number_of_parameters) - v = rand(1, 10, 1) - x = (u, v) - phi(x, init_params) + #TODO check for 'isa(chain, Vector)' + if !isa(chain, Vector) + in_dim = chain.layers.branch.layers.layer_1.in_dims + u = rand(in_dim, number_of_parameters) + v = rand(1, 10, 1) + x = (u, v) + phi(x, init_params) + end catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of input data and chain should match")) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index cebda9f681..50512ae85c 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -37,6 +37,7 @@ using NeuralPDE sol = solve(prob, alg, verbose = false, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + dt = 0.025f0 function get_trainset(bounds, tspan , number_of_parameters, dt) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = collect(reshape(p_, 1, size(p_, 1))) @@ -76,7 +77,7 @@ end opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 3000) sol.original.objective #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / @@ -144,7 +145,7 @@ end alg = PINOODE( deeponet, opt, bounds, number_of_parameters; strategy = strategy, additional_loss = additional_loss_) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = reshape(p_, 1, size(p_)[1]) ground_solution = ground_analytic.(u0, p, vec(sol.t[2])) @@ -152,17 +153,17 @@ end @test ground_solution≈sol.u rtol=0.01 end -#vector outputs and multiple parameters +#multiple parameters @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) - number_of_parameter = 3 + input_branch_size = 3 deeponet = LuxNeuralOperators.DeepONet( Chain( - Dense(number_of_parameter => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10)), + Dense(input_branch_size => 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))) @@ -176,7 +177,7 @@ end strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 3000) function get_trainset(bounds, tspan, number_of_parameters, dt) p_ = [range(start = b[1], length = number_of_parameters, stop = b[2]) @@ -203,3 +204,62 @@ end predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.01 end + +#vector output + +linear = (u, p, t) -> [cos(2pi * t), sin(2pi * t)] +tspan = (0.0f0, 1.0f0) +u0 = [0.0f0, -1.0f0 / 2pi] +linear_analytic = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] +@testset "Example ode system: du1 = cos(p * t); du2 = sin(p * t)" begin + equation = (u, p, t) -> [cos(p[1] * t), sin(p[2] * t)] + tspan = (0.0f0, 1.0f0) + u0 = 1.0f0 + prob = ODEProblem(equation, u0, tspan) + + input_branch_size = 2 + deeponet1 = LuxNeuralOperators.DeepONet( + Chain( + Dense(input_branch_size => 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))) + deeponet2 = LuxNeuralOperators.DeepONet( + Chain( + Dense(input_branch_size => 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))) + + deeponets = [deeponet1, deeponet1] + + bounds = [(1.0f0, pi), (2.0f0, 3.0f0)] + number_of_parameters = 50 + strategy = StochasticTraining(40) + opt = OptimizationOptimisers.Adam(0.03) + alg = PINOODE(deeponets, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 3000) + + 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 + + ground_solution = (u0, p, t) -> [u0[1] + sin(p * t) / (p), u0[2] - cos(p * t) / (p)] + 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 + + (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 + + p, t = get_trainset(bounds, tspan, 100, 0.01f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 +end From f6f040512ee2842053ddc5c360d7b207ecbb0f86 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:50:13 +0400 Subject: [PATCH 066/153] Update docs/src/manual/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/manual/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/manual/pino_ode.md b/docs/src/manual/pino_ode.md index 6177341dc7..c26ef79582 100644 --- a/docs/src/manual/pino_ode.md +++ b/docs/src/manual/pino_ode.md @@ -1,4 +1,4 @@ -# Physics-Informed Neural operator for solve ODEs +# Physics-Informed Neural Operator (PINO) for ODEs ```@docs PINOODE From e16b168857dff218dd9ae01e4b173b6986a3390a Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:50:33 +0400 Subject: [PATCH 067/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 0a7229d621..6bbbec8371 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -2,7 +2,7 @@ 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. +## Operator Learning for a family of parametric ODEs 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. From 596a80a9e68b6451df5a73ea76802f3323af1e55 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:51:01 +0400 Subject: [PATCH 068/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 6bbbec8371..3f7dd578d5 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -1,4 +1,4 @@ -# Physics Informed Neural Operator for ODEs Solvers +# 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). From 445dfd09b48566b3eb580fed906386ec1851c967 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:52:28 +0400 Subject: [PATCH 069/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 3f7dd578d5..d8ed639211 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -63,7 +63,7 @@ end ground_solution_ = ground_solution_f(p,t) predict = sol.interp((p,t)) -# Calculate the mean error and the standard deviation of the errors +# Calculate the mean error and the standard deviation of the errors errors = ground_solution_ - predict mean_error = mean(errors) std_error = std(errors) From d7b7a361918f09603dd48d9d1dee8e79537ec253 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:52:38 +0400 Subject: [PATCH 070/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index d8ed639211..165ba9b0f5 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -36,7 +36,7 @@ number_of_parameters = 50 strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) -sol = solve(prob, alg, verbose = true, maxiters = 3000) +sol = solve(prob, alg, verbose = false, maxiters = 3000) ``` 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 From fe40cfd51432c2340760b2c10e0223d86254ff1a Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:52:47 +0400 Subject: [PATCH 071/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 165ba9b0f5..e3f8ee470f 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -32,7 +32,7 @@ v = rand(1, 40, 1) c = deeponet((u, v), θ, st)[1] bounds = [(1.0f0, pi), (1.0f0, 2.0f0), (2.0f0, 3.0f0)] -number_of_parameters = 50 +number_of_parameter_samples = 50 strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) From e45011e9df40cf435484d7d2e7a0817c3ce75461 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:53:40 +0400 Subject: [PATCH 072/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index e3f8ee470f..3a5a44fada 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -7,7 +7,7 @@ This tutorial provides an example of how to use the Physics Informed Neural Oper 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 Test # hide using OptimizationOptimisers using Lux using Statistics, Random From 7fe32449dbdb8eab00ff3304c48cf8802578aeb4 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 2 Jul 2024 14:56:07 +0400 Subject: [PATCH 073/153] Update test/runtests.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 0990c63bb7..ace868dc78 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -48,7 +48,7 @@ end end end - if GROUP == "All" || GROUP == "ODEPINO" + if GROUP == "All" || GROUP == "PINOODE" @time @safetestset "pino ode" begin include("PINO_ode_tests.jl") end end From 839e1d50e255602d7da3738f7d809a7b63d0f8c0 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 15:22:24 +0400 Subject: [PATCH 074/153] revert "vector output" --- src/pino_ode_solve.jl | 49 +++++++++---------------------- test/PINO_ode_tests.jl | 65 ++---------------------------------------- 2 files changed, 17 insertions(+), 97 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 69b5fb7ce5..70e96d5739 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -52,9 +52,7 @@ function PINOODE(chain, strategy = nothing, additional_loss = nothing, kwargs...) - !(chain isa Lux.AbstractExplicitLayer || - any(x -> isa(x, Lux.AbstractExplicitLayer), chain)) && - (chain = Lux.transform(chain)) + !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) PINOODE(chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss, kwargs) end @@ -163,15 +161,7 @@ function get_trainset(strategy::StochasticTraining, bounds, number_of_parameters end function generate_loss( - strategy::StochasticTraining, prob::ODEProblem, phi::Vector, bounds, number_of_parameters, tspan) - function loss(θ, _) - x = get_trainset(strategy, bounds, number_of_parameters, tspan) - initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) - end -end - -function generate_loss( - strategy::StochasticTraining, prob::ODEProblem, phi::PINOPhi, bounds, number_of_parameters, tspan) + strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) function loss(θ, _) x = get_trainset(strategy, bounds, number_of_parameters, tspan) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) @@ -185,7 +175,7 @@ end (f::PINOODEInterpolation)(x) = f.phi(x, f.θ) -SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural operator interpolation" +SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation" SciMLBase.allowscomplex(::PINOODE) = true function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, @@ -197,37 +187,26 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, saveat = nothing, maxiters = nothing) @unpack tspan, u0, f = prob - @unpack chain, opt, bounds, number_of_parameters, init_params_, strategy, additional_loss = alg + @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg - if !isa(chain, CompactLuxLayer{:DeepONet,}) && !isa(chain, Vector) #!any(x -> isa(x, CompactLuxLayer{:DeepONet,}), chain) + if !isa(chain, CompactLuxLayer{:DeepONet,}) error("Only DeepONet neural networks are supported with PINO ODE") end - # !(chain isa Lux.AbstractExplicitLayer) && - # error("Only Lux.AbstractExplicitLayer neural networks are supported") - if !isa(chain, Vector) - phi, init_params = generate_pino_phi_θ(chain[1], init_params_) - else - phi, init_params = [], [] - for c in chain - p, ip = generate_pino_phi_θ(c, init_params_) - push!(phi, p) - push!(init_params, ip) - end - end + !(chain isa Lux.AbstractExplicitLayer) && + error("Only Lux.AbstractExplicitLayer neural networks are supported") + + phi, init_params = generate_pino_phi_θ(chain, init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - #TODO check for 'isa(chain, Vector)' - if !isa(chain, Vector) - in_dim = chain.layers.branch.layers.layer_1.in_dims - u = rand(in_dim, number_of_parameters) - v = rand(1, 10, 1) - x = (u, v) - phi(x, init_params) - end + in_dim = chain.layers.branch.layers.layer_1.in_dims + u = rand(in_dim, number_of_parameters) + v = rand(1, 10, 1) + x = (u, v) + phi(x, init_params) catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of input data and chain should match")) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 50512ae85c..00ad1ca430 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -3,7 +3,7 @@ using OptimizationOptimisers using Lux using Statistics, Random using LuxNeuralOperators -using NeuralPDE +# using NeuralPDE # dG(u(t, p), t) = f(G,u(t, p)) @testset "Example du = cos(p * t)" begin @@ -77,7 +77,7 @@ end opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 3000) + sol = solve(prob, alg, verbose = true, maxiters = 3000) sol.original.objective #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / @@ -145,7 +145,7 @@ end alg = PINOODE( deeponet, opt, bounds, number_of_parameters; strategy = strategy, additional_loss = additional_loss_) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + sol = solve(prob, alg, verbose = true, maxiters = 2000) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = reshape(p_, 1, size(p_)[1]) ground_solution = ground_analytic.(u0, p, vec(sol.t[2])) @@ -204,62 +204,3 @@ end predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.01 end - -#vector output - -linear = (u, p, t) -> [cos(2pi * t), sin(2pi * t)] -tspan = (0.0f0, 1.0f0) -u0 = [0.0f0, -1.0f0 / 2pi] -linear_analytic = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] -@testset "Example ode system: du1 = cos(p * t); du2 = sin(p * t)" begin - equation = (u, p, t) -> [cos(p[1] * t), sin(p[2] * t)] - tspan = (0.0f0, 1.0f0) - u0 = 1.0f0 - prob = ODEProblem(equation, u0, tspan) - - input_branch_size = 2 - deeponet1 = LuxNeuralOperators.DeepONet( - Chain( - Dense(input_branch_size => 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))) - deeponet2 = LuxNeuralOperators.DeepONet( - Chain( - Dense(input_branch_size => 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))) - - deeponets = [deeponet1, deeponet1] - - bounds = [(1.0f0, pi), (2.0f0, 3.0f0)] - number_of_parameters = 50 - strategy = StochasticTraining(40) - opt = OptimizationOptimisers.Adam(0.03) - alg = PINOODE(deeponets, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) - - 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 - - ground_solution = (u0, p, t) -> [u0[1] + sin(p * t) / (p), u0[2] - cos(p * t) / (p)] - 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 - - (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) - ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 - - p, t = get_trainset(bounds, tspan, 100, 0.01f0) - ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 -end From 53a52f84a15f40ddf8893caec6623cd53e8ad609 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 15:22:51 +0400 Subject: [PATCH 075/153] fix --- test/PINO_ode_tests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 00ad1ca430..76e29bac0e 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -3,7 +3,7 @@ using OptimizationOptimisers using Lux using Statistics, Random using LuxNeuralOperators -# using NeuralPDE +using NeuralPDE # dG(u(t, p), t) = f(G,u(t, p)) @testset "Example du = cos(p * t)" begin From a39bb69d3b9acb2e3c5014db8bef3d87e017d7be Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 16:08:49 +0400 Subject: [PATCH 076/153] Add Unregistered LuxNeuralOperators --- .github/workflows/CI.yml | 6 ++++++ Project.toml | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 0458b4cb32..2b58bf8cc3 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -59,3 +59,9 @@ jobs: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: true + - name: Add Unregistered LuxNeuralOperators Package + run: | + julie -e 'using Pkg; Pkg.add(PackageSpec(url="https://github.com/LuxDL/LuxNeuralOperators.jl"))'; + 'Pkg.instantiate()' + + \ No newline at end of file diff --git a/Project.toml b/Project.toml index c7e3634d99..80bc20b5f2 100644 --- a/Project.toml +++ b/Project.toml @@ -20,7 +20,7 @@ Integrals = "de52edbc-65ea-441a-8357-d3a637375a31" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" +#LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" while it is not registered MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" From 2983c18ae0eb31e2d8e9d94170bbb9d0ed908999 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 16:20:08 +0400 Subject: [PATCH 077/153] CI.yml fix --- .github/workflows/CI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 2b58bf8cc3..0a57350c78 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -26,7 +26,7 @@ jobs: - Logging - Forward - DGM - - ODEPINO + - PINOODE - NNODE - NeuralAdapter - IntegroDiff From 885aa29dcd20c2d28773d65e9c25789f2999209a Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 16:26:54 +0400 Subject: [PATCH 078/153] revert Add Unregistered LuxNeuralOperators --- .github/workflows/CI.yml | 4 ---- Project.toml | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 0a57350c78..42963419e4 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -59,9 +59,5 @@ jobs: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: true - - name: Add Unregistered LuxNeuralOperators Package - run: | - julie -e 'using Pkg; Pkg.add(PackageSpec(url="https://github.com/LuxDL/LuxNeuralOperators.jl"))'; - 'Pkg.instantiate()' \ No newline at end of file diff --git a/Project.toml b/Project.toml index 80bc20b5f2..c7e3634d99 100644 --- a/Project.toml +++ b/Project.toml @@ -20,7 +20,7 @@ Integrals = "de52edbc-65ea-441a-8357-d3a637375a31" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -#LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" while it is not registered +LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" From 38a872ec94dddb1a3cea887446be6ec4f0459eaa Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 16:28:54 +0400 Subject: [PATCH 079/153] fix --- src/NeuralPDE.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 2a2b1fa374..28aac16de6 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -58,7 +58,7 @@ include("PDE_BPINN.jl") include("dgm.jl") -export NNODE, NNDAE, PINOODE, DeepONet +export NNODE, NNDAE, PINOODE PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, From 28447b61e5c0ee312b4f70363dc4c73db23af0ef Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 17:21:36 +0400 Subject: [PATCH 080/153] add unregistered LuxNeuralOperators attempt 2 --- .github/workflows/CI.yml | 2 ++ Project.toml | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 42963419e4..c5cdb8c2a0 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -34,6 +34,8 @@ jobs: - "1" steps: - uses: actions/checkout@v4 + - name: Add Unregistered Package + run: julia -e 'using Pkg; Pkg.add(PackageSpec(url="https://github.com/LuxDL/LuxNeuralOperators.jl")); Pkg.instantiate()' - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} diff --git a/Project.toml b/Project.toml index c7e3634d99..729b995c69 100644 --- a/Project.toml +++ b/Project.toml @@ -20,7 +20,7 @@ Integrals = "de52edbc-65ea-441a-8357-d3a637375a31" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" +LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" #while it is not registered MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" From 129290016398b0b1eb6aaff306e315ab589ee98f Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 2 Jul 2024 17:55:27 +0400 Subject: [PATCH 081/153] undo "add unregistered LuxNeuralOperators attempt 2" --- .github/workflows/CI.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c5cdb8c2a0..42963419e4 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -34,8 +34,6 @@ jobs: - "1" steps: - uses: actions/checkout@v4 - - name: Add Unregistered Package - run: julia -e 'using Pkg; Pkg.add(PackageSpec(url="https://github.com/LuxDL/LuxNeuralOperators.jl")); Pkg.instantiate()' - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} From 10e08c2d09bc78bd311b26c9058fe9c6de8e2681 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 12 Jul 2024 16:54:34 +0400 Subject: [PATCH 082/153] FNO first shot --- src/pino_ode_solve.jl | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 70e96d5739..acf5190d12 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -91,6 +91,14 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: CompactLuxLayer{:Dee (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end +function dfdx(phi::PINOPhi{C, T}, x::Tuple, + θ) where {C <: CompactLuxLayer{:FourierNeuralOperator,}, T} + #TODO + x_right = (branch_right, trunk_right) + (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) +end + + function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { C <: CompactLuxLayer{:DeepONet,}, T} @@ -108,7 +116,7 @@ function physics_loss( norm = prod(size(du)) sum(abs2, du .- f_vec) / norm end - +#TODO Lux.AbstractExplicitLayer function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { C <: CompactLuxLayer{:DeepONet,}, T} @@ -122,6 +130,7 @@ function initial_condition_loss( sum(abs2, u .- u0) / norm end +#TODO for input FourierNeuralOperator function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspan) dt = strategy.dx if size(bounds,1) == 1 From 61fa804761bc3de5dc7ec1cd9e5be404b690df40 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 12 Jul 2024 16:55:33 +0400 Subject: [PATCH 083/153] fix --- test/PINO_ode_tests.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 76e29bac0e..d999f62fe3 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -2,7 +2,7 @@ using Test using OptimizationOptimisers using Lux using Statistics, Random -using LuxNeuralOperators +using NeuralOperators using NeuralPDE # dG(u(t, p), t) = f(G,u(t, p)) @@ -16,6 +16,9 @@ using NeuralPDE Dense(1 => 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))) + #TODO test FNO + fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) + u = rand(1, 50) v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) From 99e8e2dab39583da3268a8ba2ff0dfab0b108aa7 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:34:09 +0400 Subject: [PATCH 084/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 3a5a44fada..67d67a84f2 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -4,7 +4,7 @@ This tutorial provides an example of how to use the Physics Informed Neural Oper ## Operator Learning for a family of parametric ODEs -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. +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 # hide From c01c0ade574b6c570dc7ca7f429301ead0944c72 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:35:59 +0400 Subject: [PATCH 085/153] Update .github/workflows/CI.yml Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- .github/workflows/CI.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 42963419e4..0eb3f3de05 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -58,6 +58,4 @@ jobs: with: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: true - - \ No newline at end of file + fail_ci_if_error: true \ No newline at end of file From 1a8fd173629f62a6ab9c0fedb352427bde79de7e Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:36:14 +0400 Subject: [PATCH 086/153] Update src/NeuralPDE.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/NeuralPDE.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 28aac16de6..a610d6e2e8 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -57,7 +57,6 @@ include("BPINN_ode.jl") include("PDE_BPINN.jl") include("dgm.jl") - export NNODE, NNDAE, PINOODE PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, From c5bf9ca960da8ebb742d7afa05867d13082e7aab Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:36:32 +0400 Subject: [PATCH 087/153] Update Project.toml Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 729b995c69..df15c59bdf 100644 --- a/Project.toml +++ b/Project.toml @@ -98,4 +98,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"] \ No newline at end of file +test = ["Aqua", "Test", "CUDA", "SafeTestsets", "OptimizationOptimJL", "Pkg", "OrdinaryDiffEq", "LineSearches", "LuxCUDA", "Flux", "MethodOfLines"] From c5d2501dd52cdd3ae9e8c9550f0a39bcd41a3eb4 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:36:40 +0400 Subject: [PATCH 088/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 67d67a84f2..b46753d40a 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -92,4 +92,3 @@ plot(predict[:, i], label = "Predicted") plot!(ground_solution_[:, i], label = "Ground truth") ``` - From fc38b5044fe716e31d6056c8cb9184ab2cca508e Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:36:54 +0400 Subject: [PATCH 089/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index b46753d40a..b1d985b8a6 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -39,8 +39,7 @@ alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 3000) ``` -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. +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 From 09e55557c23d6e3f1912cd03e25ad47cbf61a4d8 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:51:21 +0400 Subject: [PATCH 090/153] Update src/pino_ode_solve.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/pino_ode_solve.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index acf5190d12..e20e1dec2f 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -1,6 +1,6 @@ """ PINOODE(chain, - OptimizationOptimisers.Adam(0.1), + opt, bounds; init_params = nothing, strategy = nothing From 8ae7b0f375e90410d05f8050109336e877c867a3 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Wed, 17 Jul 2024 16:51:43 +0400 Subject: [PATCH 091/153] Update src/pino_ode_solve.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/pino_ode_solve.jl | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index e20e1dec2f..fdfbeaf301 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -67,11 +67,8 @@ end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) - if init_params === nothing - init_params = ComponentArrays.ComponentArray(θ) - else - init_params = ComponentArrays.ComponentArray(init_params) - end + init_params = isnothing(init_params) ? θ : init_params + init_params = ComponentArrays.ComponentArray(init_params) PINOPhi(chain, st), init_params end From 15e61247ab598edab02bd37585320f34834829d8 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 22 Jul 2024 19:29:31 +0400 Subject: [PATCH 092/153] vector output --- Project.toml | 2 +- src/NeuralPDE.jl | 2 +- src/pino_ode_solve.jl | 3 ++ test/PINO_ode_tests.jl | 77 ++++++++++++++++++++++++++++++++++++++---- 4 files changed, 75 insertions(+), 9 deletions(-) diff --git a/Project.toml b/Project.toml index df15c59bdf..13b9b35db0 100644 --- a/Project.toml +++ b/Project.toml @@ -20,10 +20,10 @@ Integrals = "de52edbc-65ea-441a-8357-d3a637375a31" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -LuxNeuralOperators = "c0ba2cc5-a80b-46ec-84b3-091eb317b01d" #while it is not registered MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" +NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1" diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index a610d6e2e8..a0b96ceebf 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -33,7 +33,7 @@ using UnPack: @unpack import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor using ChainRulesCore: @non_differentiable -using LuxNeuralOperators +using NeuralOperators RuntimeGeneratedFunctions.init(@__MODULE__) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index fdfbeaf301..96f189bdc2 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -103,6 +103,9 @@ function physics_loss( f = prob.f out = phi(x, θ) if size(p,1) == 1 + if size(out)[1] == 1 + out = dropdims(out, dims = 1) + end fs = f.(out, p, vec(t)) f_vec = vec(fs) else diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index d999f62fe3..d60ebc3a9a 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -2,23 +2,33 @@ using Test using OptimizationOptimisers using Lux using Statistics, Random +#TODO remove after register Lux.NeuralOperators +# using Pkg +# Pkg.add(PackageSpec(url = "https://github.com/LuxDL/NeuralOperators.jl.git")) using NeuralOperators -using NeuralPDE +# using NeuralPDE # dG(u(t, p), t) = f(G,u(t, p)) + +# #TODO test FNO +# fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) + @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) - deeponet = LuxNeuralOperators.DeepONet( + # deeponet = NeuralOperators.DeepONet( + # Chain( + # Dense(1 => 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))) + deeponet = NeuralOperators.DeepONet( Chain( Dense(1 => 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))) - #TODO test FNO - fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - + Dense(10 => 10, Lux.tanh_fast)), + additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1))) u = rand(1, 50) v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) @@ -32,12 +42,16 @@ using NeuralPDE θ, st = Lux.setup(Random.default_rng(), trunk) t = trunk(v, θ, st)[1] + # additional = deeponet.layers.additional + # θ, st = Lux.setup(Random.default_rng(), additional) + # a = additional(rand(10,40,50), θ, st)[1] + bounds = [(pi, 2pi)] number_of_parameters = 50 strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + sol = solve(prob, alg, verbose = true, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 @@ -207,3 +221,52 @@ end predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.01 end + +#vector output +@testset "Example du = cos(p * t)" begin + equation = (u, p, t) -> [cos(p[1] * t), sin(p[2]*t)] + tspan = (0.0f0, 1.0f0) + u0 = [1.0f0, 0.0f0] + prob = ODEProblem(equation, u0, tspan) + deeponet = NeuralOperators.DeepONet( + Chain( + Dense(1 => 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)), + additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) + + bounds = [(pi, 2pi), (pi/2, 3pi/2)] + number_of_parameters = 50 + strategy = StochasticTraining(40) + opt = OptimizationOptimisers.Adam(0.01) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) + + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + dt = 0.025f0 + 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 + p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + + ground_solution = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] + 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 + + (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 + + p, t = get_trainset(bounds, tspan, 100, 0.01f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 +end From efd2a40b669a1c46d2dd45b36521619b6135dd2b Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 23 Jul 2024 18:27:28 +0400 Subject: [PATCH 093/153] wip FourierNeuralOperator --- src/pino_ode_solve.jl | 32 +++++++++++++++++++------------- test/PINO_ode_tests.jl | 17 +++++++++++++++-- 2 files changed, 34 insertions(+), 15 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 96f189bdc2..2b6a885303 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -68,7 +68,7 @@ end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) init_params = isnothing(init_params) ? θ : init_params - init_params = ComponentArrays.ComponentArray(init_params) + init_params = chain.name == "FourierNeuralOperator" ? θ : ComponentArrays.ComponentArray(init_params) PINOPhi(chain, st), init_params end @@ -95,10 +95,9 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end - function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { - C <: CompactLuxLayer{:DeepONet,}, T} + C, T} #C <: CompactLuxLayer{:DeepONet,} p, t = x f = prob.f out = phi(x, θ) @@ -119,7 +118,7 @@ end #TODO Lux.AbstractExplicitLayer function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { - C <: CompactLuxLayer{:DeepONet,}, T} + C, T} ##TODO for DeepOet and FNO C <: CompactLuxLayer{:DeepONet,} p, t = x t0 = reshape([prob.tspan[1]], (1, 1, 1)) # t[:, [1], :] x0 = (p, t0) @@ -198,24 +197,31 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, @unpack tspan, u0, f = prob @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg - if !isa(chain, CompactLuxLayer{:DeepONet,}) - error("Only DeepONet neural networks are supported with PINO ODE") - end - !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") + if !(isa(chain, CompactLuxLayer{:DeepONet,}) || chain.name == "FourierNeuralOperator") + error("Only DeepONet and FourierNeuralOperator neural networks are supported with PINO ODE") + end + phi, init_params = generate_pino_phi_θ(chain, init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - in_dim = chain.layers.branch.layers.layer_1.in_dims - u = rand(in_dim, number_of_parameters) - v = rand(1, 10, 1) - x = (u, v) - phi(x, init_params) + if chain isa CompactLuxLayer{:DeepONet,} #TODO chain.name == "DeepONet" + in_dim = chain.layers.branch.layers.layer_1.in_dims + u = rand(in_dim, number_of_parameters) + v = rand(1, 10, 1) + x = (u, v) + phi(x, init_params) + end + if chain.name == "FourierNeuralOperator" + in_dim = chain.layers.lifting.in_dims + v = rand(in_dim, number_of_parameters, 40) + phi(v, θ) + end catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of input data and chain should match")) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index d60ebc3a9a..57a8137605 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -51,7 +51,7 @@ using NeuralOperators strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 @@ -65,15 +65,28 @@ using NeuralOperators p,t = get_trainset(bounds, tspan, number_of_parameters, dt) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol.interp((p, t)) + predict_sol = dropdims(sol.interp((p, t)), dims=1) @test ground_solution≈predict_sol rtol=0.01 p, t = get_trainset(bounds, tspan, 100, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) + predict_sol = dropdims(sol.interp((p, t)), dims = 1) + + @test ground_solution≈predict_sol rtol=0.01 + + fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) + v = rand(2, 40,50) + θ, st = Lux.setup(Random.default_rng(), fno) + c = fno(v, θ, st)[1] + + alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) + predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.01 + end @testset "Example du = cos(p * t) + u" begin From 89fc2d9f7e026758ae7cf17f02699ccfb918e4d7 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 14 Aug 2024 19:04:39 +0400 Subject: [PATCH 094/153] update --- src/pino_ode_solve.jl | 77 +++++++++++++++++++++++++++++------------- test/PINO_ode_tests.jl | 50 +++++++++++++++------------ 2 files changed, 81 insertions(+), 46 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 2b6a885303..9c5ec9786f 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -68,14 +68,15 @@ end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) init_params = isnothing(init_params) ? θ : init_params - init_params = chain.name == "FourierNeuralOperator" ? θ : ComponentArrays.ComponentArray(init_params) + # init_params = chain.name == "FourierNeuralOperator" ? θ : ComponentArrays.ComponentArray(init_params) + init_params = ComponentArrays.ComponentArray(init_params) PINOPhi(chain, st), init_params end function (f::PINOPhi{C, T})( x, θ) where {C <: Lux.AbstractExplicitLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) - ChainRulesCore.@ignore_derivatives f.st = st + # ChainRulesCore.@ignore_derivatives f.st = st y end @@ -88,37 +89,60 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: CompactLuxLayer{:Dee (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end -function dfdx(phi::PINOPhi{C, T}, x::Tuple, - θ) where {C <: CompactLuxLayer{:FourierNeuralOperator,}, T} - #TODO - x_right = (branch_right, trunk_right) - (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) -end +# using Zygote +# p, t = x +# branch_left, branch_right = p, p +# trunk_left, trunk_right = t .+ sqrt(eps(Float32)), t +# x_left = (branch_left, trunk_left) +# x_right = (branch_right, trunk_right) +# du(θ) = (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(Float32)) +# #TODO error here phi(x_right, θ)) / sqrt(eps(Float32)) + + +# function dfdx(phi::PINOPhi{C, T}, x::Tuple, +# θ) where {C <: CompactLuxLayer{:FourierNeuralOperator,}, T} +# #TODO +# p, t = x +# branch_left, branch_right = p, p +# trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t +# x_left = reduce(vcat,[branch_left,trunk_left]) +# x_right = (branch_right, trunk_right) +# (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) +# end function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { - C, T} #C <: CompactLuxLayer{:DeepONet,} + C <: CompactLuxLayer{:DeepONet,}, T} p, t = x f = prob.f out = phi(x, θ) - if size(p,1) == 1 + # if size(p,1) == 1 + #TODO if size(out)[1] == 1 out = dropdims(out, dims = 1) end + # out = dropdims(out, dims = 1) fs = f.(out, p, vec(t)) f_vec = vec(fs) - else - f_vec = reduce( - vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - end + # else + # # f_vec = reduce( + # # vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + # f_vec = reduce( + # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + # end + + # norm = prod(size(out)) + # out_vec = vec(out) + # sum(abs2, du .- f_vec) / norm du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm + end #TODO Lux.AbstractExplicitLayer function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { - C, T} ##TODO for DeepOet and FNO C <: CompactLuxLayer{:DeepONet,} + C <: CompactLuxLayer{:DeepONet,}, T} p, t = x t0 = reshape([prob.tspan[1]], (1, 1, 1)) # t[:, [1], :] x0 = (p, t0) @@ -147,14 +171,6 @@ function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspa (p, t) end -function generate_loss( - strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) - x = get_trainset(strategy, bounds, number_of_parameters, tspan) - function loss(θ, _) - initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) - end -end - function get_trainset(strategy::StochasticTraining, bounds, number_of_parameters, tspan) if size(bounds,1) == 1 bound = bounds[1] @@ -168,6 +184,19 @@ function get_trainset(strategy::StochasticTraining, bounds, number_of_parameters (p, t) end +function generate_loss( + strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) + x = get_trainset(strategy, bounds, number_of_parameters, tspan) + function loss(θ, _) + initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) + end +end + +# using Zygote +# Zygote.gradient((θ) -> initial_condition_loss(phi, prob, x, θ), θ) +# Zygote.gradient((θ) -> physics_loss(phi, prob, x, θ), θ) #TODO + + function generate_loss( strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) function loss(θ, _) @@ -210,7 +239,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - if chain isa CompactLuxLayer{:DeepONet,} #TODO chain.name == "DeepONet" + if chain isa CompactLuxLayer{:DeepONet,} in_dim = chain.layers.branch.layers.layer_1.in_dims u = rand(in_dim, number_of_parameters) v = rand(1, 10, 1) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 57a8137605..0efb8a9e05 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -2,37 +2,41 @@ using Test using OptimizationOptimisers using Lux using Statistics, Random -#TODO remove after register Lux.NeuralOperators -# using Pkg -# Pkg.add(PackageSpec(url = "https://github.com/LuxDL/NeuralOperators.jl.git")) using NeuralOperators # using NeuralPDE -# dG(u(t, p), t) = f(G,u(t, p)) - -# #TODO test FNO -# fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) - # deeponet = NeuralOperators.DeepONet( - # Chain( - # Dense(1 => 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))) deeponet = NeuralOperators.DeepONet( Chain( Dense(1 => 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)), - additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1))) + Dense(10 => 10, Lux.tanh_fast))) + #TODO deeponet fail with additional layer on gradient + # deeponet = NeuralOperators.DeepONet( + # Chain( + # Dense(1 => 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)), + # additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1))) + u = rand(1, 50) v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] + # left = (u, v) + # right = (u, v .+ 1.0) + # using Zygote + # Zygote.gradient((θ) -> sum(deeponet((u, v), θ, st)[1]), θ) + + # deeponet((u .- 1.9, v), θ, st)[1] .- deeponet((u, v), θ, st)[1] + # Zygote.gradient((θ) -> sum(deeponet((u, v), θ, st)[1]), θ) + #TODO open issue DeepONet with additional fail compiler + # Zygote.gradient( + # (θ) -> sum(deeponet((u .- 1.9, v), θ, st)[1] .- deeponet((u, v), θ, st)[1]), θ) branch = deeponet.layers.branch θ, st = Lux.setup(Random.default_rng(), branch) @@ -42,16 +46,13 @@ using NeuralOperators θ, st = Lux.setup(Random.default_rng(), trunk) t = trunk(v, θ, st)[1] - # additional = deeponet.layers.additional - # θ, st = Lux.setup(Random.default_rng(), additional) - # a = additional(rand(10,40,50), θ, st)[1] - bounds = [(pi, 2pi)] number_of_parameters = 50 - strategy = StochasticTraining(40) + # strategy = StochasticTraining(40) + strategy = GridTraining(0.025f0) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + sol = solve(prob, alg, verbose = true, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 @@ -65,7 +66,8 @@ using NeuralOperators p,t = get_trainset(bounds, tspan, number_of_parameters, dt) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = dropdims(sol.interp((p, t)), dims=1) + # predict_sol = dropdims(sol.interp((p, t)), dims=1) + predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.01 @@ -75,7 +77,11 @@ using NeuralOperators @test ground_solution≈predict_sol rtol=0.01 + #TODO + #FourierNeuralOperator fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) + + fno isa CompactLuxLayer{:FourierNeuralOperator,} v = rand(2, 40,50) θ, st = Lux.setup(Random.default_rng(), fno) c = fno(v, θ, st)[1] From 5ab8a6a9c4b1bf2fcc58859a42dc415207ac974b Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 16 Aug 2024 19:31:42 +0400 Subject: [PATCH 095/153] update --- src/pino_ode_solve.jl | 138 ++++++++++++++++++++++++++--------------- test/PINO_ode_tests.jl | 51 ++++++--------- 2 files changed, 108 insertions(+), 81 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 9c5ec9786f..b3c8aa8a55 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -76,7 +76,7 @@ end function (f::PINOPhi{C, T})( x, θ) where {C <: Lux.AbstractExplicitLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) - # ChainRulesCore.@ignore_derivatives f.st = st + ChainRulesCore.@ignore_derivatives f.st = st y end @@ -89,26 +89,11 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: CompactLuxLayer{:Dee (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end -# using Zygote -# p, t = x -# branch_left, branch_right = p, p -# trunk_left, trunk_right = t .+ sqrt(eps(Float32)), t -# x_left = (branch_left, trunk_left) -# x_right = (branch_right, trunk_right) -# du(θ) = (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(Float32)) -# #TODO error here phi(x_right, θ)) / sqrt(eps(Float32)) - - -# function dfdx(phi::PINOPhi{C, T}, x::Tuple, -# θ) where {C <: CompactLuxLayer{:FourierNeuralOperator,}, T} -# #TODO -# p, t = x -# branch_left, branch_right = p, p -# trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t -# x_left = reduce(vcat,[branch_left,trunk_left]) -# x_right = (branch_right, trunk_right) -# (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) -# end +#FourierNeuralOperator and Chain +function dfdx(phi::PINOPhi{C, T}, x::Array, θ) where {C, T} + ε = [zeros(eltype(x), size(x)[1] - 1)..., sqrt(eps(eltype(x)))] + (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) +end function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { @@ -131,21 +116,38 @@ function physics_loss( # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) # end - # norm = prod(size(out)) - # out_vec = vec(out) - # sum(abs2, du .- f_vec) / norm du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm end -#TODO Lux.AbstractExplicitLayer + +#FourierNeuralOperator and Chain +function physics_loss( + phi::PINOPhi{C, T}, prob::ODEProblem, x::Array, θ) where {C, T} + p, t = x[1:(end - 1), :, :], x[[end], :, :] + f = prob.f + out = phi(x, θ) + if size(out)[1] == 1 + out = dropdims(out, dims = 1) + end + + fs = f.(vec(out), vec(p), vec(t)) + f_vec = vec(fs) + # f_vec = reduce( + # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + + du = vec(dfdx(phi, x, θ)) + norm = prod(size(du)) + sum(abs2, du .- f_vec) / norm +end + function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { C <: CompactLuxLayer{:DeepONet,}, T} p, t = x t0 = reshape([prob.tspan[1]], (1, 1, 1)) # t[:, [1], :] - x0 = (p, t0) + x0 = (p, t0) #TODO one time in get_trainset ? out = phi(x0, θ) u = vec(out) u0 = vec(fill(prob.u0, size(out))) @@ -153,33 +155,71 @@ function initial_condition_loss( sum(abs2, u .- u0) / norm end -#TODO for input FourierNeuralOperator +#FourierNeuralOperator and Chain +function initial_condition_loss( + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} + p, t = x[1:end-1,:,:], x[[end],:,:] + t0 = fill(prob.tspan[1], size(p)) + x0 = reduce(vcat, (p, t0)) #TODO one time in get_trainset + + out = phi(x0, θ) + u = vec(out) + u0 = vec(fill(prob.u0, size(out))) + norm = prod(size(u0)) + sum(abs2, u .- u0) / norm +end + +#TODO for input FourierNeuralOperator and Chain function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspan) dt = strategy.dx - if size(bounds,1) == 1 - bound = bounds[1] - p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) - p = collect(reshape(p_, 1, size(p_,1))) - else + p = collect([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(range(start = tspan[1], length = 1/dt, stop = tspan[2])) + t = collect(tspan[1]:dt:tspan[2]) + # t = reshape(t_, 1, size(t_, 1), 1) + # reduce(vcat, (p, t)) + # hcat((map(points -> collect(points), Iterators.product(p,t)))...) + combinations = collect(Iterators.product(p, t)) + N = size(p, 1) + M = size(t, 1) + x = zeros(2, N, M) + + for i in 1:N + for j in 1:M + x[:, i, j] = [combinations[(i - 1) * M + j]...] + end + end + x +end + +function get_trainset(strategy::GridTraining, chain::CompactLuxLayer{:DeepONet,}, bounds, + number_of_parameters, tspan) + dt = strategy.dx + # if size(bounds,1) == 1 + # bound = bounds[1] + # p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) + # p = collect(reshape(p_, 1, size(p_,1))) + # else 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_]...) - end + # end t_ = collect(tspan[1]:dt:tspan[2]) t = reshape(t_, 1, size(t_,1), 1) (p, t) end -function get_trainset(strategy::StochasticTraining, bounds, number_of_parameters, tspan) - if size(bounds,1) == 1 - bound = bounds[1] - p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] - else +function get_trainset(strategy::StochasticTraining, + chain::CompactLuxLayer{:DeepONet,}, bounds, number_of_parameters, tspan) + # if size(bounds,1) == 1 #TODO reduce if ? + # bound = bounds[1] + # p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] + # else p = reduce(vcat, [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] for bound in bounds]) - end + # end t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points,1) .+ tspan[1] (p, t) end @@ -192,11 +232,6 @@ function generate_loss( end end -# using Zygote -# Zygote.gradient((θ) -> initial_condition_loss(phi, prob, x, θ), θ) -# Zygote.gradient((θ) -> physics_loss(phi, prob, x, θ), θ) #TODO - - function generate_loss( strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) function loss(θ, _) @@ -226,11 +261,13 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, @unpack tspan, u0, f = prob @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg - !(chain isa Lux.AbstractExplicitLayer) && + if !(chain isa Lux.AbstractExplicitLayer) error("Only Lux.AbstractExplicitLayer neural networks are supported") - if !(isa(chain, CompactLuxLayer{:DeepONet,}) || chain.name == "FourierNeuralOperator") - error("Only DeepONet and FourierNeuralOperator neural networks are supported with PINO ODE") + if !(isa(chain, CompactLuxLayer{:DeepONet,}) || + chain.name == "FourierNeuralOperator") #TODO chain.name + error("Only DeepONet and FourierNeuralOperator neural networks are supported with PINO ODE") + end end phi, init_params = generate_pino_phi_θ(chain, init_params) @@ -245,11 +282,14 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, v = rand(1, 10, 1) x = (u, v) phi(x, init_params) - end - if chain.name == "FourierNeuralOperator" + elseif chain.name == "FourierNeuralOperator" in_dim = chain.layers.lifting.in_dims v = rand(in_dim, number_of_parameters, 40) phi(v, θ) + else #TODO identifier for simple Chain + in_dim = chain.layers.layer_1.in_dims + v = rand(in_dim, number_of_parameters, 40) + phi(v, θ) end catch err if isa(err, DimensionMismatch) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 0efb8a9e05..e3472fad22 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -15,29 +15,8 @@ using NeuralOperators Dense(1 => 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))) - #TODO deeponet fail with additional layer on gradient - # deeponet = NeuralOperators.DeepONet( - # Chain( - # Dense(1 => 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)), - # additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1))) - u = rand(1, 50) v = rand(1, 40, 1) - θ, st = Lux.setup(Random.default_rng(), deeponet) - c = deeponet((u, v), θ, st)[1] - # left = (u, v) - # right = (u, v .+ 1.0) - # using Zygote - # Zygote.gradient((θ) -> sum(deeponet((u, v), θ, st)[1]), θ) - - # deeponet((u .- 1.9, v), θ, st)[1] .- deeponet((u, v), θ, st)[1] - # Zygote.gradient((θ) -> sum(deeponet((u, v), θ, st)[1]), θ) - #TODO open issue DeepONet with additional fail compiler - # Zygote.gradient( - # (θ) -> sum(deeponet((u .- 1.9, v), θ, st)[1] .- deeponet((u, v), θ, st)[1]), θ) - branch = deeponet.layers.branch θ, st = Lux.setup(Random.default_rng(), branch) b = branch(u, θ, st)[1] @@ -73,26 +52,33 @@ using NeuralOperators p, t = get_trainset(bounds, tspan, 100, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = dropdims(sol.interp((p, t)), dims = 1) + predict_sol = sol.interp((p, t)) + # predict_sol = dropdims(sol.interp((p, t)), dims = 1) @test ground_solution≈predict_sol rtol=0.01 #TODO + #ffnn #FourierNeuralOperator fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - - fno isa CompactLuxLayer{:FourierNeuralOperator,} v = rand(2, 40,50) θ, st = Lux.setup(Random.default_rng(), fno) c = fno(v, θ, st)[1] + ffnn = Lux.Chain( + Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 1)) + θ, st = Lux.setup(Random.default_rng(), ffnn) + c = ffnn(v, θ, st)[1] - alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) + alg = PINOODE(ffnn, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 2000) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.01 + alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) + end @testset "Example du = cos(p * t) + u" begin @@ -247,13 +233,14 @@ end tspan = (0.0f0, 1.0f0) u0 = [1.0f0, 0.0f0] prob = ODEProblem(equation, u0, tspan) - deeponet = NeuralOperators.DeepONet( - Chain( - Dense(1 => 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)), - additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) - + # deeponet = NeuralOperators.DeepONet( + # Chain( + # Dense(1 => 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)), + # additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) + ffnn = Lux.Chain(Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 2)) + fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) bounds = [(pi, 2pi), (pi/2, 3pi/2)] number_of_parameters = 50 strategy = StochasticTraining(40) From 433e5296cab661c96d94ad5fc45b18f4f4443110 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 20 Aug 2024 18:36:24 +0400 Subject: [PATCH 096/153] update --- src/pino_ode_solve.jl | 46 +++++++++++++++++++++++------------------- test/PINO_ode_tests.jl | 7 +++++-- 2 files changed, 30 insertions(+), 23 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index b3c8aa8a55..a225c0de35 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -76,11 +76,11 @@ end function (f::PINOPhi{C, T})( x, θ) where {C <: Lux.AbstractExplicitLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) - ChainRulesCore.@ignore_derivatives f.st = st + #ChainRulesCore.@ignore_derivatives f.st = st #TODO open issue f.st = st fail? y end -function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: CompactLuxLayer{:DeepONet,}, T} +function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} p, t = x branch_left, branch_right = p, p trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t @@ -97,7 +97,7 @@ end function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { - C <: CompactLuxLayer{:DeepONet,}, T} + C <: DeepONet, T} p, t = x f = prob.f out = phi(x, θ) @@ -124,7 +124,8 @@ end #FourierNeuralOperator and Chain function physics_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x::Array, θ) where {C, T} + phi::PINOPhi{C, T}, prob::ODEProblem, x::Array, θ) where { + C <: Union{FourierNeuralOperator, Chain}, T} p, t = x[1:(end - 1), :, :], x[[end], :, :] f = prob.f out = phi(x, θ) @@ -144,7 +145,7 @@ end function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { - C <: CompactLuxLayer{:DeepONet,}, T} + C <: DeepONet, T} p, t = x t0 = reshape([prob.tspan[1]], (1, 1, 1)) # t[:, [1], :] x0 = (p, t0) #TODO one time in get_trainset ? @@ -157,7 +158,8 @@ end #FourierNeuralOperator and Chain function initial_condition_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where {C, T} + phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { + C <: Union{FourierNeuralOperator, Chain}, T} p, t = x[1:end-1,:,:], x[[end],:,:] t0 = fill(prob.tspan[1], size(p)) x0 = reduce(vcat, (p, t0)) #TODO one time in get_trainset @@ -170,7 +172,8 @@ function initial_condition_loss( end #TODO for input FourierNeuralOperator and Chain -function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspan) +function get_trainset(strategy::GridTraining, chain::Union{FourierNeuralOperator, Chain}, + number_of_parameters, tspan) dt = strategy.dx p = collect([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_]...) @@ -192,8 +195,7 @@ function get_trainset(strategy::GridTraining, bounds, number_of_parameters, tspa x end -function get_trainset(strategy::GridTraining, chain::CompactLuxLayer{:DeepONet,}, bounds, - number_of_parameters, tspan) +function get_trainset(strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) dt = strategy.dx # if size(bounds,1) == 1 # bound = bounds[1] @@ -211,7 +213,7 @@ function get_trainset(strategy::GridTraining, chain::CompactLuxLayer{:DeepONet,} end function get_trainset(strategy::StochasticTraining, - chain::CompactLuxLayer{:DeepONet,}, bounds, number_of_parameters, tspan) + chain::DeepONet, bounds, number_of_parameters, tspan) # if size(bounds,1) == 1 #TODO reduce if ? # bound = bounds[1] # p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] @@ -226,7 +228,7 @@ end function generate_loss( strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) - x = get_trainset(strategy, bounds, number_of_parameters, tspan) + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) function loss(θ, _) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end @@ -235,7 +237,7 @@ end function generate_loss( strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) function loss(θ, _) - x = get_trainset(strategy, bounds, number_of_parameters, tspan) + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end @@ -264,8 +266,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, if !(chain isa Lux.AbstractExplicitLayer) error("Only Lux.AbstractExplicitLayer neural networks are supported") - if !(isa(chain, CompactLuxLayer{:DeepONet,}) || - chain.name == "FourierNeuralOperator") #TODO chain.name + if !(chain isa DeepONet || chain isa FourierNeuralOperator) error("Only DeepONet and FourierNeuralOperator neural networks are supported with PINO ODE") end end @@ -276,13 +277,13 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) try - if chain isa CompactLuxLayer{:DeepONet,} - in_dim = chain.layers.branch.layers.layer_1.in_dims + if chain isa DeepONet + in_dim = chain.branch.layers.layer_1.in_dims u = rand(in_dim, number_of_parameters) v = rand(1, 10, 1) x = (u, v) phi(x, init_params) - elseif chain.name == "FourierNeuralOperator" + elseif chain isa FourierNeuralOperator in_dim = chain.layers.lifting.in_dims v = rand(in_dim, number_of_parameters, 40) phi(v, θ) @@ -332,10 +333,13 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - p, t = get_trainset(strategy, bounds, number_of_parameters, tspan) - x = (p, t) - u = phi(x, res.u) - + p, t = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) + if chain isa DeepONet + x = (p, t) + u = phi(x, res.u) + else + error("WIP") + end sol = SciMLBase.build_solution(prob, alg, x, u; k = res, dense = true, interp = PINOODEInterpolation(phi, res.u), diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index e3472fad22..31223297da 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -17,14 +17,17 @@ using NeuralOperators Dense(10 => 10, Lux.tanh_fast))) u = rand(1, 50) v = rand(1, 40, 1) - branch = deeponet.layers.branch + branch = deeponet.branch θ, st = Lux.setup(Random.default_rng(), branch) b = branch(u, θ, st)[1] - trunk = deeponet.layers.trunk + trunk = deeponet.trunk θ, st = Lux.setup(Random.default_rng(), trunk) t = trunk(v, θ, st)[1] + θ, st = Lux.setup(Random.default_rng(), deeponet) + deeponet((u,v), θ, st)[1] + bounds = [(pi, 2pi)] number_of_parameters = 50 # strategy = StochasticTraining(40) From c54d62f843f00353f3f6ba09096216f4a45e05f3 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 26 Aug 2024 19:58:58 +0400 Subject: [PATCH 097/153] update --- src/pino_ode_solve.jl | 27 +++++++++++---------------- test/PINO_ode_tests.jl | 36 ++++++++++++++++++++++++++++++++---- 2 files changed, 43 insertions(+), 20 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index a225c0de35..9d861a0b83 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -92,7 +92,7 @@ end #FourierNeuralOperator and Chain function dfdx(phi::PINOPhi{C, T}, x::Array, θ) where {C, T} ε = [zeros(eltype(x), size(x)[1] - 1)..., sqrt(eps(eltype(x)))] - (phi(t .+ ε, θ) - phi(t, θ)) ./ sqrt(eps(eltype(t))) + (phi(x .+ ε, θ) - phi(x, θ)) ./ sqrt(eps(eltype(x))) end function physics_loss( @@ -125,15 +125,14 @@ end #FourierNeuralOperator and Chain function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Array, θ) where { - C <: Union{FourierNeuralOperator, Chain}, T} + C <: Union{FourierNeuralOperator, Lux.Chain}, T} p, t = x[1:(end - 1), :, :], x[[end], :, :] f = prob.f out = phi(x, θ) - if size(out)[1] == 1 - out = dropdims(out, dims = 1) - end - - fs = f.(vec(out), vec(p), vec(t)) + # if size(out)[1] == 1 + # out = dropdims(out, dims = 1) + # end + fs = f.(out, p, t) f_vec = vec(fs) # f_vec = reduce( # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) @@ -159,7 +158,7 @@ end #FourierNeuralOperator and Chain function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { - C <: Union{FourierNeuralOperator, Chain}, T} + C <: Union{FourierNeuralOperator, Lux.Chain}, T} p, t = x[1:end-1,:,:], x[[end],:,:] t0 = fill(prob.tspan[1], size(p)) x0 = reduce(vcat, (p, t0)) #TODO one time in get_trainset @@ -172,7 +171,7 @@ function initial_condition_loss( end #TODO for input FourierNeuralOperator and Chain -function get_trainset(strategy::GridTraining, chain::Union{FourierNeuralOperator, Chain}, +function get_trainset(strategy::GridTraining, chain::Union{FourierNeuralOperator, Lux.Chain}, bounds, number_of_parameters, tspan) dt = strategy.dx p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) for b in bounds]...) @@ -333,13 +332,9 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - p, t = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) - if chain isa DeepONet - x = (p, t) - u = phi(x, res.u) - else - error("WIP") - end + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) + u = phi(x, res.u) + sol = SciMLBase.build_solution(prob, alg, x, u; k = res, dense = true, interp = PINOODEInterpolation(phi, res.u), diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 31223297da..a62c181a30 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -63,16 +63,44 @@ using NeuralOperators #TODO #ffnn #FourierNeuralOperator - fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - v = rand(2, 40,50) - θ, st = Lux.setup(Random.default_rng(), fno) - c = fno(v, θ, st)[1] + ffnn = Lux.Chain( Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 1)) θ, st = Lux.setup(Random.default_rng(), ffnn) c = ffnn(v, θ, st)[1] alg = PINOODE(ffnn, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 200) + + function get_trainset( + strategy::GridTraining, chain::Union{FourierNeuralOperator, Lux.Chain}, bounds, + number_of_parameters, tspan) + dt = strategy.dx + p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds]...) + t = collect(tspan[1]:dt:tspan[2]) + combinations = collect(Iterators.product(p, t)) + N = size(p, 1) + M = size(t, 1) + x = zeros(2, N, M) + + for i in 1:N + for j in 1:M + x[:, i, j] = [combinations[(i - 1) * M + j]...] + end + end + x + end + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) + predict_sol = sol.interp(x) + @test ground_solution≈predict_sol rtol=0.01 + + fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) + v = rand(2, 40, 50) + θ, st = Lux.setup(Random.default_rng(), fno) + c = fno(v, θ, st)[1] + + alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 2000) predict_sol = sol.interp((p, t)) From d99062db57b640901d224aa7cef573be04c50927 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 27 Aug 2024 16:53:47 +0400 Subject: [PATCH 098/153] update --- src/pino_ode_solve.jl | 128 ++++++++++++++-------------- test/PINO_ode_tests.jl | 188 +++++++++++++++++++++++++++++------------ 2 files changed, 196 insertions(+), 120 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 9d861a0b83..cd573b5c2a 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -68,7 +68,7 @@ end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) init_params = isnothing(init_params) ? θ : init_params - # init_params = chain.name == "FourierNeuralOperator" ? θ : ComponentArrays.ComponentArray(init_params) + # init_params = chain isa FourierNeuralOperator ? θ : ComponentArrays.ComponentArray(init_params) init_params = ComponentArrays.ComponentArray(init_params) PINOPhi(chain, st), init_params end @@ -76,53 +76,18 @@ end function (f::PINOPhi{C, T})( x, θ) where {C <: Lux.AbstractExplicitLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) - #ChainRulesCore.@ignore_derivatives f.st = st #TODO open issue f.st = st fail? + # ChainRulesCore.@ignore_derivatives f.st = st #TODO open issue f.st = st fail? y end -function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} - p, t = x - branch_left, branch_right = p, p - trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t - x_left = (branch_left, trunk_left) - x_right = (branch_right, trunk_right) - (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) -end - -#FourierNeuralOperator and Chain -function dfdx(phi::PINOPhi{C, T}, x::Array, θ) where {C, T} +#TODO FourierNeuralOperator and Chain +function dfdx(phi::PINOPhi{C, T}, x::Array, + θ) where {C <: Union{FourierNeuralOperator, Lux.Chain}, T} ε = [zeros(eltype(x), size(x)[1] - 1)..., sqrt(eps(eltype(x)))] (phi(x .+ ε, θ) - phi(x, θ)) ./ sqrt(eps(eltype(x))) end -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 - #TODO - if size(out)[1] == 1 - out = dropdims(out, dims = 1) - end - # out = dropdims(out, dims = 1) - fs = f.(out, p, vec(t)) - f_vec = vec(fs) - # else - # # f_vec = reduce( - # # vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - # f_vec = reduce( - # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - # end - - du = vec(dfdx(phi, x, θ)) - norm = prod(size(du)) - sum(abs2, du .- f_vec) / norm - -end - -#FourierNeuralOperator and Chain +#TODO FourierNeuralOperator and Chain function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Array, θ) where { C <: Union{FourierNeuralOperator, Lux.Chain}, T} @@ -136,32 +101,18 @@ function physics_loss( f_vec = vec(fs) # f_vec = reduce( # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm end -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)) # t[:, [1], :] - x0 = (p, t0) #TODO one time in get_trainset ? - out = phi(x0, θ) - u = vec(out) - u0 = vec(fill(prob.u0, size(out))) - norm = prod(size(u0)) - sum(abs2, u .- u0) / norm -end - -#FourierNeuralOperator and Chain +#TODO FourierNeuralOperator and Chain function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { C <: Union{FourierNeuralOperator, Lux.Chain}, T} p, t = x[1:end-1,:,:], x[[end],:,:] t0 = fill(prob.tspan[1], size(p)) - x0 = reduce(vcat, (p, t0)) #TODO one time in get_trainset + x0 = reduce(vcat, (p, t0)) out = phi(x0, θ) u = vec(out) @@ -170,30 +121,77 @@ function initial_condition_loss( sum(abs2, u .- u0) / norm end -#TODO for input FourierNeuralOperator and Chain +#TODO FourierNeuralOperator and Chain function get_trainset(strategy::GridTraining, chain::Union{FourierNeuralOperator, Lux.Chain}, bounds, number_of_parameters, tspan) - dt = strategy.dx - p = collect([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(range(start = tspan[1], length = 1/dt, stop = tspan[2])) - t = collect(tspan[1]:dt:tspan[2]) # t = reshape(t_, 1, size(t_, 1), 1) # reduce(vcat, (p, t)) # hcat((map(points -> collect(points), Iterators.product(p,t)))...) - combinations = collect(Iterators.product(p, t)) + dt = strategy.dx + p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) for b in bounds]...) + t = collect(tspan[1]:dt:tspan[2]) + combinations = (collect(Iterators.product(p, t))) N = size(p, 1) M = size(t, 1) x = zeros(2, N, M) - for i in 1:N for j in 1:M - x[:, i, j] = [combinations[(i - 1) * M + j]...] + x[:, i, j] = [combinations[i,j]...] end end x end +function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} + p, t = x + branch_left, branch_right = p, p + trunk_left, trunk_right = t .+ sqrt(eps(eltype(t))), t + x_left = (branch_left, trunk_left) + x_right = (branch_right, trunk_right) + (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) +end + +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 + #TODO + if size(out)[1] == 1 + out = dropdims(out, dims = 1) + end + # out = dropdims(out, dims = 1) + fs = f.(out, p, vec(t)) + f_vec = vec(fs) + # else + # # f_vec = reduce( + # # vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + # f_vec = reduce( + # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + # end + + du = vec(dfdx(phi, x, θ)) + norm = prod(size(du)) + sum(abs2, du .- f_vec) / norm +end + +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)) # t[:, [1], :] + x0 = (p, t0) #TODO one time in get_trainset ? + out = phi(x0, θ) + u = vec(out) + u0 = vec(fill(prob.u0, size(out))) + norm = prod(size(u0)) + sum(abs2, u .- u0) / norm +end + function get_trainset(strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) dt = strategy.dx # if size(bounds,1) == 1 @@ -283,7 +281,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, x = (u, v) phi(x, init_params) elseif chain isa FourierNeuralOperator - in_dim = chain.layers.lifting.in_dims + in_dim = chain.lifting.in_dims v = rand(in_dim, number_of_parameters, 40) phi(v, θ) else #TODO identifier for simple Chain diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index a62c181a30..be596fca8a 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -59,57 +59,6 @@ using NeuralOperators # predict_sol = dropdims(sol.interp((p, t)), dims = 1) @test ground_solution≈predict_sol rtol=0.01 - - #TODO - #ffnn - #FourierNeuralOperator - - ffnn = Lux.Chain( - Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 1)) - θ, st = Lux.setup(Random.default_rng(), ffnn) - c = ffnn(v, θ, st)[1] - - alg = PINOODE(ffnn, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 200) - - function get_trainset( - strategy::GridTraining, chain::Union{FourierNeuralOperator, Lux.Chain}, bounds, - number_of_parameters, tspan) - dt = strategy.dx - p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds]...) - t = collect(tspan[1]:dt:tspan[2]) - combinations = collect(Iterators.product(p, t)) - N = size(p, 1) - M = size(t, 1) - x = zeros(2, N, M) - - for i in 1:N - for j in 1:M - x[:, i, j] = [combinations[(i - 1) * M + j]...] - end - end - x - end - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) - predict_sol = sol.interp(x) - @test ground_solution≈predict_sol rtol=0.01 - - fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - v = rand(2, 40, 50) - θ, st = Lux.setup(Random.default_rng(), fno) - c = fno(v, θ, st)[1] - - alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) - - predict_sol = sol.interp((p, t)) - - @test ground_solution≈predict_sol rtol=0.01 - - alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) - end @testset "Example du = cos(p * t) + u" begin @@ -117,7 +66,7 @@ end tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(eq_, u0, tspan) - deeponet = LuxNeuralOperators.DeepONet( + deeponet = NeuralOperators.DeepONet( Chain( Dense(1 => 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), @@ -149,7 +98,7 @@ end u0 = 0.0f0 prob = ODEProblem(equation, u0, tspan) - deeponet = LuxNeuralOperators.DeepONet( + deeponet = NeuralOperators.DeepONet( Chain( Dense(1 => 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), @@ -214,7 +163,7 @@ end prob = ODEProblem(equation, u0, tspan) input_branch_size = 3 - deeponet = LuxNeuralOperators.DeepONet( + deeponet = NeuralOperators.DeepONet( Chain( Dense(input_branch_size => 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), @@ -258,7 +207,136 @@ end @test ground_solution_≈predict rtol=0.01 end -#vector output +#TODO FourierNeuralOperator and Chain tests +@testset "Example du = cos(p * t)" begin + equation = (u, p, t) -> cos(p * t) + tspan = (0.0f0, 1.0f0) + u0 = 1.0f0 + prob = ODEProblem(equation, u0, tspan) +## + ffnn = Lux.Chain( + Dense(2, 10, Lux.σ), + Dense(10, 10, Lux.σ), + Dense(10, 1)) + θ, st = Lux.setup(Random.default_rng(), ffnn) + v = rand(Float32, 2, 40) + c = ffnn(v, θ, st)[1] + ff = (θ) -> ffnn(v, θ, st)[1] .- 1.0f0 + using Zygote + Zygote.gradient((θ) -> sum(abs2, ff(θ)), θ) + init_params = ComponentArrays.ComponentArray(θ) + Zygote.gradient((θ) -> sum(abs2, ff(θ)), init_params) + + function total_loss(θ, _) + sum(abs2, ff(θ)) + end + init_params = ComponentArrays.ComponentArray(θ) + total_loss(init_params, 0) + + x0 = rand(Float32, 2, 50, 40) + + θ, st = Lux.setup(Random.default_rng(), chain) + init_params = ComponentArrays.ComponentArray(θ) + initial_condition_loss(phi, prob, x, θ) + initial_condition_loss(phi, prob, x, init_params) + + using Zygote + Zygote.gradient((θ) -> sum(abs2, initial_condition_loss(phi, prob, x, θ)), θ) + Zygote.gradient((θ) -> sum(abs2, initial_condition_loss(phi, prob, x, init_params)), θ) + + opt_algo = Optimization.AutoZygote() + optf = OptimizationFunction(total_loss, opt_algo) + callback = function (p, l) + @show l + false + end + maxiters = 3000 + + optprob = OptimizationProblem(optf, init_params) + res = solve(optprob, opt; callback, maxiters) +## + + bounds = [(pi, 2pi)] + number_of_parameters = 50 + # strategy = StochasticTraining(40) + strategy = GridTraining(0.025f0) + opt = OptimizationOptimisers.Adam(0.01) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) + + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + dt = 0.025f0 + + function get_trainset( + strategy::GridTraining, chain::Union{FourierNeuralOperator, Lux.Chain}, bounds, + number_of_parameters, tspan) + dt = strategy.dx + p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds]...) + t = collect(tspan[1]:dt:tspan[2]) + combinations = (collect(Iterators.product(p, t))) + N = size(p, 1) + M = size(t, 1) + x = zeros(2, N, M) + for i in 1:N + for j in 1:M + x[:, i, j] = [combinations[i, j]...] + end + end + x + end + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) + + ground_solution = ground_analytic.(u0, p, vec(t)) + # predict_sol = dropdims(sol.interp((p, t)), dims=1) + predict_sol = sol.interp((p, t)) + + @test ground_solution≈predict_sol rtol=0.01 + + p, t = get_trainset(bounds, tspan, 100, 0.01) + ground_solution = ground_analytic.(u0, p, vec(t)) + predict_sol = sol.interp((p, t)) + # predict_sol = dropdims(sol.interp((p, t)), dims = 1) + + @test ground_solution≈predict_sol rtol=0.01 + + #FourierNeuralOperator + + fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) + θ, st = Lux.setup(Random.default_rng(), fno) + v = rand(rng, Float32, 2, 40, 50) + c = fno(v, θ, st)[1] .- 1.0f0 + ff = (θ) -> abs.(fno(v, θ, st)[1] .- 1.0f0) + init_params = ComponentArrays.ComponentArray(θ) + function total_loss(θ) + sum(abs2, ff(θ)) + end + total_loss(θ) + total_loss(init_params) + + opt_algo = Optimization.AutoZygote() + optf = OptimizationFunction(total_loss, opt_algo) + callback = function (p, l) + @show l + false + end + maxiters = 1000 + + optprob = OptimizationProblem(optf, init_params) + res = solve(optprob, opt; callback, maxiters) + + alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) + + predict_sol = sol.interp((p, t)) + + @test ground_solution≈predict_sol rtol=0.01 + + alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) +end + +#TODO vector output @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> [cos(p[1] * t), sin(p[2]*t)] tspan = (0.0f0, 1.0f0) From 0bc28da57fa9c9227012f799880ab8e4811fac1a Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 29 Aug 2024 17:43:10 +0400 Subject: [PATCH 099/153] remove all unnecessary --- src/NeuralPDE.jl | 2 +- src/pino_ode_solve.jl | 125 +++++-------------------- test/PINO_ode_tests.jl | 207 +++++------------------------------------ 3 files changed, 49 insertions(+), 285 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index a0b96ceebf..8c47278b25 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -34,7 +34,7 @@ import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor using ChainRulesCore: @non_differentiable using NeuralOperators - +import NeuralOperators: DeepONet RuntimeGeneratedFunctions.init(@__MODULE__) abstract type AbstractPINN end diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index cd573b5c2a..5d542d885e 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -68,7 +68,6 @@ end function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) init_params = isnothing(init_params) ? θ : init_params - # init_params = chain isa FourierNeuralOperator ? θ : ComponentArrays.ComponentArray(init_params) init_params = ComponentArrays.ComponentArray(init_params) PINOPhi(chain, st), init_params end @@ -76,74 +75,10 @@ end function (f::PINOPhi{C, T})( x, θ) where {C <: Lux.AbstractExplicitLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) - # ChainRulesCore.@ignore_derivatives f.st = st #TODO open issue f.st = st fail? + f.st = st #? y end -#TODO FourierNeuralOperator and Chain -function dfdx(phi::PINOPhi{C, T}, x::Array, - θ) where {C <: Union{FourierNeuralOperator, Lux.Chain}, T} - ε = [zeros(eltype(x), size(x)[1] - 1)..., sqrt(eps(eltype(x)))] - (phi(x .+ ε, θ) - phi(x, θ)) ./ sqrt(eps(eltype(x))) -end - -#TODO FourierNeuralOperator and Chain -function physics_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x::Array, θ) where { - C <: Union{FourierNeuralOperator, Lux.Chain}, T} - p, t = x[1:(end - 1), :, :], x[[end], :, :] - f = prob.f - out = phi(x, θ) - # if size(out)[1] == 1 - # out = dropdims(out, dims = 1) - # end - fs = f.(out, p, t) - f_vec = vec(fs) - # f_vec = reduce( - # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - du = vec(dfdx(phi, x, θ)) - norm = prod(size(du)) - sum(abs2, du .- f_vec) / norm -end - -#TODO FourierNeuralOperator and Chain -function initial_condition_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { - C <: Union{FourierNeuralOperator, Lux.Chain}, T} - p, t = x[1:end-1,:,:], x[[end],:,:] - t0 = fill(prob.tspan[1], size(p)) - x0 = reduce(vcat, (p, t0)) - - out = phi(x0, θ) - u = vec(out) - u0 = vec(fill(prob.u0, size(out))) - norm = prod(size(u0)) - sum(abs2, u .- u0) / norm -end - -#TODO FourierNeuralOperator and Chain -function get_trainset(strategy::GridTraining, chain::Union{FourierNeuralOperator, Lux.Chain}, bounds, - number_of_parameters, tspan) - # p = vcat([collect(reshape(p_i, 1, size(p_i, 1))) for p_i in p_]...) - # t = collect(range(start = tspan[1], length = 1/dt, stop = tspan[2])) - # t = reshape(t_, 1, size(t_, 1), 1) - # reduce(vcat, (p, t)) - # hcat((map(points -> collect(points), Iterators.product(p,t)))...) - dt = strategy.dx - p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) for b in bounds]...) - t = collect(tspan[1]:dt:tspan[2]) - combinations = (collect(Iterators.product(p, t))) - N = size(p, 1) - M = size(t, 1) - x = zeros(2, N, M) - for i in 1:N - for j in 1:M - x[:, i, j] = [combinations[i,j]...] - end - end - x -end - function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} p, t = x branch_left, branch_right = p, p @@ -160,10 +95,10 @@ function physics_loss( f = prob.f out = phi(x, θ) # if size(p,1) == 1 - #TODO - if size(out)[1] == 1 - out = dropdims(out, dims = 1) - end + # #TODO + # if size(out)[1] == 1 + # out = dropdims(out, dims = 1) + # end # out = dropdims(out, dims = 1) fs = f.(out, p, vec(t)) f_vec = vec(fs) @@ -173,7 +108,6 @@ function physics_loss( # f_vec = reduce( # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) # end - du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm @@ -183,8 +117,8 @@ 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)) # t[:, [1], :] - x0 = (p, t0) #TODO one time in get_trainset ? + t0 = reshape([prob.tspan[1]], (1, 1, 1)) + x0 = (p, t0) out = phi(x0, θ) u = vec(out) u0 = vec(fill(prob.u0, size(out))) @@ -192,33 +126,28 @@ function initial_condition_loss( sum(abs2, u .- u0) / norm end +# if size(bounds,1) == 1 +# bound = bounds[1] +# p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) +# p = collect(reshape(p_, 1, size(p_,1))) +# else +# end function get_trainset(strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) dt = strategy.dx - # if size(bounds,1) == 1 - # bound = bounds[1] - # p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) - # p = collect(reshape(p_, 1, size(p_,1))) - # else - 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_]...) - # end - + 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 = reshape(t_, 1, size(t_,1), 1) (p, t) end -function get_trainset(strategy::StochasticTraining, - chain::DeepONet, bounds, number_of_parameters, tspan) - # if size(bounds,1) == 1 #TODO reduce if ? - # bound = bounds[1] - # p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] - # else - p = reduce(vcat, - [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] - for bound in bounds]) - # end +# if size(bounds,1) == 1 #TODO reduce if ? +# bound = bounds[1] +# p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] +# else +# end +function get_trainset(strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan) + p = reduce(vcat,[(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] for bound in bounds]) t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points,1) .+ tspan[1] (p, t) end @@ -263,7 +192,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, if !(chain isa Lux.AbstractExplicitLayer) error("Only Lux.AbstractExplicitLayer neural networks are supported") - if !(chain isa DeepONet || chain isa FourierNeuralOperator) + if !(chain isa DeepONet) #|| chain isa FourierNeuralOperator) error("Only DeepONet and FourierNeuralOperator neural networks are supported with PINO ODE") end end @@ -280,14 +209,6 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, v = rand(1, 10, 1) x = (u, v) phi(x, init_params) - elseif chain isa FourierNeuralOperator - in_dim = chain.lifting.in_dims - v = rand(in_dim, number_of_parameters, 40) - phi(v, θ) - else #TODO identifier for simple Chain - in_dim = chain.layers.layer_1.in_dims - v = rand(in_dim, number_of_parameters, 40) - phi(v, θ) end catch err if isa(err, DimensionMismatch) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index be596fca8a..d041a51016 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -20,45 +20,36 @@ using NeuralOperators branch = deeponet.branch θ, st = Lux.setup(Random.default_rng(), branch) b = branch(u, θ, st)[1] - trunk = deeponet.trunk θ, st = Lux.setup(Random.default_rng(), trunk) t = trunk(v, θ, st)[1] - θ, st = Lux.setup(Random.default_rng(), deeponet) - deeponet((u,v), θ, st)[1] + deeponet((u, v), θ, st)[1] bounds = [(pi, 2pi)] number_of_parameters = 50 - # strategy = StochasticTraining(40) - strategy = GridTraining(0.025f0) + strategy = StochasticTraining(40) + # strategy = GridTraining(0.025f0) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = true, maxiters = 2000) - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 - function get_trainset(bounds, tspan , number_of_parameters, dt) + function get_trainset(bounds, tspan, number_of_parameters, dt) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = collect(reshape(p_, 1, size(p_, 1))) t_ = collect(tspan[1]:dt:tspan[2]) t = collect(reshape(t_, 1, size(t_, 1), 1)) - (p,t) + (p, t) end - p,t = get_trainset(bounds, tspan, number_of_parameters, dt) - + p, t = get_trainset(bounds, tspan, number_of_parameters, dt) ground_solution = ground_analytic.(u0, p, vec(t)) - # predict_sol = dropdims(sol.interp((p, t)), dims=1) predict_sol = sol.interp((p, t)) - - @test ground_solution≈predict_sol rtol=0.01 - + @test ground_solution≈predict_sol rtol=0.05 p, t = get_trainset(bounds, tspan, 100, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) - # predict_sol = dropdims(sol.interp((p, t)), dims = 1) - - @test ground_solution≈predict_sol rtol=0.01 + @test ground_solution≈predict_sol rtol=0.05 end @testset "Example du = cos(p * t) + u" begin @@ -74,22 +65,19 @@ end bounds = [(0.1f0, 2.0f0)] number_of_parameters = 40 dt = (tspan[2] - tspan[1]) / 40 - strategy = GridTraining(dt) - + # strategy = GridTraining(dt) + strategy = StochasticTraining(50) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) sol.original.objective #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / (p^2 + 1) - - p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = collect(reshape(p_, 1, size(p_,1))) - ground_solution = ground_analytic_.(u0, p, vec(sol.t[2])) - - @test ground_solution≈sol.u rtol=0.01 + p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + ground_solution = ground_analytic_.(u0, p, vec(t)) + predict_sol = sol.interp((p, t)) + @test ground_solution≈predict_sol rtol=0.01 end @testset "Example with data du = p*t^2" begin @@ -97,62 +85,44 @@ end tspan = (0.0f0, 1.0f0) u0 = 0.0f0 prob = ODEProblem(equation, u0, tspan) - deeponet = NeuralOperators.DeepONet( Chain( Dense(1 => 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))) - bounds = [(0.0f0, 10.0f0)] number_of_parameters = 60 dt = (tspan[2] - tspan[1]) / 40 - strategy = GridTraining(dt) - + # strategy = GridTraining(dt) + strategy = StochasticTraining(60) opt = OptimizationOptimisers.Adam(0.03) - #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 - - function get_trainset(branch_size, trunk_size, bounds, tspan) - p_ = range(bounds[1][1], stop = bounds[1][2], length = branch_size) - p = reshape(p_, 1, branch_size) - t_ = collect(range(tspan[1], stop = tspan[2], length = trunk_size)) - t = reshape(t_, 1, trunk_size, 1) - (p, t) - end - function get_data() sol = ground_analytic.(u0, p, vec(t)) #x = equation.(sol, p, vec(t)) tuple_ = (p, t) sol, tuple_ end - u = rand(1, 50) v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] - - branch_size, trunk_size = 50, 40 - p, t = get_trainset(branch_size, trunk_size, bounds, tspan) data, tuple_ = get_data() - function additional_loss_(phi, θ) u = phi(tuple_, θ) norm = prod(size(u)) sum(abs2, u .- data) / norm end - alg = PINOODE( deeponet, opt, bounds, number_of_parameters; strategy = strategy, additional_loss = additional_loss_) sol = solve(prob, alg, verbose = true, maxiters = 2000) - p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = reshape(p_, 1, size(p_)[1]) - ground_solution = ground_analytic.(u0, p, vec(sol.t[2])) - @test ground_solution≈sol.u rtol=0.01 + p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + ground_solution = ground_analytic.(u0, p, vec(t)) + predict_sol = sol.interp((p, t)) + @test ground_solution≈predict_sol rtol=0.05 end #multiple parameters @@ -179,7 +149,7 @@ end strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 3000) #TODO function get_trainset(bounds, tspan, number_of_parameters, dt) p_ = [range(start = b[1], length = number_of_parameters, stop = b[2]) @@ -207,138 +177,10 @@ end @test ground_solution_≈predict rtol=0.01 end -#TODO FourierNeuralOperator and Chain tests -@testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 1.0f0) - u0 = 1.0f0 - prob = ODEProblem(equation, u0, tspan) -## - ffnn = Lux.Chain( - Dense(2, 10, Lux.σ), - Dense(10, 10, Lux.σ), - Dense(10, 1)) - θ, st = Lux.setup(Random.default_rng(), ffnn) - v = rand(Float32, 2, 40) - c = ffnn(v, θ, st)[1] - ff = (θ) -> ffnn(v, θ, st)[1] .- 1.0f0 - using Zygote - Zygote.gradient((θ) -> sum(abs2, ff(θ)), θ) - init_params = ComponentArrays.ComponentArray(θ) - Zygote.gradient((θ) -> sum(abs2, ff(θ)), init_params) - - function total_loss(θ, _) - sum(abs2, ff(θ)) - end - init_params = ComponentArrays.ComponentArray(θ) - total_loss(init_params, 0) - - x0 = rand(Float32, 2, 50, 40) - - θ, st = Lux.setup(Random.default_rng(), chain) - init_params = ComponentArrays.ComponentArray(θ) - initial_condition_loss(phi, prob, x, θ) - initial_condition_loss(phi, prob, x, init_params) - - using Zygote - Zygote.gradient((θ) -> sum(abs2, initial_condition_loss(phi, prob, x, θ)), θ) - Zygote.gradient((θ) -> sum(abs2, initial_condition_loss(phi, prob, x, init_params)), θ) - - opt_algo = Optimization.AutoZygote() - optf = OptimizationFunction(total_loss, opt_algo) - callback = function (p, l) - @show l - false - end - maxiters = 3000 - - optprob = OptimizationProblem(optf, init_params) - res = solve(optprob, opt; callback, maxiters) -## - - bounds = [(pi, 2pi)] - number_of_parameters = 50 - # strategy = StochasticTraining(40) - strategy = GridTraining(0.025f0) - opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) - - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - dt = 0.025f0 - - function get_trainset( - strategy::GridTraining, chain::Union{FourierNeuralOperator, Lux.Chain}, bounds, - number_of_parameters, tspan) - dt = strategy.dx - p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds]...) - t = collect(tspan[1]:dt:tspan[2]) - combinations = (collect(Iterators.product(p, t))) - N = size(p, 1) - M = size(t, 1) - x = zeros(2, N, M) - for i in 1:N - for j in 1:M - x[:, i, j] = [combinations[i, j]...] - end - end - x - end - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) - - ground_solution = ground_analytic.(u0, p, vec(t)) - # predict_sol = dropdims(sol.interp((p, t)), dims=1) - predict_sol = sol.interp((p, t)) - - @test ground_solution≈predict_sol rtol=0.01 - - p, t = get_trainset(bounds, tspan, 100, 0.01) - ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol.interp((p, t)) - # predict_sol = dropdims(sol.interp((p, t)), dims = 1) - - @test ground_solution≈predict_sol rtol=0.01 - - #FourierNeuralOperator - - fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - θ, st = Lux.setup(Random.default_rng(), fno) - v = rand(rng, Float32, 2, 40, 50) - c = fno(v, θ, st)[1] .- 1.0f0 - ff = (θ) -> abs.(fno(v, θ, st)[1] .- 1.0f0) - init_params = ComponentArrays.ComponentArray(θ) - function total_loss(θ) - sum(abs2, ff(θ)) - end - total_loss(θ) - total_loss(init_params) - - opt_algo = Optimization.AutoZygote() - optf = OptimizationFunction(total_loss, opt_algo) - callback = function (p, l) - @show l - false - end - maxiters = 1000 - - optprob = OptimizationProblem(optf, init_params) - res = solve(optprob, opt; callback, maxiters) - - alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) - - predict_sol = sol.interp((p, t)) - - @test ground_solution≈predict_sol rtol=0.01 - - alg = PINOODE(fno, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) -end #TODO vector output @testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> [cos(p[1] * t), sin(p[2]*t)] + equation = (u, p, t) -> [cos(p[1] * t), sin(p[2] * t)] tspan = (0.0f0, 1.0f0) u0 = [1.0f0, 0.0f0] prob = ODEProblem(equation, u0, tspan) @@ -348,9 +190,10 @@ end # Chain(Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), # Dense(10 => 10, Lux.tanh_fast)), # additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) - ffnn = Lux.Chain(Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 2)) + ffnn = Lux.Chain( + Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 2)) fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - bounds = [(pi, 2pi), (pi/2, 3pi/2)] + bounds = [(pi, 2pi), (pi / 2, 3pi / 2)] number_of_parameters = 50 strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.01) From 2b45d106599ab47705acc18d5fba73a5e1c6e13f Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 30 Aug 2024 17:26:25 +0400 Subject: [PATCH 100/153] update --- .github/workflows/CI.yml | 61 ---------------------- .github/workflows/Tests.yml | 44 ++++++++++++++++ docs/src/tutorials/pino_ode.md | 43 +++++++++------- src/pino_ode_solve.jl | 61 ++++++++-------------- test/PINO_ode_tests.jl | 93 +++++++--------------------------- 5 files changed, 110 insertions(+), 192 deletions(-) delete mode 100644 .github/workflows/CI.yml create mode 100644 .github/workflows/Tests.yml diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml deleted file mode 100644 index 0eb3f3de05..0000000000 --- a/.github/workflows/CI.yml +++ /dev/null @@ -1,61 +0,0 @@ -name: CI -on: - pull_request: - branches: - - master - paths-ignore: - - "docs/**" - push: - branches: - - master - paths-ignore: - - "docs/**" -jobs: - test: - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - group: - - QA - - ODEBPINN - - PDEBPINN - - NNPDE1 - - NNPDE2 - - AdaptiveLoss - - Logging - - Forward - - DGM - - PINOODE - - NNODE - - NeuralAdapter - - IntegroDiff - version: - - "1" - steps: - - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v2 - with: - version: ${{ matrix.version }} - - uses: actions/cache@v4 - env: - cache-name: cache-artifacts - with: - path: ~/.julia/artifacts - key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }} - restore-keys: | - ${{ runner.os }}-test-${{ env.cache-name }}- - ${{ runner.os }}-test- - ${{ runner.os }}- - - uses: julia-actions/julia-buildpkg@v1 - - uses: julia-actions/julia-runtest@v1 - env: - GROUP: ${{ matrix.group }} - - uses: julia-actions/julia-processcoverage@v1 - with: - directories: src,lib/NeuralPDELogging/src - - uses: codecov/codecov-action@v4 - with: - files: lcov.info - token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: true \ No newline at end of file diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml new file mode 100644 index 0000000000..bb6abf817c --- /dev/null +++ b/.github/workflows/Tests.yml @@ -0,0 +1,44 @@ +name: "Tests" + +on: + pull_request: + branches: + - master + - 'release-' + paths-ignore: + - 'docs/**' + push: + branches: + - master + paths-ignore: + - 'docs/**' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref_name != github.event.repository.default_branch || github.ref != 'refs/tags/v*' }} + +jobs: + tests: + name: "Tests" + strategy: + fail-fast: false + matrix: + group: + - "QA" + - "ODEBPINN" + - "PDEBPINN" + - "NNPDE1" + - "NNPDE2" + - "AdaptiveLoss" + - "Logging" + - "Forward" + - "DGM" + - "NNODE" + - "PINOODE" + - "NeuralAdapter" + - "IntegroDiff" + uses: "SciML/.github/.github/workflows/tests.yml@v1" + with: + group: "${{ matrix.group }}" + coverage-directories: "src,lib/NeuralPDELogging/src" + secrets: "inherit" \ No newline at end of file diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index b1d985b8a6..eb9c9331e3 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -7,36 +7,38 @@ This tutorial provides an example of how to use the Physics Informed Neural Oper 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 # hide +using Test using OptimizationOptimisers using Lux using Statistics, Random -using LuxNeuralOperators +using NeuralOperators using NeuralPDE +# Define the parametric ODE equation equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] -tspan = (0.0f0, 1.0f0) -u0 = 1.0f0 +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 -deeponet = LuxNeuralOperators.DeepONet( +# 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))) -u = rand(3, 50) -v = rand(1, 40, 1) -θ, st = Lux.setup(Random.default_rng(), deeponet) -c = deeponet((u, v), θ, st)[1] - -bounds = [(1.0f0, pi), (1.0f0, 2.0f0), (2.0f0, 3.0f0)] +# Define the bounds for the parameters +bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] number_of_parameter_samples = 50 -strategy = StochasticTraining(40) +# Define the training strategy +strategy = StochasticTraining(60) +# Define the optimizer opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) -sol = solve(prob, alg, verbose = false, maxiters = 3000) +# Solve the ODE problem using the PINOODE algorithm +sol = solve(prob, alg, verbose = true, maxiters = 3000) ``` 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. @@ -58,17 +60,22 @@ 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 -(p,t) = get_trainset(bounds, tspan, 50, 0.025f0) +# 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 mean error and the standard deviation of the errors +# 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) -# generate the solution with new parameters for test the model -p,t = get_trainset(bounds, tspan, 100, 0.01f0) + +p,t = get_trainset(bounds, tspan, 100, 0.01) ground_solution_ = ground_solution_f(p,t) predict = sol.interp((p,t)) @@ -84,7 +91,7 @@ plot!(ground_solution_, linetype = :contourf) ```@example pino # 'i' is the index of the parameter 'p' in the dataset -i = 5 +i = 20 # 'predict' is the predicted solution from the PINO model plot(predict[:, i], label = "Predicted") # 'ground' is the ground truth solution diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 5d542d885e..212ab5ea4d 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -57,7 +57,8 @@ function PINOODE(chain, init_params, strategy, additional_loss, kwargs) end -mutable struct PINOPhi{C, S} +#TODO imutable? +struct PINOPhi{C, S} chain::C st::S function PINOPhi(chain::Lux.AbstractExplicitLayer, st) @@ -75,7 +76,6 @@ end function (f::PINOPhi{C, T})( x, θ) where {C <: Lux.AbstractExplicitLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) - f.st = st #? y end @@ -89,25 +89,16 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} end function physics_loss( - phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { - C <: DeepONet, T} + 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 - # #TODO - # if size(out)[1] == 1 - # out = dropdims(out, dims = 1) - # end - # out = dropdims(out, dims = 1) - fs = f.(out, p, vec(t)) - f_vec = vec(fs) - # else - # # f_vec = reduce( - # # vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - # f_vec = reduce( - # vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - # end + if size(p, 1) == 1 + f_vec = vec(f.(out, p, vec(t))) + else + f_vec = reduce( + vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + end du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm @@ -126,29 +117,22 @@ function initial_condition_loss( sum(abs2, u .- u0) / norm end -# if size(bounds,1) == 1 -# bound = bounds[1] -# p_ = range(start = bound[1], length = number_of_parameters, stop = bound[2]) -# p = collect(reshape(p_, 1, size(p_,1))) -# else -# end -function get_trainset(strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) +function get_trainset( + strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) dt = strategy.dx 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_]...) + p = vcat([collect(reshape(p_i, 1, size(p_i, 1))) for p_i in p_]...) t_ = collect(tspan[1]:dt:tspan[2]) - t = reshape(t_, 1, size(t_,1), 1) + t = reshape(t_, 1, size(t_, 1), 1) (p, t) end -# if size(bounds,1) == 1 #TODO reduce if ? -# bound = bounds[1] -# p = (bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] -# else -# end -function get_trainset(strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan) - p = reduce(vcat,[(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] for bound in bounds]) - t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points,1) .+ tspan[1] +function get_trainset( + strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan) + p = reduce(vcat, + [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] + for bound in bounds]) + t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points, 1) .+ tspan[1] (p, t) end @@ -205,8 +189,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, try if chain isa DeepONet in_dim = chain.branch.layers.layer_1.in_dims - u = rand(in_dim, number_of_parameters) - v = rand(1, 10, 1) + u = rand(Float32, in_dim, number_of_parameters) + v = rand(Float32, 1, 10, 1) x = (u, v) phi(x, init_params) end @@ -219,8 +203,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end if strategy === nothing - dt = (tspan[2] - tspan[1]) / 50 - strategy = GridTraining(dt) + strategy = StochasticTraining(100) elseif !(strategy isa GridTraining || strategy isa StochasticTraining) throw(ArgumentError("Only GridTraining and StochasticTraining strategy is supported")) end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index d041a51016..0b0ccde834 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -5,6 +5,14 @@ using Statistics, Random using NeuralOperators # using NeuralPDE +function get_trainset(bounds, tspan, number_of_parameters, dt) + p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) + p = collect(reshape(p_, 1, size(p_, 1))) + t_ = collect(tspan[1]:dt:tspan[2]) + t = collect(reshape(t_, 1, size(t_, 1), 1)) + (p, t) +end + @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -29,19 +37,11 @@ using NeuralOperators bounds = [(pi, 2pi)] number_of_parameters = 50 strategy = StochasticTraining(40) - # strategy = GridTraining(0.025f0) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 - function get_trainset(bounds, tspan, number_of_parameters, dt) - p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = collect(reshape(p_, 1, size(p_, 1))) - t_ = collect(tspan[1]:dt:tspan[2]) - t = collect(reshape(t_, 1, size(t_, 1), 1)) - (p, t) - end p, t = get_trainset(bounds, tspan, number_of_parameters, dt) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @@ -65,11 +65,10 @@ end bounds = [(0.1f0, 2.0f0)] number_of_parameters = 40 dt = (tspan[2] - tspan[1]) / 40 - # strategy = GridTraining(dt) strategy = StochasticTraining(50) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 4000) sol.original.objective #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / @@ -77,7 +76,7 @@ end p, t = get_trainset(bounds, tspan, number_of_parameters, dt) ground_solution = ground_analytic_.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) - @test ground_solution≈predict_sol rtol=0.01 + @test ground_solution≈predict_sol rtol=0.5 end @testset "Example with data du = p*t^2" begin @@ -93,14 +92,12 @@ end bounds = [(0.0f0, 10.0f0)] number_of_parameters = 60 dt = (tspan[2] - tspan[1]) / 40 - # strategy = GridTraining(dt) strategy = StochasticTraining(60) opt = OptimizationOptimisers.Adam(0.03) #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 function get_data() sol = ground_analytic.(u0, p, vec(t)) - #x = equation.(sol, p, vec(t)) tuple_ = (p, t) sol, tuple_ end @@ -117,7 +114,7 @@ end alg = PINOODE( deeponet, opt, bounds, number_of_parameters; strategy = strategy, additional_loss = additional_loss_) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) p, t = get_trainset(bounds, tspan, number_of_parameters, dt) ground_solution = ground_analytic.(u0, p, vec(t)) @@ -128,8 +125,8 @@ end #multiple parameters @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] - tspan = (0.0f0, 1.0f0) - u0 = 1.0f0 + tspan = (0.0, 1.0) + u0 = 1.0 prob = ODEProblem(equation, u0, tspan) input_branch_size = 3 @@ -144,12 +141,12 @@ end θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] - bounds = [(1.0f0, pi), (1.0f0, 2.0f0), (2.0f0, 3.0f0)] + bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] number_of_parameters = 50 - strategy = StochasticTraining(40) + strategy = StochasticTraining(70) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 3000) #TODO + sol = solve(prob, alg, verbose = false, maxiters = 3000) function get_trainset(bounds, tspan, number_of_parameters, dt) p_ = [range(start = b[1], length = number_of_parameters, stop = b[2]) @@ -169,62 +166,10 @@ end (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 - - p, t = get_trainset(bounds, tspan, 100, 0.01f0) - ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 -end - - -#TODO vector output -@testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> [cos(p[1] * t), sin(p[2] * t)] - tspan = (0.0f0, 1.0f0) - u0 = [1.0f0, 0.0f0] - prob = ODEProblem(equation, u0, tspan) - # deeponet = NeuralOperators.DeepONet( - # Chain( - # Dense(1 => 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)), - # additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) - ffnn = Lux.Chain( - Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 2)) - fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) - bounds = [(pi, 2pi), (pi / 2, 3pi / 2)] - number_of_parameters = 50 - strategy = StochasticTraining(40) - opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) - - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - dt = 0.025f0 - 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 - p, t = get_trainset(bounds, tspan, number_of_parameters, dt) - - ground_solution = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] - 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 - - (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) - ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 + @test ground_solution_≈predict rtol=0.05 p, t = get_trainset(bounds, tspan, 100, 0.01f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 + @test ground_solution_≈predict rtol=0.05 end From d5b9d6dbad7e1ae24c4948a85a55656c6bc7c6d5 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 30 Aug 2024 17:26:48 +0400 Subject: [PATCH 101/153] update --- test/PINO_ode_tests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 0b0ccde834..5914987eb1 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -3,7 +3,7 @@ using OptimizationOptimisers using Lux using Statistics, Random using NeuralOperators -# using NeuralPDE +using NeuralPDE function get_trainset(bounds, tspan, number_of_parameters, dt) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) From c4092b5facbee29e57bd42c413541ab667b9b32f Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 30 Aug 2024 18:02:46 +0400 Subject: [PATCH 102/153] imutable --- src/pino_ode_solve.jl | 1 - test/PINO_ode_tests.jl | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 212ab5ea4d..30bf199767 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -57,7 +57,6 @@ function PINOODE(chain, init_params, strategy, additional_loss, kwargs) end -#TODO imutable? struct PINOPhi{C, S} chain::C st::S diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 5914987eb1..57110a1993 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -96,6 +96,7 @@ end opt = OptimizationOptimisers.Adam(0.03) #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 + function get_data() sol = ground_analytic.(u0, p, vec(t)) tuple_ = (p, t) @@ -105,6 +106,7 @@ end v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] + p, t = get_trainset(bounds, tspan, number_of_parameters, dt) data, tuple_ = get_data() function additional_loss_(phi, θ) u = phi(tuple_, θ) From 3ae25ce39ae1252a84c23c894d396852411c4702 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 12:28:31 +0400 Subject: [PATCH 103/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index eb9c9331e3..f7b2ab635e 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -97,4 +97,3 @@ plot(predict[:, i], label = "Predicted") # 'ground' is the ground truth solution plot!(ground_solution_[:, i], label = "Ground truth") ``` - From 908ca8a48e2dac6d8773c5e5a2d7188f0e67f729 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 12:28:44 +0400 Subject: [PATCH 104/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index f7b2ab635e..52e315271c 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -65,7 +65,7 @@ end # compute the ground truth solution ground_solution_ = ground_solution_f(p,t) # predict the solution with the PINO model -predict = sol.interp((p,t)) +predict = sol.interp((p, t)) # calculate the errors between the ground truth solution and the predicted solution errors = ground_solution_ - predict From 13f955bf81b56479a76366409bc076d7ae99bb4e Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 12:28:52 +0400 Subject: [PATCH 105/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 52e315271c..e426ffa4e9 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -63,7 +63,7 @@ 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) +ground_solution_ = ground_solution_f(p, t) # predict the solution with the PINO model predict = sol.interp((p, t)) From 3813b1b6cd0604867436de03bd1884a915551e46 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 12:29:02 +0400 Subject: [PATCH 106/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index e426ffa4e9..42a30af0e9 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -51,7 +51,7 @@ function get_trainset(bounds, tspan , number_of_parameters, dt) 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) + (p, t) end # Compute the ground truth solution for each parameter From 5be5db1c789a000cba7eba4e8064b0518dc9ebc9 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 12:29:13 +0400 Subject: [PATCH 107/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 42a30af0e9..c21e4e1983 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -46,7 +46,7 @@ Now let's compare the prediction from the learned operator with the ground truth ```@example pino using Plots -function get_trainset(bounds, tspan , number_of_parameters, dt) +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]) From d0665ecb3ebb616a60a0d05172587fe006b52a9d Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 18:51:44 +0400 Subject: [PATCH 108/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index c21e4e1983..7cdc1d9a01 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -61,7 +61,7 @@ function ground_solution_f(p,t) end # generate the solution with new parameters for test the model -(p,t) = get_trainset(bounds, tspan, 50, 0.025) +(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 From 5295ceb15859f70a54305e4c26a0042f9ff3ce4c Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 18:51:54 +0400 Subject: [PATCH 109/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 7cdc1d9a01..876fc16fe7 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -55,9 +55,10 @@ function get_trainset(bounds, tspan, number_of_parameters, dt) 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)]) +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 From cc67c313b2a10d0460bc7a508ee057f2211c33e6 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 18:52:02 +0400 Subject: [PATCH 110/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 876fc16fe7..b5c330ff27 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -75,10 +75,9 @@ 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)) +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) From cf5221ca6a594b6c3d2ca3a06b536709c0cd0c6f Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 18:52:10 +0400 Subject: [PATCH 111/153] Update test/runtests.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/runtests.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index ace868dc78..57e1bbe518 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -49,7 +49,8 @@ end end if GROUP == "All" || GROUP == "PINOODE" - @time @safetestset "pino ode" begin include("PINO_ode_tests.jl") + @time @safetestset "pino ode" begin + include("PINO_ode_tests.jl") end end From a3dbfc63ce56cb4cafb44efa8e3008ab8c8d3630 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Mon, 2 Sep 2024 18:52:15 +0400 Subject: [PATCH 112/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index b5c330ff27..c7a82634fd 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -48,7 +48,7 @@ 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_]...) + 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) From 0cdd8dc2bd3924ad24d39ef9fc15577ab280eff5 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Thu, 5 Sep 2024 18:19:43 +0400 Subject: [PATCH 113/153] Update src/pino_ode_solve.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/pino_ode_solve.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 30bf199767..3057527349 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -20,7 +20,7 @@ in which will be train the prediction of parametric ODE. ## Keyword Arguments -* `init_params`: The initial parameter of the neural network. By default, this is `nothing`, +* `init_params`: The initial parameters of the neural network. By default, this is `nothing`, which thus uses the random initialization provided by the neural network library. * `strategy`: The strategy for training the neural network. * `additional_loss`: additional function to the main one. For example, add training on data. From bc32c454c8592ad6236db968d00b40f98528574f Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Thu, 5 Sep 2024 18:19:51 +0400 Subject: [PATCH 114/153] Update src/pino_ode_solve.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/pino_ode_solve.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 3057527349..d4f390d488 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -14,8 +14,7 @@ neural operator, which is used as a solver for a parametrized `ODEProblem`. * `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`. `Flux.Chain` will be converted to `Lux` using `adapt(FromFluxAdaptor(false, false), chain)` * `opt`: The optimizer to train the neural network. -* `bounds`: A dictionary containing the bounds for the parameters of the neural network -in which will be train the prediction of parametric ODE. +* `bounds`: A dictionary containing the bounds for the parameters of the parametric ODE. * `number_of_parameters`: The number of points of train set in parameters boundaries. ## Keyword Arguments From 6c49dd866fc22c79a05328e0a3dcb7a678b35eb3 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Thu, 5 Sep 2024 18:20:59 +0400 Subject: [PATCH 115/153] Update src/pino_ode_solve.jl Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- src/pino_ode_solve.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index d4f390d488..e199498e31 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -22,7 +22,7 @@ neural operator, which is used as a solver for a parametrized `ODEProblem`. * `init_params`: The initial parameters of the neural network. By default, this is `nothing`, which thus uses the random initialization provided by the neural network library. * `strategy`: The strategy for training the neural network. -* `additional_loss`: additional function to the main one. For example, add training on data. +* `additional_loss`: additional loss function added to the default one. For example, add training on data. * `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. ## References From ff04565e9e4dc7e098ab487a98fa21bfd659a163 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Thu, 5 Sep 2024 18:21:23 +0400 Subject: [PATCH 116/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index c7a82634fd..f20090d3b2 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -7,7 +7,7 @@ This tutorial provides an example of how to use the Physics Informed Neural Oper 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 Test # hide using OptimizationOptimisers using Lux using Statistics, Random From 23668405b14a9727d29b115e1f1b82c198c758c7 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Thu, 5 Sep 2024 18:21:34 +0400 Subject: [PATCH 117/153] Update docs/src/tutorials/pino_ode.md Co-authored-by: Sathvik Bhagavan <35105271+sathvikbhagavan@users.noreply.github.com> --- docs/src/tutorials/pino_ode.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index f20090d3b2..98c0fba269 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -38,7 +38,7 @@ strategy = StochasticTraining(60) 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 = true, maxiters = 3000) +sol = solve(prob, alg, verbose = false, maxiters = 3000) ``` 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. From b1209aba04844a3b46ad999ec3aa5ca112d59fb9 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 19 Sep 2024 15:17:19 +0400 Subject: [PATCH 118/153] lux v1.0 --- Project.toml | 32 +++++++++++++++++++------------- src/pino_ode_solve.jl | 14 +++++++------- 2 files changed, 26 insertions(+), 20 deletions(-) diff --git a/Project.toml b/Project.toml index 201fd3dd9a..a6f64865ae 100644 --- a/Project.toml +++ b/Project.toml @@ -43,45 +43,48 @@ Adapt = "4" AdvancedHMC = "0.6.1" Aqua = "0.8" ArrayInterface = "7.9" -CUDA = "5.3" +CUDA = "5.3.2" ChainRulesCore = "1.24" -ComponentArrays = "0.15.14" +ComponentArrays = "0.15.16" Cubature = "1.5" DiffEqNoiseProcess = "5.20" Distributions = "0.25.107" DocStringExtensions = "0.9.3" DomainSets = "0.6, 0.7" -Flux = "0.14.11" +Flux = "0.14.17" ForwardDiff = "0.10.36" -Functors = "0.4.10" +Functors = "0.4.12" Integrals = "4.4" LineSearches = "7.2" -LinearAlgebra = "1" +LinearAlgebra = "1.10" LogDensityProblems = "2" -Lux = "0.5.58" +Lux = "1.0" LuxCUDA = "0.3.2" +LuxCore = "0.1.24" +LuxLib = "0.3.48" MCMCChains = "6" MethodOfLines = "0.11" ModelingToolkit = "9.9" MonteCarloMeasurements = "1.1" Optim = "1.7.8" -Optimization = "3.24" +Optimization = "3.25" OptimizationOptimJL = "0.2.1" OptimizationOptimisers = "0.2.1" -OrdinaryDiffEq = "6.74" -Pkg = "1" +OrdinaryDiffEq = "6.87" +Pkg = "1.10" +Preferences = "1.4.3" QuasiMonteCarlo = "0.3.2" Random = "1" Reexport = "1.2" RuntimeGeneratedFunctions = "0.5.12" SafeTestsets = "0.1" -SciMLBase = "2.28" +SciMLBase = "2.34" Statistics = "1.10" SymbolicUtils = "1.5, 2, 3" Symbolics = "5.27.1, 6" -Test = "1" +Test = "1.10" UnPack = "1" -Zygote = "0.6.69" +Zygote = "0.6.70" julia = "1.10" [extras] @@ -90,12 +93,15 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" +LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" +LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4" OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" 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", "CUDA", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "Preferences", "SafeTestsets", "Test"] \ No newline at end of file diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index e199498e31..03a6203151 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -11,7 +11,7 @@ neural operator, which is used as a solver for a parametrized `ODEProblem`. ## Positional Arguments -* `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`. +* `chain`: A neural network architecture, defined as a `Lux.AbstractLuxLayer` or `Flux.Chain`. `Flux.Chain` will be converted to `Lux` using `adapt(FromFluxAdaptor(false, false), chain)` * `opt`: The optimizer to train the neural network. * `bounds`: A dictionary containing the bounds for the parameters of the parametric ODE. @@ -51,7 +51,7 @@ function PINOODE(chain, strategy = nothing, additional_loss = nothing, kwargs...) - !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain)) + !(chain isa Lux.AbstractLuxLayer) && (chain = Lux.transform(chain)) PINOODE(chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss, kwargs) end @@ -59,12 +59,12 @@ end struct PINOPhi{C, S} chain::C st::S - function PINOPhi(chain::Lux.AbstractExplicitLayer, st) + function PINOPhi(chain::Lux.AbstractLuxLayer, st) new{typeof(chain), typeof(st)}(chain, st) end end -function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) +function generate_pino_phi_θ(chain::Lux.AbstractLuxLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) init_params = isnothing(init_params) ? θ : init_params init_params = ComponentArrays.ComponentArray(init_params) @@ -72,7 +72,7 @@ function generate_pino_phi_θ(chain::Lux.AbstractExplicitLayer, init_params) end function (f::PINOPhi{C, T})( - x, θ) where {C <: Lux.AbstractExplicitLayer, T} + x, θ) where {C <: Lux.AbstractLuxLayer, T} y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) y end @@ -171,8 +171,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, @unpack tspan, u0, f = prob @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg - if !(chain isa Lux.AbstractExplicitLayer) - error("Only Lux.AbstractExplicitLayer neural networks are supported") + if !(chain isa Lux.AbstractLuxLayer) + error("Only Lux.AbstractLuxLayer neural networks are supported") if !(chain isa DeepONet) #|| chain isa FourierNeuralOperator) error("Only DeepONet and FourierNeuralOperator neural networks are supported with PINO ODE") From 86677d0699f2bd60b9e56867db0b8d819863b597 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 24 Sep 2024 14:11:27 +0400 Subject: [PATCH 119/153] Update Project.toml --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 5601c2cab3..78f7cef959 100644 --- a/Project.toml +++ b/Project.toml @@ -58,7 +58,7 @@ Integrals = "4.4" LineSearches = "7.2" LinearAlgebra = "1.10" LogDensityProblems = "2" -Lux = "1.0" +Lux = "0.5.58" LuxCUDA = "0.3.2" LuxCore = "0.1.24" LuxLib = "0.3.48" @@ -103,4 +103,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Aqua", "CUDA", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "Preferences", "SafeTestsets", "Test"] \ No newline at end of file +test = ["Aqua", "CUDA", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "Preferences", "SafeTestsets", "Test"] From 6813c5d78f221fdac41d924b837ff5ad7473eb49 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 24 Sep 2024 17:08:20 +0400 Subject: [PATCH 120/153] lux 1 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 78f7cef959..ae4db5e3d6 100644 --- a/Project.toml +++ b/Project.toml @@ -58,7 +58,7 @@ Integrals = "4.4" LineSearches = "7.2" LinearAlgebra = "1.10" LogDensityProblems = "2" -Lux = "0.5.58" +Lux = "1" LuxCUDA = "0.3.2" LuxCore = "0.1.24" LuxLib = "0.3.48" From 786035f1324ea6738533ce92d39c979b7ecb9a58 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 24 Sep 2024 19:46:43 +0400 Subject: [PATCH 121/153] support Chain --- src/pino_ode_solve.jl | 78 +++++++++++++++++++++++++++++++++++++++--- test/PINO_ode_tests.jl | 66 ++++++++++++++++++++++++++++++----- 2 files changed, 130 insertions(+), 14 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 03a6203151..8d002e365a 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -71,8 +71,11 @@ function generate_pino_phi_θ(chain::Lux.AbstractLuxLayer, init_params) PINOPhi(chain, st), init_params end -function (f::PINOPhi{C, T})( - x, θ) where {C <: Lux.AbstractLuxLayer, T} +function (f::PINOPhi{C, T})(x, θ) where {C <: Lux.AbstractLuxLayer, T} + # θ_ = ComponentArrays.getdata(θ) + # eltypeθ, typeθ = eltype(θ_), parameterless_type(θ_) + # t_ = convert.(eltypeθ, adapt(typeθ, t')) + # y, st = f.chain(t_, θ, f.st) y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) y end @@ -86,6 +89,13 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end +#TODO chain +function dfdx(phi::PINOPhi{C, T}, x::Array, + θ) where {C <: Lux.Chain, T} + ε = [zeros(eltype(x), size(x)[1] - 1)..., sqrt(eps(eltype(x)))] + (phi(x .+ ε, θ) - phi(x, θ)) ./ sqrt(eps(eltype(x))) +end + function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where {C <: DeepONet, T} p, t = x @@ -102,6 +112,23 @@ function physics_loss( sum(abs2, du .- f_vec) / norm end +function physics_loss( + phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { + C <: Lux.Chain, T} + p, t = x + x_ = reduce(vcat, (p, t)) + f = prob.f + if size(p, 1) == 1 + f_vec = f.(out, p, t) + else + #TODO + # f_vec = reduce( vcat, [[f(out[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(out)) + sum(abs2, du .- f_vec) / norm +end + function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x, θ) where { C <: DeepONet, T} @@ -115,6 +142,19 @@ function initial_condition_loss( sum(abs2, u .- u0) / norm end +function initial_condition_loss( + phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { + C <: Lux.Chain, T} + p, t = x + t0 = fill(prob.tspan[1], size(p)) + x0 = reduce(vcat, (p, t0)) + out = phi(x0, θ) + u = vec(out) + u0 = vec(fill(prob.u0, size(out))) + norm = prod(size(u0)) + sum(abs2, u .- u0) / norm +end + function get_trainset( strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) dt = strategy.dx @@ -134,6 +174,25 @@ function get_trainset( (p, t) end +function get_trainset(strategy::GridTraining, chain::Lux.Chain, bounds, + number_of_parameters, tspan) + dt = strategy.dx + p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds]...) + t = collect(tspan[1]:dt:tspan[2]) + combinations = (collect(Iterators.product(p, t))) + N = size(p, 1) + M = size(t, 1) + x = zeros(2, N, M) + for i in 1:N + for j in 1:M + x[:, i, j] = [combinations[i, j]...] + end + end + p, t = x[1:(end - 1), :, :], x[[end], :, :] + (p, t) +end + function generate_loss( strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) @@ -174,8 +233,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, if !(chain isa Lux.AbstractLuxLayer) error("Only Lux.AbstractLuxLayer neural networks are supported") - if !(chain isa DeepONet) #|| chain isa FourierNeuralOperator) - error("Only DeepONet and FourierNeuralOperator neural networks are supported with PINO ODE") + if !(chain isa DeepONet) || !(chain isa Chain) + error("Only DeepONet and Chain neural networks are supported with PINO ODE") end end @@ -192,6 +251,11 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, x = (u, v) phi(x, init_params) end + if chain isa Chain + in_dim = chain.layers.layer_1.in_dims + x = rand(Float32, in_dim, number_of_parameters) + phi(x, init_params) + end catch err if isa(err, DimensionMismatch) throw(DimensionMismatch("Dimensions of input data and chain should match")) @@ -233,7 +297,11 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, res = solve(optprob, opt; callback, maxiters, alg.kwargs...) x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) - u = phi(x, res.u) + if chain isa DeepONet + u = phi(x, res.u) + elseif chain isa Chain + u = phi(reduce(vcat, x), res.u) + end sol = SciMLBase.build_solution(prob, alg, x, u; k = res, dense = true, diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 57110a1993..28dd0a476f 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -5,7 +5,7 @@ using Statistics, Random using NeuralOperators using NeuralPDE -function get_trainset(bounds, tspan, number_of_parameters, dt) +function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, dt) p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) p = collect(reshape(p_, 1, size(p_, 1))) t_ = collect(tspan[1]:dt:tspan[2]) @@ -13,6 +13,54 @@ function get_trainset(bounds, tspan, number_of_parameters, dt) (p, t) end +function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) + dt = strategy.dx + p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds]...) + t = collect(tspan[1]:dt:tspan[2]) + combinations = (collect(Iterators.product(p, t))) + N = size(p, 1) + M = size(t, 1) + x = zeros(2, N, M) + for i in 1:N + for j in 1:M + x[:, i, j] = [combinations[i, j]...] + end + end + p, t = x[1:(end - 1), :, :], x[[end], :, :] + (p, t) +end + +#Test with Chain +@testset "Example du = cos(p * t)" begin + equation = (u, p, t) -> cos(p * t) + tspan = (0.0f0, 1.0f0) + u0 = 1.0f0 + prob = ODEProblem(equation, u0, tspan) + chain = Chain(Dense(2 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) + x = rand(1, 50) + θ, st = Lux.setup(Random.default_rng(), chain) + b = chain(x, θ, st)[1] + + bounds = [(pi, 2pi)] + number_of_parameters = 50 + strategy = GridTraining(0.1f0) + opt = OptimizationOptimisers.Adam(0.01) + alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = false, maxiters = 5000) + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + dt = 0.025f0 + p, t = get_trainset(chain, bounds, number_of_parameters, tspan, dt) + ground_solution = ground_analytic.(u0, p, t) + predict_sol = sol.interp(reduce(vcat, (p, t))) + @test ground_solution≈predict_sol rtol=0.07 + p, t = get_trainset(chain, bounds, 100, tspan, 0.01) + ground_solution = ground_analytic.(u0, p, t) + predict_sol = sol.interp(reduce(vcat, (p, t))) + @test ground_solution≈predict_sol rtol=0.07 +end + +#Test with DeepONet @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -42,11 +90,11 @@ end sol = solve(prob, alg, verbose = false, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 - p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 - p, t = get_trainset(bounds, tspan, 100, 0.01) + p, t = get_trainset(deeponet, bounds, tspan, 100, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 @@ -73,7 +121,7 @@ end #if u0 == 1 ground_analytic_(u0, p, t) = (p * sin(p * t) - cos(p * t) + (p^2 + 2) * exp(t)) / (p^2 + 1) - p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic_.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.5 @@ -106,7 +154,7 @@ end v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] - p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) data, tuple_ = get_data() function additional_loss_(phi, θ) u = phi(tuple_, θ) @@ -118,7 +166,7 @@ end additional_loss = additional_loss_) sol = solve(prob, alg, verbose = false, maxiters = 2000) - p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 @@ -150,7 +198,7 @@ end alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 3000) - function get_trainset(bounds, tspan, number_of_parameters, dt) + function get_trainset(deeponet, bounds, number_of_parameters, tspan, 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_]...) @@ -165,12 +213,12 @@ end [[ground_solution(u0, p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) end - (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) + (p, t) = get_trainset(deeponet, bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(bounds, tspan, 100, 0.01f0) + p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 From 5428b01ba4f887baaace0478d6dc21c6ef8aac9a Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Tue, 24 Sep 2024 20:10:24 +0400 Subject: [PATCH 122/153] Update test/PINO_ode_tests.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/PINO_ode_tests.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 28dd0a476f..39ca6fb1ce 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -37,7 +37,8 @@ end tspan = (0.0f0, 1.0f0) u0 = 1.0f0 prob = ODEProblem(equation, u0, tspan) - chain = Chain(Dense(2 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) + chain = Chain( + Dense(2 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) x = rand(1, 50) θ, st = Lux.setup(Random.default_rng(), chain) b = chain(x, θ, st)[1] From 60c39951728e2535de3fac2332799418ee3148cf Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 25 Sep 2024 21:29:02 +0400 Subject: [PATCH 123/153] vector input Chain --- src/pino_ode_solve.jl | 30 ++++++----- test/PINO_ode_tests.jl | 120 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 132 insertions(+), 18 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 8d002e365a..970b94ee63 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -72,7 +72,7 @@ function generate_pino_phi_θ(chain::Lux.AbstractLuxLayer, init_params) end function (f::PINOPhi{C, T})(x, θ) where {C <: Lux.AbstractLuxLayer, T} - # θ_ = ComponentArrays.getdata(θ) + # θ_ = ComponentArrays.getdata(θ) #TODO eltype # eltypeθ, typeθ = eltype(θ_), parameterless_type(θ_) # t_ = convert.(eltypeθ, adapt(typeθ, t')) # y, st = f.chain(t_, θ, f.st) @@ -89,7 +89,6 @@ function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} (phi(x_left, θ) .- phi(x_right, θ)) ./ sqrt(eps(eltype(t))) end -#TODO chain function dfdx(phi::PINOPhi{C, T}, x::Array, θ) where {C <: Lux.Chain, T} ε = [zeros(eltype(x), size(x)[1] - 1)..., sqrt(eps(eltype(x)))] @@ -118,11 +117,13 @@ function physics_loss( p, t = x x_ = reduce(vcat, (p, t)) f = prob.f + out = phi(x_, θ) if size(p, 1) == 1 f_vec = f.(out, p, t) else - #TODO - # f_vec = reduce( vcat, [[f(out[i], p[i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) + f_vec = reduce(hcat, + [[f(out[1, i, j], p[:, i, j], t[1, i, j]) for j in axes(t, 3)] + for i in axes(p, 2)]) end du = dfdx(phi, x_, θ) norm = prod(size(out)) @@ -146,7 +147,7 @@ function initial_condition_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { C <: Lux.Chain, T} p, t = x - t0 = fill(prob.tspan[1], size(p)) + t0 = fill(prob.tspan[1], size(t)) x0 = reduce(vcat, (p, t0)) out = phi(x0, θ) u = vec(out) @@ -177,20 +178,21 @@ end function get_trainset(strategy::GridTraining, chain::Lux.Chain, bounds, number_of_parameters, tspan) dt = strategy.dx - p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds]...) + p = [range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds] t = collect(tspan[1]:dt:tspan[2]) - combinations = (collect(Iterators.product(p, t))) - N = size(p, 1) - M = size(t, 1) - x = zeros(2, N, M) + combinations = collect(Iterators.product(p..., t)) + N = number_of_parameters + M = length(t) + num_dims = ndims(combinations) + x = zeros(num_dims, N, M) for i in 1:N for j in 1:M - x[:, i, j] = [combinations[i, j]...] + x[:, i, j] = [combinations[(i - 1) * M + j]...] end end - p, t = x[1:(end - 1), :, :], x[[end], :, :] - (p, t) + p_, t_ = x[1:(end - 1), :, :], x[[end], :, :] + (p_, t_) end function generate_loss( diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 39ca6fb1ce..0381fb0a57 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -46,6 +46,7 @@ end bounds = [(pi, 2pi)] number_of_parameters = 50 strategy = GridTraining(0.1f0) + # strategy = StochasticTraining(70) #TODO opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 5000) @@ -173,7 +174,66 @@ end @test ground_solution≈predict_sol rtol=0.05 end -#multiple parameters +#multiple parameters chain +@testset "Example du = cos(p * t)" begin + 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) + + input_branch_size = 3 + chain = Chain( + Dense(input_branch_size+1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) + + x = rand(4, 50) + θ, st = Lux.setup(Random.default_rng(), chain) + c = chain(x, θ, st)[1] + + bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] + number_of_parameters = 50 + # strategy = StochasticTraining(70) + strategy = GridTraining(0.1f0) + opt = OptimizationOptimisers.Adam(0.03) + alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 5000) + + function get_trainset(bounds, number_of_parameters, tspan, dt) + p = [range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds] + t = collect(tspan[1]:dt:tspan[2]) + combinations = collect(Iterators.product(p..., t)) + N = number_of_parameters + M = length(t) + num_dims = ndims(combinations) + x = zeros(num_dims, N, M) + for i in 1:N + for j in 1:M + x[:, i, j] = [combinations[(i - 1) * M + j]...] + end + end + p_, t_ = x[1:(end - 1), :, :], x[[end], :, :] + (p_, t_) + end + + 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 + + (p, t) = get_trainset(bounds, 50, tspan, 0.025f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.05 + + p, t = get_trainset(bounds, 100, tspan, 0.01f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.05 +end + + +#multiple parameters DeepOnet @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0, 1.0) @@ -199,7 +259,7 @@ end alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 3000) - function get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) + function get_trainset(bounds, number_of_parameters, tspan, 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_]...) @@ -214,13 +274,65 @@ end [[ground_solution(u0, p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) end - (p, t) = get_trainset(deeponet, bounds, 50, tspan, 0.025f0) + (p, t) = get_trainset(bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01f0) + p, t = get_trainset(bounds, 100, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 end + + +#TODO vector output TODO +@testset "Example du = cos(p * t)" begin + equation = (u, p, t) -> [cos(p[1] * t), sin(p[2] * t)] + tspan = (0.0f0, 1.0f0) + u0 = [1.0f0, 0.0f0] + prob = ODEProblem(equation, u0, tspan) + # deeponet = NeuralOperators.DeepONet( + # Chain( + # Dense(1 => 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)), + # additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) + ffnn = Lux.Chain( + Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 2)) + fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) + bounds = [(pi, 2pi), (pi / 2, 3pi / 2)] + number_of_parameters = 50 + strategy = StochasticTraining(40) + opt = OptimizationOptimisers.Adam(0.01) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) + + ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + dt = 0.025f0 + 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 + p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + + ground_solution = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] + 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 + + (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 + + p, t = get_trainset(bounds, tspan, 100, 0.01f0) + ground_solution_ = ground_solution_f(p, t) + predict = sol.interp((p, t)) + @test ground_solution_≈predict rtol=0.01 +end From 90dfdb0756b4609f43f264146f0a1b6ef3e3fb9c Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Fri, 27 Sep 2024 14:37:36 +0400 Subject: [PATCH 124/153] Update test/PINO_ode_tests.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/PINO_ode_tests.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 0381fb0a57..ab54d2c5c0 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -285,7 +285,6 @@ end @test ground_solution_≈predict rtol=0.05 end - #TODO vector output TODO @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> [cos(p[1] * t), sin(p[2] * t)] From 2d7be46737e00d3fa79ef1366d8789bd3de2446a Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Fri, 27 Sep 2024 14:37:47 +0400 Subject: [PATCH 125/153] Update test/PINO_ode_tests.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/PINO_ode_tests.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ab54d2c5c0..ecca3d12f1 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -232,7 +232,6 @@ end @test ground_solution_≈predict rtol=0.05 end - #multiple parameters DeepOnet @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] From e342a7dd28f676f3ea3a22864927cce710854413 Mon Sep 17 00:00:00 2001 From: Kirill Zubov Date: Fri, 27 Sep 2024 14:37:57 +0400 Subject: [PATCH 126/153] Update test/PINO_ode_tests.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/PINO_ode_tests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ecca3d12f1..e3ff15c2d6 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -183,7 +183,7 @@ end input_branch_size = 3 chain = Chain( - Dense(input_branch_size+1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) + Dense(input_branch_size + 1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) x = rand(4, 50) θ, st = Lux.setup(Random.default_rng(), chain) From 000d8b5e962bbcb26224c1dc540b1f8993ea107e Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Sat, 28 Sep 2024 20:38:31 +0400 Subject: [PATCH 127/153] input vector chain --- src/pino_ode_solve.jl | 81 ++++++++++--------- test/PINO_ode_tests.jl | 172 ++++++++++++++++++++--------------------- 2 files changed, 125 insertions(+), 128 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 970b94ee63..e34f1f8a4e 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -71,12 +71,17 @@ function generate_pino_phi_θ(chain::Lux.AbstractLuxLayer, init_params) PINOPhi(chain, st), init_params end -function (f::PINOPhi{C, T})(x, θ) where {C <: Lux.AbstractLuxLayer, T} - # θ_ = ComponentArrays.getdata(θ) #TODO eltype - # eltypeθ, typeθ = eltype(θ_), parameterless_type(θ_) - # t_ = convert.(eltypeθ, adapt(typeθ, t')) - # y, st = f.chain(t_, θ, f.st) - y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) +function (f::PINOPhi{C, T})(x::Array, θ) where {C <: Lux.Chain, T} + eltypeθ, typeθ = eltype(θ), parameterless_type(ComponentArrays.getdata(θ)) + x = convert.(eltypeθ, adapt(typeθ, x)) + y, st = f.chain(x, θ, f.st) + y +end + +function (f::PINOPhi{C, T})(x::Tuple, θ) where {C <: DeepONet, T} + eltypeθ, typeθ = eltype(θ), parameterless_type(ComponentArrays.getdata(θ)) + x = (convert.(eltypeθ, adapt(typeθ, x[1])),convert.(eltypeθ, adapt(typeθ, x[2]))) + y, st = f.chain(x, θ, f.st) y end @@ -115,17 +120,17 @@ function physics_loss( phi::PINOPhi{C, T}, prob::ODEProblem, x::Tuple, θ) where { C <: Lux.Chain, T} p, t = x - x_ = reduce(vcat, (p, t)) + x_ = reduce(vcat, x) f = prob.f out = phi(x_, θ) if size(p, 1) == 1 - f_vec = f.(out, p, t) + f_vec = vec(f.(out, p, t)) else - f_vec = reduce(hcat, + f_vec = reduce(vcat, [[f(out[1, i, j], p[:, i, j], t[1, i, j]) for j in axes(t, 3)] for i in axes(p, 2)]) end - du = dfdx(phi, x_, θ) + du = vec(dfdx(phi, x_, θ)) norm = prod(size(out)) sum(abs2, du .- f_vec) / norm end @@ -157,56 +162,54 @@ function initial_condition_loss( end function get_trainset( - strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) + strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan, eltypeθ) dt = strategy.dx 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 = reshape(t_, 1, size(t_, 1), 1) + p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) (p, t) end function get_trainset( - strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan) + strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan, eltypeθ) p = reduce(vcat, [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] for bound in bounds]) t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points, 1) .+ tspan[1] + p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) (p, t) end -function get_trainset(strategy::GridTraining, chain::Lux.Chain, bounds, - number_of_parameters, tspan) +function get_trainset( + strategy::GridTraining, chain::Lux.Chain, bounds, number_of_parameters, tspan, eltypeθ) dt = strategy.dx - p = [range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds] - t = collect(tspan[1]:dt:tspan[2]) - combinations = collect(Iterators.product(p..., t)) - N = number_of_parameters - M = length(t) - num_dims = ndims(combinations) - x = zeros(num_dims, N, M) - for i in 1:N - for j in 1:M - x[:, i, j] = [combinations[(i - 1) * M + j]...] - end - end - p_, t_ = x[1:(end - 1), :, :], x[[end], :, :] - (p_, t_) + tspan_ = tspan[1]:dt:tspan[2] + pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds] + x_ = hcat(vec(map( + points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) + # x = reshape(x_, size(bounds, 1) + 1, size.(pspan, 1)..., size(tspan_, 1)) + x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) + p, t = x[1:(end - 1), :, :], x[[end], :, :] + p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) + (p, t) end function generate_loss( - strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) + strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan, eltypeθ) function loss(θ, _) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end +# Zygote.gradient(θ -> initial_condition_loss(phi, prob, x, θ), θ) function generate_loss( - strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) + strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) function loss(θ, _) - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan, eltypeθ) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end @@ -241,6 +244,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end phi, init_params = generate_pino_phi_θ(chain, init_params) + eltypeθ = eltype(init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) @@ -248,14 +252,14 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, try if chain isa DeepONet in_dim = chain.branch.layers.layer_1.in_dims - u = rand(Float32, in_dim, number_of_parameters) - v = rand(Float32, 1, 10, 1) + u = rand(eltypeθ, in_dim, number_of_parameters) + v = rand(eltypeθ, 1, 10, 1) x = (u, v) phi(x, init_params) end if chain isa Chain in_dim = chain.layers.layer_1.in_dims - x = rand(Float32, in_dim, number_of_parameters) + x = rand(eltypeθ, in_dim, number_of_parameters) phi(x, init_params) end catch err @@ -272,7 +276,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, throw(ArgumentError("Only GridTraining and StochasticTraining strategy is supported")) end - inner_f = generate_loss(strategy, prob, phi, bounds, number_of_parameters, tspan) + inner_f = generate_loss( + strategy, prob, phi, bounds, number_of_parameters, tspan, eltypeθ) function total_loss(θ, _) L2_loss = inner_f(θ, nothing) @@ -298,7 +303,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan) + x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan, eltypeθ) if chain isa DeepONet u = phi(x, res.u) elseif chain isa Chain diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index e3ff15c2d6..49389e404b 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -6,27 +6,21 @@ using NeuralOperators using NeuralPDE function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, dt) - p_ = range(start = bounds[1][1], length = number_of_parameters, stop = bounds[1][2]) - p = collect(reshape(p_, 1, size(p_, 1))) + 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)) + t = reshape(t_, 1, size(t_, 1), 1) (p, t) end function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) - dt = strategy.dx - p = collect([range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds]...) - t = collect(tspan[1]:dt:tspan[2]) - combinations = (collect(Iterators.product(p, t))) - N = size(p, 1) - M = size(t, 1) - x = zeros(2, N, M) - for i in 1:N - for j in 1:M - x[:, i, j] = [combinations[i, j]...] - end - end + tspan_ = tspan[1]:dt:tspan[2] + pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds] + x_ = hcat(vec(map( + points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) + # x = reshape(x_, size(bounds, 1) + 1, size.(pspan, 1)..., size(tspan_, 1)) + x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) p, t = x[1:(end - 1), :, :], x[[end], :, :] (p, t) end @@ -39,16 +33,18 @@ end prob = ODEProblem(equation, u0, tspan) chain = Chain( Dense(2 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) - x = rand(1, 50) + x = rand(1, 50, 10) θ, st = Lux.setup(Random.default_rng(), chain) b = chain(x, θ, st)[1] bounds = [(pi, 2pi)] number_of_parameters = 50 strategy = GridTraining(0.1f0) - # strategy = StochasticTraining(70) #TODO + # strategy = StochasticTraining(70) #TODO chain +StochasticTraining opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) + alg = PINOODE( + chain, opt, bounds, number_of_parameters; strategy = strategy, init_params = θ |> + f64) sol = solve(prob, alg, verbose = false, maxiters = 5000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 @@ -60,6 +56,7 @@ end ground_solution = ground_analytic.(u0, p, t) predict_sol = sol.interp(reduce(vcat, (p, t))) @test ground_solution≈predict_sol rtol=0.07 + @test eltype(sol.k) == eltype(predict_sol) end #Test with DeepONet @@ -73,8 +70,8 @@ end Dense(1 => 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))) - u = rand(1, 50) - v = rand(1, 40, 1) + u = rand(Float32, 1, 50) + v = rand(Float32, 1, 40, 1) branch = deeponet.branch θ, st = Lux.setup(Random.default_rng(), branch) b = branch(u, θ, st)[1] @@ -88,18 +85,20 @@ end number_of_parameters = 50 strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; + strategy = strategy, init_params = θ |> f64) + sol = solve(prob, alg, verbose = true, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 - p, t = get_trainset(deeponet, bounds, tspan, 100, 0.01) + p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 + @test eltype(sol.k) == eltype(predict_sol) end @testset "Example du = cos(p * t) + u" begin @@ -126,7 +125,8 @@ end p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic_.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) - @test ground_solution≈predict_sol rtol=0.5 + @test ground_solution≈predict_sol rtol=0.05 + @test eltype(sol.k) == eltype(predict_sol) end @testset "Example with data du = p*t^2" begin @@ -172,66 +172,67 @@ end ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 + @test eltype(sol.k) == eltype(predict_sol) end #multiple parameters chain @testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] + 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) - input_branch_size = 3 + input_branch_size = 2 chain = Chain( - Dense(input_branch_size + 1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) + Dense(input_branch_size + 1 => 10, Lux.tanh_fast), + Dense(10 => 10, Lux.tanh_fast), + Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) - x = rand(4, 50) + x = rand(Float32, 3, 1000, 10) θ, st = Lux.setup(Random.default_rng(), chain) c = chain(x, θ, st)[1] - bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] - number_of_parameters = 50 - # strategy = StochasticTraining(70) + bounds = [(1.0, pi), (1.0, 2.0)]#, (2.0, 3.0)] + number_of_parameters = 10 strategy = GridTraining(0.1f0) + # strategy = StochasticTraining(20) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 5000) + sol = solve(prob, alg, verbose = true, maxiters = 3000) - function get_trainset(bounds, number_of_parameters, tspan, dt) - p = [range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds] - t = collect(tspan[1]:dt:tspan[2]) - combinations = collect(Iterators.product(p..., t)) - N = number_of_parameters - M = length(t) - num_dims = ndims(combinations) - x = zeros(num_dims, N, M) - for i in 1:N - for j in 1:M - x[:, i, j] = [combinations[(i - 1) * M + j]...] - end - end - p_, t_ = x[1:(end - 1), :, :], x[[end], :, :] - (p_, t_) - end + ground_solution = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) #+ p[3] * t - 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)]) + [[ground_solution(u0, p[:, i, j], t[1, i, j]) for j in axes(t, 3)] + for i in axes(p, 2)])' end - - (p, t) = get_trainset(bounds, 50, tspan, 0.025f0) + (p, t) = get_trainset(chain, bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.05 + predict = sol.interp(reduce(vcat, (p, t)))[1, :, :] + @test ground_solution_≈predict rtol=0.4 #TODO rtol=0.05 - p, t = get_trainset(bounds, 100, tspan, 0.01f0) + p, t = get_trainset(chain, bounds, 100, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.05 + predict_sol = sol.interp(reduce(vcat, (p, t)))[1, :, :] + @test ground_solution_≈predict_sol rtol=0.05 + @test eltype(sol.k) == eltype(predict_sol) end +function plot_() + # Animate + anim = @animate for (i) in 1:100 + plot(ground_solution_[:, i], label = "Ground") + plot!(predict[:, i], label = "Predicted") + end + gif(anim, "pino.gif", fps = 10) +end + +plot_() + +plot(predict[:, :], linetype = :contourf) +plot(ground_solution_, linetype = :contourf) + #multiple parameters DeepOnet @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] @@ -253,35 +254,27 @@ end bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] number_of_parameters = 50 - strategy = StochasticTraining(70) + strategy = StochasticTraining(20) + # strategy = GridTraining(0.1f0) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 3000) - - function get_trainset(bounds, number_of_parameters, tspan, 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 - 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 - (p, t) = get_trainset(bounds, 50, tspan, 0.025f0) + (p, t) = get_trainset(deeponet, bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(bounds, 100, tspan, 0.01f0) + p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 + @test eltype(sol.k.u) == eltype(predict_sol) end #TODO vector output TODO @@ -290,15 +283,14 @@ end tspan = (0.0f0, 1.0f0) u0 = [1.0f0, 0.0f0] prob = ODEProblem(equation, u0, tspan) - # deeponet = NeuralOperators.DeepONet( - # Chain( - # Dense(1 => 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)), - # additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) - ffnn = Lux.Chain( + deeponet = NeuralOperators.DeepONet( + Chain( + Dense(1 => 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)), + additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) + chain = Lux.Chain( Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 2)) - fno = FourierNeuralOperator(gelu; chs = (2, 64, 64, 128, 1), modes = (16,)) bounds = [(pi, 2pi), (pi / 2, 3pi / 2)] number_of_parameters = 50 strategy = StochasticTraining(40) @@ -308,15 +300,15 @@ end ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 - 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 - p, t = get_trainset(bounds, tspan, number_of_parameters, dt) + # 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 + p, t = get_trainset(chain, bounds, tspan, number_of_parameters, dt) ground_solution = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] function ground_solution_f(p, t) @@ -324,7 +316,7 @@ end [[ground_solution(u0, p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) end - (p, t) = get_trainset(bounds, tspan, 50, 0.025f0) + (p, t) = get_trainset(chain, bounds, tspan, 50, 0.025f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.01 From 5e9a0255bbc6d8fd916beb64385cc7350256c4e7 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 1 Oct 2024 14:39:10 +0400 Subject: [PATCH 128/153] support chain with StochasticTraining --- src/pino_ode_solve.jl | 37 +++++++++++++++-------- test/PINO_ode_tests.jl | 66 +++++++++++++----------------------------- 2 files changed, 45 insertions(+), 58 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index e34f1f8a4e..aace54fa4a 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -80,7 +80,7 @@ end function (f::PINOPhi{C, T})(x::Tuple, θ) where {C <: DeepONet, T} eltypeθ, typeθ = eltype(θ), parameterless_type(ComponentArrays.getdata(θ)) - x = (convert.(eltypeθ, adapt(typeθ, x[1])),convert.(eltypeθ, adapt(typeθ, x[2]))) + x = (convert.(eltypeθ, adapt(typeθ, x[1])), convert.(eltypeθ, adapt(typeθ, x[2]))) y, st = f.chain(x, θ, f.st) y end @@ -172,16 +172,6 @@ function get_trainset( (p, t) end -function get_trainset( - strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan, eltypeθ) - p = reduce(vcat, - [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] - for bound in bounds]) - t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points, 1) .+ tspan[1] - p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) - (p, t) -end - function get_trainset( strategy::GridTraining, chain::Lux.Chain, bounds, number_of_parameters, tspan, eltypeθ) dt = strategy.dx @@ -197,6 +187,30 @@ function get_trainset( (p, t) end +# function get_trainset( +# strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan, eltypeθ) +# p = reduce(vcat, +# [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] +# for bound in bounds]) +# t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points, 1) .+ tspan[1] +# p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) +# (p, t) +# end + + +function get_trainset( + strategy::StochasticTraining, chain::Union{DeepONet, Lux.Chain}, + bounds, number_of_parameters, tspan, eltypeθ) + number_of_parameters != strategy.points && + throw(error("number_of_parameters should be the same strategy.points for StochasticTraining")) + p = reduce(vcat, + [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] + for bound in bounds]) + t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points, 1) .+ tspan[1] + p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) + (p, t) +end + function generate_loss( strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan, eltypeθ) @@ -204,7 +218,6 @@ function generate_loss( initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end -# Zygote.gradient(θ -> initial_condition_loss(phi, prob, x, θ), θ) function generate_loss( strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 49389e404b..1adadf8596 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -19,13 +19,12 @@ function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) for b in bounds] x_ = hcat(vec(map( points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) - # x = reshape(x_, size(bounds, 1) + 1, size.(pspan, 1)..., size(tspan_, 1)) x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) p, t = x[1:(end - 1), :, :], x[[end], :, :] (p, t) end -#Test with Chain +#Test with Chain with Float64 accuracy @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -38,9 +37,8 @@ end b = chain(x, θ, st)[1] bounds = [(pi, 2pi)] - number_of_parameters = 50 - strategy = GridTraining(0.1f0) - # strategy = StochasticTraining(70) #TODO chain +StochasticTraining + number_of_parameters = 300 + strategy = StochasticTraining(300) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE( chain, opt, bounds, number_of_parameters; strategy = strategy, init_params = θ |> @@ -51,12 +49,12 @@ end p, t = get_trainset(chain, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, t) predict_sol = sol.interp(reduce(vcat, (p, t))) + @test eltype(sol.k) == eltype(predict_sol) @test ground_solution≈predict_sol rtol=0.07 p, t = get_trainset(chain, bounds, 100, tspan, 0.01) ground_solution = ground_analytic.(u0, p, t) predict_sol = sol.interp(reduce(vcat, (p, t))) @test ground_solution≈predict_sol rtol=0.07 - @test eltype(sol.k) == eltype(predict_sol) end #Test with DeepONet @@ -87,7 +85,7 @@ end opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy, init_params = θ |> f64) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) @@ -114,7 +112,7 @@ end bounds = [(0.1f0, 2.0f0)] number_of_parameters = 40 dt = (tspan[2] - tspan[1]) / 40 - strategy = StochasticTraining(50) + strategy = GridTraining(0.1f0) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 4000) @@ -177,30 +175,29 @@ end #multiple parameters chain @testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> p[1] * cos(p[2] * t) #+ p[3] + 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) - input_branch_size = 2 + input_branch_size = 3 chain = Chain( Dense(input_branch_size + 1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) - x = rand(Float32, 3, 1000, 10) + x = rand(Float32, 4, 1000, 10) θ, st = Lux.setup(Random.default_rng(), chain) c = chain(x, θ, st)[1] - bounds = [(1.0, pi), (1.0, 2.0)]#, (2.0, 3.0)] - number_of_parameters = 10 - strategy = GridTraining(0.1f0) - # strategy = StochasticTraining(20) - opt = OptimizationOptimisers.Adam(0.03) + bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] + number_of_parameters = 200 + strategy = StochasticTraining(200) + opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 4000) - ground_solution = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) #+ p[3] * t + 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, @@ -210,29 +207,15 @@ end (p, t) = get_trainset(chain, bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp(reduce(vcat, (p, t)))[1, :, :] - @test ground_solution_≈predict rtol=0.4 #TODO rtol=0.05 + @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(chain, bounds, 100, tspan, 0.01f0) + p, t = get_trainset(chain, bounds, 60, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) predict_sol = sol.interp(reduce(vcat, (p, t)))[1, :, :] @test ground_solution_≈predict_sol rtol=0.05 @test eltype(sol.k) == eltype(predict_sol) end -function plot_() - # Animate - anim = @animate for (i) in 1:100 - plot(ground_solution_[:, i], label = "Ground") - plot!(predict[:, i], label = "Predicted") - end - gif(anim, "pino.gif", fps = 10) -end - -plot_() - -plot(predict[:, :], linetype = :contourf) -plot(ground_solution_, linetype = :contourf) - #multiple parameters DeepOnet @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] @@ -255,10 +238,9 @@ plot(ground_solution_, linetype = :contourf) bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] number_of_parameters = 50 strategy = StochasticTraining(20) - # strategy = GridTraining(0.1f0) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 4000) 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, @@ -274,7 +256,7 @@ plot(ground_solution_, linetype = :contourf) ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 - @test eltype(sol.k.u) == eltype(predict_sol) + @test eltype(sol.k.u) == eltype(predict) end #TODO vector output TODO @@ -296,18 +278,10 @@ end strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 2000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 - # 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 p, t = get_trainset(chain, bounds, tspan, number_of_parameters, dt) ground_solution = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] From e0cb52870fccfa3827c1f6e275dbbc0b112e0b96 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 1 Oct 2024 15:16:17 +0400 Subject: [PATCH 129/153] NeuralOperators v0.5.0 in project --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index ae4db5e3d6..2ab964df05 100644 --- a/Project.toml +++ b/Project.toml @@ -66,6 +66,7 @@ MCMCChains = "6" MethodOfLines = "0.11" ModelingToolkit = "9.9" MonteCarloMeasurements = "1.1" +NeuralOperators = "0.5" Optim = "1.7.8" Optimization = "3.24, 4" OptimizationOptimJL = "0.2.1" From f1b3a36672eea52b58c581cf53068a555c57b378 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 2 Oct 2024 19:51:52 +0400 Subject: [PATCH 130/153] wip output vector --- src/pino_ode_solve.jl | 42 +++++++++++++------------------ test/PINO_ode_tests.jl | 56 ++++++++++++++++++++++-------------------- 2 files changed, 46 insertions(+), 52 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index aace54fa4a..c45009b83b 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -105,12 +105,15 @@ function physics_loss( p, t = x f = prob.f out = phi(x, θ) - if size(p, 1) == 1 - f_vec = vec(f.(out, p, vec(t))) - else - f_vec = reduce( - vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - end + f_vec = reduce(vcat, + [reduce(vcat, [f.(out[j, i], p[:, i], t[j]) for j in axes(t, 2)]) + for i in axes(p, 2)]) + # if size(p, 1) == 1 + # f_vec = vec(f.(out, p, vec(t))) + # else + # end + # f_vec = reduce( + # vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) du = vec(dfdx(phi, x, θ)) norm = prod(size(du)) sum(abs2, du .- f_vec) / norm @@ -123,14 +126,14 @@ function physics_loss( x_ = reduce(vcat, x) f = prob.f out = phi(x_, θ) - if size(p, 1) == 1 - f_vec = vec(f.(out, p, t)) + if size(p, 1) == 1 && size(out, 1) == 1 + f_vec = vec(f.(out, p, t)) else f_vec = reduce(vcat, - [[f(out[1, i, j], p[:, i, j], t[1, i, j]) for j in axes(t, 3)] + [reduce(vcat, [f(out[:, i, j], p[1, i, j], t[1, i, j]) for j in axes(t, 3)]) for i in axes(p, 2)]) end - du = vec(dfdx(phi, x_, θ)) + du = vec((dfdx(phi, x_, θ))) norm = prod(size(out)) sum(abs2, du .- f_vec) / norm end @@ -143,7 +146,7 @@ function initial_condition_loss( x0 = (p, t0) out = phi(x0, θ) u = vec(out) - u0 = vec(fill(prob.u0, size(out))) + u0 = vec(reduce(vcat, [fill(u0, size(t)) for u0 in prob.u0])) norm = prod(size(u0)) sum(abs2, u .- u0) / norm end @@ -156,7 +159,8 @@ function initial_condition_loss( x0 = reduce(vcat, (p, t0)) out = phi(x0, θ) u = vec(out) - u0 = vec(fill(prob.u0, size(out))) + # u0 = vec(fill(prob.u0, size(out))) + u0 = vec(reduce(vcat, [fill(u0, size(t)) for u0 in prob.u0])) norm = prod(size(u0)) sum(abs2, u .- u0) / norm end @@ -180,28 +184,16 @@ function get_trainset( for b in bounds] x_ = hcat(vec(map( points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) - # x = reshape(x_, size(bounds, 1) + 1, size.(pspan, 1)..., size(tspan_, 1)) x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) p, t = x[1:(end - 1), :, :], x[[end], :, :] p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) (p, t) end -# function get_trainset( -# strategy::StochasticTraining, chain::DeepONet, bounds, number_of_parameters, tspan, eltypeθ) -# p = reduce(vcat, -# [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] -# for bound in bounds]) -# t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points, 1) .+ tspan[1] -# p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) -# (p, t) -# end - - function get_trainset( strategy::StochasticTraining, chain::Union{DeepONet, Lux.Chain}, bounds, number_of_parameters, tspan, eltypeθ) - number_of_parameters != strategy.points && + (number_of_parameters != strategy.points && chain isa Lux.Chain) && throw(error("number_of_parameters should be the same strategy.points for StochasticTraining")) p = reduce(vcat, [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 1adadf8596..ae5113685d 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -24,7 +24,7 @@ function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) (p, t) end -#Test with Chain with Float64 accuracy +#Test Chain with Float64 accuracy @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -32,13 +32,14 @@ end prob = ODEProblem(equation, u0, tspan) chain = Chain( Dense(2 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) - x = rand(1, 50, 10) + x = rand(2, 50, 10) θ, st = Lux.setup(Random.default_rng(), chain) b = chain(x, θ, st)[1] bounds = [(pi, 2pi)] number_of_parameters = 300 strategy = StochasticTraining(300) + # strategy = GridTraining(0.1f0) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE( chain, opt, bounds, number_of_parameters; strategy = strategy, init_params = θ |> @@ -50,14 +51,14 @@ end ground_solution = ground_analytic.(u0, p, t) predict_sol = sol.interp(reduce(vcat, (p, t))) @test eltype(sol.k) == eltype(predict_sol) - @test ground_solution≈predict_sol rtol=0.07 + @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.interp(reduce(vcat, (p, t))) - @test ground_solution≈predict_sol rtol=0.07 + @test ground_solution≈predict_sol rtol=0.05 end -#Test with DeepONet +#Test DeepONet with Float64 accuracy @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -85,7 +86,7 @@ end opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy, init_params = θ |> f64) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + sol = solve(prob, alg, verbose = false, maxiters = 3000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) @@ -191,11 +192,11 @@ end c = chain(x, θ, st)[1] bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] - number_of_parameters = 200 + number_of_parameters = 50 strategy = StochasticTraining(200) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 4000) + sol = solve(prob, alg, verbose = true, maxiters = 4000) ground_solution = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) + p[3] * t @@ -240,7 +241,7 @@ end strategy = StochasticTraining(20) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 4000) + sol = solve(prob, alg, verbose = true, maxiters = 4000) 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, @@ -259,32 +260,33 @@ end @test eltype(sol.k.u) == eltype(predict) end -#TODO vector output TODO -@testset "Example du = cos(p * t)" begin - equation = (u, p, t) -> [cos(p[1] * t), sin(p[2] * t)] +#TODO vector output +@testset "Example du = [cos(p * t), sin(p * t)]" begin + equation = (u, p, t) -> [cos(p * t), sin(p * t)] + # equation = (u, p, t) -> [cos(p * t) + u[1], sin(p * t) - u[2]] + # equation = (u, p, t) -> [cos(p[1] * t) + u[1], sin(p[2] * t) - u[2]] + # du1 = cos(p * t) + # du2 = sin(p * t) + tspan = (0.0f0, 1.0f0) u0 = [1.0f0, 0.0f0] prob = ODEProblem(equation, u0, tspan) - deeponet = NeuralOperators.DeepONet( - Chain( - Dense(1 => 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)), - additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2))) - chain = Lux.Chain( - Dense(2, 32, Lux.tanh_fast), Dense(32, 32, Lux.tanh_fast), Dense(32, 2)) - bounds = [(pi, 2pi), (pi / 2, 3pi / 2)] + chain = Chain( + Dense(2 => 10, Lux.tanh_fast), + Dense(10 => 10, Lux.tanh_fast), + Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2)) + bounds = [(pi, 2pi)] number_of_parameters = 50 - strategy = StochasticTraining(40) + # strategy = GridTraining(0.1f0) + strategy = StochasticTraining(50) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) + sol = solve(prob, alg, verbose = true, maxiters = 2000) - ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 - p, t = get_trainset(chain, bounds, tspan, number_of_parameters, dt) + (p, t) = get_trainset(deeponet, bounds, 50, tspan, dt) - ground_solution = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi] + ground_solution = (u0, p, t) -> [sin(p * t) / p, -cos(p * t) / p] 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)]) From ca24323cd140d3abca3a9697ea980d7c163a0ca8 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 3 Oct 2024 19:19:39 +0400 Subject: [PATCH 131/153] output vector --- src/pino_ode_solve.jl | 10 +++---- test/PINO_ode_tests.jl | 60 +++++++++++++++++++++++++++--------------- 2 files changed, 44 insertions(+), 26 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index c45009b83b..30ff45de8e 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -129,11 +129,11 @@ function physics_loss( if size(p, 1) == 1 && size(out, 1) == 1 f_vec = vec(f.(out, p, t)) else - f_vec = reduce(vcat, + f_vec = reduce(hcat, [reduce(vcat, [f(out[:, i, j], p[1, i, j], t[1, i, j]) for j in axes(t, 3)]) for i in axes(p, 2)]) end - du = vec((dfdx(phi, x_, θ))) + du = (dfdx(phi, x_, θ)) norm = prod(size(out)) sum(abs2, du .- f_vec) / norm end @@ -157,10 +157,10 @@ function initial_condition_loss( p, t = x t0 = fill(prob.tspan[1], size(t)) x0 = reduce(vcat, (p, t0)) - out = phi(x0, θ) - u = vec(out) + u = phi(x0, θ) + # u = vec(out) # u0 = vec(fill(prob.u0, size(out))) - u0 = vec(reduce(vcat, [fill(u0, size(t)) for u0 in prob.u0])) + u0 = (reduce(vcat, [fill(u0, size(t)) for u0 in prob.u0])) norm = prod(size(u0)) sum(abs2, u .- u0) / norm end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ae5113685d..28b74d338c 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -24,6 +24,7 @@ function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) (p, t) end + #Test Chain with Float64 accuracy @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) @@ -265,40 +266,57 @@ end equation = (u, p, t) -> [cos(p * t), sin(p * t)] # equation = (u, p, t) -> [cos(p * t) + u[1], sin(p * t) - u[2]] # equation = (u, p, t) -> [cos(p[1] * t) + u[1], sin(p[2] * t) - u[2]] - # du1 = cos(p * t) - # du2 = sin(p * t) tspan = (0.0f0, 1.0f0) - u0 = [1.0f0, 0.0f0] + u0 = [1.0f0, 1.0f0] prob = ODEProblem(equation, u0, tspan) + input_branch_size = 1 chain = Chain( - Dense(2 => 10, Lux.tanh_fast), + Dense(input_branch_size+1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2)) + + # deeponet = NeuralOperators.DeepONet( + # Chain( + # Dense(input_branch_size => 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))) + bounds = [(pi, 2pi)] - number_of_parameters = 50 - # strategy = GridTraining(0.1f0) - strategy = StochasticTraining(50) + number_of_parameters = 300 + strategy = StochasticTraining(300) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 2000) - - dt = 0.025f0 - (p, t) = get_trainset(deeponet, bounds, 50, tspan, dt) + sol = solve(prob, alg, verbose = true, maxiters = 6000) - ground_solution = (u0, p, t) -> [sin(p * t) / p, -cos(p * t) / p] + ground_solution = (u0, p, t) -> [1 + sin(p * t) / p, 1/p - cos(p * t) / p] 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 + ans_1 = reduce(hcat, + [reduce(vcat, + [ground_solution(u0, p[1, i, 1], t[1, 1, j])[1] for i in axes(p, 2)]) + for j in axes(t, 3)]) + ans_2 = reduce(hcat, + [reduce(vcat, + [ground_solution(u0, p[1, i, 1], t[1, 1, j])[2] for i in axes(p, 2)]) + for j in axes(t, 3)]) - (p, t) = get_trainset(chain, bounds, tspan, 50, 0.025f0) + ans_1 = reshape(ans_1, 1, size(ans_1)...) + ans_2 = reshape(ans_2, 1, size(ans_2)...) + vcat(ans_1, ans_2) + end + p, t = get_trainset(chain, bounds, 50, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 + predict = sol.interp(reduce(vcat, (p, t))) + @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 + @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 + @test ground_solution_≈predict rtol=0.05 + @test eltype(sol.k) == eltype(predict) - p, t = get_trainset(bounds, tspan, 100, 0.01f0) + p, t = get_trainset(chain, bounds, 300, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) - @test ground_solution_≈predict rtol=0.01 + predict = sol.interp(reduce(vcat, (p, t))) + @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 + @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 + @test ground_solution_≈predict rtol=0.3 + @test eltype(sol.k) == eltype(predict) end From f74cf63792396a8b4953ec8110c91faa008c203e Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 3 Oct 2024 19:42:59 +0400 Subject: [PATCH 132/153] update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 2ab964df05..bedb3a6918 100644 --- a/Project.toml +++ b/Project.toml @@ -66,7 +66,7 @@ MCMCChains = "6" MethodOfLines = "0.11" ModelingToolkit = "9.9" MonteCarloMeasurements = "1.1" -NeuralOperators = "0.5" +NeuralOperators = "0.5.0" Optim = "1.7.8" Optimization = "3.24, 4" OptimizationOptimJL = "0.2.1" From 885a72ec454f0b5d0aecec9da4adadc5db192600 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 4 Oct 2024 18:54:31 +0400 Subject: [PATCH 133/153] update --- src/pino_ode_solve.jl | 41 ++++++++++++++++++++++------------------- test/PINO_ode_tests.jl | 32 ++++++++++---------------------- 2 files changed, 32 insertions(+), 41 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 30ff45de8e..b10a3a880a 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -105,16 +105,16 @@ function physics_loss( p, t = x f = prob.f out = phi(x, θ) - f_vec = reduce(vcat, - [reduce(vcat, [f.(out[j, i], p[:, i], t[j]) for j in axes(t, 2)]) - for i in axes(p, 2)]) - # if size(p, 1) == 1 - # f_vec = vec(f.(out, p, vec(t))) - # else - # end - # f_vec = reduce( - # vcat, [[f(out[i], p[:, i], t[j]) for j in axes(t, 2)] for i in axes(p, 2)]) - du = vec(dfdx(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 @@ -127,13 +127,17 @@ function physics_loss( f = prob.f out = phi(x_, θ) if size(p, 1) == 1 && size(out, 1) == 1 - f_vec = vec(f.(out, p, t)) - else + f_vec = f.(out, p, t) + elseif size(p, 1) > 1 + f_vec = reduce(hcat, + [reduce(vcat, [f(out[1, i, j], p[:, i, j], t[1, i, j]) for j in axes(t, 3)]) + for i in axes(p, 2)]) + elseif size(out, 1) > 1 f_vec = reduce(hcat, [reduce(vcat, [f(out[:, i, j], p[1, i, j], t[1, i, j]) for j in axes(t, 3)]) for i in axes(p, 2)]) end - du = (dfdx(phi, x_, θ)) + du = dfdx(phi, x_, θ) norm = prod(size(out)) sum(abs2, du .- f_vec) / norm end @@ -144,9 +148,9 @@ function initial_condition_loss( p, t = x t0 = reshape([prob.tspan[1]], (1, 1, 1)) x0 = (p, t0) - out = phi(x0, θ) - u = vec(out) - u0 = vec(reduce(vcat, [fill(u0, size(t)) for u0 in prob.u0])) + 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 @@ -158,9 +162,8 @@ function initial_condition_loss( t0 = fill(prob.tspan[1], size(t)) x0 = reduce(vcat, (p, t0)) u = phi(x0, θ) - # u = vec(out) - # u0 = vec(fill(prob.u0, size(out))) - u0 = (reduce(vcat, [fill(u0, size(t)) for u0 in prob.u0])) + u0 = size(prob.u0, 1) == 1 ? fill(prob.u0, size(t)) : + reduce(vcat, [fill(u0, size(t)) for u0 in prob.u0]) norm = prod(size(u0)) sum(abs2, u .- u0) / norm end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 28b74d338c..00a2074b5b 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -24,12 +24,11 @@ function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) (p, t) end - #Test Chain with Float64 accuracy @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) - u0 = 1.0f0 + u0 = 1.0 prob = ODEProblem(equation, u0, tspan) chain = Chain( Dense(2 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) @@ -40,7 +39,6 @@ end bounds = [(pi, 2pi)] number_of_parameters = 300 strategy = StochasticTraining(300) - # strategy = GridTraining(0.1f0) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE( chain, opt, bounds, number_of_parameters; strategy = strategy, init_params = θ |> @@ -63,7 +61,7 @@ end @testset "Example du = cos(p * t)" begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) - u0 = 1.0f0 + u0 = 1.0 prob = ODEProblem(equation, u0, tspan) deeponet = NeuralOperators.DeepONet( Chain( @@ -193,11 +191,11 @@ end c = chain(x, θ, st)[1] bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] - number_of_parameters = 50 + number_of_parameters = 200 strategy = StochasticTraining(200) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 4000) + sol = solve(prob, alg, verbose = false, maxiters = 4000) ground_solution = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) + p[3] * t @@ -242,7 +240,7 @@ end strategy = StochasticTraining(20) opt = OptimizationOptimisers.Adam(0.03) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 4000) + sol = solve(prob, alg, verbose = false, maxiters = 4000) 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, @@ -261,35 +259,25 @@ end @test eltype(sol.k.u) == eltype(predict) end -#TODO vector output +#vector output @testset "Example du = [cos(p * t), sin(p * t)]" begin equation = (u, p, t) -> [cos(p * t), sin(p * t)] - # equation = (u, p, t) -> [cos(p * t) + u[1], sin(p * t) - u[2]] - # equation = (u, p, t) -> [cos(p[1] * t) + u[1], sin(p[2] * t) - u[2]] - tspan = (0.0f0, 1.0f0) - u0 = [1.0f0, 1.0f0] + u0 = [1.0f0, 0.0f0] prob = ODEProblem(equation, u0, tspan) input_branch_size = 1 chain = Chain( - Dense(input_branch_size+1 => 10, Lux.tanh_fast), + Dense(input_branch_size + 1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2)) - - # deeponet = NeuralOperators.DeepONet( - # Chain( - # Dense(input_branch_size => 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))) - bounds = [(pi, 2pi)] number_of_parameters = 300 strategy = StochasticTraining(300) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 6000) + sol = solve(prob, alg, verbose = false, maxiters = 6000) - ground_solution = (u0, p, t) -> [1 + sin(p * t) / p, 1/p - cos(p * t) / p] + ground_solution = (u0, p, t) -> [1 + sin(p * t) / p, 1 / p - cos(p * t) / p] function ground_solution_f(p, t) ans_1 = reduce(hcat, [reduce(vcat, From 3baa5a6b463fda2a6921a69d7b9c365cb999cfe1 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 9 Oct 2024 17:05:35 +0400 Subject: [PATCH 134/153] update Project.toml --- Project.toml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index bedb3a6918..6793d0bda0 100644 --- a/Project.toml +++ b/Project.toml @@ -43,7 +43,7 @@ Adapt = "4" AdvancedHMC = "0.6.1" Aqua = "0.8" ArrayInterface = "7.9" -CUDA = "5.3.2" +CUDA = "5.5.2" ChainRulesCore = "1.24" ComponentArrays = "0.15.16" Cubature = "1.5" @@ -54,14 +54,14 @@ DomainSets = "0.6, 0.7" Flux = "0.14.17" ForwardDiff = "0.10.36" Functors = "0.4.12" -Integrals = "4.4" +Integrals = "4.5" LineSearches = "7.2" LinearAlgebra = "1.10" LogDensityProblems = "2" -Lux = "1" -LuxCUDA = "0.3.2" -LuxCore = "0.1.24" -LuxLib = "0.3.48" +Lux = "1.1" +LuxCUDA = "0.3" +LuxCore = "1.0.1" +LuxLib = "1.3.2" MCMCChains = "6" MethodOfLines = "0.11" ModelingToolkit = "9.9" @@ -78,13 +78,13 @@ Random = "1" Reexport = "1.2" RuntimeGeneratedFunctions = "0.5.12" SafeTestsets = "0.1" -SciMLBase = "2.34" +SciMLBase = "2.56" Statistics = "1.10" SymbolicUtils = "1.5, 2, 3" Symbolics = "5.27.1, 6" Test = "1.10" UnPack = "1" -Zygote = "0.6.70" +Zygote = "0.6.71" julia = "1.10" [extras] From de008b3aa11c3f2507f76d659fa634db7523b014 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 9 Oct 2024 17:06:11 +0400 Subject: [PATCH 135/153] undo --- src/BPINN_ode.jl | 8 ++++---- src/advancedHMC_MCMC.jl | 24 ++++++++++++------------ src/dae_solve.jl | 8 ++++---- src/dgm.jl | 8 ++++---- src/discretize.jl | 6 +++--- src/ode_solve.jl | 22 +++++++++++----------- src/pinn_types.jl | 24 ++++++++++++------------ test/BPINN_PDE_tests.jl | 2 +- test/BPINN_Tests.jl | 14 +++++++------- test/NNODE_tests.jl | 2 +- test/NNPDE_tests.jl | 2 +- 11 files changed, 60 insertions(+), 60 deletions(-) diff --git a/src/BPINN_ode.jl b/src/BPINN_ode.jl index 9960006b18..f3a16ad61c 100644 --- a/src/BPINN_ode.jl +++ b/src/BPINN_ode.jl @@ -20,7 +20,7 @@ of the physics-informed neural network which is used as a solver for a standard ## Positional Arguments -* `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer`. +* `chain`: A neural network architecture, defined as a `Lux.AbstractLuxLayer`. * `Kernel`: Choice of MCMC Sampling Algorithm. Defaults to `AdvancedHMC.HMC` ## Keyword Arguments @@ -117,7 +117,7 @@ function BNNODE(chain, Kernel = HMC; strategy = nothing, draw_samples = 2000, numensemble = floor(Int, draw_samples / 3), estim_collocate = false, autodiff = false, progress = false, verbose = false) - !(chain isa Lux.AbstractExplicitLayer) && + !(chain isa Lux.AbstractLuxLayer) && (chain = adapt(FromFluxAdaptor(false, false), chain)) BNNODE(chain, Kernel, strategy, draw_samples, priorsNNw, param, l2std, @@ -224,7 +224,7 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem, ninv = length(param) t = collect(eltype(saveat), prob.tspan[1]:saveat:prob.tspan[2]) - if chain isa Lux.AbstractExplicitLayer + if chain isa Lux.AbstractLuxLayer θinit, st = Lux.setup(Random.default_rng(), chain) θ = [vector_to_parameters(samples[i][1:(end - ninv)], θinit) for i in 1:max(draw_samples - draw_samples ÷ 10, draw_samples - 1000)] @@ -233,7 +233,7 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem, # only need for size θinit = collect(ComponentArrays.ComponentArray(θinit)) else - throw(error("Only Lux.AbstractExplicitLayer neural networks are supported")) + throw(error("Only Lux.AbstractLuxLayer neural networks are supported")) end # constructing ensemble predictions diff --git a/src/advancedHMC_MCMC.jl b/src/advancedHMC_MCMC.jl index 7105346aa0..bba3fe72c0 100644 --- a/src/advancedHMC_MCMC.jl +++ b/src/advancedHMC_MCMC.jl @@ -43,7 +43,7 @@ mutable struct LogTargetDensity{C, S, ST <: AbstractTrainingStrategy, I, init_params, estim_collocate) end - function LogTargetDensity(dim, prob, chain::Lux.AbstractExplicitLayer, st, strategy, + function LogTargetDensity(dim, prob, chain::Lux.AbstractLuxLayer, st, strategy, dataset, priors, phystd, l2std, autodiff, physdt, extraparams, init_params::NamedTuple, estim_collocate) @@ -109,7 +109,7 @@ function L2loss2(Tar::LogTargetDensity, θ) # parameter estimation chosen or not if Tar.extraparams > 0 autodiff = Tar.autodiff - # Timepoints to enforce Physics + # Timepoints to enforce Physics t = Tar.dataset[end] u1 = Tar.dataset[2] û = Tar.dataset[1] @@ -131,12 +131,12 @@ function L2loss2(Tar::LogTargetDensity, θ) t[i]) for i in 1:length(û)] end - #form of NN output matrix output dim x n + #form of NN output matrix output dim x n deri_physsol = reduce(hcat, physsol) - + physlogprob = 0 for i in 1:length(Tar.prob.u0) - # can add phystd[i] for u[i] + # can add phystd[i] for u[i] physlogprob += logpdf(MvNormal(deri_physsol[i, :], LinearAlgebra.Diagonal(map(abs2, (Tar.l2std[i] * 4.0) .* @@ -337,12 +337,12 @@ function priorweights(Tar::LogTargetDensity, θ) end end -function generate_Tar(chain::Lux.AbstractExplicitLayer, init_params) +function generate_Tar(chain::Lux.AbstractLuxLayer, init_params) θ, st = Lux.setup(Random.default_rng(), chain) return init_params, chain, st end -function generate_Tar(chain::Lux.AbstractExplicitLayer, init_params::Nothing) +function generate_Tar(chain::Lux.AbstractLuxLayer, init_params::Nothing) θ, st = Lux.setup(Random.default_rng(), chain) return θ, chain, st end @@ -351,7 +351,7 @@ end NN OUTPUT AT t,θ ~ phi(t,θ). """ function (f::LogTargetDensity{C, S})(t::AbstractVector, - θ) where {C <: Lux.AbstractExplicitLayer, S} + θ) where {C <: Lux.AbstractLuxLayer, S} θ = vector_to_parameters(θ, f.init_params) y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t'), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st @@ -359,7 +359,7 @@ function (f::LogTargetDensity{C, S})(t::AbstractVector, end function (f::LogTargetDensity{C, S})(t::Number, - θ) where {C <: Lux.AbstractExplicitLayer, S} + θ) where {C <: Lux.AbstractLuxLayer, S} θ = vector_to_parameters(θ, f.init_params) y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), [t]), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st @@ -501,7 +501,7 @@ function ahmc_bayesian_pinn_ode(prob::SciMLBase.ODEProblem, chain; MCMCkwargs = (n_leapfrog = 30,), progress = false, verbose = false, estim_collocate = false) - !(chain isa Lux.AbstractExplicitLayer) && + !(chain isa Lux.AbstractLuxLayer) && (chain = adapt(FromFluxAdaptor(false, false), chain)) # NN parameter prior mean and variance(PriorsNN must be a tuple) if isinplace(prob) @@ -521,11 +521,11 @@ function ahmc_bayesian_pinn_ode(prob::SciMLBase.ODEProblem, chain; throw(error("Dataset Required for Parameter Estimation.")) end - if chain isa Lux.AbstractExplicitLayer + if chain isa Lux.AbstractLuxLayer # Lux-Named Tuple initial_nnθ, recon, st = generate_Tar(chain, init_params) else - error("Only Lux.AbstractExplicitLayer Neural networks are supported") + error("Only Lux.AbstractLuxLayer Neural networks are supported") end if nchains > Threads.nthreads() diff --git a/src/dae_solve.jl b/src/dae_solve.jl index 5a5ee83be3..2231f2ade6 100644 --- a/src/dae_solve.jl +++ b/src/dae_solve.jl @@ -16,7 +16,7 @@ of the physics-informed neural network which is used as a solver for a standard ## Positional Arguments -* `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.AbstractExplicitLayer`. +* `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.AbstractLuxLayer`. * `opt`: The optimizer to train the neural network. * `init_params`: The initial parameter of the neural network. By default, this is `nothing` which thus uses the random initialization provided by the neural network library. @@ -42,7 +42,7 @@ end function NNDAE(chain, opt, init_params = nothing; strategy = nothing, autodiff = false, kwargs...) - !(chain isa Lux.AbstractExplicitLayer) && + !(chain isa Lux.AbstractLuxLayer) && (chain = adapt(FromFluxAdaptor(false, false), chain)) NNDAE(chain, opt, init_params, autodiff, strategy, kwargs) end @@ -110,12 +110,12 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractDAEProblem, # A logical array which declares which variables are the differential (non-algebraic) vars differential_vars = prob.differential_vars - if chain isa Lux.AbstractExplicitLayer || chain isa Flux.Chain + if chain isa Lux.AbstractLuxLayer || chain isa Flux.Chain phi, init_params = generate_phi_θ(chain, t0, u0, init_params) init_params = ComponentArrays.ComponentArray(; depvar = ComponentArrays.ComponentArray(init_params)) else - error("Only Lux.AbstractExplicitLayer and Flux.Chain neural networks are supported") + error("Only Lux.AbstractLuxLayer and Flux.Chain neural networks are supported") end if isinplace(prob) diff --git a/src/dgm.jl b/src/dgm.jl index 40fe88134e..44e60665de 100644 --- a/src/dgm.jl +++ b/src/dgm.jl @@ -1,4 +1,4 @@ -struct dgm_lstm_layer{F1, F2} <: Lux.AbstractExplicitLayer +struct dgm_lstm_layer{F1, F2} <: Lux.AbstractLuxLayer activation1::Function activation2::Function in_dims::Int @@ -49,7 +49,7 @@ function (layer::dgm_lstm_layer)( return S_new, st end -struct dgm_lstm_block{L <: NamedTuple} <: Lux.AbstractExplicitContainerLayer{(:layers,)} +struct dgm_lstm_block{L <: NamedTuple} <: Lux.AbstractLuxContainerLayer{(:layers,)} layers::L end @@ -80,7 +80,7 @@ function (L::dgm_lstm_block)( return apply_dgm_lstm_block(L.layers, S, x, ps, st) end -struct dgm{S, L, E} <: Lux.AbstractExplicitContainerLayer{(:d_start, :lstm, :d_end)} +struct dgm{S, L, E} <: Lux.AbstractLuxContainerLayer{(:d_start, :lstm, :d_end)} d_start::S lstm::L d_end::E @@ -138,7 +138,7 @@ function dgm(in_dims::Int, out_dims::Int, modes::Int, layers::Int, end """ - DeepGalerkin(in_dims::Int, out_dims::Int, modes::Int, L::Int, activation1::Function, activation2::Function, out_activation::Function, + DeepGalerkin(in_dims::Int, out_dims::Int, modes::Int, L::Int, activation1::Function, activation2::Function, out_activation::Function, strategy::NeuralPDE.AbstractTrainingStrategy; kwargs...) returns a `discretize` algorithm for the ModelingToolkit PDESystem interface, which transforms a `PDESystem` into an `OptimizationProblem` using the Deep Galerkin method. diff --git a/src/discretize.jl b/src/discretize.jl index 9a40e0fe82..93ae038824 100644 --- a/src/discretize.jl +++ b/src/discretize.jl @@ -23,7 +23,7 @@ to end end) -for Lux.AbstractExplicitLayer. +for Lux.AbstractLuxLayer. """ function build_symbolic_loss_function(pinnrep::PINNRepresentation, eqs; eq_params = SciMLBase.NullParameters(), @@ -441,12 +441,12 @@ function SciMLBase.symbolic_discretize(pde_system::PDESystem, phi = discretization.phi - if (phi isa Vector && phi[1].f isa Lux.AbstractExplicitLayer) + if (phi isa Vector && phi[1].f isa Lux.AbstractLuxLayer) for ϕ in phi ϕ.st = adapt(parameterless_type(ComponentArrays.getdata(flat_init_params)), ϕ.st) end - elseif (!(phi isa Vector) && phi.f isa Lux.AbstractExplicitLayer) + elseif (!(phi isa Vector) && phi.f isa Lux.AbstractLuxLayer) phi.st = adapt(parameterless_type(ComponentArrays.getdata(flat_init_params)), phi.st) end diff --git a/src/ode_solve.jl b/src/ode_solve.jl index bcf9c68ebe..137cf85a88 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -14,7 +14,7 @@ of the physics-informed neural network which is used as a solver for a standard ## Positional Arguments -* `chain`: A neural network architecture, defined as a `Lux.AbstractExplicitLayer` or `Flux.Chain`. +* `chain`: A neural network architecture, defined as a `Lux.AbstractLuxLayer` or `Flux.Chain`. `Flux.Chain` will be converted to `Lux` using `adapt(FromFluxAdaptor(false, false), chain)`. * `opt`: The optimizer to train the neural network. * `init_params`: The initial parameter of the neural network. By default, this is `nothing` @@ -88,14 +88,14 @@ end function NNODE(chain, opt, init_params = nothing; strategy = nothing, autodiff = false, batch = true, param_estim = false, additional_loss = nothing, kwargs...) - !(chain isa Lux.AbstractExplicitLayer) && + !(chain isa Lux.AbstractLuxLayer) && (chain = adapt(FromFluxAdaptor(false, false), chain)) NNODE(chain, opt, init_params, autodiff, batch, strategy, param_estim, additional_loss, kwargs) end """ - ODEPhi(chain::Lux.AbstractExplicitLayer, t, u0, st) + ODEPhi(chain::Lux.AbstractLuxLayer, t, u0, st) Internal struct, used for representing the ODE solution as a neural network in a form that respects boundary conditions, i.e. `phi(t) = u0 + t*NN(t)`. @@ -105,19 +105,19 @@ mutable struct ODEPhi{C, T, U, S} t0::T u0::U st::S - function ODEPhi(chain::Lux.AbstractExplicitLayer, t::Number, u0, st) + function ODEPhi(chain::Lux.AbstractLuxLayer, t::Number, u0, st) new{typeof(chain), typeof(t), typeof(u0), typeof(st)}(chain, t, u0, st) end end -function generate_phi_θ(chain::Lux.AbstractExplicitLayer, t, u0, init_params) +function generate_phi_θ(chain::Lux.AbstractLuxLayer, t, u0, init_params) θ, st = Lux.setup(Random.default_rng(), chain) isnothing(init_params) && (init_params = θ) ODEPhi(chain, t, u0, st), init_params end function (f::ODEPhi{C, T, U})(t::Number, - θ) where {C <: Lux.AbstractExplicitLayer, T, U <: Number} + θ) where {C <: Lux.AbstractLuxLayer, T, U <: Number} y, st = f.chain( adapt(parameterless_type(ComponentArrays.getdata(θ.depvar)), [t]), θ.depvar, f.st) ChainRulesCore.@ignore_derivatives f.st = st @@ -125,7 +125,7 @@ function (f::ODEPhi{C, T, U})(t::Number, end function (f::ODEPhi{C, T, U})(t::AbstractVector, - θ) where {C <: Lux.AbstractExplicitLayer, T, U <: Number} + θ) where {C <: Lux.AbstractLuxLayer, T, U <: Number} # Batch via data as row vectors y, st = f.chain( adapt(parameterless_type(ComponentArrays.getdata(θ.depvar)), t'), θ.depvar, f.st) @@ -133,7 +133,7 @@ function (f::ODEPhi{C, T, U})(t::AbstractVector, f.u0 .+ (t' .- f.t0) .* y end -function (f::ODEPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractExplicitLayer, T, U} +function (f::ODEPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractLuxLayer, T, U} y, st = f.chain( adapt(parameterless_type(ComponentArrays.getdata(θ.depvar)), [t]), θ.depvar, f.st) ChainRulesCore.@ignore_derivatives f.st = st @@ -141,7 +141,7 @@ function (f::ODEPhi{C, T, U})(t::Number, θ) where {C <: Lux.AbstractExplicitLay end function (f::ODEPhi{C, T, U})(t::AbstractVector, - θ) where {C <: Lux.AbstractExplicitLayer, T, U} + θ) where {C <: Lux.AbstractLuxLayer, T, U} # Batch via data as row vectors y, st = f.chain( adapt(parameterless_type(ComponentArrays.getdata(θ.depvar)), t'), θ.depvar, f.st) @@ -367,8 +367,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, #train points generation init_params = alg.init_params - !(chain isa Lux.AbstractExplicitLayer) && - error("Only Lux.AbstractExplicitLayer neural networks are supported") + !(chain isa Lux.AbstractLuxLayer) && + error("Only Lux.AbstractLuxLayer neural networks are supported") phi, init_params = generate_phi_θ(chain, t0, u0, init_params) (recursive_eltype(init_params) <: Complex && alg.strategy isa QuadratureTraining) && diff --git a/src/pinn_types.jl b/src/pinn_types.jl index 59480d8a60..fc948f9fa1 100644 --- a/src/pinn_types.jl +++ b/src/pinn_types.jl @@ -105,10 +105,10 @@ struct PhysicsInformedNN{T, P, PH, DER, PE, AL, ADA, LOG, K} <: AbstractPINN kwargs...) multioutput = chain isa AbstractArray if multioutput - !all(i -> i isa Lux.AbstractExplicitLayer, chain) && + !all(i -> i isa Lux.AbstractLuxLayer, chain) && (chain = Lux.transform.(chain)) else - !(chain isa Lux.AbstractExplicitLayer) && + !(chain isa Lux.AbstractLuxLayer) && (chain = adapt(FromFluxAdaptor(false, false), chain)) end if phi === nothing @@ -119,10 +119,10 @@ struct PhysicsInformedNN{T, P, PH, DER, PE, AL, ADA, LOG, K} <: AbstractPINN end else if multioutput - all([phi.f[i] isa Lux.AbstractExplicitLayer for i in eachindex(phi.f)]) || + all([phi.f[i] isa Lux.AbstractLuxLayer for i in eachindex(phi.f)]) || throw(ArgumentError("Only Lux Chains are supported")) else - (phi.f isa Lux.AbstractExplicitLayer) || + (phi.f isa Lux.AbstractLuxLayer) || throw(ArgumentError("Only Lux Chains are supported")) end _phi = phi @@ -246,10 +246,10 @@ struct BayesianPINN{T, P, PH, DER, PE, AL, ADA, LOG, D, K} <: AbstractPINN kwargs...) multioutput = chain isa AbstractArray if multioutput - !all(i -> i isa Lux.AbstractExplicitLayer, chain) && + !all(i -> i isa Lux.AbstractLuxLayer, chain) && (chain = Lux.transform.(chain)) else - !(chain isa Lux.AbstractExplicitLayer) && + !(chain isa Lux.AbstractLuxLayer) && (chain = adapt(FromFluxAdaptor(false, false), chain)) end if phi === nothing @@ -260,10 +260,10 @@ struct BayesianPINN{T, P, PH, DER, PE, AL, ADA, LOG, D, K} <: AbstractPINN end else if multioutput - all([phi.f[i] isa Lux.AbstractExplicitLayer for i in eachindex(phi.f)]) || + all([phi.f[i] isa Lux.AbstractLuxLayer for i in eachindex(phi.f)]) || throw(ArgumentError("Only Lux Chains are supported")) else - (phi.f isa Lux.AbstractExplicitLayer) || + (phi.f isa Lux.AbstractLuxLayer) || throw(ArgumentError("Only Lux Chains are supported")) end _phi = phi @@ -493,24 +493,24 @@ value at domain points x Fields: - `f`: A representation of the chain function. -- `st`: The state of the Lux.AbstractExplicitLayer. It should be updated on each call. +- `st`: The state of the Lux.AbstractLuxLayer. It should be updated on each call. """ mutable struct Phi{C, S} f::C st::S - function Phi(chain::Lux.AbstractExplicitLayer) + function Phi(chain::Lux.AbstractLuxLayer) st = Lux.initialstates(Random.default_rng(), chain) new{typeof(chain), typeof(st)}(chain, st) end end -function (f::Phi{<:Lux.AbstractExplicitLayer})(x::Number, θ) +function (f::Phi{<:Lux.AbstractLuxLayer})(x::Number, θ) y, st = f.f(adapt(parameterless_type(ComponentArrays.getdata(θ)), [x]), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st y end -function (f::Phi{<:Lux.AbstractExplicitLayer})(x::AbstractArray, θ) +function (f::Phi{<:Lux.AbstractLuxLayer})(x::AbstractArray, θ) y, st = f.f(adapt(parameterless_type(ComponentArrays.getdata(θ)), x), θ, f.st) ChainRulesCore.@ignore_derivatives f.st = st y diff --git a/test/BPINN_PDE_tests.jl b/test/BPINN_PDE_tests.jl index 98cacb748c..21827100bc 100644 --- a/test/BPINN_PDE_tests.jl +++ b/test/BPINN_PDE_tests.jl @@ -191,7 +191,7 @@ end chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1)) discretization = BayesianPINN([chain], GridTraining([0.01])) - @test discretization.chain[1] isa Lux.AbstractExplicitLayer + @test discretization.chain[1] isa Lux.AbstractLuxLayer @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) diff --git a/test/BPINN_Tests.jl b/test/BPINN_Tests.jl index 6534e88409..aa09b08b9c 100644 --- a/test/BPINN_Tests.jl +++ b/test/BPINN_Tests.jl @@ -53,7 +53,7 @@ Random.seed!(100) @test mean(abs.(x̂ .- meanscurve)) < 0.05 @test mean(abs.(physsol1 .- meanscurve)) < 0.005 - #--------------------- solve() call + #--------------------- solve() call @test mean(abs.(x̂1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 @test mean(abs.(physsol0_1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 end @@ -117,13 +117,13 @@ end luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - # --------------------- ahmc_bayesian_pinn_ode() call + # --------------------- ahmc_bayesian_pinn_ode() call @test mean(abs.(physsol1 .- meanscurve)) < 0.15 # ESTIMATED ODE PARAMETERS (NN1 AND NN2) @test abs(p - mean([fhsamples[i][23] for i in 2000:length(fhsamples)])) < abs(0.35 * p) - #-------------------------- solve() call + #-------------------------- solve() call @test mean(abs.(physsol1_1 .- pmean(sol2lux.ensemblesol[1]))) < 8e-2 # ESTIMATED ODE PARAMETERS (NN1 AND NN2) @@ -212,7 +212,7 @@ end param1 = mean(i[62] for i in fhsampleslux22[1000:length(fhsampleslux22)]) @test abs(param1 - p) < abs(0.3 * p) - #-------------------------- solve() call + #-------------------------- solve() call # (lux chain) @test mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 0.15 # estimated parameters(lux chain) @@ -245,7 +245,7 @@ end fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( prob, chainflux, draw_samples = 2500) alg = BNNODE(chainflux, draw_samples = 2500) - @test alg.chain isa Lux.AbstractExplicitLayer + @test alg.chain isa Lux.AbstractLuxLayer end @testset "Example 3 but with the new objective" begin @@ -337,7 +337,7 @@ end @test abs(param1 - p) < abs(0.75 * p) @test abs(param2 - p) < abs(param1 - p) - #-------------------------- solve() call + #-------------------------- solve() call # (lux chain) @test mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 0.1 # estimated parameters(lux chain) @@ -409,4 +409,4 @@ end bitvec = abs.(p .- sol_pestim1.estimated_de_params) .> abs.(p .- sol_pestim2.estimated_de_params) @test bitvec == ones(size(bitvec)) -end \ No newline at end of file +end diff --git a/test/NNODE_tests.jl b/test/NNODE_tests.jl index 0cd688e310..ab5da64ee5 100644 --- a/test/NNODE_tests.jl +++ b/test/NNODE_tests.jl @@ -310,7 +310,7 @@ end u_analytical(x) = (1 / (2pi)) .* sin.(2pi .* x) fluxchain = Flux.Chain(Flux.Dense(1, 5, Flux.σ), Flux.Dense(5, 1)) alg1 = NNODE(fluxchain, opt) - @test alg1.chain isa Lux.AbstractExplicitLayer + @test alg1.chain isa Lux.AbstractLuxLayer sol1 = solve(prob, alg1, verbose = false, abstol = 1e-10, maxiters = 200) @test sol1.errors[:l2] < 0.5 end diff --git a/test/NNPDE_tests.jl b/test/NNPDE_tests.jl index 7236ac041c..78a976d165 100644 --- a/test/NNPDE_tests.jl +++ b/test/NNPDE_tests.jl @@ -426,7 +426,7 @@ end chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1)) discretization = PhysicsInformedNN(chain, QuadratureTraining()) - @test discretization.chain isa Lux.AbstractExplicitLayer + @test discretization.chain isa Lux.AbstractLuxLayer @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) prob = discretize(pde_system, discretization) From 356083be422376addccf15462c9876cb0be4239f Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 10 Oct 2024 15:16:36 +0400 Subject: [PATCH 136/153] update runtests.jl --- src/NeuralPDE.jl | 1 + test/runtests.jl | 13 +++++++++++++ 2 files changed, 14 insertions(+) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 7c34e88dc5..128aa72c26 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -34,6 +34,7 @@ import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor, recursive_eltype using ChainRulesCore: @non_differentiable using NeuralOperators +import Lux: Chain import NeuralOperators: DeepONet RuntimeGeneratedFunctions.init(@__MODULE__) diff --git a/test/runtests.jl b/test/runtests.jl index 57e1bbe518..85573f061a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,5 +1,18 @@ using Pkg using SafeTestsets +import NeuralPDE: NNODE, NNDAE, PINOODE + PhysicsInformedNN, discretize, + GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, + WeightedIntervalTraining, + build_loss_function, get_loss_function, + generate_training_sets, get_variables, get_argument, get_bounds, + get_numeric_integral, symbolic_discretize, + AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, + MiniMaxAdaptiveLoss, LogOptions, + ahmc_bayesian_pinn_ode, BNNODE, ahmc_bayesian_pinn_pde, vector_to_parameters, + BPINNsolution, BayesianPINN, + DeepGalerkin + const GROUP = get(ENV, "GROUP", "All") From f59f22526a01aa6e05da3405384b55cbe3a66d44 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 10 Oct 2024 15:45:56 +0400 Subject: [PATCH 137/153] update NeuralPDE --- src/NeuralPDE.jl | 2 +- test/runtests.jl | 13 ------------- 2 files changed, 1 insertion(+), 14 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 128aa72c26..7121ba3dd9 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -58,7 +58,7 @@ include("BPINN_ode.jl") include("PDE_BPINN.jl") include("dgm.jl") -export NNODE, NNDAE, PINOODE +export NNODE, NNDAE, PINOODE, PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, WeightedIntervalTraining, diff --git a/test/runtests.jl b/test/runtests.jl index 85573f061a..57e1bbe518 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,18 +1,5 @@ using Pkg using SafeTestsets -import NeuralPDE: NNODE, NNDAE, PINOODE - PhysicsInformedNN, discretize, - GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, - WeightedIntervalTraining, - build_loss_function, get_loss_function, - generate_training_sets, get_variables, get_argument, get_bounds, - get_numeric_integral, symbolic_discretize, - AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, - MiniMaxAdaptiveLoss, LogOptions, - ahmc_bayesian_pinn_ode, BNNODE, ahmc_bayesian_pinn_pde, vector_to_parameters, - BPINNsolution, BayesianPINN, - DeepGalerkin - const GROUP = get(ENV, "GROUP", "All") From 526fd0c7ba8d577ccd12e1954da067465d2b337f Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 10 Oct 2024 16:47:24 +0400 Subject: [PATCH 138/153] update Doc --- docs/Project.toml | 14 +++++++------- docs/src/tutorials/pino_ode.md | 6 +++--- src/NeuralPDE.jl | 1 - 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index 3e62098b0a..4071801aa7 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -29,21 +29,21 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [compat] AdvancedHMC = "0.6.1" -ComponentArrays = "0.15.8" +ComponentArrays = "0.15.16" Cubature = "1.5" DiffEqBase = "6.148" Distributions = "0.25.107" Documenter = "1" DomainSets = "0.6, 0.7" -Flux = "0.14.11" -Integrals = "4" +Flux = "0.14.17" +Integrals = "4.5" LineSearches = "7.2" -Lux = "0.5.22" -LuxCUDA = "0.3.2" +Lux = "1.1" +LuxCUDA = "0.3" MethodOfLines = "0.11" -ModelingToolkit = "9.7" +ModelingToolkit = "9.9" MonteCarloMeasurements = "1" -NeuralPDE = "5.14" +NeuralPDE = "5.16" Optimization = "3.24, 4" OptimizationOptimJL = "0.2.1, 0.3, 0.4" OptimizationOptimisers = "0.2.1, 0.3" diff --git a/docs/src/tutorials/pino_ode.md b/docs/src/tutorials/pino_ode.md index 98c0fba269..fb4d2790d9 100644 --- a/docs/src/tutorials/pino_ode.md +++ b/docs/src/tutorials/pino_ode.md @@ -7,7 +7,7 @@ This tutorial provides an example of how to use the Physics Informed Neural Oper 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 # hide +using Test using OptimizationOptimisers using Lux using Statistics, Random @@ -33,12 +33,12 @@ deeponet = NeuralOperators.DeepONet( bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] number_of_parameter_samples = 50 # Define the training strategy -strategy = StochasticTraining(60) +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 = 3000) +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. diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 7121ba3dd9..4195753db7 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -34,7 +34,6 @@ import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor, recursive_eltype using ChainRulesCore: @non_differentiable using NeuralOperators -import Lux: Chain import NeuralOperators: DeepONet RuntimeGeneratedFunctions.init(@__MODULE__) From d3ae5948afba75d378e648f319476cc4430de9f0 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 10 Oct 2024 17:07:18 +0400 Subject: [PATCH 139/153] update --- src/NeuralPDE.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 4195753db7..7121ba3dd9 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -34,6 +34,7 @@ import ChainRulesCore, Lux, ComponentArrays using Lux: FromFluxAdaptor, recursive_eltype using ChainRulesCore: @non_differentiable using NeuralOperators +import Lux: Chain import NeuralOperators: DeepONet RuntimeGeneratedFunctions.init(@__MODULE__) From 6dd9e3861bf9c4e31e75ac1dec64a166fa095e2a Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 17 Oct 2024 17:35:51 +0400 Subject: [PATCH 140/153] update --- Project.toml | 4 +- src/advancedHMC_MCMC.jl | 2 + src/pino_ode_solve.jl | 91 +++++++++++++++++++------------------- src/training_strategies.jl | 4 +- 4 files changed, 51 insertions(+), 50 deletions(-) diff --git a/Project.toml b/Project.toml index 14ba864c90..6b1d3b9e1c 100644 --- a/Project.toml +++ b/Project.toml @@ -124,12 +124,10 @@ MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4" OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -Preferences = "21216c6a-2e73-6563-6e65-726566657250" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Aqua", "CUDA", "DiffEqNoiseProcess", "ExplicitImports", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "SafeTestsets", "StochasticDiffEq", "TensorBoardLogger", "Test"] - +test = ["Aqua", "CUDA", "DiffEqNoiseProcess", "ExplicitImports", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "SafeTestsets", "StochasticDiffEq", "TensorBoardLogger", "Test"] \ No newline at end of file diff --git a/src/advancedHMC_MCMC.jl b/src/advancedHMC_MCMC.jl index 1751cbc82e..380d284f55 100644 --- a/src/advancedHMC_MCMC.jl +++ b/src/advancedHMC_MCMC.jl @@ -73,6 +73,7 @@ suggested extra loss function for ODE solver case """ @views function L2loss2(ltd::LogTargetDensity, θ) ltd.extraparams ≤ 0 && return false # XXX: type-stability? + f = ltd.prob.f t = ltd.dataset[end] u1 = ltd.dataset[2] @@ -226,6 +227,7 @@ Prior logpdf for NN parameters + ODE constants. @views function priorweights(ltd::LogTargetDensity, θ) allparams = ltd.priors nnwparams = allparams[1] # nn weights + ltd.extraparams ≤ 0 && return logpdf(nnwparams, θ) # Vector of ode parameters priors diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index b10a3a880a..b1fa2c9be6 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -11,7 +11,7 @@ neural operator, which is used as a solver for a parametrized `ODEProblem`. ## Positional Arguments -* `chain`: A neural network architecture, defined as a `Lux.AbstractLuxLayer` or `Flux.Chain`. +* `chain`: A neural network architecture, defined as a `AbstractLuxLayer` or `Flux.Chain`. `Flux.Chain` will be converted to `Lux` using `adapt(FromFluxAdaptor(false, false), chain)` * `opt`: The optimizer to train the neural network. * `bounds`: A dictionary containing the bounds for the parameters of the parametric ODE. @@ -30,17 +30,15 @@ neural operator, which is used as a solver for a parametrized `ODEProblem`. * Sifan Wang "Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets" * Zongyi Li "Physics-Informed Neural Operator for Learning Partial Differential Equations" """ -struct PINOODE{C, O, B, I, S <: Union{Nothing, AbstractTrainingStrategy}, - AL <: Union{Nothing, Function}, K} <: - SciMLBase.AbstractODEAlgorithm - chain::C - opt::O - bounds::B +@concrete struct PINOODE + chain + opt + bounds number_of_parameters::Int - init_params::I - strategy::S - additional_loss::AL - kwargs::K + init_params + strategy <: Union{Nothing, AbstractTrainingStrategy} + additional_loss <: Union{Nothing, Function} + kwargs end function PINOODE(chain, @@ -51,38 +49,37 @@ function PINOODE(chain, strategy = nothing, additional_loss = nothing, kwargs...) - !(chain isa Lux.AbstractLuxLayer) && (chain = Lux.transform(chain)) - PINOODE(chain, opt, bounds, number_of_parameters, + chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain)) + return PINOODE(chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss, kwargs) end -struct PINOPhi{C, S} - chain::C - st::S - function PINOPhi(chain::Lux.AbstractLuxLayer, st) - new{typeof(chain), typeof(st)}(chain, st) - end +@concrete struct PINOPhi + model <:AbstractLuxLayer + smodel <: StatefulLuxLayer +end + +function PINOPhi(model::AbstractLuxLayer, st) + return PINOPhi(model, StatefulLuxLayer{false}(model, nothing, st)) end -function generate_pino_phi_θ(chain::Lux.AbstractLuxLayer, init_params) - θ, st = Lux.setup(Random.default_rng(), chain) - init_params = isnothing(init_params) ? θ : init_params - init_params = ComponentArrays.ComponentArray(init_params) +function generate_pino_phi_θ(chain::AbstractLuxLayer, nothing) + θ, st = LuxCore.setup(Random.default_rng(), chain) + PINOPhi(chain, st), θ +end + +function generate_pino_phi_θ(chain::AbstractLuxLayer, init_params) + st = LuxCore.initialstates(Random.default_rng(), chain) PINOPhi(chain, st), init_params end -function (f::PINOPhi{C, T})(x::Array, θ) where {C <: Lux.Chain, T} - eltypeθ, typeθ = eltype(θ), parameterless_type(ComponentArrays.getdata(θ)) - x = convert.(eltypeθ, adapt(typeθ, x)) - y, st = f.chain(x, θ, f.st) - y +function (f::PINOPhi{C, T})(x, θ) where {C <: AbstractLuxLayer, T} + dev = safe_get_device(θ) + return f(dev, safe_expand(dev, x), θ) end -function (f::PINOPhi{C, T})(x::Tuple, θ) where {C <: DeepONet, T} - eltypeθ, typeθ = eltype(θ), parameterless_type(ComponentArrays.getdata(θ)) - x = (convert.(eltypeθ, adapt(typeθ, x[1])), convert.(eltypeθ, adapt(typeθ, x[2]))) - y, st = f.chain(x, θ, f.st) - y +function (f::PINOPhi{C, T})(dev, x, θ) where {C <: AbstractLuxLayer, T} + f.smodel(dev(x), θ) end function dfdx(phi::PINOPhi{C, T}, x::Tuple, θ) where {C <: DeepONet, T} @@ -180,7 +177,7 @@ function get_trainset( end function get_trainset( - strategy::GridTraining, chain::Lux.Chain, bounds, number_of_parameters, tspan, eltypeθ) + strategy::GridTraining, chain::Chain, bounds, number_of_parameters, tspan, eltypeθ) dt = strategy.dx tspan_ = tspan[1]:dt:tspan[2] pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) @@ -194,9 +191,9 @@ function get_trainset( end function get_trainset( - strategy::StochasticTraining, chain::Union{DeepONet, Lux.Chain}, + strategy::StochasticTraining, chain::Union{DeepONet, Chain}, bounds, number_of_parameters, tspan, eltypeθ) - (number_of_parameters != strategy.points && chain isa Lux.Chain) && + (number_of_parameters != strategy.points && chain isa Chain) && throw(error("number_of_parameters should be the same strategy.points for StochasticTraining")) p = reduce(vcat, [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] @@ -208,7 +205,8 @@ end function generate_loss( strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan, eltypeθ) + x = get_trainset( + strategy, phi.smodel.model, bounds, number_of_parameters, tspan, eltypeθ) function loss(θ, _) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end @@ -217,7 +215,8 @@ end function generate_loss( strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) function loss(θ, _) - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan, eltypeθ) + x = get_trainset( + strategy, phi.smodel.model, bounds, number_of_parameters, tspan, eltypeθ) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end @@ -240,10 +239,10 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, verbose = false, saveat = nothing, maxiters = nothing) - @unpack tspan, u0, f = prob - @unpack chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss = alg + (; tspan, u0, f) = prob + (; chain, opt, bounds, number_of_parameters, init_params, strategy, additional_loss) = alg - if !(chain isa Lux.AbstractLuxLayer) + if !(chain isa AbstractLuxLayer) error("Only Lux.AbstractLuxLayer neural networks are supported") if !(chain isa DeepONet) || !(chain isa Chain) @@ -252,7 +251,9 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end phi, init_params = generate_pino_phi_θ(chain, init_params) - eltypeθ = eltype(init_params) + + # init_params = ComponentArray(init_params) + # eltypeθ = eltype(init_params) #TODO? isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) @@ -260,14 +261,14 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, try if chain isa DeepONet in_dim = chain.branch.layers.layer_1.in_dims - u = rand(eltypeθ, in_dim, number_of_parameters) - v = rand(eltypeθ, 1, 10, 1) + u = rand(in_dim, number_of_parameters) + v = rand(1, 10, 1) x = (u, v) phi(x, init_params) end if chain isa Chain in_dim = chain.layers.layer_1.in_dims - x = rand(eltypeθ, in_dim, number_of_parameters) + x = rand(in_dim, number_of_parameters) phi(x, init_params) end catch err diff --git a/src/training_strategies.jl b/src/training_strategies.jl index 974f2529fa..ca07676f26 100644 --- a/src/training_strategies.jl +++ b/src/training_strategies.jl @@ -291,8 +291,8 @@ end """ WeightedIntervalTraining(weights, samples) -A training strategy that generates points for training based on the given inputs. -We split the timespan into equal segments based on the number of weights, +A training strategy that generates points for training based on the given inputs. +We split the timespan into equal segments based on the number of weights, then sample points in each segment based on that segments corresponding weight, such that the total number of sampled points is equivalent to the given samples From 6748febaf7fcd90fd2d5fa828d055c66aee539c9 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Thu, 17 Oct 2024 18:51:24 +0400 Subject: [PATCH 141/153] update --- Project.toml | 2 +- src/pino_ode_solve.jl | 14 +++++++------- test/PINO_ode_tests.jl | 12 ++---------- 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/Project.toml b/Project.toml index 6b1d3b9e1c..461e9601f7 100644 --- a/Project.toml +++ b/Project.toml @@ -84,7 +84,7 @@ MLDataDevices = "1.2.0" MethodOfLines = "0.11.6" ModelingToolkit = "9.46" MonteCarloMeasurements = "1.1" -NeuralOperators = "0.5.0" +NeuralOperators = "0.5" Optimisers = "0.3.3" Optimization = "4" OptimizationOptimJL = "0.4" diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index b1fa2c9be6..91f7339523 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -63,7 +63,7 @@ function PINOPhi(model::AbstractLuxLayer, st) return PINOPhi(model, StatefulLuxLayer{false}(model, nothing, st)) end -function generate_pino_phi_θ(chain::AbstractLuxLayer, nothing) +function generate_pino_phi_θ(chain::AbstractLuxLayer, ::Nothing) θ, st = LuxCore.setup(Random.default_rng(), chain) PINOPhi(chain, st), θ end @@ -252,8 +252,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, phi, init_params = generate_pino_phi_θ(chain, init_params) - # init_params = ComponentArray(init_params) - # eltypeθ = eltype(init_params) #TODO? + init_params = ComponentArray(init_params) + eltypeθ = eltype(init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) @@ -261,14 +261,14 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, try if chain isa DeepONet in_dim = chain.branch.layers.layer_1.in_dims - u = rand(in_dim, number_of_parameters) - v = rand(1, 10, 1) + u = rand(eltypeθ, in_dim, number_of_parameters) + v = rand(eltypeθ, 1, 10, 1) x = (u, v) phi(x, init_params) end if chain isa Chain in_dim = chain.layers.layer_1.in_dims - x = rand(in_dim, number_of_parameters) + x = rand(eltypeθ, in_dim, number_of_parameters) phi(x, init_params) end catch err @@ -312,7 +312,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - x = get_trainset(strategy, phi.chain, bounds, number_of_parameters, tspan, eltypeθ) + x = get_trainset(strategy, phi.smodel.model, bounds, number_of_parameters, tspan, eltypeθ) if chain isa DeepONet u = phi(x, res.u) elseif chain isa Chain diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 00a2074b5b..8cb8d5dd93 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -49,7 +49,6 @@ end p, t = get_trainset(chain, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, t) predict_sol = sol.interp(reduce(vcat, (p, t))) - @test eltype(sol.k) == eltype(predict_sol) @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) @@ -96,7 +95,6 @@ end ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 - @test eltype(sol.k) == eltype(predict_sol) end @testset "Example du = cos(p * t) + u" begin @@ -124,7 +122,6 @@ end ground_solution = ground_analytic_.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 - @test eltype(sol.k) == eltype(predict_sol) end @testset "Example with data du = p*t^2" begin @@ -170,7 +167,6 @@ end ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp((p, t)) @test ground_solution≈predict_sol rtol=0.05 - @test eltype(sol.k) == eltype(predict_sol) end #multiple parameters chain @@ -204,16 +200,15 @@ end [[ground_solution(u0, p[:, i, j], t[1, i, j]) for j in axes(t, 3)] for i in axes(p, 2)])' end - (p, t) = get_trainset(chain, bounds, 50, tspan, 0.025f0) + (p, t) = get_trainset(chain, bounds, 20, tspan, 0.1f0) ground_solution_ = ground_solution_f(p, t) predict = sol.interp(reduce(vcat, (p, t)))[1, :, :] @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(chain, bounds, 60, tspan, 0.01f0) + p, t = get_trainset(chain, bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) predict_sol = sol.interp(reduce(vcat, (p, t)))[1, :, :] @test ground_solution_≈predict_sol rtol=0.05 - @test eltype(sol.k) == eltype(predict_sol) end #multiple parameters DeepOnet @@ -256,7 +251,6 @@ end ground_solution_ = ground_solution_f(p, t) predict = sol.interp((p, t)) @test ground_solution_≈predict rtol=0.05 - @test eltype(sol.k.u) == eltype(predict) end #vector output @@ -298,7 +292,6 @@ end @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 @test ground_solution_≈predict rtol=0.05 - @test eltype(sol.k) == eltype(predict) p, t = get_trainset(chain, bounds, 300, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) @@ -306,5 +299,4 @@ end @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 @test ground_solution_≈predict rtol=0.3 - @test eltype(sol.k) == eltype(predict) end From 36226f9f8be4ef1a705fe035d38a31eeb5a5a467 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 18 Oct 2024 17:49:32 +0400 Subject: [PATCH 142/153] update --- docs/Project.toml | 2 +- src/pino_ode_solve.jl | 32 +++++++++++++------------------- test/PINO_ode_tests.jl | 17 ++++++++--------- 3 files changed, 22 insertions(+), 29 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index 8d81acf010..3e19f980c4 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -45,7 +45,7 @@ MethodOfLines = "0.11" ModelingToolkit = "9.9" MonteCarloMeasurements = "1" NeuralPDE = "5" -NeuralOperators = "0.5.0" +NeuralOperators = "0.5" Optimization = "4" OptimizationOptimJL = "0.4" OptimizationOptimisers = "0.3" diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 91f7339523..6455bb9e0a 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -55,7 +55,7 @@ function PINOODE(chain, end @concrete struct PINOPhi - model <:AbstractLuxLayer + model <: AbstractLuxLayer smodel <: StatefulLuxLayer end @@ -166,18 +166,17 @@ function initial_condition_loss( end function get_trainset( - strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan, eltypeθ) + strategy::GridTraining, chain::DeepONet, bounds, number_of_parameters, tspan) dt = strategy.dx 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 = reshape(t_, 1, size(t_, 1), 1) - p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) (p, t) end function get_trainset( - strategy::GridTraining, chain::Chain, bounds, number_of_parameters, tspan, eltypeθ) + strategy::GridTraining, chain::Chain, bounds, number_of_parameters, tspan) dt = strategy.dx tspan_ = tspan[1]:dt:tspan[2] pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) @@ -186,37 +185,33 @@ function get_trainset( points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) p, t = x[1:(end - 1), :, :], x[[end], :, :] - p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) (p, t) end function get_trainset( strategy::StochasticTraining, chain::Union{DeepONet, Chain}, - bounds, number_of_parameters, tspan, eltypeθ) + bounds, number_of_parameters, tspan) (number_of_parameters != strategy.points && chain isa Chain) && throw(error("number_of_parameters should be the same strategy.points for StochasticTraining")) p = reduce(vcat, [(bound[2] .- bound[1]) .* rand(1, number_of_parameters) .+ bound[1] for bound in bounds]) t = (tspan[2] .- tspan[1]) .* rand(1, strategy.points, 1) .+ tspan[1] - p, t = convert.(eltypeθ, p), convert.(eltypeθ, t) (p, t) end function generate_loss( - strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) - x = get_trainset( - strategy, phi.smodel.model, bounds, number_of_parameters, tspan, eltypeθ) + strategy::GridTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) + x = get_trainset(strategy, phi.smodel.model, bounds, number_of_parameters, tspan) function loss(θ, _) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end function generate_loss( - strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan, eltypeθ) + strategy::StochasticTraining, prob::ODEProblem, phi, bounds, number_of_parameters, tspan) function loss(θ, _) - x = get_trainset( - strategy, phi.smodel.model, bounds, number_of_parameters, tspan, eltypeθ) + x = get_trainset(strategy, phi.smodel.model, bounds, number_of_parameters, tspan) initial_condition_loss(phi, prob, x, θ) + physics_loss(phi, prob, x, θ) end end @@ -253,7 +248,6 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, phi, init_params = generate_pino_phi_θ(chain, init_params) init_params = ComponentArray(init_params) - eltypeθ = eltype(init_params) isinplace(prob) && throw(error("The PINOODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) @@ -261,14 +255,14 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, try if chain isa DeepONet in_dim = chain.branch.layers.layer_1.in_dims - u = rand(eltypeθ, in_dim, number_of_parameters) - v = rand(eltypeθ, 1, 10, 1) + u = rand(in_dim, number_of_parameters) + v = rand(1, 10, 1) x = (u, v) phi(x, init_params) end if chain isa Chain in_dim = chain.layers.layer_1.in_dims - x = rand(eltypeθ, in_dim, number_of_parameters) + x = rand(in_dim, number_of_parameters) phi(x, init_params) end catch err @@ -286,7 +280,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, end inner_f = generate_loss( - strategy, prob, phi, bounds, number_of_parameters, tspan, eltypeθ) + strategy, prob, phi, bounds, number_of_parameters, tspan) function total_loss(θ, _) L2_loss = inner_f(θ, nothing) @@ -312,7 +306,7 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - x = get_trainset(strategy, phi.smodel.model, bounds, number_of_parameters, tspan, eltypeθ) + x = get_trainset(strategy, phi.smodel.model, bounds, number_of_parameters, tspan) if chain isa DeepONet u = phi(x, res.u) elseif chain isa Chain diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 8cb8d5dd93..fab52b9383 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,4 +1,3 @@ -using Test using OptimizationOptimisers using Lux using Statistics, Random @@ -25,7 +24,7 @@ function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) end #Test Chain with Float64 accuracy -@testset "Example du = cos(p * t)" begin +@testitem "Example du = cos(p * t)" tags=[:pinoode] begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0 @@ -57,7 +56,7 @@ end end #Test DeepONet with Float64 accuracy -@testset "Example du = cos(p * t)" begin +@testitem "Example du = cos(p * t)" tags=[:pinoode] begin equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0 @@ -97,7 +96,7 @@ end @test ground_solution≈predict_sol rtol=0.05 end -@testset "Example du = cos(p * t) + u" begin +@testitem "Example du = cos(p * t) + u" tags=[:pinoode] begin eq_(u, p, t) = cos(p * t) + u tspan = (0.0f0, 1.0f0) u0 = 1.0f0 @@ -124,7 +123,7 @@ end @test ground_solution≈predict_sol rtol=0.05 end -@testset "Example with data du = p*t^2" begin +@testitem "Example with data du = p*t^2" tags=[:pinoode] begin equation = (u, p, t) -> p * t^2 tspan = (0.0f0, 1.0f0) u0 = 0.0f0 @@ -169,8 +168,8 @@ end @test ground_solution≈predict_sol rtol=0.05 end -#multiple parameters chain -@testset "Example du = cos(p * t)" begin +#multiple parameters Сhain +@testitem "Example multiple parameters Сhain du = p1 * cos(p2 * t) + p3" tags=[:pinoode] begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0, 1.0) u0 = 1.0 @@ -212,7 +211,7 @@ end end #multiple parameters DeepOnet -@testset "Example du = cos(p * t)" begin +@testitem "Example multiple parameters DeepOnet du = p1 * cos(p2 * t) + p3" tags=[:pinoode] begin equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0, 1.0) u0 = 1.0 @@ -254,7 +253,7 @@ end end #vector output -@testset "Example du = [cos(p * t), sin(p * t)]" begin +@testitem "Example du = [cos(p * t), sin(p * t)]" tags=[:pinoode] begin equation = (u, p, t) -> [cos(p * t), sin(p * t)] tspan = (0.0f0, 1.0f0) u0 = [1.0f0, 0.0f0] From 22f514474cf40696ef11818f25a5ca54e5df3ed8 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 18 Oct 2024 18:12:32 +0400 Subject: [PATCH 143/153] PINOODETestSetup --- test/PINO_ode_tests.jl | 67 ++++++++++++++++++++++++------------------ 1 file changed, 38 insertions(+), 29 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index fab52b9383..4ce008e925 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,30 +1,33 @@ -using OptimizationOptimisers -using Lux -using Statistics, Random -using NeuralOperators -using NeuralPDE -function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, 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 = reshape(t_, 1, size(t_, 1), 1) - (p, t) -end +@testsetup module PINOODETestSetup + using OptimizationOptimisers + using Lux + using NeuralOperators + using NeuralPDE -function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) - tspan_ = tspan[1]:dt:tspan[2] - pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds] - x_ = hcat(vec(map( - points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) - x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) - p, t = x[1:(end - 1), :, :], x[[end], :, :] - (p, t) -end + function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, 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 = reshape(t_, 1, size(t_, 1), 1) + (p, t) + end + function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) + tspan_ = tspan[1]:dt:tspan[2] + pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds] + x_ = hcat(vec(map( + points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) + x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) + p, t = x[1:(end - 1), :, :], x[[end], :, :] + (p, t) + end + export get_trainset +end #Test Chain with Float64 accuracy -@testitem "Example du = cos(p * t)" tags=[:pinoode] begin +@testitem "Example du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0 @@ -56,7 +59,8 @@ end end #Test DeepONet with Float64 accuracy -@testitem "Example du = cos(p * t)" tags=[:pinoode] begin +@testitem "Example du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0 @@ -96,7 +100,8 @@ end @test ground_solution≈predict_sol rtol=0.05 end -@testitem "Example du = cos(p * t) + u" tags=[:pinoode] begin +@testitem "Example du = cos(p * t) + u" tags=[:pinoode] setup=[PINOODETestSetup] begin + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators eq_(u, p, t) = cos(p * t) + u tspan = (0.0f0, 1.0f0) u0 = 1.0f0 @@ -123,7 +128,8 @@ end @test ground_solution≈predict_sol rtol=0.05 end -@testitem "Example with data du = p*t^2" tags=[:pinoode] begin +@testitem "Example with data du = p*t^2" tags=[:pinoode] setup=[PINOODETestSetup] begin + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators equation = (u, p, t) -> p * t^2 tspan = (0.0f0, 1.0f0) u0 = 0.0f0 @@ -169,7 +175,8 @@ end end #multiple parameters Сhain -@testitem "Example multiple parameters Сhain du = p1 * cos(p2 * t) + p3" tags=[:pinoode] begin +@testitem "Example multiple parameters Сhain du = p1 * cos(p2 * t) + p3" tags=[:pinoode] setup=[PINOODETestSetup] begin + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0, 1.0) u0 = 1.0 @@ -211,7 +218,8 @@ end end #multiple parameters DeepOnet -@testitem "Example multiple parameters DeepOnet du = p1 * cos(p2 * t) + p3" tags=[:pinoode] begin +@testitem "Example multiple parameters DeepOnet du = p1 * cos(p2 * t) + p3" tags=[:pinoode] setup=[PINOODETestSetup] begin + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0, 1.0) u0 = 1.0 @@ -253,7 +261,8 @@ end end #vector output -@testitem "Example du = [cos(p * t), sin(p * t)]" tags=[:pinoode] begin +@testitem "Example du = [cos(p * t), sin(p * t)]" tags=[:pinoode] setup=[PINOODETestSetup] begin + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators equation = (u, p, t) -> [cos(p * t), sin(p * t)] tspan = (0.0f0, 1.0f0) u0 = [1.0f0, 0.0f0] From 18e45d274a7ee18a1ef007fb9130321d097b18c2 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 18 Oct 2024 18:21:35 +0400 Subject: [PATCH 144/153] update --- test/PINO_ode_tests.jl | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 4ce008e925..00160fe0ed 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,9 +1,6 @@ @testsetup module PINOODETestSetup - using OptimizationOptimisers - using Lux - using NeuralOperators - using NeuralPDE + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, dt) p_ = [range(start = b[1], length = number_of_parameters, stop = b[2]) for b in bounds] @@ -27,7 +24,7 @@ end #Test Chain with Float64 accuracy @testitem "Example du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0 @@ -60,7 +57,7 @@ end #Test DeepONet with Float64 accuracy @testitem "Example du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) u0 = 1.0 @@ -101,7 +98,7 @@ end end @testitem "Example du = cos(p * t) + u" tags=[:pinoode] setup=[PINOODETestSetup] begin - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random eq_(u, p, t) = cos(p * t) + u tspan = (0.0f0, 1.0f0) u0 = 1.0f0 @@ -129,7 +126,7 @@ end end @testitem "Example with data du = p*t^2" tags=[:pinoode] setup=[PINOODETestSetup] begin - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> p * t^2 tspan = (0.0f0, 1.0f0) u0 = 0.0f0 @@ -176,7 +173,7 @@ end #multiple parameters Сhain @testitem "Example multiple parameters Сhain du = p1 * cos(p2 * t) + p3" tags=[:pinoode] setup=[PINOODETestSetup] begin - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0, 1.0) u0 = 1.0 @@ -219,7 +216,7 @@ end #multiple parameters DeepOnet @testitem "Example multiple parameters DeepOnet du = p1 * cos(p2 * t) + p3" tags=[:pinoode] setup=[PINOODETestSetup] begin - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] tspan = (0.0, 1.0) u0 = 1.0 @@ -262,7 +259,7 @@ end #vector output @testitem "Example du = [cos(p * t), sin(p * t)]" tags=[:pinoode] setup=[PINOODETestSetup] begin - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators + using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> [cos(p * t), sin(p * t)] tspan = (0.0f0, 1.0f0) u0 = [1.0f0, 0.0f0] From aef24e34535295dba37c690b97338f2d0fa04bdb Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Fri, 18 Oct 2024 18:48:29 +0400 Subject: [PATCH 145/153] update --- test/PINO_ode_tests.jl | 53 ++++++++++++++++++++---------------------- 1 file changed, 25 insertions(+), 28 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 00160fe0ed..eca1ef5d1c 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,29 +1,29 @@ @testsetup module PINOODETestSetup - using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random +using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random - function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, 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 = reshape(t_, 1, size(t_, 1), 1) - (p, t) - end +function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, 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 = reshape(t_, 1, size(t_, 1), 1) + (p, t) +end - function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) - tspan_ = tspan[1]:dt:tspan[2] - pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) - for b in bounds] - x_ = hcat(vec(map( - points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) - x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) - p, t = x[1:(end - 1), :, :], x[[end], :, :] - (p, t) - end - export get_trainset +function get_trainset(chain::Lux.Chain, bounds, number_of_parameters, tspan, dt) + tspan_ = tspan[1]:dt:tspan[2] + pspan = [range(start = b[1], length = number_of_parameters, stop = b[2]) + for b in bounds] + x_ = hcat(vec(map( + points -> collect(points), Iterators.product([pspan..., tspan_]...)))...) + x = reshape(x_, size(bounds, 1) + 1, prod(size.(pspan, 1)), size(tspan_, 1)) + p, t = x[1:(end - 1), :, :], x[[end], :, :] + (p, t) end -#Test Chain with Float64 accuracy -@testitem "Example du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin +export get_trainset +end +#Test Chain +@testitem "Example Chain du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -39,9 +39,7 @@ end number_of_parameters = 300 strategy = StochasticTraining(300) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE( - chain, opt, bounds, number_of_parameters; strategy = strategy, init_params = θ |> - f64) + alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 5000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 @@ -55,8 +53,8 @@ end @test ground_solution≈predict_sol rtol=0.05 end -#Test DeepONet with Float64 accuracy -@testitem "Example du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin +#Test DeepONet +@testitem "Example DeepONet du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> cos(p * t) tspan = (0.0f0, 1.0f0) @@ -82,8 +80,7 @@ end number_of_parameters = 50 strategy = StochasticTraining(40) opt = OptimizationOptimisers.Adam(0.01) - alg = PINOODE(deeponet, opt, bounds, number_of_parameters; - strategy = strategy, init_params = θ |> f64) + alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 3000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) dt = 0.025f0 From 501496ed4ca04a53909916f5accb6e6ad9ce3527 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 28 Oct 2024 15:27:15 +0400 Subject: [PATCH 146/153] update Project.toml docs --- docs/Project.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index 3e19f980c4..935421b02a 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -30,7 +30,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [compat] AdvancedHMC = "0.6.1" -ComponentArrays = "0.15.16" +ComponentArrays = "0.15.8" Cubature = "1.5" DiffEqBase = "6.148" Distributions = "0.25.107" @@ -42,7 +42,7 @@ LineSearches = "7.2" Lux = "1" LuxCUDA = "0.3.2" MethodOfLines = "0.11" -ModelingToolkit = "9.9" +ModelingToolkit = "9.7" MonteCarloMeasurements = "1" NeuralPDE = "5" NeuralOperators = "0.5" @@ -55,4 +55,4 @@ Plots = "1.36" QuasiMonteCarlo = "0.3.2" Random = "1" Roots = "2.0" -SpecialFunctions = "2.1" +SpecialFunctions = "2.1" \ No newline at end of file From e8ac7f5f8ef246dc531c72f2bbdb3e18d1ed745d Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 28 Oct 2024 19:05:25 +0400 Subject: [PATCH 147/153] add sol(t) --- src/pino_ode_solve.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 6455bb9e0a..7458165aad 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -216,7 +216,7 @@ function generate_loss( end end -struct PINOODEInterpolation{T <: PINOPhi, T2} +@concrete struct PINOODEInterpolation{T <: PINOPhi, T2} phi::T θ::T2 end @@ -226,6 +226,10 @@ end SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation" SciMLBase.allowscomplex(::PINOODE) = true +function (sol::SciMLBase.AbstractODESolution)(t::AbstractArray) + sol.interp(t) +end + function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, alg::PINOODE, args...; From 11c67dacffa6b6b97966108ac6190f96f8c1d0ee Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 28 Oct 2024 19:13:41 +0400 Subject: [PATCH 148/153] sol(t) --- test/PINO_ode_tests.jl | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index eca1ef5d1c..f63c33325b 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,6 +1,6 @@ @testsetup module PINOODETestSetup -using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random +using Lux, NeuralOperators function get_trainset(chain::DeepONet, bounds, number_of_parameters, tspan, dt) p_ = [range(start = b[1], length = number_of_parameters, stop = b[2]) for b in bounds] @@ -45,11 +45,11 @@ end dt = 0.025f0 p, t = get_trainset(chain, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, t) - predict_sol = sol.interp(reduce(vcat, (p, t))) + predict_sol = sol(reduce(vcat, (p, t))) @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.interp(reduce(vcat, (p, t))) + predict_sol = sol(reduce(vcat, (p, t))) @test ground_solution≈predict_sol rtol=0.05 end @@ -86,11 +86,11 @@ end dt = 0.025f0 p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol.interp((p, t)) + predict_sol = sol((p, t)) @test ground_solution≈predict_sol rtol=0.05 p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol.interp((p, t)) + predict_sol = sol((p, t)) @test ground_solution≈predict_sol rtol=0.05 end @@ -118,7 +118,7 @@ end (p^2 + 1) p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic_.(u0, p, vec(t)) - predict_sol = sol.interp((p, t)) + predict_sol = sol((p, t)) @test ground_solution≈predict_sol rtol=0.05 end @@ -164,7 +164,7 @@ end p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol.interp((p, t)) + predict_sol = sol((p, t)) @test ground_solution≈predict_sol rtol=0.05 end @@ -202,12 +202,12 @@ end end (p, t) = get_trainset(chain, bounds, 20, tspan, 0.1f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp(reduce(vcat, (p, t)))[1, :, :] + predict = sol(reduce(vcat, (p, t)))[1, :, :] @test ground_solution_≈predict rtol=0.05 p, t = get_trainset(chain, bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) - predict_sol = sol.interp(reduce(vcat, (p, t)))[1, :, :] + predict_sol = sol(reduce(vcat, (p, t)))[1, :, :] @test ground_solution_≈predict_sol rtol=0.05 end @@ -245,12 +245,12 @@ end (p, t) = get_trainset(deeponet, bounds, 50, tspan, 0.025f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) + predict = sol((p, t)) @test ground_solution_≈predict rtol=0.05 p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp((p, t)) + predict = sol((p, t)) @test ground_solution_≈predict rtol=0.05 end @@ -290,14 +290,14 @@ end end p, t = get_trainset(chain, bounds, 50, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp(reduce(vcat, (p, t))) + predict = sol(reduce(vcat, (p, t))) @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 @test ground_solution_≈predict rtol=0.05 p, t = get_trainset(chain, bounds, 300, tspan, 0.01f0) ground_solution_ = ground_solution_f(p, t) - predict = sol.interp(reduce(vcat, (p, t))) + predict = sol(reduce(vcat, (p, t))) @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 @test ground_solution_≈predict rtol=0.3 From 35d8346ed95696e01071b01aaf50b0531430d639 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 29 Oct 2024 15:41:47 +0400 Subject: [PATCH 149/153] update --- Project.toml | 2 +- src/pino_ode_solve.jl | 8 ++++++++ test/PINO_ode_tests.jl | 26 ++++++++++++-------------- 3 files changed, 21 insertions(+), 15 deletions(-) diff --git a/Project.toml b/Project.toml index 22256ba76b..46fb9f6804 100644 --- a/Project.toml +++ b/Project.toml @@ -80,7 +80,7 @@ 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" diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 7458165aad..483273814c 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -226,7 +226,15 @@ end SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation" SciMLBase.allowscomplex(::PINOODE) = true +#TODO function (sol::SciMLBase.AbstractODESolution)(t::AbstractArray) + # p,t = sol.t + # sol.interp(reduce(vcat, (p, t))) + sol.interp(t) +end +function (sol::SciMLBase.AbstractODESolution)(t::Tuple) + # p,t = sol.t + # sol.interp((p, t)) sol.interp(t) end diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index f63c33325b..e7db5b9810 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -42,8 +42,7 @@ end alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 5000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - dt = 0.025f0 - p, t = get_trainset(chain, bounds, number_of_parameters, tspan, dt) + p, t = get_trainset(chain, bounds, 50, tspan, 0.025) ground_solution = ground_analytic.(u0, p, t) predict_sol = sol(reduce(vcat, (p, t))) @test ground_solution≈predict_sol rtol=0.05 @@ -83,8 +82,7 @@ end alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 3000) ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - dt = 0.025f0 - p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) + p, t = get_trainset(deeponet, bounds, 50, tspan, 0.025) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol((p, t)) @test ground_solution≈predict_sol rtol=0.05 @@ -200,12 +198,12 @@ end [[ground_solution(u0, p[:, i, j], t[1, i, j]) for j in axes(t, 3)] for i in axes(p, 2)])' end - (p, t) = get_trainset(chain, bounds, 20, tspan, 0.1f0) + (p, t) = get_trainset(chain, bounds, 20, tspan, 0.1) ground_solution_ = ground_solution_f(p, t) predict = sol(reduce(vcat, (p, t)))[1, :, :] @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(chain, bounds, 50, tspan, 0.025f0) + p, t = get_trainset(chain, bounds, 50, tspan, 0.025) ground_solution_ = ground_solution_f(p, t) predict_sol = sol(reduce(vcat, (p, t)))[1, :, :] @test ground_solution_≈predict_sol rtol=0.05 @@ -233,22 +231,22 @@ end bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] number_of_parameters = 50 - strategy = StochasticTraining(20) - opt = OptimizationOptimisers.Adam(0.03) + strategy = StochasticTraining(50) + opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 4000) + sol = solve(prob, alg, verbose = false, maxiters = 5000) 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 - (p, t) = get_trainset(deeponet, bounds, 50, tspan, 0.025f0) + (p, t) = get_trainset(deeponet, bounds, 50, tspan, 0.025) ground_solution_ = ground_solution_f(p, t) predict = sol((p, t)) @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01f0) + p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01) ground_solution_ = ground_solution_f(p, t) predict = sol((p, t)) @test ground_solution_≈predict rtol=0.05 @@ -271,7 +269,7 @@ end strategy = StochasticTraining(300) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 6000) + sol = solve(prob, alg, verbose = true, maxiters = 6000) ground_solution = (u0, p, t) -> [1 + sin(p * t) / p, 1 / p - cos(p * t) / p] function ground_solution_f(p, t) @@ -288,14 +286,14 @@ end ans_2 = reshape(ans_2, 1, size(ans_2)...) vcat(ans_1, ans_2) end - p, t = get_trainset(chain, bounds, 50, tspan, 0.01f0) + p, t = get_trainset(chain, bounds, 50, tspan, 0.025) ground_solution_ = ground_solution_f(p, t) predict = sol(reduce(vcat, (p, t))) @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 @test ground_solution_≈predict rtol=0.05 - p, t = get_trainset(chain, bounds, 300, tspan, 0.01f0) + p, t = get_trainset(chain, bounds, 300, tspan, 0.01) ground_solution_ = ground_solution_f(p, t) predict = sol(reduce(vcat, (p, t))) @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 From c5a456af6693ff60413e3c3ae42b8567fafff3a4 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 29 Oct 2024 20:00:11 +0400 Subject: [PATCH 150/153] Interpolation --- src/pino_ode_solve.jl | 59 ++++++++++++++++++++++++++++------ test/PINO_ode_tests.jl | 72 +++++++++++++++++++----------------------- 2 files changed, 83 insertions(+), 48 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 483273814c..0ac03f5fd9 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -216,26 +216,67 @@ function generate_loss( end end +""" +PINOODEInterpolation(phi, θ) + +Interpolation of the solution of the ODE using a trained neural network. + +## Arguments +* `phi`: The neural network +* `θ`: The parameters of the neural network. +``` + +""" @concrete struct PINOODEInterpolation{T <: PINOPhi, T2} phi::T θ::T2 end +""" +Override interpolation method for PINOODEInterpolation + +## Arguments +* `x`: Input data on which the solution is to be interpolated. +## Example + +```jldoctest +interp = PINOODEInterpolation(phi, θ) +x = rand(2, 50, 10) +interp(x) +``` +""" (f::PINOODEInterpolation)(x) = f.phi(x, f.θ) +""" +Override interpolation method for PINOODEInterpolation + +## Arguments +# * `p`: The parameters points on which the solution is to be interpolated. +# * `t`: The time points on which the solution is to be interpolated. + +## Example +```jldoctest +interp = PINOODEInterpolation(phi, θ) +p,t = rand(1, 50, 10), rand(1, 50, 10) +interp(p, t) +``` +""" +function (f::PINOODEInterpolation)(p, t) + if f.phi.model isa DeepONet + f.phi((p, t), f.θ) + elseif f.phi.model isa Chain + f.phi(reduce(vcat, (p, t)), f.θ) + else + error("Only DeepONet and Chain neural networks are supported with PINO ODE") + end +end + SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation" SciMLBase.allowscomplex(::PINOODE) = true -#TODO function (sol::SciMLBase.AbstractODESolution)(t::AbstractArray) - # p,t = sol.t - # sol.interp(reduce(vcat, (p, t))) - sol.interp(t) -end -function (sol::SciMLBase.AbstractODESolution)(t::Tuple) - # p,t = sol.t - # sol.interp((p, t)) - sol.interp(t) + p, _ = sol.t + sol.interp(p, t) end function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index e7db5b9810..5305e03da4 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -26,7 +26,7 @@ end @testitem "Example Chain du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 1.0f0) + tspan = (0.0, 1.0) u0 = 1.0 prob = ODEProblem(equation, u0, tspan) chain = Chain( @@ -44,11 +44,12 @@ end ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) p, t = get_trainset(chain, bounds, 50, tspan, 0.025) ground_solution = ground_analytic.(u0, p, t) - predict_sol = sol(reduce(vcat, (p, t))) + predict_sol = sol.interp(p, t) + predict_sol = sol.interp(reduce(vcat, (p, t))) @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))) + predict_sol = sol.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 end @@ -56,7 +57,7 @@ end @testitem "Example DeepONet du = cos(p * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> cos(p * t) - tspan = (0.0f0, 1.0f0) + tspan = (0.0, 1.0) u0 = 1.0 prob = ODEProblem(equation, u0, tspan) deeponet = NeuralOperators.DeepONet( @@ -84,29 +85,30 @@ end ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) p, t = get_trainset(deeponet, bounds, 50, tspan, 0.025) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol((p, t)) + predict_sol = sol(t) + predict_sol = sol.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol((p, t)) + predict_sol = sol.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 end @testitem "Example du = cos(p * t) + u" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random eq_(u, p, t) = cos(p * t) + u - tspan = (0.0f0, 1.0f0) - u0 = 1.0f0 + tspan = (0.0, 1.0) + u0 = 1.0 prob = ODEProblem(eq_, u0, tspan) deeponet = NeuralOperators.DeepONet( Chain( Dense(1 => 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))) - bounds = [(0.1f0, 2.0f0)] + bounds = [(0.1, 2.0)] number_of_parameters = 40 dt = (tspan[2] - tspan[1]) / 40 - strategy = GridTraining(0.1f0) + strategy = GridTraining(0.1) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 4000) @@ -116,53 +118,45 @@ end (p^2 + 1) p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic_.(u0, p, vec(t)) - predict_sol = sol((p, t)) + predict_sol = sol.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 end @testitem "Example with data du = p*t^2" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> p * t^2 - tspan = (0.0f0, 1.0f0) - u0 = 0.0f0 + tspan = (0.0, 1.0) + u0 = 0.0 prob = ODEProblem(equation, u0, tspan) deeponet = NeuralOperators.DeepONet( Chain( Dense(1 => 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))) - bounds = [(0.0f0, 10.0f0)] + bounds = [(0.0, 10.0)] number_of_parameters = 60 dt = (tspan[2] - tspan[1]) / 40 strategy = StochasticTraining(60) - opt = OptimizationOptimisers.Adam(0.03) + opt = OptimizationOptimisers.Adam(0.01) + #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 - - function get_data() - sol = ground_analytic.(u0, p, vec(t)) - tuple_ = (p, t) - sol, tuple_ - end - u = rand(1, 50) - v = rand(1, 40, 1) - θ, st = Lux.setup(Random.default_rng(), deeponet) - c = deeponet((u, v), θ, st)[1] + sol = ground_analytic.(u0, p, vec(t)) p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) - data, tuple_ = get_data() function additional_loss_(phi, θ) - u = phi(tuple_, θ) + u = phi((p, t), θ) norm = prod(size(u)) - sum(abs2, u .- data) / norm + sum(abs2, u .- sol) / norm end + alg = PINOODE( deeponet, opt, bounds, number_of_parameters; strategy = strategy, additional_loss = additional_loss_) - sol = solve(prob, alg, verbose = false, maxiters = 2000) + sol = solve(prob, alg, verbose = true, maxiters = 3000) p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol((p, t)) + predict_sol = sol.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 end @@ -200,12 +194,12 @@ end end (p, t) = get_trainset(chain, bounds, 20, tspan, 0.1) ground_solution_ = ground_solution_f(p, t) - predict = sol(reduce(vcat, (p, t)))[1, :, :] + predict = sol.interp(p, t)[1, :, :] @test ground_solution_≈predict rtol=0.05 p, t = get_trainset(chain, bounds, 50, tspan, 0.025) ground_solution_ = ground_solution_f(p, t) - predict_sol = sol(reduce(vcat, (p, t)))[1, :, :] + predict_sol = sol.interp(p, t)[1, :, :] @test ground_solution_≈predict_sol rtol=0.05 end @@ -243,12 +237,12 @@ end (p, t) = get_trainset(deeponet, bounds, 50, tspan, 0.025) ground_solution_ = ground_solution_f(p, t) - predict = sol((p, t)) + predict = sol.interp(p, t) @test ground_solution_≈predict rtol=0.05 p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01) ground_solution_ = ground_solution_f(p, t) - predict = sol((p, t)) + predict = sol.interp(p, t) @test ground_solution_≈predict rtol=0.05 end @@ -256,8 +250,8 @@ end @testitem "Example du = [cos(p * t), sin(p * t)]" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random equation = (u, p, t) -> [cos(p * t), sin(p * t)] - tspan = (0.0f0, 1.0f0) - u0 = [1.0f0, 0.0f0] + tspan = (0.0, 1.0) + u0 = [1.0, 0.0] prob = ODEProblem(equation, u0, tspan) input_branch_size = 1 chain = Chain( @@ -269,7 +263,7 @@ end strategy = StochasticTraining(300) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = true, maxiters = 6000) + sol = solve(prob, alg, verbose = false, maxiters = 6000) ground_solution = (u0, p, t) -> [1 + sin(p * t) / p, 1 / p - cos(p * t) / p] function ground_solution_f(p, t) @@ -288,14 +282,14 @@ end end p, t = get_trainset(chain, bounds, 50, tspan, 0.025) ground_solution_ = ground_solution_f(p, t) - predict = sol(reduce(vcat, (p, t))) + predict = sol.interp(p, t) @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 @test ground_solution_≈predict rtol=0.05 p, t = get_trainset(chain, bounds, 300, tspan, 0.01) ground_solution_ = ground_solution_f(p, t) - predict = sol(reduce(vcat, (p, t))) + predict = sol.interp(p, t) @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 @test ground_solution_≈predict rtol=0.3 From b6609f9eb9a2d7c13640e19a9935ebfccd5429e9 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Tue, 29 Oct 2024 20:18:55 +0400 Subject: [PATCH 151/153] fix --- test/PINO_ode_tests.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 5305e03da4..ff33ce176c 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -141,8 +141,8 @@ end #generate data ground_analytic = (u0, p, t) -> u0 + p * t^3 / 3 - sol = ground_analytic.(u0, p, vec(t)) p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) + sol = ground_analytic.(u0, p, vec(t)) function additional_loss_(phi, θ) u = phi((p, t), θ) norm = prod(size(u)) @@ -152,7 +152,7 @@ end alg = PINOODE( deeponet, opt, bounds, number_of_parameters; strategy = strategy, additional_loss = additional_loss_) - sol = solve(prob, alg, verbose = true, maxiters = 3000) + sol = solve(prob, alg, verbose = false, maxiters = 3000) p, t = get_trainset(deeponet, bounds, number_of_parameters, tspan, dt) ground_solution = ground_analytic.(u0, p, vec(t)) From 0617f4259ee80a805117159ea4b0c0440b83f49c Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 30 Oct 2024 17:19:12 +0400 Subject: [PATCH 152/153] update PINOODEInterpolation --- src/pino_ode_solve.jl | 38 +++++++++++++++++++++++++------------- test/PINO_ode_tests.jl | 19 ++++++++++++++++--- 2 files changed, 41 insertions(+), 16 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 0ac03f5fd9..5aff8284f1 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -245,7 +245,7 @@ x = rand(2, 50, 10) interp(x) ``` """ -(f::PINOODEInterpolation)(x) = f.phi(x, f.θ) +(f::PINOODEInterpolation)(x::AbstractArray) = f.phi(x, f.θ) """ Override interpolation method for PINOODEInterpolation @@ -261,22 +261,36 @@ p,t = rand(1, 50, 10), rand(1, 50, 10) interp(p, t) ``` """ -function (f::PINOODEInterpolation)(p, t) +function (f::PINOODEInterpolation)(p::AbstractArray, t::AbstractArray) if f.phi.model isa DeepONet f.phi((p, t), f.θ) elseif f.phi.model isa Chain + if size(p, 2) != size(t, 2) + error("t should be same size as p") + end f.phi(reduce(vcat, (p, t)), f.θ) else error("Only DeepONet and Chain neural networks are supported with PINO ODE") end end +function (f::PINOODEInterpolation)(p::AbstractArray, t::Number) + if f.phi.model isa DeepONet + t_ = [t] + f.phi((p, t_), f.θ) + elseif f.phi.model isa Chain + t_ = fill(t, size(p)) + f.phi(reduce(vcat, (p, t_)), f.θ) + else + error("Only DeepONet and Chain neural networks are supported with PINO ODE") + end +end + SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation" SciMLBase.allowscomplex(::PINOODE) = true -function (sol::SciMLBase.AbstractODESolution)(t::AbstractArray) - p, _ = sol.t - sol.interp(p, t) +function (sol::SciMLBase.AbstractODESolution)(t::Union{Number, AbstractArray}) + sol.interp(sol.prob.p, t) end function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, @@ -359,16 +373,14 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem, optprob = OptimizationProblem(optf, init_params) res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - x = get_trainset(strategy, phi.smodel.model, bounds, number_of_parameters, tspan) - if chain isa DeepONet - u = phi(x, res.u) - elseif chain isa Chain - u = phi(reduce(vcat, x), res.u) - end + (p, t) = get_trainset(strategy, phi.smodel.model, bounds, number_of_parameters, tspan) + interp = PINOODEInterpolation(phi, res.u) + u = interp(p, t) + prob_sol = ODEProblem(f.f, u0, tspan, p) - sol = SciMLBase.build_solution(prob, alg, x, u; + sol = SciMLBase.build_solution(prob_sol, alg, t, u; k = res, dense = true, - interp = PINOODEInterpolation(phi, res.u), + interp = interp, calculate_error = false, retcode = ReturnCode.Success, original = res, diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index ff33ce176c..6e4545daf9 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -1,4 +1,3 @@ - @testsetup module PINOODETestSetup using Lux, NeuralOperators @@ -45,12 +44,22 @@ end p, t = get_trainset(chain, bounds, 50, tspan, 0.025) ground_solution = ground_analytic.(u0, p, t) predict_sol = sol.interp(p, t) - predict_sol = sol.interp(reduce(vcat, (p, t))) @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.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 + + p = sol.prob.p + ground_solution = ground_analytic.(u0, p, [1.0]) + predict_sol = sol(1.0) + @test ground_solution≈predict_sol rtol=0.05 + + p = sol.prob.p + t = rand(size(p)...) + ground_solution = ground_analytic.(u0, p, t) + predict_sol = sol(t) + @test ground_solution≈predict_sol rtol=0.05 end #Test DeepONet @@ -85,13 +94,17 @@ end ground_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) p, t = get_trainset(deeponet, bounds, 50, tspan, 0.025) ground_solution = ground_analytic.(u0, p, vec(t)) - predict_sol = sol(t) predict_sol = sol.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 p, t = get_trainset(deeponet, bounds, 100, tspan, 0.01) ground_solution = ground_analytic.(u0, p, vec(t)) predict_sol = sol.interp(p, t) @test ground_solution≈predict_sol rtol=0.05 + + p, t = sol.prob.p, rand(1, 20, 1) + ground_solution = ground_analytic.(u0, p, vec(t)) + predict_sol = sol(t) + @test ground_solution≈predict_sol rtol=0.05 end @testitem "Example du = cos(p * t) + u" tags=[:pinoode] setup=[PINOODETestSetup] begin From 9c857a5112c1fea63d5501a4cc460a6ec6a9795c Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Wed, 30 Oct 2024 18:36:52 +0400 Subject: [PATCH 153/153] update --- src/pino_ode_solve.jl | 1 - test/PINO_ode_tests.jl | 32 +++++++++++++++++--------------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl index 5aff8284f1..04e65ee871 100644 --- a/src/pino_ode_solve.jl +++ b/src/pino_ode_solve.jl @@ -225,7 +225,6 @@ Interpolation of the solution of the ODE using a trained neural network. * `phi`: The neural network * `θ`: The parameters of the neural network. ``` - """ @concrete struct PINOODEInterpolation{T <: PINOPhi, T2} phi::T diff --git a/test/PINO_ode_tests.jl b/test/PINO_ode_tests.jl index 6e4545daf9..4f6e7e43f1 100644 --- a/test/PINO_ode_tests.jl +++ b/test/PINO_ode_tests.jl @@ -174,31 +174,31 @@ end end #multiple parameters Сhain -@testitem "Example multiple parameters Сhain du = p1 * cos(p2 * t) + p3" tags=[:pinoode] setup=[PINOODETestSetup] begin +@testitem "Example multiple parameters Сhain du = p1 * cos(p2 * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random - equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] + 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) - input_branch_size = 3 + input_branch_size = 2 chain = Chain( Dense(input_branch_size + 1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)) - x = rand(Float32, 4, 1000, 10) + x = rand(Float32, 3, 1000, 10) θ, st = Lux.setup(Random.default_rng(), chain) c = chain(x, θ, st)[1] - bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] + bounds = [(1.0, pi), (1.0, 2.0)]#, (2.0, 3.0)] number_of_parameters = 200 strategy = StochasticTraining(200) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) - sol = solve(prob, alg, verbose = false, maxiters = 4000) + sol = solve(prob, alg, verbose = false, maxiters = 5000) - ground_solution = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) + p[3] * t + 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, @@ -217,9 +217,9 @@ end end #multiple parameters DeepOnet -@testitem "Example multiple parameters DeepOnet du = p1 * cos(p2 * t) + p3" tags=[:pinoode] setup=[PINOODETestSetup] begin +@testitem "Example multiple parameters DeepOnet du = p1 * cos(p2 * t)" tags=[:pinoode] setup=[PINOODETestSetup] begin using NeuralPDE, Lux, OptimizationOptimisers, NeuralOperators, Random - equation = (u, p, t) -> p[1] * cos(p[2] * t) + p[3] + 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) @@ -231,18 +231,18 @@ end Chain(Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast))) - u = rand(3, 50) + u = rand(2, 50) v = rand(1, 40, 1) θ, st = Lux.setup(Random.default_rng(), deeponet) c = deeponet((u, v), θ, st)[1] bounds = [(1.0, pi), (1.0, 2.0), (2.0, 3.0)] - number_of_parameters = 50 + number_of_parameters = 100 strategy = StochasticTraining(50) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(deeponet, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 5000) - ground_solution = (u0, p, t) -> u0 + p[1] / p[2] * sin(p[2] * t) + p[3] * t + 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)]) @@ -270,10 +270,12 @@ end chain = Chain( Dense(input_branch_size + 1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), + Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 2)) + bounds = [(pi, 2pi)] - number_of_parameters = 300 - strategy = StochasticTraining(300) + number_of_parameters = 100 + strategy = StochasticTraining(100) opt = OptimizationOptimisers.Adam(0.01) alg = PINOODE(chain, opt, bounds, number_of_parameters; strategy = strategy) sol = solve(prob, alg, verbose = false, maxiters = 6000) @@ -305,5 +307,5 @@ end predict = sol.interp(p, t) @test ground_solution_[1, :, :]≈predict[1, :, :] rtol=0.05 @test ground_solution_[2, :, :]≈predict[2, :, :] rtol=0.05 - @test ground_solution_≈predict rtol=0.3 + @test ground_solution_≈predict rtol=0.05 end