diff --git a/test/BPINN_PDE_tests.jl b/test/BPINN_PDE_tests.jl
new file mode 100644
index 0000000000..f2f0f27aa4
--- /dev/null
+++ b/test/BPINN_PDE_tests.jl
@@ -0,0 +1,353 @@
+@testitem "BPINN PDE I: 2D Periodic System" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+
+    Random.seed!(100)
+
+    @parameters t
+    @variables u(..)
+    Dt = Differential(t)
+    eq = Dt(u(t)) - cospi(2t) ~ 0
+    bcs = [u(0.0) ~ 0.0]
+    domains = [t ∈ Interval(0.0, 2.0)]
+
+    chainl = Chain(Dense(1, 6, tanh), Dense(6, 1))
+    initl, st = Lux.setup(Random.default_rng(), chainl)
+    @named pde_system = PDESystem(eq, bcs, domains, [t], [u(t)])
+
+    # non adaptive case
+    discretization = BayesianPINN([chainl], GridTraining([0.01]))
+
+    sol1 = ahmc_bayesian_pinn_pde(
+        pde_system, discretization; draw_samples = 1500, bcstd = [0.02],
+        phystd = [0.01], priorsNNw = (0.0, 1.0), saveats = [1 / 50.0])
+
+    analytic_sol_func(u0, t) = u0 + sinpi(2t) / (2pi)
+    ts = vec(sol1.timepoints[1])
+    u_real = [analytic_sol_func(0.0, t) for t in ts]
+    u_predict = pmean(sol1.ensemblesol[1])
+
+    @test u_predict≈u_real atol=0.5
+    @test mean(u_predict .- u_real) < 0.1
+end
+
+@testitem "BPINN PDE II: 1D ODE" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+
+    Random.seed!(100)
+
+    @parameters θ
+    @variables u(..)
+    Dθ = Differential(θ)
+
+    # 1D ODE
+    eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) -
+                    u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3)))
+
+    # Initial and boundary conditions
+    bcs = [u(0.0) ~ 1.0f0]
+
+    # Space and time domains
+    domains = [θ ∈ Interval(0.0f0, 1.0f0)]
+
+    # Discretization
+    dt = 0.1f0
+
+    # Neural network
+    chain = Chain(Dense(1, 12, σ), Dense(12, 1))
+
+    discretization = BayesianPINN([chain], GridTraining([0.01]))
+
+    @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
+
+    sol1 = ahmc_bayesian_pinn_pde(
+        pde_system, discretization; draw_samples = 500, bcstd = [0.1],
+        phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0])
+
+    analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
+    ts = sol1.timepoints[1]
+    u_real = vec([analytic_sol_func(t) for t in ts])
+    u_predict = pmean(sol1.ensemblesol[1])
+    @test u_predict≈u_real atol=0.8
+end
+
+@testitem "BPINN PDE III: 3rd Degree ODE" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+
+    Random.seed!(100)
+
+    @parameters x
+    @variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..)
+    Dxxx = Differential(x)^3
+    Dx = Differential(x)
+
+    # ODE
+    eq = Dx(Dxxu(x)) ~ cospi(x)
+
+    # Initial and boundary conditions
+    ep = (cbrt(eps(eltype(Float64))))^2 / 6
+
+    bcs = [
+        u(0.0) ~ 0.0,
+        u(1.0) ~ cospi(1.0),
+        Dxu(1.0) ~ 1.0,
+        Dxu(x) ~ Dx(u(x)) + ep * O1(x),
+        Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x)
+    ]
+
+    # Space and time domains
+    domains = [x ∈ Interval(0.0, 1.0)]
+
+    # Neural network
+    chain = [
+        Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)),
+        Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)),
+        Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)),
+        Chain(Dense(1, 4, tanh), Dense(4, 1)),
+        Chain(Dense(1, 4, tanh), Dense(4, 1))
+    ]
+
+    discretization = BayesianPINN(chain, GridTraining(0.01))
+
+    @named pde_system = PDESystem(eq, bcs, domains, [x],
+        [u(x), Dxu(x), Dxxu(x), O1(x), O2(x)])
+
+    sol1 = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200,
+        bcstd = [0.01, 0.01, 0.01, 0.01, 0.01], phystd = [0.005],
+        priorsNNw = (0.0, 10.0), saveats = [1 / 100.0])
+
+    analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sinpi(x)) / (π^3)
+
+    u_predict = pmean(sol1.ensemblesol[1])
+    xs = vec(sol1.timepoints[1])
+    u_real = [analytic_sol_func(x) for x in xs]
+    @test u_predict≈u_real atol=0.5
+end
+
+@testitem "BPINN PDE IV: 2D Poisson" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+
+    Random.seed!(100)
+
+    @parameters x y
+    @variables u(..)
+    Dxx = Differential(x)^2
+    Dyy = Differential(y)^2
+
+    # 2D PDE
+    eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y)
+
+    # Boundary conditions
+    bcs = [
+        u(0, y) ~ 0.0,
+        u(1, y) ~ 0.0,
+        u(x, 0) ~ 0.0,
+        u(x, 1) ~ 0.0
+    ]
+
+    # Space and time domains
+    domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
+
+    # Discretization
+    dt = 0.1f0
+
+    # Neural network
+    chain = Chain(Dense(2, 9, σ), Dense(9, 9, σ), Dense(9, 1))
+
+    dx = 0.04
+    discretization = BayesianPINN(chain, GridTraining(dx))
+
+    @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
+
+    sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200,
+        bcstd = [0.003, 0.003, 0.003, 0.003], phystd = [0.003],
+        priorsNNw = (0.0, 10.0), saveats = [1 / 100.0, 1 / 100.0])
+
+    xs = sol.timepoints[1]
+    analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2)
+
+    u_predict = pmean(sol.ensemblesol[1])
+    u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])]
+    @test u_predict≈u_real rtol=0.5
+end
+
+@testitem "Translating from Flux" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+    import Flux
+
+    Random.seed!(100)
+
+    @parameters θ
+    @variables u(..)
+    Dθ = Differential(θ)
+
+    # 1D ODE
+    eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) -
+                    u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3)))
+
+    # Initial and boundary conditions
+    bcs = [u(0.0) ~ 1.0f0]
+
+    # Space and time domains
+    domains = [θ ∈ Interval(0.0f0, 1.0f0)]
+
+    # Neural network
+    chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1))
+
+    discretization = BayesianPINN([chain], GridTraining([0.01]))
+    @test discretization.chain[1] isa Lux.AbstractLuxLayer
+
+    @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
+
+    sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 500,
+        bcstd = [0.1], phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0])
+
+    analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
+    ts = sol.timepoints[1]
+    u_real = [analytic_sol_func(t) for t in ts]
+    u_predict = pmean(sol.ensemblesol[1])
+
+    @test u_predict≈u_real atol=0.8
+end
+
+@testitem "BPINN PDE Inv I: 1D Periodic System" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+
+    Random.seed!(100)
+
+    @parameters t p
+    @variables u(..)
+
+    Dt = Differential(t)
+    eqs = Dt(u(t)) - cos(p * t) ~ 0
+    bcs = [u(0) ~ 0.0]
+    domains = [t ∈ Interval(0.0, 2.0)]
+
+    chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 1))
+    initl, st = Lux.setup(Random.default_rng(), chainl)
+
+    @named pde_system = PDESystem(eqs,
+        bcs,
+        domains,
+        [t],
+        [u(t)],
+        [p],
+        defaults = Dict([p => 4.0]))
+
+    analytic_sol_func1(u0, t) = u0 + sinpi(2t) / (2π)
+    timepoints = collect(0.0:(1 / 100.0):2.0)
+    u = [analytic_sol_func1(0.0, timepoint) for timepoint in timepoints]
+    u = u .+ (u .* 0.2) .* randn(size(u))
+    dataset = [hcat(u, timepoints)]
+
+    @testset "$(nameof(typeof(strategy)))" for strategy in [
+        StochasticTraining(200),
+        QuasiRandomTraining(200),
+        GridTraining([0.02])
+    ]
+        discretization = BayesianPINN([chainl], strategy; param_estim = true,
+            dataset = [dataset, nothing])
+
+        sol1 = ahmc_bayesian_pinn_pde(pde_system,
+            discretization;
+            draw_samples = 1500,
+            bcstd = [0.05],
+            phystd = [0.01], l2std = [0.01],
+            priorsNNw = (0.0, 1.0),
+            saveats = [1 / 50.0],
+            param = [LogNormal(6.0, 0.5)])
+
+        param = 2 * π
+        ts = vec(sol1.timepoints[1])
+        u_real = [analytic_sol_func1(0.0, t) for t in ts]
+        u_predict = pmean(sol1.ensemblesol[1])
+
+        @test u_predict≈u_real atol=1.5
+        @test mean(u_predict .- u_real) < 0.1
+        @test sol1.estimated_de_params[1]≈param atol=param * 0.3
+    end
+end
+
+@testitem "BPINN PDE Inv II: Lorenz System" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+
+    Random.seed!(100)
+
+    @parameters t, σ_
+    @variables x(..), y(..), z(..)
+    Dt = Differential(t)
+    eqs = [
+        Dt(x(t)) ~ σ_ * (y(t) - x(t)),
+        Dt(y(t)) ~ x(t) * (28.0 - z(t)) - y(t),
+        Dt(z(t)) ~ x(t) * y(t) - 8.0 / 3.0 * z(t)
+    ]
+
+    bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0]
+    domains = [t ∈ Interval(0.0, 1.0)]
+
+    input_ = length(domains)
+    n = 7
+    chain = [
+        Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)),
+        Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)),
+        Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1))
+    ]
+
+    # Generate Data
+    function lorenz!(du, u, p, t)
+        du[1] = 10.0 * (u[2] - u[1])
+        du[2] = u[1] * (28.0 - u[3]) - u[2]
+        du[3] = u[1] * u[2] - (8.0 / 3.0) * u[3]
+    end
+
+    u0 = [1.0; 0.0; 0.0]
+    tspan = (0.0, 1.0)
+    prob = ODEProblem(lorenz!, u0, tspan)
+    sol = solve(prob, Tsit5(), dt = 0.01, saveat = 0.05)
+    ts = sol.t
+    us = hcat(sol.u...)
+    us = us .+ ((0.05 .* randn(size(us))) .* us)
+    ts_ = hcat(sol(ts).t...)[1, :]
+    dataset = [hcat(us[i, :], ts_) for i in 1:3]
+
+    discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true,
+        dataset = [dataset, nothing])
+
+    @named pde_system = PDESystem(eqs, bcs, domains,
+        [t], [x(t), y(t), z(t)], [σ_], defaults = Dict([p => 1.0 for p in [σ_]]))
+
+    sol1 = ahmc_bayesian_pinn_pde(pde_system,
+        discretization;
+        draw_samples = 50,
+        bcstd = [0.3, 0.3, 0.3],
+        phystd = [0.1, 0.1, 0.1],
+        l2std = [1, 1, 1],
+        priorsNNw = (0.0, 1.0),
+        saveats = [0.01],
+        param = [Normal(12.0, 2)])
+
+    idealp = 10.0
+    p_ = sol1.estimated_de_params[1]
+    @test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1]
+    # @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2]
+end
diff --git a/test/BPINN_PDE_tests_wip.jl b/test/BPINN_PDE_tests_wip.jl
deleted file mode 100644
index cbb8ffa46c..0000000000
--- a/test/BPINN_PDE_tests_wip.jl
+++ /dev/null
@@ -1,187 +0,0 @@
-using Test, MCMCChains, Lux, ModelingToolkit, ForwardDiff, Distributions, OrdinaryDiffEq,
-      AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
-      ComponentArrays
-import ModelingToolkit: Interval, infimum, supremum
-import Flux
-
-Random.seed!(100)
-
-@testset "Example 1: 2D Periodic System" begin
-    # Cos(pi*t) example
-    @parameters t
-    @variables u(..)
-    Dt = Differential(t)
-    eqs = Dt(u(t)) - cos(2 * π * t) ~ 0
-    bcs = [u(0) ~ 0.0]
-    domains = [t ∈ Interval(0.0, 2.0)]
-    chainl = Chain(Dense(1, 6, tanh), Dense(6, 1))
-    initl, st = Lux.setup(Random.default_rng(), chainl)
-    @named pde_system = PDESystem(eqs, bcs, domains, [t], [u(t)])
-
-    # non adaptive case
-    discretization = BayesianPINN([chainl], GridTraining([0.01]))
-
-    sol1 = ahmc_bayesian_pinn_pde(
-        pde_system, discretization; draw_samples = 1500, bcstd = [0.02],
-        phystd = [0.01], priorsNNw = (0.0, 1.0), saveats = [1 / 50.0])
-
-    analytic_sol_func(u0, t) = u0 + sin(2 * π * t) / (2 * π)
-    ts = vec(sol1.timepoints[1])
-    u_real = [analytic_sol_func(0.0, t) for t in ts]
-    u_predict = pmean(sol1.ensemblesol[1])
-    @test u_predict≈u_real atol=0.5
-    @test mean(u_predict .- u_real) < 0.1
-end
-
-@testset "Example 2: 1D ODE" begin
-    @parameters θ
-    @variables u(..)
-    Dθ = Differential(θ)
-
-    # 1D ODE
-    eq = Dθ(u(θ)) ~ θ^3 + 2 * θ + (θ^2) * ((1 + 3 * (θ^2)) / (1 + θ + (θ^3))) -
-                    u(θ) * (θ + ((1 + 3 * (θ^2)) / (1 + θ + θ^3)))
-
-    # Initial and boundary conditions
-    bcs = [u(0.0) ~ 1.0]
-
-    # Space and time domains
-    domains = [θ ∈ Interval(0.0, 1.0)]
-
-    # Neural network
-    chain = Chain(Dense(1, 12, σ), Dense(12, 1))
-
-    discretization = BayesianPINN([chain], GridTraining([0.01]))
-
-    @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
-
-    sol1 = ahmc_bayesian_pinn_pde(
-        pde_system, discretization; draw_samples = 500, bcstd = [0.1],
-        phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0])
-
-    analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
-    ts = sol1.timepoints[1]
-    u_real = vec([analytic_sol_func(t) for t in ts])
-    u_predict = pmean(sol1.ensemblesol[1])
-    @test u_predict≈u_real atol=0.8
-end
-
-@testset "Example 3: 3rd Degree ODE" begin
-    @parameters x
-    @variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..)
-    Dxxx = Differential(x)^3
-    Dx = Differential(x)
-
-    # ODE
-    eq = Dx(Dxxu(x)) ~ cos(pi * x)
-
-    # Initial and boundary conditions
-    ep = (cbrt(eps(eltype(Float64))))^2 / 6
-
-    bcs = [u(0.0) ~ 0.0,
-        u(1.0) ~ cos(pi),
-        Dxu(1.0) ~ 1.0,
-        Dxu(x) ~ Dx(u(x)) + ep * O1(x),
-        Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x)]
-
-    # Space and time domains
-    domains = [x ∈ Interval(0.0, 1.0)]
-
-    # Neural network
-    chain = [
-        Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)),
-        Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)),
-        Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)),
-        Chain(Dense(1, 4, tanh), Dense(4, 1)),
-        Chain(Dense(1, 4, tanh), Dense(4, 1))
-    ]
-
-    discretization = BayesianPINN(chain, GridTraining(0.01))
-
-    @named pde_system = PDESystem(eq, bcs, domains, [x],
-        [u(x), Dxu(x), Dxxu(x), O1(x), O2(x)])
-
-    sol1 = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200,
-        bcstd = [0.01, 0.01, 0.01, 0.01, 0.01], phystd = [0.005],
-        priorsNNw = (0.0, 10.0), saveats = [1 / 100.0])
-
-    analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sin(π * x)) / (π^3)
-
-    u_predict = pmean(sol1.ensemblesol[1])
-    xs = vec(sol1.timepoints[1])
-    u_real = [analytic_sol_func(x) for x in xs]
-    @test u_predict≈u_real atol=0.5
-end
-
-@testset "Example 4: 2D Poissons equation" begin
-    @parameters x y
-    @variables u(..)
-    Dxx = Differential(x)^2
-    Dyy = Differential(y)^2
-
-    # 2D PDE
-    eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y)
-
-    # Boundary conditions
-    bcs = [u(0, y) ~ 0.0, u(1, y) ~ 0.0,
-        u(x, 0) ~ 0.0, u(x, 1) ~ 0.0]
-
-    # Space and time domains
-    domains = [x ∈ Interval(0.0, 1.0),
-        y ∈ Interval(0.0, 1.0)]
-
-    # Neural network
-    dim = 2 # number of dimensions
-    chain = Chain(Dense(dim, 9, σ), Dense(9, 9, σ), Dense(9, 1))
-
-    # Discretization
-    dx = 0.04
-    discretization = BayesianPINN([chain], GridTraining(dx))
-
-    @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
-
-    sol1 = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200,
-        bcstd = [0.003, 0.003, 0.003, 0.003], phystd = [0.003],
-        priorsNNw = (0.0, 10.0), saveats = [1 / 100.0, 1 / 100.0])
-
-    xs = sol1.timepoints[1]
-    analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2)
-
-    u_predict = pmean(sol1.ensemblesol[1])
-    u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])]
-    @test u_predict≈u_real atol=1.5
-end
-
-@testset "Translating from Flux" begin
-    @parameters θ
-    @variables u(..)
-    Dθ = Differential(θ)
-
-    # 1D ODE
-    eq = Dθ(u(θ)) ~ θ^3 + 2 * θ + (θ^2) * ((1 + 3 * (θ^2)) / (1 + θ + (θ^3))) -
-                    u(θ) * (θ + ((1 + 3 * (θ^2)) / (1 + θ + θ^3)))
-
-    # Initial and boundary conditions
-    bcs = [u(0.0) ~ 1.0]
-
-    # Space and time domains
-    domains = [θ ∈ Interval(0.0, 1.0)]
-
-    # Neural network
-    chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1))
-
-    discretization = BayesianPINN([chain], GridTraining([0.01]))
-    @test discretization.chain[1] isa AbstractLuxLayer
-
-    @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
-
-    sol1 = ahmc_bayesian_pinn_pde(
-        pde_system, discretization; draw_samples = 500, bcstd = [0.1],
-        phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0])
-
-    analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
-    ts = sol1.timepoints[1]
-    u_real = vec([analytic_sol_func(t) for t in ts])
-    u_predict = pmean(sol1.ensemblesol[1])
-    @test u_predict≈u_real atol=0.8
-end
diff --git a/test/BPINN_PDEinvsol_tests_wip.jl b/test/BPINN_PDEinvsol_tests_wip.jl
deleted file mode 100644
index fd64e177da..0000000000
--- a/test/BPINN_PDEinvsol_tests_wip.jl
+++ /dev/null
@@ -1,143 +0,0 @@
-using Test, MCMCChains, Lux, ModelingToolkit, ForwardDiff, Distributions, OrdinaryDiffEq,
-      AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
-      ComponentArrays
-import ModelingToolkit: Interval, infimum, supremum
-
-Random.seed!(100)
-
-@testset "Example 1: 1D Periodic System with parameter estimation" begin
-    # Cos(pi*t) periodic curve
-    @parameters t, p
-    @variables u(..)
-
-    Dt = Differential(t)
-    eqs = Dt(u(t)) - cos(p * t) ~ 0
-    bcs = [u(0) ~ 0.0]
-    domains = [t ∈ Interval(0.0, 2.0)]
-
-    chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 1))
-    initl, st = Lux.setup(Random.default_rng(), chainl)
-
-    @named pde_system = PDESystem(eqs,
-        bcs,
-        domains,
-        [t],
-        [u(t)],
-        [p],
-        defaults = Dict([p => 4.0]))
-
-    analytic_sol_func1(u0, t) = u0 + sin(2 * π * t) / (2 * π)
-    timepoints = collect(0.0:(1 / 100.0):2.0)
-    u = [analytic_sol_func1(0.0, timepoint) for timepoint in timepoints]
-    u = u .+ (u .* 0.2) .* randn(size(u))
-    dataset = [hcat(u, timepoints)]
-
-    # checking all training strategies
-    discretization = BayesianPINN([chainl], StochasticTraining(200), param_estim = true,
-        dataset = [dataset, nothing])
-
-    ahmc_bayesian_pinn_pde(pde_system,
-        discretization;
-        draw_samples = 1500,
-        bcstd = [0.05],
-        phystd = [0.01], l2std = [0.01],
-        priorsNNw = (0.0, 1.0),
-        saveats = [1 / 50.0],
-        param = [LogNormal(6.0, 0.5)])
-
-    discretization = BayesianPINN([chainl], QuasiRandomTraining(200), param_estim = true,
-        dataset = [dataset, nothing])
-
-    ahmc_bayesian_pinn_pde(pde_system,
-        discretization;
-        draw_samples = 1500,
-        bcstd = [0.05],
-        phystd = [0.01], l2std = [0.01],
-        priorsNNw = (0.0, 1.0),
-        saveats = [1 / 50.0],
-        param = [LogNormal(6.0, 0.5)])
-
-    # alternative to QuadratureTraining [WIP]
-
-    discretization = BayesianPINN([chainl], GridTraining([0.02]), param_estim = true,
-        dataset = [dataset, nothing])
-
-    sol1 = ahmc_bayesian_pinn_pde(pde_system,
-        discretization;
-        draw_samples = 1500,
-        bcstd = [0.05],
-        phystd = [0.01], l2std = [0.01],
-        priorsNNw = (0.0, 1.0),
-        saveats = [1 / 50.0],
-        param = [LogNormal(6.0, 0.5)])
-
-    param = 2 * π
-    ts = vec(sol1.timepoints[1])
-    u_real = [analytic_sol_func1(0.0, t) for t in ts]
-    u_predict = pmean(sol1.ensemblesol[1])
-
-    @test u_predict≈u_real atol=1.5
-    @test mean(u_predict .- u_real) < 0.1
-    @test sol1.estimated_de_params[1]≈param atol=param * 0.3
-end
-
-@testset "Example 2: Lorenz System with parameter estimation" begin
-    @parameters t, σ_
-    @variables x(..), y(..), z(..)
-    Dt = Differential(t)
-    eqs = [Dt(x(t)) ~ σ_ * (y(t) - x(t)),
-        Dt(y(t)) ~ x(t) * (28.0 - z(t)) - y(t),
-        Dt(z(t)) ~ x(t) * y(t) - 8 / 3 * z(t)]
-
-    bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0]
-    domains = [t ∈ Interval(0.0, 1.0)]
-
-    input_ = length(domains)
-    n = 7
-    chain = [
-        Lux.Chain(Lux.Dense(input_, n, Lux.tanh), Lux.Dense(n, n, Lux.tanh),
-            Lux.Dense(n, 1)),
-        Lux.Chain(Lux.Dense(input_, n, Lux.tanh), Lux.Dense(n, n, Lux.tanh),
-            Lux.Dense(n, 1)),
-        Lux.Chain(Lux.Dense(input_, n, Lux.tanh), Lux.Dense(n, n, Lux.tanh),
-            Lux.Dense(n, 1))
-    ]
-
-    #Generate Data
-    function lorenz!(du, u, p, t)
-        du[1] = 10.0 * (u[2] - u[1])
-        du[2] = u[1] * (28.0 - u[3]) - u[2]
-        du[3] = u[1] * u[2] - (8 / 3) * u[3]
-    end
-
-    u0 = [1.0; 0.0; 0.0]
-    tspan = (0.0, 1.0)
-    prob = ODEProblem(lorenz!, u0, tspan)
-    sol = solve(prob, Tsit5(), dt = 0.01, saveat = 0.05)
-    ts = sol.t
-    us = hcat(sol.u...)
-    us = us .+ ((0.05 .* randn(size(us))) .* us)
-    ts_ = hcat(sol(ts).t...)[1, :]
-    dataset = [hcat(us[i, :], ts_) for i in 1:3]
-
-    discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true,
-        dataset = [dataset, nothing])
-
-    @named pde_system = PDESystem(eqs, bcs, domains,
-        [t], [x(t), y(t), z(t)], [σ_], defaults = Dict([p => 1.0 for p in [σ_]]))
-
-    sol1 = ahmc_bayesian_pinn_pde(pde_system,
-        discretization;
-        draw_samples = 50,
-        bcstd = [0.3, 0.3, 0.3],
-        phystd = [0.1, 0.1, 0.1],
-        l2std = [1, 1, 1],
-        priorsNNw = (0.0, 1.0),
-        saveats = [0.01],
-        param = [Normal(12.0, 2)])
-
-    idealp = 10.0
-    p_ = sol1.estimated_de_params[1]
-    @test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1]
-    # @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2]
-end
diff --git a/test/NNPDE_cuda_tests.jl b/test/NNPDE_cuda_tests.jl
index be72b79ced..8de4163c5e 100644
--- a/test/NNPDE_cuda_tests.jl
+++ b/test/NNPDE_cuda_tests.jl
@@ -16,7 +16,7 @@ export gpud, callback
 end
 
 @testitem "1D ODE - CUDA" tags=[:cuda] setup=[CUDATestSetup] begin
