Получение dtype результирующего массива в numpy

Question

Я хочу предварительно выделить память для вывода операции массива, и мне нужно знать, какой тип dtype сделать.Ниже у меня есть функция, которая делает то, что я хочу, но ужасно уродлива.

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.result_type.

(Кроме того, понимаете ли вы, что вы можете получить доступ ко всем многомерным массивам как к 1d-массивам? Вам не нужен доступ к x[0, 0, 0, 0, 0] - вы можете получить доступ к x.flat[0].)

Answer 2

Для тех, кто использует numpy версию <1.6, вы можете использовать: </p>

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].

Для простой версии> = 1.6 используйте ответ Майка Грэма, np.result_type