fix: make 2nd order response work and add some tests (#293)

src/ExprUtils/Symbolics_utils.jl

@@ -36,8 +36,8 @@ $(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)
+Subtype = Union{Num,Equation,BasicSymbolic}
+function substitute_all(x::Subtype, rules::Dict; include_derivatives=true)
     if include_derivatives
         rules = merge(
             rules,
@@ -54,24 +54,10 @@ function substitute_all(dict::Dict, rules::Dict)::Dict
 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]
+substitute_all(x::Subtype, rules::Collections) = substitute_all(x, Dict(rules))
+function substitute_all(x::Complex{Num}, rules::Collections)
+    return substitute_all(x.re, rules) + im * substitute_all(x.im, rules)
+end
 
 get_independent(x::Num, t::Num) = get_independent(x.val, t)
 function get_independent(x::Complex{Num}, t::Num)

src/LinearResponse/response.jl

@@ -20,12 +20,12 @@ function get_response_matrix(diff_eq::DifferentialEquation, freq::Num; order=2):
     # get the response matrix by summing the orders
     M = get_Jacobian(eom.equations, get_variables(eom))
     for n in 1:order
-        M += (i * freq)^n * get_Jacobian(eom.equations, d(get_variables(eom), T, n))
+        M += (im * freq)^n * get_Jacobian(eom.equations, d(get_variables(eom), T, n))
     end
     M = substitute_all(
         M, [var => declare_variable(var_name(var)) for var in get_variables(eom)]
     )
-    return substitute_all(expand_derivatives.(M), i => im)
+    return M
 end
 
 "Get the response matrix corresponding to `res`.
@@ -37,9 +37,13 @@ function ResponseMatrix(res::Result; rules=Dict())
     M = get_response_matrix(res.problem.eom.natural_equation, Δ; order=2)
     M = substitute_all(M, merge(res.fixed_parameters, rules))
     symbols = _free_symbols(res)
-    compiled_M = [Symbolics.build_function(el, cat(symbols, [Δ]; dims=1)) for el in M]
-
-    return ResponseMatrix(eval.(compiled_M), symbols, res.problem.eom.variables)
+    compiled_M = map(M) do el
+        args = cat(symbols, [Δ]; dims=1)
+        f_re = Symbolics.build_function(el.re, args; expression=Val{false})
+        f_im = Symbolics.build_function(el.im, args; expression=Val{false})
+        (args...) -> f_re(args...) + im * f_im(args...)
+    end
+    return ResponseMatrix(compiled_M, symbols, res.problem.eom.variables)
 end
 
 """Evaluate the response matrix `resp` for the steady state `s` at (lab-frame) frequency `Ω`."""

test/linear_response.jl

@@ -1,4 +1,5 @@
-using HarmonicBalance
+using HarmonicBalance, Symbolics
+HB = HarmonicBalance
 
 @variables α, ω, ω0, F, γ, t, x(t);
 
@@ -13,31 +14,46 @@ varied = ω => range(0.9, 1.1, 10)
 
 result = get_steady_states(harmonic_eq, varied, fixed; show_progress=false, seed=SEED)
 
-plot_linear_response(
-    result, x; branch=1, Ω_range=range(0.9, 1.1, 10), order=1, logscale=true
-)
-
-plot_rotframe_jacobian_response(
-    result; Ω_range=range(0.01, 1.1, 10), branch=1, logscale=true
-)
+@testset "first order" begin
+    plot_linear_response(
+        result, x; branch=1, Ω_range=range(0.9, 1.1, 10), order=1, logscale=true
+    )
+    plot_rotframe_jacobian_response(
+        result; Ω_range=range(0.01, 1.1, 10), branch=1, logscale=true
+    )
+end
 
-plot_eigenvalues(result; branch=1)
-plot_eigenvalues(result; branch=1, type=:re, class="all")
+@testset "second order" begin
+    response_matrix = HB.LinearResponse.ResponseMatrix(result)
+    M = response_matrix.matrix
+    @test M[1](ones(4)) isa ComplexF64
 
-@testset begin
-    @variables α λ ω0 ω ωₚ F t x(t)
-    diff_eq = DifferentialEquation(
-        d(x, t, 2) + (ω0^2 - λ * cos(2 * ω * t)) * x + α * x^3 + γ * d(x, t) ~
-            F * cos(ωₚ * t),
-        x,
+    plot_linear_response(
+        result, x; branch=1, Ω_range=range(0.9, 1.1, 10), order=2, logscale=true
     )
+end
 
-    add_harmonic!(diff_eq, x, ω)
-    add_harmonic!(diff_eq, x, ωₚ)
-    harmonic_eq = get_harmonic_equations(diff_eq)
-
-    fixed = (γ => 0.008, ω0 => 1.0, α => 1.0, F => 0.0, ωₚ => 1.0, λ => 0.016)
-    varied = (ω => range(0.995, 1.005, 20))
-    result_ω = get_steady_states(harmonic_eq, varied, fixed; show_progress=false, seed=SEED)
-    @test_throws ErrorException plot_eigenvalues(result_ω; branch=1)
+@testset "eigenvalues" begin
+    plot_eigenvalues(result; branch=1)
+    plot_eigenvalues(result; branch=1, type=:re, class="all")
+
+    @testset "NaN Exception Error" begin
+        @variables α λ ω0 ω ωₚ F t x(t)
+        diff_eq = DifferentialEquation(
+            d(x, t, 2) + (ω0^2 - λ * cos(2 * ω * t)) * x + α * x^3 + γ * d(x, t) ~
+                F * cos(ωₚ * t),
+            x,
+        )
+
+        add_harmonic!(diff_eq, x, ω)
+        add_harmonic!(diff_eq, x, ωₚ)
+        harmonic_eq = get_harmonic_equations(diff_eq)
+
+        fixed = (γ => 0.008, ω0 => 1.0, α => 1.0, F => 0.0, ωₚ => 1.0, λ => 0.016)
+        varied = (ω => range(0.995, 1.005, 20))
+        result_ω = get_steady_states(
+            harmonic_eq, varied, fixed; show_progress=false, seed=SEED
+        )
+        @test_throws ErrorException plot_eigenvalues(result_ω; branch=1)
+    end
 end