Python equivalent of MATLAB's lsqr() with first argument a function

Question

As per title, I am looking for a Python equivalent of MATLAB's lsqr() (possibly in NumPy / SciPy) when the first argument is a function.

Briefly, lsqr solves numerically for x the following problem:

argmin_x || A*x - b ||_2

where x and b are vectors (of potentially different size) and A is a linear operator.

I believe that for numerical input the equivalent is numpy.linalg.lstsq().

The function scipy.optimize.least_squares() can, in principle, be used for solving the argmin problem, but it seems that it uses a different (and much slower) algorithm internally, which seems unsuited for optimization over relatively large inputs.

I believe that lsqr() internally uses A*x and A'*b and does not require an explicit rappresentation of A.

So, is there an equivalent of MATLAB's lsqr (with first argument a function)?

Answer 1

For large and sparse inputs (which would be the use case for lsqr anyway), the Python / SciPy equivalent of MATLAB's lsqr is:

scipy.sparse.linalg.lsqr()

The first argument of this function can be scipy.sparse.linalg.LinearOperator(), which is a proxy for the linear operator where A*x and A'*b (' is the transpose operator) must be provided as the callable corresponding to matvec and rmatvec (respectively).

This can ultimately be used to compute lsqr where A is not known explicitly.

For example:

def Ax(x):
    """Returns A*x"""
    ...

def Atb(b):
    """Returns A'*b"""
    ...

A = scipy.sparse.linalg.LinearOperator((m, n), matvec=Ax, rmatvec=Atb)
result = scipy.sparse.linalg.lsqr(A, b)

Note that both MATLAB and Python documentations of lsqr indicate that A'*x (more precisely A^T x in the case of Python, but with the same meaning) is computed, but this is not (and cannot be) correct. If fact they both use x as a mute variable (not related to Ax = b naming), but they both actually use b.

An important difference exists between Python and MATLAB implementations:

• MATLAB: a single function is provided and it needs to compute A*x or A'*b depending on the second argument of afun (afun(x,'notransp') or afun(x,'transp'), respectively).
• Python: the two functions are to be provided separately, and the first argument is either x or b, depending on whether the A.matvec() or A.rmatvec() (respectively) is called.

(This is based on the very informative answer from @AnderBiguri and scipy.sparse.linalg.lsqr() source code).