如何使用 Numpy 对带有多个“if”语句的嵌套“for”循环进行矢量化？

Question

我有一个简单的 2D 光线投射例程，一旦障碍物数量增加，它就会变得非常缓慢。

此例程由以下部分组成：

• 2 for 循环（外循环遍历每个射线/方向，然后内循环遍历每个线障碍物）

• 多个if 语句（检查一个值是否大于另一个值或数组是否为空）

问题：如何使用 Numpy 将所有这些操作压缩为 1 个矢量化指令块？

更具体地说，我面临两个问题：

• 我已经设法矢量化内环（光线和每个障碍物之间的交点），但我无法一次对所有光线运行此操作。

• 我发现处理if 语句的唯一解决方法是使用掩码数组。有些东西告诉我，在这种情况下处理这些语句不是正确的方法（看起来很笨拙、繁琐且不合 Python）

原始代码：

from math import radians, cos, sin
import matplotlib.pyplot as plt
import numpy as np

N = 10  # dimensions of canvas (NxN)
sides = np.array([[0, N, 0, 0], [0, N, N, N], [0, 0, 0, N], [N, N, 0, N]])
edges = np.random.rand(5, 4) * N # coordinates of 5 random segments (x1, x2, y1, y2)
edges = np.concatenate((edges, sides))
center = np.array([N/2, N/2]) # coordinates of center point
directions = np.array([(cos(radians(a)), sin(radians(a))) for a in range(0, 360, 10)]) # vectors pointing in all directions
intersections = []

# for each direction
for d in directions:
    
    min_dist = float('inf')
    
    # for each edge
    for e in edges:
        p1x, p1y = e[0], e[2]
        p2x, p2y = e[1], e[3]
        p3x, p3y = center
        p4x, p4y = center + d
        
        # find intersection point
        den = (p1x - p2x) * (p3y - p4y) - (p1y - p2y) * (p3x - p4x)
        
        if den:
            t = ((p1x - p3x) * (p3y - p4y) - (p1y - p3y) * (p3x - p4x)) / den
            u = -((p1x - p2x) * (p1y - p3y) - (p1y - p2y) * (p1x - p3x)) / den
            
            # if any:
            if t > 0 and t < 1 and u > 0:
                sx = p1x + t * (p2x - p1x)
                sy = p1y + t * (p2y - p1y)
                
                isec = np.array([sx, sy])           
                dist = np.linalg.norm(isec-center) 
                
                # make sure to select the nearest one (from center)
                if dist < min_dist:
                    min_dist = dist
                    nearest = isec
    
    # store nearest interesection point for each ray
    intersections.append(nearest)

    
# Render
plt.axis('off')

for x, y in zip(edges[:,:2], edges[:,2:]):
    plt.plot(x, y)
    
for isec in np.array(intersections):
    plt.plot((center[0], isec[0]), (center[1], isec[1]), '--', color="#aaaaaa", linewidth=.8)

矢量化版本（尝试）：

from math import radians, cos, sin
import matplotlib.pyplot as plt
from scipy import spatial
import numpy as np

N = 10  # dimensions of canvas (NxN)
sides = np.array([[0, N, 0, 0], [0, N, N, N], [0, 0, 0, N], [N, N, 0, N]])
edges = np.random.rand(5, 4) * N # coordinates of 5 random segments (x1, x2, y1, y2)
edges = np.concatenate((edges, sides))
center = np.array([N/2, N/2]) # coordinates of center point
directions = np.array([(cos(radians(a)), sin(radians(a))) for a in range(0, 360, 10)]) # vectors pointing in all directions
intersections = []

# Render edges
plt.axis('off')
for x, y in zip(edges[:,:2], edges[:,2:]):
    plt.plot(x, y)
    
    
