Векторизованная (частичная) обратные из N * M * M тензора с Numpy

Question

Я почти точно в такой же ситуации, как Аскер здесь более год назад: fast way to invert or dot kxnxn matrix

Так что есть тензор с индексами а [п, I, J] размеров (N, М, М), и я хочу, чтобы инвертировать М * М квадратная матрица часть для каждого п в N.

Например, предположим, что у меня

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]]]) 

Затем инверсия для цикла будет выглядеть следующим образом:

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, но есть еще один способ сделать это в в векторном формате образом прямо сейчас?

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 модифицирует А на месте, чтобы вычислить разложение LU. После замены каждой матрицы N x N с ее разложением LU, pylusolve может быть использован для решения массива M x Nb, представляющего правые части ваших матричных систем. Модифицирует 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 в этом случае.