Fix missing docstrings

docs/Project.toml

@@ -1,7 +1,10 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 HarmonicBalance = "e13b9ff6-59c3-11ec-14b1-f3d2cc6c135e"
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 
 [compat]
 Documenter = "1"
-julia = "1"
+julia = "1.9"

docs/make.jl

@@ -2,10 +2,18 @@ push!(LOAD_PATH, "../src/")
 
 using Documenter
 using HarmonicBalance
+using ModelingToolkit
+using OrdinaryDiffEq
+using SteadyStateDiffEq
 
 makedocs(
 	sitename="HarmonicBalance.jl",
-    modules = [HarmonicBalance],
+    modules = [
+        HarmonicBalance,
+        Base.get_extension(HarmonicBalance, :TimeEvolution),
+        Base.get_extension(HarmonicBalance, :ModelingToolkitExt),
+        Base.get_extension(HarmonicBalance, :SteadyStateDiffEqExt)
+        ],
     warnonly = true,
 	format = Documenter.HTML(
 		mathengine=MathJax(),

docs/src/examples/time_dependent.md

@@ -39,15 +39,15 @@ Variables: u1(T), v1(T)
 
 The object `harmonic_eq` encodes Eq. \eqref{eq:harmeq}.
 
-We now wish to parse this input into [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/) and use its powerful ODE solvers. The desired object here is `DifferentialEquations.ODEProblem`, which is then fed into `DifferentialEquations.solve`.
+We now wish to parse this input into [OrdinaryDiffEq.jl](https://diffeq.sciml.ai/stable/) and use its powerful ODE solvers. The desired object here is `OrdinaryDiffEq.ODEProblem`, which is then fed into `OrdinaryDiffEq.solve`.
 
 ## Evolving from an initial condition
 
 Given $\mathbf{u}(T_0)$, what is $\mathbf{u}(T)$ at future times?
 
-For constant parameters, a [`HarmonicEquation`](@ref HarmonicBalance.HarmonicEquation) object can be fed into the constructor of [`ODEProblem`](@ref HarmonicBalance.TimeEvolution.ODEProblem). The syntax is similar to DifferentialEquations.jl :
+For constant parameters, a [`HarmonicEquation`](@ref HarmonicBalance.HarmonicEquation) object can be fed into the constructor of [`ODEProblem`](@ref ODEProblem). The syntax is similar to DifferentialEquations.jl :
 ```julia
-import HarmonicBalance.TimeEvolution: ODEProblem, OrdinaryDiffEq.solve
+using OrdinaryDiffEq
 x0 = [0.0; 0.] # initial condition
 fixed = (ω0 => 1.0,γ => 1E-2, λ => 5E-2, F => 1E-3,  α => 1., η=>0.3, θ => 0, ω=>1.) # parameter values
 
@@ -62,7 +62,7 @@ u0: 2-element Vector{Float64}:
  0.0
 ```
 
-DifferentialEquations.jl takes it from here - we only need to use `solve`.
+OrdinaryDiffEq.jl takes it from here - we only need to use `solve`.
 
 ```julia
 time_evo = solve(ode_problem, saveat=1.);
@@ -76,7 +76,7 @@ Running the above code with `x0 = [0., 0.]` and `x0 = [0.2, 0.2]` gives the plot
 
 Let us compare this to the steady state diagram.
 ```julia
-varied = ω => LinRange(0.9, 1.1, 100)
+varied = ω => range(0.9, 1.1, 100)
 result = get_steady_states(harmonic_eq, varied, fixed)
 plot(result, "sqrt(u1^2 + v1^2)")
 ```
@@ -90,7 +90,7 @@ Clearly when evolving from `x0 = [0.,0.]`, the system ends up in the low-amplitu
 
 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 HarmonicBalance.TimeEvolution.ParameterSweep) specifies a sweep, which is then used as an optional `sweep` keyword in the `ODEProblem` constructor.
+The object [`ParameterSweep`](@ref ParameterSweep) specifies a sweep, which is then used as an optional `sweep` keyword in the `ODEProblem` constructor.
 ```julia
 sweep = ParameterSweep(ω => (0.9,1.1), (0, 2E4))
 ```
@@ -201,7 +201,7 @@ According to Zambon et al., a limit cycle solution exists around $F_0 \cong 0.01
 
 Let us try and simulate the limit cycle. We could in principle run a time-dependent simulation with a fixed value of $F_0$, but this would require a suitable initial condition. Instead, we will sweep $F_0$ upwards from a low starting value. To observe the dynamics just after the jump has occurred, we follow the sweep by a time interval where the system evolves under fixed parameters.
 ```julia
-import HarmonicBalance.TimeEvolution: ODEProblem, OrdinaryDiffEq.solve
+using OrdinaryDiffEq
 initial_state = result[1][1]
 
 T = 2E6

docs/src/manual/time_dependent.md

@@ -2,22 +2,21 @@
 
 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 module `TimeEvolution` is used to interface `HarmonicEquation` with the powerful solvers contained in `DifferentialEquations.jl`. Time-dependent parameter sweeps are defined using the object `ParameterSweep`.
-
+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.
 ```@docs
-HarmonicBalance.TimeEvolution.ODEProblem
-HarmonicBalance.TimeEvolution.ParameterSweep
+ODEProblem(::HarmonicEquation, ::Any; timespan::Tuple)
+ParameterSweep
 ```
 
 ## Plotting
 
 ```@docs
-HarmonicBalance.TimeEvolution.plot(::HarmonicBalance.TimeEvolution.OrdinaryDiffEq.ODESolution, ::Any, ::HarmonicEquation)
+HarmonicBalance.plot(::OrdinaryDiffEq.ODESolution, ::Any, ::HarmonicEquation)
 ```
 
 ## Miscellaneous
 Using a time-dependent simulation can verify solution stability in cases where the Jacobian is too expensive to compute.
 
 ```@docs
-HarmonicBalance.TimeEvolution.is_stable
+HarmonicBalance.is_stable
 ```

src/plotting_Plots.jl

@@ -146,7 +146,7 @@ function plot1D(res::Result; x::String="default", y::String, class="default", no
     dim(res) != 1 && error("1D plots of not-1D datasets are usually a bad idea.")
     x = x == "default" ? string(first(keys(res.swept_parameters))) : x
     X = transform_solutions(res, x, branches=branches)
-    Y = transform_solutions(res, y, branches=branches)
+    Y = transform_solutions(res, y, branches=branches, realify=true)
     Y = _apply_mask(Y, _get_mask(res, class, not_class, branches=branches))
 
     # reformat and project onto real, warning if needed
@@ -170,7 +170,7 @@ plot1D(res::Result, y::String; kwargs...) = plot1D(res; y=y, kwargs...)
 
 function plot2D(res::Result; z::String, branch::Int64, class="physical", not_class=[], add=false, kwargs...)
     X, Y = values(res.swept_parameters)
-    Z = getindex.(_apply_mask(transform_solutions(res, z, branches=branch), _get_mask(res, class, not_class, branches=branch)), 1) # there is only one branch
+    Z = getindex.(_apply_mask(transform_solutions(res, z, branches=branch, realify=true), _get_mask(res, class, not_class, branches=branch)), 1) # there is only one branch
     p = add ? Plots.plot!() : Plots.plot() # start a new plot if needed
 
     ylab, xlab = latexify.(string.(keys(res.swept_parameters)))
@@ -202,7 +202,7 @@ function plot2D_cut(res::Result; y::String, cut::Pair, class="default", not_clas
     x = swept_pars[x_index]
 
     X = res.swept_parameters[x]
-    Y =_apply_mask(transform_solutions(res, y), _get_mask(res, class, not_class)) # first transform, then filter
+    Y =_apply_mask(transform_solutions(res, y, realify=true), _get_mask(res, class, not_class)) # first transform, then filter
     branches = x_index==1 ? Y[:, cut_par_index] : Y[cut_par_index, :]
 
     branch_data = [_realify( getindex.(branches, i), warning= "branch " * string(k) ) for (i,k) in enumerate(1:branch_count(res))]

src/transform_solutions.jl

@@ -1,27 +1,28 @@
 export transform_solutions
 
-
 _parse_expression(exp) = exp isa String ? Num(eval(Meta.parse(exp))) : exp
 
-
 """
 $(TYPEDSIGNATURES)
 
 Takes a `Result` object and a string `f` representing a Symbolics.jl expression.
 Returns an array with the values of `f` evaluated for the respective solutions.
 Additional substitution rules can be specified in `rules` in the format `("a" => val)` or `(a => val)`
 """
-function transform_solutions(res::Result, func; branches = 1:branch_count(res))
+function transform_solutions(res::Result, func; branches=1:branch_count(res), realify=false)
     # preallocate an array for the numerical values, rewrite parts of it
     # when looping through the solutions
     pars = res.swept_parameters |> values |> collect
     n_vars = length(get_variables(res))
     n_pars = length(pars)
 
-    vtype = isa(Base.invokelatest(func, rand(ComplexF64, n_vars+n_pars)), Bool) ? BitVector : Vector{ComplexF64}
+    vtype = isa(Base.invokelatest(func, rand(Float64, n_vars+n_pars)), Bool) ? BitVector : Vector{ComplexF64}
     transformed = _similar(vtype, res; branches=branches)
+    f = realify ? v -> real.(v) : identity
 
-    batches = Iterators.partition(CartesianIndices(res.solutions), ceil(Int, length(res.solutions)/Threads.nthreads()))
+    batches = Iterators.partition(
+        CartesianIndices(res.solutions),
+        ceil(Int, length(res.solutions)/Threads.nthreads()))
     Threads.@threads for batch in batches |> collect
         _vals = Vector{ComplexF64}(undef, n_vars + n_pars)
         for idx in batch
@@ -30,7 +31,7 @@ function transform_solutions(res::Result, func; branches = 1:branch_count(res))
             end
             for (k, branch) in enumerate(branches)
                 _vals[1:n_vars] .= res.solutions[idx][branch]
-                transformed[idx][k] = Base.invokelatest(func, _vals) # beware, func may be mutating
+                transformed[idx][k] = Base.invokelatest(func, f(_vals)) # beware, func may be mutating
             end
         end
     end

test/hysteresis_sweep.jl

@@ -15,7 +15,7 @@ followed_branches, _ = follow_branch(1, result)
 
 @test first(followed_branches) ≠ last(followed_branches)
 
-plot_1D_solutions_branch(1, result, x="ω", y="√(u1^2+v1^2)")
+plot_1D_solutions_branch(1, result, x="ω", y="u1^2+v1^2")
 
 plot_linear_response(result, x, followed_branches, Ω_range=range(0.9, 1.1, 10), show=false);
 

test/plotting.jl

@@ -16,6 +16,7 @@ res = get_steady_states(harmonic_eq, varied, fixed, seed=SEED)
 
 # plot 1D result
 plot(res, x="ω", y="u1");
+plot(res, y="√(u1^2+v1^2)");
 plot_spaghetti(res, x="v1", y="u1", z="ω");
 plot_phase_diagram(res);
 

test/transform_solutions.jl

@@ -1,5 +1,9 @@
 using HarmonicBalance
 
+using Random
+const SEED = 0xd8e5d8df
+Random.seed!(SEED)
+
 @variables Ω γ λ F x θ η α ω0 ω t T ψ
 @variables x(t)
 
@@ -15,6 +19,9 @@ varied = ω => range(0.9, 1.1, 10)
 res = get_steady_states(harmonic_eq, varied, fixed, show_progress=false, seed=SEED);
 
 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]