Change independent variable of hierarchical systems

Author: hersle

docs/src/tutorials/change_independent_variable.md

@@ -42,7 +42,7 @@ There are at least three ways of answering this:
 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) |> complete
+M2 = change_independent_variable(M1, M1.x; dummies = true)
 @assert M2.x isa Num # hide
 M2s = structural_simplify(M2; allow_symbolic = true)
 ```
@@ -56,29 +56,38 @@ sol = solve(prob, Tsit5())
 plot(sol; idxs = M2.y) # must index by M2.y = y(x); not M1.y = y(t)!
 ```
 
-## 2. Reduce stiffness by changing to a logarithmic time axis
+## 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) Ωr(t) Ωm(t) ΩΛ(t)
-M1 = ODESystem([
+@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
-    Ω ~ Ωr + Ωm + ΩΛ
-    D(Ωm) ~ -3 * D(a)/a * Ωm
-    D(Ωr) ~ -4 * D(a)/a * Ωr
-    D(ΩΛ) ~ 0
-], t; initialization_eqs = [
-    ΩΛ + Ωr + Ωm ~ 1
-], name = :M) |> complete
+]
+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), [])
+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.Ω])
+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.
@@ -89,22 +98,22 @@ To do this, we will change the independent variable in two stages; from $t$ to $
 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) |> complete
+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)]) |> complete
+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), [])
+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.Ω])
+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

@@ -77,6 +77,11 @@ 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.
@@ -102,8 +107,6 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummi
         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 !isempty(get_systems(sys))
-        error("System $(nameof(sys)) is hierarchical. Flatten it first!") # TODO: implement and allow?
     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
@@ -112,25 +115,25 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummi
     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(var_of_iv1, iv2_of_iv1) # handle e.g. f(t) -> f(u(t)), but not u(t) -> u(u(t))
-                varname = nameof(operation(var_of_iv1)) # e.g. :f
-                var_of_iv2, = @variables $varname(iv2_of_iv1) # e.g. f(u(t))
-                ex = substitute(ex, var_of_iv1 => var_of_iv2; fold) # e.g. f(t) -> f(u(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 = [get_eqs(sys); eqs] # all equations (system-defined + user-provided) we may use
+    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!")
@@ -148,11 +151,14 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummi
         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)
-        eqs = [eqs; [div2 ~ div2_of_iv1, ddiv2 ~ ddiv2_of_iv1]] # add dummy equations
+        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).
@@ -170,18 +176,31 @@ function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummi
         end
         return ex
     end
-    eqs = map(transform, eqs) # we derived and added equations to eqs; they are not in get_eqs(sys)!
-    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))
-
-    # 5) Recreate system with transformed fields
-    return typeof(sys)(
-        eqs, iv2;
-        observed, initialization_eqs, parameter_dependencies, defaults, guesses, assertions,
-        name = nameof(sys), description = description(sys), kwargs...
-    ) |> complete # input system must be complete, so complete the output 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(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

test/basic_transformations.jl

@@ -62,25 +62,34 @@ end
     @independent_variables t
     D = Differential(t)
     @variables a(t) ȧ(t) Ω(t) ϕ(t)
+    a, ȧ = GlobalScope.([a, ȧ])
+    species(w; kw...) = ODESystem([D(Ω) ~ -3(1 + w) * D(a)/a * Ω], t, [Ω], []; kw...)
+    @named r = species(1//3)
+    @named m = species(0)
+    @named Λ = species(-1)
     eqs = [
-        Ω ~ 123
+        Ω ~ r.Ω + m.Ω + Λ.Ω
         D(a) ~ ȧ
         ȧ ~ √(Ω) * a^2
         D(D(ϕ)) ~ -3*D(a)/a*D(ϕ)
     ]
-    M1 = ODESystem(eqs, t; name = :M) |> complete
+    M1 = ODESystem(eqs, t, [Ω, a, ȧ, ϕ], []; name = :M)
+    M1 = compose(M1, r, m, Λ)
+    M1 = complete(M1; flatten = false)
 
     # Apply in two steps, where derivatives are defined at each step: first t -> a, then a -> b
     M2 = change_independent_variable(M1, M1.a; dummies = true)
-    @independent_variables a
-    @variables ȧ(a) Ω(a) ϕ(a) a_t(a) a_tt(a)
+    a, ȧ, Ω, Ωr, Ωm, ΩΛ, ϕ, a_t, a_tt = M2.a, M2.ȧ, M2.Ω, M2.r.Ω, M2.m.Ω, M2.Λ.Ω, M2.ϕ, M2.a_t, M2.a_tt
     Da = Differential(a)
     @test Set(equations(M2)) == Set([
-        a_tt*Da(ϕ) + a_t^2*(Da^2)(ϕ) ~ -3*a_t^2/a*Da(ϕ)
-        ȧ ~ √(Ω) * a^2
-        Ω ~ 123
         a_t ~ ȧ # 1st order dummy equation
         a_tt ~ Da(ȧ) * a_t # 2nd order dummy equation
+        Ω ~ Ωr + Ωm + ΩΛ
+        ȧ ~ √(Ω) * a^2
+        a_tt*Da(ϕ) + a_t^2*(Da^2)(ϕ) ~ -3*a_t^2/a*Da(ϕ)
+        a_t*Da(Ωr) ~ -4*Ωr*a_t/a
+        a_t*Da(Ωm) ~ -3*Ωm*a_t/a
+        a_t*Da(ΩΛ) ~ 0
     ])
 
     @variables b(M2.a)
@@ -143,6 +152,4 @@ end
     @test_throws "not specified" change_independent_variable(M, v)
     @test_throws "not a function of the independent variable" change_independent_variable(M, y)
     @test_throws "not a function of the independent variable" change_independent_variable(M, z)
-    M = compose(M, M)
-    @test_throws "hierarchical" change_independent_variable(M, M.x)
 end