Merge pull request #524 from yonatanwesen/yd/pass-kwargs

kwargs are passed when converting ODEProblem to SteadyStateProblem and NonlinearProblem
SciML · Oct 17, 2023 · 578352c · 578352c
2 parents 7afe983 + 2c13b5d
commit 578352c
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 4 deletions.
diff --git a/src/problems/basic_problems.jl b/src/problems/basic_problems.jl
@@ -260,8 +260,8 @@ struct NonlinearProblem{uType, isinplace, P, F, K, PT} <:
     `isinplace` optionally sets whether the function is inplace or not.
     This is determined automatically, but not inferred.
     """
-    function NonlinearProblem{iip}(f, u0, p = NullParameters()) where {iip}
-        NonlinearProblem{iip}(NonlinearFunction{iip}(f), u0, p)
+    function NonlinearProblem{iip}(f, u0, p = NullParameters(); kwargs...) where {iip}
+        NonlinearProblem{iip}(NonlinearFunction{iip}(f), u0, p; kwargs...)
     end
 end
 
@@ -309,7 +309,7 @@ this is interpreted in the form of the steady state problem, i.e.
 find the ODE's solution at time ``t = \\infty``
 """
 function NonlinearProblem(prob::AbstractODEProblem)
-    NonlinearProblem{isinplace(prob)}(prob.f, prob.u0, prob.p)
+    NonlinearProblem{isinplace(prob)}(prob.f, prob.u0, prob.p; prob.kwargs...)
 end
 
 @doc doc"""

diff --git a/src/problems/steady_state_problems.jl b/src/problems/steady_state_problems.jl
@@ -123,5 +123,5 @@ $(SIGNATURES)
 Define a steady state problem from a standard ODE problem.
 """
 function SteadyStateProblem(prob::AbstractODEProblem)
-    SteadyStateProblem{isinplace(prob)}(prob.f, prob.u0, prob.p)
+    SteadyStateProblem{isinplace(prob)}(prob.f, prob.u0, prob.p; prob.kwargs...)
 end
diff --git a/test/convert_tests.jl b/test/convert_tests.jl
@@ -33,3 +33,33 @@ end
     prob = ODEProblem(lorenz!, u0, tspan)
     nlprob = NonlinearProblem(prob)
 end
+
+@testset "Convert ODEProblem with kwargs to NonlinearProblem" begin
+    function lorenz!(du, u, p, t)
+        du[1] = p[1]*(u[2] - u[1])
+        du[2] = u[1] * (p[2] - u[3]) - u[2]
+        du[3] = u[1] * u[2] - p[3] * u[3]
+    end
+    u0 = [1.0; 0.0; 0.0]
+    tspan = (0.0, 100.0)
+    p = [10.0,28.0,8/3]
+    prob = ODEProblem(lorenz!, u0, tspan,p;a=1.0,b=2.0)
+    nlprob = NonlinearProblem(prob)
+    @test nlprob.kwargs[:a] == prob.kwargs[:a]
+    @test nlprob.kwargs[:b] == prob.kwargs[:b]
+end
+
+@testset "Convert ODEProblem with kwargs to SteadyStateProblem" begin
+    function lorenz!(du, u, p, t)
+        du[1] = p[1]*(u[2] - u[1])
+        du[2] = u[1] * (p[2] - u[3]) - u[2]
+        du[3] = u[1] * u[2] - p[3] * u[3]
+    end
+    u0 = [1.0; 0.0; 0.0]
+    tspan = (0.0, 100.0)
+    p = [10.0,28.0,8/3]
+    prob = ODEProblem(lorenz!, u0, tspan,p;a=1.0,b=2.0)
+    sprob = SteadyStateProblem(prob)
+    @test sprob.kwargs[:a] == prob.kwargs[:a]
+    @test sprob.kwargs[:b] == prob.kwargs[:b]
+end