Merge pull request JuliaLang#12681 from KristofferC/kc/sparse_condense

Condense some code duplication for sparse matrices

base/sparse/csparse.jl

@@ -127,12 +127,12 @@ function sparse{Tv,Ti<:Integer}(I::AbstractVector{Ti},
  return SparseMatrixCSC(nrow, ncol, RpT, RiT, RxT)
 end
 
-## Transpose
+## Transpose and apply f
 
 # Based on Direct Methods for Sparse Linear Systems, T. A. Davis, SIAM, Philadelphia, Sept. 2006.
 # Section 2.5: Transpose
 # http:https://www.cise.ufl.edu/research/sparse/CSparse/
-function transpose!{Tv,Ti}(T::SparseMatrixCSC{Tv,Ti}, S::SparseMatrixCSC{Tv,Ti})
+function ftranspose!{Tv,Ti}(T::SparseMatrixCSC{Tv,Ti}, S::SparseMatrixCSC{Tv,Ti}, f)
  (mS, nS) = size(S)
  nnzS = nnz(S)
  colptr_S = S.colptr
@@ -157,63 +157,37 @@ function transpose!{Tv,Ti}(T::SparseMatrixCSC{Tv,Ti}, S::SparseMatrixCSC{Tv,Ti})
  q = w[ind]
  w[ind] += 1
  rowval_T[q] = j
- nzval_T[q] = nzval_S[p]
+ nzval_T[q] = f(nzval_S[p])
  end
 
  return T
 end
 
-function transpose{Tv,Ti}(S::SparseMatrixCSC{Tv,Ti})
+function ftranspose{Tv,Ti}(S::SparseMatrixCSC{Tv,Ti}, f)
  (nT, mT) = size(S)
  nnzS = nnz(S)
  colptr_T = Array(Ti, nT+1)
  rowval_T = Array(Ti, nnzS)
  nzval_T = Array(Tv, nnzS)
 
  T = SparseMatrixCSC(mT, nT, colptr_T, rowval_T, nzval_T)
- return transpose!(T, S)
+ return ftranspose!(T, S, f)
 end
 
-function ctranspose!{Tv,Ti}(T::SparseMatrixCSC{Tv,Ti}, S::SparseMatrixCSC{Tv,Ti})
- (mS, nS) = size(S)
- nnzS = nnz(S)
- colptr_S = S.colptr
- rowval_S = S.rowval
- nzval_S = S.nzval
-
- (mT, nT) = size(T)
- colptr_T = T.colptr
- rowval_T = T.rowval
- nzval_T = T.nzval
-
- fill!(colptr_T, 0)
- colptr_T[1] = 1
- for i=1:nnzS
- @inbounds colptr_T[rowval_S[i]+1] += 1
- end
- cumsum!(colptr_T, colptr_T)
+function transpose!{Tv,Ti}(T::SparseMatrixCSC{Tv,Ti}, S::SparseMatrixCSC{Tv,Ti})
+ ftranspose!(T, S, IdFun())
+end
 
- w = copy(colptr_T)
- @inbounds for j = 1:nS, p = colptr_S[j]:(colptr_S[j+1]-1)
- ind = rowval_S[p]
- q = w[ind]
- w[ind] += 1
- rowval_T[q] = j
- nzval_T[q] = conj(nzval_S[p])
- end
+function transpose{Tv,Ti}(S::SparseMatrixCSC{Tv,Ti})
+ ftranspose(S, IdFun())
+end
 
- return T
+function ctranspose!{Tv,Ti}(T::SparseMatrixCSC{Tv,Ti}, S::SparseMatrixCSC{Tv,Ti})
+ ftranspose!(T, S, ConjFun())
 end
 
 function ctranspose{Tv,Ti}(S::SparseMatrixCSC{Tv,Ti})
- (nT, mT) = size(S)
- nnzS = nnz(S)
- colptr_T = Array(Ti, nT+1)
- rowval_T = Array(Ti, nnzS)
- nzval_T = Array(Tv, nnzS)
-
- T = SparseMatrixCSC(mT, nT, colptr_T, rowval_T, nzval_T)
- return ctranspose!(T, S)
+ ftranspose(S, ConjFun())
 end
 
 # Compute the elimination tree of A using triu(A) returning the parent vector.

base/sparse/linalg.jl

@@ -34,54 +34,39 @@ function (*){TvA,TiA,TvB,TiB}(A::SparseMatrixCSC{TvA,TiA}, B::SparseMatrixCSC{Tv
 end
 
 # In matrix-vector multiplication, the correct orientation of the vector is assumed.
-function A_mul_B!(α::Number, A::SparseMatrixCSC, B::StridedVecOrMat, β::Number, C::StridedVecOrMat)
- A.n == size(B, 1) || throw(DimensionMismatch())
- A.m == size(C, 1) || throw(DimensionMismatch())
- size(B, 2) == size(C, 2) || throw(DimensionMismatch())
- if β != 1
- β != 0 ? scale!(C, β) : fill!(C, zero(eltype(C)))
- end
- nzv = A.nzval
- rv = A.rowval
- for col = 1:A.n
- for j in 1:size(C, 2)
- αxj = α*B[col,j]
- @inbounds for k = A.colptr[col]:(A.colptr[col + 1] - 1)
- C[rv[k], j] += nzv[k]*αxj
- end
- end
- end
- C
-end
 
-function (*){TA,S,Tx}(A::SparseMatrixCSC{TA,S}, x::StridedVector{Tx})
- T = promote_type(TA,Tx)
- A_mul_B!(one(T), A, x, zero(T), similar(x, T, A.m))
-end
-function (*){TA,S,Tx}(A::SparseMatrixCSC{TA,S}, B::StridedMatrix{Tx})
- T = promote_type(TA,Tx)
- A_mul_B!(one(T), A, B, zero(T), similar(B, T, (A.m, size(B, 2))))
-end
-
-for (f, op) in ((:Ac_mul_B, :ctranspose),
- (:At_mul_B, :transpose))
+for (f, op, transp) in ((:A_mul_B, :identity, false),
+ (:Ac_mul_B, :ctranspose, true),
+ (:At_mul_B, :transpose, true))
  @eval begin
  function $(symbol(f,:!))(α::Number, A::SparseMatrixCSC, B::StridedVecOrMat, β::Number, C::StridedVecOrMat)
- A.n == size(C, 1) || throw(DimensionMismatch())
- A.m == size(B, 1) || throw(DimensionMismatch())
+ if $transp
+ A.n == size(C, 1) || throw(DimensionMismatch())
+ A.m == size(B, 1) || throw(DimensionMismatch())
+ else
+ A.n == size(B, 1) || throw(DimensionMismatch())
+ A.m == size(C, 1) || throw(DimensionMismatch())
+ end
  size(B, 2) == size(C, 2) || throw(DimensionMismatch())
  nzv = A.nzval
  rv = A.rowval
  if β != 1
  β != 0 ? scale!(C, β) : fill!(C, zero(eltype(C)))
  end
- for i = 1:A.n
+ for col = 1:A.n
  for k = 1:size(C, 2)
- tmp = zero(eltype(C))
- @inbounds for j = A.colptr[i]:(A.colptr[i + 1] - 1)
- tmp += $(op)(nzv[j])*B[rv[j],k]
+ if $transp
+ tmp = zero(eltype(C))
+ @inbounds for j = A.colptr[col]:(A.colptr[col + 1] - 1)
+ tmp += $(op)(nzv[j])*B[rv[j],k]
+ end
+ C[col,k] += α*tmp
+ else
+ αxj = α*B[col,k]
+ @inbounds for j = A.colptr[col]:(A.colptr[col + 1] - 1)
+ C[rv[j], k] += nzv[j]*αxj
+ end
  end
- C[i,k] += α*tmp
  end
  end
  C