vector output

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"

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

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

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