plotting branches, optimized sorting, deps removed (#88)

* align_pair ducktyped to simplify 1D sorting

* get_distance_matrix optimized for 2D

* Julia rolled back to 1.8.2

* fix TD plot bug

* Plots recipes (#87)

* simplified Base MD

* branches kw for transform_solutions

* transform_solutions refactored

* branches keyword for plotting

* indexing fix, version bump

* bugfixes for arrays

* realify debugged

* plotting tests to `sorting_optimization` branch (#93)

* fixed Spaghetti_plot and 2D_cut bug

* adding plotting tests

Co-authored-by: Orjan Ameye <[email protected]>

Author: jkosata

Project.toml

@@ -1,7 +1,7 @@
 name = "HarmonicBalance"
 uuid = "e13b9ff6-59c3-11ec-14b1-f3d2cc6c135e"
 authors = ["Jan Kosata <[email protected]>", "Javier del Pino <[email protected]>"]
-version = "0.6.3"
+version = "0.6.4"
 
 [deps]
 BijectiveHilbert = "91e7fc40-53cd-4118-bd19-d7fcd1de2a54"
@@ -37,7 +37,7 @@ Peaks = "0.4.1"
 Plots = "1.36.4"
 ProgressMeter = "1.7.2"
 Symbolics = "4.13.0"
-julia = "1.8.3"
+julia = "1.8.2"
 
 [extras]
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

src/HarmonicBalance.jl

@@ -26,7 +26,7 @@ module HarmonicBalance
 
    export is_real
     is_real(x) = abs(imag(x)) / abs(real(x)) < IM_TOL::Float64 || abs(x) < 1e-70
-    is_real(x::Array) = any(is_real.(x))
+    is_real(x::Array) = is_real.(x)
 
     # Symbolics does not natively support complex exponentials of variables
     import Base: exp

src/HarmonicEquation.jl

@@ -134,6 +134,7 @@ end
 
 
 get_variables(p::Problem) = get_variables(p.eom)
+get_variables(res::Result) = get_variables(res.problem)
 
 
 "Get the parameters (not time nor variables) of a HarmonicEquation"

src/HarmonicVariable.jl

@@ -1,5 +1,4 @@
 import Symbolics: get_variables; export get_variables
-import Base.isequal; export isequal
 
 # pretty-printing
 display(var::HarmonicVariable) = display(var.name)
@@ -46,7 +45,7 @@ get_variables(vars::Vector{Num}) = unique(flatten([Num.(get_variables(x)) for x
 
 get_variables(var::HarmonicVariable) = Num.(get_variables(var.symbol))
 
-isequal(v1::HarmonicVariable, v2::HarmonicVariable) = isequal(v1.symbol, v2.symbol)
+Base.isequal(v1::HarmonicVariable, v2::HarmonicVariable) = isequal(v1.symbol, v2.symbol)
 
 
 

src/modules/LinearResponse.jl

@@ -9,7 +9,7 @@ module LinearResponse
     using ..HC_wrapper
     using DocStringExtensions
 
-    import Base: *, show; export *, show
+    import Base: show; export show
 
     include("LinearResponse/types.jl")
     include("LinearResponse/utils.jl")

src/modules/LinearResponse/Lorentzian_spectrum.jl

@@ -2,12 +2,12 @@
 Here the methods to find a
 """
 # multiply a peak by a number.
- function *(number::Float64, peak::Lorentzian) # multiplication operation
+ function Base.:*(number::Float64, peak::Lorentzian) # multiplication operation
     Lorentzian(peak.ω0, peak.Γ, peak.A*number)
  end
 
 
- *(number::Float64, s::JacobianSpectrum) = JacobianSpectrum([number * peak for peak in s.peaks])
+ Base.:*(number::Float64, s::JacobianSpectrum) = JacobianSpectrum([number * peak for peak in s.peaks])
 
 
 function show(io::IO, s::JacobianSpectrum)

src/modules/TimeEvolution/ODEProblem.jl

@@ -68,7 +68,7 @@ end
 
 
 transform_solutions(soln::OrdinaryDiffEq.ODESolution, f::String, harm_eq::HarmonicEquation) = transform_solutions(soln.u, f, harm_eq)
-transform_solutions(s::OrdinaryDiffEq.ODESolution, funcs::Vector{String}, harm_eq::HarmonicEquation) = [transform_solutions(s, f, he) for f in funcs]
+transform_solutions(s::OrdinaryDiffEq.ODESolution, funcs::Vector{String}, harm_eq::HarmonicEquation) = [transform_solutions(s, f, harm_eq) for f in funcs]
 
 
 

src/modules/TimeEvolution/types.jl

@@ -1,6 +1,3 @@
-import Base: keys, getindex, +
-export +, getindex
-
 """
 
 Represents a sweep of one or more parameters of a `HarmonicEquation`.
@@ -36,12 +33,12 @@ end
 
 
 # overload so that ParameterSweep can be accessed like a Dict
-keys(s::ParameterSweep) = keys(s.functions)
-getindex(s::ParameterSweep, i) = getindex(s.functions, i)
+Base.keys(s::ParameterSweep) = keys(s.functions)
+Base.getindex(s::ParameterSweep, i) = getindex(s.functions, i)
 
 
 # overload +
-function +(s1::ParameterSweep, s2::ParameterSweep)
+function Base.:+(s1::ParameterSweep, s2::ParameterSweep)
     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))

src/plotting_Plots.jl

@@ -74,12 +74,15 @@ end
 """
 $(TYPEDSIGNATURES)
 
-Return an array of bools to mark solutions in `res` which fall into `classes` but not `not_classes`
+Return an array of bools to mark solutions in `res` which fall into `classes` but not `not_classes`.
+Only `branches` are considered.
 """
-function _get_mask(res, classes, not_classes=[])
-    classes == "all" && return fill(trues(length(res.solutions[1])), size(res.solutions))
-    bools = vcat([res.classes[c] for c in _vectorise(classes)], [map(.!, res.classes[c]) for c in _vectorise(not_classes)])
-    map(.*, bools...)
+function _get_mask(res, classes, not_classes=[]; branches=1:branch_count(res))
+    classes == "all" && return fill(trues(length(branches)), size(res.solutions))
+    bools = vcat([res.classes[c] for c in _str_to_vec(classes)], [map(.!, res.classes[c]) for c in _str_to_vec(not_classes)])
+    #m = map( x -> [getindex(x, b) for b in [branches...]], map(.*, bools...))
+
+    m = map( x -> x[[branches...]], map(.*, bools...))
 end
 
 
@@ -96,85 +99,62 @@ function _apply_mask(solns::Array{Vector{ComplexF64}},  booleans)
 end
 
 
-# convert x to Float64, raising a warning if complex
-function _realify(x, warn=true)
-    !is_real(x) && !isnan(x) && warn ? (@warn "Values with non-negligible complex parts have been projected on the real axis!", x) : nothing
-    real(x)
-end
+""" Project the array `a` into the real axis, warning if its contents are complex. """
+function _realify(a::Array{T} where T <: Number; warning="")
 
-function _realify(v::Vector, warn=true)
-    v_real = Vector{Vector{Float64}}(undef, length(v[1]))
-    for branch in eachindex(v_real)
-        values = getindex.(v, branch)
-        values_real = Vector{Float64}(undef, length(values))
-        for (j,x) in pairs(values)
-            if !is_real(x) && !isnan(x) && warn
-                (@warn "Values with non-negligible complex parts have been projected on the real axis!", x)
-                warn = false
-            end
-            values_real[j] = real(x)
+    warned = false
+    a_real = similar(a, Float64)
+    for i in eachindex(a)
+        if !isnan(a[i]) && !warned && !is_real(a[i])
+            @warn "Values with non-negligible complex parts have
+            been projected on the real axis! " * warning
+            warned = true
         end
-        v_real[branch] = values_real
-    end
-    return v_real
-end
-
-function _realify(a::Matrix, warn=true)
-    a_real = Array{Vector{Vector{Float64}}}(undef, size(a)...)
-    for idx in CartesianIndices(a)
-        v = a[idx]
-        v_real = Vector{Vector{Float64}}(undef, length(v[1]))
-        for branch in eachindex(v_real)
-            values = getindex.(v, branch)
-            values_real = Vector{Float64}(undef, length(values))
-            for (j,x) in pairs(values)
-                if !is_real(x) && !isnan(x) && warn
-                    (@warn "Values with non-negligible complex parts have been projected on the real axis!", x)
-                    warn = false
-                end
-                values_real[j] = real(x)
-            end
-            v_real[branch] = values_real
-        end
-        a_real[idx] = v_real
+        a_real[i] = real(a[i])
     end
     return a_real
 end
 
-_vectorise(s::Vector) = s
-_vectorise(s) = [s]
+
+_str_to_vec(s::Vector) = s
+_str_to_vec(s) = [s]
 
 
 # return true if p already has a label for branch index idx
 _is_labeled(p::Plots.Plot, idx::Int64) = in(string(idx), [sub[:label] for sub in p.series_list])
 
 
-function plot1D(res::Result; x::String="default", y::String, class="default", not_class=[], add=false, kwargs...)
+function plot1D(res::Result; x::String="default", y::String, class="default", not_class=[], branches=1:branch_count(res), add=false, kwargs...)
 
     if class == "default"
+        args = [:x => x, :y => y, :branches => branches]
         if not_class == [] # plot stable full, unstable dashed
-            p = plot1D(res; x=x, y=y, class=["physical", "stable"], add=add, kwargs...)
-            plot1D(res; x=x, y=y, class="physical", not_class="stable", add=true, style=:dash, kwargs...)
+            p = plot1D(res; args..., class=["physical", "stable"], add=add, kwargs...)
+            plot1D(res; args..., class="physical", not_class="stable", add=true, style=:dash, kwargs...)
             return p
         else
-            p = plot1D(res; x=x, y=y, not_class=not_class, class="physical", add=add, kwargs...)
+            p = plot1D(res; args..., not_class=not_class, class="physical", add=add, kwargs...)
             return p
         end
     end
 
     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, Y = transform_solutions(res, [x,y]) # first transform, then filter
-    Y = _apply_mask(Y, _get_mask(res, class, not_class))
-    branches = _realify(Y)
+    X = transform_solutions(res, x, branches=branches)
+    Y = transform_solutions(res, y, branches=branches)
+    Y = _apply_mask(Y, _get_mask(res, class, not_class, branches=branches))
+
+    # reformat and project onto real, warning if needed
+    branch_data = [_realify( getindex.(Y, i), warning= "branch " * string(k) ) for (i,k) in enumerate(branches)]
 
     # start a new plot if needed
     p = add ? Plots.plot!() : Plots.plot()
 
     # colouring is matched to branch index - matched across plots
-    for k in findall(x -> !all(isnan.(x)), branches[1:end]) # skip NaN branches but keep indices
-        l = _is_labeled(p, k) ? nothing : k
-        Plots.plot!(_realify.(getindex.(X, k)), branches[k];  color=k, label=l, xlabel=latexify(x), ylabel=latexify(y), kwargs...)
+    for k in findall(x -> !all(isnan.(x)), branch_data) # skip NaN branch_data
+        global_index = branches[k]
+        lab = _is_labeled(p, global_index) ? nothing : global_index
+        Plots.plot!(_realify(getindex.(X, k)), branch_data[k];  color=k, label=lab, xlabel=latexify(x), ylabel=latexify(y), kwargs...)
     end
 
     return p
@@ -185,8 +165,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), _get_mask(res, class, not_class)), branch) # first transform, then filter
-
+    Z = getindex.(_apply_mask(transform_solutions(res, z, branches=branch), _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)))
@@ -219,15 +198,17 @@ function plot2D_cut(res::Result; y::String, cut::Pair, class="default", not_clas
 
     X = res.swept_parameters[x]
     Y =_apply_mask(transform_solutions(res, y), _get_mask(res, class, not_class)) # first transform, then filter
-    branches = _realify(x_index==1 ? Y[:, cut_par_index] : Y[cut_par_index, :])
+    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))]
 
     # start a new plot if needed
     p = add ? Plots.plot!() : Plots.plot()
 
     # colouring is matched to branch index - matched across plots