-    using Lux, Optimization, OptimizationOptimisers, Random
+    using Lux, Optimization, OptimizationOptimisers, Random, ComponentArrays
     import ModelingToolkit: Interval, infimum, supremum
 
     Random.seed!(100)
@@ -60,7 +60,7 @@ end
 end
 
 @testitem "1D PDE Dirichlet BC - CUDA" tags=[:cuda] setup=[CUDATestSetup] begin
-    using Lux, Optimization, OptimizationOptimisers, Random
+    using Lux, Optimization, OptimizationOptimisers, Random, ComponentArrays
     import ModelingToolkit: Interval, infimum, supremum
 
     Random.seed!(100)
@@ -106,7 +106,8 @@ end
 end
 
 @testitem "1D PDE Neumann BC - CUDA" tags=[:cuda] setup=[CUDATestSetup] begin
-    using Lux, Optimization, OptimizationOptimisers, Random, QuasiMonteCarlo
+    using Lux, Optimization, OptimizationOptimisers, Random, QuasiMonteCarlo,
+          ComponentArrays
     import ModelingToolkit: Interval, infimum, supremum
 
     Random.seed!(100)
@@ -156,7 +157,7 @@ end
 end
 
 @testitem "2D PDE - CUDA" tags=[:cuda] setup=[CUDATestSetup] begin
