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(