如何让np.vectorize直接处理Ax2数组，避免转置操作？

问题描述 #

现有形状为Ax2的直角坐标数组（示例如下），需将每一行的(x,y)转换为极坐标(r, θ)，最终输出仍为Ax2的数组。当前实现需分割数组为x、y列，通过np.vectorize处理后转置，希望跳过分割和转置步骤，提升转换效率。

示例输入：

>>> np.arange(20).reshape(10, 2)
array([[ 0,  1],
       [ 2,  3],
       [ 4,  5],
       [ 6,  7],
       [ 8,  9],
       [10, 11],
       [12, 13],
       [14, 15],
       [16, 17],
       [18, 19]])

期望输出：

>>> toPolar(a)
[[ 1.          1.57079633]
 [ 3.60555128  0.98279372]
 [ 6.40312424  0.89605538]
 [ 9.21954446  0.86217005]
 [12.04159458  0.84415399]
 [14.86606875  0.83298127]
 [17.69180601  0.82537685]
 [20.51828453  0.81986726]
 [23.34523506  0.81569192]
 [26.17250466  0.81241861]]

当前实现代码：

def toPolar(arr):
    def rect2polar(x, y):
        r = math.sqrt(x**2 + y**2)
        ang = math.atan2(y, x)
        return (r, ang)

    vect = np.vectorize(rect2polar)
    p = vect(arr[:, 0], arr[:, 1])  # 返回2x10数组
    p = np.transpose(p)  # 转换为10x2形状
    return p

优化方案 #

方案1：使用numpy原生向量化函数（最优解） #

注意：np.vectorize并非真正的向量化实现，它只是对Python函数做了循环包装，运行效率远低于numpy原生的向量化操作

直接利用numpy内置的向量化计算函数，无需分割数组和转置，效率最高：

import numpy as np

def toPolar(arr):
    # 计算极径：每行的L2范数
    r = np.linalg.norm(arr, axis=1)
    # 计算极角：arctan2支持向量化输入
    theta = np.arctan2(arr[:, 1], arr[:, 0])
    # 按列拼接成Ax2形状的数组
    return np.column_stack((r, theta))

也可以显式计算极径，逻辑更直观：

import numpy as np

def toPolar(arr):
    x, y = arr[:, 0], arr[:, 1]
    r = np.sqrt(x**2 + y**2)
    theta = np.arctan2(y, x)
    return np.column_stack((r, theta))

方案2：调整np.vectorize的输出格式（不推荐，效率较低） #

如果一定要保留np.vectorize的写法，可以通过指定otypes参数让函数直接返回符合要求的形状，避免转置：

import numpy as np
import math

def toPolar(arr):
    def rect2polar(xy):
        x, y = xy
        r = math.sqrt(x**2 + y**2)
        ang = math.atan2(y, x)
        return r, ang

    # 指定输出为两个float类型，让vectorize按行处理输入
    vect = np.vectorize(rect2polar, otypes=[float, float])
    r, theta = vect(arr)
    return np.column_stack((r, theta))

总结 #

优先选择方案1，利用numpy原生向量化函数实现，既无需分割转置步骤，又能保证最高的运行效率。

内容的提问来源于stack exchange，提问作者Michael M.