Untitled

import torch
import scipy.sparse

# Example input matrix A
A = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)

# Number of blocks (let's say we want k blocks)
k = 4
n = A.shape[0]  # Size of each block

# Identity matrix I of size n x n
I = torch.eye(n)

# Create blocks: I on the diagonal, -A on the sub-diagonals
diagonal_blocks = [I for _ in range(k)]
off_diagonal_blocks = [-A for _ in range(k - 1)]

# Construct the block-diagonal matrix
upper_blocks = torch.block_diag(*diagonal_blocks)  # Main diagonal with I

# Add -A to sub-diagonal positions
for i in range(1, k):
    upper_blocks[i*n:(i+1)*n, (i-1)*n:i*n] = -A  # Place -A in sub-diagonals

# Print the final matrix
print(upper_blocks)

matrix = upper_blocks.numpy()
csr_matrix = scipy.sparse.csr_matrix(matrix)
print("Data:", csr_matrix.data)
print("Indices:", csr_matrix.indices)
print("Indptr:", csr_matrix.indptr)

coo_matrix = scipy.sparse.coo_matrix(matrix)
print("Row indices:", coo_matrix.row)
print("Column indices:", coo_matrix.col)
print("Data (values):", coo_matrix.data)