Stop convergence studies retaining the trajectories they are not measuring

`test_convergence` kept 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 died 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.

Trajectories are no longer retained by default. Three things compose:

  - The `EnsembleProblem` is built 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, not just 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.
  - `calculate_ensemble_errors` moves inside the solve loop, so each step
    size's ensemble is summarised and released before the next is solved
    rather than all of them being held at once.
  - `_drop_trajectories` strips the summarised ensemble to one representative,
    since the reduced objects are dead weight once the errors are computed.

`retain_solutions = true` restores the old behaviour for callers that want the
paths themselves. The reduction is skipped, and the solutions retained, where
it cannot apply: a caller-supplied `EnsembleProblem` owns its own `output_func`,
`expected_value` forms the weak error from the trajectory values directly, and
`weak_timeseries_errors`/`weak_dense_errors` need the whole timeseries.

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 weak files, 16 GiB cgroup, JULIA_NUM_THREADS=2:

  weak_srockc2.jl   OOM 15.61 GiB @3m29s   -> exit 0, 23m41s, 2.81 GiB
  oop_weak.jl       OOM 15.62 GiB @23m50s  -> exit 0, 13m08s, 1.32 GiB
  additive_weak.jl  (same group)           -> exit 0, 3m22s,  1.04 GiB

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

With no group needing more than 2.81 GiB, remove the `high-memory` runner pin
from all eight groups that carried it. The five StochasticDiffEqWeak groups
needed no 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.

BREAKING: DiffEqDevTools 3.1.5 -> 4.0.0. `ConvergenceSimulation.solutions` no
longer holds every trajectory for the Monte-Carlo method, and `sim[i, j]`
forwards into it, so `j > 1` now reaches past the representative. The field
documentation is updated and the compat pin moves "3" -> "4" in 40 Project.toml
files. The ODE/DAE methods do not take the keyword and are unchanged.

The reduce-by-default design and `drop_trajectories` are taken from #4056,
which is closed in favour of this.

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

@@ -20,7 +20,10 @@ A type which holds the data from a convergence simulation.
 
 ### Fields
 
-  - `solutions::Array{<:DESolution}`: Holds all the PdeSolutions.
+  - `solutions::Array{<:DESolution}`: Holds the solutions. For the Monte-Carlo method of
+    `test_convergence` these are reduced by default and hold a single representative
+    trajectory — see [Memory use of Monte-Carlo studies](@ref) — so pass
+    `retain_solutions = true` if the trajectories themselves are needed.
 
   - `errors`: Dictionary of the error calculations. Can contain:
 
@@ -92,10 +95,34 @@ 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. A full
+solution per trajectory carries the solver cache, the noise process, the problem and its
+interpolation, so keeping them makes a study over many trajectories retain far more than
+the error estimates it reports — at hundreds of thousands of trajectories, tens of
+gigabytes.
+
+They are therefore **not** kept by default. Each trajectory is reduced to a
+[`ConvergenceTrajectory`](@ref) by the ensemble's `output_func` as it is solved, so the
+full solutions never coexist, and once the errors have been computed the ensemble is
+stripped to a single representative trajectory. `errors`, `weak_errors`, `error_means`
+and `𝒪est` come from the full ensemble and are unaffected — what changes is that
+`solutions[i].u` holds one entry rather than `trajectories` of them.
+
+Pass `retain_solutions = true` when the trajectories themselves are the point, such as
+comparing two algorithms path by path. `sim[i, j]` forwards into `solutions[i][j]`, so
+it too reaches only the representative trajectory unless the solutions are retained.
+
+The reduction is skipped where it cannot apply, and the solutions retained: when you
+supply your own `EnsembleProblem` (set its `output_func` instead), when `expected_value`
+is given, and when `weak_timeseries_errors` or `weak_dense_errors` is set.
+
 ## API
 
 ```@docs
 DiffEqDevTools.ConvergenceSimulation
+DiffEqDevTools.ConvergenceTrajectory
 DiffEqDevTools.test_convergence
 DiffEqDevTools.analyticless_test_convergence
 ```

lib/DelayDiffEq/Project.toml

@@ -64,7 +64,7 @@ DDEProblemLibrary = "0.1"
 BinaryHeaps = "1"
 DiffEqBase = "7"
 DiffEqCallbacks = "4.17.0"
