Procesando datos con pandas

La motivación de este post es ayudarme en el futuro a mi mismo a recordar estos trucos, porque la notación no es muy amigable y siempre lo olvido.

Ejemplificaré con el siguiente dataset:

import pandas as pd
df = pd.DataFrame(
    data = {
            "color": ["negro", "negro", "negro", "azul","azul", "azul", "rojo", "rojo"],
            "x":[1, 2, 3, -2, -5, -1, -2, -5],
            "y":[1, 2, 3, -4, -5, -4, 3, 4],
           })
print(df)

   color  x  y
0  negro  1  1
1  negro  2  2
2  negro  3  3
3   azul -2 -4
4   azul -5 -5
5   azul -1 -4
6   rojo -2  3
7   rojo -5  4

Gráficamente, el set anterior es el siguiente:

from matplotlib import pyplot as plt
color_dict = {"negro":"k", "azul":"b", "rojo":"r"}
x = df.x.values
y = df.y.values
L = max(abs(x.max()), abs(x.min()), abs(y.max()), abs(y.min()), ) + 1
c = [color_dict[_] for _ in df.color.values]
plt.scatter(x, y, c=c)
plt.xlim([-L, L])
plt.ylim([-L, L])

(-6.0, 6.0)

Métodos tradicionales

Una opción es aplicar directamente los métodos:

df_g = df.groupby("color").mean().reset_index()
print(df_g)

   color         x         y
0   azul -2.666667 -4.333333
1  negro  2.000000  2.000000
2   rojo -3.500000  3.500000

o aplicandolo a una columna en específico:

df_g = df.groupby("color").y.mean().reset_index()
print(df_g)

   color         y
0   azul -4.333333
1  negro  2.000000
2   rojo  3.500000

Es desaconsejable porque: * Tendrás que renombrar manualmente los nombres de las columnas, que mantienen el nombre de la columna original y no de la función aplicada. * No puedes aplicar el método a más de una función y columna de manera simultánea * No puedes usar cualquier función, solo algunas conocidas.

Otro método que es posible usar, pero no conviene tanto es usar agg:

df_g = df.groupby("color").agg(
    {
        "x": [lambda _: _.mean()], # aplicado a y
        "y": [lambda _: _.nunique(), lambda _: _.mean()], # aplicado a z
    }
)
print(df_g)

              x          y           
       <lambda> <lambda_0> <lambda_1>
color                                
azul  -2.666667          2  -4.333333
negro  2.000000          3   2.000000
rojo  -3.500000          2   3.500000

Es desaconsejable porque: * Genera 2 niveles de columnas: es dificil de “aplanar” * Hace incluso más difícil renombrar las columnas. Pero al menos es mejor que el método anterior porque puedes usar funciones genéricas.

Método sugerido

El método que aconsejo es con agg pero con definición explícita de los nombres de las columnas.

Después del groupby, aplicas el método agg pasándole inventando un nombre del parámetro deseado y dándole como valor una tupla (columna, función).

¡Las columnas generadas obtienen inmediatamente el nombre deseado!

import pandas as pd
df_g = df.groupby("color").agg(
        promedio_x = ("x", lambda _: _.mean()), 
        conteo_y = ("y", lambda _: _.nunique()), 
        promedio_y = ("y", lambda _: _.mean()), 
).reset_index()
print(df_g)

   color  promedio_x  conteo_y  promedio_y
0   azul   -2.666667         2   -4.333333
1  negro    2.000000         3    2.000000
2   rojo   -3.500000         2    3.500000

El groupby puede recibir una lista de columnas en caso de necesitar un nivel de agregación más detallado.

Notar que las columnas a generar se pasan como parámetros de la función (promedio_y) y no strings ("promedio_y")!!!

Multicolumna

Una mejora interesante es usar lo anterior para trabajar con múltiples columnas al mismo tiempo. Veamos un ejemplo concreto: calculando distintas métricas (distancias) respecto al origen (0,0).

Para poder aplicar el truco anterior, necesitamos “empaquetar las columnas” en una variable

df["points"] = list(zip(df.x, df.y))
print(df)

   color  x  y    points
0  negro  1  1    (1, 1)
1  negro  2  2    (2, 2)
2  negro  3  3    (3, 3)
3   azul -2 -4  (-2, -4)
4   azul -5 -5  (-5, -5)
5   azul -1 -4  (-1, -4)
6   rojo -2  3   (-2, 3)
7   rojo -5  4   (-5, 4)

Y ahora podemos aplicar distintas funciones, aunque en general será mejor definirlas como funciones (más que lambda-funciones) para expresar mejor las relaciones a usar entre los elementos:

def norm_1(points):
    d = points.apply(lambda point: abs(point[0]) + abs(point[1]))
    return d.mean()

def norm_2(points):
    d = points.apply(lambda point: (point[0]**2 + point[1]**2)**.5)
    return d.mean()

def norm_inf(points):
    d = points.apply(lambda point: max(abs(point[0]), abs(point[1])))
    return d.mean()

df_agg = df.groupby("color").agg(
        norm_1_avg = ("points", norm_1), 
        norm_2_avg = ("points", norm_2), 
        norm_inf_avg = ("points", norm_inf), 
).reset_index()
print(df_agg)

   color  norm_1_avg  norm_2_avg  norm_inf_avg
0   azul         7.0    5.222103      4.333333
1  negro         4.0    2.828427      2.000000
2   rojo         7.0    5.004338      4.000000