fix: move initializeprob typeasserts to OverrideInitData

SciML · Oct 30, 2024 · 4ea5fd0 · 4ea5fd0
1 parent e4562ba
commit 4ea5fd0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 9 deletions.
diff --git a/src/initialization.jl b/src/initialization.jl
@@ -23,4 +23,10 @@ struct OverrideInitData{IProb, UIProb, IProbMap, IProbPmap}
     the parameter object of the original problem.
     """
     initializeprobpmap::IProbPmap
+
+    function OverrideInitData(initprob::I, update_initprob!::J, initprobmap::K, initprobpmap::L) where {I, J, K, L}
+        @assert initprob <: Union{NonlinearProblem, NonlinearLeastSquaresProblem}
+        return new{I, J, K, L}(initprob, update_initprob!, initprobmap, initprobpmap)
+    end
 end
+
diff --git a/src/scimlfunctions.jl b/src/scimlfunctions.jl
@@ -2483,9 +2483,6 @@ function ODEFunction{iip, specialize}(f;
         initialization_data, initializeprob, update_initializeprob!,
         initializeprobmap, initializeprobpmap)
 
-    @assert typeof(initdata.initializeprob) <:
-            Union{Nothing, NonlinearProblem, NonlinearLeastSquaresProblem}
-
     if specialize === NoSpecialize
         ODEFunction{iip, specialize,
             Any, Any, Any, Any,
@@ -2756,8 +2753,6 @@ function SplitFunction{iip, specialize}(f1, f2;
     initdata = reconstruct_initialization_data(
         initialization_data, initializeprob, update_initializeprob!,
         initializeprobmap, initializeprobpmap)
-    @assert typeof(initdata.initializeprob) <:
-            Union{Nothing, NonlinearProblem, NonlinearLeastSquaresProblem}
 
     if specialize === NoSpecialize
         SplitFunction{iip, specialize, Any, Any, Any, Any, Any, Any, Any, Any, Any,
@@ -3445,9 +3440,6 @@ function DAEFunction{iip, specialize}(f;
         initialization_data, initializeprob, update_initializeprob!,
         initializeprobmap, initializeprobpmap)
 
-    @assert typeof(initdata.initializeprob) <:
-            Union{Nothing, NonlinearProblem, NonlinearLeastSquaresProblem}
-
     if specialize === NoSpecialize
         DAEFunction{iip, specialize, Any, Any, Any,
             Any, Any, Any, Any, Any,
@@ -4455,7 +4447,7 @@ end
 function has_initializeprobpmap(f::AbstractSciMLFunction)
     __has_initializeprobpmap(f) && f.initialization_data.initializeprobpmap !== nothing
 end
-function has_initialization_data(f::AbstractSciMLFunction)
+function has_initialization_data(f)
     __has_initialization_data(f) && f.initialization_data !== nothing
 end