Remove the closure boxes

src/solve_homotopy.jl

@@ -220,8 +220,9 @@ function _get_raw_solution(
         alg_default_options(method)...,
     )
 
-    raw_solutions = reshape(result_full, size(parameter_values)...)
-    raw_solutions = HC.solutions.(getindex.(raw_solutions, 1))
+    raw_solutions = HC.solutions.(
+        getindex.(reshape(result_full, size(parameter_values)...), 1)
+    )
     # cache = HC.NewtonCache(F)
     if method.check_zero
         Threads.@threads for i in eachindex(parameter_values)
@@ -271,8 +272,9 @@ function _get_raw_solution(
         ]
     end
 
-    raw_solutions = reshape(result_full, size(parameter_values)...)
-    raw_solutions = HC.solutions.(getindex.(raw_solutions, 1))
+    raw_solutions = HC.solutions.(
+        getindex.(reshape(result_full, size(parameter_values)...), 1)
+    )
 
     # reshape back to the original shape
     return raw_solutions

src/transform_solutions.jl

@@ -1,3 +1,13 @@
+"""
+The position of `index` among `solutions`: `index` itself if it has one entry per
+dimension, otherwise a single linear index.
+"""
+function solution_position(solutions::AbstractArray, index)
+    length(index) == ndims(solutions) && return CartesianIndex(index)
+    length(index) == 1 && return CartesianIndices(solutions)[only(index)]
+    return error("Index $index undefined for a solution of size $(size(solutions))")
+end
+
 """
 $(TYPEDSIGNATURES)
 Return an ordered dictionary specifying all variables and parameters of the solution
@@ -6,24 +16,13 @@ in `result` on `branch` at the position `index`.
 function get_single_solution(
     res::Result{D,S,P}; branch::Int, index
 )::OrderedDict{Num,S} where {D,S,P}
+    position = solution_position(res.solutions, index)
 
-    # check if the dimensionality of index matches the solutions
-    if length(size(res.solutions)) !== length(index)
-        # if index is a number, use linear indexing
-        index = if length(index) == 1
-            CartesianIndices(res.solutions)[index]
-        else
-            error("Index ", index, " undefined for a solution of size ", size(res.solutions))
-        end
-    else
-        index = CartesianIndex(index)
-    end
-
-    vars = OrderedDict(zip(res.problem.variables, res.solutions[index][branch]))
+    vars = OrderedDict(zip(res.problem.variables, res.solutions[position][branch]))
 
     # collect the swept parameters required for this call
     swept_params = OrderedDict(
-        key => res.swept_parameters[key][index[i]] for
+        key => res.swept_parameters[key][position[i]] for
         (i, key) in enumerate(keys(res.swept_parameters))
     )
     full_solution = merge(vars, swept_params, res.fixed_parameters)
@@ -40,24 +39,13 @@ end
 function get_variable_solutions(
     res::Result{D,S,P}; branch::Int, index
 )::Vector{S} where {D,S,P}
+    position = solution_position(res.solutions, index)
 
-    # check if the dimensionality of index matches the solutions
-    if length(size(res.solutions)) !== length(index)
-        # if index is a number, use linear indexing
-        index = if length(index) == 1
-            CartesianIndices(res.solutions)[index]
-        else
-            error("Index ", index, " undefined for a solution of size ", size(res.solutions))
-        end
-    else
-        index = CartesianIndex(index)
-    end
-
-    sol = res.solutions[index][branch]
+    sol = res.solutions[position][branch]
 
     # collect the swept parameters required for this call
     swept_params = [
-        res.swept_parameters[key][index[i]] for
+        res.swept_parameters[key][position[i]] for
         (i, key) in enumerate(keys(res.swept_parameters))
     ]
     return vcat(sol, swept_params)
@@ -263,14 +251,15 @@ function get_branches(
     class=["physical", "stable"],
     not_class=[],
 )
-    Y = transform_solutions(res, y; branches, realify)
-    Y = _apply_mask(Y, _get_mask(res, class, not_class; branches))
+    Y = _apply_mask(
+        transform_solutions(res, y; branches, realify),
+        _get_mask(res, class, not_class; branches),
+    )
     if realify
-        branches = map(idx -> real(getindex.(Y, idx)), eachindex(first(Y)))
+        return map(idx -> real(getindex.(Y, idx)), eachindex(first(Y)))
     else
-        branches = map(idx -> getindex.(Y, idx), eachindex(first(Y)))
+        return map(idx -> getindex.(Y, idx), eachindex(first(Y)))
     end
-    return branches
 end
 
 ###

test/Project.toml

@@ -31,7 +31,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 
 [sources]
-HarmonicBalance = {rev = "refactor/harmonic-kind", url = "https://github.com/QuantumEngineeredSystems/HarmonicBalance.jl"}
+HarmonicBalance = {rev = "master", url = "https://github.com/QuantumEngineeredSystems/HarmonicBalance.jl"}
 HarmonicSteadyState = {path = ".."}
 QuestBase = {rev = "main", url = "https://github.com/QuantumEngineeredSystems/QuestBase.jl"}
 

test/steady_states/krylov.jl

@@ -57,7 +57,6 @@ res2 = get_steady_states(harmonic_eq2, varied, fixed; show_progress=false);
             solk = substitute(eqk, subs)
             solh = substitute(eqh, subs)
             @test Float64(solk - solh) ≈ 0.0 atol = 1e-10
-            # ^ different ansatz
         end
     end
 
@@ -80,7 +79,6 @@ res2 = get_steady_states(harmonic_eq2, varied, fixed; show_progress=false);
             solk = substitute(eqk, subs)
             solh = substitute(eqh, subs)
             @test Float64(solk - solh) ≈ 0.0 atol = 1e-10
-            # ^ different ansatz
         end
     end
 
@@ -109,8 +107,7 @@ res2 = get_steady_states(harmonic_eq2, varied, fixed; show_progress=false);
             subs = Dict([u1, v1, u2, v2, ω₁, ω₂, ω, F, J₂, J₁, α₁, α₂] .=> rand(12))
             solk = substitute(eqk, subs)
             solh = substitute(eqh, subs)
-            @test Float64(solk + solh) ≈ 0.0 broken = true
-            # ^ different ansatz
+            @test Float64(solk - solh) ≈ 0.0 atol = 1e-10
         end
     end
 end