test: further refactor tests and fix bugs

test/BPINN_PDE_tests.jl

@@ -148,7 +148,7 @@ end
     chain = Lux.Chain(Lux.Dense(dim, 9, Lux.σ), Lux.Dense(9, 9, Lux.σ), Lux.Dense(9, 1))
 
     # Discretization
-    dx = 0.05
+    dx = 0.04
     discretization = BayesianPINN([chain], GridTraining(dx))
 
     @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
@@ -166,7 +166,6 @@ end
 
     u_predict = pmean(sol1.ensemblesol[1])
     u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])]
-    diff_u = abs.(u_predict .- u_real)
     @test u_predict≈u_real atol=1.5
 end
 
@@ -189,7 +188,7 @@ end
     chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1))
 
     discretization = BayesianPINN([chain], GridTraining([0.01]))
-    @test discretization.chain isa Lux.AbstractExplicitLayer
+    @test discretization.chain[1] isa Lux.AbstractExplicitLayer
 
     @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
 

test/NNDAE_tests.jl

@@ -1,62 +1,64 @@
 using Test, Flux
 using Random, NeuralPDE
-using OrdinaryDiffEq, Optimisers, Statistics
+using OrdinaryDiffEq, Statistics
 import Lux, OptimizationOptimisers, OptimizationOptimJL
 
 Random.seed!(100)
 
-#Example 1
-function example1(du, u, p, t)
-    du[1] = cos(2pi * t)
-    du[2] = u[2] + cos(2pi * t)
-    nothing
+@testset "Example 1" begin
+    function example1(du, u, p, t)
+        du[1] = cos(2pi * t)
+        du[2] = u[2] + cos(2pi * t)
+        nothing
+    end
+    u₀ = [1.0, -1.0]
+    du₀ = [0.0, 0.0]
+    M = [1.0 0
+        0 0]
+    f = ODEFunction(example1, mass_matrix = M)
+    tspan = (0.0f0, 1.0f0)
+
+    prob_mm = ODEProblem(f, u₀, tspan)
+    ground_sol = solve(prob_mm, Rodas5(), reltol = 1e-8, abstol = 1e-8)
+
+    example = (du, u, p, t) -> [cos(2pi * t) - du[1], u[2] + cos(2pi * t) - du[2]]
+    differential_vars = [true, false]
+    prob = DAEProblem(example, du₀, u₀, tspan; differential_vars = differential_vars)
+    chain = Lux.Chain(Lux.Dense(1, 15, cos), Lux.Dense(15, 15, sin), Lux.Dense(15, 2))
+    opt = OptimizationOptimisers.Adam(0.1)
+    alg = NeuralPDE.NNDAE(chain, opt; autodiff = false)
+
+    sol = solve(prob,
+        alg, verbose = false, dt = 1 / 100.0f0,
+        maxiters = 3000, abstol = 1.0f-10)
+    @test ground_sol(0:(1 / 100):1)≈sol atol=0.4
 end
-u₀ = [1.0, -1.0]
-du₀ = [0.0, 0.0]
-M = [1.0 0
-    0 0]
-f = ODEFunction(example1, mass_matrix = M)
-tspan = (0.0f0, 1.0f0)
-
-prob_mm = ODEProblem(f, u₀, tspan)
-ground_sol = solve(prob_mm, Rodas5(), reltol = 1e-8, abstol = 1e-8)
-
-example = (du, u, p, t) -> [cos(2pi * t) - du[1], u[2] + cos(2pi * t) - du[2]]
-differential_vars = [true, false]
-prob = DAEProblem(example, du₀, u₀, tspan; differential_vars = differential_vars)
-chain = Flux.Chain(Dense(1, 15, cos), Dense(15, 15, sin), Dense(15, 2))
-opt = OptimizationOptimisers.Adam(0.1)
-alg = NeuralPDE.NNDAE(chain, opt; autodiff = false)
-
-sol = solve(prob,
-    alg, verbose = false, dt = 1 / 100.0f0,
-    maxiters = 3000, abstol = 1.0f-10)
-@test ground_sol(0:(1 / 100):1)≈sol atol=0.4
-
-#Example 2
-function example2(du, u, p, t)
-    du[1] = u[1] - t
-    du[2] = u[2] - t
-    nothing
+
+@testset "Example 2" begin
+    function example2(du, u, p, t)
+        du[1] = u[1] - t
+        du[2] = u[2] - t
+        nothing
+    end
+    M = [0.0 0
+        0 1]
+    u₀ = [0.0, 0.0]
+    du₀ = [0.0, 0.0]
+    tspan = (0.0f0, pi / 2.0f0)
+    f = ODEFunction(example2, mass_matrix = M)
+    prob_mm = ODEProblem(f, u₀, tspan)
+    ground_sol = solve(prob_mm, Rodas5(), reltol = 1e-8, abstol = 1e-8)
+
+    example = (du, u, p, t) -> [u[1] - t - du[1], u[2] - t - du[2]]
+    differential_vars = [false, true]
+    prob = DAEProblem(example, du₀, u₀, tspan; differential_vars = differential_vars)
+    chain = Lux.Chain(Lux.Dense(1, 15, Lux.σ), Lux.Dense(15, 2))
+    opt = OptimizationOptimisers.Adam(0.1)
+    alg = NNDAE(chain, OptimizationOptimisers.Adam(0.1); autodiff = false)
+
+    sol = solve(prob,
+        alg, verbose = false, dt = 1 / 100.0f0,
+        maxiters = 3000, abstol = 1.0f-10)
+
+    @test ground_sol(0:(1 / 100):(pi / 2))≈sol atol=0.4
 end
-M = [0.0 0
-    0 1]
-u₀ = [0.0, 0.0]
-du₀ = [0.0, 0.0]
-tspan = (0.0f0, pi / 2.0f0)
-f = ODEFunction(example2, mass_matrix = M)
-prob_mm = ODEProblem(f, u₀, tspan)
-ground_sol = solve(prob_mm, Rodas5(), reltol = 1e-8, abstol = 1e-8)
-
-example = (du, u, p, t) -> [u[1] - t - du[1], u[2] - t - du[2]]
-differential_vars = [false, true]
-prob = DAEProblem(example, du₀, u₀, tspan; differential_vars = differential_vars)
-chain = Flux.Chain(Dense(1, 15, σ), Dense(15, 2))
-opt = OptimizationOptimisers.Adam(0.1)
-alg = NNDAE(chain, OptimizationOptimisers.Adam(0.1); autodiff = false)
-
-sol = solve(prob,
-    alg, verbose = false, dt = 1 / 100.0f0,
-    maxiters = 3000, abstol = 1.0f-10)
-
-@test ground_sol(0:(1 / 100):(pi / 2))≈sol atol=0.4

test/additional_loss_tests.jl

@@ -43,8 +43,8 @@ using ComponentArrays
         function inner_f(x, θ)
             dx * phi(x, θ) .- 1
         end
-        prob = IntegralProblem(inner_f, lb, ub, θ)
-        norm2 = solve(prob, HCubatureJL(), reltol = 1e-8, abstol = 1e-8, maxiters = 10)
+        prob1 = IntegralProblem(inner_f, lb, ub, θ)
+        norm2 = solve(prob1, HCubatureJL(), reltol = 1e-8, abstol = 1e-8, maxiters = 10)
         abs(norm2[1])
     end
     discretization = PhysicsInformedNN(chain, GridTraining(dx); init_params = init_params,