Geek Answers Handbook

Getting dtype array of results in numpy

I want to preallocate memory for the output of an array operation, and I need to know which dtype to do this. Below I have a function that does what I want, but it is terribly ugly.

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)

The above is pretty ugly. Does numpy have a function that does this?

+5
source share
2 answers

You are looking for numpy.result_type.

(As an aside, do you understand that you can access all multidimensional arrays in the form of 1d arrays? You do not need to have access to x[0, 0, 0, 0, 0]- you can access x.flat[0].)

+7
source

For those using numpy version <1.6, you can use:

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

+1

More articles:


All Articles