Preserve FIRK nested solve type boundaries

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

Author: ChrisRackauckas

lib/BoundaryValueDiffEqFIRK/src/collocation.jl

@@ -12,6 +12,11 @@ function Φ!(residual, cache::FIRKCacheNested, y, u, trait, constraint)
     )
 end
 
+@inline function __nested_stage_parameter(cache::FIRKCacheNested, p)
+    length(cache.bcresid_prototype) < cache.M && return nodual_value(p)
+    return p
+end
+
 @views function Φ!(
         residual, fᵢ_cache, k_discrete, f!, TU::FIRKTableau{false}, y, u, p,
         mesh, mesh_dt, stage::Int, f_prototype, singular_term, ::DiffCacheNeeded, ::Val{true}
@@ -195,7 +200,7 @@ end
 
         K = get_tmp(k_discrete[i], u)
 
-        _nestprob = remake(nest_prob, p = nodual_value(nestprob_p))
+        _nestprob = remake(nest_prob, p = __nested_stage_parameter(cache, nestprob_p))
         nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
         @. K = nestsol.u
         @. residᵢ = yᵢ₊₁ - yᵢ
@@ -227,7 +232,7 @@ end
 
         K = get_tmp(k_discrete[i], u)
 
-        _nestprob = remake(nest_prob, p = nodual_value(nestprob_p))
+        _nestprob = remake(nest_prob, p = __nested_stage_parameter(cache, nestprob_p))
         nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
         @. K = nestsol.u
         @. residᵢ = yᵢ₊₁ - yᵢ
@@ -385,7 +390,7 @@ end
         nestprob_p[2] = T(mesh_dt[i])
         nestprob_p[3:end] = yᵢ
 
-        _nestprob = remake(nest_prob, p = nestprob_p)
+        _nestprob = remake(nest_prob, p = __nested_stage_parameter(cache, nestprob_p))
         nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
 
         @. residᵢ = yᵢ₊₁ - yᵢ

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -1393,15 +1393,23 @@ function __construct_problem(
     )
 end
 
+@inline function __firk_eval_sol!(eval_sol::EvalSol, y_, cache)
+    u = __restructure_sol(y_, cache.in_size)
+    if eltype(first(u)) <: eltype(first(eval_sol.u))
+        eval_sol.u[1:end] .= u
+        return eval_sol
+    end
+    return EvalSol(u, eval_sol.t, cache)
+end
+
 @views function __firk_loss!(
         resid, u, p, y, pt::StandardBVProblem, bc!::BC, residual, mesh,
         cache, eval_sol, trait::DiffCacheNeeded, constraint
     ) where {BC}
     y_ = recursive_unflatten!(y, u)
     resids = [get_tmp(r, u) for r in residual]
     Φ!(resids[2:end], cache, y_, u, trait, constraint)
-    eval_sol.u[1:end] .= y_
-    eval_bc_residual!(resids[1], pt, bc!, eval_sol, p, mesh)
+    eval_bc_residual!(resids[1], pt, bc!, __firk_eval_sol!(eval_sol, y_, cache), p, mesh)
     recursive_flatten!(resid, resids)
     return nothing
 end
@@ -1413,8 +1421,7 @@ end
     y_ = recursive_unflatten!(y, u)
     resids = [r for r in residual]
     Φ!(resids[2:end], cache, y_, u, trait, constraint)
-    eval_sol.u[1:end] .= y_
-    eval_bc_residual!(resids[1], pt, bc!, eval_sol, p, mesh)
+    eval_bc_residual!(resids[1], pt, bc!, __firk_eval_sol!(eval_sol, y_, cache), p, mesh)
     recursive_flatten!(resid, resids)
     return nothing
 end
@@ -1473,8 +1480,7 @@ end
         u, p, y, pt::StandardBVProblem, bc::BC, mesh, cache, eval_sol, trait
     ) where {BC}
     y_ = recursive_unflatten!(y, u)
-    eval_sol.u[1:end] .= y_
-    resid_bc = eval_bc_residual(pt, bc, eval_sol, p, mesh)
+    resid_bc = eval_bc_residual(pt, bc, __firk_eval_sol!(eval_sol, y_, cache), p, mesh)
     resid_co = Φ(cache, y_, u, trait)
     return vcat(resid_bc, mapreduce(vec, vcat, resid_co))
 end