对于使用标量或向量作为输入的函数,Numpy 最佳实践是什么?

What is the Numpy best practice for a function that works with scalars or vectors as inputs?

我经常使用 numpy 编写方程式,将标量作为输入,returns 另一个标量或向量。后来我发现我想做同样的事情,但使用一个或多个向量作为输入。我正在尝试找出一种方法,使一个函数在两种情况下都能正常工作,而不会在各种 if 测试中调用 np.isscalar()np.atleast1d()(除非那是唯一的方法)。

是否可以仅使用一个函数来处理标量和向量输入,还是我会被多个实现所困?

例如,这里有一些函数可以将 x、y 角转换为北、东、下单位向量(假设角度已经以弧度为单位)。我包括了一个标量版本、矢量化版本和一个使用 np.meshgrid() 调用标量版本的版本。我试图避免对 np.vectorize().

进行显式循环或调用

标量示例

import numpy as np

def xy_to_nez_scalar(x, y):
    
    n = np.sin(y)
    e = np.sin(x) * np.cos(y)
    z = np.cos(x) * np.cos(y)
    
    return np.array([n, e, z])

x = 1
y = 1

nez1 = xy_to_nez_scalar(x, y)
print(f'{nez1=}')
print(f'{nez1.shape=}\n')

生产

nez1=array([0.84147098, 0.45464871, 0.29192658])
nez1.shape=(3,)

矢量化示例

# X and Y are arrays, so we'll follow the convention that X
# is a column of N rows and Y is a row of M columns. Then we would
# return an array of shape (N, M, 3).

def xy_to_nez_vector(x, y):
    
    n = np.sin(y)
    e = np.sin(x[:, np.newaxis]) * np.cos(y)
    z = np.cos(x[:, np.newaxis]) * np.cos(y)
    
    # make sure n has the same shape as e and z
    nn = np.broadcast_to(n, e.shape)
                 
    nez = np.stack([nn, e, z], axis=-1)
    return nez

x_array = np.arange(4)
y_array = np.arange(2)

nez2 = xy_to_nez_vector(x_array, y_array)
print(f'{nez2=}')
print(f'{nez2.shape=}\n')

生产

nez2=array([[[ 0.        ,  0.        ,  1.        ],
        [ 0.84147098,  0.        ,  0.54030231]],

       [[ 0.        ,  0.84147098,  0.54030231],
        [ 0.84147098,  0.45464871,  0.29192658]],

       [[ 0.        ,  0.90929743, -0.41614684],
        [ 0.84147098,  0.4912955 , -0.2248451 ]],

       [[ 0.        ,  0.14112001, -0.9899925 ],
        [ 0.84147098,  0.07624747, -0.53489523]]])
nez2.shape=(4, 2, 3)

网格示例

# note this produces a (3, M, N) result.

xv, yv = np.meshgrid(x_array, y_array)
nez3 = xy_to_nez_scalar(xv, yv)
print(f'{nez3=}')
print(f'{nez3.shape=}\n')

# fix things by transposing.

nez4 = nez3.T
print(f'{nez4=}')
print(f'{nez4.shape=}\n')
print(np.allclose(nez2, nez4))

生产

nez3=array([[[ 0.        ,  0.        ,  0.        ,  0.        ],
        [ 0.84147098,  0.84147098,  0.84147098,  0.84147098]],

       [[ 0.        ,  0.84147098,  0.90929743,  0.14112001],
        [ 0.        ,  0.45464871,  0.4912955 ,  0.07624747]],

       [[ 1.        ,  0.54030231, -0.41614684, -0.9899925 ],
        [ 0.54030231,  0.29192658, -0.2248451 , -0.53489523]]])
nez3.shape=(3, 2, 4)

nez4=array([[[ 0.        ,  0.        ,  1.        ],
        [ 0.84147098,  0.        ,  0.54030231]],

       [[ 0.        ,  0.84147098,  0.54030231],
        [ 0.84147098,  0.45464871,  0.29192658]],

       [[ 0.        ,  0.90929743, -0.41614684],
        [ 0.84147098,  0.4912955 , -0.2248451 ]],

       [[ 0.        ,  0.14112001, -0.9899925 ],
        [ 0.84147098,  0.07624747, -0.53489523]]])
nez4.shape=(4, 2, 3)

True

这可能与框架挑战接壤,但我建议稍微更改您的实现理念以符合大多数 numpy 函数已经执行的操作。这有两个优点:(1) 有经验的 numpy 用户会知道您的函数有什么期望,以及 (2) 标量-向量问题消失了。

通常,如果遇到像 xy_to_nez(x, y) 这样的函数,我希望它采用数组 xy,以及 return 具有 [=25] 的数组=] 两个的形状,第一个或最后一个维度为 3。将 3 放在最后一个维度的选择在这里完全没问题。然而,神奇地将阵列啮合在一起而不是广播它们是一件相当不常见的事情pythonic。

让用户明确告诉你你想要什么(python 的核心原则,import this). For example, given a broadcasting interface as suggested above, your scalar function is nearly complete. The only change you would need to make is to stack along the last axis instead of the first, as np.array 中的第 2 项:

def xy_to_nez(x, y):

    n = np.sin(y)
    e = np.sin(x) * np.cos(y)
    z = np.cos(x) * np.cos(y)

    return np.stack(np.broadcast_arrays(n, e, z), -1)

我希望以下三个示例可以作为 nez1nez2nez4:

>>> xy_to_nez(1, 1)  # shape: 3
array([0.84147098, 0.45464871, 0.29192658])
>>> xy_to_nez(np.arange(4)[:, None], np.arange(2))  # shape: 4, 2
array([[[ 0.        ,  0.        ,  1.        ],
        [ 0.84147098,  0.        ,  0.54030231]],
       [[ 0.        ,  0.84147098,  0.54030231],
        [ 0.84147098,  0.45464871,  0.29192658]],
       [[ 0.        ,  0.90929743, -0.41614684],
        [ 0.84147098,  0.4912955 , -0.2248451 ]],
       [[ 0.        ,  0.14112001, -0.9899925 ],
        [ 0.84147098,  0.07624747, -0.53489523]]])
>>> xy_to_nez(*np.meshgrid(np.arange(4), np.arange(2), indexing='ij'))
array([[[ 0.        ,  0.        ,  1.        ],
        [ 0.84147098,  0.        ,  0.54030231]],
       [[ 0.        ,  0.84147098,  0.54030231],
        [ 0.84147098,  0.45464871,  0.29192658]],
       [[ 0.        ,  0.90929743, -0.41614684],
        [ 0.84147098,  0.4912955 , -0.2248451 ]],
       [[ 0.        ,  0.14112001, -0.9899925 ],
        [ 0.84147098,  0.07624747, -0.53489523]]])

在第三个示例中,我通过将 indexing='ij' 传递给 np.meshgrid 来更正 nez3 的问题。问题不在于矢量化,而在于您传入的网格的形状。

我根本不希望 xy_to_nez(np.arange(4), np.arange(2)) 工作:数组不会一起广播,我们不应该试图弄清楚如何组合它们。要了解原因,请假装它们各有三个随机维度。你交织吗?你是不是一套接一套地放?如果某些维度广播而其他维度不广播怎么办?将这些问题留给用户考虑。

同时:

>>> xy_to_nez(np.arange(4), np.arange(4))  # Shape: 4, 3
array([[ 0.        ,  0.        ,  1.        ],
       [ 0.84147098,  0.45464871,  0.29192658],
       [ 0.90929743, -0.37840125,  0.17317819],
       [ 0.14112001, -0.13970775,  0.98008514]])