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

Probably there is a high level API for it somewhere, but, writing a for loop should work.

function change_zero!(a)
 lo = 1
 for j in 1:size(a, 2)
   if isodd(j)
     lo += 2
   end
   for i in 1:lo-3
     a[i,j]=0
   end
   for i in lo:size(a,1)
      a[i,j]=0
   end
 end
 a
end

change_zero!(A)

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.)

The shortest code to achieve this is by using BlockBandedMatrices as here:

julia> BlockBandedMatrix(A,repeat([2],4),repeat([2],4),(0,0))
4×4-blocked 8×8 BlockBandedMatrix{Int64}:
 1  2  │  ⋅  ⋅   │  ⋅  ⋅  │  ⋅  ⋅
 1  2  │  ⋅  ⋅   │  ⋅  ⋅  │  ⋅  ⋅
 ──────┼────────┼───────┼──────
  ⋅  ⋅  │  3  4  │  ⋅  ⋅  │  ⋅  ⋅
  ⋅  ⋅  │  3  4  │  ⋅  ⋅  │  ⋅  ⋅
 ──────┼────────┼───────┼──────
  ⋅  ⋅  │  ⋅  ⋅   │  5  6 │  ⋅  ⋅
  ⋅  ⋅  │  ⋅  ⋅   │  5  6 │  ⋅  ⋅
 ──────┼────────┼───────┼──────
  ⋅  ⋅  │  ⋅  ⋅   │  ⋅  ⋅  │  7  8
  ⋅  ⋅  │  ⋅  ⋅   │  ⋅  ⋅  │  7  8

Another thing worth looking is BandedMatrices package that provides such functionality along with a set of dedicated linear algebra functions for efficient handling of such data structures.

julia> using BandedMatrices

julia> BandedMatrix(A, (1,0))
8×8 BandedMatrix{Int64} with bandwidths (1, 0):
 1  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅
 1  2  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅
 ⋅  2  3  ⋅  ⋅  ⋅  ⋅  ⋅
 ⋅  ⋅  3  4  ⋅  ⋅  ⋅  ⋅
 ⋅  ⋅  ⋅  4  5  ⋅  ⋅  ⋅
 ⋅  ⋅  ⋅  ⋅  5  6  ⋅  ⋅
 ⋅  ⋅  ⋅  ⋅  ⋅  6  7  ⋅
 ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  7  8

The following is, for completeness, an idiomatic answer to the question in your (original) title:

julia> function zeronondiag!(A)
           di = diagind(A)
           for i in eachindex(A)
               i ∉ di && (A[i] = zero(A[i]))
           end
           return A
       end
zeronondiag! (generic function with 1 method)

julia> zeronondiag!(copy(A))
8×8 Matrix{Int64}:
 1  0  0  0  0  0  0  0
 0  2  0  0  0  0  0  0
 0  0  3  0  0  0  0  0
 0  0  0  4  0  0  0  0
 0  0  0  0  5  0  0  0
 0  0  0  0  0  6  0  0
 0  0  0  0  0  0  7  0
 0  0  0  0  0  0  0  8

Note that diagind returns a range of the linear indices, so ∉ checking is reasonably efficient.

julia> diagind(A)
1:9:64

You should be able to use very similar logic with BlockArrays.jl to get your block diagonal form.

Like one of the answers, I prefer to write a simple function with loops inside for this kind of manipulation. A slightly more general function that allows you to specify the value of the off-diagonal elements and the size of the block diagonals:

function setoffdiag!(A::AbstractMatrix{T}, value::T = zero(T); block_size::Integer = 1) where {T}

    m, n = size(A)
    k = 1
    r = 1

    @inbounds for j = 1:n

        @inbounds for i = 1:(k - 1)
            A[i, j] = value
        end

        @inbounds for i = (k + block_size):m
            A[i, j] = value
        end

        k += (r == block_size) * block_size
        r += 1 - (r == block_size) * block_size
    end

    return A
end

This quick fix might do the trick: (Assuming the input is a squared matrix)

function two_by_two_diag(A)
    B = zeros(Int64,size(A))
    for i in 1:2:size(A,1)
        B[i,i] = A[i,i]
        if i != size(A,1)
            B[i,i+1] = A[i,i+1]
            B[i+1,i] = A[i+1,i]
            B[i+1, i+1] = A[i+1, i+1]
        end
    end
    return B
end

using LinearAlgebra

d = diag(CartesianIndices(A))
blocksize = 2;
mapreduce(p->A[first(p):last(p)],(x,y)->cat(x,y,dims=(1,2)), Iterators.partition(d, blocksize))

below a much more performing version ( > 10X )

function blockmat(A,bs)
    n=size(A, 1)   
    B = zero(A);
    doff=1
    soff=1
    for g in 1:div(n,bs)
        for i in 1:bs
            copyto!(B,doff,A,soff,bs)
            doff += n
            soff += n
        end
        doff+=bs
        soff+=bs
    end
    B
