Change of variable for non-simplified SDE

Author: fchen121

src/systems/diffeqs/basic_transformations.jl

@@ -98,9 +98,6 @@ function changeofvariables(
     sys::System, iv, forward_subs, backward_subs;
     simplify=true, t0=missing, isSDE=false
 )
-    if !iscomplete(sys)
-        sys = mtkcompile(sys)
-    end
     t = iv
 
     old_vars = first.(backward_subs)
@@ -110,15 +107,32 @@ function changeofvariables(
     # use: dY = (∂f/∂t + μ∂f/∂x + (1/2)*σ^2*∂2f/∂x2)dt + σ∂f/∂xdW
     old_eqs = equations(sys)
     neqs = get_noise_eqs(sys)
-    if neqs !== nothing
+    brownvars = brownians(sys)
+    
+    
+    if neqs === nothing && length(brownvars) === 0
+        neqs = ones(1, length(old_eqs))
+    elseif neqs !== nothing
         isSDE = true
         neqs = [neqs[i,:] for i in 1:size(neqs,1)]
 
         brownvars = map([Symbol(:B, :_, i) for i in 1:length(neqs[1])]) do name
             unwrap(only(@brownian $name))
         end
     else
-        neqs = ones(1, length(old_eqs))
+        isSDE = true
+        neqs = Vector{Any}[]
+        for (i, eq) in enumerate(old_eqs)
+            neq = Any[]
+            right = eq.rhs
+            for Bv in brownvars
+                lin_exp = linear_expansion(right, Bv)
+                right = lin_exp[2]
+                push!(neq, lin_exp[1])
+            end
+            push!(neqs, neq)
+            old_eqs[i] = eq.lhs ~ right
+        end
     end
 
     # df/dt = ∂f/∂x dx/dt + ∂f/∂t

test/changeofvariables.jl

@@ -127,9 +127,10 @@ new_sys = changeofvariables(sys, t, forward_subs, backward_subs)
 D = Differential(t)
 eqs = [D(x) ~ μ*x + σ*x*Bx, D(y) ~ α*By, D(u) ~ μ*u + σ*u*Bx + α*u*By]
 def = [x=>0., y=> 0., u=>0., μ => 2., σ=>1., α=>3.]
-@mtkcompile sys = System(eqs, t; defaults=def)
 forward_subs = [log(x) => z, y^2 => w, log(u) => v]
 backward_subs = [x => exp(z), y => w^.5, u => exp(v)]
+
+@mtkcompile sys = System(eqs, t; defaults=def)
 new_sys = changeofvariables(sys, t, forward_subs, backward_subs)
 @test equations(new_sys)[1] == (D(z) ~ μ - 1/2*σ^2)
 @test equations(new_sys)[2] == (D(w) ~ α^2)
@@ -139,4 +140,11 @@ new_sys = changeofvariables(sys, t, forward_subs, backward_subs)
 @test noise_eqs(new_sys)[2,1] === value(0)
 @test noise_eqs(new_sys)[2,2] === value(substitute(2*α*y, backward_subs[2]))
 @test noise_eqs(new_sys)[3,1] === value(σ)
-@test noise_eqs(new_sys)[3,2] === value(α)
+@test noise_eqs(new_sys)[3,2] === value(α)
+
+# Test for  Brownian instead of noise
+@named sys = System(eqs, t; defaults=def)
+new_sys = changeofvariables(sys, t, forward_subs, backward_subs; simplify=false)
+@test simplify(equations(new_sys)[1]) == simplify((D(z) ~ μ - 1/2*σ^2 + σ*Bx))
+@test simplify(equations(new_sys)[2]) == simplify((D(w) ~ α^2 + 2*α*w^.5*By))
+@test simplify(equations(new_sys)[3]) == simplify((D(v) ~ μ - 1/2*(α^2 + σ^2) + σ*Bx + α*By))