Efficient calculation of frequencies of multiindexed category variables

Viewed 40

I'm considering a multiindexed (i,j,k) DataFrame with one column that contains three categorical variables A, B or C.

I want to compute the frequency of the variables for each i over all (j,k). I have a solution, but I think there exists a more pythonic and efficient way of doing it.

The code for a MWE reads (in reality len(I)*len(J)*len(K) is large, in the millions, say):

import pandas as pd
import numpy as np

I = range(10)
J = range(3)
K = range(2)

data = pd.DataFrame(
    np.random.randint(0, 3, size=len(I)*len(J)*len(K)),
    index=pd.MultiIndex.from_product([I, J, K]),
    columns=['cat']
)

data.index.names = ['i', 'j', 'k']

data[data['cat'] == 0] = 'A'
data[data['cat'] == 1] = 'B'
data[data['cat'] == 2] = 'C'


data = data.unstack(['j', 'k'])
result = data.apply(lambda x: x.value_counts(), axis=1).fillna(0) / (len(J)*len(K))
2 Answers

You could use groupby, and also normalize your value_counts:

data.groupby(level=0)['cat'].value_counts(normalize=True).unstack(level=1).fillna(0)

To compare, first let's make the dummy data big (60 million rows):

import pandas as pd
import numpy as np

I = range(100000)
J = range(30)
K = range(20)

data = pd.DataFrame(
    np.random.randint(0, 3, size=len(I)*len(J)*len(K)),
    index=pd.MultiIndex.from_product([I, J, K]),
    columns=['cat']
)

data.index.names = ['i', 'j', 'k']

data[data['cat'] == 0] = 'A'
data[data['cat'] == 1] = 'B'
data[data['cat'] == 2] = 'C'

Timing your original method:

data_interim = data.unstack(['j', 'k'])
data_interim.apply(lambda x: x.value_counts(), axis=1).fillna(0) / (len(J)*len(K))

gives (on my machine) 1min 24s ± 1.98 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

the alternative:

data.groupby(level=0)['cat'].value_counts(normalize=True).unstack(level=1).fillna(0)

gives (on my machine) 8.86 s ± 216 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Baseline for s_pike's method:

%%timeit 
(data.groupby(level=0)['cat']
     .value_counts(normalize=True)
     .unstack(level=1)
     .fillna(0))

6.41 s ± 243 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

If they're truly categorical, you can get a lot of benefit out of explicitly making the column categorical, and then using either of these methods.

They're both still about twice as fast without being categorical, but become about 3x as fast when made categorical.

data['cat'] = data['cat'].astype('category')
%%timeit
(data.groupby(level=0, as_index=False)['cat']
     .value_counts(normalize=True)
     .pivot(index='i', columns='cat', values='proportion'))

1.82 s ± 91 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%%timeit
(x := data.pivot_table(index='i', columns='cat', aggfunc='value_counts')).div(x.sum(1), 0)

1.8 s ± 107 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Outputs:

cat           A         B         C
i                                  
0      0.341667  0.318333  0.340000
1      0.311667  0.388333  0.300000
2      0.351667  0.350000  0.298333
3      0.363333  0.333333  0.303333
4      0.326667  0.350000  0.323333
...         ...       ...       ...
99995  0.315000  0.313333  0.371667
99996  0.323333  0.351667  0.325000
99997  0.305000  0.353333  0.341667
99998  0.318333  0.341667  0.340000
99999  0.331667  0.340000  0.328333

[100000 rows x 3 columns]
Related