Clean up independent variable change implementation

Author: hersle

docs/src/tutorials/change_independent_variable.md

@@ -26,7 +26,7 @@ M1 = ODESystem([
 ], t; defaults = [
     y => 0.0
 ], initialization_eqs = [
-    #x ~ 0.0, # TODO: handle?
+    #x ~ 0.0, # TODO: handle? # hide
     D(x) ~ D(y) # equal initial horizontal and vertical velocity (45 °)
 ], name = :M) |> complete
 M1s = structural_simplify(M1)

src/systems/diffeqs/basic_transformations.jl

@@ -51,7 +51,7 @@ function liouville_transform(sys::AbstractODESystem; kwargs...)
 end
 
 """
-    function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummies = false, simplify = true, verbose = false, kwargs...)
+    function 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. ``f(t)``).
 An equation in `sys` must define the rate of change of the new independent variable (e.g. ``df(t)/dt``).
@@ -64,7 +64,7 @@ Keyword arguments
 =================
 If `dummies`, derivatives of the new independent variable 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 `verbose`, the function prints intermediate transformations of equations to aid debugging.
+If `fold`, internal substitutions will evaluate numerical expressions.
 Any additional keyword arguments `kwargs...` are forwarded to the constructor that rebuilds the system.
 
 Usage before structural simplification
@@ -89,118 +89,94 @@ julia> unknowns(M′)
  y(x)
 ```
 """
-function change_independent_variable(sys::AbstractODESystem, iv, eqs = []; dummies = false, simplify = true, verbose = false, kwargs...)
+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("Cannot change independent variable of incomplete system $(nameof(sys))")
+        error("System $(nameof(sys)) is incomplete. Complete it first!")
     elseif isscheduled(sys)
-        error("Cannot change independent variable of structurally simplified system $(nameof(sys))")
+        error("System $(nameof(sys)) is structurally simplified. Change independent variable before structural simplification!")
     elseif !isempty(get_systems(sys))
-        error("Cannot change independent variable of hierarchical system $(nameof(sys)). Flatten it first.") # TODO: implement
+        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
 
-    iv = unwrap(iv)
-    iv1 = get_iv(sys) # e.g. t
-    if !iscall(iv) || !isequal(only(arguments(iv)), iv1)
-        error("New independent variable $iv is not a function of the independent variable $iv1 of the system $(nameof(sys))")
-    elseif !isautonomous(sys) && isempty(findall(eq -> isequal(eq.lhs, iv1), eqs))
-        error("System $(nameof(sys)) is autonomous in $iv1. An equation of the form $iv1 ~ F($iv) must be provided.")
+    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)
+    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))
+            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
 
-    iv2func = iv # e.g. a(t)
-    iv2name = nameof(operation(iv))
-    iv2, = @independent_variables $iv2name # e.g. a
-    D1 = Differential(iv1)
-
-    iv1name = nameof(iv1) # e.g. t
-    iv1func, = @variables $iv1name(iv2) # e.g. t(a)
-
-    eqs = [get_eqs(sys); eqs] # copies system equations to avoid modifying original system
-
-    # 1) Find and compute all necessary expressions for e.g. df/dt, d²f/dt², ...
-    # 1.1) Find the 1st order derivative of the new independent variable (e.g. da(t)/dt = ...), ...
-    div2_div1_idxs = findall(eq -> isequal(eq.lhs, D1(iv2func)), eqs) # index of e.g. da/dt = ...
-    if length(div2_div1_idxs) != 1
-        error("Exactly one equation for $D1($iv2func) was not specified.")
+    # 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
+    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_div1_eq = popat!(eqs, only(div2_div1_idxs)) # get and remove e.g. df/dt = ... (may be added back later)
-    div2_div1 = div2_div1_eq.rhs
-    if isequal(div2_div1, 0)
-        error("Cannot change independent variable from $iv1 to $iv2 with singular transformation $div2_div1_eq.")
+    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
-    # 1.2) ... then compute the 2nd order derivative of the new independent variable
-    div1_div2 = 1 / div2_div1 # TODO: URL reference for clarity
-    ddiv2_ddiv1 = expand_derivatives(-Differential(iv2func)(div1_div2) / div1_div2^3, simplify) # e.g. https://math.stackexchange.com/questions/249253/second-derivative-of-the-inverse-function # TODO: higher orders # TODO: pass simplify here
-    # 1.3) # TODO: handle higher orders (3+) derivatives ...
+    ddiv2_of_iv1 = chain_rule(D1(div2_of_iv1)) # TODO: implement higher orders (order >= 3) derivatives with a loop
 
