test: add KB tests

Project.toml

@@ -53,7 +53,7 @@ Latexify = "0.16"
 ModelingToolkit = "9.34"
 NonlinearSolve = "3.14"
 OrderedCollections = "1.6"
-OrdinaryDiffEqRosenbrock = "1.1"
+OrdinaryDiffEqRosenbrock = "1.3"
 OrdinaryDiffEqTsit5 = "1.1"
 Peaks = "0.5"
 Plots = "1.39"

src/ExprUtils/Symbolics_utils.jl

@@ -9,6 +9,16 @@ function expand_all(x)
 end
 expand_all(x::Complex{Num}) = expand_all(x.re) + im * expand_all(x.im)
 
+function expand_fraction(x::BasicSymbolic)
+    @compactified x::BasicSymbolic begin
+        Add => _apply_termwise(expand_fraction, x)
+        Mul => _apply_termwise(expand_fraction, x)
+        Div => sum([arg / x.den for arg in arguments(x.num)])
+        _   => x
+    end
+end
+expand_fraction(x::Num) = Num(expand_fraction(x.val))
+
 "Apply a function f on every member of a sum or a product"
 function _apply_termwise(f, x::BasicSymbolic)
     @compactified x::BasicSymbolic begin

src/KrylovBogoliubov/KrylovEquation.jl

@@ -111,6 +111,7 @@ function van_der_Pol(eom::DifferentialEquation, t::Num)
             nvar, ω, t, "v"; new_symbol="v" * string(uv_idx)
         )
         rule = rule_u - rule_v
+        # ~ this is a choice, we use u*cos(ωt) + v*sin(ωt) as the ansatz
         rules[nvar] = rule
 
         D = Differential(t)

test/krylov.jl

@@ -1,5 +1,6 @@
 using HarmonicBalance;
 HB = HarmonicBalance;
+using Test
 
 @variables t T x(t) y(t) # symbolic variables
 @variables ω ω0 γ F α λ ψ θ η
@@ -29,3 +30,82 @@ varied = (ω => range(0.99, 1.01, 5), λ => range(1e-6, 0.05, 5))
 
 res1 = get_steady_states(harmonic_eq1, varied, fixed; show_progress=false);
 res2 = get_steady_states(harmonic_eq2, varied, fixed; show_progress=false);
+
+@testset "not damped" begin
+    using HarmonicBalance.ExprUtils: expand_fraction
+    using Symbolics: substitute
+
+    @testset "single resonator" begin
+        @variables t x(t) y(t) ω0 ω F α # symbolic variables
+        eq1 = d(d(x, t), t) + ω0^2 * x + α * x^3 ~ F * cos(ω * t)
+        EOM = DifferentialEquation(eq1, x)
+        add_harmonic!(EOM, x, ω)
+        krylov_eq = get_krylov_equations(EOM; order=1)
+        harmonic_eq = get_harmonic_equations(EOM)
+        rearranged = HarmonicBalance.rearrange_standard(harmonic_eq)
+
+        @testset for i in 1:2
+            eqk = expand_fraction(krylov_eq.equations[i].lhs)
+            eqh = expand_fraction(rearranged.equations[i].lhs)
+            @variables T u1(T) v1(T) ω0 ω F α # symbolic variables
+            subs = Dict(u1 => 1, v1 => 1, α => 1, F => 1, ω0 => 1, ω => 1)
+            solk = substitute(eqk, subs)
+            solh = substitute(eqh, subs)
+            @test Float64(solk + solh) ≈ 0.0 atol = 1e-10
+            # ^ different ansatz
+        end
+    end
+
+    @testset "multiple resonators" begin
+        @variables t x(t) y(t) ω0 ω F α J # symbolic variables
+        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
+        EOM = DifferentialEquation([eq1, eq2], [x, y])
+        add_harmonic!(EOM, x, ω)
+        add_harmonic!(EOM, y, ω)
+        krylov_eq = get_krylov_equations(EOM; order=1)
+        harmonic_eq = get_harmonic_equations(EOM)
+        rearranged = HarmonicBalance.rearrange_standard(harmonic_eq)
+
+        @testset for i in 1:4
+            eqk = expand_fraction(krylov_eq.equations[i].lhs)
+            eqh = expand_fraction(rearranged.equations[i].lhs)
+            @variables T u1(T) v1(T) ω0 ω F α # symbolic variables
+            subs = Dict([u1, v1, α, F, ω0, ω, J] .=> rand(7))
+            solk = substitute(eqk, subs)
+            solh = substitute(eqh, subs)
+            @test Float64(solk + solh) ≈ 0.0 atol = 1e-10
+            # ^ different ansatz
+        end
+    end
+
+    @testset "two resonators at different frequencies" begin
+        @variables ω₁, ω₂, t, ω, F, γ₁, γ₂
+        @variables α₁, J₂, J₁, α₂, η₁, η₂
+        @variables x(t), y(t)
+        # eq1 = d(d(x, t), t) + ω₁^2 * x + α₁ * x^3 + 3 * J₁ * x^2 * y + J₂ * y^2 * x
+        # eq2 = d(d(y, t), t) + ω₂^2 * y + α₂ * y^3 + J₁ * x^3 + J₂ * x^2 * y
+        eq1 = d(d(x, t), t) + ω₁^2 * x + α₁ * x^3 + 3 * J₁ * x^2 * y + J₂ * y^2 * x
+        eq2 = d(d(y, t), t) + ω₂^2 * y + α₂ * y^3 + J₁ * x^3 + J₂ * x^2 * y
+        forces = [F * cos(ω * t), 0]
+        dEOM_temp = DifferentialEquation([eq1, eq2] - forces, [x, y])
+
+        add_harmonic!(dEOM_temp, x, ω) # x will rotate at ω
+        add_harmonic!(dEOM_temp, y, 3 * ω) # y will rotate at 3*ω
+
+        krylov_eq = get_krylov_equations(dEOM_temp; order=1)
+        harmonic_eq = get_harmonic_equations(dEOM_temp)
+        rearranged = HarmonicBalance.rearrange_standard(harmonic_eq)
+
+        @testset for i in 1:4
+            eqk = expand_fraction(krylov_eq.equations[1].lhs)
+            eqh = expand_fraction(rearranged.equations[1].lhs)
+            @variables T u1(T) v1(T) u2(T) v2(T)
+            subs = Dict([u1, v1, u2, v2, ω₁, ω₂, ω, F, J₂, J₁, α₁, α₂] .=> rand(12))
+            solk = substitute(eqk, subs)
+            solh = substitute(eqh, subs)
+            @test Float64(solk + solh) ≈ 0.0 broken = true
+            # ^ different ansatz
+        end
+    end
+end

test/symbolics.jl

@@ -178,6 +178,12 @@ end
     @eqtest get_independent(cos(t)^2 + 5, t) == 5
 end
 
+@testset "expand_fraction" begin
+    using HarmonicBalance.ExprUtils: expand_fraction
+    @variables a, b, c
+
+    @eqtest expand_fraction((a + b) / c) == a / c + b / c
+end
 @testset "count_derivatives" begin
     using HarmonicBalance.ExprUtils: count_derivatives
     @variables t x(t) y(t)