feat: to_lab_frame for multiple variables ansatz (#287)

src/ExprUtils/Symbolics_utils.jl

@@ -126,3 +126,14 @@ is_harmonic(x, t) = is_harmonic(Num(x), Num(t))
 
 "Return true if `f` is a function of `var`."
 is_function(f, var) = any(isequal.(get_variables(f), var))
+
+"""
+Counts the number of derivatives of a symbolic variable.
+"""
+function count_derivatives(x::Symbolics.BasicSymbolic)
+    (Symbolics.isterm(x) || Symbolics.issym(x)) ||
+        error("The input is not a single term or symbol")
+    bool = Symbolics.is_derivative(x)
+    return bool ? 1 + count_derivatives(first(arguments(x))) : 0
+end
+count_derivatives(x::Num) = count_derivatives(Symbolics.unwrap(x))

src/HC_wrapper/homotopy_interface.jl

@@ -61,7 +61,7 @@ function HarmonicBalance.Problem(
     system = HomotopyContinuation.System(eqs_HC; variables=conv_vars, parameters=conv_para)
     J = HarmonicBalance.get_Jacobian(equations, variables) #all derivatives are assumed to be in the left hand side;
     return Problem(variables, parameters, system, J)
-end # is this funciton still needed/used?
+end # TODO is this funciton still needed/used?
 
 function System(eom::HarmonicEquation)
     eqs = expand_derivatives.(_remove_brackets(eom))

src/HarmonicBalance.jl

@@ -23,11 +23,19 @@ using Plots: Plots, plot, plot!, savefig, heatmap, Plot
 using Latexify: Latexify, latexify
 
 using Symbolics:
-    Symbolics, Num, Equation, @variables, expand_derivatives, get_variables, Differential
+    Symbolics,
+    Num,
+    Equation,
+    @variables,
+    expand_derivatives,
+    get_variables,
+    Differential,
+    unwrap,
+    wrap
 using SymbolicUtils: SymbolicUtils
 
 include("ExprUtils/ExprUtils.jl")
-using .ExprUtils: is_harmonic, substitute_all, drop_powers
+using .ExprUtils: is_harmonic, substitute_all, drop_powers, count_derivatives
 
 include("extention_functions.jl")
 include("utils.jl")

src/HarmonicEquation.jl

@@ -266,7 +266,9 @@ function get_harmonic_equations(
     slow_time = isnothing(slow_time) ? (@variables T; T) : slow_time
     fast_time = isnothing(fast_time) ? get_independent_variables(diff_eom)[1] : fast_time
 
-    all(isempty.(values(diff_eom.harmonics))) && error("No harmonics specified!")
+    for pair in diff_eom.harmonics
+        isempty(pair[2]) && error("No harmonics specified for the variable $(pair[1])!")
+    end
     eom = harmonic_ansatz(diff_eom, fast_time) # substitute trig functions into the differential equation
     eom = slow_flow(eom; fast_time=fast_time, slow_time=slow_time, degree=degree) # drop 2nd order time derivatives
     fourier_transform!(eom, fast_time) # perform averaging over the frequencies originally specified in dEOM

src/transform_solutions.jl

@@ -96,39 +96,39 @@ end
 # TRANSFORMATIONS TO THE LAB frame
 ###
 
-function to_lab_frame(soln, res, times)::Vector{AbstractFloat}
-    timetrace = zeros(length(times))
-
-    for var in res.problem.eom.variables
-        val = real(substitute_all(Symbolics.unwrap(_remove_brackets(var)), soln))
-        ω = real(Symbolics.unwrap(substitute_all(var.ω, soln)))
-        if var.type == "u"
-            timetrace .+= val * cos.(ω * times)
-        elseif var.type == "v"
-            timetrace .+= val * sin.(ω * times)
-        elseif var.type == "a"
-            timetrace .+= val
-        end
-    end
-    return timetrace
-end
-
 """
-    to_lab_frame(res::Result, times; index::Int, branch::Int)
-    to_lab_frame(soln::OrderedDict, res::Result, times)
+    to_lab_frame(res::Result, x::Num, times; index::Int, branch::Int)
+    to_lab_frame(soln::OrderedDict, res::Result, nat_var::Num, times)
 
-Transform a solution into the lab frame (i.e., invert the harmonic ansatz) for `times`.
+Transform a solution into the lab frame (i.e., invert the harmonic ansatz) for the natural variable `x` for `times`. You can also compute the velocity by passing `d(x,t)` for the natural variable `x`.
 Either extract the solution from `res::Result` by `index` and `branch` or input `soln::OrderedDict` explicitly.
 """
-to_lab_frame(res::Result, times; index::Int, branch::Int) =
-    to_lab_frame(res[index][branch], res, times)
+function to_lab_frame(soln::OrderedDict, res::Result, nat_var::Num, times)
+    count_derivatives(nat_var) > 2 &&
+        throw(ArgumentError("Only the first derivative is supported"))
 
-function to_lab_frame_velocity(soln, res, times)
-    timetrace = zeros(length(times))
+    var = if Symbolics.is_derivative(unwrap(nat_var))
+        wrap(first(Symbolics.arguments(unwrap(nat_var))))
+    else
+        nat_var
+    end
+    vars = filter(x -> isequal(x.natural_variable, var), res.problem.eom.variables)
+
+    return if Symbolics.is_derivative(unwrap(nat_var))
+        _to_lab_frame_velocity(soln, vars, times)
+    else
+        _to_lab_frame(soln, vars, times)
+    end
+end
+function to_lab_frame(res::Result, nat_var::Num, times; index::Int, branch::Int)
+    return to_lab_frame(res[index][branch], res, nat_var, times)
+end
 
-    for var in res.problem.eom.variables
-        val = real(substitute_all(Symbolics.unwrap(_remove_brackets(var)), soln))
-        ω = real(real(Symbolics.unwrap(substitute_all(var.ω, soln))))
+function _to_lab_frame_velocity(soln, vars, times)
+    timetrace = zeros(length(times))
+    for var in vars
+        val = real(substitute_all(unwrap(_remove_brackets(var)), soln))
+        ω = real(real(unwrap(substitute_all(var.ω, soln))))
         if var.type == "u"
             timetrace .+= -ω * val * sin.(ω * times)
         elseif var.type == "v"
@@ -138,12 +138,19 @@ function to_lab_frame_velocity(soln, res, times)
     return timetrace
 end
 
-"""
-    to_lab_frame_velocity(res::Result, times; index::Int, branch::Int)
-    to_lab_frame_velocity(soln::OrderedDict, res::Result, times)
+function _to_lab_frame(soln, vars, times)::Vector{AbstractFloat}
+    timetrace = zeros(length(times))
 
-Transform a solution's velocity into the lab frame (i.e., invert the harmonic ansatz for dx/dt ) for `times`.
-Either extract the solution from `res::Result` by `index` and `branch` or input `soln::OrderedDict` explicitly.
-"""
-to_lab_frame_velocity(res::Result, times; index, branch) =
-    to_lab_frame_velocity(res[index][branch], res, times)
+    for var in vars
+        val = real(substitute_all(unwrap(_remove_brackets(var)), soln))
+        ω = real(unwrap(substitute_all(var.ω, soln)))
+        if var.type == "u"
+            timetrace .+= val * cos.(ω * times)
+        elseif var.type == "v"
+            timetrace .+= val * sin.(ω * times)
+        elseif var.type == "a"
+            timetrace .+= val
+        end
+    end
+    return timetrace
+end

test/API.jl

@@ -1,6 +1,6 @@
 using HarmonicBalance
 
-@testset begin
+@testset "Equation{Vector}" begin
     # define equation of motion
     @variables ω1, ω2, t, ω, F, γ, α1, α2, k, x(t), y(t)
     rhs = [
@@ -27,3 +27,20 @@ end
     @test_throws MethodError get_steady_states(harmonic_eq, varied, threading=true)
     @test_throws ArgumentError get_steady_states(harmonic_eq, Dict(varied), threading=true)
 end
+
+@testset "forgot variable" begin
+    @variables Ω γ λ F x θ η α ω0 ω t T ψ
+    @variables x(t) y(t)
+
+    natural_equation = [
+        d(d(x, t), t) + γ * d(x, t) + Ω^2 * x + α * x^3 ~ F * cos(ω * t),
+        d(d(y, t), t) + γ * d(y, t) + Ω^2 * y + α * y^3 ~ 0,
+    ]
+    dEOM = DifferentialEquation(natural_equation, [x, y])
+
+    @test_throws ErrorException get_harmonic_equations(dEOM)
+
+    add_harmonic!(dEOM, x, ω)
+    # add_harmonic!(dEOM, y, ω)
+    @test_throws ErrorException get_harmonic_equations(dEOM)
+end

test/symbolics.jl

@@ -177,3 +177,12 @@ end
     @eqtest get_independent(cos(t), t) == 0
     @eqtest get_independent(cos(t)^2 + 5, t) == 5
 end
+
+@testset "count_derivatives" begin
+    using HarmonicBalance.ExprUtils: count_derivatives
+    @variables t x(t) y(t)
+    @test count_derivatives(x) == 0
+    @test count_derivatives(d(x, t)) == 1
+    @test count_derivatives(d(d(x, t), t)) == 2
+    @test_throws ErrorException count_derivatives(d(d(5 * x, t), t))
+end

test/transform_solutions.jl

@@ -5,14 +5,17 @@ const SEED = 0xd8e5d8df
 Random.seed!(SEED)
 
 @variables Ω γ λ F x θ η α ω0 ω t T ψ
-@variables x(t)
+@variables x(t) y(t)
 
-natural_equation = d(d(x, t), t) + γ * d(x, t) + Ω^2 * x + α * x^3
-forces = F * cos(ω * t)
-dEOM = DifferentialEquation(natural_equation ~ forces, x)
+natural_equation = [
+    d(d(x, t), t) + γ * d(x, t) + Ω^2 * x + α * x^3 ~ F * cos(ω * t),
+    d(d(y, t), t) + γ * d(y, t) + Ω^2 * y + α * y^3 ~ 0,
+]
+dEOM = DifferentialEquation(natural_equation, [x, y])
 
 add_harmonic!(dEOM, x, ω)
-harmonic_eq = get_harmonic_equations(dEOM; slow_time=T, fast_time=t);
+add_harmonic!(dEOM, y, ω)
+harmonic_eq = get_harmonic_equations(dEOM);
 
 fixed = (Ω => 1.0, γ => 1e-2, F => 1e-3, α => 1.0)
 varied = ω => range(0.9, 1.1, 10)
@@ -21,9 +24,12 @@ res = get_steady_states(harmonic_eq, varied, fixed; show_progress=false, seed=SE
 transform_solutions(res, "u1^2+v1^2")
 transform_solutions(res, "√(u1^2+v1^2)"; realify=true)
 
-# transform_solutions(res.solutions[1][1], "√(u1^2+v1^2)", harmonic_eq)
-
-times = 0:1:10
-soln = res[5][1]
-HarmonicBalance.to_lab_frame(soln, res, times)
-HarmonicBalance.to_lab_frame_velocity(soln, res, times)
+@testset "to_lab_frame" begin
+    using HarmonicBalance: to_lab_frame
+    @variables z(t)
+    times = 0:1:10
+    @test to_lab_frame(res, x, times; index=1, branch=1) != zeros(length(times))
+    @test all(isapprox.(to_lab_frame(res, y, times; index=1, branch=1), 0.0, atol=1e-10))
+    @test all(to_lab_frame(res, z, times; index=1, branch=1) .≈ zeros(length(times)))
+    @test to_lab_frame(res, d(x, t), times; index=1, branch=1) != zeros(length(times))
+end