-    for k in findall(branch -> !all(isnan.(branch)), branches[1:end]) # skip NaN branches but keep indices
+    for k in findall(branch -> !all(isnan.(branch)), branch_data) # skip NaN branches but keep indices
         l = _is_labeled(p, k) ? nothing : k
-        Plots.plot!(X, branches[k]; color=k, label=l, xlabel=latexify(string(x)), ylabel=latexify(y), kwargs...)
+        Plots.plot!(X, _realify(getindex.(branches, k)); color=k, label=l, xlabel=latexify(string(x)), ylabel=latexify(y), kwargs...)
     end
 
     return p
@@ -302,6 +283,9 @@ Class selection done by passing `String` or `Vector{String}` as kwarg:
 Other kwargs are passed onto Plots.gr()
 """
 function plot_spaghetti(res::Result; x::String, y::String, z::String, class="default", not_class=[], add=false, kwargs...)::Plots.Plot
+    if dim(res) == 2
+        error("Data dimension ", dim(res), " not supported")
+    end
 
     if class == "default"
         if not_class == [] # plot stable full, unstable dashed
@@ -325,16 +309,18 @@ function plot_spaghetti(res::Result; x::String, y::String, z::String, class="def
     isnothing(z_index) && error("The variable $z was not swept over.")
 
     Z = res.swept_parameters.vals[z_index]
-    X = _apply_mask(transform_solutions(res, x), _get_mask(res, class, not_class)) |> _realify
-    Y = _apply_mask(transform_solutions(res, y), _get_mask(res, class, not_class)) |> _realify
+    X = _apply_mask(transform_solutions(res, x), _get_mask(res, class, not_class))
+    Y = _apply_mask(transform_solutions(res, y), _get_mask(res, class, not_class))
 
     # start a new plot if needed
     p = add ? Plots.plot!() : Plots.plot()
 
+    branch_data = [_realify( getindex.(X, i), warning= "branch " * string(k) ) for (i,k) in enumerate(1:branch_count(res))]
+
     # colouring is matched to branch index - matched across plots
-    for k in findall(x -> !all(isnan.(x)), X[1:end]) # skip NaN branches but keep indices
+    for k in findall(x -> !all(isnan.(x)), branch_data) # skip NaN branches but keep indices
         l = _is_labeled(p, k) ? nothing : k
-        Plots.plot!(X[k], Y[k], Z; _set_Plots_default...,
+        Plots.plot!(_realify(getindex.(X, k)), _realify(getindex.(Y, k)), Z; _set_Plots_default...,
          color=k, label=l, xlabel=latexify(x), ylabel=latexify(y), zlabel=latexify(z), xlim=:symmetric, ylim=:symmetric, kwargs...)
     end
     return p

src/sorting.jl

@@ -66,39 +66,38 @@ end
 
 "Match each solution from to_sort to a closest partner from refs"
 function get_distance_matrix(refs::Vector{Vector{SteadyState}}, to_sort::Vector{SteadyState})
-    distances = [get_distance_matrix(ref, to_sort) for ref in refs]
+    distances = map( ref -> get_distance_matrix(ref, to_sort), refs)
     lowest_distances = similar(distances[1])
-    for (i, el) in enumerate(lowest_distances)
-        lowest_distances[i] = min([d[i] for d in distances]...)
+    for idx in CartesianIndices(lowest_distances)
+        lowest_distances[idx] = minimum( x[idx] for x in distances )
     end
     lowest_distances
 end
 
 
 
 """
