numpy.linalg.solve with right-hand side of more than three dimensions

问题

I'm trying to solve an equation system with a 3x3 matrix a and a right hand side b of arbitrary shape (3, ...). If b has one or two dimensions, numpy.linalg.solve does the trick. It breaks down for more dimensions though:

import numpy

a = numpy.random.rand(3, 3)

b = numpy.random.rand(3)
numpy.linalg.solve(a, b)  # okay

b = numpy.random.rand(3, 4)
numpy.linalg.solve(a, b)  # okay

b = numpy.random.rand(3, 4, 5)
numpy.linalg.solve(a, b)  # ERR

ValueError: solve: Input operand 1 has a mismatch in its core 
dimension 0, with gufunc signature (m,m),(m,n)->(m,n) (size 5 is 
different from 3)

I would have expected an output array sol of shape (3, 4, 5) with the solution corresponding to the right-hand side b[:, i, j] is sol[:, i, j].

Any hint on how to best work around this?

回答1:

Temporarily reshape b to (3, 20), solve the linear system, and then reshape the resultant array to original shape of b (3, 4, 5):

In [34]: a = numpy.random.rand(3, 3)
In [35]: b = numpy.random.rand(3, 4, 5)

In [36]: x = numpy.linalg.solve(a, b.reshape(b.shape[0], -1)).reshape(b.shape)

Swap the first axis of b with the second using np.swapaxes, solve the linear system, and then restore the axes:

In [58]: x = np.swapaxes(np.linalg.solve(a, np.swapaxes(b, 0, 1)), 0, 1)

Sanity Check:

In [38]: np.einsum('ij,jkl', a, x)
Out[38]: 
array([[[ 0.44859955,  0.22967928,  0.74336067,  0.47440575,  0.53798895],
        [ 0.80045696,  0.54138958,  0.89870834,  0.56862419,  0.28217437],
        [ 0.02093982,  0.78534718,  0.77208236,  0.41568151,  0.95100661],
        [ 0.03820421,  0.47067312,  0.71928294,  0.30852615,  0.64454321]],

       [[ 0.31757072,  0.30527186,  0.36768759,  0.95869289,  0.86601996],
        [ 0.60616508,  0.69927063,  0.53470332,  0.88906606,  0.76066344],
        [ 0.95411847,  0.51116677,  0.29338398,  0.04418815,  0.96210206],
        [ 0.23449429,  0.64159963,  0.7732404 ,  0.4314741 ,  0.81279619]],

       [[ 0.6399571 ,  0.57640652,  0.0186913 ,  0.66304489,  0.83372239],
        [ 0.28426522,  0.62367363,  0.37163699,  0.78217433,  0.90573787],
        [ 0.91066088,  0.06699638,  0.43079394,  0.00263537,  0.399102  ],
        [ 0.17711441,  0.48724858,  0.05526752,  0.34251648,  0.94059739]]])

In [39]: b
Out[39]: 
array([[[ 0.44859955,  0.22967928,  0.74336067,  0.47440575,  0.53798895],
        [ 0.80045696,  0.54138958,  0.89870834,  0.56862419,  0.28217437],
        [ 0.02093982,  0.78534718,  0.77208236,  0.41568151,  0.95100661],
        [ 0.03820421,  0.47067312,  0.71928294,  0.30852615,  0.64454321]],

       [[ 0.31757072,  0.30527186,  0.36768759,  0.95869289,  0.86601996],
        [ 0.60616508,  0.69927063,  0.53470332,  0.88906606,  0.76066344],
        [ 0.95411847,  0.51116677,  0.29338398,  0.04418815,  0.96210206],
        [ 0.23449429,  0.64159963,  0.7732404 ,  0.4314741 ,  0.81279619]],

       [[ 0.6399571 ,  0.57640652,  0.0186913 ,  0.66304489,  0.83372239],
        [ 0.28426522,  0.62367363,  0.37163699,  0.78217433,  0.90573787],
        [ 0.91066088,  0.06699638,  0.43079394,  0.00263537,  0.399102  ],
        [ 0.17711441,  0.48724858,  0.05526752,  0.34251648,  0.94059739]]])

Use np.allclose() so that you don't have to manually going through the numbers and check, particularly for large arrays:

In [32]: b_ = np.einsum('ij,jkl', a, x)

In [33]: np.allclose(b, b_)
Out[33]: True

回答2:

I'd like to add that the manual clearly states:

a : (..., M, M) array_like

Coefficient matrix.

b : {(..., M,), (..., M, K)}, array_like

Ordinate or “dependent variable” values.

So the before last dimension must be the same as the last two dimensions of a (M). Other than that it behaves as you expect - more dimensions are possible, returning a result with the same dimension as B. This way the solution of Ax=B is naturally calculated, and the dimensions automatically converted - just solving many systems of equations with solution of dimension (M,K), and embedding them in the outer dimensions. In your case having 3 at the beginning and not in the middle is confusing the algorithm. Example with 3 dimensions;

>>> a=np.random.rand(9).reshape(3,3)
>>> b=np.random.rand(12).reshape(2,3,2)
>>> np.linalg.solve(a,b)
array([[[-0.63673083,  0.57508091],
        [ 0.87653408,  0.46092677],
        [ 0.61128222, -0.19641607]],

       [[-0.91645601,  1.30939652],
        [ 0.83591936, -0.17006344],
        [ 0.19086912,  0.29082206]]])

来源：https://stackoverflow.com/questions/48387261/numpy-linalg-solve-with-right-hand-side-of-more-than-three-dimensions

标签

python

numpy

linear-algebra