add fkeep! and children for sparse vectors

base/sparse/sparsematrix.jl

@@ -732,22 +732,22 @@ ctranspose{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}) = qftranspose(A, 1:A.n, Base.ConjFu
 ## fkeep! and children tril!, triu!, droptol!, dropzeros[!]
 
 """
-    fkeep!{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}, f, other, trim::Bool = true)
+    fkeep!(A::AbstractSparseArray, f, other, trim::Bool = true)
 
 Keep elements of `A` for which test `f` returns `true`. `f`'s signature should be
 
-    f{Tv,Ti}(i::Ti, j::Ti, x::Tv, other::Any) -> Bool
+    f(i::Integer, [j::Integer,] x, other::Any) -> Bool
 
 where `i` and `j` are an element's row and column indices, `x` is the element's value,
 and `other` is passed in from the call to `fkeep!`. This method makes a single sweep
-through `A`, requiring `O(A.n, nnz(A))`-time and no space beyond that passed in. If `trim`
-is `true`, this method trims `A.rowval` and `A.nzval` to length `nnz(A)` after dropping
-elements.
+through `A`, requiring `O(A.n, nnz(A))`-time for matrices and `O(nnz(A))`-time for vectors
+and no space beyond that passed in. If `trim` is `true`, this method trims `A.rowval` or `A.nzind` and
+`A.nzval` to length `nnz(A)` after dropping elements.
 
 Performance note: As of January 2016, `f` should be a functor for this method to perform
 well. This caveat may disappear when the work in `jb/functions` lands.
 """
-function fkeep!{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}, f, other, trim::Bool = true)
+function fkeep!(A::SparseMatrixCSC, f, other, trim::Bool = true)
     An = A.n
     Acolptr = A.colptr
     Arowval = A.rowval
@@ -811,7 +811,7 @@ droptol!(A::SparseMatrixCSC, tol, trim::Bool = true) = fkeep!(A, DroptolFunc(),
 
 immutable DropzerosFunc <: Base.Func{4} end
 (::DropzerosFunc){Tv,Ti}(i::Ti, j::Ti, x::Tv, other) = x != 0
-dropzeros!(A::SparseMatrixCSC, trim::Bool = true) = fkeep!(A, DropzerosFunc(), Void, trim)
+dropzeros!(A::SparseMatrixCSC, trim::Bool = true) = fkeep!(A, DropzerosFunc(), nothing, trim)
 dropzeros(A::SparseMatrixCSC, trim::Bool = true) = dropzeros!(copy(A), trim)
 
 

base/sparse/sparsevector.jl

@@ -1678,3 +1678,43 @@ function sort{Tv,Ti}(x::SparseVector{Tv,Ti}; kws...)
     newnzvals = allvals[deleteat!(sinds[1:k],z)]
     SparseVector(n,newnzind,newnzvals)
 end
+
+function fkeep!(x::SparseVector, f, other, trim::Bool = true)
+    n = x.n
+    nzind = x.nzind
+    nzval = x.nzval
+
+    x_writepos = 1
+    @inbounds for xk in 1:nnz(x)
+        xi = nzind[xk]
+        xv = nzval[xk]
+        # If this element should be kept, rewrite in new position
+        if f(xi, xv, other)
+            if x_writepos != xk
+                nzind[x_writepos] = xi
+                nzval[x_writepos] = xv
+            end
+            x_writepos += 1
+        end
+    end
+
+    # Trim x's storage if necessary and desired
+    if trim
+        x_nnz = x_writepos - 1
+        if length(nzind) != x_nnz
+            resize!(nzval, x_nnz)
+            resize!(nzind, x_nnz)
+        end
+    end
+
+    x
+end
+
+immutable DroptolFuncVec <: Base.Func{3} end
+(::DroptolFuncVec){Tv,Ti}(i::Ti, x::Tv, tol::Real) = abs(x) > tol
+droptol!(x::SparseVector, tol, trim::Bool = true) = fkeep!(x, DroptolFuncVec(), tol, trim)
+
+immutable DropzerosFuncVec <: Base.Func{3} end
+(::DropzerosFuncVec){Tv,Ti}(i::Ti, x::Tv, other) = x != 0
+dropzeros!(x::SparseVector, trim::Bool = true) = fkeep!(x, DropzerosFuncVec(), nothing, trim)
+dropzeros(x::SparseVector, trim::Bool = true) = dropzeros!(copy(x), trim)

test/sparsedir/sparsevector.jl

@@ -869,6 +869,30 @@ let m = 10
     end
 end
 
+# fkeep!
+let x = sparsevec(1:7, [3., 2., -1., 1., -2., -3., 3.], 7)
+    # droptol
+    xdrop = Base.droptol!(copy(x), 1.5)
+    @test exact_equal(xdrop, SparseVector(7, [1, 2, 5, 6, 7], [3., 2., -2., -3., 3.]))
+    Base.droptol!(xdrop, 2.5)
+    @test exact_equal(xdrop, SparseVector(7, [1, 6, 7], [3., -3., 3.]))
+    Base.droptol!(xdrop, 3.)
+    @test exact_equal(xdrop, SparseVector(7, Int[], Float64[]))
+
+    # dropzeros
+    xdrop = copy(x)
+    xdrop.nzval[[2, 4, 6]] = 0.0
+    Base.SparseArrays.dropzeros!(xdrop)
+    @test exact_equal(xdrop, SparseVector(7, [1, 3, 5, 7], [3, -1., -2., 3.]))
+
+    xdrop = copy(x)
+    # This will keep index 1, 3, 4, 7 in xdrop
+    f_drop(i, x, other) = (abs(x) == 1.) || (i in [1, 7])
+    Base.SparseArrays.fkeep!(xdrop, f_drop, Void)
+    @test exact_equal(xdrop, SparseVector(7, [1, 3, 4, 7], [3., -1., 1., 3.]))
+end
+
+
 # It's tempting to share data between a SparseVector and a SparseArrays,
 # but if that's done, then modifications to one or the other will cause
 # an inconsistent state: