Refactoring ODE sweeps

.github/workflows/Documentation.yml

@@ -1,5 +1,5 @@
-
 name: Documenter
+
 on:
   push:
     branches: master

.github/workflows/test.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: ['1.8']
+        julia-version: ['1.9']
         julia-arch: [x64]
         os: [ubuntu-latest]
         exclude:

Project.toml

@@ -15,21 +15,23 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Peaks = "18e31ff7-3703-566c-8e60-38913d67486b"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
+[weakdeps]
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+
 [compat]
 BijectiveHilbert = "0.3.0"
 DSP = "0.7.4"
 Distances = "0.10.7"
 FFTW = "1.5.0"
 HomotopyContinuation = "2.6.4"
 JLD2 = "0.4.24"
-Latexify = "0.15.16"
 OrderedCollections = "1.4.1"
 OrdinaryDiffEq = "v6.33.1"
 Peaks = "0.4.0"
@@ -39,8 +41,14 @@ Symbolics = "5.0.0"
 julia = "1.8.2"
 
 [extras]
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Pkg", "Test"]
+
+[test]
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/HarmonicBalance.jl

@@ -8,6 +8,7 @@ module HarmonicBalance
     using BijectiveHilbert
     using LinearAlgebra
     using Plots, Latexify
+    using Requires
     import HomotopyContinuation
     import Distances
     # using SnoopPrecompile
@@ -17,27 +18,27 @@ module HarmonicBalance
     export @variables
     export d
 
-
     import Base: ComplexF64, Float64; export ComplexF64, Float64
     ComplexF64(x::Complex{Num}) = ComplexF64(Float64(x.re) + im*Float64(x.im))
     Float64(x::Complex{Num}) = Float64(ComplexF64(x))
     Float64(x::Num) = Float64(x.val)
 
-   # default global settings
-   export IM_TOL
-   IM_TOL::Float64 = 1E-6
-   function set_imaginary_tolerance(x::Float64)
-       @eval(IM_TOL::Float64 = $x)
-   end
+    # default global settings
+    export IM_TOL
+    IM_TOL::Float64 = 1E-6
+    function set_imaginary_tolerance(x::Float64)
+        @eval(IM_TOL::Float64 = $x)
+    end
 
-   export is_real
+    export is_real
     is_real(x) = abs(imag(x)) / abs(real(x)) < IM_TOL::Float64 || abs(x) < 1e-70
     is_real(x::Array) = is_real.(x)
 
     # Symbolics does not natively support complex exponentials of variables
     import Base: exp
     exp(x::Complex{Num}) = x.re.val == 0 ? exp(im*x.im.val) : exp(x.re.val + im*x.im.val)
 
+
     include("types.jl")
 
     include("utils.jl")
@@ -52,18 +53,13 @@ module HarmonicBalance
     include("saving.jl")
     include("transform_solutions.jl")
     include("plotting_Plots.jl")
-    include("hysteresis_sweep.jl")
 
     include("modules/HC_wrapper.jl")
     using .HC_wrapper
 
     include("modules/LinearResponse.jl")
     using .LinearResponse
 
-    include("modules/TimeEvolution.jl")
-    using .TimeEvolution
-    export ParameterSweep, ODEProblem, solve
-
     include("modules/LimitCycles.jl")
     using .LimitCycles
 
@@ -72,6 +68,14 @@ module HarmonicBalance
     export first_order_transform!, is_rearranged_standard, rearrange_standard!, get_equations
     export get_krylov_equations
 
+    function __init__()
+        @require OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" begin
+            include("modules/TimeEvolution.jl")
+            # using .TimeEvolution
+        end
+    end
+    # export ParameterSweep, ODEProblem, solve, follow_branch
+
     # precomp_path = (@__DIR__) * "/../test/"
     # @precompile_all_calls include(precomp_path * "parametron.jl")
     # @precompile_all_calls include(precomp_path * "plotting.jl")

src/modules/TimeEvolution.jl

