pythonnumpynumpy-ndarraynumpy-dtype

numpy: determine the corresponding float datatype from a complex datatype


I'm using two numpy arrays of floating point numbers

Rs = np.linspace(*realBounds, realResolution, dtype=fdtype)
Is = np.linspace(*imagBounds, imagResolution, dtype=fdtype)

to make a grid Zs of complex numbers with shape (realResolution, imagResolution). I'd like to only specify the datatype of Zs and use this to determine the float datatype fdtype.

As I understand it from this page, complex datatypes are always represented as two floats with the same datatype, hence why Rs and Is both have dtype=fdtype. If the datatype of Zs is specified as np.csingle then I'd like fdtype to be np.single, if it's specified as np.dtype("complex128") then I'd like fdtype to be np.dtype("float64"), and so on. Is there a nice way to do this?

Edit: I'm equally interested in the inverse function which sends, for example, np.single to np.csingle.


Solution

  • I don't think there is such a function in numpy but you can easily roll your own using num, which uniquely identifies each of the built-in types:

    def get_corresonding_dtype(dt):
        return {11: np.csingle,
                12: np.cdouble,
                13: np.clongdouble,
                14: np.single,
                15: np.double,
                16: np.longdouble}[np.dtype(dt).num]
    

    The dict was obtained from {np.dtype(dt).num: dt for dt in (np.single, np.double, np.longdouble, np.csingle, np.cdouble, np.clongdouble)} which returns

    {11: numpy.float32,
     12: numpy.float64,
     13: numpy.longdouble,
     14: numpy.complex64,
     15: numpy.complex128,
     16: numpy.clongdouble}
    

    on win-amd64 (sysconfig.get_platform()). The names may vary, e.g. on linux-x86_64 it shows numpy.complex256 instead of numpy.clongdouble, but thanks to num you don't need to pay attention to it.

    So for instance for all of np.double, np.float_, np.float64, float, 'float64' etc., get_corresonding_dtype will return numpy.complex128.

    If you like you can also add np.half as 23: np.csingle to the dict, as there's no complex half/float16 type in numpy.