Merge pull request #739 from oameye/master

add `g` kwarg to `remake(::SDEProblem)`

Author: ChrisRackauckas

src/remake.jl

@@ -225,14 +225,15 @@ function remake(prob::BVProblem{uType, tType, iip, nlls}; f = missing, bc = miss
 end
 
 """
-    remake(prob::SDEProblem; f = missing, u0 = missing, tspan = missing,
+    remake(prob::SDEProblem; f = missing, g = missing, u0 = missing, tspan = missing,
            p = missing, noise = missing, noise_rate_prototype = missing,
            seed = missing, kwargs = missing, _kwargs...)
 
 Remake the given `SDEProblem`.
 """
 function remake(prob::SDEProblem;
         f = missing,
+        g = missing,
         u0 = missing,
         tspan = missing,
         p = missing,
@@ -261,8 +262,23 @@ function remake(prob::SDEProblem;
         seed = prob.seed
     end
 
-    if f === missing #TODO: Need more features, e.g. remake `g`
+    if f === missing && g === missing
         f = prob.f
+        g = prob.g
+    elseif f !== missing && g === missing
+        g = prob.g
+    elseif f === missing && g !== missing
+        if prob.f isa SDEFunction
+            f = remake(prob.f; g = g)
+        else
+            f = SDEFunction(prob.f, g)
+        end
+    else
+        if f isa SDEFunction
+            f = remake(f; g = g)
+        else
+            f = SDEFunction(f, g)
+        end
     end
 
     iip = isinplace(prob)
@@ -282,6 +298,37 @@ function remake(prob::SDEProblem;
     end
 end
 
+"""
+    remake(func::SDEFunction; f = missing, g = missing,
+           mass_matrix = missing, analytic = missing, kwargs...)
+
+Remake the given `SDEFunction`.
+"""
+function remake(func::SDEFunction;
+        f = missing,
+        g = missing,
+        mass_matrix = missing,
+        analytic = missing,
+        kwargs...)
+    if f === missing
+        f = func.f
+    end
+
+    if g === missing
+        g = func.g
+    end
+
+    if mass_matrix === missing
+        mass_matrix = func.mass_matrix
+    end
+
+    if analytic === missing
+        analytic = func.analytic
+    end
+
+    return SDEFunction(f, g; mass_matrix, analytic, kwargs...)
+end
+
 """
     remake(prob::OptimizationProblem; f = missing, u0 = missing, p = missing,
         lb = missing, ub = missing, int = missing, lcons = missing, ucons = missing,

test/remake_tests.jl

@@ -273,3 +273,14 @@ uspan = (1.0, 2.0)
 interval_prob = IntervalNonlinearProblem(interval_f, uspan)
 new_prob = @inferred IntervalNonlinearProblem remake(interval_prob; p = [0])
 @test new_prob.p == [0]
+
+# SDEProblem specific
+function noise2!(du, u, p, t)
+    du .= 0.2u
+end
+fn = SDEFunction(lorenz!, noise!; sys)
+sdeprob = SDEProblem(fn, u0, tspan, Tuple(p))
+newprob = remake(sdeprob; g = noise2!)
+@test newprob.f isa SDEFunction
+tmp = newprob.g([0.0, 0.0, 0.0], [1.0, 2.0, 3.0], nothing, 0.0)
+@test tmp≈[0.2, 0.4, 0.6] atol=1e-6