Merge branch 'master' into x0_to_u0

docs/src/manual/time_dependent.md

@@ -2,10 +2,10 @@
 
 Generally, solving the ODE of oscillatory systems in time requires numerically tracking the oscillations. This is a computationally expensive process; however, using the harmonic ansatz removes the oscillatory time-dependence. Simulating instead the harmonic variables of a `HarmonicEquation` is vastly more efficient - a steady state of the system appears as a fixed point in multidimensional space rather than an oscillatory function.
 
-The Extention `TimeEvolution` is used to interface `HarmonicEquation` with the solvers contained in `OrdinaryDiffEq.jl`. Time-dependent parameter sweeps are defined using the object `ParameterSweep`. To use the `TimeEvolution` extension, one must first load the `OrdinaryDiffEq.jl` package.
+The Extention `TimeEvolution` is used to interface `HarmonicEquation` with the solvers contained in `OrdinaryDiffEq.jl`. Time-dependent parameter sweeps are defined using the object `AdiabaticSweep`. To use the `TimeEvolution` extension, one must first load the `OrdinaryDiffEq.jl` package.
 ```@docs
 ODEProblem(::HarmonicEquation, ::Any; timespan::Tuple)
-ParameterSweep
+AdiabaticSweep
 ```
 
 ## Plotting

docs/src/tutorials/limit_cycles.md

@@ -105,7 +105,7 @@ using OrdinaryDiffEqTsit5
 initial_state = result[1][1]
 
 T = 2e6
-sweep = ParameterSweep(F0 => (0.002, 0.011), (0,T))
+sweep = AdiabaticSweep(F0 => (0.002, 0.011), (0,T))
 
 # start from initial_state, use sweep, total time is 2*T
 time_problem = ODEProblem(harmonic_eq, initial_state, sweep=sweep, timespan=(0,2*T))

docs/src/tutorials/time_dependent.md

@@ -76,9 +76,9 @@ Clearly when evolving from `u0 = [0., 0.]`, the system ends up in the low-amplit
 
 Experimentally, the primary means of exploring the steady state landscape is an adiabatic sweep one or more of the system parameters. This takes the system along a solution branch. If this branch disappears or becomes unstable, a jump occurs.
 
-The object [`ParameterSweep`](@ref ParameterSweep) specifies a sweep, which is then used as an optional `sweep` keyword in the `ODEProblem` constructor.
+The object [`AdiabaticSweep`](@ref AdiabaticSweep) specifies a sweep, which is then used as an optional `sweep` keyword in the `ODEProblem` constructor.
 ```@example time_dependent
-sweep = ParameterSweep(ω => (0.9,1.1), (0, 2e4))
+sweep = AdiabaticSweep(ω => (0.9,1.1), (0, 2e4))
 ```
 The sweep linearly interpolates between $\omega = 0.9$ at time 0 and $\omega  = 1.1$ at time 2e4. For earlier/later times, $\omega$ is constant.
 

ext/TimeEvolution/TimeEvolution.jl

@@ -28,7 +28,7 @@ include("ODEProblem.jl")
 include("hysteresis_sweep.jl")
 include("plotting.jl")
 
-export ParameterSweep
+export AdiabaticSweep
 export transform_solutions, plot, plot!
 export ODEProblem, solve
 export follow_branch

ext/TimeEvolution/sweeps.jl

@@ -1,18 +1,18 @@
-using HarmonicBalance: ParameterSweep
+using HarmonicBalance: AdiabaticSweep
 
-function HarmonicBalance.ParameterSweep(functions::Dict, timespan::Tuple)
+function HarmonicBalance.AdiabaticSweep(functions::Dict, timespan::Tuple)
     t0, t1 = timespan[1], timespan[2]
     sweep_func = Dict{Num,Any}([])
     for swept_p in keys(functions)
         bounds = functions[swept_p]
         tfunc = swept_function(bounds, timespan)
         setindex!(sweep_func, tfunc, swept_p)
     end
-    return ParameterSweep(sweep_func)
+    return AdiabaticSweep(sweep_func)
 end
 
