python - 以 float 或 ndarray 作为参数的 numpy 函数

Question

考虑以下代码。

import numpy as np

def f(x, y):
    return (x[:, np.newaxis] - y[:]).sum(axis = 1)

x1 = np.linspace(1, 2, 5)
y1 = np.linspace(3, 4, 7)

print(f(x1, y1))

# x2 = 2. # (won't work)
# x2 = [2.] # (won't work either)
x2 = np.asarray([2.])
print(f(x2, y1))

# x3, _ = np.meshgrid(x1, x1) # won't work
# print(f(x3, y1))

第一次调用f(x1, y1)按预期工作。

第二次调用， float 在工作之前需要进行“转换”(我希望转换对用户透明，因此转换应该包含在函数的定义中)。

第三次调用根本不起作用，我不知道如何调整函数的定义，使其与第一次调用的功能相同，但对于网格上的某些点。

有什么想法吗？谢谢。

Answer 1

您可以尝试沿另一个轴进行减法，并强制 x 成为兼容维度的数组:

import numpy as np

def g(x,y):
    x = np.asarray(x)
    if len(x.shape)==2:
        return (x - y[:,np.newaxis, np.newaxis]).sum(axis=0)
    return (x - y[:,np.newaxis]).sum(axis=0)

x1 = np.linspace(1, 2, 5)
y1 = np.linspace(3, 4, 7)

print('x1 =', x1)
print('y1 =', y1)
print('g(x1,y1) =', g(x1, y1))

x2 = 2.
print('x2 =', x2)
print('y1 =', y1)
print('g(x2, y1) =', g(x2, y1))

x3, _ = np.meshgrid(x1, x1)

print('x3 =', x3)
print('g(x3,y1) =', g(x3, y1))

输出:

x1 = [ 1.    1.25  1.5   1.75  2.  ]
y1 = [ 3.          3.16666667  3.33333333  3.5         3.66666667  3.83333333
  4.        ]
g(x1,y1) = [-17.5  -15.75 -14.   -12.25 -10.5 ]
x2 = 2.0
y1 = [ 3.          3.16666667  3.33333333  3.5         3.66666667  3.83333333
  4.        ]
g(x2, y1) = [-10.5]
x3 = [[ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]]
g(x3,y1) = [[-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]]