Implement change of variables for sde

Author: fchen121

src/systems/diffeqs/basic_transformations.jl

@@ -141,8 +141,10 @@ function changeofvariables(sys::System, iv, forward_subs, backward_subs; simplif
     return new_sys
 end
 
-function change_of_variable_SDE(sys::System, iv, nvs, forward_subs, backward_subs; simplify=false, t0=missing)
+function change_of_variable_SDE(sys::System, iv, forward_subs, backward_subs; simplify=true, t0=missing)
+    sys = mtkcompile(sys)
     t = iv
+    @brownian B
 
     old_vars = first.(backward_subs)
     new_vars = last.(forward_subs)
@@ -151,13 +153,17 @@ function change_of_variable_SDE(sys::System, iv, nvs, forward_subs, backward_sub
     # use: dY = (∂f/∂t + μ∂f/∂x + (1/2)*σ^2*∂2f/∂x2)dt + σ∂f/∂xdW
     old_eqs = equations(sys)
     old_noise = ModelingToolkit.get_noise_eqs(sys)
+
+    # Is there a function to find partial derivative?
     ∂f∂t = Symbolics.derivative( first.(forward_subs), t )
-    ∂f∂x = [Symbolics.derivative( first(f_sub), old_var )]
-    ∂2f∂x2 = Symbolics.derivative( ∂f∂x, old_vars )
+    ∂f∂t = [substitute(f_t, Differential(t)(old_var) => 0) for (f_t, old_var) in zip(∂f∂t, old_vars)]
+
+    ∂f∂x = [Symbolics.derivative( first(f_sub), old_var ) for (f_sub, old_var) in zip(forward_subs, old_vars)]
+    ∂2f∂x2 = Symbolics.derivative.( ∂f∂x, old_vars )
     new_eqs = Equation[]
 
-    for (new_var, eq, noise, nv, first, second, third) in zip(new_vars, old_eqs, old_noise, nvs, ∂f∂t, ∂f∂x, ∂2f∂x2)
-        ex = first + eq.rhs * second + 1/2 * noise^2 * third + noise*second*nv
+    for (new_var, eq, noise, first, second, third) in zip(new_vars, old_eqs, old_noise, ∂f∂t, ∂f∂x, ∂2f∂x2)
+        ex = first + eq.rhs * second + 1/2 * noise^2 * third + noise*second*B
         for eqs in old_eqs
             ex = substitute(ex, eqs.lhs => eqs.rhs)
         end
@@ -169,9 +175,28 @@ function change_of_variable_SDE(sys::System, iv, nvs, forward_subs, backward_sub
         push!(new_eqs, Differential(t)(new_var) ~ ex)
     end
 
-    @named new_sys = System(new_eqs;
+    defs = get_defaults(sys)
+    new_defs = Dict()
+    for f_sub in forward_subs
+        ex = substitute(first(f_sub), defs)
+        if !ismissing(t0)
+            ex = substitute(ex, t => t0)
+        end
+        new_defs[last(f_sub)] = ex
+    end
+    for para in parameters(sys)
+        if haskey(defs, para)
+            new_defs[para] = defs[para]
+        end
+    end
+
+    @named new_sys = System(new_eqs, t;
+            defaults=new_defs,
             observed=vcat(observed(sys),first.(backward_subs) .~ last.(backward_subs))
             )
+    if simplify
+        return mtkcompile(new_sys)
+    end
     return new_sys
 end
 

test/changeofvariables.jl

@@ -5,11 +5,11 @@ using Test, LinearAlgebra
 # Change of variables: z = log(x)
 # (this implies that x = exp(z) is automatically non-negative)
 @independent_variables t
-# @variables z(t)[1:2, 1:2]
-# D = Differential(t)
-# eqs = [D(D(z)) ~ ones(2, 2)]
-# @mtkcompile sys = System(eqs, t)
-# @test_nowarn ODEProblem(sys, [z => zeros(2, 2), D(z) => ones(2, 2)], (0.0, 10.0))
+@variables z(t)[1:2, 1:2]
+D = Differential(t)
+eqs = [D(D(z)) ~ ones(2, 2)]
+@mtkcompile sys = System(eqs, t)
+@test_nowarn ODEProblem(sys, [z => zeros(2, 2), D(z) => ones(2, 2)], (0.0, 10.0))
 
 @parameters α
 @variables x(t)
@@ -94,17 +94,19 @@ new_sol = solve(new_prob, Tsit5())
 # @test isapprox( new_sol[x[1],end], sol[x[1],end], rtol=1e-4)
 
 # Change of variables for sde
-# @independent_variables t
-# @brownian B
-# @parameters μ σ
-# @variables x(t) y(t)
-# D = Differential(t)
-# eqs = [D(x) ~ μ*x + σ*x*B]
+@independent_variables t
+@brownian B
+@parameters μ σ
+@variables x(t) y(t)
+D = Differential(t)
+eqs = [D(x) ~ μ*x + σ*x*B]
 
-# def = [x=>0., μ => 2., σ=>1.]
-# @mtkcompile sys = System(eqs, t; defaults=def)
-# forward_subs = [log(x) => y]
-# backward_subs = [x => exp(y)]
-# new_sys = change_of_variable_SDE(sys, t, [B], forward_subs, backward_subs)
+def = [x=>0., μ => 2., σ=>1.]
+@mtkcompile sys = System(eqs, t; defaults=def)
+forward_subs = [log(x) => y]
+backward_subs = [x => exp(y)]
+new_sys = change_of_variable_SDE(sys, t, forward_subs, backward_subs)
+@test equations(new_sys)[1] == (D(y) ~ μ - 1/2*σ^2)
+@test ModelingToolkit.get_noise_eqs(new_sys)[1] === ModelingToolkit.value(σ)