I am trying to use Python/Numpy vectorized functions to reduce for loops.
My function call looks like this
out_vectors = v_calculation(
in_vectors,
p
)
The vectorized function definition is
v_calculation = np.vectorize(
my_calculation,
signature='(j,i),(i)->()'
)
in_vectors is an array of shape
(3,6,200,3)
But the first 2 dimensions (3,6) can be anything. These are the loop dimensions.
p is an array of shape
(3,)
My calculation is this
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 this function, in_vector is an array of shape (200, 3) and p is the same shape (3,). out_vector shape is (3,). This is correct.
out_vectors, the result of the vectorized function should be (6,3). This should be the results of my_calculation summed over the first dimension of the input_vectors (3 in this case) for each of the second dimension of the input_vectors (6 in this case). The second dimension of the result is 3 (x, y, z components for the vector), same as the dimension of p and the fourth dimension of input_vectors. I hope this is all clear.
My code is failing in the vectorized function call
Stacktrace
~/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.
This works for me. Note I changed the return signature, to match the shared final dimension of both inputs.
In [54]: A = np.arange(12).reshape(4,3); b = np.arange(3)
In [55]: my_calculation(A,b)
Out[55]: array([0., 0., 0.])
In [56]: f = np.vectorize(my_calculation, signature='(j,i),(i)->(i)')
In [57]: f(A,b)
Out[57]: array([0., 0., 0.])
In [58]: f([A,A,A],b)
Out[58]:
array([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]])