DE Transformation (Change of Variables)

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,12 @@ 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"
+TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
+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 +164,8 @@ StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.39"
 SymbolicUtils = "3.26.1"
 Symbolics = "6.40"
+TermInterface = "2.0.0"
+Test = "1.11.0"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"

src/ModelingToolkit.jl

@@ -297,7 +297,7 @@ export isinput, isoutput, getbounds, hasbounds, getguess, hasguess, isdisturbanc
        hasunit, getunit, hasconnect, getconnect,
        hasmisc, getmisc, state_priority
 export liouville_transform, change_independent_variable, substitute_component,
-       add_accumulations, noise_to_brownians, Girsanov_transform
+       add_accumulations, noise_to_brownians, Girsanov_transform, changeofvariables, change_of_variable_SDE
 export PDESystem
 export Differential, expand_derivatives, @derivatives
 export Equation, ConstrainedEquation

src/systems/diffeqs/basic_transformations.jl

@@ -53,6 +53,138 @@ function liouville_transform(sys::System; kwargs...)
     )
 end
 
+"""
+$(TYPEDSIGNATURES)
+
+Generates the set of ODEs after change of variables.
+
+
+Example:
+
+```julia
+using ModelingToolkit, OrdinaryDiffEq, Test
+
+# Change of variables: z = log(x)
+# (this implies that x = exp(z) is automatically non-negative)
+
+@parameters t α
+@variables x(t)
+D = Differential(t)
+eqs = [D(x) ~ α*x]
+
+tspan = (0., 1.)
+u0 = [x => 1.0]
+p = [α => -0.5]
+
+@named sys = ODESystem(eqs; defaults=u0)
+prob = ODEProblem(sys, [], tspan, p)
+sol = solve(prob, Tsit5())
+
+@variables z(t)
+forward_subs  = [log(x) => z]
+backward_subs = [x => exp(z)]
+
+@named new_sys = changeofvariables(sys, forward_subs, backward_subs)
+@test equations(new_sys)[1] == (D(z) ~ α)
+
+new_prob = ODEProblem(new_sys, [], tspan, p)
+new_sol = solve(new_prob, Tsit5())
+
+@test isapprox(new_sol[x][end], sol[x][end], atol=1e-4)
+```
+
+"""
+function changeofvariables(
+    sys::System, iv, forward_subs, backward_subs;
+    simplify=true, t0=missing, isSDE=false
+)
+    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)
+    neqs = get_noise_eqs(sys)
+    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(@brownians $name))
+        end
+    else
+        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
+    dfdt = Symbolics.derivative( first.(forward_subs), t )
+    ∂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, ex, first, second) in zip(new_vars, dfdt, ∂f∂x, ∂2f∂x2)
+        for (eqs, neq) in zip(old_eqs, neqs)
+            if occursin(value(eqs.lhs), value(ex))
+                ex = substitute(ex, eqs.lhs => eqs.rhs)
+                if isSDE
+                    for (noise, B) in zip(neq, brownvars)
+                        ex = ex + 1/2 * noise^2 * second + noise*first*B
+                    end
+                end
+            end
+        end
+        ex = substitute(ex, Dict(forward_subs))
+        ex = substitute(ex, Dict(backward_subs))
+        if simplify
+            ex = Symbolics.simplify(ex, expand=true)
+        end
+        push!(new_eqs, Differential(t)(new_var) ~ ex)
+    end
+
+    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(vcat(new_eqs, first.(backward_subs) .~ last.(backward_subs)), t;
+            defaults=new_defs,
+            observed=observed(sys)
+            )
+    if simplify
+        return mtkcompile(new_sys)
+    end
+    return new_sys
+end
+
 """
     change_independent_variable(
         sys::System, iv, eqs = [];

test/changeofvariables.jl

@@ -0,0 +1,150 @@
+using ModelingToolkit, OrdinaryDiffEq, StochasticDiffEq
+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))
+
+@parameters α
+@variables x(t)
+D = Differential(t)
+eqs = [D(x) ~ α*x]
+
+tspan = (0., 1.)
+def = [x => 1.0, α => -0.5]
+
+@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, t, forward_subs, backward_subs)
+@test equations(new_sys)[1] == (D(z) ~ α)
+
+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 α
+@variables x(t)
+D = Differential(t)
+eqs = [D(x) ~ t^2 + α - x^2]
+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, 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.)
+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
+@independent_variables t
+@variables x(t)[1:3]
+x = reshape(x, 3, 1)
+D = Differential(t)
+A = [0. -1. 0.; -0.5 0.5 0.; 0. 0. -1.]
+right = A*x
+eqs = vec(D.(x) .~ right)
+
+tspan = (0., 10.)
+u0 = [x[1] => 1.0, x[2] => 2.0, x[3] => -1.0]
+
+@mtkcompile sys = System(eqs, t; defaults=u0)
+prob = ODEProblem(sys,[],tspan)
+sol = solve(prob, Tsit5())
+
+T = eigen(A).vectors
+T_inv = inv(T)
+
+@variables z(t)[1:3]
+z = reshape(z, 3, 1)
+forward_subs  = vec(T_inv*x .=> z)
+backward_subs = vec(x .=> T*z)
+
+new_sys = changeofvariables(sys, t, forward_subs, backward_subs; simplify=true)
+
+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))
+A = diagm(eigen(A).values)
+A = sortslices(A, dims=1)
+new_A = sortslices(new_A, dims=1)
+@test isapprox(A, new_A, rtol = 1e-10)
+@test isapprox( new_sol[x[1],end], sol[x[1],end], rtol=1e-4)
+
+# Change of variables for sde
+noise_eqs = ModelingToolkit.get_noise_eqs
+value = ModelingToolkit.value
+
+@independent_variables t
+@brownians 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 = changeofvariables(sys, t, forward_subs, backward_subs)
+@test equations(new_sys)[1] == (D(y) ~ μ - 1/2*σ^2)
+@test noise_eqs(new_sys)[1] === value(σ)
+
+#Multiple Brownian and equations
+@independent_variables t
+@brownians Bx By
+@parameters μ σ α
+@variables x(t) y(t) z(t) w(t) u(t) v(t)
+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.]
+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)
+@test equations(new_sys)[3] == (D(v) ~ μ - 1/2*(α^2 + σ^2))
+@test noise_eqs(new_sys)[1,1] === value(σ)
+@test noise_eqs(new_sys)[1,2] === value(0)
+@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 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))