numpy vectorize一个函数以接受不同长度的向量并返回张量结果 [英] numpy vectorize a function to accepts vectors of different lengths and return the tensor result

问题描述

我想向量化函数f(a, b)，以便当我将a和b输入为两个向量时，返回组合的张量.这是一个说明性示例:

I want to vectorize a function f(a, b) so that, when I enter a and b as two vectors, the tensor of combinations is returned. Here is an illustrative example:

import numpy as np
def tester(a, b):
   mysumm = 0.
   for ii in range(a):
       for jj in range(b):
           mysumm += a * b
   return mysumm
tester = np.vectorize(tester)
x, y = [2, 4], [3, 5, 8] 
print(tester(x, 3)) # [ 36. 144.]
print(tester(x, 5)) # [100. 400.]
print(tester(x, 8)) # [ 256. 1024.]
print(tester(2, y)) # [ 36. 100. 256.]
print(tester(4, y)) # [ 144.  400. 1024.]
print(tester(x, y)) # ValueError: operands could not be broadcast together with shapes (2,) (3,) 

我希望tester(x, y)调用返回一个2x3矩阵，类似于[[ 36. 100. 256.], [ 144. 400. 1024.]]，我很惊讶这不是默认行为.

I expected the tester(x, y) call to return a 2x3 matrix, something like [[ 36. 100. 256.], [ 144. 400. 1024.]] and I was surprised that this isn't the default behavior.

如何使vecotirzed函数返回输入向量的可能组合的张量?

How can I make the vecotirzed function return the tensor of possible combinations of the input vectors?

推荐答案

您可以与np.ix_链接:

>>> import functools
>>> 
>>> def tensorize(f):
...     fv = np.vectorize(f)
...     @functools.wraps(f)
...     def ft(*args):
...         return fv(*np.ix_(*map(np.ravel, args)))
...     return ft
... 
>>> tester = tensorize(tester)
>>> tester(np.arange(3), np.arange(2))
array([[0., 0.],
       [0., 1.],
       [0., 4.]])

这篇关于numpy vectorize一个函数以接受不同长度的向量并返回张量结果的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！