minimal symbolic clean up

ext/SteadyStateDiffEqExt.jl

@@ -5,7 +5,7 @@ using HarmonicBalance: HarmonicBalance, steady_state_sweep
 
 using SteadyStateDiffEq: solve, NonlinearProblem, SteadyStateProblem, DynamicSS, remake
 using LinearAlgebra: norm, eigvals
-using SteadyStateDiffEq.SciMLBase.SciMLStructures: isscimlstructure, Tunable, replace
+using SteadyStateDiffEq.SciMLBase.SciMLStructures: Tunable, replace
 
 function HarmonicBalance.steady_state_sweep(
     prob::SteadyStateProblem, alg::DynamicSS; varied::Pair, kwargs...

ext/TimeEvolution/TimeEvolution.jl

@@ -28,10 +28,9 @@ include("ODEProblem.jl")
 include("hysteresis_sweep.jl")
 include("plotting.jl")
 
-export FFT
 export ParameterSweep
-export transform_solutions, plot, plot!, is_stable
+export transform_solutions, plot, plot!
 export ODEProblem, solve
-export plot_1D_solutions_branch, follow_branch
+export follow_branch
 
 end

src/HarmonicBalance.jl

@@ -24,8 +24,44 @@ function set_imaginary_tolerance(x::Float64)
     @eval(IM_TOL::Float64 = $x)
 end
 
-include("Symbolics_customised.jl")
-include("Symbolics_utils.jl")
+using SymbolicUtils:
+    SymbolicUtils,
+    Postwalk,
+    Sym,
+    BasicSymbolic,
+    isterm,
+    ispow,
+    isadd,
+    isdiv,
+    ismul,
+    add_with_div,
+    frac_maketerm,
+    @compactified,
+    issym
+
+using Symbolics:
+    Symbolics,
+    Num,
+    unwrap,
+    wrap,
+    get_variables,
+    simplify,
+    expand_derivatives,
+    build_function,
+    Equation,
+    Differential,
+    @variables,
+    arguments,
+    simplify_fractions,
+    substitute,
+    term,
+    expand,
+    operation
+
+include("Symbolics/Symbolics_utils.jl")
+include("Symbolics/exponentials.jl")
+include("Symbolics/fourier.jl")
+include("Symbolics/drop_powers.jl")
 
 include("modules/extention_functions.jl")
 include("utils.jl")
@@ -74,14 +110,14 @@ export get_krylov_equations
 include("modules/FFTWExt.jl")
 using .FFTWExt
 
-@setup_workload begin
-    # Putting some things in `@setup_workload` instead of `@compile_workload` can reduce the size of the
-    # precompile file and potentially make loading faster.
-    @compile_workload begin
-        # all calls in this block will be precompiled, regardless of whether
-        # they belong to your package or not (on Julia 1.8 and higher)
-        include("precompilation.jl")
-    end
-end
+# @setup_workload begin
+#     # Putting some things in `@setup_workload` instead of `@compile_workload` can reduce the size of the
+#     # precompile file and potentially make loading faster.
+#     @compile_workload begin
+#         # all calls in this block will be precompiled, regardless of whether
+#         # they belong to your package or not (on Julia 1.8 and higher)
+#         include("precompilation.jl")
+#     end
+# end
 
 end # module

src/HarmonicEquation.jl

@@ -63,7 +63,8 @@ function slow_flow!(eom::HarmonicEquation; fast_time::Num, slow_time::Num, degre
     drop = [d(var, fast_time, degree) => 0 for var in get_variables(eom)]
 
     eom.equations = substitute_all(substitute_all(eom.equations, drop), replace)
-    return eom.variables = substitute_all(eom.variables, replace)
+    eom.variables = substitute_all(eom.variables, replace)
+    return nothing
 end
 
 """

src/HarmonicVariable.jl

@@ -51,3 +51,33 @@ Symbolics.get_variables(var::HarmonicVariable)::Num = Num(first(get_variables(va
 
 Base.isequal(v1::HarmonicVariable, v2::HarmonicVariable)::Bool =
     isequal(v1.symbol, v2.symbol)
+
+"The derivative of f w.r.t. x of degree deg"
+function d(f::Num, x::Num, deg=1)::Num
+    return isequal(deg, 0) ? f : (Differential(x)^deg)(f)
+end
+d(funcs::Vector{Num}, x::Num, deg=1) = Num[d(f, x, deg) for f in funcs]
+
+"Declare a variable in the the currently active namespace"
+function declare_variable(name::String)
+    var_sym = Symbol(name)
+    @eval($(var_sym) = first(Symbolics.@variables $var_sym))
+    return eval(var_sym)
+end
+
+"Declare a variable that is a function of another variable in the the current namespace"
+function declare_variable(name::String, independent_variable::Num)
+    # independent_variable = declare_variable(independent_variable) convert string into Num
+    var_sym = Symbol(name)
+    new_var = Symbolics.@variables $var_sym(independent_variable)
+    @eval($(var_sym) = first($new_var)) # store the variable under "name" in this namespace
+    return eval(var_sym)
+end
+
+"Return the name of a variable (excluding independent variables)"
+function var_name(x::Num)
+    var = Symbolics._toexpr(x)
+    return var isa Expr ? String(var.args[1]) : String(var)
+end
+#  var_name(x::Term) = String(Symbolics._toexpr(x).args[1])
+var_name(x::Sym) = String(x.name)

src/Symbolics/Symbolics_utils.jl

@@ -0,0 +1,128 @@
+
+expand_all(x::Num) = Num(expand_all(x.val))
+_apply_termwise(f, x::Num) = wrap(_apply_termwise(f, unwrap(x)))
+
+"Expands using SymbolicUtils.expand and expand_exp_power (changes exp(x)^n to exp(x*n)"
+function expand_all(x)
+    result = Postwalk(expand_exp_power)(SymbolicUtils.expand(x))
+    return isnothing(result) ? x : result
+end
+expand_all(x::Complex{Num}) = expand_all(x.re) + im * expand_all(x.im)
+
+"Apply a function f on every member of a sum or a product"
+function _apply_termwise(f, x::BasicSymbolic)
+    @compactified x::BasicSymbolic begin
+        Add => sum([f(arg) for arg in arguments(x)])
+        Mul => prod([f(arg) for arg in arguments(x)])
+        Div => _apply_termwise(f, x.num) / _apply_termwise(f, x.den)
+        _   => f(x)
+    end
+end
+
+simplify_complex(x::Complex) = isequal(x.im, 0) ? x.re : x.re + im * x.im
+simplify_complex(x) = x
+function simplify_complex(x::BasicSymbolic)
+    @compactified x::BasicSymbolic begin
+        Add => _apply_termwise(simplify_complex, x)
+        Mul => _apply_termwise(simplify_complex, x)
+        Div => _apply_termwise(simplify_complex, x)
+        _   => x
+    end
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Perform substitutions in `rules` on `x`.
+`include_derivatives=true` also includes all derivatives of the variables of the keys of `rules`.
+"""
+subtype = Union{Num,Equation,BasicSymbolic}
+function substitute_all(x::subtype, rules::Dict; include_derivatives=true)
+    if include_derivatives
+        rules = merge(
+            rules,
+            Dict([Differential(var) => Differential(rules[var]) for var in keys(rules)]),
+        )
+    end
+    return substitute(x, rules)
+end
+"Variable substitution - dictionary"
+function substitute_all(dict::Dict, rules::Dict)::Dict
+    new_keys = substitute_all.(keys(dict), rules)
+    new_values = substitute_all.(values(dict), rules)
+    return Dict(zip(new_keys, new_values))
+end
+Collections = Union{Dict,Pair,Vector,OrderedDict}
+substitute_all(v::AbstractArray, rules) = [substitute_all(x, rules) for x in v]
+substitute_all(x::subtype, rules::Collections) = substitute_all(x, Dict(rules))
+# Collections = Union{Dict,OrderedDict}
+# function substitute_all(x, rules::Collections; include_derivatives=true)
+#     if include_derivatives
+#         rules = merge(
+#             rules,
+#             Dict([Differential(var) => Differential(rules[var]) for var in keys(rules)]),
+#         )
+#     end
+#     return substitute(x, rules)
+# end
+# "Variable substitution - dictionary"
+# function substitute_all(dict::Dict, rules::Dict)::Dict
+#     new_keys = substitute_all.(keys(dict), rules)
+#     new_values = substitute_all.(values(dict), rules)
+#     return Dict(zip(new_keys, new_values))
+# end
+# substitute_all(v::AbstractArray, rules::Collections) = [substitute_all(x, rules) for x in v]
+
+get_independent(x::Num, t::Num) = get_independent(x.val, t)
+function get_independent(x::Complex{Num}, t::Num)
+    return get_independent(x.re, t) + im * get_independent(x.im, t)
+end
+get_independent(v::Vector{Num}, t::Num) = [get_independent(el, t) for el in v]
+get_independent(x, t::Num) = x
+
+function get_independent(x::BasicSymbolic, t::Num)
+    @compactified x::BasicSymbolic begin
+        Add  => sum([get_independent(arg, t) for arg in arguments(x)])
+        Mul  => prod([get_independent(arg, t) for arg in arguments(x)])
+        Div  => !is_function(x.den, t) ? get_independent(x.num, t) / x.den : 0
+        Pow  => !is_function(x.base, t) && !is_function(x.exp, t) ? x : 0
+        Term => !is_function(x, t) ? x : 0
+        Sym  => !is_function(x, t) ? x : 0
+        _    => x
+    end
+end
+
+"Return all the terms contained in `x`"
+get_all_terms(x::Num) = unique(_get_all_terms(Symbolics.expand(x).val))
+function get_all_terms(x::Equation)
+    return unique(cat(get_all_terms(Num(x.lhs)), get_all_terms(Num(x.rhs)); dims=1))
+end
+function _get_all_terms(x::BasicSymbolic)
+    @compactified x::BasicSymbolic begin
+        Add => vcat([_get_all_terms(term) for term in SymbolicUtils.arguments(x)]...)
+        Mul => Num.(SymbolicUtils.arguments(x))
+        Div => Num.([_get_all_terms(x.num)..., _get_all_terms(x.den)...])
+        _   => Num(x)
+    end
+end
+_get_all_terms(x) = Num(x)
+
+function is_harmonic(x::Num, t::Num)::Bool
+    all_terms = get_all_terms(x)
+    t_terms = setdiff(all_terms, get_independent(all_terms, t))
+    isempty(t_terms) && return true
+    trigs = is_trig.(t_terms)
+
+    if !prod(trigs)
+        return false
+    else
+        powers = [max_power(first(term.val.arguments), t) for term in t_terms[trigs]]
+        return all(isone, powers)
+    end
+end
+
+is_harmonic(x::Equation, t::Num) = is_harmonic(x.lhs, t) && is_harmonic(x.rhs, t)
+is_harmonic(x, t) = is_harmonic(Num(x), Num(t))
+
+"Return true if `f` is a function of `var`."
+is_function(f, var) = any(isequal.(get_variables(f), var))

src/Symbolics/drop_powers.jl

@@ -0,0 +1,70 @@
+"""
+$(SIGNATURES)
+Remove parts of `expr` where the combined power of `vars` is => `deg`.
+
+# Example
+```julia-repl
+julia> @variables x,y;
+julia>drop_powers((x+y)^2, x, 2)
+y^2 + 2*x*y
+julia>drop_powers((x+y)^2, [x,y], 2)
+0
+julia>drop_powers((x+y)^2 + (x+y)^3, [x,y], 3)
+x^2 + y^2 + 2*x*y
+```
+"""
+function drop_powers(expr::Num, vars::Vector{Num}, deg::Int)
+    Symbolics.@variables ϵ
+    subs_expr = deepcopy(expr)
+    rules = Dict([var => ϵ * var for var in unique(vars)])
+    subs_expr = Symbolics.expand(substitute_all(subs_expr, rules))
+    max_deg = max_power(subs_expr, ϵ)
+    removal = Dict([ϵ^d => Num(0) for d in deg:max_deg])
+    res = substitute_all(substitute_all(subs_expr, removal), Dict(ϵ => Num(1)))
+    return Symbolics.expand(res)
+end
+
+function drop_powers(expr::Vector{Num}, var::Vector{Num}, deg::Int)
+    return [drop_powers(x, var, deg) for x in expr]
+end
+
+# calls the above for various types of the first argument
+function drop_powers(eq::Equation, var::Vector{Num}, deg::Int)
+    return drop_powers(eq.lhs, var, deg) .~ drop_powers(eq.lhs, var, deg)
+end
+function drop_powers(eqs::Vector{Equation}, var::Vector{Num}, deg::Int)
+    return [
+        Equation(drop_powers(eq.lhs, var, deg), drop_powers(eq.rhs, var, deg)) for eq in eqs
+    ]
+end
+drop_powers(expr, var::Num, deg::Int) = drop_powers(expr, [var], deg)
+drop_powers(x, vars, deg::Int) = drop_powers(Num(x), vars, deg)
+
+"Return the highest power of `y` occuring in the term `x`."
+function max_power(x::Num, y::Num)
+    terms = get_all_terms(x)
+    powers = power_of.(terms, y)
+    return maximum(powers)
+end
+
+max_power(x::Vector{Num}, y::Num) = maximum(max_power.(x, y))
+max_power(x::Complex, y::Num) = maximum(max_power.([x.re, x.im], y))
+max_power(x, t) = max_power(Num(x), Num(t))
+
+"Return the power of `y` in the term `x`"
+function power_of(x::Num, y::Num)
+    issym(y.val) ? nothing : error("power of " * string(y) * " is ambiguous")
+    return power_of(x.val, y.val)
+end
+
+function power_of(x::BasicSymbolic, y::BasicSymbolic)
+    if ispow(x) && issym(y)
+        return isequal(x.base, y) ? x.exp : 0
+    elseif issym(x) && issym(y)
+        return isequal(x, y) ? 1 : 0
+    else
+        return 0
+    end
+end
+
+power_of(x, y) = 0

src/Symbolics/exponentials.jl

@@ -0,0 +1,41 @@
+expand_exp_power(expr::Num) = expand_exp_power(expr.val)
+simplify_exp_products(x::Num) = simplify_exp_products(x.val)
+
+"Returns true if expr is an exponential"
+isexp(expr) = isterm(expr) && expr.f == exp
+
+"Expand powers of exponential such that exp(x)^n => exp(x*n) "
+function expand_exp_power(expr::BasicSymbolic)
+    @compactified expr::BasicSymbolic begin
+        Add => sum([expand_exp_power(arg) for arg in arguments(expr)])
+        Mul => prod([expand_exp_power(arg) for arg in arguments(expr)])
+        _   => ispow(expr) && isexp(expr.base) ? exp(expr.base.arguments[1] * expr.exp) : expr
+    end
+end
+expand_exp_power(expr) = expr
+
+"Simplify products of exponentials such that exp(a)*exp(b) => exp(a+b)
+This is included in SymbolicUtils as of 17.0 but the method here avoid other simplify calls"
+function simplify_exp_products(expr::BasicSymbolic)
+    @compactified expr::BasicSymbolic begin
+        Add => _apply_termwise(simplify_exp_products, expr)
+        Div => _apply_termwise(simplify_exp_products, expr)
+        Mul => simplify_exp_products_mul(expr)
+        _   => expr
+    end
+end
+function simplify_exp_products(x::Complex{Num})
+    return Complex{Num}(simplify_exp_products(x.re.val), simplify_exp_products(x.im.val))
+end
+function simplify_exp_products_mul(expr)
+    ind = findall(x -> isexp(x), arguments(expr))
+    rest_ind = setdiff(1:length(arguments(expr)), ind)
+    rest = isempty(rest_ind) ? 1 : prod(arguments(expr)[rest_ind])
+    total = isempty(ind) ? 0 : sum(getindex.(arguments.(arguments(expr)[ind]), 1))
+    if SymbolicUtils.is_literal_number(total)
+        (total == 0 && return rest)
+    else
+        return rest * exp(total)
+    end
+end
+simplify_exp_products(x) = x