Merge pull request #546 from ErikQQY/qqy/refactor_sde

Refactor SDEProblem constructor
SciML · Sep 29, 2023 · f2d7e98 · f2d7e98
2 parents 0e396bf + 500a725
commit f2d7e98
Show file tree

Hide file tree

Showing 17 changed files with 47 additions and 47 deletions.
diff --git a/Project.toml b/Project.toml
@@ -35,7 +35,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Adapt = "3"
 ArrayInterface = "6, 7"
 DataStructures = "0.18"
-DiffEqBase = "6.122"
+DiffEqBase = "6.130.1"
 DiffEqNoiseProcess = "5.13"
 DocStringExtensions = "0.8, 0.9"
 FillArrays = "0.6, 0.7, 0.8, 0.9, 0.10, 0.11, 0.12, 0.13, 1"
@@ -49,7 +49,7 @@ OrdinaryDiffEq = "6.52"
 RandomNumbers = "1.5.3"
 RecursiveArrayTools = "2"
 Reexport = "0.2, 1.0"
-SciMLBase = "1.70"
+SciMLBase = "2.0.6"
 SciMLOperators = "0.2.9, 0.3"
 SparseDiffTools = "2"
 StaticArrays = "0.11, 0.12, 1.0"

diff --git a/test/adaptive/sde_linearadaptive_tests.jl b/test/adaptive/sde_linearadaptive_tests.jl
@@ -11,8 +11,8 @@ add_probs[2] = prob_sde_additivesystem
 
 for i in 1:2
   global sol,sol2,err1
-  bigprob = SDEProblem(probs[i].f,probs[i].g,big.(probs[i].u0),(big.(probs[i].tspan[1]),big.(probs[i].tspan[2])),noise=probs[i].noise)
-  add_bigprob = SDEProblem(add_probs[i].f,add_probs[i].g,big.(add_probs[i].u0),(big.(add_probs[i].tspan[1]),big.(add_probs[i].tspan[2])),noise=add_probs[i].noise)
+  bigprob = SDEProblem(probs[i].f,big.(probs[i].u0),(big.(probs[i].tspan[1]),big.(probs[i].tspan[2])),noise=probs[i].noise)
+  add_bigprob = SDEProblem(add_probs[i].f,big.(add_probs[i].u0),(big.(add_probs[i].tspan[1]),big.(add_probs[i].tspan[2])),noise=add_probs[i].noise)
   ## SRIW1
 
   Random.seed!(100)

diff --git a/test/commutative_tests.jl b/test/commutative_tests.jl
@@ -44,7 +44,7 @@ end
 
 ff_commute = SDEFunction(f_commute,σ,analytic=f_commute_analytic)
 
-prob = SDEProblem(ff_commute,σ,u0,(0.0,1.0),noise_rate_prototype=rand(2,4))
+prob = SDEProblem(ff_commute,u0,(0.0,1.0),noise_rate_prototype=rand(2,4))
 
 sol = solve(prob,RKMilCommute(),dt=1/2^(8))
 sol = solve(prob,RKMilGeneral(ii_approx=IICommutative()),dt=1/2^(8))
@@ -67,7 +67,7 @@ sim2 = test_convergence(dts,prob,RKMilGeneral(p=2),trajectories=Int(2e2))
 
 ff_commute_oop = SDEFunction(f_commute_oop,σ_oop,analytic=f_commute_analytic)
 
-proboop = SDEProblem(ff_commute_oop,σ_oop,u0,(0.0,1.0),noise_rate_prototype=rand(2,4))
+proboop = SDEProblem(ff_commute_oop,u0,(0.0,1.0),noise_rate_prototype=rand(2,4))
 
 sol = solve(proboop,RKMilCommute(),dt=1/2^(8))
 sol = solve(proboop,RKMilGeneral(ii_approx=IICommutative()),dt=1/2^(8))

diff --git a/test/iif_methods.jl b/test/iif_methods.jl
@@ -12,9 +12,9 @@ f1_no_noise(u,p,t) = μ
 f1_no_noise_analytic(u0,p,t,W) = u0.*exp.(2μ*t)
 
 ff1_μ = SplitSDEFunction(f1_μ,f2,σ,analytic=f1_μ_analytic)
-prob = SDEProblem(ff1_μ,σ,1/2,(0.0,1.0))
+prob = SDEProblem(ff1_μ,1/2,(0.0,1.0))
 ff1_no_noise = SplitSDEFunction(f1_no_noise,f2,no_noise,analytic=f1_no_noise_analytic)