@@ -1,8 +1,7 @@
-module TimeEvolution
-
-    using ..HarmonicBalance
+# module TimeEvolution
+    # using .OrdinaryDiffEq
+    # using ..HarmonicBalance; HB = HarmonicBalance
     using Symbolics
-    using OrdinaryDiffEq
     using DSP
     using FFTW
     using Peaks
@@ -12,5 +11,5 @@ module TimeEvolution
     include("TimeEvolution/ODEProblem.jl")
     include("TimeEvolution/FFT_analysis.jl")
     include("TimeEvolution/sweeps.jl")
-
-end
+    include("TimeEvolution/hysteresis_sweep.jl")
+# end

src/modules/TimeEvolution/ODEProblem.jl

@@ -1,6 +1,6 @@
 using LinearAlgebra, Latexify
 import HarmonicBalance: transform_solutions, plot, plot!
-import OrdinaryDiffEq: ODEProblem, solve
+# import OrdinaryDiffEq: ODEProblem, solve
 export transform_solutions, plot, plot!
 export ODEProblem, solve
 
@@ -18,7 +18,7 @@ Creates an ODEProblem object used by OrdinaryDiffEq.jl from the equations in `eo
 `fixed_parameters` must be a dictionary mapping parameters+variables to numbers (possible to use a solution index, e.g. solutions[x][y] for branch y of solution x).
 If `x0` is specified, it is used as an initial condition; otherwise the values from `fixed_parameters` are used.
 """
-function ODEProblem(eom::HarmonicEquation, fixed_parameters; sweep::ParameterSweep=ParameterSweep(), x0::Vector=[], timespan::Tuple, perturb_initial=0.0, kwargs...)
+function OrdinaryDiffEq.ODEProblem(eom::HarmonicEquation, fixed_parameters; sweep::ParameterSweep=ParameterSweep(), x0::Vector=[], timespan::Tuple, perturb_initial=0.0, kwargs...)
 
     if !is_rearranged(eom) # check if time-derivatives of the variable are on the right hand side
         eom = HarmonicBalance.rearrange_standard(eom)
@@ -105,4 +105,4 @@ function plot(soln::OrdinaryDiffEq.ODESolution, funcs, harm_eq::HarmonicEquation
 end
 
 
-plot!(soln::OrdinaryDiffEq.ODESolution, varargs...; kwargs...) = plot(soln, varargs...; add=true, kwargs...)
+plot!(soln::OrdinaryDiffEq.ODESolution, varargs...; kwargs...) = plot(soln, varargs...; add=true, kwargs...)

src/modules/TimeEvolution/hysteresis_sweep.jl

@@ -0,0 +1,80 @@
+export plot_1D_solutions_branch, follow_branch
+
+
+"""Calculate distance between a given state and a stable branch"""
+function _closest_branch_index(res::Result, state::Vector{Float64}, index::Int64)
+    #search only among stable solutions
+    stable = _apply_mask(res.solutions, _get_mask(res, ["physical", "stable"], []))
+
+    steadystates = reduce(hcat, stable[index])
+    distances    = vec(sum(abs2.(steadystates .- state), dims=1))
+    return argmin(replace(distances, NaN => Inf))
+end
+
+"""Return the indexes and values following stable branches along a 1D sweep.
+When a no stable solutions are found (e.g. in a bifurcation), the next stable solution is calculated by time evolving the previous solution (quench).
+
+Keyword arguments
+- `y`:  Dependent variable expression (parsed into Symbolics.jl) to evaluate the followed solution branches on .
+- `sweep`: Direction for the sweeping of solutions. A `right` (`left`) sweep proceeds from the first (last) solution, ordered as the sweeping parameter.
+- `tf`: time to reach steady
+- `ϵ`: small random perturbation applied to quenched solution, in a bifurcation in order to favour convergence in cases where multiple solutions are identically accessible (e.g. symmetry breaking into two equal amplitude states)
+"""
+function follow_branch(starting_branch::Int64, res::Result; y="u1^2+v1^2", sweep="right", tf=10000, ϵ=1e-4)
+    sweep_directions = ["left", "right"]
+    sweep ∈ sweep_directions  || error("Only the following (1D) sweeping directions are allowed:  ", sweep_directions)
+
+    # get stable solutions
+    Y  = transform_solutions(res, y)
+    Ys = _apply_mask(Y, _get_mask(res, ["physical", "stable"], []))
+    Ys = sweep == "left" ? reverse(Ys) : Ys
+
+    followed_branch    = zeros(Int64,length(Y))  # followed branch indexes
+    followed_branch[1] = starting_branch
+
+    p1 = first(keys(res.swept_parameters)) # parameter values
+
+    for i in 2:length(Ys)
+        s = Ys[i][followed_branch[i-1]] # solution amplitude in the current branch and current parameter index
+        if !isnan(s) # the solution is not unstable or unphysical
+            followed_branch[i] = followed_branch[i-1]
+        else # bifurcation found
+            next_index = sweep == "right" ? i : length(Y)-i+1
+
+            # create a synthetic starting point out of an unphysical solution: quench and time evolve
+            # the actual solution is complex there, i.e. non physical. Take real part for the quench.
+            sol_dict  = get_single_solution(res, branch=followed_branch[i-1], index=next_index)
+            print("bifurcation @ ", p1 ," = ", real(sol_dict[p1])," ")
+
+            values_noise = real.(values(sol_dict)) .+ 0.0im .+ ϵ*rand(length(values(sol_dict)))
+            sol_dict_noise  = Dict(zip(keys(sol_dict), values_noise))
+
+            problem_t = OrdinaryDiffEq.ODEProblem(res.problem.eom, sol_dict_noise, timespan=(0, tf))
+            res_t     = OrdinaryDiffEq.solve(problem_t, OrdinaryDiffEq.Tsit5(), saveat=tf)
+
+            followed_branch[i] = _closest_branch_index(res, res_t.u[end], next_index) # closest branch to final state
+
+            print("switched branch ", followed_branch[i-1] ," -> ", followed_branch[i],"\n")
+        end
+    end
+    if sweep == "left"
+        Ys = reverse(Ys)
+        followed_branch = reverse(followed_branch)
+    end
+
+    return followed_branch, Ys
+end
+
+
+"""1D plot with the followed branch highlighted"""
+function plot_1D_solutions_branch(starting_branch::Int64, res::Result;
+    x::String, y::String, sweep="right", tf=10000, ϵ=1e-4, class="default", not_class=[], kwargs...)
+    p = plot(res; x=x, y=y, class=class, not_class=not_class, kwargs...)
+
+    followed_branch, Ys = follow_branch(starting_branch, res, y=y, sweep=sweep, tf=tf, ϵ=ϵ)
+    Y_followed = [Ys[param_idx][branch] for (param_idx,branch) in enumerate(followed_branch)]
+    X = res.swept_parameters[_parse_expression(x)]
+
+    plot!(p, X, real.(Y_followed); linestyle=:dash, c=:gray, label = sweep*" sweep", _set_Plots_default..., kwargs...)
+    return p
+end

src/plotting_Plots.jl

@@ -94,6 +94,15 @@ function _apply_mask(solns::Array{Vector{ComplexF64}},  booleans)
     factors = replace.(booleans, 0 => NaN)
     map(.*, solns, factors)
 end
+function _apply_mask(solns::Array{Vector{Vector{ComplexF64}}}, booleans)
+    new_solns = deepcopy(solns)
+    for i in 1:length(solns)
+        for j in 1:length(solns[i])
+            new_solns[i][j] = booleans[i][j] ? solns[i][j] : [NaN for _ in 1:length(solns[i][j])]
+        end
+    end
+    solns
+end
 
 
 """ Project the array `a` into the real axis, warning if its contents are complex. """

src/utils.jl

@@ -9,4 +9,4 @@ macro maybethread(loop)
       # @warn "running single threaded"
       quote $(esc(loop)); end
     end
-end
+end