Skip to content

Commit 9bdd00a

Browse files
Add initializeprob to splitfunction
1 parent 8b20758 commit 9bdd00a

File tree

1 file changed

+16
-6
lines changed

1 file changed

+16
-6
lines changed

src/scimlfunctions.jl

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -518,7 +518,7 @@ information on generating the SplitFunction from this symbolic engine.
518518
struct SplitFunction{
519519
iip, specialize, F1, F2, TMM, C, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt,
520520
TPJ, O,
521-
TCV, SYS} <: AbstractODEFunction{iip}
521+
TCV, SYS, IProb, IProbMap} <: AbstractODEFunction{iip}
522522
f1::F1
523523
f2::F2
524524
mass_matrix::TMM
@@ -536,6 +536,8 @@ struct SplitFunction{
536536
observed::O
537537
colorvec::TCV
538538
sys::SYS
539+
initializeprob::IProb
540+
initializeprobmap::IProbMap
539541
end
540542

541543
@doc doc"""
@@ -2642,28 +2644,36 @@ function SplitFunction{iip, specialize}(f1, f2;
26422644
DEFAULT_OBSERVED,
26432645
colorvec = __has_colorvec(f1) ? f1.colorvec :
26442646
nothing,
2645-
sys = __has_sys(f1) ? f1.sys : nothing) where {iip,
2647+
sys = __has_sys(f1) ? f1.sys : nothing,
2648+
initializeprob = __has_initializeprob(f) ? f.initializeprob : nothing,
2649+
initializeprobmap = __has_initializeprobmap(f) ? f.initializeprobmap : nothing
2650+
) where {iip,
26462651
specialize
26472652
}
26482653
sys = sys_or_symbolcache(sys, syms, paramsyms, indepsym)
2654+
@assert typeof(initializeprob) <:
2655+
Union{Nothing, NonlinearProblem, NonlinearLeastSquaresProblem}
2656+
26492657
if specialize === NoSpecialize
26502658
SplitFunction{iip, specialize, Any, Any, Any, Any, Any, Any, Any, Any, Any,
26512659
Any, Any, Any, Any, Any,
2652-
Any, Any, Any}(f1, f2, mass_matrix, _func_cache,
2660+
Any, Any, Any, Any, Any}(f1, f2, mass_matrix, _func_cache,
26532661
analytic,
26542662
tgrad, jac, jvp, vjp, jac_prototype,
26552663
sparsity, Wfact, Wfact_t, paramjac,
2656-
observed, colorvec, sys)
2664+
observed, colorvec, sys, initializeprob, initializeprobmap)
26572665
else
26582666
SplitFunction{iip, specialize, typeof(f1), typeof(f2), typeof(mass_matrix),
26592667
typeof(_func_cache), typeof(analytic),
26602668
typeof(tgrad), typeof(jac), typeof(jvp), typeof(vjp),
26612669
typeof(jac_prototype), typeof(sparsity),
26622670
typeof(Wfact), typeof(Wfact_t), typeof(paramjac), typeof(observed),
26632671
typeof(colorvec),
2664-
typeof(sys)}(f1, f2, mass_matrix, _func_cache, analytic, tgrad, jac,
2672+
typeof(sys), typeof(initializeprob), typeof(initializeprobmap)}(f1, f2,
2673+
mass_matrix, _func_cache, analytic, tgrad, jac,
26652674
jvp, vjp, jac_prototype,
2666-
sparsity, Wfact, Wfact_t, paramjac, observed, colorvec, sys)
2675+
sparsity, Wfact, Wfact_t, paramjac, observed, colorvec, sys,
2676+
initializeprob, initializeprobmap)
26672677
end
26682678
end
26692679

0 commit comments

Comments
 (0)