-no_noise_prob = SDEProblem(ff1_no_noise,no_noise,1/2,(0.0,1.0))
+no_noise_prob = SDEProblem(ff1_no_noise,1/2,(0.0,1.0))
 
 sol = solve(prob,IIF1M(),dt=1/10)
 
@@ -74,7 +74,7 @@ end
 f2(du,u,p,t) = du .= μ .* u
 
 ff1_A = SplitSDEFunction(f1_A,f2,σ,analytic=f1_A_analytic)
-prob = SDEProblem(ff1_A,σ,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
+prob = SDEProblem(ff1_A,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
 
 f1_no_noise(du,u,p,t) = A
 f2(du,u,p,t) = (du .= μ .* u)
@@ -86,7 +86,7 @@ function f1_no_noise_analytic(u0,p,t,W)
  exp(tmp)*u0
 end
 ff1_A = SplitSDEFunction(f1_no_noise,f2,σ22,analytic=f1_no_noise_analytic)
-prob_no_noise = SDEProblem(ff1_A,σ22,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
+prob_no_noise = SDEProblem(ff1_A,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
 
 
 sol = solve(prob,IIF1M(),dt=1/10)

diff --git a/test/mass_matrix_tests.jl b/test/mass_matrix_tests.jl
@@ -23,8 +23,8 @@ function g!(du,u,p,t)
   @. du = 0.0
 end
 
-prob2 = SDEProblem(SDEFunction(mm_g,g!;analytic=mm_analytic),g!,ones(3),(0.0,1.0))
-prob = SDEProblem(SDEFunction(mm_f,g!;analytic=mm_analytic,mass_matrix=mm_A),g!,
+prob2 = SDEProblem(SDEFunction(mm_g,g!;analytic=mm_analytic),ones(3),(0.0,1.0))
+prob = SDEProblem(SDEFunction(mm_f,g!;analytic=mm_analytic,mass_matrix=mm_A),
                   ones(3),(0.0,1.0))
 
 sol = solve(prob, ImplicitRKMil(theta=1), dt = 0.01, adaptive = false)
@@ -60,7 +60,7 @@ function mm_g2(du,u,p,t)
     mul!(du,mm_A,u)
 end
 prob2 = SDEProblem(no_mm_f2,no_mm_g2,ones(3),(0.0,1.0))
-prob = SDEProblem(SDEFunction(mm_f2,no_mm_g2;mass_matrix=mm_A),no_mm_g2,ones(3),(0.0,1.0))
+prob = SDEProblem(SDEFunction(mm_f2,no_mm_g2;mass_matrix=mm_A),ones(3),(0.0,1.0))
 
 Random.seed!(1)
 sol = solve(prob, ImplicitEM(theta=1), dt = 0.01, adaptive = false)

diff --git a/test/multivariate_geometric.jl b/test/multivariate_geometric.jl
@@ -19,7 +19,7 @@ function f_analytic(u0,p,t,W)
   exp(tmp)*u0
 end
 
-prob2 = SDEProblem(SDEFunction(f,σ,analytic=f_analytic),σ,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
+prob2 = SDEProblem(SDEFunction(f,σ,analytic=f_analytic),u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
 
 sol2 = solve(prob2,EM(),dt=1/100)
 
@@ -51,7 +51,7 @@ function f_analytic_iip(u0,p,t,W)
  exp(tmp)*u0
 end
 
-prob2 = SDEProblem(SDEFunction(f_iip,σ_iip,analytic=f_analytic_iip),σ_iip,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
+prob2 = SDEProblem(SDEFunction(f_iip,σ_iip,analytic=f_analytic_iip),u0,(0.0,1.0),noise_rate_prototype=rand(2,2))
 
 sol2 = solve(prob2,EM(),dt=1/100)
 

diff --git a/test/noncommutative_tests.jl b/test/noncommutative_tests.jl
@@ -44,7 +44,7 @@ end
 
 ff_noncommute = SDEFunction(f_noncommute,g_noncommute,analytic=f_noncommute_analytic)
 
-prob = SDEProblem(ff_noncommute,g_noncommute,u0,(0.0,1.0),noise_rate_prototype=rand(4,m))
+prob = SDEProblem(ff_noncommute,u0,(0.0,1.0),noise_rate_prototype=rand(4,m))
 
 sol = solve(prob,EM(),dt=1/2^(8))
 sol = solve(prob,RKMilGeneral(p=10),dt=1/2^(8))
@@ -58,7 +58,7 @@ sim3 = test_convergence(dts,prob,RKMilGeneral(p=2),trajectories=Int(1e2))
 @test abs(sim3.𝒪est[:final] - 1) < 0.2
 
 ff_noncommute_stratonovich = SDEFunction(f_noncommute,g_noncommute,analytic=f_noncommute_analytic_stratonovich)
-prob_stratonovich = SDEProblem(ff_noncommute_stratonovich,g_noncommute,u0,(0.0,1.0),noise_rate_prototype=rand(4,m))
+prob_stratonovich = SDEProblem(ff_noncommute_stratonovich,u0,(0.0,1.0),noise_rate_prototype=rand(4,m))
 
 sim4 = test_convergence(dts,prob_stratonovich,EulerHeun(),trajectories=Int(1e2))
 @test abs(sim4.𝒪est[:final] - 1.0) < 0.2

diff --git a/test/nondiagonal_tests.jl b/test/nondiagonal_tests.jl
@@ -47,8 +47,8 @@ function ggprime(du, u, p, t)
   du .= coeff*u
 end
 
-prob = SDEProblem(SDEFunction(f_nondiag,g_nondiag,analytic=f_analytic_nondiag),g_nondiag,u0,(0.0,1.0),noise_rate_prototype=zeros(2,2))
-probiip = SDEProblem(SDEFunction(f_nondiag_iip,g_nondiag_iip,analytic=f_analytic_nondiag),g_nondiag_iip,u0,(0.0,1.0),noise_rate_prototype=zeros(2,2))
+prob = SDEProblem(SDEFunction(f_nondiag,g_nondiag,analytic=f_analytic_nondiag),u0,(0.0,1.0),noise_rate_prototype=zeros(2,2))
+probiip = SDEProblem(SDEFunction(f_nondiag_iip,g_nondiag_iip,analytic=f_analytic_nondiag),u0,(0.0,1.0),noise_rate_prototype=zeros(2,2))
 
 ## Just solve to test compatibility
 IEM = solve(probiip,ImplicitEM())

diff --git a/test/ode_convergence_regression.jl b/test/ode_convergence_regression.jl
@@ -3,7 +3,7 @@ using StochasticDiffEq, DiffEqDevTools, Test
 linear = (u,p,t) -> (p*u)
 g = (u,p,t) -> zero(u)
 linear_analytic = (u0,p,t,W) -> u0*exp(p*t)
-prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic),g,
+prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic),
                   1/2,(0.0,1.0),1.01)
 
 dts = (1/2) .^ (7:-1:4) #14->7 good plot
@@ -14,7 +14,7 @@ sim2 = test_convergence(dts,prob,SKenCarp(),trajectories=20)
 linear = (du,u,p,t) -> (du.=p.*u)
 g = (du,u,p,t) -> (du.=0)
 linear_analytic = (u0,p,t,W) -> u0*exp.(p.*t)
-prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic),g,
+prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic),
                   rand(4,2),(0.0,1.0),1.01)
 
 dts = (1/2) .^ (7:-1:4) #14->7 good plot
@@ -28,7 +28,7 @@ sim2 = test_convergence(dts,prob,SKenCarp(),trajectories=20)
 linear = (u,p,t) -> (p*u)
 g = (u,p,t) -> zero(u)
 linear_analytic = (u0,p,t,W) -> u0*exp(p*t)
-prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic),g, 1/2,(0.0,1.0),1.01)
+prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic), 1/2,(0.0,1.0),1.01)
 dts = 1 .//2 .^(10:-1:2)
 
 sim2 = test_convergence(dts,prob, DRI1(),trajectories=20)

