使用float或ndarray作为参数的numpy函数

时间:2015-12-01 08:45:37

标签: python numpy numpy-broadcasting

请考虑以下代码。

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需要在它工作之前被“转换”(我希望转换对用户是透明的,所以转换应该包含在函数的定义中)。

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

有什么想法吗?感谢。

2 个答案:

答案 0 :(得分:1)

要回答浮动的透明转换问题,至少可以在定义中使用np.atleast_1d()方法:

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

答案 1 :(得分: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 ]]