Panda: crea due nuove colonne in un dataframe con valori calcolati da una colonna preesistente

Sto lavorando con la libreria panda e voglio aggiungere due nuove colonne a un dataframe dfcon n colonne (n> 0).
Queste nuove colonne risultano dall'applicazione di una funzione a una delle colonne nel dataframe.

La funzione da applicare è come:

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

Un metodo per creare una nuova colonna per una funzione che restituisce solo un valore è:

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

Quindi, quello che voglio, e ho provato senza successo (*), è qualcosa del tipo:

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

Quale potrebbe essere il modo migliore per ottenere ciò? Ho scansionato la documentazione senza indizi.

** df['column_A'].map(calculate)restituisce una serie panda ogni elemento costituito da una tupla z, y. E provare ad assegnarlo a due colonne del dataframe produce un ValueError. *

Userei solo 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 risposta migliore è imperfetta secondo me. Si spera che nessuno stia importando in massa tutti i panda nel proprio spazio dei nomi con from pandas import *. Inoltre, il mapmetodo dovrebbe essere riservato per quelle volte in cui gli si passa un dizionario o una serie. Può assumere una funzione, ma questo è ciò per cui applyviene utilizzato.

Quindi, se devi usare l'approccio sopra, lo scriverei in questo modo

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

In realtà non c'è motivo di usare zip qui. Puoi semplicemente fare questo:

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

Questo secondo metodo è anche molto più veloce su DataFrame più grandi

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

DataFrame creato con 300.000 righe

%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 volte più veloce di zip

In generale, evitare di utilizzare applica

Applicare generalmente non è molto più veloce dell'iterazione su un elenco Python. Testiamo le prestazioni di un ciclo for per fare la stessa cosa di cui sopra

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

Quindi questo è due volte più lento, il che non è una terribile regressione delle prestazioni, ma se citiamo quanto sopra, otteniamo prestazioni molto migliori. Supponendo che tu stia 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)

Assegnazione diretta senza applicazione

È possibile ottenere miglioramenti della velocità ancora maggiori se si utilizzano le operazioni vettorializzate dirette.

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

Questo sfrutta le operazioni vettorializzate estremamente veloci di NumPy invece dei nostri cicli. Ora abbiamo una velocità 30 volte superiore rispetto all'originale.

Il test di velocità più semplice con apply

L'esempio sopra dovrebbe mostrare chiaramente quanto applypuò essere lento , ma solo così è più chiaro diamo un'occhiata all'esempio più semplice. Quadriamo una serie di 10 milioni di numeri con e senza applicazione

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)

Senza applicare è 50 volte più veloce

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