enforce type inferribility and stability with Jet

Project.toml

@@ -34,7 +34,6 @@ SteadyStateDiffEqExt = "SteadyStateDiffEq"
 TimeEvolution = "OrdinaryDiffEq"
 
 [compat]
-JET = "0.9"
 Aqua = "0.8"
 BijectiveHilbert = "0.3.0"
 DSP = "0.7.4"
@@ -44,13 +43,14 @@ DocStringExtensions = "0.9"
 ExplicitImports = "1.6"
 FFTW = "1.5.0"
 HomotopyContinuation = "2.6.4"
+JET = "0.9"
 JLD2 = "0.4.24"
 Latexify = "0.15.16, 0.16"
 ModelingToolkit = "9"
 NonlinearSolve = "3.12"
 OrderedCollections = "1.4.1"
 OrdinaryDiffEq = "v6.33.1"
-Peaks = "0.4"
+Peaks = "0.5"
 Plots = "1.35.0"
 ProgressMeter = "1.7.2"
 SteadyStateDiffEq = "1, 2"
@@ -59,9 +59,10 @@ Symbolics = "5.0.0"
 julia = "1.10.0"
 
 [extras]
-JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
@@ -70,4 +71,4 @@ SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Pkg", "Test", "OrdinaryDiffEq", "ModelingToolkit", "SteadyStateDiffEq", "NonlinearSolve", "ExplicitImports", "Aqua"]
+test = ["Pkg", "Test", "OrdinaryDiffEq", "ModelingToolkit", "SteadyStateDiffEq", "NonlinearSolve", "ExplicitImports", "Aqua", "JET", "Documenter"]

src/DifferentialEquation.jl

@@ -26,7 +26,7 @@ end
 $(TYPEDSIGNATURES)
 Return the dependent variables of `diff_eom`.
 """
-Symbolics.get_variables(diff_eom::DifferentialEquation) = collect(keys(diff_eom.equations))
+Symbolics.get_variables(diff_eom::DifferentialEquation)::Vector{Num} = collect(keys(diff_eom.equations))
 
 is_harmonic(diff_eom::DifferentialEquation, t::Num)::Bool =
     all([is_harmonic(eq, t) for eq in values(diff_eom.equations)])

src/HarmonicEquation.jl

@@ -133,12 +133,12 @@ end
 $(TYPEDSIGNATURES)
 Get the internal symbols of the independent variables of `eom`.
 """
-function Symbolics.get_variables(eom::HarmonicEquation)
-    return flatten(get_variables.(eom.variables))
+function Symbolics.get_variables(eom::HarmonicEquation)::Vector{Num}
+    return get_variables.(eom.variables)
 end
 
-Symbolics.get_variables(p::Problem) = get_variables(p.eom)
-Symbolics.get_variables(res::Result) = get_variables(res.problem)
+Symbolics.get_variables(p::Problem)::Vector{Num} = get_variables(p.eom)
+Symbolics.get_variables(res::Result)::Vector{Num} = get_variables(res.problem)
 
 "Get the parameters (not time nor variables) of a HarmonicEquation"
 function _parameters(eom::HarmonicEquation)

src/HarmonicVariable.jl

@@ -47,7 +47,7 @@ end
 get_variables_nums(vars::Vector{Num}) =
     unique(flatten([Num.(get_variables(x)) for x in vars]))
 
-Symbolics.get_variables(var::HarmonicVariable) = Num.(get_variables(var.symbol))
+Symbolics.get_variables(var::HarmonicVariable)::Num = Num(first(get_variables(var.symbol)))
 
 Base.isequal(v1::HarmonicVariable, v2::HarmonicVariable)::Bool =
     isequal(v1.symbol, v2.symbol)

src/Symbolics_utils.jl

@@ -1,8 +1,8 @@
 "The derivative of f w.r.t. x of degree deg"
