Convert 2d convolution kernel into linear operator over vectorized image

convolution_matrix.py

import numpy as np
import itertools

def Amat(kernel, size):
    assert kernel.shape[0] == kernel.shape[1]
    ksize = kernel.shape[0]
    assert ksize % 2
    osize = (size + 1 - ksize)
    A = np.zeros((osize ** 2, size ** 2))
    pos = lambda _1, _2, s: _1 * s + _2
    for p in range(size ** 2):
        i = p // size
        j = p % size
        if not (((ksize // 2) <= j < (size - ksize // 2)) and ((ksize // 2) <= i < (size - ksize // 2))):
            continue
        k = i - ksize // 2
        l = j - ksize // 2
        for off1, off2 in itertools.product(range(-(ksize // 2), ksize // 2 + 1),
                                            range(-(ksize // 2), ksize // 2 + 1)):
            A[pos(k, l, osize), pos(i + off1, j + off2, size)] = kernel[off1 + (ksize // 2), off2 + (ksize // 2)]
    return A