-    using Lux, Optimization, OptimizationOptimisers, Random
+    using Lux, Optimization, OptimizationOptimisers, Random, ComponentArrays
     import ModelingToolkit: Interval, infimum, supremum
 
     Random.seed!(100)
diff --git a/test/NNPDE_tests_wip.jl b/test/NNPDE_tests.jl
similarity index 60%
rename from test/NNPDE_tests_wip.jl
rename to test/NNPDE_tests.jl
index 888179b561..82d503623d 100644
--- a/test/NNPDE_tests_wip.jl
+++ b/test/NNPDE_tests.jl
@@ -1,51 +1,14 @@
-using NeuralPDE, Test, Optimization, OptimizationOptimJL, OptimizationOptimisers, Integrals,
-      Cubature, QuasiMonteCarlo, DomainSets, Lux, LineSearches, Random
-import ModelingToolkit: Interval, infimum, supremum
-import Flux
+@testsetup module NNPDE1TestSetup
 
-Random.seed!(100)
+using NeuralPDE, Cubature, Integrals, QuasiMonteCarlo
 
-callback = function (p, l)
-    println("Current loss is: $l")
+function callback(p, l)
+    if p.iter == 1 || p.iter % 250 == 0
+        println("Current loss is: $l after $(p.iter) iterations")
+    end
     return false
 end
 
