fix first ODE perturbation example

[karlwessel]

Checklist

• Appropriate tests were added (see very last bullet point below)
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

This PR makes the first part from the Perturbation theory for ODEs example work. This fixes #1994.
The following things needed to be changed:

• collect powers included zero order terms in each order term,
• link to algebraic perturbation example from Symbolics.jl was wrong,
• indexing in definition and usage of def_taylor did not match,
• use enumerate instead of eachindex for symbolic arrays to get a linear index instead of a CartesianIndex (eachindex(IndexLinear(), ...) would also work).

The following things were not changed but would probably improve the example:

• Use Symbolics.coeff in collect_powers, but coeff currently can't be used to get the zeroth order term correctly (see second point in Symbolics.coeff silently returns incorrect results for products that need expansion, unbounded recursion on fractions.  JuliaSymbolics/Symbolics.jl#910).
• Start indexing of y and its derivative at zero, to match the order they represent, but non-one-based indexing isn't well supported yet (see Structural simplification fails on non-1-indexed variable array  #2670).
• It would probably be good to use the @example ... feature from Documenter.jl to add the output of commands to the documentation. This way it would be easier to detect when some change breaks the example.

[ChrisRackauckas]

It would probably be good to use the @example ... feature from Documenter.jl to add the output of commands to the documentation. This way it would be easier to detect when some change breaks the example.

We are doing that over time. Most examples are already updated. We shouldn't merge this one until it's using @example.

[hersle]

Nice! Since I am working on something somewhat similar to this example, just a few suggestions:

1. Agreed, definitely use @example.
2. Can you get away with using Symbolics instead of using Symbolics, SymbolicUtils?
3. To simplify collect_powers(), I have used this solution:

   def_power(x, ϵ) = sum([xi * ϵ^(i - 1) for (i, xi) in enumerate(x)]) # define power series expansion of x
   get_power(x, ϵ, p) = p == 0 ? substitute(x, ϵ => 0) : get_power(expand_derivatives(Differential(ϵ)(x)) / p, ϵ, p - 1) # get O(ϵ^p) term of x

   get_power() is basically a recursion implementation of that the $O(\epsilon^p)$ term of $x$ is $(1/p!) \cdot \partial^p x / \partial^p \epsilon |_{\epsilon = 0}$ (i.e. get the factor in front of $\epsilon^p$ by differentiating away $p$ factors of $\epsilon$ and removing the $p!$ factor that "comes down with it"). I think it should be equivalent to collect_powers() and (IMO) maybe a little more symmetrical to the def_power definition, and you avoid the annoying max_power variable. You can then do eqs = get_power.(eq, ϵ, 0:order). Maybe this is less obvious, though.
4. I agree that 0-based indices would make this much nicer. I reported Structural simplification fails on non-1-indexed variable array  #2670 and tried to make a fix. I will see if I can "revive" that.

[karlwessel]

It would probably be good to use the @example ... feature from Documenter.jl to add the output of commands to the documentation. This way it would be easier to detect when some change breaks the example.

We are doing that over time. Most examples are already updated. We shouldn't merge this one until it's using @example.

I'll also update the second part of the example and can change the whole example to use the @example tag if you want.

[ChrisRackauckas]

If you've just got this working, mark the whole thing as @example so it won't regress.

[karlwessel]

3. To simplify collect_powers(), I have used this solution:

Very nice! But why the recursive implementation and not directly implement the mathematical equation you gave:

function get_power(poly, sym, p)
    x = sym
    ∂ₓ = Differential(x)
    p! = factorial(p)
    substitute(expand_derivatives((∂ₓ^p)(poly)) / p!, x => 0)
end

If there aren't any performance concerns, this probably could/should be used as implementation for an zeroth order respecting Symbolics.coeff.

[hersle]

Beautiful, that's even better! I didn't think about the Differential(ϵ)^p syntax.

If there aren't any performance concerns, this probably could/should be used as implementation for an zeroth order respecting Symbolics.coeff.

