From e4cdfc33942f30f16c5ca2cb6ab3133f9615c5ac Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 28 Oct 2024 23:56:33 -0400 Subject: [PATCH] feat: add unwrapped_f for NonlinearProblem --- src/scimlfunctions.jl | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/src/scimlfunctions.jl b/src/scimlfunctions.jl index faaafe9db..ecb49592e 100644 --- a/src/scimlfunctions.jl +++ b/src/scimlfunctions.jl @@ -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!,