我正在尝试使用 Python/Numpy 矢量化函数来减少 for 循环。
我的函数调用看起来像这样
out_vectors = v_calculation(
in_vectors,
p
)
矢量化函数定义为
v_calculation = np.vectorize(
my_calculation,
signature='(j,i),(i)->()'
)
in_vectors 是一个形状数组
(3,6,200,3)
但是前两个维度 (3,6) 可以是任何东西。这些是循环尺寸。
p 是一个形状数组
(3,)
我的计算是这样的
def my_calculation(in_vector, p):
"""
total magnetic field from Biot-Savart's law
"""
out_vector = np.zeros((3,))
l_vector = in_vector[1:, :] - in_vector[:-1, :]
r_vector = (in_vector[:-1, :] + l_vector / 2) - p
out_vector = np.sum(np.cross(l_vector, r_vector) / \
np.linalg.norm(r_vector) ** 3,
axis=0
)
return out_vector
在这个函数中,in_vector 是一个形状为 (200, 3) 的数组,p 是相同的形状 (3,)。out_vector 形状为 (3,)。这是对的。
out_vectors,向量化函数的结果应该是(6,3)。对于 input_vectors 的每个第二维(在本例中为 6),这应该是对 input_vectors 的第一维(在本例中为 3)求和的 my_calculation 的结果。结果的第二维是 3(向量的 x、y、z 分量),与 p 的维数和 input_vectors 的第四维相同。我希望这一切都清楚了。
我的代码在矢量化函数调用中失败
堆栈跟踪
~/path/to/my/code.py in calculate_vectors(mgr)
588 out_vectors = v_calculation(
589 in_vectors,
--> 590 p
591 )
~/miniconda/lib/python3.7/site-packages/numpy/lib/function_base.py in __call__(self, *args, **kwargs)
1970 vargs.extend([kwargs[_n] for _n in names])
1971
-> 1972 return self._vectorize_call(func=func, args=vargs)
1973
1974 def _get_ufunc_and_otypes(self, func, args):
~/miniconda/lib/python3.7/site-packages/numpy/lib/function_base.py in _vectorize_call(self, func, args)
2036 """Vectorized call to `func` over positional `args`."""
2037 if self.signature is not None:
-> 2038 res = self._vectorize_call_with_signature(func, args)
2039 elif not args:
2040 res = func()
~/miniconda/lib/python3.7/site-packages/numpy/lib/function_base.py in _vectorize_call_with_signature(self, func, args)
2100
2101 for output, result in zip(outputs, results):
-> 2102 output[index] = result
2103
2104 if outputs is None:
ValueError: setting an array element with a sequence.