Building a frequency dictionary from a Pandas dataframe

Question

I need to make a frequency dictionary from a pandas series (from the 'amino_acid' column in dataframe below) that also adds an adjacent row for each entry in the dictionary (from 'templates' column).

    templates   amino_acid
0   118         CAWSVGQYSNQPQHF
1   635         CASSLRGNQPQHF
2   468         CASSHGTAYEQYF
3   239         CASSLDRLSSGEQYF
4   51          CSVEDGPRGTQYF

My current approach of iterating through the dataframe seems to be inefficient and even an anti-pattern. How can I improve the efficiency/use best practice for doing this?

My current approach:

sequence_counts = {}
seqs= list(zip(df.amino_acid, df.templates))

for seq in seqs:
    if seq[0] not in sequence_counts:
        sequence_counts[seq[0]] = 0
    sequence_counts[seq[0]] += seq[1]

I've seen people the below way, but can't figure out how to adjust it to add each respective 'templates' entry:

sequence_counts = df['amino_acid'].value_counts().to_dict()

Any help/feedback would be greatly appreciated! :)

Answer 1

Yes, iterating through a dataframe usually is at odds with the spirit of dataframes and numpy arrays. They are best suited for vectorized operations, which are operations applied in bulk to rows/columns of the data structure.

Assuming that you've cast templates as an int64 or similar, you can use pd.groupby to do the same sort of grouping as you are doing with your dictionary:

import pandas as pd

df
        amino_acid  templates
0  CAWSVGQYSNQPQHF        118
1    CASSLRGNQPQHF        635
2    CASSHGTAYEQYF        468
3  CASSLDRLSSGEQYF        239
4    CSVEDGPRGTQYF         51
5  CASSLDRLSSGEQYF         66    # I've added this extra row here to show the effect

# these act as Series objects, so you can add together the
# grouped templates values
df.groupby('amino_acid')
# pd.groupby object

# use pd.Series.sum() to do this:
df.groupby('amino_acid').sum()

                 templates
amino_acid
CASSHGTAYEQYF          468
CASSLDRLSSGEQYF        305    # this was added for the two amino acids
CASSLRGNQPQHF          635
CAWSVGQYSNQPQHF        118
CSVEDGPRGTQYF           51

So what's going on in df.groupby? Well, you give it what to group on first. In this case, you group on the value for amino_acid. This creates a data structure that looks quite familiar if you've used itertools.groupby: tuples of (grouping key, DataFrame) pairs, where the DataFrame is a subset that matches the key. For example:

tmp = [(x, y) for x,y in df.groupby('amino_acid')]

[('CASSHGTAYEQYF',       amino_acid  templates
2  CASSHGTAYEQYF        468), 
('CASSLDRLSSGEQYF',         amino_acid  templates
3  CASSLDRLSSGEQYF        239
5  CASSLDRLSSGEQYF         66), 
('CASSLRGNQPQHF',       amino_acid  templates
1  CASSLRGNQPQHF        635), 
('CAWSVGQYSNQPQHF',         amino_acid  templates
0  CAWSVGQYSNQPQHF        118), 
('CSVEDGPRGTQYF',       amino_acid  templates
4  CSVEDGPRGTQYF         51)]

And per the docs, [df.sum] will return the sum on the specified axis (1 by default). So, for tmp[1], which contains two rows:

tmp[1][1].sum()

amino_acid    CASSLDRLSSGEQYFCASSLDRLSSGEQYF
templates                                305
dtype: object

Where 305 is the sum. The pd.groupby object as a whole supports the .sum call, so we are able to call it like we did above.

# Now, using `to_dict`, you can see that you want what's inside templates
df.groupby("amino_acid").sum().to_dict()

# {'templates': {'CASSHGTAYEQYF': 468, 'CASSLDRLSSGEQYF': 305, 'CASSLRGNQPQHF': 635, 'CAWSVGQYSNQPQHF': 118, 'CSVEDGPRGTQYF': 51}}

# so use the templates attribute to grab it
df.groupby("amino_acid").sum().templates.to_dict()
# {'CASSHGTAYEQYF': 468, 'CASSLDRLSSGEQYF': 305, 'CASSLRGNQPQHF': 635, 'CAWSVGQYSNQPQHF': 118, 'CSVEDGPRGTQYF': 51}

This applies the operations within the dataframe, which is more efficient than a loop. The analogue that you were trying to use could leverage defaultdict from the collections module, as well. It prevents the need to check for a key's existence, speeding up the loop immensely:

from collections import defaultdict

# specify your input type here
sequence_counts = defaultdict(int)

# it is more pythonic to use tuple-unpacking in loops
# as indexing is less readable
for amino, template in zip(df.amino_acid, df.templates):
    sequence_counts[amino] += template

Also, by iterating over the zip object directly, you don't copy data into memory. list will aggregate all of the members into a data structure, whereas zip is a generator that will just produce the members one at a time until exhausted. It's like the difference between for x in range and for x in list(range):

# This will just run for a really really really long time
for i in range(1000000000000): 
    print(i)

# This will crash your computer, it will never
# get to the print statement, because it must evaluate
# list(range) before it starts the loop
for i in list(range(1000000000000)):
    print(i)