Fix the solution bookkeeping for save_everystep=false (#428)

With `save_everystep=false`, `savevalues!` skips all of its
ProbNumDiffEq-specific saving, so `sol.x_filt` was left with only the
initial state and `sol.diffusions` was left empty. Two things broke as a
result:

- `calibrate_solution!` zips `sol.pu` (2 entries) against `sol.x_filt`
  (1 entry), so with a calibrated static diffusion the final `sol.pu[end].Σ`
  never got rescaled by the quasi-MLE sigma^2, and sigma^2 was not recoverable
  from the solution either.
- Interpolating such a solution threw a `BoundsError`, since the
  interpolation indexes into the empty `sol.diffusions`. This is what the
  `WorkPrecisionSet` test with a dense `TestSolution` was marked broken for.

`pn_solution_endpoint_match_cur_integrator!` now always saves the endpoint
into `sol.x_filt`, and saves the last step's diffusion when nothing was saved
during the solve. The endpoint used `integ.saveiter_dense` for `x_filt`, which
stays at 1 whenever `dense=false`; it now uses `integ.saveiter` like
`savevalues!` does.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: nathanaelbosch

src/integrator_utils.jl

@@ -30,7 +30,7 @@ function OrdinaryDiffEqCore.postamble!(
         smooth_solution!(integ)
     end
 
-    @assert (length(integ.sol.u) == length(integ.sol.pu))
+    @assert (length(integ.sol.u) == length(integ.sol.pu) == length(integ.sol.x_filt))
 
     return nothing
 end
@@ -58,8 +58,9 @@ function calibrate_solution!(integ, mle_diffusion)
     end
 
     # Re-write into the solution estimates
-    for (pu, x) in zip(integ.sol.pu, integ.sol.x_filt)
-        _gaussian_mul!(pu, integ.cache.SolProj, x)
+    # two-argument `eachindex` errors on a length mismatch instead of truncating
+    for i in eachindex(integ.sol.pu, integ.sol.x_filt)
+        _gaussian_mul!(integ.sol.pu[i], integ.cache.SolProj, integ.sol.x_filt[i])
     end
     # [(su[:] .= pu) for (su, pu) in zip(integ.sol.u, integ.sol.pu.μ)]
 end
@@ -123,22 +124,33 @@ end
 "Inspired by `OrdinaryDiffEqCore.solution_match_cur_integrator!`"
 function pn_solution_endpoint_match_cur_integrator!(integ)
     if integ.opts.save_end
-        if integ.alg.smooth
-            copyat_or_push!(
-                integ.sol.x_filt,
-                integ.saveiter_dense,
-                integ.cache.x,
-            )
+        i = integ.saveiter
+
+        # `savevalues!` saved nothing here, so the last step's diffusion is still missing
+        if !integ.opts.save_everystep && i > 1
+            save_diffusion!(integ.sol, i, integ.cache.local_diffusion)
         end
 
+        copyat_or_push!(integ.sol.x_filt, i, integ.cache.x)
+
         copyat_or_push!(
             integ.sol.pu,
-            integ.saveiter,
+            i,
             _gaussian_mul!(integ.cache.pu_tmp, integ.cache.SolProj, integ.cache.x),
         )
     end
 end
 
+"Save `diffusion` as the `i`-th entry of `sol.diffusions`, appending if necessary."
+function save_diffusion!(sol, i, diffusion)
+    if i <= length(sol.diffusions)
+        sol.diffusions[i] = copy(diffusion)
+    else
+        push!(sol.diffusions, copy(diffusion))
+    end
+    return nothing
+end
+
 "Extends `OrdinaryDiffEqCore._savevalues!` to save ProbNumDiffEq.jl-specific things."
 function DiffEqBase.savevalues!(
     integ::OrdinaryDiffEqCore.ODEIntegrator{<:AbstractEK},
@@ -152,11 +164,7 @@ function DiffEqBase.savevalues!(
     # Save our custom stuff that we need for the posterior
     if integ.opts.save_everystep
         i = integ.saveiter
-        if i <= length(integ.sol.diffusions)
-            integ.sol.diffusions[i] = copy(integ.cache.local_diffusion)
-        else
-            push!(integ.sol.diffusions, copy(integ.cache.local_diffusion))
-        end
+        save_diffusion!(integ.sol, i, integ.cache.local_diffusion)
         copyat_or_push!(integ.sol.x_filt, i, integ.cache.x)
         _gaussian_mul!(integ.cache.pu_tmp, integ.cache.SolProj, integ.cache.x)
         copyat_or_push!(integ.sol.pu, i, integ.cache.pu_tmp)

test/diffeqdevtools.jl

@@ -76,14 +76,16 @@ end
     @test plot(wps) isa AbstractPlot
 end
 
-@testset "WorkPrecisionSet with TestSolution is broken" begin
+@testset "WorkPrecisionSet with TestSolution" begin
     abstols = 1.0 ./ 10.0 .^ (6:7)
     reltols = 1.0 ./ 10.0 .^ (3:4)
     setups = [
         Dict(:alg => EK0(smooth=false))
         Dict(:alg => EK1(smooth=false))
     ]
-    @test_broken wp = WorkPrecisionSet(
+    # This used to throw a `BoundsError`: with `save_everystep=false` the solution was
+    # left without any diffusions, which the interpolation needs; see issue #428
+    wp = WorkPrecisionSet(
         prob, abstols, reltols, setups;
         appxsol=test_sol,
         dense=false,
@@ -92,6 +94,7 @@ end
         maxiters=Int(1e7),
         timeseries_errors=false,
     )
+    @test wp isa WorkPrecisionSet
     @test_nowarn WorkPrecisionSet(
         prob, abstols, reltols, setups;
         appxsol=appxsol_nondense,

test/diffusions.jl

@@ -81,4 +81,28 @@ import ODEProblemLibrary: prob_ode_fitzhughnagumo
         appxsol = appxtrue(sol, true_sol, dense_errors=false)
         @test appxsol.errors[:final] < 1e-5
     end
+
+    # Fixes https://github.com/nathanaelbosch/ProbNumDiffEq.jl/issues/428
+    @testset "`save_everystep=false` returns the same endpoint: $D" for D in (
+        FixedDiffusion(),
+        FixedMVDiffusion(),
+        FixedDiffusion(1e3, false),
+        DynamicDiffusion(),
+        DynamicMVDiffusion(),
+    )
+        alg = EK0(diffusionmodel=D, smooth=false)
+        kwargs = (dense=false, adaptive=false, dt=1e-2)
+        sol_all = solve(prob, alg; save_everystep=true, kwargs...)
+        sol_end = solve(prob, alg; save_everystep=false, kwargs...)
+
+        @test length(sol_end.u) == length(sol_end.pu) == length(sol_end.x_filt) == 2
+        @test length(sol_end.diffusions) == 1
+
+        @test sol_end.pu[end].μ ≈ sol_all.pu[end].μ
+        @test Matrix(sol_end.pu[end].Σ) ≈ Matrix(sol_all.pu[end].Σ)
+        @test sol_end.diffusions[end] ≈ sol_all.diffusions[end]
+
+        # the interpolation reads `sol.diffusions`, which used to be empty here
+        @test length(sol_end(0.5).μ) == length(prob.u0)
+    end
 end