Change independent variable of ODE systems

docs/pages.jl

@@ -10,6 +10,7 @@ pages = [
         "tutorials/stochastic_diffeq.md",
         "tutorials/discrete_system.md",
         "tutorials/parameter_identifiability.md",
+        "tutorials/change_independent_variable.md",
         "tutorials/bifurcation_diagram_computation.md",
         "tutorials/SampledData.md",
         "tutorials/domain_connections.md",

docs/src/systems/ODESystem.md

@@ -30,6 +30,7 @@ ODESystem
 structural_simplify
 ode_order_lowering
 dae_index_lowering
+change_independent_variable
 liouville_transform
 alias_elimination
 tearing

docs/src/tutorials/change_independent_variable.md

@@ -0,0 +1,119 @@
+# Changing the independent variable of ODEs
+
+Ordinary differential equations describe the rate of change of some dependent variables with respect to one independent variable.
+For the modeler it is often most natural to write down the equations with a particular independent variable, say time $t$.
+In many cases there are good reasons for reparametrizing ODEs in terms of a different independent variable:
+
+1. One may want $y(x)$ as a function of $x$ instead of $(x(t), y(t))$ as a function of $t$
+2. Some differential equations vary more nicely (e.g. less stiff or better behaved) with respect to one independent variable than another.
+3. It can reduce the number of equations that must be solved (e.g. $y(x)$ is one equation, while $(x(t), y(t))$ are two).
+
+To manually change the independent variable of an ODE, one must rewrite all equations in terms of a new variable and transform differentials with the chain rule.
+This is mechanical and error-prone.
+ModelingToolkit provides the utility function [`change_independent_variable`](@ref) that automates this process.
+
+## 1. Get one dependent variable as a function of another
+
+Consider a projectile shot with some initial velocity in a gravitational field.
+```@example changeivar
+using ModelingToolkit
+@independent_variables t
+D = Differential(t)
+@variables x(t) y(t)
+@parameters g = 9.81 v # gravitational acceleration and constant horizontal velocity
+M1 = ODESystem([
+    D(D(y)) ~ -g, D(x) ~ v # constant horizontal velocity
+], t; defaults = [
+    y => 0.0
+], initialization_eqs = [
+    #x ~ 0.0, # TODO: handle? # hide
+    D(x) ~ D(y) # equal initial horizontal and vertical velocity (45 °)
+], name = :M) |> complete
+M1s = structural_simplify(M1)
+```
+This is the standard parametrization that arises naturally from kinematics and Newton's laws.
+It expresses the position $(x(t), y(t))$ as a function of time $t$.
+But suppose we want to determine whether the projectile hits a target 10 meters away.
+There are at least three ways of answering this:
+* Solve the ODE for $(x(t), y(t))$ and use a callback to terminate when $x$ reaches 10 meters, and evaluate $y$ at the final time.
+* Solve the ODE for $(x(t), y(t))$ and use root finding to find the time when $x$ reaches 10 meters, and evaluate $y$ at that time.
+* Solve the ODE for $y(x)$ and evaluate it at 10 meters.
+
+We will demonstrate the last method by changing the independent variable from $t$ to $x$.
+This transformation is well-defined for any non-zero horizontal velocity $v$.
+```@example changeivar
+M2 = change_independent_variable(M1, M1.x; dummies = true)
+@assert M2.x isa Num # hide
+M2s = structural_simplify(M2; allow_symbolic = true)
+```
+The derivatives are now with respect to the new independent variable $x$, which can be accessed with `M2.x`.
+Notice how the number of equations has also decreased from three to two, as $\mathrm{d}x/\mathrm{d}t$ has been turned into an observed equation.
+It is straightforward to evolve the ODE for 10 meters and plot the resulting trajectory $y(x)$:
+```@example changeivar
+using OrdinaryDiffEq, Plots
+prob = ODEProblem(M2s, [], [0.0, 10.0], [v => 8.0]) # throw 10 meters with x-velocity 8 m/s
+sol = solve(prob, Tsit5())
+plot(sol; idxs = M2.y) # must index by M2.y = y(x); not M1.y = y(t)!
+```
+
+## 2. Alleviating stiffness by changing to logarithmic time
+
+In cosmology, the [Friedmann equations](https://en.wikipedia.org/wiki/Friedmann_equations) describe the expansion of the universe.
+In terms of conformal time $t$, they can be written
+```@example changeivar
+@variables a(t) Ω(t)
+a = GlobalScope(a) # global var needed by all species
+function species(w; kw...)
+    eqs = [D(Ω) ~ -3(1 + w) * D(a)/a * Ω]
+    return ODESystem(eqs, t, [Ω], []; kw...)
+end
+@named r = species(1//3) # radiation
+@named m = species(0) # matter
+@named Λ = species(-1) # dark energy / cosmological constant
+eqs = [
+    Ω ~ r.Ω + m.Ω + Λ.Ω # total energy density
+    D(a) ~ √(Ω) * a^2
+]
+initialization_eqs = [
+    Λ.Ω + r.Ω + m.Ω ~ 1
+]
+M1 = ODESystem(eqs, t, [Ω, a], []; initialization_eqs, name = :M)
+M1 = compose(M1, r, m, Λ)
+M1 = complete(M1; flatten = false)
+M1s = structural_simplify(M1)
+```
+Of course, we can solve this ODE as it is:
+```@example changeivar
+prob = ODEProblem(M1s, [M1.a => 1.0, M1.r.Ω => 5e-5, M1.m.Ω => 0.3], (0.0, -3.3), [])
+sol = solve(prob, Tsit5(); reltol = 1e-5)
+@assert Symbol(sol.retcode) == :Unstable # surrounding text assumes this was unstable # hide
+plot(sol, idxs = [M1.a, M1.r.Ω/M1.Ω, M1.m.Ω/M1.Ω, M1.Λ.Ω/M1.Ω])
+```
+The solver becomes unstable due to stiffness.
+Also notice the interesting dynamics taking place towards the end of the integration (in the early universe), which gets compressed into a very small time interval.
+These ODEs would benefit from being defined with respect to a logarithmic "time" that better captures the evolution of the universe through *orders of magnitude* of time.
+
+It is therefore common to write these ODEs in terms of $b = \ln a$.
+To do this, we will change the independent variable in two stages; from $t$ to $a$ to $b$.
+Notice that $\mathrm{d}a/\mathrm{d}t > 0$ provided that $\Omega > 0$, and $\mathrm{d}b/\mathrm{d}a > 0$, so the transformation is well-defined.
+First, we transform from $t$ to $a$:
+```@example changeivar
+M2 = change_independent_variable(M1, M1.a; dummies = true)
+```
+Unlike the original, notice that this system is *non-autonomous* because the independent variable $a$ appears explicitly in the equations!
+This means that to change the independent variable from $a$ to $b$, we must provide not only the rate of change relation $db(a)/da = \exp(-b)$, but *also* the equation $a(b) = \exp(b)$ so $a$ can be eliminated in favor of $b$:
+```@example changeivar
+a = M2.a
+Da = Differential(a)
+@variables b(a)
+M3 = change_independent_variable(M2, b, [Da(b) ~ exp(-b), a ~ exp(b)])
+```
+We can now solve and plot the ODE in terms of $b$:
+```@example changeivar
+M3s = structural_simplify(M3; allow_symbolic = true)
+prob = ODEProblem(M3s, [M3.r.Ω => 5e-5, M3.m.Ω => 0.3], (0, -15), [])
+sol = solve(prob, Tsit5())
+@assert Symbol(sol.retcode) == :Success # surrounding text assumes the solution was successful # hide
+plot(sol, idxs = [M3.r.Ω/M3.Ω, M3.m.Ω/M3.Ω, M3.Λ.Ω/M3.Ω])
+```
+Notice that the variables vary "more nicely" with respect to $b$ than $t$, making the solver happier and avoiding numerical problems.

src/ModelingToolkit.jl

@@ -255,7 +255,7 @@ export isinput, isoutput, getbounds, hasbounds, getguess, hasguess, isdisturbanc
        tunable_parameters, isirreducible, getdescription, hasdescription,
        hasunit, getunit, hasconnect, getconnect,
        hasmisc, getmisc, state_priority
-export ode_order_lowering, dae_order_lowering, liouville_transform
+export ode_order_lowering, dae_order_lowering, liouville_transform, change_independent_variable
 export PDESystem
 export Differential, expand_derivatives, @derivatives
 export Equation, ConstrainedEquation

src/systems/diffeqs/basic_transformations.jl

@@ -14,26 +14,20 @@ the final value of `trJ` is the probability of ``u(t)``.
 Example:
 
 ```julia
-using ModelingToolkit, OrdinaryDiffEq, Test
+using ModelingToolkit, OrdinaryDiffEq
 
 @independent_variables t
 @parameters α β γ δ
 @variables x(t) y(t)
 D = Differential(t)
+eqs = [D(x) ~ α*x - β*x*y, D(y) ~ -δ*y + γ*x*y]
+@named sys = ODESystem(eqs, t)
 
-eqs = [D(x) ~ α*x - β*x*y,
-       D(y) ~ -δ*y + γ*x*y]
-
-sys = ODESystem(eqs)
 sys2 = liouville_transform(sys)
-@variables trJ
-
-u0 = [x => 1.0,
-      y => 1.0,
-      trJ => 1.0]
-
-prob = ODEProblem(complete(sys2),u0,tspan,p)
-sol = solve(prob,Tsit5())
+sys2 = complete(sys2)
+u0 = [x => 1.0, y => 1.0, sys2.trJ => 1.0]
+prob = ODEProblem(sys2, u0, tspan, p)
+sol = solve(prob, Tsit5())
 ```
 
 Where `sol[3,:]` is the evolution of `trJ` over time.
@@ -46,12 +40,167 @@ Optimal Transport Approach
 Abhishek Halder, Kooktae Lee, and Raktim Bhattacharya
 https://abhishekhalder.bitbucket.io/F16ACC2013Final.pdf
 """
-function liouville_transform(sys::AbstractODESystem)
+function liouville_transform(sys::AbstractODESystem; kwargs...)
     t = get_iv(sys)
     @variables trJ
     D = ModelingToolkit.Differential(t)
     neweq = D(trJ) ~ trJ * -tr(calculate_jacobian(sys))
     neweqs = [equations(sys); neweq]
     vars = [unknowns(sys); trJ]
-    ODESystem(neweqs, t, vars, parameters(sys), checks = false)
+    ODESystem(neweqs, t, vars, parameters(sys); checks = false, name = nameof(sys), kwargs...)
+end
+
+"""
+    change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummies = false, simplify = true, fold = false, kwargs...)
+
+Transform the independent variable (e.g. ``t``) of the ODE system `sys` to a dependent variable `iv` (e.g. ``u(t)``).
+An equation in `sys` must define the rate of change of the new independent variable (e.g. ``du(t)/dt``).
+This or other additional equations can also be specified through `eqs`.
+
+The transformation is well-defined when the mapping between the new and old independent variables are one-to-one.
+This is satisfied if one is a strictly increasing function of the other (e.g. ``du(t)/dt > 0`` or ``du(t)/dt < 0``).
+
+Keyword arguments
+=================
+If `dummies`, derivatives of the new independent variable with respect to the old one are expressed through dummy equations; otherwise they are explicitly inserted into the equations.
+If `simplify`, these dummy expressions are simplified and often give a tidier transformation.
+If `fold`, internal substitutions will evaluate numerical expressions.
+Additional keyword arguments `kwargs...` are forwarded to the constructor that rebuilds `sys`.
+
+Usage before structural simplification
+======================================
+The variable change must take place before structural simplification.
+Subsequently, consider passing `allow_symbolic = true` to `structural_simplify(sys)` to reduce the number of unknowns, with the understanding that the transformation is well-defined.
+
+Usage with non-autonomous systems
+=================================
+If `sys` is autonomous (i.e. ``t`` appears explicitly in its equations), it is often desirable to also pass an algebraic equation relating the new and old independent variables (e.g. ``t = f(u(t))``).
+Otherwise the transformed system will be underdetermined and cannot be structurally simplified without additional changes.
+
+Usage with hierarchical systems
+===============================
+It is recommended that `iv` is a non-namespaced variable in `sys`.
+This means it can belong to the top-level system or be a variable in a subsystem declared with `GlobalScope`.
+
+Example
+=======
+Consider a free fall with constant horizontal velocity.
+Physics naturally describes position as a function of time.
+By changing the independent variable, it can be reformulated for vertical position as a function of horizontal position:
+```julia
+julia> @variables x(t) y(t);
+
+julia> @named M = ODESystem([D(D(y)) ~ -9.81, D(x) ~ 10.0], t);
+
+julia> M′ = change_independent_variable(complete(M), x);
+
+julia> unknowns(M′)
+1-element Vector{SymbolicUtils.BasicSymbolic{Real}}:
+ y(x)
+```
+"""
+function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummies = false, simplify = true, fold = false, kwargs...)
+    iv2_of_iv1 = unwrap(iv) # e.g. u(t)
+    iv1 = get_iv(sys) # e.g. t
+
+    if !iscomplete(sys)
+        error("System $(nameof(sys)) is incomplete. Complete it first!")
+    elseif isscheduled(sys)
+        error("System $(nameof(sys)) is structurally simplified. Change independent variable before structural simplification!")
+    elseif !iscall(iv2_of_iv1) || !isequal(only(arguments(iv2_of_iv1)), iv1)
+        error("Variable $iv is not a function of the independent variable $iv1 of the system $(nameof(sys))!")
+    end
+
+    iv1name = nameof(iv1) # e.g. :t
+    iv2name = nameof(operation(iv2_of_iv1)) # e.g. :u
+    iv2, = @independent_variables $iv2name # e.g. u
+    iv1_of_iv2, = @variables $iv1name(iv2) # inverse in case sys is autonomous; e.g. t(u)
+    iv1_of_iv2 = GlobalScope(iv1_of_iv2) # do not namespace old independent variable as new dependent variable
+    D1 = Differential(iv1) # e.g. d/d(t)
+    D2 = Differential(iv2_of_iv1) # e.g. d/d(u(t))
+
+    # 1) Utility that performs the chain rule on an expression, e.g. (d/dt)(f(t)) -> (d/dt)(f(u(t))) -> df(u(t))/du(t) * du(t)/dt
+    function chain_rule(ex)
+        vars = get_variables(ex)
+        for var_of_iv1 in vars # loop through e.g. f(t)
+            if iscall(var_of_iv1) && isequal(only(arguments(var_of_iv1)), iv1) && !isequal(var_of_iv1, iv2_of_iv1) # handle e.g. f(t) -> f(u(t)), but not u(t) -> u(u(t))
+                var_of_iv2 = substitute(var_of_iv1, iv1 => iv2_of_iv1) # e.g. f(t) -> f(u(t))
+                ex = substitute(ex, var_of_iv1 => var_of_iv2)
+            end
+        end
+        ex = expand_derivatives(ex, simplify) # expand chain rule, e.g. (d/dt)(f(u(t)))) -> df(u(t))/du(t) * du(t)/dt
+        return ex
+    end
+
+    # 2) Find e.g. du/dt in equations, then calculate e.g. d²u/dt², ...
+    eqs = [eqs; get_eqs(sys)] # all equations (system-defined + user-provided) we may use
+    idxs = findall(eq -> isequal(eq.lhs, D1(iv2_of_iv1)), eqs)
+    if length(idxs) != 1
+        error("Exactly one equation for $D1($iv2_of_iv1) was not specified!")
+    end
+    div2_of_iv1_eq = popat!(eqs, only(idxs)) # get and remove e.g. du/dt = ... (may be added back later as a dummy)
+    div2_of_iv1 = chain_rule(div2_of_iv1_eq.rhs)
+    if isequal(div2_of_iv1, 0) # e.g. du/dt ~ 0
+        error("Independent variable transformation $(div2_of_iv1_eq) is singular!")
+    end
+    ddiv2_of_iv1 = chain_rule(D1(div2_of_iv1)) # TODO: implement higher orders (order >= 3) derivatives with a loop
+
+    # 3) If requested, insert extra dummy equations for e.g. du/dt, d²u/dt², ...
+    #    Otherwise, replace all these derivatives by their explicit expressions
+    if dummies
+        div2name = Symbol(iv2name, :_, iv1name) # e.g. :u_t # TODO: customize
+        ddiv2name = Symbol(div2name, iv1name) # e.g. :u_tt # TODO: customize
+        div2, ddiv2 = @variables $div2name(iv2) $ddiv2name(iv2) # e.g. u_t(u), u_tt(u)
+        div2, ddiv2 = GlobalScope.([div2, ddiv2]) # do not namespace dummies in new system
+        eqs = [[div2 ~ div2_of_iv1, ddiv2 ~ ddiv2_of_iv1]; eqs] # add dummy equations
+        @set! sys.unknowns = [get_unknowns(sys); [div2, ddiv2]] # add dummy variables
+    else
+        div2 = div2_of_iv1
+        ddiv2 = ddiv2_of_iv1
+    end
+    @set! sys.eqs = eqs # add extra equations we derived before starting transformation process
+
+    # 4) Transform everything from old to new independent variable, e.g. t -> u.
+    #    Substitution order matters! Must begin with highest order to get D(D(u(t))) -> u_tt(u).
+    #    If we had started with the lowest order, we would get D(D(u(t))) -> D(u_t(u)) -> 0!
+    iv1_to_iv2_subs = [ # a vector ensures substitution order
+        D1(D1(iv2_of_iv1)) => ddiv2 # order 2, e.g. D(D(u(t))) -> u_tt(u)
+        D1(iv2_of_iv1) => div2 # order 1, e.g. D(u(t)) -> u_t(u)
+        iv2_of_iv1 => iv2 # order 0, e.g. u(t) -> u
+        iv1 => iv1_of_iv2 # in case sys was autonomous, e.g. t -> t(u)
+    ]
+    function transform(ex)
+        ex = chain_rule(ex)
+        for sub in iv1_to_iv2_subs
+            ex = substitute(ex, sub; fold)
+        end
+        return ex
+    end
+    function transform(sys::AbstractODESystem)
+        eqs = map(transform, get_eqs(sys))
+        unknowns = map(transform, get_unknowns(sys))
+        unknowns = filter(var -> !isequal(var, iv2), unknowns) # remove e.g. u
+        ps = map(transform, get_ps(sys))
+        ps = filter(!isinitial, ps) # remove Initial(...) # # TODO: shouldn't have to touch this
+        observed = map(transform, get_observed(sys))
+        initialization_eqs = map(transform, get_initialization_eqs(sys))
+        parameter_dependencies = map(transform, get_parameter_dependencies(sys))
+        defaults = Dict(transform(var) => transform(val) for (var, val) in get_defaults(sys))
+        guesses = Dict(transform(var) => transform(val) for (var, val) in get_guesses(sys))
+        assertions = Dict(transform(condition) => msg for (condition, msg) in get_assertions(sys))
+        systems = get_systems(sys) # save before reconstructing system
+        wascomplete = iscomplete(sys) # save before reconstructing system
+        sys = typeof(sys)( # recreate system with transformed fields
+            eqs, iv2, unknowns, ps; observed, initialization_eqs, parameter_dependencies, defaults, guesses,
+            assertions, name = nameof(sys), description = description(sys), kwargs...
+        )
+        systems = map(transform, systems) # recurse through subsystems
+        sys = compose(sys, systems) # rebuild hierarchical system
+        if wascomplete
+            wasflat = isempty(systems)
+            sys = complete(sys; flatten = wasflat) # complete output if input was complete
+        end
+        return sys
+    end
+    return transform(sys)
 end