Assign length of / number of items in group to new column

Viewed 65

TLDR: I need to add the length or the number of items in each group to a new column. How can that be achieved?

I am dealing with experimental results containing the results of several repetitions (repeating the same experiment with same settings to improve data's statistical resilience). Each experiment has an identifier, whereas each single repetition/run has an index 'inside' each experiment (see code snippet for clarification).

Due to upcoming data processing and to show the run number in relation to the total number of runs per experiment (e.g. run 1/3, run 2/3, run 3/3) I need to add two columns containing

  • a 'human-readable' run number (basically a one-based run index) and
  • the total number of runs per experiment.

The first can be easily achieved by incrementing the run_id given for each run:

import io
import pandas as pd
import numpy as np
    
DATA_STRING = """
experiment_id   run_id    value     other_data        
9tfc6d          0         0.448     0.883 
9tfc6d          1         0.963     0.230
9tfc6d          2         0.711     0.724
q9tqjq          0         0.748     0.959
q9tqjq          1         0.662     0.772
q9tqjq          2         0.530     0.834
jsxp2m          0         0.087     0.346
jsxp2m          1         0.362     0.569
jsxp2m          2         0.124     0.206
"""    

file_like = io.StringIO(DATA_STRING)
df = pd.read_csv(file_like, sep='\s+')

df['run_number'] = df['run_id'] + 1

However, I am struggeling with generating the second column. The conceptional approach should be as follows:

  1. Group dataframe df by experiment_id to access each experiment block individually.
  2. Apply a function to each group which determines the number of runs in each group (equivalent to the length of the group or the number of rows).
  3. Return a Series with the same length as the group containing the group length as an integer value.
  4. Combine/tack all series to a single series which can then be assigned to the dataframe df as a new column.

Despite of creating the new column, using an ugly for-loop would look like this:

for name, group in df.groupby('experiment_id'):
    group.loc[:, 'total_runs'] = group['run_id'].count()
    
    print(group, end='\n\n')

As this is an ugly approach, I do not want to go down that rabbit hole. Especially with the warning given by Pandas:

C:\Users\Albert\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\indexing.py:376: SettingWithCopyWarning: A value is trying to be set on a copy of a slice from a DataFrame. Try using .loc[row_indexer,col_indexer] = value instead

As my idea was to call a function on each group and return a series with the relevant information, I had a look at the documentation. Reading the guide Groupby: Split, Apply, Combine from the docs, I stumbled upon .transform() which I had a more detailed look at.

Calling .transform() like

df.groupby('experiment_id').transform(lambda x: len(x))

produces the desired output:

    run_id  value  other_data  run_number
0   3       3      3           3
1   3       3      3           3
2   3       3      3           3
3   3       3      3           3
4   3       3      3           3
5   3       3      3           3
6   4       4      4           4
7   4       4      4           4
8   4       4      4           4
9   4       4      4           4

However, calling this exact same line to create a new column

df['total_runs'] = df.groupby('experiment_id').transform(lambda x: len(x))

raises several KeyError and a ValueError:

KeyError                                  Traceback (most recent call last) ~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\indexes\base.py in get_loc(self, key, method, tolerance)    2896             try:
-> 2897                 return self._engine.get_loc(key)    2898             except KeyError:

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'total_runs'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last) ~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\internals\managers.py in set(self, item, value)    1068         try:
-> 1069             loc = self.items.get_loc(item)    1070         except KeyError:

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\indexes\base.py in get_loc(self, key, method, tolerance)    2898             except KeyError:
-> 2899                 return self._engine.get_loc(self._maybe_cast_indexer(key))    2900         indexer = self.get_indexer([key], method=method, tolerance=tolerance)

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'total_runs'

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last) <ipython-input-11-874dd354de5d> in <module>
----> 1 df['total_runs'] = df.groupby('experiment_id').transform(lambda x: len(x))

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\frame.py in __setitem__(self, key, value)    3485         else:    3486        
# set column
-> 3487             self._set_item(key, value)    3488     3489     def _setitem_slice(self, key, value):

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\frame.py in _set_item(self, key, value)    3563         self._ensure_valid_index(value)    3564         value = self._sanitize_column(key, value)
-> 3565         NDFrame._set_item(self, key, value)    3566     3567         # check if we are modifying a copy

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\generic.py in _set_item(self, key, value)    3379     3380     def
_set_item(self, key, value):
-> 3381         self._data.set(key, value)    3382         self._clear_item_cache()    3383 

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\internals\managers.py in set(self, item, value)    1070         except KeyError:    1071    
# This item wasn't present, just insert at end
-> 1072             self.insert(len(self.items), item, value)    1073             return    1074 

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\internals\managers.py in insert(self, loc, item, value, allow_duplicates)    1179         new_axis = self.items.insert(loc, item)    1180 
-> 1181         block = make_block(values=value, ndim=self.ndim, placement=slice(loc, loc + 1))    1182     1183         for blkno, count in _fast_count_smallints(self._blknos[loc:]):

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\internals\blocks.py in make_block(values, placement, klass, ndim, dtype, fastpath)    3282 values = DatetimeArray._simple_new(values, dtype=dtype)    3283 
-> 3284     return klass(values, ndim=ndim, placement=placement)    3285     3286 

