Let convergence studies drop the trajectories they are not measuring

`test_convergence` retains one full solution object per trajectory per step
size. A `RODESolution` carries the solver cache, the noise process, the problem
and its interpolation, none of which a convergence study needs — it wants the
per-trajectory errors and the endpoints the weak error is formed from. At the
trajectory counts the weak-convergence tests use, that gap is the difference
between passing and being OOM-killed: `SROCKC2WeakConvergence`,
`IIPWeakConvergence` and `OOPWeakConvergence` all die at 15.6 GiB in a 16 GiB
cgroup, with no test output, which is the "runner lost communication with the
server" signature reported in #4036.

Add a `save_solutions` keyword, default `true` so existing behaviour is
unchanged. With `false`, build the `EnsembleProblem` with an `output_func` that
reduces each trajectory to a `ConvergenceTrajectory` as it is solved, so the
full solutions never coexist and the peak is bounded rather than only the
retention. `ConvergenceTrajectory` stores one-tuples and a `NamedTuple` rather
than one-element `Vector`s and the solver's error `Dict`; indexing and key
lookup are unchanged, but the `Dict` alone dominated everything else retained
once a study runs to millions of trajectories.

Reject the combinations that cannot work with an explicit `ArgumentError`:
`weak_timeseries_errors` and `weak_dense_errors` need the full timeseries,
`expected_value` averages the trajectory values themselves, and an
`EnsembleProblem` supplied by the caller has its own `output_func` to set.

Drop `uEltype` from the `ConvergenceSimulation` constructor — it was assigned
and never used, and was the only thing requiring `solutions[i].u[j]` to be a
full solution.

Measured on the four weak files, 16 GiB cgroup, JULIA_NUM_THREADS=2:

  weak_srockc2.jl   OOM 15.61 GiB @3m29s   -> exit 0, 24m31s, 4.69 GiB
  oop_weak.jl       OOM 15.62 GiB @23m50s  -> exit 0, 12m48s, 1.58 GiB
  additive_weak.jl  (same group)           -> exit 0, 3m10s,  1.03 GiB
  iip_weak.jl       OOM 15.61 GiB @21m24s  -> runs fully, 10m13s, 2.22 GiB

Order estimates are bit-identical with and without the reduction.

With no group needing more than 4.69 GiB, remove the `high-memory` runner pin
from all eight groups that carried it. The five StochasticDiffEqWeak groups
needed no code change — they already reduce through their own `output_func` and
peak at 0.71-2.09 GiB, so the label was simply wrong for them.

Fixes #4037

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UwLXp5WY1uiqPun7qhwxeU

Author: ChrisRackauckas

docs/src/devtools/alg_dev/convergence.md

@@ -92,10 +92,26 @@ log2(error[i+1]/error[i])
 
 Returns the mean of the convergence estimates.
 
+## Memory use of Monte-Carlo studies
+
+For an SDE or RODE problem, `test_convergence` runs an ensemble per step size and by
+default keeps every trajectory's full solution in `solutions`. Each of those carries the
+solver cache, the noise process and the problem, so a study over many trajectories
+retains far more than the error estimates it reports — at hundreds of thousands of
+trajectories it reaches tens of gigabytes.
+
+Pass `save_solutions = false` to reduce each trajectory to a
+[`ConvergenceTrajectory`](@ref) as soon as it is solved, keeping only the errors and
+endpoints the order estimate is computed from. The order estimates are identical; only
+the contents of `solutions` change. This is incompatible with `weak_timeseries_errors`
+and `weak_dense_errors`, which need the whole timeseries, and it does not apply when you
+supply your own `EnsembleProblem` — set that problem's `output_func` instead.
+
 ## API
 
 ```@docs
 DiffEqDevTools.ConvergenceSimulation
+DiffEqDevTools.ConvergenceTrajectory
 DiffEqDevTools.test_convergence
 DiffEqDevTools.analyticless_test_convergence
 ```

lib/DiffEqDevTools/Project.toml

@@ -1,7 +1,7 @@
 name = "DiffEqDevTools"
 uuid = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "3.1.5"
+version = "3.2.0"
 
 [deps]
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"

lib/DiffEqDevTools/src/DiffEqDevTools.jl

@@ -57,7 +57,7 @@ include("test_solution.jl")
 include("ode_tableaus.jl")
 include("tableau_info.jl")
 
-export ConvergenceSimulation, Shootout, ShootoutSet, TestSolution
+export ConvergenceSimulation, ConvergenceTrajectory, Shootout, ShootoutSet, TestSolution
 
 #Benchmark Functions
 export Shootout, ShootoutSet, WorkPrecision, WorkPrecisionSet

lib/DiffEqDevTools/src/convergence.jl

@@ -24,7 +24,6 @@ function ConvergenceSimulation(
         expected_value = nothing
     )
     N = size(solutions, 1)