-function test_ode(strategy_)
-    println("Example 1, 1D ode: strategy: $(nameof(typeof(strategy_)))")
-    @parameters θ
-    @variables u(..)
-    Dθ = Differential(θ)
-
-    # 1D ODE
-    eq = Dθ(u(θ)) ~ θ^3 + 2 * θ + (θ^2) * ((1 + 3 * (θ^2)) / (1 + θ + (θ^3))) -
-                    u(θ) * (θ + ((1 + 3 * (θ^2)) / (1 + θ + θ^3)))
-
-    # Initial and boundary conditions
-    bcs = [u(0.0) ~ 1.0]
-
-    # Space and time domains
-    domains = [θ ∈ Interval(0.0, 1.0)]
-
-    # Neural network
-    chain = Chain(Dense(1, 12, σ), Dense(12, 1))
-
-    discretization = PhysicsInformedNN(chain, strategy_)
-    @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
-    prob = discretize(pde_system, discretization)
-
-    res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.1); maxiters = 1000)
-    prob = remake(prob, u0 = res.u)
-    res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.01); maxiters = 500)
-    prob = remake(prob, u0 = res.u)
-    res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.001); maxiters = 500)
-    phi = discretization.phi
-    analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
-    ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
-    u_real = [analytic_sol_func(t) for t in ts]
-    u_predict = [first(phi(t, res.u)) for t in ts]
-    @test u_predict≈u_real atol=0.1
-end
-
 grid_strategy = GridTraining(0.1)
 quadrature_strategy = QuadratureTraining(quadrature_alg = CubatureJLh(),
     reltol = 1e3, abstol = 1e-3, maxiters = 50, batch = 100)
