Complex number operation on DataFrame

[davidtrem]

Hello,

Would it be possible to add complex function operation on DataFrame ?

Function like np.sin, np.abs works properly on Pandas DataFrame containing complex numbers and
returns a DataFrame (as expected).
However, np.real, np.imag, np.angle operation return an numpy.ndarray array instead of DataFrame...

import pandas as pd
import numpy as np
arr = pd.DataFrame(np.random.randn(2, 3) + 1j*np.random.randn(2, 3))
print(type(np.abs(arr)))  # <class 'pandas.core.frame.DataFrame'>
print(type(np.angle(arr)))  # <class 'numpy.ndarray'>

[davidtrem]

Same issue is inherited by xray. Where I first hit the issue.
pydata/xarray#553

[jreback]

real/image are defined on a Series
not sure these make sense for a DataFrame

[jreback]

In [40]: df  = pd.DataFrame(np.random.randn(2, 3) + 1j*np.random.randn(2, 3))

In [41]: df
Out[41]: 
                                   0                                 1                                 2
0  (-0.605037211028-0.777421542176j)  (0.456047518198-0.236148034977j)  (0.0203707536042-1.75589595197j)
1     (1.37752497265-2.01959737672j)   (0.504932940188+1.97880585951j)   (0.33416891278+0.847363479932j)

In [42]: df.dtypes
Out[42]: 
0    complex128
1    complex128
2    complex128
dtype: object

In [43]: df.apply(np.imag)
Out[43]: 
          0         1         2
0 -0.777422 -0.236148 -1.755896
1 -2.019597  1.978806  0.847363

In [44]: df.apply(np.angle)
Out[44]: 
          0         1         2
0 -2.232141 -0.477797 -1.559196
1 -0.972199  1.320957  1.195159

[jreback]

you can also use .applymap. or simply

DataFrame(np.angle(df),index=df.index,columns=df.columns)

.imag,.real,.angle don't call __array_wrap__ after their execution (I suppose because its a bug and they should, but you'd have to ask on the numpy list), so this is not really possible to do with special casing these functions (and I am not sure we can even intercept them w/o a lot of work).

Further you have to make sure that you have the appropriate dtypes (e.g. complex) or these don't make sense, so generally applying to an entire frame is not what you want, and you should use .select_dtypes first.

[davidtrem]

@jreback

real/image are defined on a Series
not sure these make sense for a DataFrame

Good point.
But for consistency abs, real, imag, angle should work the same way when applied to a DataFrame which is not the case right now.

[davidtrem]

Thanks for pointing the .apply and .applymap workaround by the way. 😄

[davidtrem]

@jreback I see your point now about why applying real/imag does not make sense for DataFrame.
Indeed, it does not really make sense to apply this on possibly non homogeneous data !
Actually xray.DataArray corresponds to pandas.Series. But as xray.DataArray can
be multidimensional (but homogeneous) I was incorrectly making a parallel with pandas.DataFrame.
Sorry for this misunderstanding and thanks again.