fix tests

Author: SebastianM-C

test/lotka_volterra.jl

@@ -48,18 +48,18 @@ end
 
 rbf(x) = exp.(-(x .^ 2))
 
-chain = Lux.Chain(
-    Lux.Dense(2, 5, rbf), Lux.Dense(5, 5, rbf), Lux.Dense(5, 5, rbf),
-    Lux.Dense(5, 2))
+chain = multi_layer_feed_forward(2, 2, width = 5, initial_scaling_factor = 1)
 ude_sys = lotka_ude(chain)
 
-sys = mtkcompile(ude_sys, allow_symbolic = true)
+sys = mtkcompile(ude_sys)
+
+@test length(equations(sys)) == 2
 
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0, 5.0))
 
 model_true = mtkcompile(lotka_true())
 prob_true = ODEProblem{true, SciMLBase.FullSpecialize}(model_true, [], (0, 5.0))
-sol_ref = solve(prob_true, Vern9(), abstol = 1e-12, reltol = 1e-12)
+sol_ref = solve(prob_true, Vern9(), abstol = 1e-8, reltol = 1e-8)
 
 ts = range(0, 5.0, length = 21)
 data = reduce(hcat, sol_ref(ts, idxs = [model_true.x, model_true.y]).u)
@@ -70,10 +70,9 @@ get_vars = getu(sys, [sys.x, sys.y])
 set_x = setsym_oop(sys, sys.nn.p)
 
 function loss(x, (prob, sol_ref, get_vars, data, ts, set_x))
-    # new_u0, new_p = set_x(prob, 1, x)
     new_u0, new_p = set_x(prob, x)
     new_prob = remake(prob, p = new_p, u0 = new_u0)
-    new_sol = solve(new_prob, Vern9(), abstol = 1e-10, reltol = 1e-8, saveat = ts)
+    new_sol = solve(new_prob, Vern9(), abstol = 1e-8, reltol = 1e-8, saveat = ts)
 
     if SciMLBase.successful_retcode(new_sol)
         mean(abs2.(reduce(hcat, get_vars(new_sol)) .- data))
@@ -106,30 +105,32 @@ op = OptimizationProblem(of, x0, ps)
 # oh = []
 
 # plot_cb = (opt_state, loss) -> begin
+#     opt_state.iter % 500 ≠ 0 && return false
 #     @info "step $(opt_state.iter), loss: $loss"
 #     push!(oh, opt_state)
 #     new_p = SciMLStructures.replace(Tunable(), prob.p, opt_state.u)
 #     new_prob = remake(prob, p = new_p)
-#     sol = solve(new_prob, Rodas4())
+#     sol = solve(new_prob, Vern9(), abstol = 1e-8, reltol = 1e-8)
 #     display(plot(sol))
 #     false
 # end
 
-res = solve(op, Adam(), maxiters = 10000)#, callback = plot_cb)
+res = solve(op, Adam(1e-3), maxiters = 25_000, callback = plot_cb)
 
 display(res.stats)
-@test res.objective < 1
+@test res.objective < 1e-4
+
+u0, p = set_x(prob, res.u)
+res_prob = remake(prob; u0, p)
+res_sol = solve(res_prob, Vern9(), abstol = 1e-8, reltol = 1e-8, saveat = ts)
 
-res_p = set_x(prob, res.u)
-res_prob = remake(prob, p = res_p)
-res_sol = solve(res_prob, Vern9())
+@test SciMLBase.successful_retcode(res_sol)
+@test mean(abs2.(reduce(hcat, get_vars(res_sol)) .- data)) ≈ res.objective
 
 # using Plots
 # plot(sol_ref, idxs = [model_true.x, model_true.y])
 # plot!(res_sol, idxs = [sys.x, sys.y])
 
-@test SciMLBase.successful_retcode(res_sol)
-
 function lotka_ude2()
     @variables t x(t)=3.1 y(t)=1.5 pred(t)[1:2]
     @parameters α=1.3 [tunable = false] δ=1.8 [tunable = false]