Search code examples
pythonnumpyscipyleast-squares

Multiple coefficient sets for least squares fitting in numpy/scipy


Is there a way to perform multiple simultaneous (but unrelated) least-squares fits with different coefficient matrices in either numpy.linalg.lstsq or scipy.linalg.lstsq? For example, here is a trivial linear fit that I would like to be able to do with different x-values but the same y-values. Currently, I have to write a loop:

x = np.arange(12.0).reshape(4, 3)
y = np.arange(12.0, step=3.0)
m = np.stack((x, np.broadcast_to(1, x.shape)), axis=0)

fit = np.stack(tuple(np.linalg.lstsq(w, y, rcond=-1)[0] for w in m), axis=-1)

This results in a set of fits with the same slope and different intercepts, such that fit[n] corresponds to coefficients m[n].

Linear least squares is not a great example since it is invertible, and both functions have an option for multiple y-values. However, it serves to illustrate my point.

Ideally, I would like to extend this to any "broadcastable" combination of a and b, where a.shape[-2] == b.shape[0] exactly, and the last dimensions have to either match or be one (or missing). I am not really hung up on which dimension of a is the one representing the different matrices: it was just convenient to make it the first one to shorten the loop.

Is there a built in method in numpy or scipy to avoid the Python loop? I am very much interested in using lstsq rather than manually transposing, multiplying and inverting the matrices.


Solution

  • You could use scipy.sparse.linalg.lsqr together with scipy.sparse.block_diag. I'm just not sure it will be any faster.

    Example:

    >>> import numpy as np
    >>> from scipy.sparse import block_diag
    >>> from scipy.sparse import linalg as sprsla
    >>> 
    >>> x = np.random.random((3,5,4))
    >>> y = np.random.random((3,5))
    >>> 
    >>> for A, b in zip(x, y):
    ...     print(np.linalg.lstsq(A, b))
    ... 
    (array([-0.11536962,  0.22575441,  0.03597646,  0.52014899]), array([0.22232195]), 4, array([2.27188101, 0.69355384, 0.63567141, 0.21700743]))
    (array([-2.36307163,  2.27693405, -1.85653264,  3.63307554]), array([0.04810252]), 4, array([2.61853881, 0.74251282, 0.38701194, 0.06751288]))
    (array([-0.6817038 , -0.02537582,  0.75882223,  0.03190649]), array([0.09892803]), 4, array([2.5094637 , 0.55673403, 0.39252624, 0.18598489]))
    >>> 
    >>> sprsla.lsqr(block_diag(x), y.ravel())
    (array([-0.11536962,  0.22575441,  0.03597646,  0.52014899, -2.36307163,
            2.27693405, -1.85653264,  3.63307554, -0.6817038 , -0.02537582,
            0.75882223,  0.03190649]), 2, 15, 0.6077437777160813, 0.6077437777160813, 6.226368324510392, 106.63227777368986, 1.3277892240815807e-14, 5.36589277249043, array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]))