Search code examples
pythonnumpycomplex-numbersnumpy-ndarray

Unpacking a numpy ndarray of tuples and add it in a new dimension


I have a numpy ndarray x of of shape (M,N) where each element in N is a 2 element tuple containing real and imaginary parts of a complex number.

I would like to unpack the tuple into another dimension such that my x would now have the shape (M,N,2)

My first idea is to seperate the real and imaginary parts into two numpy arrays of shape (M,N) and concatenate them over the second axis so that I can get the desired output. To do so, I tried to access only the real and imaginary parts as follows

x[:,:][0]
x[:,:][1]

However, it returns me the whole array instead of only the real/imaginary parts. Any pointers on how I can proceed?

An example array is as follows

array([[( 4.82641562, -1.84333128), (-1.29724314,  2.06545141),
        ( 0.02074953, -1.32518303)],
       [(-0.81477981, -3.34148138), ( 0.231383  ,  1.35495823),
        ( 1.66341462, -1.33603417)],
       [( 2.2809137 , -1.17720382), (-1.05115026, -2.09556016),
        ( 1.57932884, -0.7438237 )],
       [(-2.2170903 ,  0.81147059), ( 0.17489367,  0.15123805),
        (-1.95151614, -0.5669619 )],
       [(-2.16413158,  0.44555251), ( 1.2545366 , -0.06933326),
        ( 1.08469658,  1.71108236)],
       [(-0.99776606, -0.0607623 ), (-0.11785624, -1.90687216),
        ( 1.28964911, -1.03083365)],
       [( 1.65550362, -1.26383881), ( 0.13404437,  1.4402823 ),
        ( 0.04223597, -1.00992713)],
       [(-0.4696885 ,  2.66014003), (-1.12157114, -1.03019938),
        ( 2.6312301 ,  0.22770433)],
       [( 1.74457246, -0.91478885), (-0.18861201,  0.75360561),
        ( 0.27141119,  0.71356903)],
       [(-0.15940566,  1.26613812), ( 1.28238354,  1.31858431),
        ( 0.31217745,  0.13244299)]],
      dtype=[('real', '<f8'), ('imag', '<f8')])

Solution

  • You conveniently have your array setup as a structured array, which means you can just access the real and imaginary parts using each field's string name:

    arr['real']
    arr['imag']
    

    For example:

    >>> arr['real']
    array([[ 4.82641562, -1.29724314,  0.02074953],
           [-0.81477981,  0.231383  ,  1.66341462],
           [ 2.2809137 , -1.05115026,  1.57932884],
           [-2.2170903 ,  0.17489367, -1.95151614],
           [-2.16413158,  1.2545366 ,  1.08469658],
           [-0.99776606, -0.11785624,  1.28964911],
           [ 1.65550362,  0.13404437,  0.04223597],
           [-0.4696885 , -1.12157114,  2.6312301 ],
           [ 1.74457246, -0.18861201,  0.27141119],
           [-0.15940566,  1.28238354,  0.31217745]])
    
    >>> arr['imag']
    array([[-1.84333128,  2.06545141, -1.32518303],
           [-3.34148138,  1.35495823, -1.33603417],
           [-1.17720382, -2.09556016, -0.7438237 ],
           [ 0.81147059,  0.15123805, -0.5669619 ],
           [ 0.44555251, -0.06933326,  1.71108236],
           [-0.0607623 , -1.90687216, -1.03083365],
           [-1.26383881,  1.4402823 , -1.00992713],
           [ 2.66014003, -1.03019938,  0.22770433],
           [-0.91478885,  0.75360561,  0.71356903],
           [ 1.26613812,  1.31858431,  0.13244299]])