Skip to content

Commit

Permalink
style: Better docstrings (#326)
Browse files Browse the repository at this point in the history
  • Loading branch information
@oameye
oameye authored Nov 20, 2024
1 parent a99a6ca commit 58bffe7
Show file tree
Hide file tree
Showing 20 changed files with 370 additions and 222 deletions.
39 changes: 24 additions & 15 deletions src/DifferentialEquation.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@
$(TYPEDEF)
Holds differential equation(s) of motion and a set of harmonics to expand each variable.
This is the primary input for `HarmonicBalance.jl` ; after inputting the equations, the harmonics
ansatz needs to be specified using `add_harmonic!`.
This is the primary input for `HarmonicBalance.jl`. After inputting the equations, the harmonics
ansatz needs to be specified using `add_harmonic!`.
# Fields
$(TYPEDFIELDS)
Expand All @@ -17,7 +17,9 @@ julia> DifferentialEquation(d(x,t,2) + ω0^2 * x - F * cos(ω*t), x);
julia> DifferentialEquation(d(x,t,2) + ω0^2 * x ~ F * cos(ω*t), x);
# two coupled oscillators, one of them driven
julia> DifferentialEquation([d(x,t,2) + ω0^2 * x - k*y, d(y,t,2) + ω0^2 * y - k*x] .~ [F * cos(ω*t), 0], [x,y]);
julia> DifferentialEquation(
[d(x,t,2) + ω0^2 * x - k*y, d(y,t,2) + ω0^2 * y - k*x] .~ [F * cos(ω*t), 0], [x,y]
);
```
"""
mutable struct DifferentialEquation
Expand Down Expand Up @@ -46,7 +48,8 @@ mutable struct DifferentialEquation
if eq.rhs isa AbstractVector || eq.lhs isa AbstractVector
throw(
ArgumentError(
"The equation is of the form $(typerhs)~$(typelhs) is not supported. Commenly one forgot to broadcast the equation symbol `~`.",
"The equation is of the form $(typerhs)~$(typelhs) is not supported.
Commenly one forgot to broadcast the equation symbol `~`."
),
)
end
Expand All @@ -57,13 +60,15 @@ mutable struct DifferentialEquation
typelhs = typeof(eq.lhs)
throw(
ArgumentError(
"The variables are of type $(typeof(vars)). Whereas you gave one equation of type $(typerhs)~$(typelhs). Commenly one forgot to broadcast the equation symbol `~`.",
"The variables are of type $(typeof(vars)). Whereas you gave one equation of
type $(typerhs)~$(typelhs). Commenly one forgot to broadcast the equation symbol `~`.",
),
)
end
DifferentialEquation(lhs::Num, var::Num) = DifferentialEquation([lhs ~ Int(0)], [var])
end

"show method of the type `DifferentialEquation`"
function Base.show(io::IO, diff_eq::DifferentialEquation)
println(io, "System of ", length(keys(diff_eq.equations)), " differential equations")
println(io, "Variables: ", join(keys(diff_eq.equations), ", "))
Expand All @@ -75,9 +80,14 @@ function Base.show(io::IO, diff_eq::DifferentialEquation)
println(io, "\n")
return [println(io, eq) for eq in values(diff_eq.equations)]
end
"
Displays the fields of the differential equation object.
"
Base.show(eom::DifferentialEquation) = show_fields(eom)

"""
$(TYPEDSIGNATURES)
Add the harmonic `ω` to the harmonic ansatz used to expand the variable `var` in `diff_eom`.
## Example
Expand All @@ -102,28 +112,27 @@ end

"""
$(TYPEDSIGNATURES)
Return the dependent variables of `diff_eom`.
"""
Symbolics.get_variables(diff_eom::DifferentialEquation)::Vector{Num} =
collect(keys(diff_eom.equations))

"""
$(TYPEDSIGNATURES)
Check if all equations in `diff_eom` are harmonic with respect to `t`. The function takes a
differential equation system `diff_eom` and a variable `t`, and returns `true` if all equations
are harmonic with respect to `t`, otherwise it returns `false`.
"""
ExprUtils.is_harmonic(diff_eom::DifferentialEquation, t::Num)::Bool =
all([is_harmonic(eq, t) for eq in values(diff_eom.equations)])

"Pretty printing of the newly defined types"
function show_fields(object)
for field in fieldnames(typeof(object)) # display every field
display(string(field))
display(getfield(object, field))
end
end

"""
$(TYPEDSIGNATURES)
Return the independent dependent variables of `diff_eom`.
"""
function get_independent_variables(diff_eom::DifferentialEquation)
return Num.(flatten(unique([x.val.arguments for x in keys(diff_eom.equations)])))
end

Base.show(eom::DifferentialEquation) = show_fields(eom)
30 changes: 0 additions & 30 deletions src/FFTWExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,36 +24,6 @@ function FFT(soln_u, soln_t; window=DSP.Windows.hanning)
return (fft_u / length(fft_f), 2 * pi * fft_f)
end

# function HarmonicBalance.FFT(soln::OrdinaryDiffEqTsit5.ODESolution; window=DSP.Windows.hanning)
# return HarmonicBalance.FFT(soln.u, soln.t; window=window)
# end

# function FFT_analyze(fft_u::Vector{ComplexF64}, fft_f)
# "finds peaks in the spectrum and returns corresponding frequency, amplitude and phase.
# Frequency and phase are corrected according to Huang Dishan, Mechanical Systems and Signal Processing (1995) 9(2), 113–118
# This correction works for a rectangular window."

# # retaining more sigdigits gives more ''spurious'' peaks
# max_indices, mxval = Peaks.peakproms(round.(abs.(fft_u), sigdigits=3); minprom=1)
# Del = fft_f[2] - fft_f[1] # frequency spacing
# A1 = abs.(fft_u)[max_indices]
# df = zeros(length(max_indices))

# # correction to the amplitude and phase of the peak
# for i in 1:length(max_indices)
# if abs.(fft_u)[max_indices[i] - 1] < abs.(fft_u)[max_indices[i] + 1]
# A2 = abs.(fft_u)[max_indices[i] + 1]
# df[i] = -Del / (A1[i] / A2 + 1)
# else
# A2 = abs.(fft_u)[max_indices[i] - 1]
# df[i] = Del / (A1[i] / A2 + 1)
# end
# end
# return 2 * pi * (fft_f[max_indices] - df),
# A1 .* 2,
# angle.(fft_u)[max_indices] + pi * df / Del
# end

function u_of_t(omegas_peak, As_peak, phis_peak, t)
"Calculate us or vs as a function of time from the Fourier components."
N = length(omegas_peak)
Expand Down
16 changes: 0 additions & 16 deletions src/HC_wrapper.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,22 +32,6 @@ function parse_equations(eqs::Vector{Num})
return [Expression(eval(symbol)) for symbol in parsed_strings]
end

"A constructor for Problem from explicitly entered equations, variables and parameters."
function HarmonicBalance.Problem(
equations::Vector{Num}, variables::Vector{Num}, parameters::Vector{Num}
)
conv_vars = Num_to_Variable.(variables)
conv_para = Num_to_Variable.(parameters)

eqs_HC = [
Expression(eval(symbol)) for
symbol in [Meta.parse(s) for s in [string(eq) for eq in equations]]
] #note in polar coordinates there could be imaginary factors, requiring the extra replacement "I"=>"1im"
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 # TODO is this function still needed/used?

function System(eom::HarmonicEquation)
eqs = expand_derivatives.(_remove_brackets(eom))
conv_vars = Num_to_Variable.(get_variables(eom))
Expand Down
91 changes: 71 additions & 20 deletions src/HarmonicBalance.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ using BijectiveHilbert: BijectiveHilbert, Simple2D, decode_hilbert!, encode_hilb
using HomotopyContinuation: HomotopyContinuation
const HC = HomotopyContinuation

using Plots: Plots, plot, plot!, savefig, heatmap, Plot
using Plots: Plots, plot, plot!, heatmap, Plot
using Latexify: Latexify, latexify

using Symbolics:
Expand All @@ -38,6 +38,60 @@ using SymbolicUtils: SymbolicUtils
include("ExprUtils/ExprUtils.jl")
using .ExprUtils: is_harmonic, substitute_all, drop_powers, count_derivatives, is_identity

# symbolics equations
export @variables, d
export DifferentialEquation
export HarmonicVariable
export HarmonicEquation

export rearrange_standard
export rearrange_standard!
export first_order_transform!
export is_rearranged_standard
export get_equations

export get_harmonic_equations
export get_krylov_equations
export add_harmonic!

export get_independent_variables
export get_variables

# handle solutions
export get_steady_states
export classify_solutions!
export get_class
export get_single_solution
export transform_solutions
export IM_TOL

# methods
export WarmUp
export TotalDegree
export Polyhedral

# Limit cycles
export get_cycle_variables, get_limit_cycles, add_pairs!

# LinearResponse
export get_Jacobian

# plotting
export plot_linear_response
export plot_phase_diagram
export plot_rotframe_jacobian_response
export plot_eigenvalues
export plot, plot!
export plot_spaghetti

# extension functions
export AdiabaticSweep
export steady_state_sweep
export plot_1D_solutions_branch
export follow_branch

# src code

include("extension_functions.jl")
include("utils.jl")
include("types.jl")
Expand All @@ -57,41 +111,23 @@ include("saving.jl")
include("transform_solutions.jl")
include("plotting_Plots.jl")

export show, *, @variables, d, IM_TOL
export WarmUp, TotalDegree, Polyhedral

export DifferentialEquation, HarmonicVariable, HarmonicEquation
export get_steady_states, get_single_solution, get_harmonic_equations, add_harmonic!
export get_variables, get_independent_variables, classify_branch, classify_solutions!
export rearrange_standard

export plot, plot!, plot_phase_diagram, savefig, plot_spaghetti

export AdiabaticSweep, steady_state_sweep
export plot_1D_solutions_branch, follow_branch

include("HC_wrapper.jl")
using .HC_wrapper

include("LinearResponse/LinearResponse.jl")
using .LinearResponse
export plot_linear_response, plot_rotframe_jacobian_response, get_Jacobian, plot_eigenvalues
export transform_solutions

include("LimitCycles/LimitCycles.jl")
using .LimitCycles
export get_cycle_variables, get_limit_cycles, add_pairs!

include("KrylovBogoliubov/KrylovBogoliubov.jl")
using .KrylovBogoliubov
export first_order_transform!, is_rearranged_standard, rearrange_standard!, get_equations
export get_krylov_equations

include("FFTWExt.jl")
using .FFTWExt

# Precompilation setup
using PrecompileTools: @setup_workload, @compile_workload

@setup_workload begin
# Putting some things in `@setup_workload` instead of `@compile_workload` can reduce the size of the
# precompile file and potentially make loading faster.
Expand All @@ -102,4 +138,19 @@ using PrecompileTools: @setup_workload, @compile_workload
end
end

# Error hint for extensions stubs
function __init__()
Base.Experimental.register_error_hint(
_error_hinter("OrdinaryDiffEq", :TimeEvolution, follow_branch), MethodError
)
Base.Experimental.register_error_hint(
_error_hinter("OrdinaryDiffEq", :TimeEvolution, plot_1D_solutions_branch),
MethodError,
)
return Base.Experimental.register_error_hint(
_error_hinter("SteadyStateDiffEq", :SteadyStateDiffEqExt, steady_state_sweep),
MethodError,
)
end

end # module
35 changes: 18 additions & 17 deletions src/HarmonicEquation.jl
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@

"""
$(TYPEDEF)
Expand Down Expand Up @@ -57,9 +56,9 @@ Base.show(eom::HarmonicEquation) = show_fields(eom)
"""
harmonic_ansatz(eom::DifferentialEquation, time::Num; coordinates="Cartesian")
Expand each variable of `diff_eom` using the harmonics assigned to it with `time` as the time variable.
For each harmonic of each variable, instance(s) of `HarmonicVariable` are automatically created and named.
Expand each variable of `diff_eom` using the harmonics assigned to it with `time` as the
time variable. For each harmonic of each variable, instance(s) of `HarmonicVariable` are
automatically created and named.
"""
function harmonic_ansatz(diff_eom::DifferentialEquation, time::Num)
!is_harmonic(diff_eom, time) &&
Expand Down Expand Up @@ -203,9 +202,7 @@ function _parameters(eom::HarmonicEquation)
return setdiff(all_symbols, get_variables(eom), get_independent_variables(eom))
end

###
# Extending Symbolics.jl's simplify and substitute
###
### Extending Symbolics.jl's simplify and substitute ###

"Apply `rules` to both `equations` and `variables` field of `eom`"
function ExprUtils.substitute_all(
Expand Down Expand Up @@ -273,15 +270,15 @@ end
get_harmonic_equations(diff_eom::DifferentialEquation; fast_time=nothing, slow_time=nothing)
Apply the harmonic ansatz, followed by the slow-flow, Fourier transform and dropping
higher-order derivatives to obtain
a set of ODEs (the harmonic equations) governing the harmonics of `diff_eom`.
higher-order derivatives to obtain a set of ODEs (the harmonic equations) governing the
harmonics of `diff_eom`.
The harmonics evolve in `slow_time`, the oscillating terms themselves in `fast_time`.
If no input is used, a variable T is defined for `slow_time` and `fast_time` is taken as the independent variable
of `diff_eom`.
If no input is used, a variable T is defined for `slow_time` and `fast_time` is taken as
the independent variable of `diff_eom`.
By default, all products of order > 1 of `slow_time`-derivatives are dropped,
which means the equations are linear in the time-derivatives.
By default, all products of order > 1 of `slow_time`-derivatives are dropped, which means
the equations are linear in the time-derivatives.
# Example
```julia-repl
Expand Down Expand Up @@ -320,10 +317,14 @@ function get_harmonic_equations(
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
ft_eom_simplified = drop_powers(eom, d(get_variables(eom), slow_time), 2) # drop higher powers of the first-order derivatives
# substitute trig functions into the differential equation
eom = harmonic_ansatz(diff_eom, fast_time)
# drop 2nd order time derivatives
eom = slow_flow(eom; fast_time=fast_time, slow_time=slow_time, degree=degree)
# perform averaging over the frequencies originally specified in dEOM
fourier_transform!(eom, fast_time)
# drop higher powers of the first-order derivatives
ft_eom_simplified = drop_powers(eom, d(get_variables(eom), slow_time), 2)
return ft_eom_simplified
end

Expand Down
Loading

0 comments on commit 58bffe7

Please sign in to comment.