fix: give each natural variable its own harmonic in the van der Pol ansatz

src/HarmonicBalance.jl

@@ -55,16 +55,16 @@ using Reexport: @reexport
 @reexport using HarmonicSteadyState
 
 # Precompilation setup
-# using PrecompileTools: @setup_workload, @compile_workload
-# @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
+using PrecompileTools: @setup_workload, @compile_workload
+@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
 
 # symbolics equations
 export @variables

src/krylov-bogoliubov.jl

@@ -59,9 +59,10 @@ function get_krylov_equations(
     eom = slow_flow(eom; fast_time=fast_time, slow_time=slow_time, degree=2)
 
     rearrange!(eom, d(get_variables(eom), slow_time))
-    eom.equations = Symbolics.expand.(Symbolics.simplify.(eom.equations))
-    eom.equations = Symbolics.expand.(Symbolics.simplify.(eom.equations))
-    #^ need it two times to get it completely simplified due to some weird bug in Symbolics.jl
+    # `average` works term by term, so a flat sum of products is enough. A full
+    # `Symbolics.simplify` here also cancels the fraction, but its cost explodes with the
+    # number of harmonics; `simplify_fractions` below does the cancelling instead.
+    eom.equations = Symbolics.expand.(eom.equations)
 
     if order == 1
         average!(eom, fast_time)
@@ -83,6 +84,10 @@ function get_krylov_equations(
     end
 
     change_convention!(eom, slow_time)
+    # averaging has removed the fast time, so cancelling `rearrange!`'s denominator is cheap here
+    eom.equations = [
+        Symbolics.simplify_fractions(Num(eq.lhs)) ~ eq.rhs for eq in eom.equations
+    ]
     return eom
 end
 function proper_krylov_system(diff_eom::QuestBase.DifferentialEquation, order::Int)
@@ -117,11 +122,13 @@ function van_der_Pol(eom::QuestBase.DifferentialEquation, t::Num)
 
     # keep count to label new variables
     uv_idx = 1
-    ω = first(flatten(unique(values(dEOM.harmonics))))
     nvars = get_variables(dEOM)
     nvars = nvars[(length(nvars) ÷ 2 + 1):end]
 
     for nvar in nvars # sum over natural variables
+        # each variable rotates at its own harmonic; the harmonics may be
+        # commensurate (e.g. ω and 3ω) but need not be equal
+        ω = first(dEOM.harmonics[nvar])
         rule_u, hvar_u = _create_harmonic_variable(
             nvar, ω, t, "u"; new_symbol="u" * string(uv_idx)
         )
@@ -165,10 +172,10 @@ end
 
 function take_trig_integral(x::BasicSymbolic, ω, t)
     if isdiv(x)
+        # termwise integration leaves the denominator alone, so it goes back as is; `expand`
+        # is enough to flatten for the averaging that follows
         arg_num = Symbolics.arguments(x.num)
-        return Symbolics.simplify(
-            Symbolics.expand(sum(take_trig_integral.(arg_num, ω, t)) * ω)
-        ) / (x.den * ω)
+        return Symbolics.expand(sum(take_trig_integral.(arg_num, ω, t))) / x.den
     else
         all_terms = get_all_terms(Num(x))
         trigs = filter(z -> is_trig(z), all_terms)

src/precompilation.jl

@@ -13,6 +13,9 @@ add_harmonic!(dEOM, x, ω)
 
 harmonic_eq = get_harmonic_equations(dEOM; slow_time=T, fast_time=t);
 
+get_krylov_equations(dEOM; order=1)
+get_krylov_equations(dEOM; order=2)
+
 # fixed = (Ω => 1.0, γ => 1e-2, λ => 5e-2, F => 0, α => 1.0, η => 0.3, θ => 0, ψ => 0)
 # varied = ω => range(0.9, 1.1, 20)
 # res = get_steady_states(harmonic_eq, varied, fixed; show_progress=false)

test/krylov.jl

@@ -1,6 +1,7 @@
 using HarmonicBalance
 using HarmonicBalance: get_krylov_equations
-using Symbolics: Symbolics
+using QuestBase: QuestBase
+using Symbolics: Symbolics, substitute
 using Test
 
 # Regression for the Symbolics 7 slow-flow bug. `trig_reduce` linearises the
@@ -50,6 +51,43 @@ end
     assert_nondegenerate(get_krylov_equations(diff_eom; order=2))
 end
 
+@testset "commensurate frequencies (issue #251)" begin
+    # `van_der_Pol` used to pick a single frequency off the front of the harmonics
+    # dict and build the ansatz with it for *every* natural variable, so a system
+    # with x at ω and y at 3ω got y = u2*cos(ωt) + v2*sin(ωt) and the whole 3ω
+    # sector came out wrong. Each variable must rotate at its own harmonic.
+    @variables t x(t) y(t) ω0 ω F α J
+    eq1 = d(d(x, t), t) + ω0^2 * x + α * x^3 ~ F * cos(ω * t) + J * y
+    eq2 = d(d(y, t), t) + ω0^2 * y + α * y^3 ~ F * cos(ω * t) + J * x
+    diff_eom = DifferentialEquation([eq1, eq2], [x, y])
+    add_harmonic!(diff_eom, x, ω)
+    add_harmonic!(diff_eom, y, 3 * ω)
+
+    krylov_eq = get_krylov_equations(diff_eom; order=1)
+    rearranged = HarmonicBalance.rearrange_standard(get_harmonic_equations(diff_eom))
+
+    # the ansatz itself must carry the two distinct harmonics
+    @test Set(Symbolics.unwrap.(getfield.(krylov_eq.variables, :ω))) ==
+        Set(Symbolics.unwrap.([ω, ω, 3 * ω, 3 * ω]))
+
+    # KB and harmonic balance must agree up to the opposite sign convention
+    lhss = [
+        QuestBase.expand_fraction.(getfield.(eom.equations, :lhs)) for
+        eom in (krylov_eq, rearranged)
+    ]
+    symbols = unique(
+        reduce(vcat, [collect(Symbolics.get_variables(e)) for e in reduce(vcat, lhss)])
+    )
+    for _ in 1:3
+        # keep the sample away from 0 so no denominator blows up
+        subs = Dict(symbols .=> rand(length(symbols)) .+ 0.5)
+        for (k, h) in zip(lhss...)
+            residual = Symbolics.value(substitute(k, subs) + substitute(h, subs))
+            @test Float64(residual) ≈ 0.0 atol = 1e-10
+        end
+    end
+end
+
 @testset "three-wave mixing (large order-2 expressions)" begin
     # The example that exposed the bug: quadratic + cubic nonlinearity driven at 2ω.
     # Its order-2 equations are large enough that a naive `simplify`-the-whole-thing