diff --git a/test/outofplace_arrays.jl b/test/outofplace_arrays.jl
@@ -5,7 +5,7 @@ A = [-1.0 0.0; 0.0 -0.5]
 u0 = [1.0, 1.0]; tspan = (0.0,1.0)
 _f = (u,p,t) -> t*(A*u)
 _g = (u,p,t) -> 1.0
-prob = SDEProblem(SDEFunction(_f, _g), _g, u0, tspan)
+prob = SDEProblem(SDEFunction(_f, _g), u0, tspan)
 integrator = init(prob, SKenCarp(); adaptive=false, dt=0.01)
 step!(integrator)
 @test_broken solve(prob, SOSRI(); adaptive=false, dt=0.01) isa RODESolution
@@ -18,7 +18,7 @@ solve(prob, SOSRA2(); adaptive=false, dt=0.01)
 
 println("Vector g")
 _g = (u,p,t) -> [1.0, 1.0]
-prob = SDEProblem(SDEFunction(_f, _g), _g, u0, tspan)
+prob = SDEProblem(SDEFunction(_f, _g), u0, tspan)
 println("Implicit EM")
 integrator = init(prob, ImplicitEM(); adaptive=false, dt=0.01)
 step!(integrator)

diff --git a/test/sde/sde_additive_tests.jl b/test/sde/sde_additive_tests.jl
@@ -7,7 +7,7 @@ println("Bunch of additive solves")
 f_bm(u,p,t) = 0.0
 f_analytic_bm(u0,p,t,W) = W
 g_bm(u,p,t) = 1.0