I think it primarily depends on what the function should do. Suppose you have a non-polynomial expression, like x = sin(ϵ*t). If I want its O(ϵ) perturbation, I probably want t (i.e. the derivative trick, effectively treating sin with its Taylor series expansion). But if I want the coefficient in front of ϵ of the same expression, I probably want 0 (i.e. what Symbolics.coeff currently gives).

Also, #2671 was merged, so 0-based arrays should work now (on the master branch), if you would like to try it.

[karlwessel]

Suppose you have a non-polynomial expression, like x = sin(ϵ*t). If I want its O(ϵ) perturbation, I probably want t (i.e. the derivative trick, effectively treating sin with its Taylor series expansion). But if I want the coefficient in front of ϵ of the same expression, I probably want 0 (i.e. what Symbolics.coeff currently gives).

You are right. But this also means get_power could be used in the algebraic perturbation example to automatically linearize sin(x) without having to expand it manually, right?

[karlwessel]

Also, #2671 was merged, so 0-based arrays should work now (on the master branch), if you would like to try it.

This still fails:

julia> using Symbolics

julia> def_taylor(x, ps) = sum([a * x^(i - 1) for (i, a) in enumerate(ps)])
def_taylor (generic function with 1 method)

julia> function collect_powers(eq, x, ns; max_power = 100)
           eq = substitute(expand(eq), Dict(x^j => 0 for j in (last(ns) + 1):max_power))

           eqs = []
           for i in ns
               powers = Dict(x^j => (i == j ? 1 : 0) for j in 1:last(ns))
               e = substitute(eq, powers)

               # manually remove zeroth order from higher orders
               if 0 in ns && i != 0
                   e = e - eqs[1]
               end

               push!(eqs, e)
           end
           eqs
       end
collect_powers (generic function with 1 method)

julia> function solve_coef(eqs, ps)
           vals = Dict()

           for (i, p) in enumerate(ps)
               eq = substitute(eqs[i], vals)
               vals[p] = Symbolics.symbolic_linear_solve(eq ~ 0, p)
           end
           vals
       end
solve_coef (generic function with 1 method)

julia> using ModelingToolkit: t_nounits as t, D_nounits as D

julia> order = 2
2

julia> @variables ϵ (y(t))[0:order] (∂∂y(t))[0:order]
3-element Vector{Any}:
 ϵ
  (y(t))[0:2]
  (∂∂y(t))[0:2]

julia> x = def_taylor(ϵ, y)
(y(t))[0] + (y(t))[1]*ϵ + (y(t))[2]*(ϵ^2)

julia> ∂∂x = def_taylor(ϵ, ∂∂y)
(∂∂y(t))[0] + ϵ*(∂∂y(t))[1] + (ϵ^2)*(∂∂y(t))[2]

julia> eq = ∂∂x * (1 + ϵ * x)^2 + 1
1 + ((1 + ((y(t))[0] + (y(t))[1]*ϵ + (y(t))[2]*(ϵ^2))*ϵ)^2)*((∂∂y(t))[0] + ϵ*(∂∂y(t))[1] + (ϵ^2)*(∂∂y(t))[2])

julia> eqs = collect_powers(eq, ϵ, 0:order)
3-element Vector{Any}:
                                                                           1 + (∂∂y(t))[0]
                                                      (∂∂y(t))[1] + 2(y(t))[0]*(∂∂y(t))[0]
 (∂∂y(t))[2] + 2(y(t))[0]*(∂∂y(t))[1] + 2(y(t))[1]*(∂∂y(t))[0] + ((y(t))[0]^2)*(∂∂y(t))[0]

julia> vals = solve_coef(eqs, ∂∂y)
Dict{Any, Any} with 3 entries:
  (∂∂y(t))[2] => 2.0(y(t))[1] - 3.0((y(t))[0]^2)
  (∂∂y(t))[1] => 2.0(y(t))[0]
  (∂∂y(t))[0] => -1.0

julia> subs = Dict(∂∂y[i] => D(D(y[i])) for i in eachindex(y))
Dict{Num, Num} with 3 entries:
  (∂∂y(t))[2] => Differential(t)(Differential(t)((y(t))[2]))
  (∂∂y(t))[1] => Differential(t)(Differential(t)((y(t))[1]))
  (∂∂y(t))[0] => Differential(t)(Differential(t)((y(t))[0]))

