Store the role of a harmonic variable as a HarmonicKind enum

Project.toml

@@ -1,6 +1,6 @@
 name = "QuestBase"
 uuid = "7e80f742-43d6-403d-a9ea-981410111d43"
-version = "0.5.3"
+version = "0.6.0"
 authors = ["Orjan Ameye <orjan.ameye@hotmail.com>", "Jan Kosata <kosataj@phys.ethz.ch>", "Javier del Pino <jdelpino@phys.ethz.ch>"]
 
 [deps]

docs/src/API.md

@@ -28,6 +28,8 @@ QuestBase.get_variables(::QuestBase.DifferentialEquation)
 
 ```@docs
 QuestBase.HarmonicVariable
+QuestBase.HarmonicKind
+QuestBase.kind_label
 QuestBase.add_harmonic!
 ```
 

src/HarmonicEquation.jl

@@ -113,23 +113,22 @@ function Base.show(io::IO, eom::HarmonicEquation)
     println(io, "A set of ", length(eom.equations), " harmonic equations")
     println(io, "Variables: ", join(string.(get_variables(eom)), ", "))
     println(io, "Parameters: ", join(string.(eom.parameters), ", "))
-    println(io, "\nHarmonic ansatz: ", _show_ansatz(eom))
+    println(io, "\nHarmonic ansatz: ", show_ansatz(eom))
     println(io, "\nHarmonic equations:")
     return [println(io, "\n", eq) for eq in eom.equations]
 end
 
 """Gives the full harmonic ansatz used to construct `eom`."""
-function _show_ansatz(eom::HarmonicEquation)
-    output = ""
-    vars = unique(getfield.(eom.variables, :natural_variable))
-    for nat_var in vars
-        harm_vars = filter(
-            x -> isequal(nat_var, x.natural_variable) && x.type !== "Hopf", eom.variables
-        )
-        ansatz = join([_show_ansatz(var) for var in harm_vars], " + ")
-        output *= "\n" * string(nat_var) * " = " * ansatz
+function show_ansatz(eom::HarmonicEquation)
+    natural_variables = unique(Num[hv.natural_variable for hv in eom.variables])
+    lines = map(natural_variables) do nat_var
+        terms = String[
+            show_ansatz(hv) for hv in eom.variables if
+            isequal(nat_var, hv.natural_variable) && hv.kind != LimitCycleFrequency
+        ]
+        return "\n" * string(nat_var) * " = " * join(terms, " + ")
     end
-    return output
+    return join(lines)
 end
 
 Base.show(eom::HarmonicEquation) = show_fields(eom)
@@ -177,19 +176,19 @@ function get_independent_variables(eom::HarmonicEquation)
     )
 end
 
-_remove_brackets(var::Num) = declare_variable(var_name(var))
-_remove_brackets(vars::Vector{Num}) = _remove_brackets.(vars)
-_remove_brackets(hv::HarmonicVariable) = _remove_brackets(hv.symbol)
+remove_brackets(var::Num) = declare_variable(var_name(var))
+remove_brackets(vars::Vector{Num}) = remove_brackets.(vars)
+remove_brackets(hv::HarmonicVariable) = remove_brackets(hv.symbol)
 
 "Returns the equation system in `eom`, dropping all argument brackets (i.e., u(T) becomes u)."
-function _remove_brackets(eom::HarmonicEquation)
-    variable_rules = [var => _remove_brackets(var) for var in get_variables(eom)]
+function remove_brackets(eom::HarmonicEquation)
+    variable_rules = [var => remove_brackets(var) for var in get_variables(eom)]
     equations_lhs = Num.(getfield.(eom.equations, :lhs) - getfield.(eom.equations, :rhs))
     return substitute_all(equations_lhs, variable_rules)
 end
 
-function _remove_brackets(eqs::Vector{Num}, vars::Vector{Num})
-    vars_ = _remove_brackets.(vars)
+function remove_brackets(eqs::Vector{Num}, vars::Vector{Num})
+    vars_ = remove_brackets.(vars)
     variable_rules = Dict(zip(vars, vars_))
     return substitute_all(eqs, variable_rules), vars_
 end
@@ -245,7 +244,7 @@ function is_rearranged(eom::HarmonicEquation)
     rhs = getfield.(eom.equations, :rhs)
 
     HB_bool = isequal(rhs, dvar)
