test: fix NaN in gradients

The initial conditions for the UDE are the same as the reference solution,
leading to taking the derivative of sqrt at 0, which is NaN.
Changing the loss function to squared l2loss fixes this.

Author: SebastianM-C

Project.toml

@@ -27,6 +27,7 @@ OrdinaryDiffEq = "6.74"
 Random = "1"
 SafeTestsets = "0.1"
 SciMLStructures = "1.1.0"
+StableRNGs = "1"
 SymbolicIndexingInterface = "0.3.15"
 Symbolics = "5.27"
 Test = "1"
@@ -41,8 +42,9 @@ Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "JET", "Test", "OrdinaryDiffEq", "ForwardDiff", "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures", "SymbolicIndexingInterface"]
+test = ["Aqua", "JET", "Test", "OrdinaryDiffEq", "ForwardDiff", "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures", "StableRNGs", "SymbolicIndexingInterface"]

test/lotka_volterra.jl

@@ -10,6 +10,7 @@ using OptimizationOptimisers: Adam
 using SciMLStructures
 using SciMLStructures: Tunable
 using ForwardDiff
+using StableRNGs
 
 function lotka_ude()
     @variables t x(t)=3.1 y(t)=1.5
@@ -41,7 +42,9 @@ function lotka_true()
 end
 
 model = lotka_ude()
-nn = create_ude_component(2, 2)
+
+chain = multi_layer_feed_forward(2, 2)
+nn = create_ude_component(2, 2; chain, rng = StableRNG(42))
 
 eqs = [connect(model.nn_in, nn.output)
        connect(model.nn_out, nn.input)]
@@ -71,7 +74,7 @@ function loss(x, (prob, sol_ref, get_vars, get_refs))
     loss = zero(eltype(x))
 
     for i in eachindex(new_sol.u)
-        loss += sum(sqrt.(abs2.(get_vars(new_sol, i) .- get_refs(sol_ref, i))))
+        loss += sum(abs2.(get_vars(new_sol, i) .- get_refs(sol_ref, i)))
     end
 
     if SciMLBase.successful_retcode(new_sol)
@@ -106,12 +109,16 @@ op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs))
 #     false
 # end
 
-res = solve(op, Adam(), maxiters = 2000)#, callback = plot_cb)
+res = solve(op, Adam(), maxiters = 5000)#, callback = plot_cb)
 
 @test res.objective < 1
 
 res_p = SciMLStructures.replace(Tunable(), prob.p, res)
 res_prob = remake(prob, p = res_p)
 res_sol = solve(res_prob, Rodas4())
 
+# using Plots
+# plot(sol_ref, idxs = [model_true.x, model_true.y])
+# plot!(res_sol, idxs = [sys.lotka.x, sys.lotka.y])
+
 @test SciMLBase.successful_retcode(res_sol)