@@ -63,13 +26,60 @@ strategies = [
     quadrature_strategy
 ]
 
-@testset "Test ODE/Heterogeneous" begin
+export callback, strategies
+
+end
+
+@testitem "Test Heterogeneous ODE" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using Cubature, Integrals, QuasiMonteCarlo, DomainSets, Lux, Random, Optimisers
+
+    function simple_1d_ode(strategy)
+        @parameters θ
+        @variables u(..)
+        Dθ = Differential(θ)
+
+        # 1D ODE
+        eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ +
+                        (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) -
+                        u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3)))
+
+        # Initial and boundary conditions
+        bcs = [u(0.0) ~ 1.0f0]
+
+        # Space and time domains
+        domains = [θ ∈ Interval(0.0f0, 1.0f0)]
+
+        # Neural network
+        chain = Chain(Dense(1, 12, σ), Dense(12, 1))
+
+        discretization = PhysicsInformedNN(chain, strategy)
+        @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
+        prob = discretize(pde_system, discretization)
+
+        res = solve(prob, Adam(0.1); maxiters = 1000)
+        prob = remake(prob, u0 = res.u)
+        res = solve(prob, Adam(0.01); maxiters = 500)
+        prob = remake(prob, u0 = res.u)
+        res = solve(prob, Adam(0.001); maxiters = 500)
+        phi = discretization.phi
+
+        analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
+        ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
+        u_real = [analytic_sol_func(t) for t in ts]
+        u_predict = [first(phi([t], res.u)) for t in ts]
+        @test u_predict≈u_real atol=0.8
+    end
+
     @testset "$(nameof(typeof(strategy)))" for strategy in strategies
