I have a pandas Dataframe with a column of peptide sequences and I want to know how many times each each amino acid appears at each position. I have written the following code to create the position frequency matrix:

import pandas as pd
from itertools import chain

def frequency_matrix(df):
    # Empty position frequency matrix
    freq_matrix_df = pd.DataFrame(
        columns =  sorted(set(chain.from_iterable(df.peptide_alpha))),
        index=range(df.peptide_len.max()),
    ).fillna(0)

    for _, row in df.iterrows():
      for idx, aa in enumerate(row["peptide_alpha"]):
        freq_matrix_df.loc[idx, aa]  = 1
    
    return freq_matrix_df
    

which for the following sample DataFrame:

mini_df = pd.DataFrame(["YTEGDALDALGLKRY", 
                        "LTEIYGERLYETSY",
                        "PVEEFNELLSKY", 
                        "TVDIQNPDITSSRY", 
                        "ASDKETYELRY"], 
                       columns=["peptide_alpha"])
mini_df["peptide_len"] = mini_df["peptide_alpha"].str.len()

gives the following output:

This works for small DataFrames but because of the for loop becomes too slow for bigger datasets. Is there a way to rewrite this in a faster/vectorized way?

CodePudding user response：

Solution

mini_df['peptide_len'] = mini_df.peptide_len.map(lambda x: range(x))
mini_df['peptide_alpha'] = mini_df.peptide_alpha.map(list)
mini_df = mini_df.explode(["peptide_alpha", "peptide_len"])

pd.crosstab(mini_df.peptide_len, mini_df.peptide_alpha)

Performance

With the dataframe

mini_df = pd.concat([mini_df] * 10000)

On my machine, my solution solves the problem within 0.5s, whereas the solution of the OP takes 1m8.6s. Consequently, I believe that my solution can be useful for him.

Output

peptide_alpha  A  D  E  F  G  I  K  L  N  P  Q  R  S  T  V  Y
peptide_len                                                  
0              1  0  0  0  0  0  0  1  0  1  0  0  0  1  0  1
1              0  0  0  0  0  0  0  0  0  0  0  0  1  2  2  0
2              0  2  3  0  0  0  0  0  0  0  0  0  0  0  0  0
3              0  0  1  0  1  2  1  0  0  0  0  0  0  0  0  0
4              0  1  1  1  0  0  0  0  0  0  1  0  0  0  0  1
5              1  0  0  0  1  0  0  0  2  0  0  0  0  1  0  0
6              0  0  2  0  0  0  0  1  0  1  0  0  0  0  0  1
7              0  2  1  0  0  0  0  1  0  0  0  1  0  0  0  0
8              1  0  0  0  0  1  0  3  0  0  0  0  0  0  0  0
9              0  0  0  0  0  0  0  1  0  0  0  1  1  1  0  1
10             0  0  1  0  1  0  1  0  0  0  0  0  1  0  0  1
11             0  0  0  0  0  0  0  1  0  0  0  0  1  1  0  1
12             0  0  0  0  0  0  1  0  0  0  0  1  1  0  0  0
13             0  0  0  0  0  0  0  0  0  0  0  1  0  0  0  2
14             0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  1