Recover sparse matrix from linear operator

Question

For many linear operators, it is easy to express them as a function. For example, to blur an image, simply compute a weighted sum of neighboring pixels (or even easier, import gaussian_filter from scipy.ndimage). To reverse this process, it is useful to represent this function/linear operator as a matrix, because matrices can be inverted easily. Such a matrix can be obtained by applying the linear operator to all unit vectors and stacking the results row-wise, which is slow with large matrices. Is there a faster way to recover sparse matrices from linear operators, perhaps by using the sparsity property? If not, how about an approximation of the sparse matrix?

For simple problems like Gaussian blur, it is of course possible to construct the linear operator by hand, but for more complex operators, it might be very difficult. I am interested in a general solution for sparse operators.

The following code illustrates deblurring of an image. The problem is that the function recover_matrix_from_operator is excruciatingly slow for large images.

import numpy as np
import scipy.ndimage
import scipy.linalg
import matplotlib.pyplot as plt

def lin_op(x):
    # Linear operator which blurs an image
    return scipy.ndimage.gaussian_filter(x.reshape(25, 17), 1.0).flatten()

def recover_matrix_from_operator(lin_op, n):
    # Recover matrix from linear operator using unit vectors (slow)
    return np.array([lin_op(unit_vector) for unit_vector in np.eye(n)])

# Draw a stick figure
img = np.bincount(np.cumsum([
    58,1,1,14,4,13,4,9,4,4,4,6,4,1,1,4,8,4,4,10,3,3,12,
    2,2,14,1,1,16,17,17,17,17,16,2,14,4,12,6,10,8,8,10,
]), minlength=25*17).reshape(25, 17) * 1.0

# Apply linear operator to image to blur it
blurred = lin_op(img).reshape(img.shape)

# Recover matrix A from linear operator
A = recover_matrix_from_operator(lin_op, img.size)

# Verify that lin_op(x) == A x
blurred_using_A = (A @ img.ravel()).reshape(img.shape)
assert np.allclose(blurred, blurred_using_A)

# Deblur image
deblurred = scipy.linalg.lstsq(A, blurred.ravel())[0].reshape(img.shape)

# Plot all the images
plt.figure(figsize=(10, 4))
for i, tmp_img in enumerate([img, blurred, blurred_using_A, deblurred]):
    plt.subplot(1, 4, 1 + i)
    plt.title(["image", "blurred", "blurred using A", "deblurred"][i])
    plt.imshow(tmp_img, cmap='gray', vmin=0, vmax=1)
    plt.axis("off")
plt.show()

Answer 1

If the goal is to construct the sparse matrix from any arbitrary linear operator treated as a black box which only accepts dense vectors (which is the case for ndimage stuff), then it's possible to just convert as we go

def recover_matrix_from_operator(lin_op, n):
    # Recover matrix from linear operator using unit vectors (slow)
    return sp.sparse.vstack(sp.sparse.csc(lin_op(unit_vector.toarray())) for unit_vector in sp.sparse.eye(n))

though this is likely still awfully slow, but at least memory consumption will be kept in check. If the algorithm isn't a black box, then we could just assemble it into a sparse matrix directly, which wouldn't be nearly as bad.

Or matrix free like @lightxbulb suggests;

import scipy.sparse.linalg

Aop = scipy.sparse.linalg.LinearOperator((25*17, 25*17), matvec=lin_op)
# A = Aop @ np.eye(25*17) # you can actually get the dense form out like this
# A_sparse = Aop @ scipy.sparse.eye(25*17) # would work if the operator could handle sparse inputs

blurred_using_Aop = (Aop @ img.ravel()).reshape(img.shape)
deblurred_using_Aop = scipy.sparse.linalg.cg(A, blurred.ravel())[0].reshape(img.shape)

