Skip to content

Commit

Permalink
feat: add unwrapped_f for NonlinearProblem
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Oct 29, 2024
1 parent 0f8ec1f commit e4cdfc3
Showing 1 changed file with 29 additions and 0 deletions.
29 changes: 29 additions & 0 deletions src/scimlfunctions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2661,6 +2661,35 @@ function NonlinearFunction{iip}(f::ODEFunction) where {iip}
colorvec = f.colorvec)
end

function unwrapped_f(f::NonlinearFunction, newf = unwrapped_f(f.f))
if specialization(f) === NoSpecialize
return NonlinearFunction{isinplace(f), specialization(f), Any, Any,
Any, Any, Any, Any, Any, Any, Any,
Any, Any, Any, Any, Any,
typeof(f.colorvec), Any, Any}(newf, f.mass_matrix,
f.analytic, f.tgrad, f.jac,
f.jvp, f.vjp, f.jac_prototype,
f.sparsity, f.Wfact,
f.Wfact_t, f.paramjac,
f.observed, f.colorvec, f.sys,
f.resid_prototype)
else
return NonlinearFunction{isinplace(f), specialization(f), typeof(newf),
typeof(f.mass_matrix), typeof(f.analytic), typeof(f.tgrad),
typeof(f.jac), typeof(f.jvp), typeof(f.vjp), typeof(f.jac_prototype),
typeof(f.sparsity), typeof(f.Wfact), typeof(f.Wfact_t),
typeof(f.paramjac),
typeof(f.observed), typeof(f.colorvec),
typeof(f.sys), typeof(f.resid_prototype)}(newf, f.mass_matrix,
f.analytic, f.tgrad, f.jac,
f.jvp, f.vjp, f.jac_prototype,
f.sparsity, f.Wfact,
f.Wfact_t, f.paramjac,
f.observed, f.colorvec, f.sys,
f.resid_prototype)
end
end

@add_kwonly function SplitFunction(f1, f2, mass_matrix, cache, analytic, tgrad, jac, jvp,
vjp, jac_prototype, W_prototype, sparsity, Wfact, Wfact_t, paramjac,
observed, colorvec, sys, initializeprob, update_initializeprob!,
Expand Down

0 comments on commit e4cdfc3

Please sign in to comment.