# for each direction
for d in directions:

    p1x, p1y = edges[:,0], edges[:,2]
    p2x, p2y = edges[:,1], edges[:,3]
    p3x, p3y = center
    p4x, p4y = center + d

    # denominator
    den = (p1x - p2x) * (p3y - p4y) - (p1y - p2y) * (p3x - p4x)

    # first 'if' statement -> if den > 0
    mask = den > 0
    den = den[mask]
    p1x = p1x[mask]
    p1y = p1y[mask]
    p2x = p2x[mask]
    p2y = p2y[mask]

    t = ((p1x - p3x) * (p3y - p4y) - (p1y - p3y) * (p3x - p4x)) / den
    u = -((p1x - p2x) * (p1y - p3y) - (p1y - p2y) * (p1x - p3x)) / den

    # second 'if' statement -> if (t>0) & (t<1) & (u>0)
    mask2 = (t > 0) & (t < 1) & (u > 0)
    t = t[mask2]
    p1x = p1x[mask2]
    p1y = p1y[mask2]
    p2x = p2x[mask2]
    p2y = p2y[mask2]
    
    # x, y coordinates of all intersection points in the current direction
    sx = p1x + t * (p2x - p1x)
    sy = p1y + t * (p2y - p1y)
    pts = np.c_[sx, sy]
    
    # if any:
    if pts.size > 0:
        
        # find nearest intersection point
        tree = spatial.KDTree(pts)
        nearest = pts[tree.query(center)[1]]

        # Render
        plt.plot((center[0], nearest[0]), (center[1], nearest[1]), '--', color="#aaaaaa", linewidth=.8)

Answer 1

问题的重新表述——寻找线段和线射线的交点

让q 和q2 成为段的端点（障碍物）。为方便起见，让我们定义一个类来表示平面中的点和向量。除了通常的操作，向量乘法由u × v = u.x * v.y - u.y * v.x定义。

注意：这里Coord(2, 1) * 3 返回Coord(6, 3)，而Coord(2, 1) * Coord(-1, 4) 输出9。为了避免这种混淆，可以将* 限制为标量乘法，并通过__xor__ 使用^ 作为向量乘法。

class Coord:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    @property
    def radius(self):
        return np.sqrt(self.x ** 2 + self.y ** 2)

    def _cross_product(self, other):
        assert isinstance(other, Coord)
        return self.x * other.y - self.y * other.x

    def __mul__(self, other):
        if isinstance(other, Coord):
            # 2D "cross"-product
            return self._cross_product(other)
        elif isinstance(other, int) or isinstance(other, float):
            # scalar multiplication
            return Coord(self.x * other, self.y * other)
    
    def __rmul__(self, other):
        return self * other

    def __sub__(self, other):
        return Coord(self.x - other.x, self.y - other.y)

    def __add__(self, other):
        return Coord(self.x + other.x, self.y + other.y)

    def __repr__(self):
        return f"Coord({self.x}, {self.y})"

现在，我发现在极坐标中处理射线更容易：对于给定角度theta（方向），目标是确定它是否与线段相交，如果是，则确定相应的半径。这是一个找到它的函数。请参阅here，了解原因和方式。我尝试使用与上一个链接中相同的变量名称。

def find_intersect_btw_ray_and_sgmt(q, q2, theta):
    """
    Args:
        q (Coord): first endpoint of the segment
        q2 (Coord): second endpoint of the segment
        theta (float): angle of the ray

    Returns:
        (float): np.inf if the ray does not intersect the segment,
                 the distance from the origin of the intersection otherwise
   """
    assert isinstance(q, Coord) and isinstance(q2, Coord)
    s = q2 - q
    r = Coord(np.cos(theta), np.sin(theta))
    cross = r * s  # 2d cross-product
    t_num = q * s
    u_num = q * r
    ## the intersection point is roughly at a distance t_num / cross
    ## from the origin. But some cases must be checked beforehand.

    ## (1) the segment [PQ2] is aligned with the ray
    if np.isclose(cross, 0) and np.isclose(u_num, 0):
        return min(q.radius, q2.radius)
    ## (2) the segment [PQ2] is parallel with the ray
    elif np.isclose(cross, 0):
        return np.inf
   
    t, u = t_num / cross, u_num / cross
    ## There is actually an intersection point
    if t >= 0 and 0 <= u <= 1:
        return t
   
    ## (3) No intersection point
    return np.inf

例如find_intersect_btw_ray_and_sgmt(Coord(1, 2), Coord(-1, 2), np.pi / 2) 应该返回2。

请注意，这里为简单起见，我只考虑了光线原点位于Coord(0, 0) 的情况。这可以通过设置t_num = (q - origin) * s 和u_num = (q - origin) * r 轻松扩展到一般情况。

让我们矢量化它！

这里非常有趣的是Coord 类中定义的操作也适用于 x 和 y 是 numpy 数组的情况！因此，对Coord(np.array([1, 2, 0]), np.array([2, -1, 3])) 应用任何定义的操作，将其逐元素应用于点(1, 2), (2, -1) and (0, 3)。 Coord 的操作因此已经被向量化了。构造函数可以修改为：

    def __init__(self, x, y):
        x, y = np.array(x), np.array(y)
        assert x.shape == y.shape
        self.x, self.y = x, y
        self.shape = x.shape

