python - 通过 NumPy 归一化矢量场

Question

我有形状为 [height width 2] 的数组 field我想让每个 [i, j] 的长度小于或等于 1我使用以下程序来做

def projectionOntoDisc(var):
    res = np.zeros_like(var)
    for i in xrange(var.shape[0]):
        for j in xrange(var.shape[1]):
            norm = max(1.0, np.linalg.norm(var[i, j]))
            res[i, j] = var[i, j] / norm
    return res

有什么方法可以更快吗？

PS 抱歉我的英语不好

Answer 1

您可以向量化此操作，这应该会提高执行速度几个数量级:

norm = numpy.fmax(1.0, numpy.linalg.norm(var, axis=2))
res = var / norm[:, :, numpy.newaxis]