-        test_ode(strategy)
+        simple_1d_ode(strategy)
     end
 end
 
-@testset "Example 1: Heterogeneous system" begin
+@testitem "PDE I: Heterogeneous system" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using DomainSets, Lux, Random, Optimisers, Integrals
+    import ModelingToolkit: Interval, infimum, supremum
+    import OptimizationOptimJL: BFGS
+
     @parameters x, y, z
     @variables u(..), v(..), h(..), p(..)
     Dz = Differential(z)
@@ -82,7 +92,7 @@ end
                                                     exp(x) * exp(z)
     ]
 
-    bcs = [u(0, 0, 0) ~ 0.0]
+    bcs = [u(0.0, 0.0, 0.0) ~ 0.0]
 
     domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0), z ∈ Interval(0.0, 1.0)]
 
@@ -104,12 +114,7 @@ end
 
     prob = discretize(pde_system, discretization)
 
-    callback = function (p, l)
-        println("Current loss is: $l")
-        return false
-    end
-
-    res = Optimization.solve(prob, OptimizationOptimJL.BFGS(); maxiters = 2000)
+    res = solve(prob, BFGS(); maxiters = 2000, callback)
 
     phi = discretization.phi
 
@@ -134,82 +139,98 @@ end
     v_predict = [phi[2]([y, x], res.u.depvar.v)[1] for y in ys for x in xs]
     h_predict = [phi[3]([z], res.u.depvar.h)[1] for z in zs]
     p_predict = [phi[4]([x, z], res.u.depvar.p)[1] for x in xs for z in zs]
+
     predict = [u_predict, v_predict, h_predict, p_predict]
+
     for i in 1:4
         @test predict[i]≈real_[i] rtol=10^-2
     end
 end
 
-function test_2d_poisson_equation(chain_, strategy_)
-    println("Example 2, 2D Poisson equation, chain: $(nameof(typeof(chain_))), strategy: $(nameof(typeof(strategy_)))")
-    @parameters x y
-    @variables u(..)
-    Dxx = Differential(x)^2
-    Dyy = Differential(y)^2
+@testitem "PDE II: 2D Poisson" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using Lux, Random, Optimisers, DomainSets, Cubature, QuasiMonteCarlo, Integrals
+    import ModelingToolkit: Interval, infimum, supremum
+    import OptimizationOptimJL: BFGS
 
-    # 2D PDE
-    eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y)
+    function test_2d_poisson_equation(chain, strategy)
+        @parameters x y
+        @variables u(..)
+        Dxx = Differential(x)^2
+        Dyy = Differential(y)^2
 
-    # Initial and boundary conditions
-    bcs = [u(0, y) ~ 0.0, u(1, y) ~ -sin(pi * 1) * sin(pi * y),
-        u(x, 0) ~ 0.0, u(x, 1) ~ -sin(pi * x) * sin(pi * 1)]
-    # Space and time domains
-    domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
-    ps = Lux.setup(Random.default_rng(), chain_)[1]
-    discretization = PhysicsInformedNN(chain_, strategy_; init_params = ps)
-    @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
-    prob = discretize(pde_system, discretization)
-    res = solve(prob, OptimizationOptimisers.Adam(0.1); maxiters = 500, cb = callback)
-    phi = discretization.phi
+        # 2D PDE
+        eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y)
 
