带多边形掩码的2D向量场构建：求点到最近顶点的实现方法

我正在尝试创建一个带多边形“掩码”的2D向量场，要求每个点的向量角度和幅值由该点到多边形最近顶点的距离决定。目前已经能通过matplotlib生成多边形、获取顶点并构建基础插值器，但卡在计算任意坐标点对应的多边形最近顶点这一步，猜测可以用numpy/scipy的相关工具实现。

我预想的核心逻辑伪代码如下：

for point in all_points:
    nearest = NearestPoint(polygon_vertices, point)
    value = point - nearest  # 向量分量，或基于距离的自定义函数
    angle = Angle(nearest, point)  # 向量角度

以下是我已完成的基础代码：

import numpy as np
import scipy.interpolate
import matplotlib.pyplot as plt

def CircleVertices():
    circ = plt.Circle((0.5, 0.5), radius=0.4, edgecolor='b', facecolor='None')
    return circ.get_path().vertices

def NearestNeighborInterpolate():
    vertices = CircleVertices()
    x, y = vertices[:,0], vertices[:,1]
    z = [np.hypot(x, y) for x,y in vertices]

    X = np.linspace(min(x)-0.5, max(x)+0.5, num=100)
    Y = np.linspace(min(y)-0.5, max(y)+0.5, num=100)

    X, Y = np.meshgrid(X, Y)
    interp = scipy.interpolate.NearestNDInterpolator(list(zip(x, y)), z)
    Z = interp(X, Y)

    plt.pcolormesh(X, Y, Z, shading='auto')
    plt.plot(x, y, "ok", label="input point")

    plt.legend()
    plt.colorbar()
    plt.axis("equal")
    plt.show()

def GriddataInterpolate():
    vertices = CircleVertices()
    x, y = vertices[:,0], vertices[:,1]

    xx = np.linspace(-1, 1, 20)
    yy = np.linspace(-1, 1, 20)
    xx, yy = np.meshgrid(xx, yy)

    u = [np.hypot(x, y) for x,y in vertices]
    v = [0.5 for _ in x]

    u_interp = scipy.interpolate.griddata(vertices, u, (xx, yy), method='linear')
    v_interp = scipy.interpolate.griddata(vertices, v, (xx, yy), method='linear')

    plt.quiver(xx, yy, u_interp, v_interp)
    plt.show()

以上代码已生成两张可视化图：

• 最近邻插值的颜色分布图
• 线性插值的向量场图

解决方案：用scipy.spatial.cKDTree快速查询最近顶点 #

核心难点任意点到多边形最近顶点的批量计算可以通过scipy.spatial.cKDTree高效实现，它能一次性处理所有网格点的近邻查询，比循环遍历效率高得多。

完整实现代码

import numpy as np
import scipy.spatial
import matplotlib.pyplot as plt
from matplotlib.path import Path

def CircleVertices():
    circ = plt.Circle((0.5, 0.5), radius=0.4, edgecolor='b', facecolor='None')
    return circ.get_path().vertices

def CreateMaskedVectorField():
    # 1. 获取多边形顶点并修正坐标（plt.Circle的path顶点为相对坐标）
    polygon_vertices = CircleVertices()
    polygon_vertices += np.array([0.5, 0.5])

    # 2. 构建KD树用于快速近邻查询
    kdtree = scipy.spatial.cKDTree(polygon_vertices)

    # 3. 生成目标网格点
    xx, yy = np.meshgrid(np.linspace(-0.2, 1.2, 50), np.linspace(-0.2, 1.2, 50))
    grid_points = np.column_stack((xx.ravel(), yy.ravel()))

    # 4. 批量查询每个网格点的最近顶点及距离
    distances, nearest_indices = kdtree.query(grid_points, k=1)
    nearest_vertices = polygon_vertices[nearest_indices]

    # 5. 计算向量场分量（方向：从最近顶点指向当前点；幅值：到顶点的距离）
    u = grid_points[:, 0] - nearest_vertices[:, 0]
    v = grid_points[:, 1] - nearest_vertices[:, 1]

    # 6. 实现多边形掩码：仅保留内部点的向量
    polygon_path = Path(polygon_vertices)
    mask = polygon_path.contains_points(grid_points).reshape(xx.shape)
    u = u.reshape(xx.shape)
    v = v.reshape(yy.shape)
    u[~mask] = 0
    v[~mask] = 0

    # 7. 可视化向量场
    plt.figure(figsize=(8,8))
    plt.quiver(xx, yy, u, v, distances.reshape(xx.shape), cmap='viridis', scale=20)
    plt.plot(polygon_vertices[:,0], polygon_vertices[:,1], 'ok', label='多边形顶点')
    plt.legend()
    plt.axis('equal')
    plt.colorbar(label='到最近顶点的距离')
    plt.title('带多边形掩码的向量场（向量指向远离最近顶点的方向）')
    plt.show()

if __name__ == '__main__':
    CreateMaskedVectorField()

关键说明：

• KD树查询：kdtree.query()一次性返回所有网格点的最近顶点距离和索引，时间复杂度远低于循环遍历
• 掩码实现：用matplotlib.path.Path.contains_points()判断点是否在多边形内部，实现掩码效果
• 自定义向量规则：当前向量方向为从顶点指向当前点、幅值等于距离，你可以根据需求修改为自定义函数（比如距离的倒数、指数函数等）

内容的提问来源于stack exchange，提问作者offbyone