Merge pull request #3695 from fchen121/master

DE Transformation (Change of Variables)

Author: ChrisRackauckas

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
 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))

test/runtests.jl

@@ -47,6 +47,7 @@ end
             @safetestset "Error Handling" include("error_handling.jl")
             @safetestset "StructuralTransformations" include("structural_transformation/runtests.jl")
             @safetestset "Basic transformations" include("basic_transformations.jl")
+            @safetestset "Change of variables" include("changeofvariables.jl")
             @safetestset "State Selection Test" include("state_selection.jl")
             @safetestset "Symbolic Event Test" include("symbolic_events.jl")
             @safetestset "Stream Connect Test" include("stream_connectors.jl")