Stop convergence studies retaining the trajectories they are not measuring

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 hold a single representative trajectory unless
+    `retain_solutions = true` is passed — see
+    [Memory use of Monte-Carlo studies](@ref).
 
   - `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/DiffEqDevTools/Project.toml

@@ -1,7 +1,7 @@
 name = "DiffEqDevTools"
 uuid = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "3.1.4"
+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,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/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
 

lib/StochasticDiffEqROCK/test/test_groups.toml

@@ -6,7 +6,6 @@ versions = ["lts", "1", "pre"]
 # compute_affected_sublibraries.jl.
 [SROCKC2WeakConvergence]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 240
 local_only = true
 

lib/StochasticDiffEqWeak/test/test_groups.toml

@@ -13,7 +13,6 @@ local_only = true
 
 [WeakConvergence2]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 300
 local_only = true
 
@@ -23,7 +22,6 @@ local_only = true
 
 [WeakConvergence3]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 300
 local_only = true
 
@@ -33,19 +31,16 @@ local_only = true
 
 [WeakConvergence4]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 420
 local_only = true
 
 [WeakConvergence5]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 300
 local_only = true
 
 [WeakConvergence6]
 versions = ["1"]
-runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 300
 local_only = true