Pandas: créez deux nouvelles colonnes dans un dataframe avec des valeurs calculées à partir d'une colonne préexistante

Je travaille avec la bibliothèque pandas et je souhaite ajouter deux nouvelles colonnes à un dataframedf avec n colonnes (n> 0).
Ces nouvelles colonnes résultent de l'application d'une fonction à l'une des colonnes du dataframe.

La fonction à appliquer est comme:

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

Une méthode pour créer une nouvelle colonne pour une fonction ne renvoyant qu'une valeur est:

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

Donc, ce que je veux, et j'ai essayé sans succès (*), c'est quelque chose comme:

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

Quelle pourrait être la meilleure façon d'y parvenir? J'ai scanné le documentation sans aucune idée.

** df['column_A'].map(calculate)renvoie une série pandas chaque élément consistant en un tuple z, y. Et essayer d'attribuer ceci à deux colonnes de dataframe produit une ValueError. *

J'utiliserais juste 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 meilleure réponse est imparfaite à mon avis. Espérons que personne n'importe en masse tous les pandas dans leur espace de noms avec from pandas import *. En outre, la mapméthode doit être réservée à ces moments lorsque vous lui passez un dictionnaire ou une série. Cela peut prendre une fonction mais c'est ce queapply est utilisé.

Donc, si vous devez utiliser l'approche ci-dessus, je l'écrirais comme ceci

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

Il n'y a en fait aucune raison d'utiliser zip ici. Vous pouvez simplement faire ceci:

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

Cette deuxième méthode est également beaucoup plus rapide sur des DataFrames plus volumineux

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

DataFrame créé avec 300 000 lignes

%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)

60x plus rapide que zip

En général, évitez d'utiliser apply

Apply n'est généralement pas beaucoup plus rapide que d'itérer sur une liste Python. Testons les performances d'une boucle for pour faire la même chose que ci-dessus

%%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)

C'est donc deux fois plus lent, ce qui n'est pas une terrible régression des performances, mais si nous cythonisons ce qui précède, nous obtenons de bien meilleures performances. En supposant que vous utilisez 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)

Attribuer directement sans appliquer

Vous pouvez obtenir des améliorations de vitesse encore plus importantes si vous utilisez les opérations vectorisées directes.

%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)

Cela tire parti des opérations vectorisées extrêmement rapides de NumPy au lieu de nos boucles. Nous avons maintenant une accélération 30x par rapport à l'original.

Le test de vitesse le plus simple avec apply

L'exemple ci-dessus devrait clairement montrer à quel point il applypeut être lent , mais pour qu'il soit très clair, examinons l'exemple le plus basique. Mettons au carré une série de 10 millions de nombres avec et sans application

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)

Sans appliquer est 50 fois plus rapide

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