-    hopf_bool = any(hv -> hv.type == "Hopf", eom.variables)
+    hopf_bool = any(hv -> hv.kind == LimitCycleFrequency, eom.variables)
     MF_bool = !any(occurs_in(dv, l) for dv in dvar, l in lhs)
 
     return HB_bool || hopf_bool || MF_bool

src/HarmonicVariable.jl

@@ -1,3 +1,33 @@
+"""
+Role of a [`HarmonicVariable`](@ref) in the ansatz of its natural variable.
+
+- `CosineQuadrature`: the variable `u` multiplying `cos(ωt)`.
+- `SineQuadrature`: the variable `v` multiplying `sin(ωt)`.
+- `ZeroHarmonic`: the constant part `a`.
+- `LimitCycleFrequency`: the limit-cycle frequency that replaces a gauge-fixed quadrature.
+- `Unassigned`: a bare symbol outside any harmonic ansatz.
+"""
+@enum HarmonicKind begin
+    CosineQuadrature
+    SineQuadrature
+    ZeroHarmonic
+    LimitCycleFrequency
+    Unassigned
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Short label of `kind` used in variable names: `u`, `v`, `a`, `Hopf` or the empty string.
+"""
+function kind_label(kind::HarmonicKind)
+    kind == CosineQuadrature && return "u"
+    kind == SineQuadrature && return "v"
+    kind == ZeroHarmonic && return "a"
+    kind == LimitCycleFrequency && return "Hopf"
+    return ""
+end
+
 """
 $(TYPEDEF)
 
@@ -11,24 +41,24 @@ mutable struct HarmonicVariable
     symbol::Num
     """Human-readable labels of the variable, used for plotting."""
     name::String
-    """Type of the variable (u or v for quadratures, a for a constant, Hopf for Hopf etc.)"""
-    type::String
+    """Role of the variable in the ansatz."""
+    kind::HarmonicKind
     """The harmonic being described."""
     ω::Num
     """The natural variable whose harmonic is being described."""
     natural_variable::Num
 end
 
 function HarmonicVariable(symbol::Num)
-    return HarmonicVariable(symbol, "", "", Num(1), Num(0))
+    return HarmonicVariable(symbol, "", Unassigned, Num(1), Num(0))
 end
 
 function Base.copy(hv::HarmonicVariable)
-    return HarmonicVariable(hv.symbol, hv.name, hv.type, hv.ω, hv.natural_variable)
+    return HarmonicVariable(hv.symbol, hv.name, hv.kind, hv.ω, hv.natural_variable)
 end
 
 function Base.show(io::IO, hv::HarmonicVariable)
-    if isempty(hv.type)
+    if hv.kind == Unassigned
         s = "Harmonic variable " * string.(hv.symbol)
     else
         s =
@@ -43,43 +73,48 @@ function Base.show(io::IO, hv::HarmonicVariable)
 end
 
 """Gives the relation between `var` and the underlying natural variable."""
-function _show_ansatz(var::HarmonicVariable)
-    if isempty(var.type)
-        return string(var.symbol)
-    end
+function show_ansatz(var::HarmonicVariable)
+    var.kind in (CosineQuadrature, SineQuadrature) || return string(var.symbol)
     t = arguments(unwrap(var.natural_variable))
     t = length(t) == 1 ? string(t[1]) : error("more than 1 independent variable")
-    ω = string(var.ω)
-    terms = Dict("u" => "*cos(" * ω * t * ")", "v" => "*sin(" * ω * t * ")", "a" => "")
-    return string(string(var.symbol) * terms[var.type])
+    trig = var.kind == CosineQuadrature ? "cos" : "sin"
+    return string(var.symbol) * "*" * trig * "(" * string(var.ω) * t * ")"
 end
 
 Base.display(var::HarmonicVariable) = display(var.name)
 Base.display(var::Vector{HarmonicVariable}) = display.(getfield.(var, Symbol("name")))
 
-function coordinate_transform(new_var, ω, t, type)::Num
-    coords = Dict([
-        "u" => new_var * cos(ω * t), "v" => new_var * sin(ω * t), "a" => new_var
-    ])
-    return coords[type]
+function coordinate_transform(new_var::Num, ω::Num, t::Num, kind::HarmonicKind)
+    kind == CosineQuadrature && return new_var * cos(ω * t)
+    kind == SineQuadrature && return new_var * sin(ω * t)
+    kind == ZeroHarmonic && return new_var
+    return throw(
+        ArgumentError("no coordinate transform for a harmonic variable of kind $kind")
+    )
 end
 
 """
 $(TYPEDSIGNATURES)
 
 Creates a new harmonic variable and its corresponding transformation rule.
 This function takes a natural variable (`nat_var`), a harmonic frequency (`ω`),
-an independent variable (`t`), and a type (`type`), and creates a new harmonic variable with
+an independent variable (`t`), and a [`HarmonicKind`](@ref) (`kind`), and creates a new harmonic variable with
 the specified `new_symbol`. It returns a tuple containing the transformation rule and
 the new harmonic variable.
 """
