From c741b894b70c374a0df1e6c69a1e03a07ac1c877 Mon Sep 17 00:00:00 2001 From: Orjan Ameye Date: Sun, 27 Oct 2024 12:30:15 +0100 Subject: [PATCH] fix: make 2nd order response work and add some tests --- src/ExprUtils/Symbolics_utils.jl | 26 +++---------- src/LinearResponse/response.jl | 14 ++++--- test/linear_response.jl | 64 ++++++++++++++++++++------------ 3 files changed, 55 insertions(+), 49 deletions(-) diff --git a/src/ExprUtils/Symbolics_utils.jl b/src/ExprUtils/Symbolics_utils.jl index a5560a1f..ff467066 100644 --- a/src/ExprUtils/Symbolics_utils.jl +++ b/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) diff --git a/src/LinearResponse/response.jl b/src/LinearResponse/response.jl index ac854834..312dcb68 100644 --- a/src/LinearResponse/response.jl +++ b/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 `Ω`.""" diff --git a/test/linear_response.jl b/test/linear_response.jl index c0d86f55..a862cb91 100644 --- a/test/linear_response.jl +++ b/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