Merge pull request #1143 from AayushSabharwal/as/pre-solve-hook

fix: handle `u0` and `p` in `kwargs` of `get_concrete_problem`

Author: ChrisRackauckas

src/solve.jl

@@ -523,7 +523,13 @@ an updated `prob` to be used for solving. All implementations should accept arbi
 keyword arguments.
 
 Should be called before the problem is solved, after performing type-promotion on the
-problem.
+problem. If the returned problem is not `===` the provided `prob`, it is assumed to
+contain the `u0` and `p` passed as keyword arguments.
+
+# Keyword Arguments
+
+- `u0`, `p`: Override values for `state_values(prob)` and `parameter_values(prob)` which
+  should be used instead of the ones in `prob`.
 """
 function get_updated_symbolic_problem(indp, prob; kw...)
     return prob
@@ -1239,27 +1245,36 @@ function checkkwargs(kwargshandle; kwargs...)
 end
 
 function get_concrete_problem(prob::AbstractJumpProblem, isadapt; kwargs...)
-    get_updated_symbolic_problem(_get_root_indp(prob), prob)
+    get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)
 end
 
 function get_concrete_problem(prob::SteadyStateProblem, isadapt; kwargs...)
-    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob)
+    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)
+    if prob !== prob
+        kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
+    end
     p = get_concrete_p(prob, kwargs)
     u0 = get_concrete_u0(prob, isadapt, Inf, kwargs)
     u0 = promote_u0(u0, p, nothing)
     remake(prob; u0 = u0, p = p)
 end
 
 function get_concrete_problem(prob::NonlinearProblem, isadapt; kwargs...)
-    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob)
+    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)
+    if prob !== prob
+        kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
+    end
     p = get_concrete_p(prob, kwargs)
     u0 = get_concrete_u0(prob, isadapt, nothing, kwargs)
     u0 = promote_u0(u0, p, nothing)
     remake(prob; u0 = u0, p = p)
 end
 
 function get_concrete_problem(prob::NonlinearLeastSquaresProblem, isadapt; kwargs...)
-    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob)
+    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)
+    if prob !== prob
+        kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
+    end
     p = get_concrete_p(prob, kwargs)
     u0 = get_concrete_u0(prob, isadapt, nothing, kwargs)
     u0 = promote_u0(u0, p, nothing)
@@ -1281,7 +1296,10 @@ function init(prob::PDEProblem, alg::AbstractDEAlgorithm, args...;
 end
 
 function get_concrete_problem(prob, isadapt; kwargs...)
-    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob)
+    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)
+    if prob !== prob
+        kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
+    end
     p = get_concrete_p(prob, kwargs)
     tspan = get_concrete_tspan(prob, isadapt, kwargs, p)
     u0 = get_concrete_u0(prob, isadapt, tspan[1], kwargs)
@@ -1300,7 +1318,10 @@ function get_concrete_problem(prob, isadapt; kwargs...)
 end
 
 function get_concrete_problem(prob::DAEProblem, isadapt; kwargs...)
-    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob)
+    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)
+    if prob !== prob
+        kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
+    end
     p = get_concrete_p(prob, kwargs)
     tspan = get_concrete_tspan(prob, isadapt, kwargs, p)
     u0 = get_concrete_u0(prob, isadapt, tspan[1], kwargs)
@@ -1324,7 +1345,10 @@ function get_concrete_problem(prob::DAEProblem, isadapt; kwargs...)
 end
 
 function get_concrete_problem(prob::DDEProblem, isadapt; kwargs...)
-    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob)
+    prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)
+    if prob !== prob
+        kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
+    end
     p = get_concrete_p(prob, kwargs)
     tspan = get_concrete_tspan(prob, isadapt, kwargs, p)
     u0 = get_concrete_u0(prob, isadapt, tspan[1], kwargs)