numpy的polyfit在3D数组上如何快速处理？

Question

我有两个numpy数组，形状为(400, 1000, 1000)。 我想使用这两个数组作为输入，并在轴0上进行拟合，应用numpy.polyfit。我尝试了以下代码，但如果要拟合更大的数组，可能会非常慢。

``` rows = arr1.shape[1] cols = arr1.shape[2] for i in range(rows): for j in range(cols): idx = np.isfinite(arr1[:, i, j]) & np.isfinite(arr2[:, i, j]) # 获取斜率 slope = np.polyfit(arr1[:, i, j][idx], arr2[:, i, j][idx], 1)[0] ... ```

有什么办法可以加快这个过程吗？ 提前感谢您！

- Piotr De

Answer 1

线性回归

关于你的代码示例，我假设你只想进行简单的线性回归，需要线性函数的斜率。

这比一般的多项式拟合问题简单得多。https://en.wikipedia.org/wiki/Simple_linear_regression

在我的电脑上，你的示例大约需要170秒。我想展示两个版本（一个仅使用numpy，需要45秒，另一个使用numba，在我的机器上持续2.5秒（Core i7-4771）。

仅使用numpy（不包括测试数组创建的45秒）

#Creating the Test Arrays
import numpy as np
import time

arr1=np.random.rand(400, 1000, 1000)
arr2=np.random.rand(400, 1000, 1000)

#Test function
def Test(arr1,arr2):
  Slope=np.empty((arr1.shape[1],arr1.shape[2]))
  for i in range(arr1.shape[1]):
    for j in range(arr1.shape[2]):
      idx = np.isfinite(arr1[:, i, j]) & np.isfinite(arr2[:, i, j])
      # Get slope
      Slope[i,j]=lin_regression(arr1[idx, i, j],arr2[idx, i, j])

#Linear regression
def lin_regression(arr1,arr2):
  m_x=np.mean(arr1)
  m_y=np.mean(arr2)

  slope=np.sum((arr1-m_x)*(arr2-m_y))/np.sum((arr1-m_x)**2)
  return slope

使用 numba（2.5秒）

#Creating the Test Arrays
import numba as nb
import numpy as np
import time

arr1=np.random.rand(400, 1000, 1000)
arr2=np.random.rand(400, 1000, 1000)

@nb.njit(fastmath=True)
def lin_regression(arr1,arr2,slope):
  m_x=np.mean(arr1)
  m_y=np.mean(arr2)

  slope=np.sum((arr1-m_x)*(arr2-m_y))/np.sum((arr1-m_x)**2)

@nb.njit(fastmath=True,parallel=True)
def Test(arr1,arr2):
  Slope=np.empty((arr1.shape[1],arr1.shape[2]))
  for i in nb.prange(arr1.shape[1]):
    for j in range(arr1.shape[2]):
      idx = np.isfinite(arr1[:, i, j]) & np.isfinite(arr2[:, i, j])
      # Get slope
      Slope[i,j]=lin_regression(arr1[idx, i, j],arr2[idx, i, j])

- max9111