diff --git a/src/algorithms.jl b/src/algorithms.jl index 11ca0024..c829b746 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -87,6 +87,15 @@ This introduces the following limitations on its usage: Callbacks with `terminate!` do not work well with `EnsembleGPUArray` because the entire integration halts when any trajectory halts. Use with caution. +# Step-size control + +All trajectories of a batch advance with one shared step. The error estimate is measured +per trajectory (the RMS over its own components) and the largest of these decides +acceptance and the next step, so every trajectory meets `abstol`/`reltol` as it would in a +solve of its own. The shared step is therefore the one the hardest trajectory needs: a batch +mixing easy and hard trajectories takes as many steps as the hard ones. Pass +`internalnorm` to `solve` to replace this norm; it receives the batched state array. + # Examples ```julia diff --git a/src/ensemblegpuarray/lowerlevel_solve.jl b/src/ensemblegpuarray/lowerlevel_solve.jl index 6a7dc7e9..ae6e59f2 100644 --- a/src/ensemblegpuarray/lowerlevel_solve.jl +++ b/src/ensemblegpuarray/lowerlevel_solve.jl @@ -100,7 +100,7 @@ function vectorized_map_solve_up(prob, alg, ensemblealg, I, u0, p; kwargs...) end return sol = solve( - prob, _alg; kwargs..., callback = _callback, merge_callbacks = false, - internalnorm = diffeqgpunorm + prob, _alg; internalnorm = TrajectoryNorm(len), kwargs..., + callback = _callback, merge_callbacks = false ) end diff --git a/src/solve.jl b/src/solve.jl index a700d788..ccbd2ebc 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -465,8 +465,8 @@ function batch_solve_up(ensembleprob, probs, alg, ensemblealg, I, u0, p; kwargs. end sol = solve( - prob, _alg; kwargs..., callback = _callback, merge_callbacks = false, - internalnorm = diffeqgpunorm + prob, _alg; internalnorm = TrajectoryNorm(len), kwargs..., + callback = _callback, merge_callbacks = false ) us = Array.(sol.u) @@ -568,8 +568,8 @@ function ChainRulesCore.rrule( end sol = solve( - prob, _alg; kwargs..., callback = _callback, merge_callbacks = false, - internalnorm = diffeqgpunorm + prob, _alg; internalnorm = TrajectoryNorm(len), kwargs..., + callback = _callback, merge_callbacks = false ) us = Array.(sol.u) diff --git a/src/utils.jl b/src/utils.jl index b9ae524c..e97bd2e5 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -5,6 +5,32 @@ function diffeqgpunorm(u::AbstractArray{<:ForwardDiff.Dual}, t) end diffeqgpunorm(u::ForwardDiff.Dual, t) = abs(ForwardDiff.value(u)) +# Error norm for the batched `EnsembleGPUArray` problem, whose state stacks trajectories of +# `len` components each as columns (so a vectorized state is trajectory-blocked). The RMS +# over the whole batch lets one hard trajectory be outvoted by the easy ones and miss its +# own tolerance; the largest per-trajectory RMS makes the shared step the one the hardest +# trajectory needs. +struct TrajectoryNorm + len::Int +end + +# OrdinaryDiffEq broadcasts the norm over residual arrays; treat it as a scalar like a function. +Base.broadcastable(n::TrajectoryNorm) = Ref(n) + +_norm_value(x) = x +_norm_value(x::ForwardDiff.Dual) = ForwardDiff.value(x) + +(::TrajectoryNorm)(u::Union{AbstractFloat, Complex}, t) = abs(u) +(::TrajectoryNorm)(u::ForwardDiff.Dual, t) = abs(ForwardDiff.value(u)) +function (n::TrajectoryNorm)(u::AbstractArray, t) + isempty(u) && return zero(real(_norm_value(zero(eltype(u))))) + # An array not made of whole trajectories (no state-shaped solver array is) falls back + # to the RMS over all entries. + length(u) % n.len == 0 || return diffeqgpunorm(u, t) + sq = sum(abs2 ∘ _norm_value, reshape(u, n.len, :); dims = 1) + return sqrt(maximum(sq) / n.len) +end + """ make_prob_compatible(prob) diff --git a/test/ensemblegpuarray_trajectory_norm.jl b/test/ensemblegpuarray_trajectory_norm.jl new file mode 100644 index 00000000..b03616f0 --- /dev/null +++ b/test/ensemblegpuarray_trajectory_norm.jl @@ -0,0 +1,47 @@ +using DiffEqGPU, ForwardDiff, Test +using Adapt: adapt +using OrdinaryDiffEq: Tsit5 +using SciMLBase: EnsembleProblem, ODEProblem, remake, solve + +include("utils.jl") + +@testset "TrajectoryNorm" begin + norm = DiffEqGPU.TrajectoryNorm(2) + U = [1.0 0.0 0.0; 1.0 0.0 0.0] + # The worst trajectory's RMS, not the RMS over the batch. + @test norm(U, 0.0) ≈ 1.0 + @test norm(vec(U), 0.0) ≈ 1.0 + @test norm(adapt(backend, Float32.(U)), 0.0f0) ≈ 1.0f0 + @test norm(-3.0, 0.0) == 3.0 + @test norm(ForwardDiff.Dual(-2.0, 1.0), 0.0) == 2.0 + @test norm(ForwardDiff.Dual.(U, 1.0), 0.0) ≈ 1.0 + @test norm(Float64[], 0.0) == 0.0 + # Not a whole number of trajectories: RMS over all entries. + @test norm([3.0, 4.0, 0.0], 0.0) ≈ sqrt(25 / 3) +end + +# u' = ω cos(ω t), u(0) = 0, so u = sin(ω t). Trajectory 1 needs small steps (ω = 50); +# the others are easy (ω = 1). With a batch-wide RMS the easy trajectories outvote the hard +# one and its error grows with the batch size; each trajectory has to meet its own tolerance. +f_osc!(du, u, p, t) = (du[1] = p[1] * cos(p[1] * t); nothing) + +@testset "EnsembleGPUArray meets each trajectory's tolerance" begin + prob = ODEProblem(f_osc!, [0.0], (0.0, 10.0), [1.0]) + tol = 1.0e-8 + alone = solve(remake(prob; p = [50.0]), Tsit5(); abstol = tol, reltol = tol) + err_alone = abs(alone.u[end][1] - sin(500.0)) + ep = EnsembleProblem( + prob; safetycopy = false, + prob_func = (pr, ctx) -> remake(pr; p = [ctx.sim_id == 1 ? 50.0 : 1.0]) + ) + for ntraj in (2, 1000) + sol = solve( + ep, Tsit5(), EnsembleGPUArray(backend, 0.0); trajectories = ntraj, + abstol = tol, reltol = tol, save_everystep = false + ) + err = abs(sol.u[1].u[end][1] - sin(500.0)) + @test err <= 5 * max(err_alone, tol / 10) + # The easy trajectories ride along on the hard one's steps. + @test abs(sol.u[end].u[end][1] - sin(10.0)) <= 10 * tol + end +end diff --git a/test/runtests.jl b/test/runtests.jl index ca932140..5b7b41e7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -100,6 +100,9 @@ if GROUP in SUPPORTS_DOUBLE_PRECISION @time @safetestset "EnsembleGPUArray scalar batch parameters" begin include("ensemblegpuarray_scalar_batch.jl") end + @time @safetestset "EnsembleGPUArray per-trajectory error norm" begin + include("ensemblegpuarray_trajectory_norm.jl") + end @time @safetestset "Reduction" begin include("reduction.jl") end @@ -130,6 +133,9 @@ if GROUP == "CPU" @time @safetestset "EnsembleGPUArray scalar batch parameters" begin include("ensemblegpuarray_scalar_batch.jl") end + @time @safetestset "EnsembleGPUArray per-trajectory error norm" begin + include("ensemblegpuarray_trajectory_norm.jl") + end @time @safetestset "Adaptive endpoint termination" begin include("gpu_kernel_de/adaptive_endpoint.jl") end