From 15e61247ab598edab02bd37585320f34834829d8 Mon Sep 17 00:00:00 2001 From: KirillZubov Date: Mon, 22 Jul 2024 19:29:31 +0400 Subject: [PATCH] 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 df15c59bd..13b9b35db 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 a610d6e2e..a0b96ceeb 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 fdfbeaf30..96f189bdc 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 d999f62fe..d60ebc3a9 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