From 0257163a01da7fc00a1e57c0a8357666470b3d6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?= Date: Fri, 2 Oct 2026 03:54:29 +0300 Subject: [PATCH] fix: per-trajectory error norm for EnsembleGPUArray MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batched solve used `diffeqgpunorm`, an RMS over all N·ntraj components, as `internalnorm`, and passed it after `kwargs...` so it could not be overridden. A trajectory that needs small steps is outvoted by the easy ones and silently misses its tolerance, by more the larger the batch. For u' = ω cos(ω t) on [0, 10] with Tsit5 at abstol = reltol = 1e-8, trajectory 1 at ω = 50 and the rest at ω = 1, its error at t = 10 is 1.6e-9 solved alone and 5.0e-7 inside a batch of 1000. `TrajectoryNorm(len)` takes the RMS of each trajectory's own components and returns the largest, for N×ntraj state arrays and their vectorized form, on host and device arrays, and for Dual-valued arrays. With it the error above is 1.6e-9 for any batch size. The default is set before `kwargs...`, so `internalnorm` passed to `solve` replaces it. The shared step is now the one the hardest trajectory needs, so batches take more steps. Lorenz, the EnsembleGPUArray docstring example (Float32, random p, saveat = 1, CUDA): - Tsit5, 10_000 trajectories on [0, 100]: naccept 1076 -> 2092, nreject 0 -> 6 - Rosenbrock23, 1_000 trajectories on [0, 10]: naccept 265 -> 797, nreject 15 -> 128 Co-Authored-By: Claude Opus 5.5 (1M context) --- src/algorithms.jl | 9 +++++ src/ensemblegpuarray/lowerlevel_solve.jl | 4 +- src/solve.jl | 8 ++-- src/utils.jl | 26 +++++++++++++ test/ensemblegpuarray_trajectory_norm.jl | 47 ++++++++++++++++++++++++ test/runtests.jl | 6 +++ 6 files changed, 94 insertions(+), 6 deletions(-) create mode 100644 test/ensemblegpuarray_trajectory_norm.jl 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