-function d(f::Num, x::Num, deg=1)
+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) = [d(f, x, deg) for f in funcs]
+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)
@@ -64,7 +64,7 @@ end
 function substitute_all(x::Complex{Num}, rules::Union{Pair,Vector,OrderedDict,Dict})
     return substitute_all(Num(x.re.val.arguments[1]), rules)
 end
-substitute_all(x, rules) = substitute_all(Num(x), rules::Dict)
+substitute_all(x, rules) = substitute_all(Num(x), rules)
 
 """
 $(SIGNATURES)

src/modules/FFTWExt.jl

@@ -28,31 +28,31 @@ end
 #     return HarmonicBalance.FFT(soln.u, soln.t; window=window)
 # end
 
-function FFT_analyze(fft_u::Vector{ComplexF64}, fft_f)
-    "finds peaks in the spectrum and returns corresponding frequency, amplitude and phase.
-    Frequency and phase are corrected according to Huang Dishan, Mechanical Systems and Signal Processing (1995) 9(2), 113–118
-    This correction works for a rectangular window."
+# function FFT_analyze(fft_u::Vector{ComplexF64}, fft_f)
+#     "finds peaks in the spectrum and returns corresponding frequency, amplitude and phase.
+#     Frequency and phase are corrected according to Huang Dishan, Mechanical Systems and Signal Processing (1995) 9(2), 113–118
+#     This correction works for a rectangular window."
 
-    # retaining more sigdigits gives more ''spurious'' peaks
-    max_indices, mxval = Peaks.peakproms(round.(abs.(fft_u), sigdigits=3); minprom=1)
-    Del = fft_f[2] - fft_f[1] # frequency spacing
-    A1 = abs.(fft_u)[max_indices]
-    df = zeros(length(max_indices))
+#     # retaining more sigdigits gives more ''spurious'' peaks
+#     max_indices, mxval = Peaks.peakproms(round.(abs.(fft_u), sigdigits=3); minprom=1)
+#     Del = fft_f[2] - fft_f[1] # frequency spacing
+#     A1 = abs.(fft_u)[max_indices]
+#     df = zeros(length(max_indices))
 
-    # correction to the amplitude and phase of the peak
-    for i in 1:length(max_indices)
-        if abs.(fft_u)[max_indices[i] - 1] < abs.(fft_u)[max_indices[i] + 1]
-            A2 = abs.(fft_u)[max_indices[i] + 1]
-            df[i] = -Del / (A1[i] / A2 + 1)
-        else
-            A2 = abs.(fft_u)[max_indices[i] - 1]
-            df[i] = Del / (A1[i] / A2 + 1)
-        end
-    end
-    return 2 * pi * (fft_f[max_indices] - df),
-    A1 .* 2,
-    angle.(fft_u)[max_indices] + pi * df / Del
-end
+#     # correction to the amplitude and phase of the peak
+#     for i in 1:length(max_indices)
+#         if abs.(fft_u)[max_indices[i] - 1] < abs.(fft_u)[max_indices[i] + 1]
+#             A2 = abs.(fft_u)[max_indices[i] + 1]
+#             df[i] = -Del / (A1[i] / A2 + 1)
+#         else
+#             A2 = abs.(fft_u)[max_indices[i] - 1]
+#             df[i] = Del / (A1[i] / A2 + 1)
+#         end
+#     end
+#     return 2 * pi * (fft_f[max_indices] - df),
+#     A1 .* 2,
+#     angle.(fft_u)[max_indices] + pi * df / Del
+# end
 
 function u_of_t(omegas_peak, As_peak, phis_peak, t)
     "Calculate us or vs as a function of time from the Fourier components."

src/modules/HC_wrapper.jl

@@ -3,6 +3,7 @@ module HC_wrapper
 using DocStringExtensions
 using Symbolics: Num, @variables
 using Symbolics.SymbolicUtils: isterm
+using LinearAlgebra: LinearAlgebra
 
 using HarmonicBalance:
     HarmonicBalance,

src/modules/HC_wrapper/homotopy_interface.jl

@@ -37,7 +37,7 @@ function HarmonicBalance.Problem(eom::HarmonicEquation; Jacobian=true)
     if Jacobian == true || Jacobian == "explicit"
         J = HarmonicBalance.get_Jacobian(eom)
     elseif Jacobian == "false" || Jacobian == false
-        dummy_J(arg) = I(1)
+        dummy_J(arg) = LinearAlgebra.I(1)
         J = dummy_J
     else
         J = Jacobian
@@ -59,9 +59,9 @@ function HarmonicBalance.Problem(
         symbol in [Meta.parse(s) for s in [string(eq) for eq in equations]]
     ] #note in polar coordinates there could be imaginary factors, requiring the extra replacement "I"=>"1im"
     system = HomotopyContinuation.System(eqs_HC; variables=conv_vars, parameters=conv_para)
-    J = get_Jacobian(equations, variables) #all derivatives are assumed to be in the left hand side;
+    J = HarmonicBalance.get_Jacobian(equations, variables) #all derivatives are assumed to be in the left hand side;
     return Problem(variables, parameters, system, J)
-end
+end # is this funciton still needed/used?
 
 function System(eom::HarmonicEquation)
     eqs = expand_derivatives.(_remove_brackets(eom))

src/modules/LimitCycles/gauge_fixing.jl

@@ -23,25 +23,29 @@ function get_cycle_variables(eom::HarmonicEquation, ω_lc::Num)
     return vars = filter(x -> any(isequal.(ω_lc, get_all_terms(x.ω))), eom.variables)
 end
 
-"""
-    Obtain the Jacobian of `eom` with a gauge-fixed variable `fixed_var`.
-    `fixed_var` marks the variable which is fixed to zero due to U(1) symmetry.
-    This leaves behind a finite degeneracy of solutions (see Chapter 6 of Jan's thesis).
-
-    For limit cycles, we always use an 'implicit' Jacobian - a function which only returns the numerical Jacobian when a numerical solution
-    is inserted. Finding the analytical Jacobian is usually extremely time-consuming.
-"""
+# """
+#     Obtain the Jacobian of `eom` with a gauge-fixed variable `fixed_var`.
+#     `fixed_var` marks the variable which is fixed to zero due to U(1) symmetry.
+#     This leaves behind a finite degeneracy of solutions (see Chapter 6 of Jan's thesis).
+
+#     For limit cycles, we always use an 'implicit' Jacobian - a function which only returns the numerical Jacobian when a numerical solution
+#     is inserted. Finding the analytical Jacobian is usually extremely time-consuming.
+# """
 function _gaugefixed_Jacobian(
     eom::HarmonicEquation, fixed_var::HarmonicVariable; explicit=false, sym_order, rules
 )
     if explicit
-        _fix_gauge!(get_Jacobian(eom), fixed_var) # get a symbolic explicit J, compile later
+        error("Explicit Jacobian not implemented for limit cycles yet!")
+        # _fix_gauge!(get_Jacobian(eom), fixed_var) # get a symbolic explicit J, compile later
     else
         rules = Dict(rules)
         setindex!(rules, 0, _remove_brackets(fixed_var))
-        get_implicit_Jacobian(eom; rules=rules, sym_order=sym_order)
+        jac = get_implicit_Jacobian(eom; rules=rules, sym_order=sym_order)
     end
+    return jac
 end
+# ^ The above function is not finished?
+# no matching method found `_fix_gauge!(::Matrix{Symbolics.Num}, ::HarmonicBalance.HarmonicVariable)`
 
 """
 $(TYPEDSIGNATURES)

src/modules/LinearResponse.jl

@@ -9,11 +9,12 @@ using Latexify: Latexify, latexify, @L_str
 using Latexify.LaTeXStrings: LaTeXStrings
 
 using Symbolics: Num, build_function, Equation, substitute, unwrap
-using LinearAlgebra: norm, eigvals, eigen
+using LinearAlgebra: norm, eigvals, eigen, eigvecs
 using OrderedCollections: OrderedDict
 
 using HarmonicBalance
 using HarmonicBalance:
+    var_name,
     rearrange_standard,
     _remove_brackets,
     expand_derivatives,

src/modules/LinearResponse/Lorentzian_spectrum.jl

@@ -40,14 +40,10 @@ function add_peak(s1::JacobianSpectrum, s2::JacobianSpectrum)
 end
 
 "Gives the numerical value of a peak at ω."
-evaluate(peak::Lorentzian, ω::Float64) = peak.A / sqrt(((peak.ω0 - ω)^2 + (peak.Γ)^2))
+evaluate(peak::Lorentzian, ω::Float64)::Float64 = peak.A / sqrt(((peak.ω0 - ω)^2 + (peak.Γ)^2))
 
 "Gives the numerical value of a JacobianSpectrum at ω"
-evaluate(s::JacobianSpectrum, ω::Float64) = sum([evaluate(p, ω) for p in s.peaks])
-
-function evaluate(spectra::Dict{Num,JacobianSpectrum}, ω::Float64)
-    return Dict([[var, evaluate(spectra[var], ω)] for var in keys(spectra)])
-end
+evaluate(s::JacobianSpectrum, ω::Float64)::Float64 = sum([evaluate(p, ω) for p in s.peaks])
 
 "Take a pair of harmonic variable u,v and an eigenvalue λ and eigenvector eigvec_2d of the Jacobian to generate corresponding Lorentzians.
     IMPORTANT: The eigenvetor eigen_2d contains only the two components of the full eigenvector which correspond to the u,v pair in question."

src/modules/LinearResponse/plotting.jl

@@ -83,7 +83,7 @@ function get_linear_response(
 end
 
 function get_rotframe_jacobian_response(
-    res::Result, Ω_range, branch::Int; show_progress=show_progress, damping_mod::Float64
+    res::Result, Ω_range, branch::Int; show_progress=true, damping_mod::Float64
 )
     stable = classify_branch(res, branch, "stable")
     !any(stable) && error("Cannot generate a spectrum - no stable solutions!")

test/HarmonicVariable.jl

@@ -1,62 +1,74 @@
 using Test
+using HarmonicBalance
+@variables α ω ω0 F γ η t x(t); # declare constant variables and a function x(t)
 
-@testset "HarmonicVariable Tests" begin
-    # Test display function
-    @testset "Display" begin
-        var = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
-        @test display(var) == :x
-        vars = [
-            HarmonicVariable(:x, "x", "u", 1.0, 2.0),
-            HarmonicVariable(:y, "y", "v", 2.0, 3.0),
-        ]
-        @test display(vars) == [:x, :y]
-    end
-
-    # Test _coordinate_transform function
-    @testset "_coordinate_transform" begin
-        @test _coordinate_transform(2.0, 1.0, 0.0, "u") == 2.0
-        @test _coordinate_transform(2.0, 1.0, 0.0, "v") == 0.0
-        @test _coordinate_transform(2.0, 1.0, 0.0, "a") == 2.0
-    end
-
-    # Test _create_harmonic_variable function
-    @testset "_create_harmonic_variable" begin
-        rule, hvar = _create_harmonic_variable(2.0, 1.0, 0.0, "u"; new_symbol="x")
-        @test rule == 2.0
-        @test hvar == HarmonicVariable(:x, "u_{2.0}", "u", 1.0, 2.0)
-    end
-
-    # Test substitute_all function
-    @testset "substitute_all" begin
-        eq = 2.0 * :x + 3.0 * :y
-        rules = Dict(:x => 1.0, :y => 2.0)
-        @test substitute_all(eq, rules) == 2.0 + 6.0
-        var = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
-        @test substitute_all(var, rules) == HarmonicVariable(1.0, "x", "u", 1.0, 2.0)
-        vars = [
-            HarmonicVariable(:x, "x", "u", 1.0, 2.0),
-            HarmonicVariable(:y, "y", "v", 2.0, 3.0),
-        ]
-        @test substitute_all(vars, rules) == [
-            HarmonicVariable(1.0, "x", "u", 1.0, 2.0),
-            HarmonicVariable(2.0, "y", "v", 2.0, 3.0),
-        ]
-    end
-
-    # Test Symbolics.get_variables function
-    @testset "Symbolics.get_variables" begin
-        vars = [1.0, 2.0, 3.0]
-        @test Symbolics.get_variables(vars) == [1.0, 2.0, 3.0]
-        var = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
-        @test Symbolics.get_variables(var) == [:x]
-    end
-
-    # Test isequal function
-    @testset "isequal" begin
-        var1 = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
-        var2 = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
-        var3 = HarmonicVariable(:y, "y", "v", 2.0, 3.0)
-        @test isequal(var1, var2) == true
-        @test isequal(var1, var3) == false
-    end
-end
+diff_eq = DifferentialEquation(
+    d(x, t, 2) + ω0 * x + α * x^3 + γ * d(x, t) + η * x^2 * d(x, t) ~ F * cos(ω * t), x
+) # define ODE
+
+add_harmonic!(diff_eq, x, ω) # specify the ansatz x = u(T) cos(ωt) + v(T) sin(ωt)
+
+
+harmonic_eq = get_harmonic_equations(diff_eq) # implement ansatz to get harmonic equations
+get_variables(harmonic_eq)
+
+# @testset "HarmonicVariable Tests" begin
+#     # Test display function
+#     @testset "Display" begin
+#         var = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
+#         @test display(var) == :x
+#         vars = [
+#             HarmonicVariable(:x, "x", "u", 1.0, 2.0),
+#             HarmonicVariable(:y, "y", "v", 2.0, 3.0),
+#         ]
+#         @test display(vars) == [:x, :y]
+#     end
+
+#     # Test _coordinate_transform function
+#     @testset "_coordinate_transform" begin
+#         @test _coordinate_transform(2.0, 1.0, 0.0, "u") == 2.0
+#         @test _coordinate_transform(2.0, 1.0, 0.0, "v") == 0.0
+#         @test _coordinate_transform(2.0, 1.0, 0.0, "a") == 2.0
+#     end
+
+#     # Test _create_harmonic_variable function
+#     @testset "_create_harmonic_variable" begin
+#         rule, hvar = _create_harmonic_variable(2.0, 1.0, 0.0, "u"; new_symbol="x")
+#         @test rule == 2.0
+#         @test hvar == HarmonicVariable(:x, "u_{2.0}", "u", 1.0, 2.0)
+#     end
+
+#     # Test substitute_all function
+#     @testset "substitute_all" begin
+#         eq = 2.0 * :x + 3.0 * :y
+#         rules = Dict(:x => 1.0, :y => 2.0)
+#         @test substitute_all(eq, rules) == 2.0 + 6.0
+#         var = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
+#         @test substitute_all(var, rules) == HarmonicVariable(1.0, "x", "u", 1.0, 2.0)
+#         vars = [
+#             HarmonicVariable(:x, "x", "u", 1.0, 2.0),
+#             HarmonicVariable(:y, "y", "v", 2.0, 3.0),
+#         ]
+#         @test substitute_all(vars, rules) == [
+#             HarmonicVariable(1.0, "x", "u", 1.0, 2.0),
+#             HarmonicVariable(2.0, "y", "v", 2.0, 3.0),
+#         ]
+#     end
+
+#     # Test Symbolics.get_variables function
+#     @testset "Symbolics.get_variables" begin
+#         vars = [1.0, 2.0, 3.0]
+#         @test Symbolics.get_variables(vars) == [1.0, 2.0, 3.0]
+#         var = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
+#         @test Symbolics.get_variables(var) == [:x]
+#     end
+
+#     # Test isequal function
+#     @testset "isequal" begin
+#         var1 = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
+#         var2 = HarmonicVariable(:x, "x", "u", 1.0, 2.0)
+#         var3 = HarmonicVariable(:y, "y", "v", 2.0, 3.0)
+#         @test isequal(var1, var2) == true
+#         @test isequal(var1, var3) == false
+#     end
+# end