Merge pull request #45 from SciML/smc/ad

Add reverse mode AD tests

Author: ChrisRackauckas

Project.toml

@@ -15,36 +15,42 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 [compat]
 Aqua = "0.8"
 ComponentArrays = "0.15.11"
+DifferentiationInterface = "0.6"
 ForwardDiff = "0.10.36"
 JET = "0.8, 0.9"
 Lux = "1"
 LuxCore = "1"
-ModelingToolkit = "9.9"
+ModelingToolkit = "9.64"
 ModelingToolkitStandardLibrary = "2.7"
-Optimization = "3.24"
-OptimizationOptimisers = "0.2.1"
+Optimization = "3.24, 4"
+OptimizationOptimisers = "0.2.1, 0.3"
 OrdinaryDiffEq = "6.74"
 Random = "1.10"
 SafeTestsets = "0.1"
+SciMLSensitivity = "7.72"
 SciMLStructures = "1.1.0"
 StableRNGs = "1"
 SymbolicIndexingInterface = "0.3.15"
-Symbolics = "6.9"
+Symbolics = "6.22"
 Test = "1.10"
+Zygote = "0.6.73"
 julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "JET", "Test", "OrdinaryDiffEq", "ForwardDiff", "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures", "StableRNGs", "SymbolicIndexingInterface"]
+test = ["Aqua", "JET", "Test", "OrdinaryDiffEq", "DifferentiationInterface", "SciMLSensitivity", "Zygote", "ForwardDiff", "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures", "StableRNGs", "SymbolicIndexingInterface"]

test/lotka_volterra.jl

@@ -11,6 +11,9 @@ using SciMLStructures
 using SciMLStructures: Tunable, canonicalize
 using ForwardDiff
 using StableRNGs
+using DifferentiationInterface
+using SciMLSensitivity
+using Zygote: Zygote
 
 function lotka_ude()
     @variables t x(t)=3.1 y(t)=1.5
@@ -59,7 +62,7 @@ prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0, 1.0), [])
 
 model_true = structural_simplify(lotka_true())
 prob_true = ODEProblem{true, SciMLBase.FullSpecialize}(model_true, [], (0, 1.0), [])
-sol_ref = solve(prob_true, Rodas4())
+sol_ref = solve(prob_true, Rodas5P(), abstol = 1e-10, reltol = 1e-8)
 
 x0 = default_values(sys)[nn.p]
 
@@ -71,7 +74,7 @@ function loss(x, (prob, sol_ref, get_vars, get_refs, set_x))
     new_p = set_x(prob, x)
     new_prob = remake(prob, p = new_p, u0 = eltype(x).(prob.u0))
     ts = sol_ref.t
-    new_sol = solve(new_prob, Rodas4(), saveat = ts)
+    new_sol = solve(new_prob, Rodas5P(), abstol = 1e-10, reltol = 1e-8, saveat = ts)
 
     loss = zero(eltype(x))
 
@@ -86,14 +89,22 @@ function loss(x, (prob, sol_ref, get_vars, get_refs, set_x))
     end
 end
 
-of = OptimizationFunction{true}(loss, AutoForwardDiff())
+of = OptimizationFunction{true}(loss, AutoZygote())
 
 ps = (prob, sol_ref, get_vars, get_refs, set_x);
 
 @test_call target_modules=(ModelingToolkitNeuralNets,) loss(x0, ps)
 @test_opt target_modules=(ModelingToolkitNeuralNets,) loss(x0, ps)
 
-@test all(.!isnan.(ForwardDiff.gradient(Base.Fix2(of, ps), x0)))
+∇l1 = DifferentiationInterface.gradient(Base.Fix2(of, ps), AutoForwardDiff(), x0)
+∇l2 = DifferentiationInterface.gradient(Base.Fix2(of, ps), AutoFiniteDiff(), x0)
+∇l3 = DifferentiationInterface.gradient(Base.Fix2(of, ps), AutoZygote(), x0)
+
+@test all(.!isnan.(∇l1))
+@test !iszero(∇l1)
+
+@test ∇l1≈∇l2 rtol=1e-2
+@test ∇l1≈∇l3 rtol=1e-5
 
 op = OptimizationProblem(of, x0, ps)
 
@@ -111,7 +122,7 @@ op = OptimizationProblem(of, x0, ps)
 #     false
 # end
 
-res = solve(op, Adam(), maxiters = 5000)#, callback = plot_cb)
+res = solve(op, Adam(), maxiters = 10000)#, callback = plot_cb)
 
 @test res.objective < 1