|
1 | 1 | using OrdinaryDiffEq, GTPSA, Test |
2 | | -using SciMLSensitivity # Required for reverse-mode AD |
3 | 2 | using DifferentiationInterface |
4 | | -using ADTypes: AutoForwardDiff, AutoMooncake |
5 | | - |
6 | | -# Load backends for all versions (required for DifferentiationInterface extensions) |
| 3 | +using ADTypes: AutoForwardDiff |
7 | 4 | using ForwardDiff |
8 | | -using Mooncake |
9 | | - |
10 | | -# Version-dependent imports |
11 | | -if VERSION <= v"1.11" |
12 | | - using Zygote |
13 | | - using ADTypes: AutoZygote |
14 | | -end |
15 | | -if VERSION <= v"1.11" |
16 | | - using Enzyme |
17 | | - using ADTypes: AutoEnzyme |
18 | | -end |
19 | | - |
20 | | -# Define backends based on Julia version |
21 | | -# ForwardDiff: All versions |
22 | | -# Mooncake: All versions |
23 | | -# Zygote: Julia <= 1.11 |
24 | | -# Enzyme: Julia <= 1.11 |
25 | | -function get_test_backends() |
26 | | - backends = Pair{String, Any}[] |
27 | | - # ForwardDiff on all versions |
28 | | - push!(backends, "ForwardDiff" => AutoForwardDiff()) |
29 | | - # Mooncake on all versions |
30 | | - push!(backends, "Mooncake" => AutoMooncake(; config = nothing)) |
31 | | - # Zygote only on Julia <= 1.11 |
32 | | - if VERSION <= v"1.11" |
33 | | - push!(backends, "Zygote" => AutoZygote()) |
34 | | - end |
35 | | - # Enzyme only on Julia <= 1.11 |
36 | | - if VERSION <= v"1.11" |
37 | | - push!(backends, "Enzyme" => AutoEnzyme()) |
38 | | - end |
39 | | - return backends |
40 | | -end |
41 | 5 |
|
42 | | -backends = get_test_backends() |
43 | | -# Note: All backends except ForwardDiff are excluded due to issues with GTPSA differentiation |
44 | | -# Mooncake: issues with GTPSA jacobian/hessian compilation |
45 | | -# Zygote: Jacobian test fails, Hessian tests error |
46 | | -# Enzyme: All tests error (missing rules for GTPSA types) |
47 | | -backends_for_gtpsa = filter(b -> b[1] == "ForwardDiff", backends) |
| 6 | +# GTPSA is itself an AD engine - these tests compare GTPSA jacobian/hessian |
| 7 | +# results against ForwardDiff as a reference implementation |
48 | 8 |
|
49 | 9 | # ODEProblem 1 ======================= |
50 | 10 |
|
@@ -74,28 +34,19 @@ function sol_end_problem1(t) |
74 | 34 | end |
75 | 35 |
|
76 | 36 | @testset "GTPSA Problem 1 Jacobian tests" begin |
77 | | - for (name, backend) in backends_for_gtpsa |
78 | | - @testset "Jacobian comparison with $name" begin |
79 | | - J_AD = DifferentiationInterface.jacobian(sol_end_problem1, backend, [x..., p...]) |
80 | | - @test J_AD ≈ GTPSA.jacobian(sol_GTPSA.u[end], include_params = true) |
81 | | - end |
82 | | - end |
| 37 | + J_AD = DifferentiationInterface.jacobian(sol_end_problem1, AutoForwardDiff(), [x..., p...]) |
| 38 | + @test J_AD ≈ GTPSA.jacobian(sol_GTPSA.u[end], include_params = true) |
83 | 39 | end |
84 | 40 |
|
85 | | -# Compare Hessians against AD backends using DifferentiationInterface |
86 | 41 | @testset "GTPSA Problem 1 Hessian tests" begin |
87 | | - for (name, backend) in backends_for_gtpsa |
88 | | - @testset "Hessian comparison with $name" begin |
89 | | - for i in 1:3 |
90 | | - function sol_end_i_problem1(t) |
91 | | - prob = ODEProblem(f!, t[1:3], (0.0, 1.0), t[4:6]) |
92 | | - sol = solve(prob, Tsit5(), reltol = 1.0e-16, abstol = 1.0e-16) |
93 | | - return sol.u[end][i] |
94 | | - end |
95 | | - Hi_AD = DifferentiationInterface.hessian(sol_end_i_problem1, backend, [x..., p...]) |
96 | | - @test Hi_AD ≈ GTPSA.hessian(sol_GTPSA.u[end][i], include_params = true) |
97 | | - end |
| 42 | + for i in 1:3 |
| 43 | + function sol_end_i_problem1(t) |
| 44 | + prob = ODEProblem(f!, t[1:3], (0.0, 1.0), t[4:6]) |
| 45 | + sol = solve(prob, Tsit5(), reltol = 1.0e-16, abstol = 1.0e-16) |
| 46 | + return sol.u[end][i] |
98 | 47 | end |
| 48 | + Hi_AD = DifferentiationInterface.hessian(sol_end_i_problem1, AutoForwardDiff(), [x..., p...]) |
| 49 | + @test Hi_AD ≈ GTPSA.hessian(sol_GTPSA.u[end][i], include_params = true) |
99 | 50 | end |
100 | 51 | end |
101 | 52 |
|
@@ -127,27 +78,18 @@ function sol_end_problem2(t) |
127 | 78 | end |
128 | 79 |
|
129 | 80 | @testset "GTPSA Problem 2 Jacobian tests" begin |
130 | | - for (name, backend) in backends_for_gtpsa |
131 | | - @testset "Jacobian comparison with $name" begin |
132 | | - J_AD = DifferentiationInterface.jacobian(sol_end_problem2, backend, zeros(6)) |
133 | | - @test J_AD ≈ GTPSA.jacobian(sol_GTPSA2.u[end], include_params = true) |
134 | | - end |
135 | | - end |
| 81 | + J_AD = DifferentiationInterface.jacobian(sol_end_problem2, AutoForwardDiff(), zeros(6)) |
| 82 | + @test J_AD ≈ GTPSA.jacobian(sol_GTPSA2.u[end], include_params = true) |
136 | 83 | end |
137 | 84 |
|
138 | | -# Compare Hessians against AD backends using DifferentiationInterface |
139 | 85 | @testset "GTPSA Problem 2 Hessian tests" begin |
140 | | - for (name, backend) in backends_for_gtpsa |
141 | | - @testset "Hessian comparison with $name" begin |
142 | | - for i in 1:6 |
143 | | - function sol_end_i_problem2(t) |
144 | | - prob = DynamicalODEProblem(pdot!, qdot!, t[1:3], t[4:6], (0.0, 25.0)) |
145 | | - sol = solve(prob, Yoshida6(), dt = 1.0, reltol = 1.0e-16, abstol = 1.0e-16) |
146 | | - return sol.u[end][i] |
147 | | - end |
148 | | - Hi_AD = DifferentiationInterface.hessian(sol_end_i_problem2, backend, zeros(6)) |
149 | | - @test Hi_AD ≈ GTPSA.hessian(sol_GTPSA2.u[end][i], include_params = true) |
150 | | - end |
| 86 | + for i in 1:6 |
| 87 | + function sol_end_i_problem2(t) |
| 88 | + prob = DynamicalODEProblem(pdot!, qdot!, t[1:3], t[4:6], (0.0, 25.0)) |
| 89 | + sol = solve(prob, Yoshida6(), dt = 1.0, reltol = 1.0e-16, abstol = 1.0e-16) |
| 90 | + return sol.u[end][i] |
151 | 91 | end |
| 92 | + Hi_AD = DifferentiationInterface.hessian(sol_end_i_problem2, AutoForwardDiff(), zeros(6)) |
| 93 | + @test Hi_AD ≈ GTPSA.hessian(sol_GTPSA2.u[end][i], include_params = true) |
152 | 94 | end |
153 | 95 | end |
0 commit comments