Precompilation (#94)

* add_harmonic type stable

* return types for functions

* SnoopPrecompile

* more type stability

* parametron file used for precompilation

* plotting precompiled

Project.toml

@@ -20,6 +20,7 @@ Peaks = "18e31ff7-3703-566c-8e60-38913d67486b"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
@@ -36,6 +37,7 @@ OrdinaryDiffEq = "v6.33.1"
 Peaks = "0.4.1"
 Plots = "1.36.4"
 ProgressMeter = "1.7.2"
+SnoopPrecompile = "1.0.1"
 Symbolics = "4.13.0"
 julia = "1.8.2"
 

src/DifferentialEquation.jl

@@ -22,7 +22,8 @@ Harmonic ansatz: x(t) => ω;
 ```
 """
 function add_harmonic!(diff_eom::DifferentialEquation, var::Num, ω)
-    diff_eom.harmonics[var] = unique(cat(diff_eom.harmonics[var], ω, dims=1))
+    push!(diff_eom.harmonics[var], ω)
+    diff_eom.harmonics[var] = unique(diff_eom.harmonics[var])
     diff_eom
 end
 
@@ -34,7 +35,7 @@ Return the dependent variables of `diff_eom`.
 get_variables(diff_eom::DifferentialEquation) = collect(keys(diff_eom.equations))
 
 
-is_harmonic(diff_eom::DifferentialEquation, t::Num) = all([is_harmonic(eq, t) for eq in values(diff_eom.equations)])
+is_harmonic(diff_eom::DifferentialEquation, t::Num)::Bool = all([is_harmonic(eq, t) for eq in values(diff_eom.equations)])
 
 "Pretty printing of the newly defined types"
 function show_fields(object)

src/HarmonicBalance.jl

@@ -11,6 +11,7 @@ module HarmonicBalance
     using ProgressMeter
     import Symbolics.SymbolicUtils: Term, Add, Div, Mul, Pow, Sym
     using DocStringExtensions
+    using SnoopPrecompile
 
     import Base: ComplexF64, Float64; export ComplexF64, Float64
     ComplexF64(x::Complex{Num}) = ComplexF64(Float64(x.re) + im*Float64(x.im))
@@ -61,6 +62,8 @@ module HarmonicBalance
     include("modules/LimitCycles.jl")
     using .LimitCycles
 
-
+    precomp_path = (@__DIR__) * "/../test/"
+    @precompile_all_calls include(precomp_path * "parametron.jl")
+    @precompile_all_calls include(precomp_path * "plotting.jl")
 
 end # module

src/HarmonicEquation.jl

@@ -24,18 +24,18 @@ function harmonic_ansatz(diff_eom::DifferentialEquation, time::Num)
     a_idx = 1
 
     for nvar in get_variables(diff_eom) # sum over natural variables
-        to_substitute = 0 # combine all the subtitution rules for var
+        to_substitute = Num(0) # combine all the subtitution rules for var
         for ω in diff_eom.harmonics[nvar]
             if !isequal(ω, 0) # nonzero harmonic - create u,v
                 rule_u, hvar_u = _create_harmonic_variable(nvar, ω, time, "u", new_symbol="u"*string(uv_idx))
                 rule_v, hvar_v = _create_harmonic_variable(nvar, ω, time, "v", new_symbol="v"*string(uv_idx))
                 rule = rule_u + rule_v
                 uv_idx += 1
-                append!(vars, [hvar_u, hvar_v])
+                push!(vars, hvar_u, hvar_v)
             else # zero harmonic - create a
                 rule, hvar = _create_harmonic_variable(nvar, ω, time, "a", new_symbol="a"*string(a_idx))
                 a_idx += 1
-                append!(vars, [hvar])
+                push!(vars, hvar)
             end
             to_substitute += rule
         end
@@ -150,7 +150,7 @@ end
 ###
 
 "Apply `rules` to both `equations` and `variables` field of `eom`"
-function substitute_all(eom::HarmonicEquation, rules::Union{Dict, Pair})
+function substitute_all(eom::HarmonicEquation, rules::Union{Dict, Pair})::HarmonicEquation
     new_eom = deepcopy(eom)
     new_eom.equations = expand_derivatives.(substitute_all(eom.equations, rules))
     new_eom
@@ -196,8 +196,8 @@ function fourier_transform!(eom::HarmonicEquation, time::Num)
     # loop over the HarmonicVariables, each generates one equation
     for (i, hvar) in enumerate(eom.variables)
         # find the equation belonging to this variable
-        eq_idx = findall(x -> isequal(x, hvar.natural_variable), collect(keys(eom.natural_equation.equations)))
-        eq = eom.equations[eq_idx...]
+        eq_idx = findfirst(x -> isequal(x, hvar.natural_variable), collect(keys(eom.natural_equation.equations)))
+        eq = eom.equations[eq_idx]
         # "type" is usually "u" or "v" (harmonic) or ["a"] (zero-harmonic)
         if hvar.type == "u"
             avg_eqs[i] = fourier_cos_term(eq, hvar.ω, time)

src/HarmonicVariable.jl

@@ -5,7 +5,7 @@ display(var::HarmonicVariable) = display(var.name)
 display(var::Vector{HarmonicVariable}) = display.(getfield.(var, Symbol("name")))
 
 
-function _coordinate_transform(new_var, ω, t, type)
+function _coordinate_transform(new_var, ω, t, type)::Num
     coords = Dict([
         "u" =>  new_var * cos(ω*t),
         "v" => new_var * sin(ω*t),
@@ -14,7 +14,7 @@ function _coordinate_transform(new_var, ω, t, type)
 end
 
 
-function _create_harmonic_variable(nat_var::Num, ω::Num, t::Num, type::String; new_symbol::String)
+function _create_harmonic_variable(nat_var::Num, ω::Num, t::Num, type::String; new_symbol::String)::Tuple{Num, HarmonicVariable}
     new_var = declare_variable(new_symbol, t) # this holds the internal symbol
     name = type * "_{" * var_name(nat_var) * "," * Base.replace(string(ω),"*"=>"") * "}"
     rule = _coordinate_transform(new_var, ω, t, type) # contribution of this harmonic variable to the natural variable
@@ -45,7 +45,7 @@ get_variables(vars::Vector{Num}) = unique(flatten([Num.(get_variables(x)) for x
 
 get_variables(var::HarmonicVariable) = Num.(get_variables(var.symbol))
 
-Base.isequal(v1::HarmonicVariable, v2::HarmonicVariable) = isequal(v1.symbol, v2.symbol)
+Base.isequal(v1::HarmonicVariable, v2::HarmonicVariable)::Bool = isequal(v1.symbol, v2.symbol)
 
 
 

src/Symbolics_customised.jl

@@ -95,7 +95,7 @@ exp_to_trig(x::Complex{Num}) = exp_to_trig(x.re)+im* exp_to_trig(x.im)
 
 # sometimes, expressions get stored as Complex{Num} with no way to decode what real(x) and imag(x)
 # this overloads the Num constructor to return a Num if x.re and x.im have similar arguments
-Num(x::Complex{Num}) = isequal(x.re.val.arguments, x.im.val.arguments) ? Num(first(x.re.val.arguments)) : error("Cannot convert Complex{Num} " * string(x) * " to Num")
+Num(x::Complex{Num})::Num = isequal(x.re.val.arguments, x.im.val.arguments) ? Num(first(x.re.val.arguments)) : error("Cannot convert Complex{Num} " * string(x) * " to Num")
 
 
 #=

src/Symbolics_utils.jl

@@ -44,15 +44,15 @@ d(funcs::Vector{Num}, x::Num, deg=1) = [d(f, x, deg) for f in funcs]
  Perform substitutions in `rules` on `x`.
  `include_derivatives=true` also includes all derivatives of the variables of the keys of `rules`.
  """
- function substitute_all(x::Union{Num, Equation}, rules::Dict; include_derivatives=true)
+ function substitute_all(x::T, rules::Dict; include_derivatives=true)::T where {T<:Union{Equation, Num}}
     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)
+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))
@@ -97,7 +97,7 @@ drop_powers(expr::Vector{Num}, var::Num, deg::Int) = [drop_powers(x, var, deg) f
 
 # calls the above for various types of the first argument
 drop_powers(eq::Equation, var, deg) = drop_powers(eq.lhs, var, deg) .~ drop_powers(eq.lhs, var, deg)
-drop_powers(eqs::Vector{Equation}, var, deg) = [drop_powers(eq.lhs, var, deg) .~ drop_powers(eq.rhs, var, deg) for eq in eqs]
+drop_powers(eqs::Vector{Equation}, var, deg) = [Equation(drop_powers(eq.lhs, var, deg), drop_powers(eq.rhs, var, deg)) for eq in eqs]
 drop_powers(expr, var::Num, deg::Int) = drop_powers(expr, [var], deg)
 drop_powers(x, vars, deg) = drop_powers(Num(x), vars, deg)
 
@@ -150,7 +150,7 @@ function _get_all_terms(x::Add)::Vector{Num}
 end
 
 
-function is_harmonic(x::Num, t::Num)
+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
@@ -160,7 +160,7 @@ function is_harmonic(x::Num, t::Num)
         return false
     else
         powers = [max_power(first(term.val.arguments), t) for term in t_terms[trigs]]
-        return unique(powers) == [1]
+        return all(isone, powers)
     end
 end
 
@@ -233,7 +233,7 @@ end
 
 
 function _fourier_term(x::Equation, ω, t, f)
-    _fourier_term(x.lhs, ω, t, f) .~ _fourier_term(x.rhs, ω, t, f)
+    Equation(_fourier_term(x.lhs, ω, t, f) , _fourier_term(x.rhs, ω, t, f))
 end
 
 "Return the coefficient of f(ωt) in `x` where `f` is a cos or sin."

test/parametron.jl

@@ -1,6 +1,6 @@
 using HarmonicBalance
 using Symbolics
-using Test
+#using Test # do not use Test as this file is used for precompilation
 
 @variables Ω,γ,λ,F, x,θ,η,α, ω0, ω,t,T, ψ
 @variables x(t)
@@ -15,7 +15,7 @@ p = HarmonicBalance.Problem(harmonic_eq);
 
 fixed = (Ω => 1.0,γ => 1E-2, λ => 5E-2, F => 1E-3,  α => 1.,  η=>0.3, θ => 0, ψ => 0)
 varied = ω => LinRange(0.9, 1.1, 100)
-res = get_steady_states(p, varied, fixed)
+res = get_steady_states(p, varied, fixed, show_progress=false)
 
 # save the result, try and load in the next step
 #current_path = @__DIR__
@@ -24,7 +24,7 @@ res = get_steady_states(p, varied, fixed)
 # try to run a 2D calculation
 fixed = (Ω => 1.0,γ => 1E-2, F => 1E-3,  α => 1.,  η=>0.3, θ => 0, ψ => 0)
 varied = (ω => LinRange(0.9, 1.1, 10), λ => LinRange(0.01, 0.05, 10))
-res = get_steady_states(p, varied, fixed)
+res = get_steady_states(p, varied, fixed, show_progress=false)
 
 
 ###
@@ -35,5 +35,5 @@ ref1 = (Ω^2)*u1 + F*cos(θ) + γ*Differential(T)(u1) + (3//4)*α*(u1^3) + γ*ω
 ref2 = γ*Differential(T)(v1) + (Ω^2)*v1 + (3//4)*α*(v1^3) + (3//4)*α*(u1^2)*v1 + (1//4)*η*(u1^2)*Differential(T)(v1) + (3//4)*η*(v1^2)*Differential(T)(v1) + (1//2)*λ*(Ω^2)*v1*cos(ψ) + (1//2)*η*u1*v1*Differential(T)(u1) + (1//2)*λ*(Ω^2)*u1*sin(ψ) - F*sin(θ) - (ω^2)*v1 - (2//1)*ω*Differential(T)(u1) - γ*ω*u1 - (1//4)*η*ω*(u1^3) - (1//4)*η*ω*(v1^2)*u1
 
 averaged = HarmonicBalance._remove_brackets(harmonic_eq)
-@test isequal(simplify(expand(averaged[1]-ref1)), 0)
-@test isequal(simplify(expand(averaged[2]-ref2)), 0)
+@assert isequal(simplify(expand(averaged[1]-ref1)), 0)
+@assert isequal(simplify(expand(averaged[2]-ref2)), 0)

test/plotting.jl

@@ -1,6 +1,6 @@
 using HarmonicBalance, Plots
 default(show=false)
-using Test
+#using Test #
 
 @variables γ, λ, x, η, α, ω0, ω
 @variables t x(t)