update gpu test

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, θ)

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

test/PINO_ode_tests_gpu.jl

@@ -33,32 +33,79 @@ 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.σ),
         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)
+    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