Search code examples
arraysmatrixjuliavectorizationarray-broadcasting

How to get a list of submatrices given a huge matrix and two vector of vector as indices in Julia


I have a matrix W (of some size), and two Vector{Vector} xList, and yList. I know I can access a submatrix of W using each element of xList and yList (which are Vectors themselves). I want to know a way I can vectorize this over all elements of these two Vectors.

Eg. let W matrix be

W=rand(1:10, 3, 3)
3×3 Matrix{Int64}:
 6   2   5
 1   5  10
 7  10   6

and

xList=[[1, 2], [1, 3]]
2-element Vector{Vector{Int64}}:
 [1, 2]
 [1, 3]

yList=[[2, 2], [3, 1]]
2-element Vector{Vector{Int64}}:
 [2, 2]
 [3, 1]

I basically want to vectorize the following for loop :

 Wsub=[]
 for i in 1:size(xList)[1]
      push!(Wsub, W[xList[i], yList[i]])
 end

julia> Wsub
2-element Vector{Any}:
 [2 2; 5 5]
 [5 6; 6 7]

The actual matrix I have is roughly 1000 x 500, and the xList and yList will have around 25000 Vectors each. Any help will be appreciated!


Solution

  • Instead of the loop, you have getindex.((W,), xList, yList) to get the same result. This can also be written as getindex.(Ref(W), xList, yList) if you prefer. Broadcasting is used here to apply the getindex function elementwise to the pairs of indices in xList and yList, using W as the matrix being indexed.

    julia> xList = [[1, 2], [1, 3]]
    2-element Vector{Vector{Int64}}:
     [1, 2]
     [1, 3]
    
    julia> yList = [[2, 2], [3, 1]]
    2-element Vector{Vector{Int64}}:
     [2, 2]
     [3, 1]
    
    julia> W = rand(1:10, 3, 3)
    3×3 Matrix{Int64}:
     6  2  1
     5  8  7
     6  1  5
    
    julia> getindex.((W,), xList, yList)
    2-element Vector{Matrix{Int64}}:
     [2 2; 8 8]
     [1 6; 5 6]