python - 带有 numpy 的 N*M*M 张量的矢量化(部分)逆

Question

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

所以我有一个索引为 (N,M,M) 的张量，我想为 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 中实现，根据 this issue在 github 上...

我想我需要按照 github 问题线程中提到的 seberg 安装开发版本，但是现在是否有另一种方法可以矢量化方式来完成此操作？

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 一样快。