Skip to content

Commit c2b7485

Browse files
ChrisRackauckasChrisRackauckas
committed
Split MIRK AD CI shards
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
1 parent a62a555 commit c2b7485

3 files changed

Lines changed: 147 additions & 103 deletions

File tree

‎lib/BoundaryValueDiffEqMIRK/test/AD/ad_tests.jl‎

Lines changed: 106 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -1,118 +1,127 @@
11
using BoundaryValueDiffEqMIRK
22
using Test
33

4+
mirk_ad_group() = get(ENV, "BOUNDARYVALUEDIFFEQ_MIRK_AD_GROUP", "ALL")
5+
run_mirk_ad_group(group) = mirk_ad_group() in ("ALL", group)
6+
47
@testset "Different AD compatibility" begin
58
using BoundaryValueDiffEqMIRK
69
using ForwardDiff, Enzyme, Mooncake
710

8-
@testset "Test different AD on multipoint BVP" begin
9-
function simplependulum!(du, u, p, t)
10-
θ = u[1]
11-
dθ = u[2]
12-
du[1] = dθ
13-
du[2] = -9.81 * sin(θ)
14-
end
15-
function bc!(residual, u, p, t)
16-
residual[1] = u[:, end ÷ 2][1] + pi / 2
17-
residual[2] = u[:, end][1] - pi / 2
18-
end
19-
u0 = [pi / 2, pi / 2]
20-
tspan = (0.0, pi / 2)
21-
prob = BVProblem(simplependulum!, bc!, u0, tspan)
22-
jac_alg_forwarddiff = BVPJacobianAlgorithm(
23-
bc_diffmode = AutoSparse(AutoForwardDiff()), nonbc_diffmode = AutoForwardDiff()
24-
)
25-
jac_alg_enzyme = BVPJacobianAlgorithm(
26-
bc_diffmode = AutoSparse(
27-
AutoEnzyme(
28-
mode = Enzyme.Reverse, function_annotation = Enzyme.Duplicated
29-
)
30-
),
31-
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
32-
)
33-
jac_alg_mooncake = BVPJacobianAlgorithm(
34-
bc_diffmode = AutoSparse(AutoMooncake(; config = nothing)),
35-
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
36-
)
37-
for jac_alg in [jac_alg_forwarddiff, jac_alg_enzyme, jac_alg_mooncake]
38-
sol = solve(prob, MIRK4(; jac_alg = jac_alg), dt = 0.05)
39-
@test SciMLBase.successful_retcode(sol)
11+
if run_mirk_ad_group("MULTIPOINT_GRID")
12+
@testset "Test different AD on multipoint BVP" begin
13+
function simplependulum!(du, u, p, t)
14+
θ = u[1]
15+
dθ = u[2]
16+
du[1] = dθ
17+
du[2] = -9.81 * sin(θ)
18+
end
19+
function bc!(residual, u, p, t)
20+
residual[1] = u[:, end ÷ 2][1] + pi / 2
21+
residual[2] = u[:, end][1] - pi / 2
22+
end
23+
u0 = [pi / 2, pi / 2]
24+
tspan = (0.0, pi / 2)
25+
prob = BVProblem(simplependulum!, bc!, u0, tspan)
26+
jac_alg_forwarddiff = BVPJacobianAlgorithm(
27+
bc_diffmode = AutoSparse(AutoForwardDiff()), nonbc_diffmode = AutoForwardDiff()
28+
)
29+
jac_alg_enzyme = BVPJacobianAlgorithm(
30+
bc_diffmode = AutoSparse(
31+
AutoEnzyme(
32+
mode = Enzyme.Reverse, function_annotation = Enzyme.Duplicated
33+
)
34+
),
35+
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
36+
)
37+
jac_alg_mooncake = BVPJacobianAlgorithm(
38+
bc_diffmode = AutoSparse(AutoMooncake(; config = nothing)),
39+
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
40+
)
41+
for jac_alg in [jac_alg_forwarddiff, jac_alg_enzyme, jac_alg_mooncake]
42+
sol = solve(prob, MIRK4(; jac_alg = jac_alg), dt = 0.05)
43+
@test SciMLBase.successful_retcode(sol)
44+
end
4045
end
4146
end
4247

