Skip to content

Commit 2977825

Browse files
MIRK: cache the flat nlprob u0 buffer instead of reallocating each iteration
`__perform_mirk_iteration` previously called `__construct_problem(cache, copy(vec(cache.y₀)), copy(cache.y₀))`. The `copy(vec(...))` allocated a fresh `Vector{T}` every outer mesh-refinement step purely to materialize the flat representation NonlinearSolve / LinearSolve require — `vec(::VectorOfArray)` under RAT v4 returns a `Base.ReshapedArray` which `LinearSolve.LinearCache`'s declared `Vector{T}` field doesn't accept. Add a `y₀_flat::Vector{T}` field on `MIRKCache` that mirrors `cache.y₀`, allocate it once in `__init` (`collect(vec(y₀))`), keep it in sync with the mesh in `__expand_cache!` (`resize!(..., Nₙ * cache.M)`), and refresh it in place via `copyto!(cache.y₀_flat, vec(cache.y₀))` at the start of each outer iteration. The flat mirror is then handed to `__construct_problem` as `nlprob.u0`, so NonlinearSolve / LinearSolve still see a concrete `Vector{T}` while the per-iteration allocation disappears. Refs #486. This is the minimal step in that direction that doesn't require upstream changes to LinearSolve / NonlinearSolveBase to accept `AbstractVector{T}` (or to RAT to change what `vec(::VOA)` returns); the deeper refactor — passing `cache.y₀::VectorOfArray` directly through to NonlinearSolve and dropping the `recursive_unflatten!` round-trips inside the `__mirk_loss!` family — needs that upstream work first (see issue #486 comment). Verified by running `Pkg.test()` for `BoundaryValueDiffEqMIRK`: all 25 test items / 308 tests pass. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
1 parent b27fede commit 2977825

1 file changed

Lines changed: 11 additions & 2 deletions

File tree

  • lib/BoundaryValueDiffEqMIRK/src

lib/BoundaryValueDiffEqMIRK/src/mirk.jl

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,10 @@
2121
k_interp # Stage information associated with the discrete Runge-Kutta method
2222
y
2323
y₀
24+
y₀_flat # Flat Vector{T} mirror of y₀ used as nlprob u0 to keep
25+
# LinearSolve / NonlinearSolveBase happy (they require a
26+
# concrete `Vector{T}`, not the `Base.ReshapedArray` that
27+
# `vec(::VectorOfArray)` returns under RAT v4).
2428
residual
2529
# The following 2 caches are never resized
2630
fᵢ_cache
@@ -75,6 +79,7 @@ function SciMLBase.__init(
7579

7680
# Don't flatten this here, since we need to expand it later if needed
7781
y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p; tune_parameters = tune_parameters)
82+
y₀_flat = collect(vec(y₀))
7883

7984
y = __alloc.(copy.(y₀.u))
8085
TU, ITU = constructMIRK(alg, T)
@@ -235,7 +240,7 @@ function SciMLBase.__init(
235240
return MIRKCache{iip, T, use_both, typeof(diffcache), tune_parameters}(
236241
alg_order(alg), stage, N, size(u0), f, bc, prob_, prob.problem_type, prob.p, alg,
237242
TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete, k_interp, y,
238-
y₀, residual, fᵢ_cache, fᵢ₂_cache, errors, new_stages, resid₁_size, prob.singular_term
243+
y₀, y₀_flat, residual, fᵢ_cache, fᵢ₂_cache, errors, new_stages, resid₁_size, prob.singular_term
239244
, nlsolve_kwargs, optimize_kwargs, (; abstol, dt, adaptive, controller, tune_parameters, kwargs...), verbose_spec
240245
)
241246
end
@@ -252,6 +257,7 @@ function __expand_cache!(cache::MIRKCache{iip, T, use_both}) where {iip, T, use_
252257
__resize!(cache.k_interp.u, Nₙ - 1, cache.M)
253258
__resize!(cache.y, Nₙ, cache.M)
254259
__resize!(cache.y₀.u, Nₙ, cache.M)
260+
resize!(cache.y₀_flat, Nₙ * cache.M)
255261
__resize!(cache.residual, Nₙ, cache.M)
256262
__resize!(cache.errors.u, ifelse(use_both, 2 * (Nₙ - 1), (Nₙ - 1)), cache.M)
257263
__resize!(cache.new_stages.u, Nₙ - 1, cache.M)
@@ -307,7 +313,10 @@ function SciMLBase.solve!(
307313
end
308314

309315
function __perform_mirk_iteration(cache::MIRKCache, abstol, adaptive::Bool, controller::AbstractErrorControl)
310-
nlprob = __construct_problem(cache, copy(vec(cache.y₀)), copy(cache.y₀))
316+
# Refresh the flat mirror from the structured guess (in-place; no fresh allocation
317+
# per outer iteration). NonlinearSolve / LinearSolve still see a `Vector{T}`.
318+
copyto!(cache.y₀_flat, vec(cache.y₀))
319+
nlprob = __construct_problem(cache, cache.y₀_flat, copy(cache.y₀))
311320
solve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve, cache.alg.optimize)
312321
kwargs = __concrete_kwargs(
313322
cache.alg.nlsolve, cache.alg.optimize, cache.nlsolve_kwargs, cache.optimize_kwargs,

0 commit comments

Comments
 (0)