带有 numpy 的 N*M*M 张量的矢量化（部分）逆

Question

一年多前，我和这里的提问者几乎完全一样： fast way to invert or dot kxnxn matrix

所以我有一个具有维度 (N,M,M) 的索引 a[n,i,j] 的张量，我想为 N 中的每个 n 反转 M*M 方阵部分。

例如，假设我有

In [1]:    a = np.arange(12)
           a.shape = (3,2,2)
           a

Out[1]: array([[[ 0,  1],
                  [ 2,  3]],

                  [[ 4,  5],
                  [ 6,  7]],

                  [[ 8,  9],
                  [10, 11]]])

然后 for 循环反转将如下所示：

In [2]: inv_a = np.zeros([3,2,2])
        for m in xrange(0,3):
            inv_a[m] = np.linalg.inv(a[m])
        inv_a

Out[2]: array([[[-1.5,  0.5],
                  [ 1. ,  0. ]],

                  [[-3.5,  2.5],
                  [ 3. , -2. ]],

                  [[-5.5,  4.5],
                 [ 5. , -4. ]]])

这显然将在 NumPy 2.0 中实现，根据 github 上的this issue...

我想我需要按照 github 问题线程中提到的 seberg 安装开发版本，但是现在有没有其他方法可以以 vectorized 方式执行此操作？

Answer 1

更新： 在 NumPy 1.8 及更高版本中，numpy.linalg 中的函数是泛化的通用函数。 这意味着您现在可以执行以下操作：

import numpy as np
a = np.random.rand(12, 3, 3)
np.linalg.inv(a)

这将反转每个 3x3 数组并将结果作为 12x3x3 数组返回。 请参阅numpy 1.8 release notes。

---

原答案：

由于N 相对较小，我们如何一次手动计算所有矩阵的 LU 分解。 这确保了所涉及的 for 循环相对较短。

以下是使用普通 NumPy 语法的方法：

import numpy as np
from numpy.random import rand

def pylu3d(A):
    N = A.shape[1]
    for j in xrange(N-1):
        for i in xrange(j+1,N):
            #change to L
            A[:,i,j] /= A[:,j,j]
            #change to U
            A[:,i,j+1:] -= A[:,i,j:j+1] * A[:,j,j+1:]

def pylusolve(A, B):
    N = A.shape[1]
    for j in xrange(N-1):
        for i in xrange(j+1,N):
            B[:,i] -= A[:,i,j] * B[:,j]
    for j in xrange(N-1,-1,-1):
        B[:,j] /= A[:,j,j]
        for i in xrange(j):
            B[:,i] -= A[:,i,j] * B[:,j]

#usage
A = rand(1000000,3,3)
b = rand(3)
b = np.tile(b,(1000000,1))
pylu3d(A)
# A has been replaced with the LU decompositions
pylusolve(A, b)
# b has been replaced to the solutions of
# A[i] x = b[i] for each A[i] and b[i]

正如我所写，pylu3d 修改 A 以计算 LU 分解。 在用其 LU 分解替换每个 NxN 矩阵后，pylusolve 可用于求解代表矩阵系统右侧的 MxN 数组 b。 它会修改b 并进行适当的反向替换以解决系统问题。 正如它所写的那样，这个实现不包括旋转，所以它在数值上不稳定，但在大多数情况下它应该可以很好地工作。

根据您的数组在内存中的排列方式，使用 Cython 可能仍然要快一些。 这里有两个 Cython 函数做同样的事情，但它们首先迭代 M。 它没有矢量化，但速度相对较快。

from numpy cimport ndarray as ar
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def lu3d(ar[double,ndim=3] A):
    cdef int n, i, j, k, N=A.shape[0], h=A.shape[1], w=A.shape[2]
    for n in xrange(N):
        for j in xrange(h-1):
            for i in xrange(j+1,h):
                #change to L
                A[n,i,j] /= A[n,j,j]
                #change to U
                for k in xrange(j+1,w):
                    A[n,i,k] -= A[n,i,j] * A[n,j,k]

@cython.boundscheck(False)
@cython.wraparound(False)
def lusolve(ar[double,ndim=3] A, ar[double,ndim=2] b):
    cdef int n, i, j, N=A.shape[0], h=A.shape[1]
    for n in xrange(N):
        for j in xrange(h-1):
            for i in xrange(j+1,h):
                b[n,i] -= A[n,i,j] * b[n,j]
        for j in xrange(h-1,-1,-1):
            b[n,j] /= A[n,j,j]
            for i in xrange(j):
                b[n,i] -= A[n,i,j] * b[n,j]

您也可以尝试使用 Numba，但在这种情况下我无法让它像 Cython 一样快。