Pandas: cree dos nuevas columnas en un marco de datos con valores calculados a partir de una columna preexistente

Estoy trabajando con la biblioteca de pandas y quiero agregar dos columnas nuevas a un marco de datos dfcon n columnas (n> 0).
Estas nuevas columnas son el resultado de la aplicación de una función a una de las columnas del marco de datos.

La función a aplicar es como:

def calculate(x):
    ...operate...
    return z, y

Un método para crear una nueva columna para una función que devuelve solo un valor es:

df['new_col']) = df['column_A'].map(a_function)

Entonces, lo que quiero, y lo intenté sin éxito (*), es algo como:

(df['new_col_zetas'], df['new_col_ys']) = df['column_A'].map(calculate)

¿Cuál podría ser la mejor manera de lograr esto? Escaneé la documentación sin ninguna pista.

** df['column_A'].map(calculate)devuelve una serie de pandas, cada elemento que consta de una tupla z, y. Y tratar de asignar esto a dos columnas de marco de datos produce un ValueError. *

Solo usaría zip:

In [1]: from pandas import *

In [2]: def calculate(x):
   ...:     return x*2, x*3
   ...: 

In [3]: df = DataFrame({'a': [1,2,3], 'b': [2,3,4]})

In [4]: df
Out[4]: 
   a  b
0  1  2
1  2  3
2  3  4

In [5]: df["A1"], df["A2"] = zip(*df["a"].map(calculate))

In [6]: df
Out[6]: 
   a  b  A1  A2
0  1  2   2   3
1  2  3   4   6
2  3  4   6   9

La respuesta principal es defectuosa en mi opinión. Con suerte, nadie está importando en masa todos los pandas en su espacio de nombres con from pandas import *. Además, el mapmétodo debe reservarse para esos momentos en los que se le pasa un diccionario o una Serie. Puede tomar una función, pero para eso applyse usa.

Entonces, si debe usar el enfoque anterior, lo escribiría así

df["A1"], df["A2"] = zip(*df["a"].apply(calculate))

En realidad, no hay razón para usar zip aquí. Simplemente puede hacer esto:

df["A1"], df["A2"] = calculate(df['a'])

Este segundo método también es mucho más rápido en DataFrames más grandes

df = pd.DataFrame({'a': [1,2,3] * 100000, 'b': [2,3,4] * 100000})

DataFrame creado con 300.000 filas

%timeit df["A1"], df["A2"] = calculate(df['a'])
2.65 ms ± 92.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit df["A1"], df["A2"] = zip(*df["a"].apply(calculate))
159 ms ± 5.24 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

60 veces más rápido que zip

En general, evite usar aplicar

Aplicar generalmente no es mucho más rápido que iterar sobre una lista de Python. Probemos el rendimiento de un bucle for para hacer lo mismo que el anterior

%%timeit
A1, A2 = [], []
for val in df['a']:
    A1.append(val**2)
    A2.append(val**3)

df['A1'] = A1
df['A2'] = A2

298 ms ± 7.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Entonces, esto es dos veces más lento, lo que no es una regresión de rendimiento terrible, pero si citonizamos lo anterior, obtenemos un rendimiento mucho mejor. Suponiendo que está usando ipython:

%load_ext cython

%%cython
cpdef power(vals):
    A1, A2 = [], []
    cdef double val
    for val in vals:
        A1.append(val**2)
        A2.append(val**3)

    return A1, A2

%timeit df['A1'], df['A2'] = power(df['a'])
72.7 ms ± 2.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Asignación directa sin aplicar

Puede obtener mejoras de velocidad aún mayores si utiliza las operaciones vectorizadas directas.

%timeit df['A1'], df['A2'] = df['a'] ** 2, df['a'] ** 3
5.13 ms ± 320 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Esto aprovecha las operaciones vectorizadas extremadamente rápidas de NumPy en lugar de nuestros bucles. Ahora tenemos una aceleración 30 veces superior al original.

La prueba de velocidad más sencilla con apply

El ejemplo anterior debería mostrar claramente lo lento que applypuede ser, pero para que sea más claro, veamos el ejemplo más básico. Cuadremos una serie de 10 millones de números con y sin aplicar

s = pd.Series(np.random.rand(10000000))

%timeit s.apply(calc)
3.3 s ± 57.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Sin aplicar es 50 veces más rápido

%timeit s ** 2
66 ms ± 2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)