-    xs, ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
-    analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2)
+        # Boundary conditions
+        bcs = [
+            u(0, y) ~ 0.0,
+            u(1, y) ~ 0.0,
+            u(x, 0) ~ 0.0,
+            u(x, 1) ~ 0.0
+        ]
 
-    u_predict = reshape([first(phi([x, y], res.u)) for x in xs for y in ys],
-        (length(xs), length(ys)))
-    u_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys],
-        (length(xs), length(ys)))
-    @test u_predict≈u_real atol=2.0
-end
+        # Space and time domains
+        domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
+
+        ps = Lux.initialparameters(Random.default_rng(), chain)
+
+        discretization = PhysicsInformedNN(chain, strategy; init_params = ps)
+        @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
+        prob = discretize(pde_system, discretization)
+        res = solve(prob, BFGS(); maxiters = 500, callback)
+        phi = discretization.phi
+
+        xs, ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
+        analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2)
+
+        u_predict = [first(phi([x, y], res.u)) for x in xs for y in ys]
+        u_real = [analytic_sol_func(x, y) for x in xs for y in ys]
+
+        @test u_predict≈u_real atol=2.0
+    end
 
-@testset "Example 2, 2D Poisson equation" begin
-    grid_strategy = GridTraining(0.1)
     chain = Chain(Dense(2, 12, σ), Dense(12, 12, σ), Dense(12, 1))
-    test_2d_poisson_equation(chain, grid_strategy)
 
     @testset "$(nameof(typeof(strategy)))" for strategy in strategies
-        chain_ = Chain(Dense(2, 12, σ), Dense(12, 12, σ), Dense(12, 1))
-        test_2d_poisson_equation(chain_, strategy)
+        test_2d_poisson_equation(chain, strategy)
     end
 
     algs = [CubatureJLp()]
     @testset "$(nameof(typeof(alg)))" for alg in algs
-        chain_ = Chain(Dense(2, 12, σ), Dense(12, 12, σ), Dense(12, 1))
-        strategy_ = NeuralPDE.QuadratureTraining(quadrature_alg = alg, reltol = 1e-4,
+        strategy = QuadratureTraining(quadrature_alg = alg, reltol = 1e-4,
             abstol = 1e-3, maxiters = 30, batch = 10)
-        test_2d_poisson_equation(chain_, strategy_)
+        test_2d_poisson_equation(chain, strategy)
     end
 end
 
-@testset "Example 3, 3rd-order ode" begin
+@testitem "PDE III: 3rd-order ODE" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using Lux, Random, Optimisers, DomainSets, Cubature, QuasiMonteCarlo, Integrals
+    import ModelingToolkit: Interval, infimum, supremum
+    import OptimizationOptimJL: BFGS
+
     @parameters x
     @variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..)
     Dxxx = Differential(x)^3
     Dx = Differential(x)
 
     # ODE
-    eq = Dx(Dxxu(x)) ~ cos(pi * x)
+    eq = Dx(Dxxu(x)) ~ cospi(x)
 
     # Initial and boundary conditions
-    bcs_ = [u(0.0) ~ 0.0,
-        u(1.0) ~ cos(pi),
-        Dxu(1.0) ~ 1.0]
+    bcs_ = [
+        u(0.0) ~ 0.0,
+        u(1.0) ~ cospi(1.0),
+        Dxu(1.0) ~ 1.0
+    ]
     ep = (cbrt(eps(eltype(Float64))))^2 / 6
 
-    der = [Dxu(x) ~ Dx(u(x)) + ep * O1(x),
-        Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x)]
+    der = [
+        Dxu(x) ~ Dx(u(x)) + ep * O1(x),
+        Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x)
+    ]
 
     bcs = [bcs_; der]
+
     # Space and time domains
     domains = [x ∈ Interval(0.0, 1.0)]
 
@@ -229,14 +250,16 @@ end
     pde_inner_loss_functions = sym_prob.loss_functions.pde_loss_functions
     bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
-    cb_ = function (p, l)
-        println("loss: ", l)
-        println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
-        println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
+    callback = function (p, l)
+        if p.iter % 100 == 0 || p.iter == 1
+            println("loss: ", l)
+            println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+            println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
+        end
         return false
     end
 
-    res = solve(prob, OptimizationOptimJL.BFGS(); maxiters = 1000)
+    res = solve(prob, BFGS(); maxiters = 1000, callback)
     phi = discretization.phi[1]
 
     analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sin(π * x)) / (π^3)
@@ -248,19 +271,27 @@ end
     @test u_predict≈u_real atol=10^-4
 end
 
-@testset "Example 4, system of pde" begin
+@testitem "PDE IV: System of PDEs" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using Lux, Random, Optimisers, DomainSets, Cubature, QuasiMonteCarlo, Integrals
+    import ModelingToolkit: Interval, infimum, supremum
+    import OptimizationOptimJL: BFGS
+
     @parameters x, y
     @variables u1(..), u2(..)
     Dx = Differential(x)
     Dy = Differential(y)
 
     # System of pde
-    eqs = [Dx(u1(x, y)) + 4 * Dy(u2(x, y)) ~ 0,
-        Dx(u2(x, y)) + 9 * Dy(u1(x, y)) ~ 0]
-    # 3*u1(x,0) ~ 2*u2(x,0)]
+    eqs = [
+        Dx(u1(x, y)) + 4 * Dy(u2(x, y)) ~ 0,
+        Dx(u2(x, y)) + 9 * Dy(u1(x, y)) ~ 0
+    ]
 
     # Initial and boundary conditions
-    bcs = [u1(x, 0) ~ 2 * x, u2(x, 0) ~ 3 * x]
+    bcs = [
+        u1(x, 0) ~ 2 * x,
+        u2(x, 0) ~ 3 * x
+    ]
 
     # Space and time domains
     domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
