Merge pull request #84 from AayushSabharwal/as/initialization

feat: add initialization support
SciML · Dec 14, 2024 · 4c1ae36 · 4c1ae36
2 parents 4912067 + 662e455
commit 4c1ae36
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 18 deletions.
diff --git a/Project.toml b/Project.toml
@@ -36,7 +36,7 @@ Reexport = "1.0"
 SciMLBase = "2.59.2"
 SparseArrays = "1.9"
 StaticArrays = "1.0"
-StochasticDiffEq = "6.19"
+StochasticDiffEq = "6.72.1"
 UnPack = "0.1, 1.0"
 julia = "1.9"
 

diff --git a/src/functionwrapper.jl b/src/functionwrapper.jl
@@ -9,7 +9,7 @@ end
 (g::SDEDiffusionTermWrapper{false})(u, p, t) = g.g(u, g.h, p, t)
 
 struct SDEFunctionWrapper{iip, F, G, H, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt, TPJ, GG,
-                          TCV} <: DiffEqBase.AbstractRODEFunction{iip}
+                          TCV, ID, S} <: DiffEqBase.AbstractRODEFunction{iip}
     f::F
     g::G
     h::H
@@ -26,6 +26,8 @@ struct SDEFunctionWrapper{iip, F, G, H, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, T
     paramjac::TPJ
     ggprime::GG
     colorvec::TCV
+    initialization_data::ID
+    sys::S
 end
 
 (f::SDEFunctionWrapper{true})(du, u, p, t) = f.f(du, u, f.h, p, t)
@@ -53,17 +55,9 @@ function wrap_functions_and_history(f::SDDEFunction, g, h)
                        typeof(f.analytic), typeof(f.tgrad), typeof(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.ggprime), typeof(f.colorvec)}(f.f, gwh, h,
-                                                                              f.mass_matrix,
-                                                                              f.analytic,
-                                                                              f.tgrad, jac,
-                                                                              f.jvp, f.vjp,
-                                                                              f.jac_prototype,
-                                                                              f.sparsity,
-                                                                              f.Wfact,
-                                                                              f.Wfact_t,
-                                                                              f.paramjac,
-                                                                              f.ggprime,
-                                                                              f.colorvec),
+                       typeof(f.ggprime), typeof(f.colorvec), typeof(f.initialization_data),
+                       typeof(f.sys)}(f.f, gwh, h, f.mass_matrix, f.analytic, f.tgrad, jac,
+                       f.jvp, f.vjp, f.jac_prototype, f.sparsity, f.Wfact, f.Wfact_t, f.paramjac,
+                       f.ggprime, f.colorvec, f.initialization_data, f.sys),
     gwh
 end
diff --git a/src/integrators/type.jl b/src/integrators/type.jl
@@ -26,7 +26,7 @@ end
 mutable struct SDDEIntegrator{algType, IIP, uType, uEltype, tType, P, eigenType,
                               tTypeNoUnits, uEltypeNoUnits, randType, randType2, rateType,
                               solType, cacheType, F, G, F6, OType, noiseType,
-                              EventErrorType, CallbackCacheType, H, IType} <:
+                              EventErrorType, CallbackCacheType, H, IType, IA} <:
                AbstractSDDEIntegrator{algType, IIP, uType, tType}
     f::F
     g::G
@@ -81,6 +81,7 @@ mutable struct SDDEIntegrator{algType, IIP, uType, uEltype, tType, P, eigenType,
     history::H
     stats::DiffEqBase.Stats
     integrator::IType # history integrator
+    initializealg::IA
 end
 
 function (integrator::SDDEIntegrator)(t, deriv::Type = Val{0}; idxs = nothing)

diff --git a/src/solve.jl b/src/solve.jl
@@ -66,7 +66,8 @@ function DiffEqBase.__init(prob::AbstractSDDEProblem,# TODO DiffEqBasee.Abstract
                            #  Keywords for Delay problems (from DDE)
                            discontinuity_interp_points::Int = 10,
                            discontinuity_abstol = eltype(prob.tspan)(1 // Int64(10)^12),
-                           discontinuity_reltol = 0, kwargs...) where {recompile_flag}
+                           discontinuity_reltol = 0,
+                           initializealg = StochasticDiffEq.SDEDefaultInit(), kwargs...) where {recompile_flag}
 
     # alg = getalg(alg0);
     if prob.f isa Tuple
@@ -468,7 +469,7 @@ function DiffEqBase.__init(prob::AbstractSDDEProblem,# TODO DiffEqBasee.Abstract
                                 typeof(c),
                                 typeof(opts), typeof(noise), typeof(last_event_error),
                                 typeof(callback_cache), typeof(history),
-                                typeof(sde_integrator)}(f_with_history,
+                                typeof(sde_integrator), typeof(initializealg)}(f_with_history,
                                                         g_with_history, c, noise, uprev,
                                                         tprev,
                                                         order_discontinuity_t0,
@@ -486,9 +487,10 @@ function DiffEqBase.__init(prob::AbstractSDDEProblem,# TODO DiffEqBasee.Abstract
                                                         P,
                                                         opts, iter, success_iter, eigen_est,
                                                         EEst, q, QT(qoldinit), q11, history,
-                                                        stats, sde_integrator)
+                                                        stats, sde_integrator, initializealg)
 
     if initialize_integrator
+        DiffEqBase.initialize_dae!(integrator)
         StochasticDiffEq.initialize_callbacks!(integrator, initialize_save)
         initialize!(integrator, integrator.cache)