现在，我们希望函数find_intersect_btw_ray_and_sgmt 能够处理参数q 和q2 包含端点序列的情况。在进行完整性检查之前，所有操作都正常工作，因为正如我们所提到的，它们已经被矢量化了。正如您提到的，条件语句可以使用掩码“矢量化”。这是我的建议：

def find_intersect_btw_ray_and_sgmts(q, q2, theta):
    assert isinstance(q, Coord) and isinstance(q2, Coord)
    assert q.shape == q2.shape
    EPS = 1e-14
    s = q2 - q
    r = Coord(np.cos(theta), np.sin(theta))
    cross = r * s
   
    cross_sign = np.sign(cross)
    cross = cross * cross_sign
    t_num = (q * s) * cross_sign
    u_num = (q * r) * cross_sign

    radii = np.zeros_like(t_num)
    mask = ~np.isclose(cross, 0) & (t_num >= -EPS) & (-EPS <= u_num) & (u_num <= cross + EPS)

    radii[~mask] = np.inf  # no intersection
    radii[mask] = t_num[mask] / cross[mask]  # intersection
    return radii

请注意，cross、t_num 和 u_num 乘以 cross 的符号，以确保除以 cross 时保持除法的符号。因此((t_num >= 0) & (cross >= 0)) | ((t_num <= 0) & (cross <= 0)) 形式的条件可以替换为(t_num >= 0)。

为简单起见，我们省略了半径和线段对齐的情况 (1) ((cross == 0) & (u_num == 0))。这可以通过仔细添加第二个掩码来合并。

对于给定的 theta 值，我们能够确定相应的光线是否同时与多个线段相交。

## Some useful functions
def polar_to_cartesian(r, theta):
    return Coord(r * np.cos(theta), r * np.sin(theta))

def plot_segments(p, q, *args, **kwargs):
    plt.plot([p.x, q.x], [p.y, q.y], *args, **kwargs)

def plot_rays(radii, thetas, *args, **kwargs):
    endpoints = polar_to_cartesian(radii, thetas)
    n = endpoints.shape
    origin = Coord(np.zeros(n), np.zeros(n))
    plot_segments(origin, endpoints, *args, **kwargs)

## Data generation
M = 5  # size of the canvas
N = 10  # number of segments
K = 16  # number of rays

q = Coord(*np.random.uniform(-M/2, M/2, size=(2, N)))
p = q + Coord(*np.random.uniform(-M/2, M/2, size=(2, N)))
thetas = np.linspace(0, 2 * np.pi, K, endpoint=False)

## For each ray, find the minimal distance of intersection
## with all segments
plt.figure(figsize=(5, 5))
plot_segments(p, q, "royalblue", marker=".")

for theta in thetas:
    radii = find_intersect_btw_ray_and_sgmts(p, q, theta)
    radius = np.min(radii)
    if not np.isinf(radius):
        plot_rays(radius, theta, color="orange")
    else:
        plot_rays(2*M, theta, ':', c='orange')

plt.plot(0, 0, 'kx')
plt.xlim(-M, M)
plt.ylim(-M, M)

这还不是全部！由于 python 的广播，可以避免对 theta 值的迭代。例如，回想一下np.array([1, 2, 3]) * np.array([[1], [2], [3], [4]]) 生成大小为 4 × 3 的成对乘积矩阵。以同样的方式Coord([[5],[7]], [[5],[1]]) * Coord([2, 4, 6], [-2, 4, 0]) 输出一个 2 × 3 矩阵，其中包含向量 (5, 5)、(7, 1) 和 (2, -2)、(4, 4)、(6, 0) 之间的所有成对叉积)。

最后，可以通过以下方式确定交点：

radii_all = find_intersect_btw_ray_and_sgmts(p, q, np.vstack(thetas))
# p and q have a shape of (N,) and np.vstack(thetas) of (K, 1)
# this radii_all have a shape of (K, N)
# radii_all[k, n] contains the distance from the origin of the intersection
# between k-th ray and n-th segment (or np.inf if there is no intersection point)
radii = np.min(radii_all, axis=1)
# radii[k] contains the distance from the origin of the closest intersection
# between k-th ray and all segments


do_intersect = ~np.isinf(radii)
plot_rays(radii[do_intersect], thetas[do_intersect], color="orange")
plot_rays(2*M, thetas[~do_intersect], ":", color="orange")