@@ -270,8 +301,7 @@ end
     chain2 = Chain(Dense(2, 15, tanh), Dense(15, 1))
 
     quadrature_strategy = QuadratureTraining(quadrature_alg = CubatureJLh(),
-        reltol = 1e-3, abstol = 1e-3,
-        maxiters = 50, batch = 100)
+        reltol = 1e-3, abstol = 1e-3, maxiters = 50, batch = 100)
     chain = [chain1, chain2]
 
     discretization = PhysicsInformedNN(chain, quadrature_strategy)
@@ -280,7 +310,7 @@ end
 
     prob = discretize(pde_system, discretization)
 
-    res = solve(prob, OptimizationOptimJL.BFGS(); maxiters = 1000)
+    res = solve(prob, BFGS(); maxiters = 1000, callback)
     phi = discretization.phi
 
     analytic_sol_func(x, y) = [1 / 3 * (6x - y), 1 / 2 * (6x - y)]
@@ -295,8 +325,12 @@ end
     @test u_predict[2]≈u_real[2] atol=0.1
 end
 
-@testset "Example 5, 2d wave equation, neumann boundary condition" begin
-    # here we use low level api for build solution
+@testitem "PDE V: 2D Wave Equation" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using Lux, Random, Optimisers, DomainSets, Cubature, QuasiMonteCarlo, Integrals,
+          LineSearhces, Integrals
+    import ModelingToolkit: Interval, infimum, supremum
+    import OptimizationOptimJL: BFGS
+
     @parameters x, t
     @variables u(..)
     Dxx = Differential(x)^2
@@ -337,7 +371,7 @@ end
         return false
     end
 
-    res = solve(prob, OptimizationOptimJL.BFGS(linesearch = BackTracking()); maxiters = 500)
+    res = solve(prob, BFGS(linesearch = BackTracking()); maxiters = 500, callback)
 
     dx = 0.1
     xs, ts = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
@@ -352,7 +386,11 @@ end
     @test u_predict≈u_real atol=0.1
 end
 
-@testset "Example 6, pde with mixed derivative" begin
+@testitem "PDE VI: PDE with mixed derivative" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using Lux, Random, Optimisers, DomainSets, Cubature, QuasiMonteCarlo, Integrals
+    import ModelingToolkit: Interval, infimum, supremum
+    import OptimizationOptimJL: BFGS
+
     @parameters x y
     @variables u(..)
     Dxx = Differential(x)^2
@@ -363,9 +401,11 @@ end
     eq = Dxx(u(x, y)) + Dx(Dy(u(x, y))) - 2 * Dyy(u(x, y)) ~ -1.0
 
     # Initial and boundary conditions
-    bcs = [u(x, 0) ~ x,
+    bcs = [
+        u(x, 0) ~ x,
         Dy(u(x, 0)) ~ x,
-        u(x, 0) ~ Dy(u(x, 0))]
+        u(x, 0) ~ Dy(u(x, 0))
+    ]
 
     # Space and time domains
     domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
@@ -379,7 +419,7 @@ end
 
     prob = discretize(pde_system, discretization)
 
-    res = solve(prob, OptimizationOptimJL.BFGS(); maxiters = 1500)
+    res = solve(prob, BFGS(); maxiters = 1500)
     @show res.original
 
     phi = discretization.phi
@@ -394,7 +434,12 @@ end
     @test u_predict≈u_real rtol=0.1
 end
 
-@testset "Translating from Flux" begin
+@testitem "Translating from Flux" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
+    using Lux, Random, Optimisers, DomainSets, Cubature, QuasiMonteCarlo, Integrals
+    import ModelingToolkit: Interval, infimum, supremum
+    import OptimizationOptimJL: BFGS
+    import Flux
+
     @parameters θ
     @variables u(..)
     Dθ = Differential(θ)
@@ -409,11 +454,11 @@ end
 
     @named pde_system = PDESystem(eq, bcs, domains, [θ], [u])
     prob = discretize(pde_system, discretization)
-    res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.1); maxiters = 1000)
+    res = solve(prob, Adam(0.1); maxiters = 1000)
     prob = remake(prob, u0 = res.u)
-    res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.01); maxiters = 500)
+    res = solve(prob, Adam(0.01); maxiters = 500)
     prob = remake(prob, u0 = res.u)
-    res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.001); maxiters = 500)
+    res = solve(prob, Adam(0.001); maxiters = 500)
     phi = discretization.phi
     analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
     ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
diff --git a/test/direct_function_tests.jl b/test/direct_function_tests.jl
index 3e101afa15..022cbae2ad 100644
--- a/test/direct_function_tests.jl
+++ b/test/direct_function_tests.jl
@@ -63,7 +63,7 @@ end
     prob = discretize(pde_system, discretization)
     res = solve(prob, OptimizationOptimisers.Adam(0.01), maxiters = 500)
     prob = remake(prob, u0 = res.u)
-    res = solve(prob, OptimizationOptimJL.BFGS(), maxiters = 1000)
+    res = solve(prob, BFGS(), maxiters = 1000)
     dx = 0.01
     xs = collect(x0:dx:x_end)
     func_s = func(xs)
@@ -101,9 +101,9 @@ end
     symprob.loss_functions.full_loss_function(symprob.flat_init_params, nothing)
     res = solve(prob, OptimizationOptimisers.Adam(0.01), maxiters = 500)
     prob = remake(prob, u0 = res.u)
-    res = solve(prob, OptimizationOptimJL.BFGS(), maxiters = 1000)
+    res = solve(prob, BFGS(), maxiters = 1000)
     prob = remake(prob, u0 = res.u)
-    res = solve(prob, OptimizationOptimJL.BFGS(), maxiters = 500)
+    res = solve(prob, BFGS(), maxiters = 500)
     phi = discretization.phi
     xs = collect(x0:0.1:x_end)
     ys = collect(y0:0.1:y_end)