-prob = SDEProblem(SDEFunction(f_bm,g_bm,analytic=f_analytic_bm),g_bm,0.0,(0.0,1.0))
+prob = SDEProblem(SDEFunction(f_bm,g_bm,analytic=f_analytic_bm),0.0,(0.0,1.0))
 sol1 = solve(prob,SRA1(),dt=1/2^(3),adaptive=false)
 sol2 = solve(prob,SOSRA(),dt=1/2^(3),adaptive=false)
 sol3 = solve(prob,SKenCarp(),dt=1/2^(3),adaptive=false)

diff --git a/test/sde/sde_dynamical.jl b/test/sde/sde_dynamical.jl
@@ -11,7 +11,7 @@ g(u,p,t) = 1
     v0 = 1
 
     ff_harmonic = DynamicalSDEFunction(f1_harmonic,f2_harmonic,g)
-    prob1 = DynamicalSDEProblem(ff_harmonic,g,v0,u0,(0.0,5.0))
+    prob1 = DynamicalSDEProblem(ff_harmonic,v0,u0,(0.0,5.0))
 
     dts = (1/2) .^ (8:-1:4)
 
@@ -31,7 +31,7 @@ end
     g_iip(du,u,p,t) = du .= g(u,p,t)
 
     ff_harmonic = DynamicalSDEFunction(f1_harmonic,f2_harmonic,g)
-    prob1 = DynamicalSDEProblem(ff_harmonic,g,v0,u0,(0.0,5.0))
+    prob1 = DynamicalSDEProblem(ff_harmonic,v0,u0,(0.0,5.0))
     sol1 = solve(prob1,BAOAB(gamma=γ);dt=1/10,save_noise=true)
 
     prob2 = DynamicalSDEProblem(f1_harmonic_iip,f2_harmonic_iip,g_iip,v0,u0,(0.0,5.0); noise=NoiseWrapper(sol1.W))

diff --git a/test/sparsediff_tests.jl b/test/sparsediff_tests.jl
@@ -40,7 +40,7 @@ colorvec = repeat(1:3,10)[1:10]
 u0=[1.,2.,3,4,5,5,4,3,2,1]
 tspan=(0.,10.)
 sdefun_sp= SDEFunction(f,g,colorvec=colorvec,jac_prototype=jac_sp)
-prob_sp = SDEProblem(sdefun_sp,g,u0,tspan)
+prob_sp = SDEProblem(sdefun_sp,u0,tspan)
 prob_std = SDEProblem(f,g,u0,tspan)
 
 sol_sp=solve(prob_sp,SKenCarp(autodiff=false))

diff --git a/test/split_tests.jl b/test/split_tests.jl
@@ -17,7 +17,7 @@ sol2 = solve(prob,EM(),dt=1/10)
 u0 = rand(4)
 
 ff_split = SplitSDEFunction(f1,f2,σ,analytic=f_split_analytic)
-prob = SplitSDEProblem(ff_split,σ,u0,(0.0,1.0))
+prob = SplitSDEProblem(ff_split,u0,(0.0,1.0))
 
 sol = solve(prob,SplitEM(),dt=1/10,save_noise=true)
 
@@ -40,7 +40,7 @@ ff2 = (u,p,t) -> 0.0
 σ2 = (u,p,t) -> α*β./sqrt.(1+t)
 ff1_analytic(u0,p,t,W) = @. u0/sqrt(1+t) + β*(t+α*W)/sqrt(1+t)
 f_ff1 = SplitSDEFunction(ff1,ff2,σ2,analytic=ff1_analytic)
-prob = SplitSDEProblem(f_ff1,σ2,1.,(0.0,1.0))
+prob = SplitSDEProblem(f_ff1,1.,(0.0,1.0))
 
 
 sol = solve(prob,EM(),dt=1/10)