end



julia> @btime begin
       for b in Iterators.partition(1:size(A, 1), blocksize)
           B[b,b] .= @view A[b,b]
       end
       B
       end
  2.156 μs (37 allocations: 2.41 KiB)
15×15 Matrix{Int64}:
 1  16  31   0   0   0   0    0    0    0    0    0    0    0    0     
 2  17  32   0   0   0   0    0    0    0    0    0    0    0    0     
 3  18  33   0   0   0   0    0    0    0    0    0    0    0    0     
 0   0   0  49  64  79   0    0    0    0    0    0    0    0    0     
 0   0   0  50  65  80   0    0    0    0    0    0    0    0    0     
 0   0   0  51  66  81   0    0    0    0    0    0    0    0    0     
 0   0   0   0   0   0  97  112  127    0    0    0    0    0    0     
 0   0   0   0   0   0  98  113  128    0    0    0    0    0    0     
 0   0   0   0   0   0  99  114  129    0    0    0    0    0    0
 0   0   0   0   0   0   0    0    0  145  160  175    0    0    0     
 0   0   0   0   0   0   0    0    0  146  161  176    0    0    0     
 0   0   0   0   0   0   0    0    0  147  162  177    0    0    0     
 0   0   0   0   0   0   0    0    0    0    0    0  193  208  223     
 0   0   0   0   0   0   0    0    0    0    0    0  194  209  224     
 0   0   0   0   0   0   0    0    0    0    0    0  195  210  225     

julia> @btime blockmat(A,3)
  207.407 ns (1 allocation: 1.98 KiB)
15×15 Matrix{Int64}:
 1  16  31   0   0   0   0    0    0    0    0    0    0    0    0     
 2  17  32   0   0   0   0    0    0    0    0    0    0    0    0     
 3  18  33   0   0   0   0    0    0    0    0    0    0    0    0     
 0   0   0  49  64  79   0    0    0    0    0    0    0    0    0
 0   0   0  50  65  80   0    0    0    0    0    0    0    0    0     
 0   0   0  51  66  81   0    0    0    0    0    0    0    0    0     
 0   0   0   0   0   0  97  112  127    0    0    0    0    0    0     
 0   0   0   0   0   0  98  113  128    0    0    0    0    0    0     
 0   0   0   0   0   0  99  114  129    0    0    0    0    0    0     
 0   0   0   0   0   0   0    0    0  145  160  175    0    0    0     
 0   0   0   0   0   0   0    0    0  146  161  176    0    0    0     
 0   0   0   0   0   0   0    0    0  147  162  177    0    0    0     
 0   0   0   0   0   0   0    0    0    0    0    0  193  208  223     
 0   0   0   0   0   0   0    0    0    0    0    0  194  209  224     
 0   0   0   0   0   0   0    0    0    0    0    0  195  210  225     

here the version with any residues

function blockmatwr(A,bs)
    n=size(A, 1)   
    B = zero(A);
    doff=1
    soff=1
    foreach(div(n,bs)) do
        for i in 1:bs
            copyto!(B,doff,A,soff,bs)
            doff += n
            soff += n
        end
        doff+=bs
        soff+=bs
    end
    r=%(n,bs)
    foreach(r) do
        copyto!(B,doff,A,soff,r)
        doff += n
        soff += n
    end
    B
end


julia> A=reshape(1:225,15,15);

julia> @btime blockmatwr(A,4)
  201.155 ns (1 allocation: 1.98 KiB)
15×15 Matrix{Int64}:
 1  16  31  46   0   0   0    0    0    0    0    0    0    0    0
 2  17  32  47   0   0   0    0    0    0    0    0    0    0    0
 3  18  33  48   0   0   0    0    0    0    0    0    0    0    0
 4  19  34  49   0   0   0    0    0    0    0    0    0    0    0
 0   0   0   0  65  80  95  110    0    0    0    0    0    0    0
 0   0   0   0  66  81  96  111    0    0    0    0    0    0    0
 0   0   0   0  67  82  97  112    0    0    0    0    0    0    0
 0   0   0   0  68  83  98  113    0    0    0    0    0    0    0
 0   0   0   0   0   0   0    0  129  144  159  174    0    0    0
 0   0   0   0   0   0   0    0  130  145  160  175    0    0    0
 0   0   0   0   0   0   0    0  131  146  161  176    0    0    0
 0   0   0   0   0   0   0    0  132  147  162  177    0    0    0
 0   0   0   0   0   0   0    0    0    0    0    0  193  208  223
 0   0   0   0   0   0   0    0    0    0    0    0  194  209  224
 0   0   0   0   0   0   0    0    0    0    0    0  195  210  225