-function HarmonicBalance.ParameterSweep(functions, timespan::Tuple)
-    return ParameterSweep(Dict(functions), timespan)
+function HarmonicBalance.AdiabaticSweep(functions, timespan::Tuple)
+    return AdiabaticSweep(Dict(functions), timespan)
 end
 
 function swept_function(bounds, timespan)
@@ -29,15 +29,15 @@ function swept_function(bounds, timespan)
     return f
 end
 
-# overload so that ParameterSweep can be accessed like a Dict
-Base.keys(s::ParameterSweep) = keys(s.functions)
-Base.getindex(s::ParameterSweep, i) = getindex(s.functions, i)
+# overload so that AdiabaticSweep can be accessed like a Dict
+Base.keys(s::AdiabaticSweep) = keys(s.functions)
+Base.getindex(s::AdiabaticSweep, i) = getindex(s.functions, i)
 
 # overload +
-function Base.:+(s1::ParameterSweep, s2::ParameterSweep)
+function Base.:+(s1::AdiabaticSweep, s2::AdiabaticSweep)
     common_params = intersect(keys(s1), keys(s2))
     !isempty(common_params) && error("cannot combine sweeps of the same parameter")
-    return ParameterSweep(merge(s1.functions, s2.functions))
+    return AdiabaticSweep(merge(s1.functions, s2.functions))
 
     # combine sweeps of the same parameter
     #interval_overlap(s1.timespan, s2.timespan) && error("cannot combine sweeps with overlapping timespans")

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")
@@ -53,7 +61,7 @@ export rearrange_standard
 
 export plot, plot!, plot_phase_diagram, savefig, plot_spaghetti
 
-export ParameterSweep, steady_state_sweep
+export AdiabaticSweep, steady_state_sweep
 export plot_1D_solutions_branch, follow_branch
 
 include("HC_wrapper/HC_wrapper.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

src/types.jl

@@ -284,36 +284,36 @@ $(TYPEDFIELDS)
 ```julia-repl
 # create a sweep of parameter a from 0 to 1 over time 0 -> 100
 julia> @variables a,b;
-julia> sweep = ParameterSweep(a => [0., 1.], (0, 100));
+julia> sweep = AdiabaticSweep(a => [0., 1.], (0, 100));
 julia> sweep[a](50)
 0.5
 julia> sweep[a](200)
 1.0
 
 # do the same, varying two parameters simultaneously
-julia> sweep = ParameterSweep([a => [0.,1.], b => [0., 1.]], (0,100))
+julia> sweep = AdiabaticSweep([a => [0.,1.], b => [0., 1.]], (0,100))
 ```
 Successive sweeps can be combined,
 ```julia-repl
-sweep1 = ParameterSweep(ω => [0.95, 1.0], (0, 2e4))
-sweep2 = ParameterSweep(λ => [0.05, 0.01], (2e4, 4e4))
+sweep1 = AdiabaticSweep(ω => [0.95, 1.0], (0, 2e4))
+sweep2 = AdiabaticSweep(λ => [0.05, 0.01], (2e4, 4e4))
 sweep = sweep1 + sweep2
 ```
 multiple parameters can be swept simultaneously,
 ```julia-repl
-sweep = ParameterSweep([ω => [0.95;1.0], λ => [5e-2;1e-2]], (0, 2e4))
+sweep = AdiabaticSweep([ω => [0.95;1.0], λ => [5e-2;1e-2]], (0, 2e4))
 ```
 and custom sweep functions may be used.
 ```julia-repl
 ωfunc(t) = cos(t)
-sweep = ParameterSweep(ω => ωfunc)
+sweep = AdiabaticSweep(ω => ωfunc)
 ```
 
 """
-struct ParameterSweep
+struct AdiabaticSweep
     """Maps each swept parameter to a function."""
     functions::Dict{Num,Function}
 
-    ParameterSweep(functions...) = new(Dict(functions...))
-    ParameterSweep() = ParameterSweep([])
+    AdiabaticSweep(functions...) = new(Dict(functions...))
+    AdiabaticSweep() = AdiabaticSweep([])
 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