Skip to content

Commit

Permalink
Clean up independent variable change implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
@hersle
hersle committed Mar 4, 2025
1 parent 98125b7 commit 701ac0e
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 101 deletions.
2 changes: 1 addition & 1 deletion docs/src/tutorials/change_independent_variable.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
146 changes: 61 additions & 85 deletions src/systems/diffeqs/basic_transformations.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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``).
Expand All @@ -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
Expand All @@ -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
55 changes: 40 additions & 15 deletions test/basic_transformations.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down

0 comments on commit 701ac0e

Please sign in to comment.