Implement and fix change of variable for ODE

Author: fchen121

Project.toml

@@ -44,7 +44,9 @@ NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
+OrdinaryDiffEqDefault = "50262376-6c5a-4cf5-baba-aaf4f84d72d7"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -59,9 +61,11 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 URIs = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
@@ -159,6 +163,7 @@ StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.39"
 SymbolicUtils = "3.26.1"
 Symbolics = "6.40"
+Test = "1.11.0"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"

src/systems/diffeqs/basic_transformations.jl

@@ -94,13 +94,13 @@ new_sol = solve(new_prob, Tsit5())
 ```
 
 """
-function changeofvariables(sys::System, forward_subs, backward_subs; simplify=false, t0=missing)
-    t = independent_variable(sys)
+function changeofvariables(sys::System, iv, forward_subs, backward_subs; simplify=true, t0=missing)
+    t = iv
 
     old_vars = first.(backward_subs)
     new_vars = last.(forward_subs)
-    kept_vars = setdiff(states(sys), old_vars)
-    rhs = [eq.rhs for eq in equations(sys)]
+    # kept_vars = setdiff(states(sys), old_vars)
+    # rhs = [eq.rhs for eq in equations(sys)]
 
     # use: dz/dt = ∂z/∂x dx/dt + ∂z/∂t
     dzdt = Symbolics.derivative( first.(forward_subs), t )
@@ -120,36 +120,44 @@ function changeofvariables(sys::System, forward_subs, backward_subs; simplify=fa
     defs = get_defaults(sys)
     new_defs = Dict()
     for f_sub in forward_subs
-        #TODO call value(...)?
         ex = substitute(first(f_sub), defs)
         if !ismissing(t0)
             ex = substitute(ex, t => t0)
         end
         new_defs[last(f_sub)] = ex
     end
-    return ODESystem(new_eqs;
+    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
 
-function change_of_variable_SDE(sys::System, forward_subs, backward_subs, iv; simplify=false, t0=missing)
-    t = independent_variable(sys)
+function change_of_variable_SDE(sys::System, iv, nvs, forward_subs, backward_subs; simplify=false, t0=missing)
+    t = iv
 
     old_vars = first.(backward_subs)
     new_vars = last.(forward_subs)
 
     # use: f = Y(t, X)
     # use: dY = (∂f/∂t + μ∂f/∂x + (1/2)*σ^2*∂2f/∂x2)dt + σ∂f/∂xdW
     old_eqs = equations(sys)
-    old_noise = get_noiseeqs(sys)
+    old_noise = ModelingToolkit.get_noise_eqs(sys)
     ∂f∂t = Symbolics.derivative( first.(forward_subs), t )
-    ∂f∂x = Symbolics.derivative( first.(forward_subs), old_vars )
+    ∂f∂x = [Symbolics.derivative( first(f_sub), old_var )]
     ∂2f∂x2 = Symbolics.derivative( ∂f∂x, old_vars )
     new_eqs = Equation[]
 
-    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
+    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 eqs in old_eqs
             ex = substitute(ex, eqs.lhs => eqs.rhs)
         end
@@ -160,11 +168,11 @@ function change_of_variable_SDE(sys::System, forward_subs, backward_subs, iv; si
         end
         push!(new_eqs, Differential(t)(new_var) ~ ex)
     end
-    new_noise = [noise * div for (noise, div) in zip(old_noise, ∂f∂x)]
 
-    return SDESystem(new_eqs, new_noise;
+    @named new_sys = System(new_eqs;
             observed=vcat(observed(sys),first.(backward_subs) .~ last.(backward_subs))
             )
+    return new_sys
 end
 
 """

test/changeofvariables.jl

@@ -1,96 +1,110 @@
-using ModelingToolkit, OrdinaryDiffEq
+using ModelingToolkit, OrdinaryDiffEq, StochasticDiffEq
 using Test, LinearAlgebra
 
 
 # Change of variables: z = log(x)
 # (this implies that x = exp(z) is automatically non-negative)
-
-@parameters t α
+@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))
+
+@parameters α
 @variables x(t)
 D = Differential(t)
 eqs = [D(x) ~ α*x]
 
 tspan = (0., 1.)
-u0 = [x => 1.0]
-p = [α => -0.5]
+def = [x => 1.0, α => -0.5]
 
-sys = ODESystem(eqs; defaults=u0)
-prob = ODEProblem(sys, [], tspan, p)
+@mtkcompile sys = System(eqs, t;defaults=def)
+prob = ODEProblem(sys, [], tspan)
 sol = solve(prob, Tsit5())
 
 @variables z(t)
 forward_subs  = [log(x) => z]
 backward_subs = [x => exp(z)]
-new_sys = changeofvariables(sys, forward_subs, backward_subs)
+new_sys = changeofvariables(sys, t, forward_subs, backward_subs)
 @test equations(new_sys)[1] == (D(z) ~ α)
 
-new_prob = ODEProblem(new_sys, [], tspan, p)
+new_prob = ODEProblem(new_sys, [], tspan)
 new_sol = solve(new_prob, Tsit5())
 
 @test isapprox(new_sol[x][end], sol[x][end], atol=1e-4)
 
 
 
 # Riccati equation
-@parameters t α
+@parameters α
 @variables x(t)
 D = Differential(t)
 eqs = [D(x) ~ t^2 + α - x^2]
-sys = ODESystem(eqs, defaults=[x=>1.])
+def = [x=>1., α => 1.]
+@mtkcompile sys = System(eqs, t; defaults=def)
 
 @variables z(t)
 forward_subs  = [t + α/(x+t) =>    z       ]
 backward_subs = [       x    => α/(z-t) - t]
 
-new_sys = changeofvariables(sys, forward_subs, backward_subs; simplify=true, t0=0.)
+new_sys = changeofvariables(sys, t, forward_subs, backward_subs; simplify=true, t0=0.)
 # output should be equivalent to
 # t^2 + α - z^2 + 2   (but this simplification is not found automatically)
 
 tspan = (0., 1.)
-p = [α => 1.]
-prob = ODEProblem(sys,[],tspan,p)
-new_prob = ODEProblem(new_sys,[],tspan,p)
+prob = ODEProblem(sys,[],tspan)
+new_prob = ODEProblem(new_sys,[],tspan)
 
 sol = solve(prob, Tsit5())
 new_sol = solve(new_prob, Tsit5())
 
 @test isapprox(sol[x][end], new_sol[x][end], rtol=1e-4)
 
 
-# Linear transformation to diagonal system
-@parameters t
-@variables x[1:3](t)
-D = Differential(t)
-A = [0. -1. 0.; -0.5 0.5 0.; 0. 0. -1.]
-eqs = D.(x) .~ A*x
+# # Linear transformation to diagonal system
+# @variables x(t)[1:3]
+# D = Differential(t)
+# A = [0. -1. 0.; -0.5 0.5 0.; 0. 0. -1.]
+# right = A.*transpose(x)
+# eqs = [D(x[1]) ~ sum(right[1, 1:3]), D(x[2]) ~ sum(right[2, 1:3]), D(x[3]) ~ sum(right[3, 1:3])]
 
-tspan = (0., 10.)
-u0 = x .=> [1.0, 2.0, -1.0]
+# tspan = (0., 10.)
+# u0 = [x[1] => 1.0, x[2] => 2.0, x[3] => -1.0]
 
-sys = ODESystem(eqs; defaults=u0)
-prob = ODEProblem(sys,[],tspan)
-sol = solve(prob, Tsit5())
+# @mtkcompile sys = System(eqs, t; defaults=u0)
+# prob = ODEProblem(sys,[],tspan)
+# sol = solve(prob, Tsit5())
 
-T = eigen(A).vectors
+# T = eigen(A).vectors
 
-@variables z[1:3](t)
-forward_subs  = T \ x .=> z
-backward_subs = x     .=> T*z
+# @variables z(t)[1:3]
+# forward_subs  = T \ x .=> z
+# backward_subs = x     .=> T*z
 
-new_sys = changeofvariables(sys, forward_subs, backward_subs; simplify=true)
+# new_sys = changeofvariables(sys, t, forward_subs, backward_subs; simplify=true)
 
-new_prob = ODEProblem(new_sys, [], tspan, p)
-new_sol = solve(new_prob, Tsit5())
+# new_prob = ODEProblem(new_sys, [], tspan)
+# new_sol = solve(new_prob, Tsit5())
 
-# test RHS
-new_rhs = [eq.rhs for eq in equations(new_sys)]
-new_A = Symbolics.value.(Symbolics.jacobian(new_rhs, z))
-@test isapprox(diagm(eigen(A).values), new_A, rtol = 1e-10)
-@test isapprox( new_sol[x[1],end], sol[x[1],end], rtol=1e-4)
+# # test RHS
+# new_rhs = [eq.rhs for eq in equations(new_sys)]
+# new_A = Symbolics.value.(Symbolics.jacobian(new_rhs, z))
+# @test isapprox(diagm(eigen(A).values), new_A, rtol = 1e-10)
+# @test isapprox( new_sol[x[1],end], sol[x[1],end], rtol=1e-4)
 
 # Change of variables for sde
-@Browian B
-@parameters μ σ
-@variables x(t) y(t)
+# @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)