How to efficiently create sparse band matrices?

Hi all,

I want to construct the precision matrix for a model that has an AR(1)-like structure. This means that the precision matrix is a sparse band-like matrix and if you were using R, you would use functionality such as bandSparse from the Matrix library.

The matrix should be of the form (sorry, can’t find any LaTeX support):
$$
\begin{bmatrix}
1 & -\alpha & & & \
-\alpha & 1+\alpha^2 & -\alpha & & \
& & \ddots & & \
& & -\alpha & 1+\alpha^2 & -\alpha \
& & & -\alpha & 1
\end{bmatrix}
$$

Below is working code, but as I am sure that you can tell, won’t be the most direct way of creating this matrix. Any thoughts on how to do it better would be appreciated. Thanks!

N = 100
alpha = 0.6

dependence = torch.cat([torch.tensor([1., -alpha, 0]).reshape(1, 3),
                        torch.tensor([-alpha, 1+alpha**2, -alpha]).reshape(1, 3).expand(N-2, 3),
                        torch.tensor([0., -alpha, 1.]).reshape(1, 3)], 0)

band_matrix = torch.zeros(N, N)

total_pad = N - 3
band_matrix[0] = torch.cat((dependence[0], torch.zeros(total_pad)), 0)
for i in range(1, N-1):
    left_pad = i - 1
    right_pad = total_pad - left_pad
    band_matrix[i] = torch.cat((torch.zeros(left_pad), 
                                dependence[i], 
                                torch.zeros(right_pad)), 0)
band_matrix[N-1] = torch.cat((torch.zeros(total_pad), dependence[N-1]), 0)