fix: per-trajectory error norm for EnsembleGPUArray

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

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

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)

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)
 

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

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