plt.figure(figsize=(10, 4))
for i, tmp_img in enumerate([img, blurred, blurred_using_Aop, deblurred_using_Aop]):
    plt.subplot(1, 4, 1 + i)
    plt.title(["image", "blurred", "blurred using Aop", "deblurred using Aop"][i])
    plt.imshow(tmp_img, cmap='gray', vmin=0, vmax=1)
    plt.axis("off")

plt.show()

For GPU, i did spot the existance of cupyx.scipy.sparse.linalg.LinearOperator, so it seems like it could be trivial to port it to GPUs.

Answer 2

The following solution will not work for unstructured sparse matrices, but when each row of the linear operator corresponds to a relatively small filter kernel, multiple rows of the sparse matrix can be extracted at once. The trick is to multiply with the sum of multiple unit vectors such that the resulting filter kernels do not overlap. For the blur example above, 6 kernels can be extracted at once without loss of accuracy (or 12 with a negligible loss of accuracy):

import numpy as np
import scipy.ndimage
import scipy.sparse.linalg
import matplotlib.pyplot as plt

def matvec(x):
    # Linear operator which blurs an image
    return scipy.ndimage.gaussian_filter(x.reshape(25, 17), 1.0).flatten()

def recover_matrix(matvec, image_shape, radius=None):
    # Recover sparse matrix from linear operator assuming that it represents a filter kernel of fixed size.
    # Size of filter kernel is 2 * radius + 1. For example, for a 5x5 blur, the radius is 2.
    h, w = image_shape
    if radius is None:
        radius = estimate_kernel_radius(matvec, image_shape)
    step = 2 * radius + 1
    i_inds, j_inds, values = [], [], []
    dy, dx = np.mgrid[-radius:radius+1, -radius:radius+1]

    for y in range(step):
        for x in range(step):
            # Create unit vectors so that the kernels do not overlap
            unit_vectors = np.zeros((h, w))
            unit_vectors[y:y+h:step, x:x+w:step] = 1

            # Extract multiple kernels at once
            kernels = matvec(unit_vectors.ravel()).reshape(h, w)

            # Coordinates of 1s in unit vectors
            y1, x1 = np.mgrid[y:h:step, x:w:step]

            # Coordinates of neighboring pixels within radius
            x2 = (x1.reshape(-1, 1) + dx.reshape(1, -1)).ravel()
            y2 = (y1.reshape(-1, 1) + dy.reshape(1, -1)).ravel()

            # Only consider indices which are inside the image
            is_valid = (0 <= x2) & (x2 < w) & (0 <= y2) & (y2 < h)

            x1 = np.repeat(x1.ravel(), dx.size)[is_valid]
            y1 = np.repeat(y1.ravel(), dy.size)[is_valid]

            # Indices of coordinates within matrix
            i = x1 + y1 * w
            j = x2[is_valid] + y2[is_valid] * w

            i_inds.append(i)
            j_inds.append(j)
            values.append(kernels.ravel()[j])

    # Assemble sparse matrices from coordinates and values
    i_inds = np.concatenate(i_inds)
    j_inds = np.concatenate(j_inds)
    values = np.concatenate(values)
    return scipy.sparse.csr_matrix((values, (i_inds, j_inds)), shape=(w * h, w * h))

def estimate_kernel_radius(matvec, image_shape):
    unit_vector = np.zeros(image_shape)
    y = image_shape[0] // 2
    x = image_shape[1] // 2
    unit_vector[y, x] = 1
    kernel = matvec(unit_vector.ravel()).reshape(image_shape)

    for radius in range(1, min(image_shape)):
        if x - radius < 0 or y - radius < 0: break
        if np.any(kernel[y - radius, x-radius:x+radius+1]): continue
        if np.any(kernel[y + radius, x-radius:x+radius+1]): continue
        if np.any(kernel[y-radius:y+radius+1, x - radius]): continue
        if np.any(kernel[y-radius:y+radius+1, x + radius]): continue
        return radius - 1

    raise ValueError(f"Kernel too large")