-    uEltype = eltype(solutions[1].u[1])
     errors = Dict() #Should add type information
     if expected_value == nothing
         if isnothing(solutions[1].errors) || isempty(solutions[1].errors)
@@ -54,6 +53,46 @@ function ConvergenceSimulation(
     return (ConvergenceSimulation(solutions, errors, N, auxdata, 𝒪est, convergence_axis))
 end
 
+"""
+    ConvergenceTrajectory(t, u, u_analytic, errors)
+
+The part of a trajectory's solution that a convergence study actually consumes: its
+error measurements and the endpoint values the weak errors are formed from.
+
+[`test_convergence`](@ref) stores one of these per trajectory in place of the full
+solution when `save_solutions = false`, which is what keeps a Monte-Carlo study's
+memory proportional to the errors it reports rather than to the solver state it
+happened to allocate along the way.
+"""
+struct ConvergenceTrajectory{tType, uType, EType <: NamedTuple}
+    t::Tuple{tType}
+    u::Tuple{uType}
+    u_analytic::Tuple{uType}
+    errors::EType
+end
+
+"""
+    _convergence_output_func(sol, ctx) -> (ConvergenceTrajectory, false)
+
+`output_func` that reduces a trajectory to its errors and endpoints as soon as it is
+solved, so the ensemble never holds the full solutions at once.
+
+Only valid when the weak timeseries and dense errors are switched off, since those
+need the whole timeseries and the noise process; [`test_convergence`](@ref) checks
+that before installing this.
+"""
+function _convergence_output_func(sol, ctx)
+    # One-tuples rather than one-element vectors, and a NamedTuple rather than the
+    # solver's error `Dict`: `[end]` and key lookup work the same, but a `Dict` alone
+    # costs several hundred bytes per trajectory, which dominates everything else
+    # retained here once the study runs to millions of trajectories.
+    u_analytic = sol.u_analytic === nothing ? sol.u[end] : sol.u_analytic[end]
+    reduced = ConvergenceTrajectory(
+        (sol.t[end],), (sol.u[end],), (u_analytic,), NamedTuple(sol.errors)
+    )
+    return (reduced, false)
+end
+
 """
     test_convergence(dts, prob, alg[, ensemblealg]; kwargs...)
     test_convergence(probs, convergence_axis, alg; kwargs...)
@@ -66,6 +105,19 @@ is passed to the solver as `dt`. Remaining keyword arguments are forwarded to `s
 The computed solutions must contain error measurements, usually obtained from an
 analytic solution. Use [`analyticless_test_convergence`](@ref) when only a numerical
 reference solution is available.
+
+## Keyword Arguments
+
+  - `save_solutions`: whether `ConvergenceSimulation.solutions` keeps the full
+    trajectory solutions (default `true`). A Monte-Carlo study over many trajectories
+    retains one complete solution object per trajectory per step size, which for large
+    `trajectories` is orders of magnitude more memory than the error estimates it is
+    computing. Pass `false` to store a [`ConvergenceTrajectory`](@ref) per trajectory
+    instead, holding only the errors and endpoints the order estimate is derived from.
+    Only applies when this function builds the `EnsembleProblem` itself — pass your own
+    `output_func` to control the reduction when supplying an `EnsembleProblem` — and it
+    is incompatible with `weak_timeseries_errors` and `weak_dense_errors`, which need
+    the whole timeseries.
 """
 function test_convergence(
         dts::AbstractArray,
@@ -77,12 +129,45 @@ function test_convergence(
         trajectories, save_start = true, save_everystep = true,
         timeseries_errors = save_everystep, adaptive = false,
         weak_timeseries_errors = false, weak_dense_errors = false,
-        expected_value = nothing, kwargs...
+        expected_value = nothing, save_solutions = true, kwargs...
     )
     N = length(dts)
 
+    reduce_trajectories = !save_solutions && !(prob isa AbstractEnsembleProblem)
+    if !save_solutions && !reduce_trajectories
+        throw(
+            ArgumentError(
+                "`save_solutions = false` reduces each trajectory through the " *
+                    "`EnsembleProblem`'s `output_func`, so it cannot be applied to an " *
+                    "`EnsembleProblem` you supplied yourself. Set `output_func` on your " *
+                    "`EnsembleProblem` instead."
+            )
+        )
+    end
+    if reduce_trajectories && (weak_timeseries_errors || weak_dense_errors)
+        throw(
+            ArgumentError(
+                "`save_solutions = false` keeps only each trajectory's errors and " *
+                    "endpoints, but `weak_timeseries_errors` and `weak_dense_errors` are " *
+                    "computed from the full timeseries. Use one or the other."
+            )
+        )
+    end
+    if reduce_trajectories && expected_value !== nothing
+        throw(
+            ArgumentError(
+                "`save_solutions = false` replaces each trajectory with its errors and " *
+                    "endpoints, but `expected_value` forms the weak error from the " *
+                    "trajectory values directly. Supply an `EnsembleProblem` whose " *
+                    "`output_func` reduces to the quantity being averaged instead."
+            )
+        )
+    end
+
     if prob isa AbstractEnsembleProblem
         ensemble_prob = prob
+    elseif reduce_trajectories
+        ensemble_prob = EnsembleProblem(prob; output_func = _convergence_output_func)
     else
         ensemble_prob = EnsembleProblem(prob)
     end

lib/DiffEqDevTools/test/runtests.jl

@@ -30,6 +30,9 @@ if TEST_GROUP == "Core" || TEST_GROUP == "ALL"
     @time @testset "Analyticless Stochastic WP" begin
         include("analyticless_stochastic_wp.jl")
     end
+    @time @testset "Convergence save_solutions" begin
+        include("save_solutions_tests.jl")
+    end
     @time @testset "Stability Region Tests" begin
         include("stability_region_test.jl")
     end

lib/DiffEqDevTools/test/save_solutions_tests.jl

@@ -0,0 +1,74 @@
+using StochasticDiffEq, DiffEqDevTools, Random, Test
+
+# `test_convergence` keeps one full solution object per trajectory per step size, which
+# on a Monte-Carlo study is orders of magnitude more than the error estimates need.
+# `save_solutions = false` reduces each trajectory to a `ConvergenceTrajectory` through
+# the ensemble's `output_func`. The order estimates must be unaffected.
+
+f_iip(du, u, p, t) = @.(du = 1.01 * u)
+σ_iip(du, u, p, t) = @.(du = 0.87 * u)
+linear_analytic(u0, p, t, W) = @.(u0 * exp(0.63155t + 0.87W))
+prob = SDEProblem(
+    SDEFunction(f_iip, σ_iip, analytic = linear_analytic), [1 / 2], (0.0, 1.0)
+)
+
+dts = 1 .// 2 .^ (5:-1:3)
+ntraj = 400
+
+function run_sim(save_solutions)
+    Random.seed!(100)
+    return test_convergence(
+        dts, prob, EM(); save_everystep = false, trajectories = ntraj,
+        weak_timeseries_errors = false, save_solutions = save_solutions
+    )
+end
+
+full = run_sim(true)
+reduced = run_sim(false)
+
+@testset "order estimates are unchanged" begin
+    @test keys(full.𝒪est) == keys(reduced.𝒪est)
+    for k in keys(full.𝒪est)
+        @test full.𝒪est[k] ≈ reduced.𝒪est[k]
+    end
+    @test keys(full.errors) == keys(reduced.errors)
+    for k in keys(full.errors)
+        @test collect(full.errors[k]) ≈ collect(reduced.errors[k])
+    end
+end
+
+@testset "trajectories are reduced" begin
+    @test reduced.solutions[1].u[1] isa ConvergenceTrajectory
+    @test full.solutions[1].u[1] isa SciMLBase.AbstractRODESolution
+    # the endpoints the weak errors are formed from must survive the reduction
+    for i in eachindex(dts), j in 1:ntraj
+        @test reduced.solutions[i].u[j].u[end] == full.solutions[i].u[j].u[end]
+        @test reduced.solutions[i].u[j].u_analytic[end] ==
+            full.solutions[i].u[j].u_analytic[end]
+    end
+    @test Base.summarysize(reduced.solutions) < Base.summarysize(full.solutions) / 5
+end
+
+@testset "incompatible options are rejected" begin
+    # these need the whole timeseries, which the reduction throws away
+    @test_throws ArgumentError test_convergence(
+        dts, prob, EM(); save_everystep = false, trajectories = 10,
+        save_solutions = false, weak_timeseries_errors = true
+    )
+    @test_throws ArgumentError test_convergence(
+        dts, prob, EM(); save_everystep = false, trajectories = 10,
+        save_solutions = false, weak_dense_errors = true
+    )
+    # the reduction goes through the ensemble's output_func, so it cannot be applied
+    # on top of an EnsembleProblem the caller built themselves
+    @test_throws ArgumentError test_convergence(
+        dts, EnsembleProblem(prob), EM(); save_everystep = false, trajectories = 10,
+        save_solutions = false
+    )
+    # `expected_value` averages the trajectory values themselves, which the reduction
+    # replaces with a ConvergenceTrajectory
+    @test_throws ArgumentError test_convergence(
+        dts, prob, EM(); save_everystep = false, trajectories = 10,
+        save_solutions = false, expected_value = 0.5
+    )
+end

lib/StochasticDiffEq/test/test_groups.toml

@@ -32,13 +32,11 @@ local_only = true
 
 [OOPWeakConvergence]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 300
 local_only = true
 
 [IIPWeakConvergence]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 300
 local_only = true