Merge pull request #490 from avik-pal/ap/resid_prototype

Add resid_prototype to NonlinearFunction
SciML · Sep 8, 2023 · 03cdaed · 03cdaed
2 parents 519cb4f + 8719f47
commit 03cdaed
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "SciMLBase"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 authors = ["Chris Rackauckas <[email protected]> and contributors"]
-version = "1.96.2"
+version = "1.97.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/src/scimlfunctions.jl b/src/scimlfunctions.jl
@@ -1884,10 +1884,7 @@ For more details on this argument, see the ODEFunction documentation.
 The fields of the NonlinearFunction type directly match the names of the inputs.
 """
 struct NonlinearFunction{iip, specialize, F, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt,
-    TPJ,
-    S, S2, O, TCV,
-    SYS,
-} <: AbstractNonlinearFunction{iip}
+    TPJ, S, S2, O, TCV, SYS, RP} <: AbstractNonlinearFunction{iip}
     f::F
     mass_matrix::TMM
     analytic::Ta
@@ -1905,6 +1902,7 @@ struct NonlinearFunction{iip, specialize, F, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP,
     observed::O
     colorvec::TCV
     sys::SYS
+    resid_prototype::RP
 end
 
 TruncatedStacktraces.@truncate_stacktrace NonlinearFunction 1 2
@@ -3648,7 +3646,8 @@ function NonlinearFunction{iip, specialize}(f;
                DEFAULT_OBSERVED_NO_TIME,
     colorvec = __has_colorvec(f) ? f.colorvec :
                nothing,
-    sys = __has_sys(f) ? f.sys : nothing) where {
+    sys = __has_sys(f) ? f.sys : nothing,
+    resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing) where {
     iip,
     specialize,
 }
@@ -3687,14 +3686,14 @@ function NonlinearFunction{iip, specialize}(f;
             Any, Any, Any, Any, Any,
             Any, Any, Any, Any, Any,
             Any, Any, typeof(syms), typeof(paramsyms), Any,
-            typeof(_colorvec), Any}(f, mass_matrix,
+            typeof(_colorvec), Any, Any}(f, mass_matrix,
             analytic, tgrad, jac,
             jvp, vjp,
             jac_prototype,
             sparsity, Wfact,
             Wfact_t, paramjac,
             syms, paramsyms, observed,
-            _colorvec, sys)
+            _colorvec, sys, resid_prototype)
     else
         NonlinearFunction{iip, specialize,
             typeof(f), typeof(mass_matrix), typeof(analytic), typeof(tgrad),
@@ -3703,13 +3702,13 @@ function NonlinearFunction{iip, specialize}(f;
             typeof(Wfact_t), typeof(paramjac), typeof(syms),
             typeof(paramsyms),
             typeof(observed),
-            typeof(_colorvec), typeof(sys)}(f, mass_matrix, analytic, tgrad,
-            jac,
+            typeof(_colorvec), typeof(sys), typeof(resid_prototype)}(f, mass_matrix,
+            analytic, tgrad, jac,
             jvp, vjp, jac_prototype, sparsity,
             Wfact,
             Wfact_t, paramjac, syms,
             paramsyms,
-            observed, _colorvec, sys)
+            observed, _colorvec, sys, resid_prototype)
     end
 end
 
@@ -3978,6 +3977,7 @@ __has_analytic(f) = isdefined(f, :analytic)
 __has_colorvec(f) = isdefined(f, :colorvec)
 __has_sys(f) = isdefined(f, :sys)
 __has_analytic_full(f) = isdefined(f, :analytic_full)
+__has_resid_prototype(f) = isdefined(f, :resid_prototype)
 
 # compatibility
 has_invW(f::AbstractSciMLFunction) = false