export filter_result! (#402)

Author: oameye

docs/src/manual/API.md

@@ -61,6 +61,7 @@ transform_solutions
 ```@docs
 classify_solutions!
 get_class
+filter_result!
 ```
 
 ### Plotting

docs/src/manual/analyse_solutions.md

@@ -17,3 +17,11 @@ get_single_solution
 attractors
 phase_diagram
 ```
+
+## classifying solutions
+
+```@docs; canonical=false
+classify_solutions!
+get_class
+filter_result!
+```

src/HarmonicBalance.jl

@@ -59,6 +59,7 @@ export get_variables
 export get_steady_states
 export classify_solutions!
 export get_class
+export filter_result!
 export get_single_solution
 export transform_solutions
 export IM_TOL

src/Result.jl

@@ -7,7 +7,7 @@ Stores the steady states of a HarmonicEquation.
 $(TYPEDFIELDS)
 
 """
-struct Result{D,SolType<:Number,ParType<:Number,F<:JacobianFunction(SolType)}
+mutable struct Result{D,SolType<:Number,ParType<:Number,F<:JacobianFunction(SolType)}
     "The variable values of steady-state solutions."
     solutions::Array{Vector{Vector{SolType}},D}
     "Values of all parameters for all solutions."
@@ -81,7 +81,7 @@ function _free_symbols(res::Result)
 end
 
 # overload to use [] for indexing
-Base.getindex(r::Result, idx::Int...) = get_single_solution(r, idx)
+Base.getindex(r::Result, idx) = get_single_solution(r, idx)
 Base.size(r::Result) = size(r.solutions)
 
 branch_count(r::Result) = length(r.solutions[1])

src/classification.jl

@@ -194,7 +194,7 @@ Removes all solution branches from `res` where NONE of the solution falls into `
 Typically used to filter out unphysical solutions to prevent huge file sizes.
 """
 function filter_result!(res::Result, class::String)
-    bools = [any(getindex.(res.classes[class], i)) for i in 1:length(res[1])]
+    bools = [any(getindex.(res.classes[class], i)) for i in 1:branch_count(res)]
     res.solutions = [s[bools] for s in res.solutions]
     for c in keys(res.classes)
         res.classes[c] = [s[bools] for s in res.classes[c]]

test/classification.jl

@@ -0,0 +1,29 @@
+using HarmonicBalance, Test
+
+@variables ω₀ γ λ α ω t x(t)
+
+natural_equation = d(d(x, t), t) + γ * d(x, t) + (ω₀^2 - λ * cos(2 * ω * t)) * x + α * x^3
+diff_eq = DifferentialEquation(natural_equation, x)
+
+add_harmonic!(diff_eq, x, ω);
+
+harmonic_eq = get_harmonic_equations(diff_eq)
+
+fixed = (ω₀ => 1.0, γ => 0.002, α => 1.0)
+varied = (ω => range(0.99, 1.01, 100), λ => range(1e-6, 0.03, 100))
+
+result_2D = get_steady_states(harmonic_eq, varied, fixed; verbose=true)
+
+@testset "filter solutions" begin
+    result = deepcopy(result_2D)
+    @test length(result.solutions[1]) == 5
+    filter_result!(result, "stable")
+    @test length(result.solutions[1]) == 3
+end
+
+@testset "Classifying solutions" begin
+    result = deepcopy(result_2D)
+    filter_result!(result, "stable")
+    classify_solutions!(result, "sqrt(u1^2 + v1^2) >0.0", "not_zero")
+    @test sum(any.(get_class(result, "not_zero"))) == 2
+end

test/runtests.jl

@@ -30,6 +30,7 @@ end
     include("Jacobian.jl")
     include("transform_solutions.jl")
     include("sorting.jl")
+    include("classification.jl")
 end
 
 @testset "Plotting" begin