-Match a to_sort vector of solutions to a set of reference vector of solutions.
+Match a to_sort vector of solutions to a set of reference vectors of solutions.
 Returns a list of Tuples of the form (1, i1), (2, i2), ... such that
 reference[1] and to_sort[i1] belong to the same branch
 """
-function align_pair(reference::Vector{Vector{SteadyState}}, to_sort::Vector{SteadyState})
-
+function align_pair(reference, to_sort::Vector{SteadyState})
+    
     distances = get_distance_matrix(reference, to_sort)
     n = length(to_sort)
-    cartesians = [(j,i) for i in 1:n for j in 1:n]
-    sorted_cartesians = cartesians[sortperm(vec(distances))]
+    sorted_cartesians = CartesianIndices(distances)[sortperm(vec(distances))]
 
     matched = falses(n)
     matched_ref = falses(n)
 
-    sorted = Array{Tuple{Int64, Int64}, 1}(undef, n)
+    sorted = Vector{CartesianIndex}(undef, n)
 
-    for i in 1:length(cartesians)
-        j,k = sorted_cartesians[i]
+    for idx in sorted_cartesians
+        j,k = idx[1], idx[2]
         if !matched[k] && !matched_ref[j]
             matched[k] = true
             matched_ref[j] = true
-            sorted[j] = (j,k)
+            sorted[j] = idx
         end
     end
 
@@ -107,9 +106,6 @@ function align_pair(reference::Vector{Vector{SteadyState}}, to_sort::Vector{Stea
 end
 
 
-align_pair(ref::Vector{SteadyState}, to_sort::Vector{SteadyState}) = align_pair([ref], to_sort)
-
-
 """
 Go through a vector of solution and sort each according to Euclidean norm.
 """