feat: vector helper functions

Signed-off-by: Anirudh <[email protected]>

src/SemioticOpt.jl

@@ -9,6 +9,7 @@ using Zygote
 
 abstract type OptAlgorithm end
 
+include("vector.jl")
 include("hook/core.jl")
 include("hook/stop.jl")
 include("hook/postiteration.jl")

src/vector.jl

@@ -0,0 +1,18 @@
+# Copyright 2022-, Semiotic AI, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+export nonzeroixs, klargestixs
+
+"""
+    nonzeroixs(x::AbstractVector{T}) where {T<:Real}
+
+Returns the indices of the non-zero elements of `x`.
+"""
+nonzeroixs(x::AbstractVector{T}) where {T<:Real} = findall(!iszero, x) 
+
+"""
+    klargestixs(x::AbstractVector{T}, k::I) where {T<:Real, I<:Integer}
+
+Returns the indices of the `k` largest elements of `x`.
+"""
+klargestixs(x::AbstractVector{T}, k::I) where {T<:Real, I<:Integer} = partialsortperm(x, 1:k; rev=true)

test/testgroups

@@ -1,3 +1,4 @@
+vector
 hook
 project
 gradientdescent

test/vector.jl

@@ -0,0 +1,12 @@
+# Copyright 2022-, Semiotic AI, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+@testset "vector" begin
+    @testset "nonzerosixs" begin
+        @test nonzeroixs([1.0, 0.0, 2.0, 0.0, 3.0]) == [1, 3, 5]
+    end
+
+    @testset "klargestixs" begin
+        @test klargestixs([1.0, 0.0, 2.0, 0.0, 3.0], 2) == [5, 3]
+    end
+end