Search code examples
pythonpandascomplex-numbers

Splitting a dataframe with complex numbers into real and imaginary numbers with python


I would like to split the following complex dataframe into two columns,

df = pd.DataFrame({"AB": ['0.316227766016838-0.316227766016838i',
                    '0.316227766016838-0.316227766016838i',
                    '0.316227766016838-0.316227766016838i',
                    '0.316227766016838-0.316227766016838i',
                    '0.316227766016838+0.316227766016838i',
                    '0.3162277660168380+.316227766016838i']})

I tried in the following way but it works either for - or +

df1=df['AB'].str.split('-', n=1, expand=True)

How can I get two new columns for real and imaginary values only? Thanks!


Solution

  • Convert values to numpy array with repalce i to j with casting by np.complex128 and then get real and imaginary parts to new columns of DataFrame:

    df = pd.DataFrame({"AB": ['0.316227766016838-0.316227766016838i',
                        '0.316227766016838-0.316227766016838i',
                        '0.316227766016838-0.316227766016838i',
                        '0.316227766016838-0.316227766016838i',
                        '0.316227766016838+0.316227766016838i',
                        '0.316227766016838+0.316227766016838i']})
    
    arr = np.complex128(df['AB'].str.replace('i', 'j').to_numpy())
    print (arr)
    [0.31622777-0.31622777j 0.31622777-0.31622777j 0.31622777-0.31622777j
     0.31622777-0.31622777j 0.31622777+0.31622777j 0.31622777+0.31622777j]
    
    df = df.assign(real = arr.real, imag = arr.imag)
    print (df)
                                         AB      real      imag
    0  0.316227766016838-0.316227766016838i  0.316228 -0.316228
    1  0.316227766016838-0.316227766016838i  0.316228 -0.316228
    2  0.316227766016838-0.316227766016838i  0.316228 -0.316228
    3  0.316227766016838-0.316227766016838i  0.316228 -0.316228
    4  0.316227766016838+0.316227766016838i  0.316228  0.316228
    5  0.316227766016838+0.316227766016838i  0.316228  0.316228