Cómo acceder a submarcos de datos de pandas groupby por clave
¿Cómo accedo al marco de datos groupby correspondiente en un objeto groupby mediante la clave?
Con el siguiente grupo:
rand = np.random.RandomState(1)
df = pd.DataFrame({'A': ['foo', 'bar'] * 3,
'B': rand.randn(6),
'C': rand.randint(0, 20, 6)})
gb = df.groupby(['A'])
Puedo iterarlo para obtener las claves y los grupos:
In [11]: for k, gp in gb:
print 'key=' + str(k)
print gp
key=bar
A B C
1 bar -0.611756 18
3 bar -1.072969 10
5 bar -2.301539 18
key=foo
A B C
0 foo 1.624345 5
2 foo -0.528172 11
4 foo 0.865408 14
Me gustaría poder acceder a un grupo por su clave:
In [12]: gb['foo']
Out[12]:
A B C
0 foo 1.624345 5
2 foo -0.528172 11
4 foo 0.865408 14
Pero cuando intento hacer eso gb[('foo',)]me sale esto raro.pandas.core.groupby.DataFrameGroupBy objeto extraño que no parece tener ningún método que corresponda al DataFrame que quiero.
Lo mejor que se me ocurre es:
In [13]: def gb_df_key(gb, key, orig_df):
ix = gb.indices[key]
return orig_df.ix[ix]
gb_df_key(gb, 'foo', df)
Out[13]:
A B C
0 foo 1.624345 5
2 foo -0.528172 11
4 foo 0.865408 14
pero esto es un poco desagradable, considerando lo buenos que suelen ser los pandas en estas cosas.
¿Cuál es la forma integrada de hacer esto?
Puedes utilizar el get_groupmétodo:
In [21]: gb.get_group('foo')
Out[21]:
A B C
0 foo 1.624345 5
2 foo -0.528172 11
4 foo 0.865408 14
Nota: Esto no requiere la creación de un diccionario intermedio/copia de cada submarco de datos para cada grupo, por lo que consumirá mucha más memoria que crear el diccionario ingenuo con dict(iter(gb)). Esto se debe a que utiliza estructuras de datos que ya están disponibles en el objeto groupby.
Puede seleccionar diferentes columnas usando el grupo mediante división:
In [22]: gb[["A", "B"]].get_group("foo")
Out[22]:
A B
0 foo 1.624345
2 foo -0.528172
4 foo 0.865408
In [23]: gb["C"].get_group("foo")
Out[23]:
0 5
2 11
4 14
Name: C, dtype: int64
Wes McKinney (autor de pandas) en Python para análisis de datos proporciona la siguiente receta:
groups = dict(list(gb))
que devuelve un diccionario cuyas claves son las etiquetas de su grupo y cuyos valores son DataFrames, es decir
groups['foo']
te dará lo que estás buscando:
A B C
0 foo 1.624345 5
2 foo -0.528172 11
4 foo 0.865408 14
En vez de
gb.get_group('foo')
prefiero usargb.groups
df.loc[gb.groups['foo']]
Porque de esta manera también puedes elegir varias columnas. Por ejemplo:
df.loc[gb.groups['foo'],('A','B')]