43-
@testset "Test different AD on multipoint BVP using Interpolation BC" begin
44-
function simplependulum!(du, u, p, t)
45-
θ = u[1]
46-
dθ = u[2]
47-
du[1] = dθ
48-
du[2] = -9.81 * sin(θ)
49-
end
50-
function bc!(residual, u, p, t)
51-
residual[1] = u(pi / 4)[1] + pi / 2
52-
residual[2] = u(pi / 2)[1] - pi / 2
53-
end
54-
u0 = [pi / 2, pi / 2]
55-
tspan = (0.0, pi / 2)
56-
prob = BVProblem(simplependulum!, bc!, u0, tspan)
57-
jac_alg_forwarddiff = BVPJacobianAlgorithm(
58-
bc_diffmode = AutoSparse(AutoForwardDiff()), nonbc_diffmode = AutoForwardDiff()
59-
)
60-
jac_alg_enzyme = BVPJacobianAlgorithm(
61-
bc_diffmode = AutoSparse(
62-
AutoEnzyme(
63-
mode = Enzyme.Reverse, function_annotation = Enzyme.Duplicated
64-
)
65-
),
66-
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
67-
)
68-
jac_alg_mooncake = BVPJacobianAlgorithm(
69-
bc_diffmode = AutoSparse(AutoMooncake(; config = nothing)),
70-
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
71-
)
72-
for jac_alg in [jac_alg_forwarddiff, jac_alg_enzyme, jac_alg_mooncake]
73-
sol = solve(prob, MIRK4(; jac_alg = jac_alg), dt = 0.05)
74-
@test SciMLBase.successful_retcode(sol)
48+
if run_mirk_ad_group("MULTIPOINT_INTERPOLATION")
49+
@testset "Test different AD on multipoint BVP using Interpolation BC" begin
50+
function simplependulum!(du, u, p, t)
51+
θ = u[1]
52+
dθ = u[2]
53+
du[1] = dθ
54+
du[2] = -9.81 * sin(θ)
55+
end
56+
function bc!(residual, u, p, t)
57+
residual[1] = u(pi / 4)[1] + pi / 2
58+
residual[2] = u(pi / 2)[1] - pi / 2
59+
end
60+
u0 = [pi / 2, pi / 2]
61+
tspan = (0.0, pi / 2)
62+
prob = BVProblem(simplependulum!, bc!, u0, tspan)
63+
jac_alg_forwarddiff = BVPJacobianAlgorithm(
64+
bc_diffmode = AutoSparse(AutoForwardDiff()), nonbc_diffmode = AutoForwardDiff()
65+
)
66+
jac_alg_enzyme = BVPJacobianAlgorithm(
67+
bc_diffmode = AutoSparse(
68+
AutoEnzyme(
69+
mode = Enzyme.Reverse, function_annotation = Enzyme.Duplicated
70+
)
71+
),
72+
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
73+
)
74+
jac_alg_mooncake = BVPJacobianAlgorithm(
75+
bc_diffmode = AutoSparse(AutoMooncake(; config = nothing)),
76+
nonbc_diffmode = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated)
77+
)
78+
for jac_alg in [jac_alg_forwarddiff, jac_alg_enzyme, jac_alg_mooncake]
79+
sol = solve(prob, MIRK4(; jac_alg = jac_alg), dt = 0.05)
80+
@test SciMLBase.successful_retcode(sol)
81+
end
7582
end
7683
end
7784

78-
@testset "Test different AD on twopoint BVP" begin
79-
function f!(du, u, p, t)
80-
du[1] = u[2]
81-
du[2] = 0
82-
end
83-
function boundary_two_point_a!(resida, ua, p)
84-
resida[1] = ua[1] - 5
85-
end
86-
function boundary_two_point_b!(residb, ub, p)
87-
residb[1] = ub[1]
88-
end
85+
if run_mirk_ad_group("TWOPOINT")
86+
@testset "Test different AD on twopoint BVP" begin
87+
function f!(du, u, p, t)
88+
du[1] = u[2]
89+
du[2] = 0
90+
end
91+
function boundary_two_point_a!(resida, ua, p)
92+
resida[1] = ua[1] - 5
93+
end
94+
function boundary_two_point_b!(residb, ub, p)
95+
residb[1] = ub[1]
96+
end
8997

