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,146 @@
+# 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$.
+However, in many cases there are good reasons for changing the 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) 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 vertical gravitational field with constant horizontal velocity.
+
+```@example changeivar
+using ModelingToolkit
+using ModelingToolkit: t_nounits as t, D_nounits as D
+@variables x(t) y(t)
+@parameters g=9.81 v # gravitational acceleration and horizontal velocity
+eqs = [D(D(y)) ~ -g, D(x) ~ v]
+initialization_eqs = [D(x) ~ D(y)] # 45° initial angle
+M1 = ODESystem(eqs, t; initialization_eqs, name = :M)
+M1s = structural_simplify(M1)
+@assert length(equations(M1s)) == 3 # hide
+M1s # hide
+```
+
+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$, so $x$ and $t$ are one-to-one.
+
+```@example changeivar
+M2 = change_independent_variable(M1, x)
+M2s = structural_simplify(M2; allow_symbolic = true)
+# a sanity test on the 10 x/y variables that are accessible to the user # hide
+@assert allequal([x, M1s.x]) # hide
+@assert allequal([M2.x, M2s.x]) # hide
+@assert allequal([y, M1s.y]) # hide
+@assert allunique([M1.x, M1.y, M2.y, M2s.y]) # hide
+@assert length(equations(M2s)) == 2 # hide
+M2s # display this # hide
+```
+
+The derivatives are now with respect to the new independent variable $x$, which can be accessed with `M2.x`.
+
+!!! warn
+
+    At this point `x`, `M1.x`, `M1s.x`, `M2.x`, `M2s.x` are *three* different variables.
+    Meanwhile `y`, `M1.y`, `M1s.y`, `M2.y` and `M2s.y` are *four* different variables.
+    It can be instructive to inspect these yourself to see their subtle differences.
+
+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, [M2s.y => 0.0], [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)!
+```
+
+!!! tip "Usage tips"
+
+    Look up the documentation of [`change_independent_variable`](@ref) for tips on how to use it.
+
+    For example, if you also need $t(x)$, you can tell it to add a differential equation for the old independent variable in terms of the new one using the [inverse function rule](https://en.wikipedia.org/wiki/Inverse_function_rule) (here $\mathrm{d}t/\mathrm{d}x = 1 / (\mathrm{d}x/\mathrm{d}t)$). If you know an analytical expression between the independent variables (here $t = x/v$), you can also pass it directly to the function to avoid the extra differential equation.
+
+## 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.Ω + Λ.Ω, D(a) ~ √(Ω) * a^2]
+initialization_eqs = [Λ.Ω + r.Ω + m.Ω ~ 1]
+M1 = ODESystem(eqs, t, [Ω, a], []; initialization_eqs, name = :M)
+M1 = compose(M1, r, m, Λ)
+M1s = structural_simplify(M1)
+```
+
+Of course, we can solve this ODE as it is:
+
+```@example changeivar
+prob = ODEProblem(M1s, [M1s.a => 1.0, M1s.r.Ω => 5e-5, M1s.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.Ω])
+```
+
+But 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, rather than linear 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; first from $t$ to $a$, and then from $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 since $t \leftrightarrow a \leftrightarrow b$ are one-to-one.
+First, we transform from $t$ to $a$:
+
+```@example changeivar
+M2 = change_independent_variable(M1, M1.a)
+@assert !ModelingToolkit.isautonomous(M2) # hide
+M2 # hide
+```
+
+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 # get independent variable of M2
+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, [M3s.r.Ω => 5e-5, M3s.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,8 @@ 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,159 @@ 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 = [];
+        add_old_diff = false, simplify = true, fold = false
+    )
+
+Transform the independent variable (e.g. ``t``) of the ODE system `sys` to a dependent variable `iv` (e.g. ``u(t)``).
+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``).
+
+Any extra equations `eqs` involving the new and old independent variables will be taken into account in the transformation.
+
+# Keyword arguments
+
+- `add_old_diff`: Whether to add a differential equation for the old independent variable in terms of the new one using the inverse function rule ``dt/du = 1/(du/dt)``.
+- `simplify`: Whether expanded derivative expressions are simplified. This can give a tidier transformation.
+- `fold`: Whether internal substitutions will evaluate numerical expressions.
+
+# Usage before structural simplification
+
+The variable change must take place before structural simplification.
+In following calls to `structural_simplify`, consider passing `allow_symbolic = true` to avoid undesired constraint equations between between dummy variables.
+
+# Usage with non-autonomous systems
+
+If `sys` is non-autonomous (i.e. ``t`` appears explicitly in its equations), consider passing an algebraic equation relating the new and old independent variables (e.g. ``t = f(u(t))``).
+Otherwise the transformed system can be underdetermined.
+If an algebraic relation is not known, consider using `add_old_diff` instead.
+
+# 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(D(x)) ~ 0.0], t);
+
+julia> M = change_independent_variable(M, x);
+
+julia> M = structural_simplify(M; allow_symbolic = true);
+
+julia> unknowns(M)
+3-element Vector{SymbolicUtils.BasicSymbolic{Real}}:
+ xˍt(x)
+ y(x)
+ yˍx(x)
+```
+"""
+function change_independent_variable(
+        sys::AbstractODESystem, iv, eqs = [];
+        add_old_diff = false, simplify = true, fold = false
+)
+    iv2_of_iv1 = unwrap(iv) # e.g. u(t)
+    iv1 = get_iv(sys) # e.g. t
+
+    if is_dde(sys)
+        error("System $(nameof(sys)) contains delay differential equations (DDEs). This is currently not supported!")
+    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
+
+    # Set up intermediate and final variables for the transformation
+    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, = GlobalScope.(@variables $iv1name(iv2)) # inverse, e.g. t(u), global because iv1 has no namespacing in sys
+    D1 = Differential(iv1) # e.g. d/d(t)
+    div2_of_iv1 = GlobalScope(default_toterm(D1(iv2_of_iv1))) # e.g. uˍt(t)
+    div2_of_iv2 = substitute(div2_of_iv1, iv1 => iv2) # e.g. uˍt(u)
+    div2_of_iv2_of_iv1 = substitute(div2_of_iv2, iv2 => iv2_of_iv1) # e.g. uˍt(u(t))
+
+    # If requested, add a differential equation for the old independent variable as a function of the old one
+    if add_old_diff
+        eqs = [eqs; Differential(iv2)(iv1_of_iv2) ~ 1 / div2_of_iv2] # e.g. dt(u)/du ~ 1 / uˍt(u) (https://en.wikipedia.org/wiki/Inverse_function_rule)
+    end
+    @set! sys.eqs = [get_eqs(sys); eqs] # add extra equations we derived
+    @set! sys.unknowns = [get_unknowns(sys); [iv1, div2_of_iv1]] # add new variables, will be transformed to e.g. t(u) and uˍt(u)
+
+    # Create a utility that performs the chain rule on an expression, followed by insertion of the new independent variable:
+    # e.g. (d/dt)(f(t)) -> (d/dt)(f(u(t))) -> df(u(t))/du(t) * du(t)/dt -> df(u)/du * uˍt(u)
+    function transform(ex)
+        # 1) Replace the argument of every function; e.g. f(t) -> f(u(t))
+        for var in vars(ex; op = Nothing) # loop over all variables in expression (op = Nothing prevents interpreting "D(f(t))" as one big variable)
+            is_function_of_iv1 = iscall(var) && isequal(only(arguments(var)), iv1) # of the form f(t)?
+            if is_function_of_iv1 && !isequal(var, iv2_of_iv1) # prevent e.g. u(t) -> u(u(t))
+                var_of_iv1 = var # e.g. f(t)
+                var_of_iv2_of_iv1 = substitute(var_of_iv1, iv1 => iv2_of_iv1) # e.g. f(u(t))
+                ex = substitute(ex, var_of_iv1 => var_of_iv2_of_iv1; fold)
+            end
+        end
+        # 2) Repeatedly expand chain rule until nothing changes anymore
+        orgex = nothing
+        while !isequal(ex, orgex)
+            orgex = ex # save original
+            ex = expand_derivatives(ex, simplify) # expand chain rule, e.g. (d/dt)(f(u(t)))) -> df(u(t))/du(t) * du(t)/dt
+            ex = substitute(ex, D1(iv2_of_iv1) => div2_of_iv2_of_iv1; fold) # e.g. du(t)/dt -> uˍt(u(t))
+        end
+        # 3) Set new independent variable
+        ex = substitute(ex, iv2_of_iv1 => iv2; fold) # set e.g. u(t) -> u everywhere
+        ex = substitute(ex, iv1 => iv1_of_iv2; fold) # set e.g. t -> t(u) everywhere
+        return ex
+    end
+
+    # Use the utility function to transform everything in the system!
+    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(ass) => msg for (ass, 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)
+        )
+        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