-function _create_harmonic_variable(
-    nat_var::Num, ω::Num, t::Num, type::String; new_symbol::String
+function create_harmonic_variable(
+    nat_var::Num, ω::Num, t::Num, kind::HarmonicKind; new_symbol::String
 )::Tuple{Num,HarmonicVariable}
     new_var = declare_variable(new_symbol, t)
-    name = type * "_{" * var_name(nat_var) * "," * Base.replace(string(ω), "*" => "") * "}"
-    rule = coordinate_transform(new_var, ω, t, type)
-    hvar = HarmonicVariable(new_var, name, type, ω, nat_var)
+    name =
+        kind_label(kind) *
+        "_{" *
+        var_name(nat_var) *
+        "," *
+        Base.replace(string(ω), "*" => "") *
+        "}"
+    rule = coordinate_transform(new_var, ω, t, kind)
+    hvar = HarmonicVariable(new_var, name, kind, ω, nat_var)
     return rule, hvar
 end
 
@@ -95,7 +130,7 @@ function QuestBase.substitute_all(var::HarmonicVariable, rules)
     return HarmonicVariable(
         substitute_all(sym, rules),
         var.name,
-        var.type,
+        var.kind,
         substitute_all(freq, rules),
         var.natural_variable,
     )

src/Variables.jl

@@ -16,7 +16,7 @@ $(TYPEDSIGNATURES)
 Declare a symbolic variable named `name`.
 
 `Symbolics.@variables` needs the name as a literal, so names only known at runtime (`u1`,
-`v2`, the bracket-free copies [`_remove_brackets`](@ref) makes) have to go through
+`v2`, the bracket-free copies [`remove_brackets`](@ref) makes) have to go through
 `Symbolics.variable` instead.
 
 This used to build the variable with `@eval` and bind the result inside this module. Nothing

test/DifferentialEquations.jl

@@ -16,7 +16,7 @@ using QuestBase:
     is_rearranged_standard,
     rearrange,
     substitute_all,
-    _remove_brackets,
+    remove_brackets,
     declare_variables,
     dummy_symbolic_Jacobian,
     get_all_terms

test/HarmonicEquation.jl

@@ -6,14 +6,16 @@ using QuestBase
 using QuestBase:
     DifferentialEquation,
     HarmonicVariable,
+    CosineQuadrature,
+    SineQuadrature,
     HarmonicEquation,
     harmonic_parameters,
     get_variables,
     get_independent_variables,
     rearrange_standard,
     substitute_all,
     is_rearranged,
-    _remove_brackets,
+    remove_brackets,
     declare_variables,
     dummy_symbolic_Jacobian,
     @eqtest,
@@ -33,8 +35,8 @@ eq2 = D(v) ~ -u + v
 nat_eq = DifferentialEquation([eq1, eq2], [x, y])
 
 # Create test harmonic variables
-hv1 = HarmonicVariable(u, "test", "u", Num(1.0), x)
-hv2 = HarmonicVariable(v, "test", "v", Num(1.0), y)
+hv1 = HarmonicVariable(u, "test", CosineQuadrature, Num(1.0), x)
+hv2 = HarmonicVariable(v, "test", SineQuadrature, Num(1.0), y)
 
 # Test constructor
 @testset "Construction" begin
@@ -114,8 +116,8 @@ end
 @testset "Utility functions" begin
     heq = HarmonicEquation([eq1, eq2], [hv1, hv2], nat_eq)
 
-    # Test _remove_brackets
-    no_brackets = _remove_brackets(heq)
+    # Test remove_brackets
+    no_brackets = remove_brackets(heq)
 
     list = get_all_terms.(Symbolics.expand_derivatives.(no_brackets))
     list = unique(filter(x -> !(x isa Real), Symbolics.unwrap.(reduce(vcat, list))))

test/HarmonicVariable.jl