90-
odef! = ODEFunction(f!, analytic = (u0, p, t) -> [5 - t, -1])
91-
bcresid_prototype = (Array{Float64}(undef, 1), Array{Float64}(undef, 1))
92-
tspan = (0.0, 5.0)
93-
u0 = [5.0, -3.5]
94-
prob = TwoPointBVProblem(
95-
odef!, (boundary_two_point_a!, boundary_two_point_b!),
96-
u0, tspan; bcresid_prototype, nlls = Val(false)
97-
)
98-
jac_alg_forwarddiff = BVPJacobianAlgorithm(AutoSparse(AutoForwardDiff()))
99-
jac_alg_enzyme = BVPJacobianAlgorithm(
100-
AutoSparse(
101-
AutoEnzyme(
102-
mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated
98+
odef! = ODEFunction(f!, analytic = (u0, p, t) -> [5 - t, -1])
99+
bcresid_prototype = (Array{Float64}(undef, 1), Array{Float64}(undef, 1))
100+
tspan = (0.0, 5.0)
101+
u0 = [5.0, -3.5]
102+
prob = TwoPointBVProblem(
103+
odef!, (boundary_two_point_a!, boundary_two_point_b!),
104+
u0, tspan; bcresid_prototype, nlls = Val(false)
105+
)
106+
jac_alg_forwarddiff = BVPJacobianAlgorithm(AutoSparse(AutoForwardDiff()))
107+
jac_alg_enzyme = BVPJacobianAlgorithm(
108+
AutoSparse(
109+
AutoEnzyme(
110+
mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated
111+
)
103112
)
104113
)
105-
)
106-
jac_alg_mooncake = BVPJacobianAlgorithm(
107-
AutoSparse(
108-
AutoMooncake(;
109-
config = nothing
114+
jac_alg_mooncake = BVPJacobianAlgorithm(
115+
AutoSparse(
116+
AutoMooncake(;
117+
config = nothing
118+
)
110119
)
111120
)
112-
)
113-
for jac_alg in [jac_alg_forwarddiff, jac_alg_enzyme, jac_alg_mooncake]
114-
sol = solve(prob, MIRK4(; jac_alg = jac_alg), dt = 0.01)
115-
@test SciMLBase.successful_retcode(sol)
121+
for jac_alg in [jac_alg_forwarddiff, jac_alg_enzyme, jac_alg_mooncake]
122+
sol = solve(prob, MIRK4(; jac_alg = jac_alg), dt = 0.01)
123+
@test SciMLBase.successful_retcode(sol)
124+
end
116125
end
117126
end
118127
end

‎lib/BoundaryValueDiffEqMIRK/test/runtests.jl‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,25 @@
11
using SafeTestsets, Test
22
using SciMLTesting
33

4+
const MIRK_AD_ENV = joinpath(@__DIR__, "AD")
5+
6+
mirk_ad_all() = @time @safetestset "MIRK AD Tests" include("AD/ad_tests.jl")
7+
mirk_ad_multipoint_grid() = @time @safetestset "MIRK AD Multipoint Grid Tests" begin
8+
withenv("BOUNDARYVALUEDIFFEQ_MIRK_AD_GROUP" => "MULTIPOINT_GRID") do
9+
include("AD/ad_tests.jl")
10+
end
11+
end
12+
mirk_ad_multipoint_interpolation() = @time @safetestset "MIRK AD Multipoint Interpolation Tests" begin
13+
withenv("BOUNDARYVALUEDIFFEQ_MIRK_AD_GROUP" => "MULTIPOINT_INTERPOLATION") do
14+
include("AD/ad_tests.jl")
15+
end
16+
end
17+
mirk_ad_twopoint() = @time @safetestset "MIRK AD TwoPoint Tests" begin
18+
withenv("BOUNDARYVALUEDIFFEQ_MIRK_AD_GROUP" => "TWOPOINT") do
19+
include("AD/ad_tests.jl")
20+
end
21+
end
22+
423
run_tests(;
524
env = "BOUNDARYVALUEDIFFEQ_TEST_GROUP",
625
core = function ()
@@ -18,10 +37,20 @@ run_tests(;
1837
# large joint at-floor resolve on the Downgrade lane); they live in this
1938
# group's own test/AD/Project.toml, auto-activated before the body runs.
2039
"AD" => (;
21-
env = joinpath(@__DIR__, "AD"),
22-
body = function ()
23-
return @time @safetestset "MIRK AD Tests" include("AD/ad_tests.jl")
24-
end,
40+
env = MIRK_AD_ENV,
41+
body = mirk_ad_all,
42+
),
43+
"AD_MULTIPOINT_GRID" => (;
44+
env = MIRK_AD_ENV,
45+
body = mirk_ad_multipoint_grid,
46+
),
47+
"AD_MULTIPOINT_INTERPOLATION" => (;
48+
env = MIRK_AD_ENV,
49+
body = mirk_ad_multipoint_interpolation,
50+
),
51+
"AD_TWOPOINT" => (;
52+
env = MIRK_AD_ENV,
53+
body = mirk_ad_twopoint,
2554
),
2655
),
2756
qa = (;
@@ -32,5 +61,5 @@ run_tests(;
3261
return @time @safetestset "Quality Assurance" include("qa/qa.jl")
3362
end,
3463
),
35-
all = ["Core", "AD", "QA"],
64+
all = ["Core", "AD_MULTIPOINT_GRID", "AD_MULTIPOINT_INTERPOLATION", "AD_TWOPOINT", "QA"],
3665
)

‎lib/BoundaryValueDiffEqMIRK/test/test_groups.toml‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,13 @@ versions = ["lts", "1", "pre"]
1212
# AD: different-AD-backend (Enzyme/Mooncake) compatibility tests. The heavy
1313
# optional backends live in test/AD/Project.toml (kept out of the main test env so
1414
# the Downgrade lane does not force a joint at-floor resolve of them).
15-
[AD]
15+
[AD_MULTIPOINT_GRID]
16+
versions = ["lts", "1", "pre"]
17+
18+
[AD_MULTIPOINT_INTERPOLATION]
19+
versions = ["lts", "1", "pre"]
20+
21+
[AD_TWOPOINT]
1622
versions = ["lts", "1", "pre"]
1723

1824
[QA]

0 commit comments

Comments
 (0)