remove dep NeuralOperator

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"]

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__)
 

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)

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_

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(