如何使用 Numba 堆叠数组并计算逆？

Question

我需要计算一些向量，将它们垂直堆叠在一个数组中，最后计算堆叠向量的逆。我可以使用 numpy 来做到这一点，但为了更好的性能（并将其分配给单独的线程），我想使用 Numba 进行计算。

下面的代码可以在没有

@njit

装饰器的情况下工作，但需要较大的

nRows x ncols x nFrames

值，例如 151 x 151 x 24。

import numpy as np
from numba import njit


@njit
def compute_inverse_numba():
    nRows = 15
    nCols = 15
    nFrames = 2
    result_list = []
    for frame in range(nFrames):
        for row in range(nRows):
            for col in range(nCols):        
                array_of_args = np.random.normal(3, 2.5, size=(10, 3)) #dummy array
                vectors_list = []
                for arg in array_of_args:
                    vec = np.zeros((4, 10), dtype=np.float64)

                    vec[0, 0] = arg[0]
                    vec[0, 1] = arg[1] 
                    vec[0, 2] = arg[2]
                    vec[0, 3] =  arg[0] * arg[2]
                    vec[0, 4] = arg[0] * arg[2]
                    vec[0, 5] = arg[1]
                    vec[0, 6] = arg[0]
                    vec[0, 7] = arg[1]
                    vec[0, 8] = arg[2]
                    vec[0, 9] = 2.0
                    
                    vec[1, 0] = arg[0]
                    vec[1, 3] = 2.0 * arg[1]
                    vec[1, 4] = arg[2]
                    vec[1, 6] = 1.0
                    
                    vec[2, 1] = arg[1]
                    vec[2, 3] =  arg[0]
                    vec[2, 5] = 2.0 * arg[2]
                    vec[2, 7] = 1.0
                    
                    vec[3, 2] =  arg[2]
                    vec[3, 4] =  arg[0]
                    vec[3, 5] = 3.0 * arg[1]
                    vec[3, 8] = 1.0

                    vectors_list.append(vec)
                
                # vertically stack the results
                vectors_list = np.vstack(vectors_list)
                
                # compute inverse matrix
                inv = np.linalg.pinv(vectors_list)
                result_list.append(inv)

    return result_list


###########----main()
compute_inverse_numba()

print()

问题：使用

@njit

装饰器，我不断收到与堆叠相关的步骤的异常。我尝试了不同的选择，但没有任何效果。

Answer 1

您可以以与 Numba 兼容的方式手动堆叠数组

import numpy as np
from numba import njit

@njit
def compute_inverse_numba():
nRows = 15
nCols = 15
nFrames = 2
result_list = []
for frame in range(nFrames):
    for row in range(nRows):
        for col in range(nCols):        
            array_of_args = np.random.normal(3, 2.5, size=(10, 3)) #dummy array
            vectors_list = []
            for arg in array_of_args:
                vec = np.zeros((4, 10), dtype=np.float64)

                # ... [Fill vec as before] ...

                vectors_list.append(vec)

            # Manually stack arrays
            total_rows = sum(v.shape[0] for v in vectors_list)
            stacked_array = np.empty((total_rows, 10), dtype=np.float64)
            row_start = 0
            for v in vectors_list:
                row_end = row_start + v.shape[0]
                stacked_array[row_start:row_end, :] = v
                row_start = row_end
            
            # compute inverse matrix
            inv = np.linalg.pinv(stacked_array)
            result_list.append(inv)

return result_list

 result = compute_inverse_numba()
 print(result)