-    # 2) If requested, insert extra dummy equations for e.g. df/dt, d²f/dt², ...
+    # 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, :_t) # TODO: not always t
-        div2, = @variables $div2name(iv2) # e.g. a_t(a)
-        ddiv2name = Symbol(iv2name, :_tt) # TODO: not always t
-        ddiv2, = @variables $ddiv2name(iv2) # e.g. a_tt(a)
-        eqs = [eqs; [div2 ~ div2_div1, ddiv2 ~ ddiv2_ddiv1]] # add dummy equations
-        derivsubs = [D1(D1(iv2func)) => ddiv2, D1(iv2func) => div2] # order is crucial!
+        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
     else
-        derivsubs = [D1(D1(iv2func)) => ddiv2_ddiv1, D1(iv2func) => div2_div1] # order is crucial!
-    end
-    derivsubs = [derivsubs; [iv2func => iv2, iv1 => iv1func]]
-
-    if verbose
-        # Explain what we just did
-        println("Order 1 (found):    $div2_div1_eq")
-        println("Order 2 (computed): $(D1(div2_div1_eq.lhs) ~ ddiv2_ddiv1)")
-        println()
-        println("Substitutions will be made in this order:")
-        for (n, sub) in enumerate(derivsubs)
-            println("$n: $(sub[1]) => $(sub[2])")
-        end
-        println()
+        div2 = div2_of_iv1
+        ddiv2 = ddiv2_of_iv1
     end
 
-    # 3) Define a transformation function that performs the change of variable on any expression/equation
+    # 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)
-        verbose && println("Step 0: ", ex)
-
-        # Step 1: substitute f(t₁) => f(t₂(t₁)) in all variables in the expression
-        vars = Symbolics.get_variables(ex)
-        for var1 in vars
-            if Symbolics.iscall(var1) && !isequal(var1, iv2func) # && isequal(only(arguments(var1)), iv1) # skip e.g. constants
-                name = nameof(operation(var1))
-                var2, = @variables $name(iv2func)
-                ex = substitute(ex, var1 => var2; fold = false)
-            end
-        end
-        verbose && println("Step 1: ", ex)
-
-        # Step 2: expand out all chain rule derivatives
-        ex = expand_derivatives(ex) # expand out with chain rule to get d(iv2)/d(iv1)
-        verbose && println("Step 2: ", ex)
-
-        # Step 3: substitute d²f/dt², df/dt, ... (to dummy variables or explicit expressions, depending on dummies)
-        for sub in derivsubs
-            ex = substitute(ex, sub)
+        ex = chain_rule(ex)
+        for sub in iv1_to_iv2_subs
+            ex = substitute(ex, sub; fold)
         end
-        verbose && println("Step 3: ", ex)
-        verbose && println()
-
         return ex
     end
-
-    # 4) Transform all fields
-    eqs = map(transform, eqs)
+    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))
-    # TODO: handle subsystems
 
     # 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 # original system had to be complete
+    ) |> complete # input system must be complete, so complete the output system
 end

test/basic_transformations.jl

@@ -59,34 +59,44 @@ end
 end
 
 @testset "Change independent variable (Friedmann equation)" begin
+    @independent_variables t
     D = Differential(t)
