fix: handle u0 and p in kwargs of get_concrete_problem

Author: AayushSabharwal

src/solve.jl

@@ -524,6 +524,11 @@ keyword arguments.
 
 Should be called before the problem is solved, after performing type-promotion on the
 problem.
+
+# 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 +1244,30 @@ 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...)
+    kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
     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...)
+    kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
     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...)
+    kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
     p = get_concrete_p(prob, kwargs)
     u0 = get_concrete_u0(prob, isadapt, nothing, kwargs)
     u0 = promote_u0(u0, p, nothing)
@@ -1281,7 +1289,8 @@ 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...)
+    kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
     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 +1309,8 @@ 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...)
+    kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
     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 +1334,8 @@ 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...)
+    kwargs = (; kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
     p = get_concrete_p(prob, kwargs)
     tspan = get_concrete_tspan(prob, isadapt, kwargs, p)
     u0 = get_concrete_u0(prob, isadapt, tspan[1], kwargs)