@@ -7,10 +7,17 @@ using QuestBase:
     substitute_all,
     var_name,
     HarmonicVariable,
-    _show_ansatz,
+    HarmonicKind,
+    CosineQuadrature,
+    SineQuadrature,
+    ZeroHarmonic,
+    LimitCycleFrequency,
+    Unassigned,
+    kind_label,
+    show_ansatz,
     coordinate_transform,
     @eqtest,
-    _create_harmonic_variable,
+    create_harmonic_variable,
     get_variables_nums
 
 # Setup
@@ -20,53 +27,65 @@ rot_var = declare_variable("u", T)
 ω = declare_variable("ω")
 
 @testset "Constructor" begin
-    hv = HarmonicVariable(rot_var, "test_name", "u", ω, nat_var)
+    hv = HarmonicVariable(rot_var, "test_name", CosineQuadrature, ω, nat_var)
     @eqtest hv.symbol == rot_var
     @test hv.name == "test_name"
-    @test hv.type == "u"
+    @test hv.kind == CosineQuadrature
     @eqtest hv.ω == ω
     @eqtest hv.natural_variable == nat_var
 end
 
 @testset "Show and Display" begin
-    hv = HarmonicVariable(rot_var, "test_name", "u", ω, nat_var)
+    hv = HarmonicVariable(rot_var, "test_name", CosineQuadrature, ω, nat_var)
     # Test show output
     io = IOBuffer()
     show(io, hv)
     @test String(take!(io)) == "Harmonic variable u(T) for harmonic ω of x(t)\n"
 
-    # Test _show_ansatz
-    @test contains(_show_ansatz(hv), "cos(ωt)")
+    # Test show_ansatz
+    @test contains(show_ansatz(hv), "cos(ωt)")
 
     @variables x
     hv = HarmonicVariable(x)
     @test repr(hv) == "Harmonic variable x\n"
 
-    @test _show_ansatz(hv) == "x"
+    @test show_ansatz(hv) == "x"
 end
 
 @testset "Coordinate Transforms" begin
     new_var = declare_variable("y")
     # Test u-type transform
-    @eqtest coordinate_transform(new_var, ω, t, "u") == new_var * cos(ω * t)
+    @eqtest coordinate_transform(new_var, ω, t, CosineQuadrature) == new_var * cos(ω * t)
     # Test v-type transform
-    @eqtest coordinate_transform(new_var, ω, t, "v") == new_var * sin(ω * t)
+    @eqtest coordinate_transform(new_var, ω, t, SineQuadrature) == new_var * sin(ω * t)
     # Test a-type transform
-    @eqtest coordinate_transform(new_var, ω, t, "a") == new_var
+    @eqtest coordinate_transform(new_var, ω, t, ZeroHarmonic) == new_var
 end
 
 @testset "Create Harmonic Variable" begin
-    rule, hvar = _create_harmonic_variable(nat_var, ω, t, "u"; new_symbol="a")
+    rule, hvar = create_harmonic_variable(nat_var, ω, t, CosineQuadrature; new_symbol="a")
     a = declare_variable("a", t)
     @eqtest rule == a * cos(ω * t)
     @test hvar isa HarmonicVariable
     @test rule isa Num
-    @test hvar.type == "u"
+    @test hvar.kind == CosineQuadrature
     @eqtest hvar.ω == ω
 end
 
+@testset "HarmonicKind" begin
+    @test kind_label.(instances(HarmonicKind)) == ("u", "v", "a", "Hopf", "")
+    @test HarmonicVariable(declare_variable("z")).kind == Unassigned
+    @test show_ansatz(HarmonicVariable(rot_var, "s", SineQuadrature, ω, nat_var)) ==
+        "u(T)*sin(ωt)"
+    @test show_ansatz(HarmonicVariable(rot_var, "a", ZeroHarmonic, ω, nat_var)) == "u(T)"
+    @test_throws ArgumentError coordinate_transform(rot_var, ω, t, LimitCycleFrequency)
+
+    _, hvar = create_harmonic_variable(nat_var, ω, t, SineQuadrature; new_symbol="v1")
+    @test hvar.name == "v_{x,ω}"
+end
+
 @testset "Substitution" begin
-    hv = HarmonicVariable(nat_var, "test_name", "u", ω, nat_var)
+    hv = HarmonicVariable(nat_var, "test_name", CosineQuadrature, ω, nat_var)
     new_var = declare_variable("z")
     rules = Dict(hv => new_var)