CSCスパース行列から、あえて特定行を取り出す

訳あって、CSC(列志向)形式で保存されたスパース行列の特定の行だけをJuliaで取り出す必要があったため、書いたコード。

現状、JuliaではCSR(行志向)に対応していないらしい。

CSRであれば行の切り出しはスムーズだが、それをCSCでやると割と大変だった。

CSC、CSRなどの説明は以下の通り。

http://hamukazu.com/2014/12/03/internal-data-structure-scipy-sparse/

using SparseArrays

# こういう疎行列を列志向のスパース形式として持っていたとする
# col1, col2, col3, col4
# row1 | 3 4 3 |
# row2 | 5 2 |
# row3 | 6 1 |
# row4 | 1 2 |
# row5 | 2 1 3 3 |

# 下の二つのベクトルのどこからどこまでの要素がどのcolに属するのかを保存したポインター
indptr = [1,4,7,12,14]
# 何行目の要素か
indices = [2,4,5,1,3,5,1,2,3,4,5,1,5]
# 要素の値
data = [5,1,2,3,6,1,4,2,1,2,3,3,3]
A = SparseMatrixCSC(5, 4, indptr, indices, data)

# 今、row2からrow4までを切り出したCSCを新たに作成したい
# Aがかなり巨大でメモリに乗り切らないので、A[2:4,:]みたいに切り出すのではなく、
# indptr, indices, dataから作成する

function ExtractedCSC(indptr, indices, data, ncol, startp, endp)
# new objects
newindptr = ones(Int64, 1)
newindices = zeros(Int64, 0)
newdata = zeros(Int64, 0)
# Each column
l = length(indptr) - 1
for i=1:l
lo = indptr[i]
hi = indptr[i+1]
# Extract
extindices = indices[lo:hi-1]
extdata = data[lo:hi-1]
lower = searchsortedfirst(extindices, startp)
upper = searchsortedlast(extindices, endp)
# update
append!(newindices, extindices[lower:upper])
append!(newdata, extdata[lower:upper])
append!(newindptr, length(newdata)+1)
end
@assert minimum(newindices) == startp
@assert maximum(newindices) == endp
newindices = newindices .- (startp - 1)
return SparseMatrixCSC(endp - startp + 1, ncol,
newindptr, newindices, newdata)
end
ExtractedCSC(indptr, indices, data, 4, 2, 4)
# 3×4 SparseMatrixCSC{Int64,Int64} with 6 stored entries:
# [1, 1] = 5
# [3, 1] = 1
# [2, 2] = 6
# [1, 3] = 2
# [2, 3] = 1
# [3, 3] = 2