-    @variables a(t) ȧ(t) ρr(t) ρm(t) ρΛ(t) ρ(t) P(t) ϕ(t)
-    @parameters Ωr0 Ωm0 ΩΛ0
+    @variables a(t) ȧ(t) Ω(t) ϕ(t)
     eqs = [
-        ρr ~ 3/(8*Num(π)) * Ωr0 / a^4
-        ρm ~ 3/(8*Num(π)) * Ωm0 / a^3
-        ρΛ ~ 3/(8*Num(π)) * ΩΛ0
-        ρ ~ ρr + ρm + ρΛ
-        ȧ ~ √(8*Num(π)/3*ρ*a^4)
+        Ω ~ 123
         D(a) ~ ȧ
+        ȧ ~ √(Ω) * a^2
         D(D(ϕ)) ~ -3*D(a)/a*D(ϕ)
     ]
     M1 = ODESystem(eqs, t; name = :M) |> complete
 
     # Apply in two steps, where derivatives are defined at each step: first t -> a, then a -> b
-    M2 = change_independent_variable(M1, M1.a) |> complete #, D(b) ~ D(a)/a; verbose = true)
+    M2 = change_independent_variable(M1, M1.a; dummies = true)
+    @independent_variables a
+    @variables ȧ(a) Ω(a) ϕ(a) a_t(a) a_tt(a)
+    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
+    ])
+
     @variables b(M2.a)
     M3 = change_independent_variable(M2, b, [Differential(M2.a)(b) ~ exp(-b), M2.a ~ exp(b)])
+
+    M1 = structural_simplify(M1)
     M2 = structural_simplify(M2; allow_symbolic = true)
     M3 = structural_simplify(M3; allow_symbolic = true)
-    @test length(unknowns(M2)) == 2 && length(unknowns(M3)) == 2
+    @test length(unknowns(M3)) == length(unknowns(M2)) == length(unknowns(M1)) - 1
 end
 
 @testset "Change independent variable (simple)" begin
     @variables x(t)
-    Mt = ODESystem([D(x) ~ 2*x], t; name = :M) |> complete # TODO: avoid complete. can avoid it if passing defined $variable directly to change_independent_variable
+    Mt = ODESystem([D(x) ~ 2*x], t; name = :M) |> complete
     Mx = change_independent_variable(Mt, Mt.x; dummies = true)
-    @test (@variables x x_t(x) x_tt(x); Set(equations(Mx)) == Set([x_t ~ 2x, x_tt ~ 4x]))
+    @test (@variables x x_t(x) x_tt(x); Set(equations(Mx)) == Set([x_t ~ 2*x, x_tt ~ 2*x_t]))
 end
 
 @testset "Change independent variable (free fall)" begin
@@ -101,10 +111,25 @@ end
     @test all(isapprox.(sol[Mx.y], sol[Mx.x - g*(Mx.x/v)^2/2]; atol = 1e-10)) # compare to analytical solution (x(t) = v*t, y(t) = v*t - g*t^2/2)
 end
 
-@testset "Change independent variable (autonomous system)" begin
-    M = ODESystem([D(x) ~ t], t; name = :M) |> complete # non-autonomous
-    @test_throws "t ~ F(x(t)) must be provided" change_independent_variable(M, M.x)
-    @test_nowarn change_independent_variable(M, M.x, [t ~ 2*x])
+@testset "Change independent variable (crazy analytical example)" begin
+    @independent_variables t
+    D = Differential(t)
+    @variables x(t) y(t)
+    M1 = ODESystem([ # crazy non-autonomous non-linear 2nd order ODE
+        D(D(y)) ~ D(x)^2 + D(y^3) |> expand_derivatives # expand D(y^3) # TODO: make this test 3rd order
+        D(x) ~ x^4 + y^5 + t^6
+    ], t; name = :M) |> complete
+    M2 = change_independent_variable(M1, M1.x; dummies = true)
+
+    # Compare to pen-and-paper result
+    @independent_variables x
+    Dx = Differential(x)
+    @variables x_t(x) x_tt(x) y(x) t(x)
+    @test Set(equations(M2)) == Set([
+        x_t^2*(Dx^2)(y) + x_tt*Dx(y) ~ x_t^2 + 3*y^2*Dx(y)*x_t # from D(D(y))
+        x_t ~ x^4 + y^5 + t^6 # 1st order dummy equation
+        x_tt ~ 4*x^3*x_t + 5*y^4*Dx(y)*x_t + 6*t^5 # 2nd order dummy equation
+    ])
 end
 
 @testset "Change independent variable (errors)" begin