-DiffEqDevTools = "3"
+DiffEqDevTools = "4"
 DiffEqNoiseProcess = "5.30.0"
 FastBroadcast = "1.3"
 FiniteDiff = "2.27"

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 = "4.0.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,66 @@ 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 unless `retain_solutions = true`, 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
+
+"""
+    _drop_trajectories(sim::EnsembleTestSolution) -> EnsembleTestSolution
+
+Discard every trajectory but the first, keeping the error statistics that were already
+computed from the full ensemble.
+
+Reducing at `output_func` bounds the peak, but the reduced trajectories still add up
+across step sizes, and once `calculate_ensemble_errors` has run they are dead weight —
+a [`ConvergenceSimulation`](@ref) reads only the error dictionaries. One is kept rather
+than none so that anything reaching for a representative trajectory still finds one.
+"""
+function _drop_trajectories(sim::SciMLBase.EnsembleTestSolution{T, N}) where {T, N}
+    u = sim.u[1:min(1, length(sim.u))]
+    return SciMLBase.EnsembleTestSolution{T, N, typeof(u)}(
+        u, sim.errors, sim.weak_errors, sim.error_means, sim.error_medians,
+        sim.elapsedTime, sim.converged
+    )
+end
+
 """
     test_convergence(dts, prob, alg[, ensemblealg]; kwargs...)
     test_convergence(probs, convergence_axis, alg; kwargs...)
@@ -66,6 +125,31 @@ 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
+
+  - `retain_solutions`: whether `ConvergenceSimulation.solutions` keeps the trajectory
+    solutions (default `false`). A Monte-Carlo study over many trajectories otherwise
+    retains one complete solution object — solver cache, noise process, problem and
+    interpolation — per trajectory per step size, which for large `trajectories` is
+    orders of magnitude more memory than the error estimates it is computing, and is
+    what put the weak-convergence test groups past a 16 GB runner.
+
+    By default each trajectory is reduced to a [`ConvergenceTrajectory`](@ref) by the
+    ensemble's `output_func` as it is solved, so the full solutions never coexist, and
+    the reduced ensemble is stripped to one representative trajectory once the errors
+    have been computed from it. `errors`, `weak_errors`, `error_means` and `𝒪est` are
+    computed from the full ensemble either way and are unaffected; what changes is that
+    `solutions[i].u` holds one entry rather than `trajectories` of them.
+
+    Pass `true` when the trajectories themselves are the point, for example to compare
+    two algorithms path by path.
+
+    The reduction is skipped, and the solutions retained, where it cannot apply: when
+    you supply your own `EnsembleProblem` (set its `output_func` instead), when
+    `expected_value` is given (the weak error is formed from the trajectory values
+    directly), and when `weak_timeseries_errors` or `weak_dense_errors` is set (both
+    need the whole timeseries).
 """
 function test_convergence(
         dts::AbstractArray,
@@ -77,12 +161,19 @@ 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, retain_solutions = false, kwargs...
     )
     N = length(dts)
 
+    reduce_trajectories = !retain_solutions &&
+        !(prob isa AbstractEnsembleProblem) &&
+        expected_value === nothing &&
+        !weak_timeseries_errors && !weak_dense_errors
+
     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
@@ -98,20 +189,24 @@ function test_convergence(
             kwargs...
         )
         @info "dt: $(dts[i]) ($i/$N)"
-        _solutions[i] = sol
+        # Summarise each ensemble before solving the next, so that only one step size's
+        # trajectories are ever reachable rather than the whole study's.
+        _solutions[i] = if expected_value === nothing
+            summarised = SciMLBase.calculate_ensemble_errors(
+                sol;
+                weak_timeseries_errors = weak_timeseries_errors,
+                weak_dense_errors = weak_dense_errors
+            )
+            reduce_trajectories ? _drop_trajectories(summarised) : summarised
+        else
+            sol
+        end
     end
 
     auxdata = Dict("dts" => dts)
 
     if expected_value == nothing
-        solutions = [
-            SciMLBase.calculate_ensemble_errors(
-                    sim;
-                    weak_timeseries_errors = weak_timeseries_errors,
-                    weak_dense_errors = weak_dense_errors
-                )
-                for sim in _solutions
-        ]
+        solutions = _solutions
         # Now Calculate Weak Errors
         additional_errors = Dict()
         for k in keys(solutions[1].weak_errors)

lib/DiffEqDevTools/test/retain_solutions_tests.jl

