如何用NumPy高效计算大量2D三角形的signed area？

我目前需要批量计算大量2D三角形的有符号面积（signed area），输入是形状为(2, 3, n)的NumPy数组——维度依次对应x/y坐标、三角形的三个顶点、三角形的数量。我已经实现了几种基础计算方法，现在想找到更高效的优化思路。

已实现的计算方法 #

import numpy
import perfplot

def six(p):
    return (
        +p[0][2] * p[1][0]
        + p[0][0] * p[1][1]
        + p[0][1] * p[1][2]
        - p[0][2] * p[1][1]
        - p[0][0] * p[1][2]
        - p[0][1] * p[1][0]
    ) / 2

def mix(p):
    return (
        +p[0][2] * (p[1][0] - p[1][1])
        + p[0][0] * (p[1][1] - p[1][2])
        + p[0][1] * (p[1][2] - p[1][0])
    ) / 2

def mix2(p):
    p1 = p[1] - p[1][[1, 2, 0]]
    return (+p[0][2] * p1[0] + p[0][0] * p1[1] + p[0][1] * p1[2]) / 2

def cross(p):
    e1 = p[:, 1] - p[:, 0]
    e2 = p[:, 2] - p[:, 0]
    return (e1[0] * e2[1] - e1[1] * e2[0]) / 2

def einsum(p):
    return (
        +numpy.einsum("ij,ij->j", p[0][[2, 0, 1]], p[1][[0, 1, 2]])
        - numpy.einsum("ij,ij->j", p[0][[2, 0, 1]], p[1][[1, 2, 0]])
    ) / 2

def einsum2(p):
    return numpy.einsum("ij,ij->j", p[0][[2, 0, 1]], p[1] - p[1][[1, 2, 0]]) / 2

def einsum3(p):
    return (
        numpy.einsum(
            "ij,ij->j", numpy.roll(p[0], 1, axis=0), p[1] - numpy.roll(p[1], 2, axis=0)
        ) / 2
    )

# 性能测试代码
perfplot.save(
    "out.png",
    setup=lambda n: numpy.random.rand(2, 3, n),
    kernels=[six, mix, mix2, cross, einsum, einsum2, einsum3],
    n_range=[2 ** k for k in range(19)],
)

现有方法性能对比 #

更高效的优化方案 #

从性能图能看到，cross和einsum系列的方法已经表现不错，但针对大规模计算，还可以尝试以下几种优化方向：

1. 简化广播实现

利用NumPy的广播特性直接展开有符号面积公式，代码简洁且性能优异：

def broadcast_opt(p):
    x = p[0]
    y = p[1]
    # 核心公式：(x0(y1-y2) + x1(y2-y0) + x2(y0-y1)) / 2
    return (x[0] * (y[1] - y[2]) + x[1] * (y[2] - y[0]) + x[2] * (y[0] - y[1])) / 2

2. 利用Roll简化索引

用numpy.roll循环顶点索引，代码更紧凑，同时保持向量化效率：

def roll_opt(p):
    y_roll_next = numpy.roll(p[1], 1, axis=0)
    y_roll_prev = numpy.roll(p[1], -1, axis=0)
    return numpy.sum(p[0] * (y_roll_next - y_roll_prev), axis=0) / 2

3. Numba JIT并行加速

对于超大规模的三角形计算（n极大），用Numba的JIT编译+并行计算能显著提升速度：

from numba import njit, prange

@njit(parallel=True)
def numba_parallel_opt(p):
    n = p.shape[2]
    res = numpy.empty(n, dtype=p.dtype)
    # 并行遍历每个三角形
    for i in prange(n):
        x0, x1, x2 = p[0, :, i]
        y0, y1, y2 = p[1, :, i]
        res[i] = (x0 * (y1 - y2) + x1 * (y2 - y0) + x2 * (y0 - y1)) / 2
    return res

4. 单Einsum调用优化

把两次Einsum合并成一次，减少函数调用开销：

def einsum_single_opt(p):
    return numpy.einsum(
        "ij,ij->j", 
        p[0], 
        numpy.roll(p[1], 1, axis=0) - numpy.roll(p[1], -1, axis=0)
    ) / 2

这些优化方案中，broadcast_opt和roll_opt保持了纯NumPy向量化的优势，代码简洁且性能接近现有最优的cross方法；numba_parallel_opt在n超过百万级别的场景下会展现出明显的并行加速效果。

内容的提问来源于stack exchange，提问作者Nico Schlömer