Clean up and format

Author: hersle

docs/src/tutorials/change_independent_variable.md

@@ -2,42 +2,46 @@
 
 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:
+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 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).
+    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 gravitational field.
+Consider a projectile shot with some initial velocity in a vertical gravitational field with constant horizontal velocity.
+
 ```@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; initialization_eqs = [
-    D(x) ~ D(y) # equal initial horizontal and vertical velocity (45°)
-], name = :M)
+@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.
+
+    - 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$.
+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)
@@ -46,17 +50,21 @@ M2s = structural_simplify(M2; allow_symbolic = true)
 @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
@@ -65,67 +73,75 @@ 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 * Ω]
+    eqs = [D(Ω) ~ -3(1 + w) * D(a) / a * Ω]
     return ODESystem(eqs, t, [Ω], []; kw...)
 end
-@named r = species(1//3) # radiation
+@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
-]
+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.Ω])
+plot(sol, idxs = [M1.a, M1.r.Ω / M1.Ω, M1.m.Ω / M1.Ω, M1.Λ.Ω / M1.Ω])
 ```
-The solver becomes unstable due to stiffness.
+
+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.
+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; 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.
+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
 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.Ω])
+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/systems/diffeqs/basic_transformations.jl

@@ -47,11 +47,17 @@ function liouville_transform(sys::AbstractODESystem; kwargs...)
     neweq = D(trJ) ~ trJ * -tr(calculate_jacobian(sys))
     neweqs = [equations(sys); neweq]
     vars = [unknowns(sys); trJ]
-    ODESystem(neweqs, t, vars, parameters(sys); checks = false, name = nameof(sys), kwargs...)
+    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, kwargs...)
+    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.
@@ -68,13 +74,13 @@ Any extra equations `eqs` involving the new and old independent variables will b
 # 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.
+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), 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.
-If an algebraic relation is not known, consider using `add_old_diff`.
+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
 
@@ -102,7 +108,10 @@ julia> unknowns(M)
  yˍx(x)
 ```
 """
-function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; add_old_diff = false, simplify = true, fold = false, kwargs...)
+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
 
@@ -114,6 +123,7 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; add_o
         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
@@ -127,23 +137,22 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; add_o
     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 before starting transformation process
-    @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) # add dummy variables and old independent variable as a function of the new one
+    @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)
-    # Then use it to transform everything in the system!
+    # 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) # is the expression of the form f(t)?
+            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) Expand with chain rule until nothing changes anymore
+        # 2) Repeatedly expand chain rule until nothing changes anymore
         orgex = nothing
         while !isequal(ex, orgex)
             orgex = ex # save original
@@ -156,22 +165,25 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; add_o
         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
+        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))
+        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))
+        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,
+            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
@@ -182,6 +194,5 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; add_o
         end
         return sys
     end
-
     return transform(sys)
 end