ITADN

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

#486OpenChrisRackauckas-Claude 创建于 2026-05-01
## 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: ```julia 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: ```julia @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 #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](https://claude.com/claude-code) Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
0 条评论