Merge pull request #8 from SciML/smc/fix

Fix NaN in gradients

Author: ChrisRackauckas

Project.toml

@@ -9,7 +9,6 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
@@ -22,13 +21,13 @@ Lux = "0.5.32"
 LuxCore = "0.1.14"
 ModelingToolkit = "9.9.0"
 ModelingToolkitStandardLibrary = "2.6"
-NNlib = "0.9"
 Optimization = "3.22"
 OptimizationOptimisers = "0.2"
 OrdinaryDiffEq = "6.74"
 Random = "1"
 SafeTestsets = "0.1"
 SciMLStructures = "1.1.0"
+StableRNGs = "1"
 SymbolicIndexingInterface = "0.3.15"
 Symbolics = "5.27"
 Test = "1"
@@ -43,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"]

src/UDEComponents.jl

@@ -6,7 +6,6 @@ using Symbolics: Symbolics, @register_array_symbolic, @wrapped
 using LuxCore: stateless_apply
 using Lux: Lux
 using Random: Xoshiro
-using NNlib: softplus
 using ComponentArrays: ComponentArray
 
 export create_ude_component, multi_layer_feed_forward

src/utils.jl

@@ -1,5 +1,5 @@
 function multi_layer_feed_forward(input_length, output_length; width::Int = 5,
-        depth::Int = 1, activation = softplus, disable_optimizations = false)
+        depth::Int = 1, activation = tanh, disable_optimizations = false)
     Lux.Chain(Lux.Dense(input_length, width, activation),
         [Lux.Dense(width, width, activation) for _ in 1:(depth)]...,
         Lux.Dense(width, output_length); disable_optimizations)

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)