julia> eqs = [substitute(first(v), subs) ~ substitute(last(v), subs) for v in vals]
3-element Vector{Equation}:
 Differential(t)(Differential(t)((y(t))[2])) ~ 2.0(y(t))[1] - 3.0((y(t))[0]^2)
 Differential(t)(Differential(t)((y(t))[1])) ~ 2.0(y(t))[0]
 Differential(t)(Differential(t)((y(t))[0])) ~ -1.0

julia> using ModelingToolkit, DifferentialEquations

julia> @mtkbuild sys = ODESystem(eqs, t)
Model sys with 6 equations
Unknowns (6):
  var"y(t)ˍt"[2]
  var"y(t)ˍt"[1]
  var"y(t)ˍt"[0]
  (y(t))[2]
  (y(t))[1]
  (y(t))[0]
Parameters (0):

julia> unknowns(sys)
6-element Vector{SymbolicUtils.BasicSymbolic{Real}}:
 var"y(t)ˍt"[2]
 var"y(t)ˍt"[1]
 var"y(t)ˍt"[0]
 (y(t))[2]
 (y(t))[1]
 (y(t))[0]

julia> u0 = zeros(2order + 2)
6-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0

julia> u0[3] = 1.0
1.0

julia> prob = ODEProblem(sys, u0, (0, 3.0))
ERROR: BoundsError: attempt to access 3-element Vector{SymbolicUtils.BasicSymbolic{Real}} at index [0]
Stacktrace:
  [1] getindex(A::Vector{SymbolicUtils.BasicSymbolic{Real}}, i1::Int64)
    @ Base ./essentials.jl:13
  [2] fast_substitute(expr::SymbolicUtils.BasicSymbolic{Real}, subs::Dict{Any, Any}; operator::Type)
    @ Symbolics ~/.julia/packages/Symbolics/rtkf9/src/variable.jl:577
  [3] fast_substitute(eq::Equation, subs::Dict{Any, Any}; operator::Type)
    @ Symbolics ~/.julia/packages/Symbolics/rtkf9/src/variable.jl:543
  [4] tearing_reassemble(state::TearingState{…}, var_eq_matching::ModelingToolkit.BipartiteGraphs.Matching{…}, full_var_eq_matching::ModelingToolkit.BipartiteGraphs.Matching{…}; simplify::Bool, mm::ModelingToolkit.SparseMatrixCLIL{…})
    @ ModelingToolkit.StructuralTransformations ~/.julia/packages/ModelingToolkit/5OzIt/src/structural_transformation/symbolics_tearing.jl:594
  [5] tearing_reassemble
    @ ~/.julia/packages/ModelingToolkit/5OzIt/src/structural_transformation/symbolics_tearing.jl:230 [inlined]
  [6] tearing(sys::NonlinearSystem, state::TearingState{…}; mm::ModelingToolkit.SparseMatrixCLIL{…}, simplify::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit.StructuralTransformations ~/.julia/packages/ModelingToolkit/5OzIt/src/structural_transformation/symbolics_tearing.jl:641
  [7] _structural_simplify!(state::TearingState{…}, io::Nothing; simplify::Bool, check_consistency::Bool, fully_determined::Bool, warn_initialize_determined::Bool, dummy_derivative::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/systemstructure.jl:706
  [8] _structural_simplify!
    @ ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/systemstructure.jl:676 [inlined]
  [9] #structural_simplify!#1221
    @ ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/systemstructure.jl:669 [inlined]
 [10] __structural_simplify(sys::NonlinearSystem, io::Nothing; simplify::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/systems.jl:94
 [11] __structural_simplify
    @ ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/systems.jl:75 [inlined]
 [12] structural_simplify(sys::NonlinearSystem, io::Nothing; simplify::Bool, split::Bool, allow_symbolic::Bool, allow_parameter::Bool, conservative::Bool, fully_determined::Bool, kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/systems.jl:33
 [13] structural_simplify (repeats 2 times)
    @ ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/systems.jl:28 [inlined]
 [14] ModelingToolkit.InitializationProblem{…}(sys::ODESystem, t::Int64, u0map::Vector{…}, parammap::SciMLBase.NullParameters; guesses::Dict{…}, check_length::Bool, warn_initialize_determined::Bool, initialization_eqs::Vector{…}, fully_determined::Bool, check_units::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:1475
 [15] (ModelingToolkit.InitializationProblem{…})(::ODESystem, ::Int64, ::Vararg{…}; kwargs::@Kwargs{…})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:1434
 [16] ModelingToolkit.InitializationProblem(::ODESystem, ::Int64, ::Vararg{…}; kwargs::@Kwargs{…})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:1422
 [17] process_DEProblem(constructor::Type, sys::ODESystem, u0map::Vector{…}, parammap::SciMLBase.NullParameters; implicit_dae::Bool, du0map::Nothing, version::Nothing, tgrad::Bool, jac::Bool, checkbounds::Bool, sparse::Bool, simplify::Bool, linenumbers::Bool, parallel::Symbolics.SerialForm, eval_expression::Bool, eval_module::Module, use_union::Bool, tofloat::Bool, symbolic_u0::Bool, u0_constructor::typeof(identity), guesses::Dict{…}, t::Int64, warn_initialize_determined::Bool, build_initializeprob::Bool, initialization_eqs::Vector{…}, fully_determined::Bool, check_units::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:852
 [18] process_DEProblem
    @ ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:792 [inlined]
 [19] (ODEProblem{…})(sys::ODESystem, u0map::Vector{…}, tspan::Tuple{…}, parammap::SciMLBase.NullParameters; callback::Nothing, check_length::Bool, warn_initialize_determined::Bool, eval_expression::Bool, eval_module::Module, kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:1022
 [20] ODEProblem
    @ ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:1010 [inlined]
 [21] (ODEProblem{true, SciMLBase.AutoSpecialize})(sys::ODESystem, u0map::Vector{Float64}, tspan::Tuple{Int64, Float64})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:1010
 [22] (ODEProblem{true})(::ODESystem, ::Vector{Float64}, ::Vararg{Any}; kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:997
 [23] (ODEProblem{true})(::ODESystem, ::Vector{Float64}, ::Vararg{Any})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:996
 [24] ODEProblem(::ODESystem, ::Vector{Float64}, ::Vararg{Any}; kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:986
 [25] ODEProblem(::ODESystem, ::Vector{Float64}, ::Vararg{Any})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/5OzIt/src/systems/diffeqs/abstractodesystem.jl:985
 [26] top-level scope
    @ REPL[22]:1
Some type information was truncated. Use `show(err)` to see complete types.

[karlwessel]

Both examples work now (with one-based arrays only) and I added the outputs using the @example feature.

I am not completely happy with the definition of the initial condition with

u0 = zeros(2order + 2)
u0[3] = 1.0

The order of the equations defined by @mtkbuild might change and the correct index might change with it. But I don't know how to define it in a more robust way.

I think making the indices zero based and using a different implementation for collect_powers can be done later with a separate PR.

[hersle]

You are right. But this also means get_power could be used in the algebraic perturbation example to automatically linearize sin(x) without having to expand it manually, right?

I think so, yes! I will try to add it in JuliaSymbolics/Symbolics.jl#1298.

This still fails:

Ah shit. If I get time I will have a look.

I think making the indices zero based and using a different implementation for collect_powers can be done later with a separate PR.

Agreed, making the example correct is most important! :)

[hersle]

I am not completely happy with the definition of the initial condition with

I think you can do something like (with 1-based y[1:3])

u0 = Dict(unknowns(sys) .=> 0.0) # set all initial conditions
u0[D(y[1])] => 1.0 # ... except this one

[karlwessel]

I am not completely happy with the definition of the initial condition with

I think you can do something like (with 1-based y[1:3])

u0 = Dict(unknowns(sys) .=> 0.0) # set all initial conditions
u0[D(y[1])] => 1.0 # ... except this one

Thank you, that worked.

Now everything is good to merge from my side.