@@ -60,7 +60,7 @@ ff1 = (u,p,t) -> 0.0
 ff2 = (u,p,t) -> β./sqrt.(1+t) - u./(2*(1+t))
 σ2 = (u,p,t) -> α*β./sqrt.(1+t)
 f_ff1 = SplitSDEFunction(ff1,ff2,σ2,analytic=ff1_analytic)
-prob = SplitSDEProblem(f_ff1,σ2,1.,(0.0,1.0))
+prob = SplitSDEProblem(f_ff1,1.,(0.0,1.0))
 
 sol = solve(prob,EM(),dt=1/10)
 sol2 = solve(prob,SKenCarp(),dt=1/10,seed=1)
@@ -79,7 +79,7 @@ ff1 = (u,p,t) -> β./sqrt.(1+t)
 ff2 = (u,p,t) -> - u./(2*(1+t))
 σ2 = (u,p,t) -> α*β./sqrt.(1+t)
 f_ff1 = SplitSDEFunction(ff1,ff2,σ2,analytic=ff1_analytic)
-prob = SplitSDEProblem(f_ff1,σ2,1.,(0.0,1.0))
+prob = SplitSDEProblem(f_ff1,1.,(0.0,1.0))
 
 sol = solve(prob,EM(),dt=1/10)
 sol2 = solve(prob,SKenCarp(),dt=1/10)
@@ -102,7 +102,7 @@ ff1 = (du,u,p,t) -> @. du = β/sqrt(1+t) - u/(2*(1+t))
 ff2 = (du,u,p,t) -> @. du = 0.0
 σ2 = (du,u,p,t) -> @. du = α*β/sqrt(1+t)
 f_ff1 = SplitSDEFunction(ff1,ff2,σ2,analytic=ff1_analytic)
-prob = SplitSDEProblem(f_ff1,σ2,[1.],(0.0,1.0))
+prob = SplitSDEProblem(f_ff1,[1.],(0.0,1.0))
 
 sol = solve(prob,EM(),dt=1/10)
 sol2 = solve(prob,SKenCarp(),dt=1/10)
@@ -121,7 +121,7 @@ ff1 = (du,u,p,t) -> @. du = 0.0
 ff2 = (du,u,p,t) -> @. du = β/sqrt(1+t) - u/(2*(1+t))
 σ2 = (du,u,p,t) -> @. du = α*β/sqrt(1+t)
 f_ff1 = SplitSDEFunction(ff1,ff2,σ2,analytic=ff1_analytic)
-prob = SplitSDEProblem(f_ff1,σ2,[1.],(0.0,1.0))
+prob = SplitSDEProblem(f_ff1,[1.],(0.0,1.0))
 
 sol = solve(prob,EM(),dt=1/10)
 sol2 = solve(prob,SKenCarp(),dt=1/10)
@@ -140,7 +140,7 @@ ff1 = (du,u,p,t) -> @. du = β/sqrt(1+t)
 ff2 = (du,u,p,t) -> @. du = - u/(2*(1+t))
 σ2 = (du,u,p,t) -> @. du = α*β/sqrt(1+t)
 f_ff1 = SplitSDEFunction(ff1,ff2,σ2,analytic=ff1_analytic)
-prob = SplitSDEProblem(f_ff1,σ2,[1.],(0.0,1.0))
+prob = SplitSDEProblem(f_ff1,[1.],(0.0,1.0))
 
 sol = solve(prob,EM(),dt=1/10)
 sol2 = solve(prob,SKenCarp(),dt=1/10)

diff --git a/test/utility_tests.jl b/test/utility_tests.jl
@@ -15,7 +15,7 @@ using StochasticDiffEq.SciMLOperators: MatrixOperator
     fun = SDEFunction(_f, _g;
                       mass_matrix=mm,
                       jac=(u,p,t) -> A)
-    prob = SDEProblem(fun, _g, u0, tspan)
+    prob = SDEProblem(fun, u0, tspan)
     integrator = init(prob, ImplicitEM(theta=1); adaptive=false, dt=dt)
     W = integrator.cache.nlsolver.cache.W
     calc_W!(W, integrator, integrator.cache.nlsolver, integrator.cache, dtgamma, false)
@@ -27,7 +27,7 @@ using StochasticDiffEq.SciMLOperators: MatrixOperator
     fun = SDEFunction(_f, _g;
                       mass_matrix=mm,
                       jac_prototype=MatrixOperator(A))