# Draw a stick figure
img = np.bincount(np.cumsum([
    58,1,1,14,4,13,4,9,4,4,4,6,4,1,1,4,8,4,4,10,3,3,12,
    2,2,14,1,1,16,17,17,17,17,16,2,14,4,12,6,10,8,8,10,
]), minlength=25*17).reshape(25, 17) * 1.0

# Apply linear operator to image to blur it
blurred = matvec(img).reshape(img.shape)

# Recover matrix A from linear operator
A = recover_matrix(matvec, img.shape)

# Verify that matvec(x) == A x
blurred_using_A = (A @ img.ravel()).reshape(img.shape)
#assert np.allclose(blurred, blurred_using_A)

# Deblur image
deblurred = scipy.sparse.linalg.cg(A, blurred.ravel())[0].reshape(img.shape)

# Plot all the images
plt.figure(figsize=(10, 4))
for i, tmp_img in enumerate([img, blurred, blurred_using_A, deblurred]):
    plt.subplot(1, 4, 1 + i)
    plt.title(["image", "blurred", "blurred using A", "deblurred"][i])
    plt.imshow(tmp_img, cmap='gray', vmin=0, vmax=1)
    plt.axis("off")
plt.show()

When the kernel has uniform values everywhere, life is much easier. It is sufficient to extract the kernel for any unit vector and then shift it around a bunch.

import numpy as np
import scipy.ndimage
import scipy.sparse.linalg
import matplotlib.pyplot as plt

# Draw a stick figure
img = np.bincount(np.cumsum([
    58,1,1,14,4,13,4,9,4,4,4,6,4,1,1,4,8,4,4,10,3,3,12,
    2,2,14,1,1,16,17,17,17,17,16,2,14,4,12,6,10,8,8,10,
]), minlength=25*17).reshape(25, 17) * 1.0

def matvec(x):
    # Linear operator which blurs an image
    return scipy.ndimage.gaussian_filter(x.reshape(img.shape), 1.0).flatten()

def recover_matrix_assuming_uniform_kernel(matvec, image_shape):
    h, w = image_shape

    # Measure kernel at center of image
    cx = h // 2
    cy = w // 2
    unit_vector = np.zeros((h, w))
    unit_vector[cy, cx] = 1
    kernel = matvec(unit_vector.ravel()).reshape(h, w)

    # Find coordinates of non-zero values
    mask = kernel != 0.0

    dy, dx = np.mgrid[:h, :w]
    dx = dx[mask] - cx
    dy = dy[mask] - cy

    y, x = np.mgrid[:h, :w]

    x = (x.reshape(1, -1) + dx.reshape(-1, 1)).ravel()
    y = (y.reshape(1, -1) + dy.reshape(-1, 1)).ravel()

    # Reflective boundary conditions
    x[x < 0] = -1 - x[x < 0]
    y[y < 0] = -1 - y[y < 0]
    y[y >= h] = 2 * h - 1 - y[y >= h]
    x[x >= w] = 2 * w - 1 - x[x >= w]

    # Build sparse matrix from coordinates and values
    i = np.tile(np.arange(w * h), dx.size)
    j = x + y * w
    data = np.repeat(kernel[mask].ravel(), w * h)
    return scipy.sparse.csr_matrix((data, (i, j)), shape=(w * h, w * h))

def main():
    A = recover_matrix_assuming_uniform_kernel(matvec, img.shape)

    # Apply linear operator to image to blur it
    blurred = matvec(img).reshape(img.shape)

    # Deblur image
    deblurred = scipy.sparse.linalg.cg(A, blurred.ravel())[0].reshape(img.shape)

    plt.subplot(1, 3, 1)
    plt.imshow(img, cmap="gray")
    plt.subplot(1, 3, 2)
    plt.imshow(blurred, cmap="gray")
    plt.subplot(1, 3, 3)
    plt.imshow(deblurred, cmap="gray", vmin=0, vmax=1)
    plt.show()

if __name__ == "__main__":
    main()