~\AppData\Local\Programs\Python\Python37\lib\site-packages\pandas\core\internals\blocks.py in __init__(self, values, placement, ndim)
    126             raise ValueError(
    127                 "Wrong number of items passed {val}, placement implies "
--> 128                 "{mgr}".format(val=len(self.values), mgr=len(self.mgr_locs))
    129             )
    130 

ValueError: Wrong number of items passed 4, placement implies 1

As a KeyError is raised at the very beginning and in order to move on, I implemented a small hack (really ugly):

df['total_runs'] = np.zeros_like(df['run_id'])
df['total_runs'] = df.groupby('experiment_id').transform(lambda x: len(x))

Finally, that did the trick. However, I want to remove that ugly hack and generate the required column based on Pandas' great (GroupBy object) features. How can I achieve this?

3 Answers

Command df.groupby('experiment_id').transform(lambda x: len(x)) returns 4 columns.

Hence, when you are trying to save the above command's output in total_runs which is just 1 column, it naturally fails.

df['total_runs'] = df.groupby('experiment_id').transform(lambda x: len(x))

Instead, do this:

In [1517]: df['total_runs'] = df.groupby('experiment_id')['run_number'].transform('count')

In [1518]: df
Out[1518]: 
  experiment_id  run_id  value  other_data  run_number  total_runs
0        9tfc6d       0  0.448       0.883           1           3
1        9tfc6d       1  0.963       0.230           2           3
2        9tfc6d       2  0.711       0.724           3           3
3        q9tqjq       0  0.748       0.959           1           3
4        q9tqjq       1  0.662       0.772           2           3
5        q9tqjq       2  0.530       0.834           3           3
6        jsxp2m       0  0.087       0.346           1           3
7        jsxp2m       1  0.362       0.569           2           3
8        jsxp2m       2  0.124       0.206           3           3

You can use groupby, count and merge:

df.merge(
        df[['experiment_id', 'run_id']].groupby('experiment_id').count().rename(columns={'run_id': 'total_runs'}),
        how='left',
        left_on='experiment_id',
        right_on='experiment_id',
    )

Result:

  experiment_id  run_id  value  other_data  run_number  total_runs
0        9tfc6d       0  0.448       0.883           1           3
1        9tfc6d       1  0.963       0.230           2           3
2        9tfc6d       2  0.711       0.724           3           3
3        q9tqjq       0  0.748       0.959           1           3
4        q9tqjq       1  0.662       0.772           2           3
5        q9tqjq       2  0.530       0.834           3           3
6        jsxp2m       0  0.087       0.346           1           3
7        jsxp2m       1  0.362       0.569           2           3
8        jsxp2m       2  0.124       0.206           3           3

Your are assigning four columns to a single column:

df['total_runs'] = df.groupby('experiment_id').transform(lambda x: len(x))

This is also what the error message says at the very bottom: Wrong number of items passed 4, placement implies 1

Limit this to a single column:

df['total_runs'] = df.groupby('experiment_id').iloc[:, 0].transform(lambda x: len(x))
Related