Add ode integrator rng

Project.toml

@@ -175,6 +175,7 @@ SimpleNonlinearSolve = "2.7"
 Static = "1.2"
 StaticArrayInterface = "1.8"
 StaticArrays = "1.9.14"
+StableRNGs = "1"
 StructArrays = "0.6, 0.7"
 Symbolics = "6, 7"
 TruncatedStacktraces = "1.4"
@@ -197,6 +198,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -206,4 +208,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["ComponentArrays", "AlgebraicMultigrid", "IncompleteLU", "DiffEqCallbacks", "DifferentiationInterface", "DiffEqDevTools", "ExplicitImports", "ODEProblemLibrary", "ElasticArrays", "JLArrays", "Random", "SafeTestsets", "StructArrays", "Test", "Unitful", "Pkg", "NLsolve", "RecursiveFactorization", "SparseConnectivityTracer", "SparseMatrixColorings", "Statistics"]
+test = ["ComponentArrays", "AlgebraicMultigrid", "IncompleteLU", "DiffEqCallbacks", "DifferentiationInterface", "DiffEqDevTools", "ExplicitImports", "ODEProblemLibrary", "ElasticArrays", "JLArrays", "Random", "SafeTestsets", "StableRNGs", "StructArrays", "Test", "Unitful", "Pkg", "NLsolve", "RecursiveFactorization", "SparseConnectivityTracer", "SparseMatrixColorings", "Statistics"]

lib/OrdinaryDiffEqCore/Project.toml

@@ -4,6 +4,7 @@ authors = ["ParamThakkar123 <paramthakkar864@gmail.com>"]
 version = "3.6.0"
 
 [deps]
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
@@ -43,7 +44,6 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
@@ -101,7 +101,7 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [targets]
-test = ["DiffEqDevTools", "Random", "SafeTestsets", "SparseArrays", "Test", "Pkg"]
+test = ["DiffEqDevTools", "SafeTestsets", "SparseArrays", "Test", "Pkg"]
 
 [extensions]
 OrdinaryDiffEqCoreMooncakeExt = "Mooncake"

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -42,6 +42,8 @@ import SciMLOperators: AbstractSciMLOperator, AbstractSciMLScalarOperator,
 
 using DiffEqBase: DEIntegrator
 
+import Random
+
 import RecursiveArrayTools: chain, recursivecopy!, recursivecopy, recursive_bottom_eltype, recursive_unitless_bottom_eltype, recursive_unitless_eltype, copyat_or_push!, DiffEqArray, recursivefill!
 
 import RecursiveArrayTools

lib/OrdinaryDiffEqCore/src/integrators/integrator_interface.jl

@@ -360,6 +360,23 @@ end
 
 const EMPTY_ARRAY_OF_PAIRS = Pair[]
 
+SciMLBase.has_rng(::ODEIntegrator) = true
+SciMLBase.get_rng(integrator::ODEIntegrator) = integrator.rng
+function SciMLBase.set_rng!(integrator::ODEIntegrator, rng)
+    R = typeof(integrator.rng)
+    if !isa(rng, R)
+        throw(
+            ArgumentError(
+                "Cannot set RNG of type $(typeof(rng)) on an integrator " *
+                    "whose RNG type parameter is $R. " *
+                    "Construct a new integrator via `init(prob, alg; rng = your_rng)` instead."
+            )
+        )
+    end
+    integrator.rng = rng
+    return nothing
+end
+
 SciMLBase.has_reinit(integrator::ODEIntegrator) = true
 function SciMLBase.reinit!(
         integrator::ODEIntegrator, u0 = integrator.sol.prob.u0;
@@ -374,7 +391,8 @@ function SciMLBase.reinit!(
         reinit_dae = true,
         reinit_callbacks = true, initialize_save = true,
         reinit_cache = true,
-        reinit_retcode = true
+        reinit_retcode = true,
+        rng = nothing
     )
     if reinit_dae && SciMLBase.has_initializeprob(integrator.sol.prob.f)
         # This is `remake` infrastructure. `reinit!` is somewhat like `remake` for
@@ -461,6 +479,10 @@ function SciMLBase.reinit!(
     integrator.erracc = typeof(integrator.erracc)(1)
     integrator.dtacc = typeof(integrator.dtacc)(1)
 
+    if rng !== nothing
+        SciMLBase.set_rng!(integrator, rng)
+    end
+
     if reset_dt
         auto_dt_reset!(integrator)
     end

lib/OrdinaryDiffEqCore/src/integrators/type.jl

@@ -89,7 +89,7 @@ mutable struct ODEIntegrator{
         algType <: Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm}, IIP,
         uType, duType, tType, pType, eigenType, EEstT, QT, tdirType,
         ksEltype, SolType, F, CacheType, O, FSALType, EventErrorType,
-        CallbackCacheType, IA, DV, CC,
+        CallbackCacheType, IA, DV, CC, RNGType,
     } <:
     SciMLBase.AbstractODEIntegrator{algType, IIP, uType, tType}
     sol::SolType
@@ -144,4 +144,5 @@ mutable struct ODEIntegrator{
     differential_vars::DV
     fsalfirst::FSALType
     fsallast::FSALType
+    rng::RNGType
 end

lib/OrdinaryDiffEqCore/src/solve.jl

@@ -79,6 +79,7 @@ function SciMLBase.__init(
         initialize_integrator = true,
         alias = ODEAliasSpecifier(),
         initializealg = DefaultInit(),
+        rng = nothing,
         kwargs...
     )
     if prob isa SciMLBase.AbstractDAEProblem && alg isa OrdinaryDiffEqAlgorithm
@@ -650,6 +651,8 @@ function SciMLBase.__init(
     saveiter_dense = 0
     fsalfirst, fsallast = get_fsalfirstlast(cache, rate_prototype)
 
+    _rng = rng === nothing ? Random.default_rng() : rng
+
     integrator = ODEIntegrator{
         typeof(_alg), isinplace(prob), uType, typeof(du),
         tType, typeof(p),
@@ -659,7 +662,7 @@ function SciMLBase.__init(
         typeof(opts), typeof(fsalfirst),
         typeof(last_event_error), typeof(callback_cache),
         typeof(initializealg), typeof(differential_vars),
-        typeof(controller_cache),
+        typeof(controller_cache), typeof(_rng),
     }(
         sol, u, du, k, t, tType(_dt), f, p,
         uprev, uprev2, duprev, tprev,
@@ -682,7 +685,7 @@ function SciMLBase.__init(
         isout, reeval_fsal,
         u_modified, reinitiailize, isdae,
         opts, stats, initializealg, differential_vars,
-        fsalfirst, fsallast
+        fsalfirst, fsallast, _rng
     )
 
     if initialize_integrator