@@ -0,0 +1,92 @@
+using StochasticDiffEq, DiffEqDevTools, Random, Test
+
+# The Monte-Carlo method of `test_convergence` used to keep one full solution object per
+# trajectory per step size, which is orders of magnitude more than the error estimates
+# need and is what put the weak-convergence groups past a 16 GB runner. Trajectories are
+# now reduced by default: an `output_func` turns each into a `ConvergenceTrajectory` as
+# it is solved, and the ensemble is stripped to one representative once the errors have
+# been computed. The statistics must be identical either way.
+
+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(; kwargs...)
+    Random.seed!(100)
+    return test_convergence(
+        dts, prob, EM(); save_everystep = false, trajectories = ntraj,
+        weak_timeseries_errors = false, kwargs...
+    )
+end
+
+reduced = run_sim()                            # the new default
+full = run_sim(retain_solutions = true)
+
+@testset "statistics 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
+    for i in eachindex(dts)
+        @test full.solutions[i].weak_errors == reduced.solutions[i].weak_errors
+        @test full.solutions[i].error_means == reduced.solutions[i].error_means
+        # the per-trajectory error vectors are computed from the full ensemble either way
+        for k in keys(full.solutions[i].errors)
+            @test full.solutions[i].errors[k] ≈ reduced.solutions[i].errors[k]
+            @test length(reduced.solutions[i].errors[k]) == ntraj
+        end
+    end
+end
+
+@testset "trajectories are dropped by default" begin
+    for i in eachindex(dts)
+        @test length(reduced.solutions[i].u) == 1
+        @test length(full.solutions[i].u) == ntraj
+    end
+    @test reduced.solutions[1].u[1] isa ConvergenceTrajectory
+    @test full.solutions[1].u[1] isa SciMLBase.AbstractRODESolution
+    # the representative keeps the endpoints the weak error is formed from
+    @test reduced.solutions[1].u[1].u[end] == full.solutions[1].u[1].u[end]
+    @test reduced.solutions[1].u[1].u_analytic[end] == full.solutions[1].u[1].u_analytic[end]
+    @test Base.summarysize(reduced.solutions) < Base.summarysize(full.solutions) / 20
+end
+
+@testset "reduction is skipped where it cannot apply" begin
+    # these need the full timeseries, so the solutions are retained rather than the
+    # call failing
+    for kwargs in ((; weak_timeseries_errors = true), (; weak_dense_errors = true))
+        sim = run_sim(; kwargs...)
+        @test length(sim.solutions[1].u) == ntraj
+    end
+    # a caller-supplied EnsembleProblem owns its own output_func
+    Random.seed!(100)
+    sim = test_convergence(
+        dts, EnsembleProblem(prob), EM(); save_everystep = false,
+        trajectories = ntraj, weak_timeseries_errors = false
+    )
+    @test length(sim.solutions[1].u) == ntraj
+
+    # `expected_value` averages the trajectory values themselves, so it is used with an
+    # EnsembleProblem that reduces to the quantity being averaged (as PL1WM.jl does).
+    # That path must keep every reduced value, not one representative.
+    Random.seed!(100)
+    ens = EnsembleProblem(prob; output_func = (sol, ctx) -> (sol.u[end][1], false))
+    sim = test_convergence(
+        dts, ens, EM(); save_everystep = false, save_start = false,
+        trajectories = ntraj, weak_timeseries_errors = false,
+        expected_value = 1 / 2 * exp(1.01)
+    )
+    @test length(sim.solutions[1].u) == ntraj
+    @test sim.solutions[1].u[1] isa Float64
+    @test haskey(sim.𝒪est, :weak_final)
+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 retain_solutions" begin
+        include("retain_solutions_tests.jl")
+    end
     @time @testset "Stability Region Tests" begin
         include("stability_region_test.jl")
     end

lib/OrdinaryDiffEqAdamsBashforthMoulton/Project.toml

@@ -34,7 +34,7 @@ Pkg = "1"
 Test = "<0.0.1, 1"
 FastBroadcast = "1.3"
 Random = "<0.0.1, 1"
-DiffEqDevTools = "3"
+DiffEqDevTools = "4"
 MuladdMacro = "0.2.4"
 SciMLBase = "3.39"
 OrdinaryDiffEqCore = "4"

lib/OrdinaryDiffEqBDF/Project.toml

@@ -53,7 +53,7 @@ ForwardDiff = "1.3.3"
 Test = "<0.0.1, 1"
 FastBroadcast = "1.3"
 Random = "<0.0.1, 1"
-DiffEqDevTools = "3"
+DiffEqDevTools = "4"
 DifferentiationInterface = "0.7.18"
 MuladdMacro = "0.2.4"
 LinearSolve = "5.1"

lib/OrdinaryDiffEqCore/Project.toml

@@ -58,7 +58,7 @@ ArrayInterface = "7.28"
 BinaryHeaps = "1"
 CommonSolve = "0.2.6"
 DiffEqBase = "7.8"
-DiffEqDevTools = "3"
+DiffEqDevTools = "4"
 DifferentiationInterface = "0.7.18"
 DocStringExtensions = "0.9.5"
 EnumX = "1.0.4"