init: add jac_prototype

Author: vyudu

src/systems/diffeqs/sdesystem.jl

@@ -593,6 +593,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         u0 = nothing;
         version = nothing, tgrad = false, sparse = false,
         jac = false, Wfact = false, eval_expression = false,
+        sparsity = false, analytic = nothing,
         eval_module = @__MODULE__,
         checkbounds = false, initialization_data = nothing,
         cse = true, kwargs...) where {iip, specialize}
@@ -641,6 +642,17 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         _Wfact, _Wfact_t = nothing, nothing
     end
 
+    jac_prototype = if sparse
+        uElType = u0 === nothing ? Float64 : eltype(u0)
+        if jac
+            similar(calculate_jacobian(sys, sparse = sparse), uElType)
+        else
+            similar(jacobian_sparsity(sys), uElType)
+        end
+    else
+        nothing
+    end
+
     M = calculate_massmatrix(sys)
     _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0', M)
 
@@ -651,10 +663,14 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         tgrad = _tgrad === nothing ? nothing : _tgrad,
+        mass_matrix = _M
+        jac_prototype = jac_prototype,
+        observed = observedfun,
+        sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
+        analytic = analytic,
         Wfact = _Wfact === nothing ? nothing : _Wfact,
         Wfact_t = _Wfact_t === nothing ? nothing : _Wfact_t,
-        mass_matrix = _M, initialization_data,
-        observed = observedfun)
+        initialization_data)
 end
 
 """
@@ -724,6 +740,17 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         _jac = :nothing
     end
 
+    jac_prototype = if sparse
+        uElType = u0 === nothing ? Float64 : eltype(u0)
+        if jac
+            similar(calculate_jacobian(sys, sparse = sparse), uElType)
+        else
+            similar(jacobian_sparsity(sys), uElType)
+        end
+    else
+        nothing
+    end
+
     if Wfact
         tmp_Wfact, tmp_Wfact_t = generate_factorized_W(
             sys, dvs, ps; expression = Val{true},
@@ -743,11 +770,13 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         g = $g
         tgrad = $_tgrad
         jac = $_jac
+        jac_prototype = $jac_prototype
         Wfact = $_Wfact
         Wfact_t = $_Wfact_t
         M = $_M
         SDEFunction{$iip}(f, g,
             jac = jac,
+            jac_prototype = jac_prototype,
             tgrad = tgrad,
             Wfact = Wfact,
             Wfact_t = Wfact_t,