Skip to content

Commit 83891bf

Browse files
committedNov 4, 2020
Improve quadprod with columnar-only implementation
First, this removes the option to do row-wise quadratic products since they aren't being used within this package anyway. That allows removing the "keyword" argument for choosing which direction to apply. Second, the original implementation was optimized for the fast vector outer product (that I got added to SparseArrays in JuliaLang/julia#24980 and made it into Julia v1.2), but when scaling up to multiple columns the performance was disastrous because the dispatch of a transposed view led to generic matrix multiplication which did the full dense-dense style loops. By not using views, we get the desired sparse matrix multiplication instead.
1 parent 8dc6d4d commit 83891bf

File tree

1 file changed

+19
-9
lines changed

1 file changed

+19
-9
lines changed
 

‎src/numerics.jl

+19-9
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
using Compat # Compat@v3.23 for sincospi()
33
end
44
using SparseArrays
5+
using SparseArrays: AbstractSparseMatrix
56

67
# COV_EXCL_START
78

@@ -23,18 +24,27 @@ unchecked_sqrt(x::T) where {T <: Integer} = unchecked_sqrt(float(x))
2324
unchecked_sqrt(x) = Base.sqrt(x)
2425

2526
"""
26-
quadprod(A, b, n, dir=:col)
27+
quadprod(A::AbstractSparseMatrixCSC, b::AbstractVecOrMat, n::Integer)
2728
2829
Computes the quadratic product ``ABA^\\top`` efficiently for the case where ``B`` is all zero
29-
except for the `n`th column or row vector `b`, for `dir = :col` or `dir = :row`,
30-
respectively.
30+
except for a small number of columns `b` starting at the `n`th.
3131
"""
32-
@inline function quadprod(A, b, n, dir::Symbol=:col)
33-
if dir == :col
34-
return (A * sparse(b)) * view(A, :, n)'
35-
elseif dir == :row
36-
return view(A, :, n) * (A * sparse(b))'
32+
function quadprod(A::AbstractSparseMatrix, b::AbstractVecOrMat, n::Integer)
33+
size(b, 1) == size(A, 2) || throw(DimensionMismatch())
34+
35+
# sparse * dense naturally returns dense, but we want to dispatch to
36+
# a sparse-sparse matrix multiplication, so forceably sparsify.
37+
# - Tests with a few example matrices A show that `sparse(A * b)` is faster than
38+
# `A * sparse(b)`.
39+
w = sparse(A * b)
40+
p = n + size(b, 2) - 1
41+
42+
if ndims(w) == 1
43+
# vector outer product using column view into matrix is fast
44+
C = w * transpose(view(A, :, n))
3745
else
38-
error("Unrecognized direction `dir = $(repr(dir))`.")
46+
# views are not fast for multiple columns; subset copies are faster
47+
C = w * transpose(A[:, n:p])
3948
end
49+
return C
4050
end

0 commit comments

Comments
 (0)
Please sign in to comment.