Drop vec/flatten/unflatten roundtrips in MIRK/FIRK/MIRKN nlprob construction

[ChrisRackauckas-Claude]

Background

In __perform_mirk_iteration / __perform_firk_iteration / __perform_mirkn_iteration, the structured initial guess cache.y₀::VectorOfArray{Float64, 2, Vector{Vector{Float64}}} is flattened to a 1D vector before being handed to NonlinearSolve:

nlprob = __construct_problem(cache, copy(vec(cache.y₀)), copy(cache.y₀))

(copy(vec(...)) rather than vec(...) is necessary because under RecursiveArrayTools v4, vec(::VectorOfArray) returns Base.ReshapedArray{T, 1, VectorOfArray{...}, …}, which causes NonlinearSolve's polyalg to fail to infer T, N, uType, R of the resulting NonlinearSolution — see #484 / #473 / #485.)

The copy is essentially papering over a deeper structural choice: the entire NL solve pipeline below this point is written to take a flat AbstractVector and round-trip it through the structured VectorOfArray representation:

@views function __mirk_loss!(resid, u, p, y, pt, bc!, residual, mesh, cache, eval_sol, trait, constraint) where {BC}
    y_ = recursive_unflatten!(y, u)            # u (flat) → y_ (VOA)
    resids = [get_tmp(r, u) for r in residual]
    Φ!(resids[2:end], cache, y_, u, trait, constraint)
    update_eval_sol!(eval_sol, y_, cache)
    eval_bc_residual!(resids[1], pt, bc!, eval_sol, p, mesh)
    recursive_flatten!(resid, resids)           # resids (VOA) → resid (flat)
    return nothing
end

Each NL iteration: VOA → flat → solve → flat → VOA via recursive_unflatten! / recursive_flatten!. With RAT v4 making VectorOfArray <: AbstractArray, NonlinearSolve in principle accepts cache.y₀ directly as u0, so the round-trip is no longer load-bearing — it's just historical from when VOA wasn't an AbstractArray.

Proposal

Skip the flatten/unflatten and pass cache.y₀ (the VOA) directly as nlprob.u0. Concretely:

• __perform_*_iteration: __construct_*problem(cache, cache.y₀, copy(cache.y₀)) (drop copy(vec(...))).
• __mirk_loss! / __mirk_loss_bc! / __mirk_loss_collocation! (and MIRKN/FIRK analogues): receive u::AbstractVectorOfArray directly, drop the recursive_unflatten!(y, u) / recursive_flatten!(resid, resids) calls. The structured form is the input.
• safe_similar(y, ...) calls building residual buffers: keep returning a flat Vector{Float64} (residuals stay 1D — recursive_flatten_twopoint! etc. pack into the flat layout the NL solver expects).
• DI prepare_jacobian(loss, resid, diffmode, y, Constant(cache.p)): confirm DI handles a 2D VectorOfArray input correctly — Jacobian shape is (length(resid_flat), length(y_voa)) under linear indexing, which should match what the existing vcat(J_bc, J_c) jac_prototype assembly expects. Some adapter work may be required.

Why bother

1. Eliminates a per-iteration allocation+copy (copy(vec(cache.y₀)) allocates a fresh Vector each iteration; with the round-trip closures also allocating internally).
2. Removes the RAT-v4 inference cliff at the source rather than per-call-site. Right now any new code path that hands a vec(VOA) to NonlinearSolve will silently lose T, N, uType, R inference — as MIRKN did until Fix MIRK/FIRK adaptive mesh refinement under RAT v4 (closes #484) #485.
3. Removes a source of confusion: it's not obvious from reading __perform_mirk_iteration why copy(vec(...)) is required and not vec(...). Comments help, but not having the round-trip is cleaner.

Scope

Probably one well-scoped PR per sublibrary (MIRK, FIRK, MIRKN). Each touches the iteration entry point + the loss function family. Shooting/MultipleShooting separately if useful.

Out of scope here

This issue is about removing the round-trip on the post-cache-construction side. The choice to store cache.y₀ as a VectorOfArray of per-timestep vectors (vs. e.g. a contiguous Matrix{Float64} with column views) is independent and not changed by this proposal.

🤖 Generated with Claude Code

Co-Authored-By: Chris Rackauckas accounts@chrisrackauckas.com

[ChrisRackauckas-Claude]

Took a look at implementing this for MIRK and hit a wall worth flagging here.

The minimal change — drop copy(vec(...)) in __perform_mirk_iteration and replace recursive_unflatten!(y, u) with y_ = u.u in every loss method — compiles fine but blows up during BoundaryValueDiffEqMIRK's @compile_workload (precompile time) the moment solve(prob, MIRK4(); dt = 0.1) reaches NonlinearSolve's polyalg → LinearSolve:

nested task error: TypeError: in setfield!, expected Vector{Float64},
got a value of type Base.ReshapedArray{Float64, 1, RecursiveArrayTools.VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, …}
  [1] setproperty!(cache::LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, NonlinearSolveBase.LinearSolveParameters{…}, …}, …, x::Base.ReshapedArray{Float64, 1, VectorOfArray{Float64, 2, …}, …})

The chain:

1. We pass cache.y₀::VectorOfArray{Float64, 2, Vector{Vector{Float64}}} as u0.
2. NonlinearSolve / NonlinearSolveBase, somewhere on the way to the linear subproblem, calls vec(u) to flatten.
3. Under RAT v4, vec(::VectorOfArray{Float64, 2, Vector{Vector{Float64}}}) returns a Base.ReshapedArray{Float64, 1, VOA, …} (it aliases — modifying the ReshapedArray writes back through to voa.u, verified locally).
4. That ReshapedArray is assigned to a LinearSolve.LinearCache field whose declared type is the concrete Vector{Float64} — setproperty! rejects it.

So the proposal as written can't be done purely inside BVDE. It needs one of:

• Upstream: LinearSolve.LinearCache (and the b/u plumbing it goes through in NonlinearSolveBase) needs to accept any AbstractVector{T}, not specifically Vector{T}. That covers ReshapedArray, plus any other lazy view a downstream caller might pass.
• RAT: change vec(::VectorOfArray) to materialize a fresh Vector{T} rather than wrap. Wider blast radius — anyone currently relying on the aliasing behavior would break.
• Status quo in BVDE: keep copy(vec(cache.y₀)). The copy on its own is what materializes to Vector{Float64}, which is exactly why it was added — without it the same setfield! failure shows up at runtime instead of a different inference cliff.

The middle option seems wrong. The first option is the cleanest match for what this issue is actually asking for: the round-trip is only needed because the rest of the pipeline can't yet take a structured u0. Making LinearCache polymorphic in b/u removes the reason to flatten at all.

I think it's worth resolving the upstream LinearSolve/NonlinearSolveBase typing first; once a VectorOfArray u0 round-trips through solve(::NonlinearProblem) cleanly, the BVDE-side changes here become a fairly mechanical drop-the-vec-and-recursive_unflatten PR per sublibrary as proposed.

Happy to send that follow-up once the upstream side is unblocked. (Did not push the speculative MIRK branch since the precompile workload fails on it — keeping master green.)