How to keep a banded diagonal matrix and replace other elements by 0 in a large matrix for julia

Question

I would like to keep the diagonal matrix and replace other elements by 0 in a large matrix for julia. For example, A is the matrix I have, I want to only keep the 2 by 2 diagonal elements in A and replace all other elements by 0. B matrix is what I want. I am just wondering is there an eleglant way to do it.

A = [1 2 3 4 5 6 7 8; 
     1 2 3 4 5 6 7 8; 
     1 2 3 4 5 6 7 8; 
     1 2 3 4 5 6 7 8; 
     1 2 3 4 5 6 7 8;
     1 2 3 4 5 6 7 8; 
     1 2 3 4 5 6 7 8; 
     1 2 3 4 5 6 7 8]

B = [1 2 0 0 0 0 0 0; 
     1 2 0 0 0 0 0 0; 
     0 0 3 4 0 0 0 0; 
     0 0 3 4 0 0 0 0; 
     0 0 0 0 5 6 0 0; 
     0 0 0 0 5 6 0 0; 
     0 0 0 0 0 0 7 8; 
     0 0 0 0 0 0 7 8]

Answer 1

Method 1:
Here's an interesting way using CartesianIndices:

julia> B = zero(A);
julia> blocksize = 2; 
julia> d = diag(CartesianIndices(A))
8-element Vector{CartesianIndex{2}}:
 CartesianIndex(1, 1)
 CartesianIndex(2, 2)
 CartesianIndex(3, 3)
 CartesianIndex(4, 4)
 CartesianIndex(5, 5)
 CartesianIndex(6, 6)
 CartesianIndex(7, 7)
 CartesianIndex(8, 8)

julia> for p in Iterators.partition(d, blocksize)
         block = first(p):last(p)
         B[block] .= @view A[block]
       end

In each iteration, Iterators.partition returns blocksize number of diagonal elements, so all the diagonal elements that belong in a block.

A useful thing about CartesianIndices is that ranges operate blockwise already: CartesianIndex(1,1):CartesianIndex(2,2) returns CartesianIndex values of (1,1),(2,1),(1,2), and (2,2) automatically. So first(p):last(p) returns the indices of all the elements in the block we want, in each iteration.

---

Method 2:
In this case, because things are symmetrical, the non-CartesianIndices way is is also pretty neat and simple:

julia> B = zero(A);
julia> for b in Iterators.partition(1:size(A, 1), blocksize)
         B[b,b] .= @view A[b,b]
       end
julia> B
8×8 Matrix{Int64}:
 1  2  0  0  0  0  0  0
 1  2  0  0  0  0  0  0
 0  0  3  4  0  0  0  0
 0  0  3  4  0  0  0  0
 0  0  0  0  5  6  0  0
 0  0  0  0  5  6  0  0
 0  0  0  0  0  0  7  8
 0  0  0  0  0  0  7  8

In the first iteration (as an example), partition returns 1:2 (assuming blocksize = 2), so we assign to B[1:2, 1:2] which is the block we want.

To generalize that to allow non-standard indexing (eg. OffsetArrays):

julia> for (r, c) in zip(Iterators.partition.(axes(A), blocksize)...)
         B[r, c] .= @view A[r, c] 
       end

---

(Thanks to @phipsgabler for the .= @view suggestion which avoids unnecessary allocations, and for the axes(A) method.)