在numpy中获取结果数组的数据类型

Question

我想为数组操作的输出预先分配内存，我需要知道要生成什么数据类型。下面我有一个函数，它可以做我想做的事情，但是非常难看。

import numpy as np

def array_operation(arr1, arr2):
    out_shape = arr1.shape
    # Get the dtype of the output, these lines are the ones I want to replace.
    index1 = ([0],) * arr1.ndim
    index2 = ([0],) * arr2.ndim
    tmp_arr = arr1[index1] * arr2[index2]
    out_dtype = tmp_arr.dtype
    # All so I can do the following.
    out_arr = np.empty(out_shape, out_dtype)

上面的很难看。numpy有这样的功能吗？

Answer 1

对于使用numpy版本<；1.6的用户，可以使用：

def result_type(arr1, arr2):
    x1 = arr1.flat[0]
    x2 = arr2.flat[0]
    return (x1 * x2).dtype

def array_operation(arr1, arr2):
    return np.empty(arr1.shape, result_type(arr1, arr2))

这与您发布的代码没有太大区别，尽管我认为arr1.flat[0]比index1 = ([0],) * arr1.ndim; arr1[index1]稍有改进。

对于numpy版本>；=1.6，使用Mike Graham的答案，np.result_type