在 3d 数组上快速处理 numpy polyfit？

Question

我有两个形状为 (400, 1000, 1000) 的 numpy 数组。 我想使用两个数组作为输入并在轴 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]) # Get slope slope = np.polyfit(arr1[:, i, j][idx], arr2[:, i, j][idx], 1)[0] ...

知道如何加快速度吗？ 提前致谢！

Answer 1

线性回归

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

这是一个比一般 polyfit 简单得多的问题。 https://en.wikipedia.org/wiki/Simple_linear_regression

您的示例在我的 PC 上使用了大约 170s。我不想展示两个版本（只有 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])