python - 向量化 numpy 的函数，将一个向量的每个元素与另一个向量的每个元素相关联

Question

我有一个函数，它接受两个向量的两个元素并计算标量值。如何使用 numpy 工具向量化这个函数，以便我可以编写

A = my_func(vec_a, vec_b)

其中 A 是维度 len(vec_a) x len(vec_b) 的矩阵。我怎样才能做到这一点？或者我必须明确地迭代 my_func 吗？ 额外的好处是:矩阵将非常稀疏，即 my_func 中的许多计算值都为零。是否可以在实现中包含这种稀疏性？

<小时/>

根据要求，举个例子:

假设我有两个向量a和b:

a = numpy.array([...]) # length n
b = numpy.array([...]) # length m

现在，我想调用 my_func(a,b) 并让它返回一个稀疏矩阵，其密集表示为

A = [
        [my_func(a[0], b[0]), my_func(a[0], b[1]), ..., my_func(a[0], b[n])],
        [my_func(a[1], b[0]), my_func(a[1], b[1]), ..., my_func(a[1], b[n])],
        ...
        [my_func(a[m], b[0]), my_func(a[m], b[1]), ..., my_func(a[m], b[n])]
]

当然，很多条目都是零。

<小时/>

根据要求，my_func 函数。

# note, that each element of the above vectors is a 
# list itself, with 4 elements. 
def my_func(a, b):
    distance = sp.sqrt(sp.sum((a[1:] - b[1:])**2))
    rate = sp.exp(-2*distance/loclength)

    if a[0] < b[0]:
        rate *= sp.exp((a[0] - b[0])/kT)

    return rate if rate > cutoff else 0

Answer 1

您可以适本地使用广播来做到这一点:

def my_func_vec(a, b):
    a = np.array(a, copy=False, ndmin=2)
    b = np.array(b, copy=False, ndmin=2)
    a = a[..., np.newaxis, :]
    b = b[..., np.newaxis, :, :]
    distance = np.sqrt(np.sum((a[..., 1:] - b[..., 1:])**2, axis=-1))
    rate = np.exp(-2*distance / loclength)
    mask = a[..., 0] < b[..., 0]
    rate[mask] *= np.exp((a[..., 0] - b[..., 0])[mask] / kT)
    mask = rate <= cutoff
    rate[mask] = 0
    return rate

要测试它，请设置一些虚拟值:

loclength = 1
kT = 1
cutoff = 0.25
a = np.random.rand(3, 5)
b = np.random.rand(4, 5)

现在:

>>> my_func_vec(a, b)
array([[ 0.34220076,  0.        ,  0.25392478,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.25953994,  0.        ,  0.        ]])

而不是:

>>> out = np.empty((3, 4))
>>> for r, j in enumerate(a):
...     for c, k in enumerate(b):
...         out[r, c] = my_func(j, k)
... 
>>> out
array([[ 0.34220076,  0.        ,  0.25392478,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.25953994,  0.        ,  0.        ]])