-    prob = SDEProblem(fun, _g, u0, tspan)
+    prob = SDEProblem(fun, u0, tspan)
     integrator = init(prob, ImplicitEM(theta=1); adaptive=false, dt=dt)
     W = integrator.cache.nlsolver.cache.W
     calc_W!(W, integrator, integrator.cache.nlsolver, integrator.cache, dtgamma, false)
@@ -50,11 +50,11 @@ using StochasticDiffEq.SciMLOperators: MatrixOperator
 
     _f = (u,p,t) -> t*(A*u); _f_ip = (du,u,p,t) -> lmul!(t,mul!(du,A,u))
     _g = (u,p,t) -> σ*u; _g_ip = (du,u,p,t) -> mul!(du,σ,u)
-    prob1 = SDEProblem(SDEFunction(_f, _g; mass_matrix=mm), _g, u0, tspan)
-    prob2 = SDEProblem(SDEFunction(_f, _g; mass_matrix=mm, jac=(u,p,t) -> t*A), _g, u0, tspan)
-    prob1_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; mass_matrix=mm), _g_ip, u0, tspan)
+    prob1 = SDEProblem(SDEFunction(_f, _g; mass_matrix=mm), u0, tspan)
+    prob2 = SDEProblem(SDEFunction(_f, _g; mass_matrix=mm, jac=(u,p,t) -> t*A), u0, tspan)
+    prob1_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; mass_matrix=mm), u0, tspan)
     jac_prototype=MatrixOperator(similar(A); update_func! = (J,u,p,t) -> (J .= t .* A; J))
-    prob2_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; mass_matrix=mm, jac_prototype=jac_prototype), _g_ip, u0, tspan)
+    prob2_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; mass_matrix=mm, jac_prototype=jac_prototype), u0, tspan)
 
     for Alg in [ImplicitEM, ISSEM]
       println(Alg)
@@ -68,10 +68,10 @@ using StochasticDiffEq.SciMLOperators: MatrixOperator
 
     σ = 1.0
     _g = (u,p,t) -> σ; _g_ip = (du,u,p,t) -> (du .= σ)
-    prob1 = SDEProblem(SDEFunction(_f, _g), _g, u0, tspan)
-    prob2 = SDEProblem(SDEFunction(_f, _g; jac=(u,p,t) -> t*A), _g, u0, tspan)
-    prob1_ip = SDEProblem(SDEFunction(_f_ip, _g_ip), _g_ip, u0, tspan)
-    prob2_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; jac_prototype=jac_prototype), _g_ip, u0, tspan)
+    prob1 = SDEProblem(SDEFunction(_f, _g), u0, tspan)
+    prob2 = SDEProblem(SDEFunction(_f, _g; jac=(u,p,t) -> t*A), u0, tspan)
+    prob1_ip = SDEProblem(SDEFunction(_f_ip, _g_ip), u0, tspan)
+    prob2_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; jac_prototype=jac_prototype), u0, tspan)
 
     println(SKenCarp)
     Random.seed!(0); sol1 = solve(prob1, SKenCarp(); adaptive=false, dt=0.01)

diff --git a/test/weak_convergence/iip_weak.jl b/test/weak_convergence/iip_weak.jl
@@ -7,7 +7,7 @@ f_linear_iip(du,u,p,t) = @.(du = 1.01*u)
 linear_analytic(u0,p,t,W) = @.(u0*exp(0.63155t+0.87W))
 
 prob_sde_linear_iip = SDEProblem(SDEFunction(f_linear_iip,σ_linear_iip,
-                             analytic=linear_analytic),σ_linear_iip,[1/2],(0.0,1.0))
+                             analytic=linear_analytic),[1/2],(0.0,1.0))
 
 Random.seed!(100)
 dts = 1 .//2 .^(7:-1:3) #14->7 good plot

diff --git a/test/weak_convergence/weak_srockc2.jl b/test/weak_convergence/weak_srockc2.jl
@@ -8,7 +8,7 @@ f_linear_iip(du,u,p,t) = @.(du = 1.01*u)
 linear_analytic(u0,p,t,W) = @.(u0*exp(0.63155t+0.87W))
 
 prob_sde_linear_iip = SDEProblem(SDEFunction(f_linear_iip,σ_linear_iip,
-                             analytic=linear_analytic),σ_linear_iip,[1/2],(0.0,1.0))
+                             analytic=linear_analytic),[1/2],(0.0,1.0))
 
 Random.seed!(100)
 dts = 1 .//2 .^(7:-1:3) #14->7 good plot