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:

One may want $y(x)$ as a function of $x$ instead of $(x(t), y(t))$ as a function of $t$
Some differential equations vary more nicely (e.g. less stiff) with respect to one independent variable than another.
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 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.

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)

\[ \begin{align} \frac{\mathrm{d} \mathtt{yˍt}\left( t \right)}{\mathrm{d}t} &= - g \\ \frac{\mathrm{d} x\left( t \right)}{\mathrm{d}t} &= v \\ \frac{\mathrm{d} y\left( t \right)}{\mathrm{d}t} &= \mathtt{yˍt}\left( t \right) \end{align} \]

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.

M2 = change_independent_variable(M1, x)
M2s = structural_simplify(M2; allow_symbolic = true)

\[ \begin{align} \frac{\mathrm{d} y\left( x \right)}{\mathrm{d}x} &= \mathtt{yˍx}\left( x \right) \\ \frac{\mathrm{d} \mathtt{yˍx}\left( x \right)}{\mathrm{d}x} &= \frac{g + \mathtt{yˍx}\left( x \right) \mathtt{xˍt}\left( x \right) \mathtt{xˍtx}\left( x \right)}{ - \left( \mathtt{xˍt}\left( x \right) \right)^{2}} \end{align} \]

The derivatives are now with respect to the new independent variable $x$, which can be accessed with M2.x.

Info

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

using OrdinaryDiffEq, Plots
prob = ODEProblem(M2s, [M2s.y => 0.0], [0.0, 10.0], [v => 8.0]) # throw 10 meters
sol = solve(prob, Tsit5())
plot(sol; idxs = M2.y) # must index by M2.y = y(x); not M1.y = y(t)!

Usage tips

Look up the documentation of change_independent_variable 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 (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 describe the expansion of the universe. In terms of conformal time $t$, they can be written

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

\[ \begin{align} \frac{\mathrm{d} a\left( t \right)}{\mathrm{d}t} &= \left( a\left( t \right) \right)^{2} \sqrt{\Omega\left( t \right)} \\ \frac{\mathrm{d} \mathtt{m.\Omega}\left( t \right)}{\mathrm{d}t} &= - 3 \sqrt{\Omega\left( t \right)} a\left( t \right) \mathtt{m.\Omega}\left( t \right) \\ \frac{\mathrm{d} \mathtt{r.\Omega}\left( t \right)}{\mathrm{d}t} &= - 4 \sqrt{\Omega\left( t \right)} a\left( t \right) \mathtt{r.\Omega}\left( t \right) \\ \frac{\mathrm{d} \mathtt{\Lambda.\Omega}\left( t \right)}{\mathrm{d}t} &= -0 \end{align} \]

Of course, we can solve this ODE as it is:

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

M2 = change_independent_variable(M1, M1.a)

\[ \begin{align} \Omega\left( a \right) &= \mathtt{r.\Omega}\left( a \right) + \mathtt{\Lambda.\Omega}\left( a \right) + \mathtt{m.\Omega}\left( a \right) \\ \mathtt{aˍt}\left( a \right) &= a^{2} \sqrt{\Omega\left( a \right)} \\ \mathtt{aˍt}\left( a \right) \frac{\mathrm{d} \mathtt{r.\Omega}\left( a \right)}{\mathrm{d}a} &= \frac{ - 4 \mathtt{aˍt}\left( a \right) \mathtt{r.\Omega}\left( a \right)}{a} \\ \mathtt{aˍt}\left( a \right) \frac{\mathrm{d} \mathtt{m.\Omega}\left( a \right)}{\mathrm{d}a} &= \frac{ - 3 \mathtt{aˍt}\left( a \right) \mathtt{m.\Omega}\left( a \right)}{a} \\ \mathtt{aˍt}\left( a \right) \frac{\mathrm{d} \mathtt{\Lambda.\Omega}\left( a \right)}{\mathrm{d}a} &= 0 \end{align} \]

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$:

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

\[ \begin{align} \Omega\left( b \right) &= \mathtt{\Lambda.\Omega}\left( b \right) + \mathtt{m.\Omega}\left( b \right) + \mathtt{r.\Omega}\left( b \right) \\ \mathtt{aˍt}\left( b \right) &= \left( a\left( b \right) \right)^{2} \sqrt{\Omega\left( b \right)} \\ \mathtt{bˍa}\left( b \right) &= e^{ - b} \\ a\left( b \right) &= e^{b} \\ \mathtt{bˍa}\left( b \right) \frac{\mathrm{d} \mathtt{r.\Omega}\left( b \right)}{\mathrm{d}b} \mathtt{aˍt}\left( b \right) &= \frac{ - 4 \mathtt{r.\Omega}\left( b \right) \mathtt{aˍt}\left( b \right)}{a\left( b \right)} \\ \mathtt{bˍa}\left( b \right) \frac{\mathrm{d} \mathtt{m.\Omega}\left( b \right)}{\mathrm{d}b} \mathtt{aˍt}\left( b \right) &= \frac{ - 3 \mathtt{m.\Omega}\left( b \right) \mathtt{aˍt}\left( b \right)}{a\left( b \right)} \\ \frac{\mathrm{d} \mathtt{\Lambda.\Omega}\left( b \right)}{\mathrm{d}b} \mathtt{bˍa}\left( b \right) \mathtt{aˍt}\left( b \right) &= 0 \end{align} \]

We can now solve and plot the ODE in terms of $b$:

M3s = structural_simplify(M3; allow_symbolic = true)
prob = ODEProblem(M3s, [M3s.r.Ω => 5e-5, M3s.m.Ω => 0.3], (0, -15), [])
sol = solve(prob, Tsit5())
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.