python - 更有效的循环方式？

Question

我有一小段代码来自一个更大的脚本。我发现当函数 t_area 被调用时，它负责大部分的运行时间。我自己测试了这个功能，它并不慢，我相信它需要运行很多次，因为它需要很多时间。这是调用该函数的代码:

tri_area = np.zeros((numx,numy),dtype=float)
for jj in range(0,numy-1):
    for ii in range(0,numx-1):
      xp = x[ii,jj]
      yp = y[ii,jj]
      zp = surface[ii,jj]
      ap = np.array((xp,yp,zp))

      xp = xp+dx
      zp = surface[ii+1,jj]
      bp = np.array((xp,yp,zp))

      yp = yp+dx
      zp = surface[ii+1,jj+1]
      dp = np.array((xp,yp,zp))

      xp = xp-dx
      zp = surface[ii,jj+1]
      cp = np.array((xp,yp,zp))

      tri_area[ii,jj] = t_area(ap,bp,cp,dp)

此处使用的数组大小为 216 x 217，x 和 y 的值也是如此。我是 python 编码的新手，我过去使用过 MATLAB。所以我的问题是，有没有一种方法可以绕过这两个 for 循环，或者有一种更有效的方法来运行这个代码块？寻找任何帮助加快速度!谢谢!

编辑:

感谢大家的帮助，解决了很多困惑。我被问及循环中使用的函数 t_area，这是下面的代码:

def t_area(a,b,c,d):
ab=b-a
ac=c-a
tri_area_a = 0.5*linalg.norm(np.cross(ab,ac))

db=b-d
dc=c-d
tri_area_d = 0.5*linalg.norm(np.cross(db,dc))

ba=a-b
bd=d-b
tri_area_b = 0.5*linalg.norm(np.cross(ba,bd))

ca=a-c
cd=d-c
tri_area_c = 0.5*linalg.norm(np.cross(ca,cd))

av_area = (tri_area_a + tri_area_b + tri_area_c + tri_area_d)*0.5
return(av_area)

很抱歉让我混淆了符号，当时它是有道理的，现在回想起来我可能会改变它。谢谢!

Answer 1

开始前的警告。 range(0, numy-1) 等于 range(numy-1)，生成从 0 到 numy-2 的数字，不包括 numy-1。那是因为你有从 0 到 numy-2 的 numy-1 值。虽然 MATLAB 的索引是从 1 开始的，但 Python 的索引是从 0 开始的，所以在转换时要小心你的索引。考虑到您有 tri_area = np.zeros((numx, numy), dtype=float)，tri_area[ii,jj] 永远不会访问最后一行或最后一列你已经设置了你的循环。因此，我怀疑正确的意图是编写 range(numy)。

由于函数 t_area() 是可矢量化的，因此您可以完全取消循环。矢量化意味着 numpy 通过处理引擎盖下的循环同时对整个数组应用一些操作，这样它们会更快。

首先，我们将每个 (i, j) 元素的所有 ap 堆叠在一个 (m, n, 3) 数组中，其中 (m, n) 是 的大小x。如果我们取两个 (m, n, 3) 数组的叉积，则默认情况下该操作将应用于最后一个轴。这意味着 np.cross(a, b) 将为每个元素 (i, j) 取 a[i,j]< 中的 3 个数字的叉积 和 b[i,j]。类似地，np.linalg.norm(a, axis=2) 将为每个元素 (i, j) 计算 a[i,j 中的 3 个数的范数]。这也将有效地将我们的数组减少到大小 (m, n)。不过这里要小心一点，因为我们需要明确声明我们希望在第二个轴上完成此操作。

请注意，在以下示例中，我的索引关系可能与您的不一致。使这项工作的最低限度是 surface 从 x 和 y 中有一个额外的行和列。

import numpy as np

def _t_area(a, b, c):
    ab = b - a
    ac = c - a
    return 0.5 * np.linalg.norm(np.cross(ab, ac), axis=2)

def t_area(x, y, surface, dx):
    a = np.zeros((x.shape[0], y.shape[0], 3), dtype=float)
    b = np.zeros_like(a)
    c = np.zeros_like(a)
    d = np.zeros_like(a)

    a[...,0] = x
    a[...,1] = y
    a[...,2] = surface[:-1,:-1]

    b[...,0] = x + dx
    b[...,1] = y
    b[...,2] = surface[1:,:-1]

    c[...,0] = x
    c[...,1] = y + dx
    c[...,2] = surface[:-1,1:]

    d[...,0] = bp[...,0]
    d[...,1] = cp[...,1]
    d[...,2] = surface[1:,1:]

    # are you sure you didn't mean 0.25???
    return 0.5 * (_t_area(a, b, c) + _t_area(d, b, c) + _t_area(b, a, d) + _t_area(c, a, d))

nx, ny = 250, 250

dx = np.random.random()
x = np.random.random((nx, ny))
y = np.random.random((nx, ny))
surface = np.random.random((nx+1, ny+1))

tri_area = t_area(x, y, surface, dx)

此示例中的

x 支持索引 0-249，而 surface 支持索引 0-250。 surface[:-1]，surface[0:-1] 的简写，将返回从 0 开始到最后一行的所有行，但不包括它. -1 与 MATLAB 中的 end 具有相同的功能。因此，surface[:-1] 将返回索引 0-249 的行。同样，surface[1:] 将返回索引 1-250 的行，这与您的 surface[ii+1] 相同。

---

注意:在人们知道t_area() 可以完全矢量化之前，我已经写了这一部分。因此，尽管出于此答案的目的，此处的内容已过时，但我会将其保留为遗产，以展示如果函数不可矢量化，本可以进行哪些优化。

不是为每个元素调用函数，这很昂贵，您应该传递给它 x、y,、surface 和 dx 并在内部迭代。这意味着只有一个函数调用和更少的开销。

此外，你不应该为 ap、bp、cp 和 dp 每次循环都创建一个数组，这又增加了开销。在循环外分配它们一次，然后更新它们的值。

最后一个变化应该是循环的顺序。 Numpy 数组默认以行为主(而 MATLAB 以列为主)，因此 ii 作为外层循环表现更好。您不会注意到您的大小数组的差异，但是嘿，为什么不呢？

总的来说，修改后的函数应该是这样的。

def t_area(x, y, surface, dx):
    # I assume numx == x.shape[0]. If not, pass it as an extra argument.
    tri_area = np.zeros(x.shape, dtype=float)

    ap = np.zeros((3,), dtype=float)
    bp = np.zeros_like(ap)
    cp = np.zeros_like(ap)
    dp = np.zeros_like(ap)

    for ii in range(x.shape[0]-1): # do you really want range(numx-1) or just range(numx)?
        for jj in range(x.shape[1]-1):
            xp = x[ii,jj]
            yp = y[ii,jj]
            zp = surface[ii,jj]
            ap[:] = (xp, yp, zp)

            # get `bp`, `cp` and `dp` in a similar manner and compute `tri_area[ii,jj]`