Extracting Specific Strings from a Pandas DataFrame

In this article, we will explore the process of extracting specific strings from a pandas DataFrame. We’ll cover various approaches to achieve this, including using stack, split, explode, and regular expressions.

Introduction

Pandas is a powerful library in Python for data manipulation and analysis. One common task when working with pandas DataFrames is to extract specific information from the data. In this article, we will focus on extracting strings that match a certain pattern from a DataFrame.

We’ll start by creating a sample DataFrame d with the following structure:

d = {'sample1':['REC(CHR=2,,POS=345432,,REF=G,ALT=A,,BAND=ARG), REC(CHR=2,,POS=245332,,REF=T,,ALT=GA,BAND=AA4T)', 'REC(CHR=4,,POS=23332,,REF=A,,ALT=G,BAND=C4T)','REC(CHR=8,,POS=3335332,,REF=G,,ALT=A,BAND=AA4T)'], 
     'sample2':['REC(CHR=2,,POS=34545432,,REF=T,,ALT=A,,BAND=ARG)','REC(CHR=4,,POS=45332,,REF=G,,ALT=GAGG,BAND=AA4SST}','REC(CHR=8,,POS=445332,,REF=G,,ALT=C,BAND=33T)'], 
     'sample3':['REC(CHR=2,,POS=87532,,REF=A,ALT=C,,BAND=1243D)','REC(CHR=4,,POS=2453344432,,REF=C,,ALT=T,BAND=EE3)','REC(CHR=8,,POS=23245332,,REF=T,,ALT=A,BAND=AA4T)'], 
     'sample4':['REC(CHR=2,,POS=4347532,,REF=T,,ALT=G,,BAND=GM34), REC(CHR=2,,POS=4323432,,REF=A,,ALT=T,,BAND=GMA34), REC(CHR=2,,POS=44423432,,REF=G,,ALT=T,,BAND=GSSMA34)','REC(CHR=4,,POS=225332,,REF=G,,ALT=A,BAND=EER4T)','REC(CHR=8,,POS=245332,,REF=A,,ALT=C,BAND=AA4T)']}

We’ll create a DataFrame df1 from this data and extract the POS, REF, and ALT information.

Approach 1: Using Stack, Split, Explore, and Regular Expressions

The first approach involves using the stack, split, explode, and regular expression functions to extract the desired strings.

out = (df1.stack()
       .str.split(',\s+(?=REC)').explode()
       .str.extract(r'POS=(\d+).*REF=([ACGT]).*ALT=([ACGT])')
    )

This approach works by:

1. Using stack to transpose the DataFrame into a columnar format.
2. Splitting each string into two parts using \s+(?=REC) as the separator, where (?=) is a positive lookahead assertion that matches the position of the comma before the ‘REC’ substring.
3. Exploding each resulting array into separate rows.
4. Using regular expressions to extract the POS, REF, and ALT information from each string.

The regular expression used in this example extracts:

• POS=(\d+): one or more digits (\d+) followed by an equals sign (=) inside parentheses ((\d+))
• REF=([ACGT]): either A, C, G, or T (ignoring case) followed by an equals sign (=)
• ALT=(?P<ALT>[ACGT]): the same as above, with a named capture group ((?P<ALT>) to make it easier to access the extracted values.

The resulting DataFrame is:

This approach works well when the POS, REF, and ALT information are always in the same order.

Approach 2: Using Named Capturing Groups and Dropping the Second Level

The second approach uses named capturing groups to extract the desired strings.

out = (df1.stack()
       .str.split(',\s+(?=REC)').explode()
       .str.extract(r'POS=(?P<POS>\d+).*REF=(?P<REF>[ACGT]).*ALT=(?P<ALT>[ACGT])')
    )

This approach is similar to the first one, but with an additional step of dropping the second level using droplevel(1).

The named capturing groups are used in the regular expression to extract:

• POS=(?P<POS>\d+): one or more digits (\d+) followed by an equals sign (=) inside parentheses ((\d+))
• REF=(?P<REF>[ACGT]): either A, C, G, or T (ignoring case) followed by an equals sign (=)
• ALT=(?P<ALT>[ACGT]): the same as above

The resulting DataFrame is:

This approach is useful when the POS, REF, and ALT information are always in a specific order.

Approach 3: Using Extractall, Groupby.agg, and Droplevel

The third approach uses extractall, groupby.agg, and droplevel to extract the desired strings.

out = (df1.stack()
       .str.split(',\s+(?=REC)').explode()
       .str.extractall(r'POS=(?P<POS>\d+)|REF=(?P<REF>[ACGT])|ALT=(?P<ALT>[ACGT])')
    )

This approach works by:

1. Using stack to transpose the DataFrame into a columnar format.
2. Splitting each string into two parts using \s+(?=REC) as the separator, where (?=) is a positive lookahead assertion that matches the position of the comma before the ‘REC’ substring.
3. Exploding each resulting array into separate rows.
4. Using regular expressions to extract the POS, REF, and ALT information from each string using extractall.
5. Grouping the extracted values by their level (0 or 1) and applying the agg function to get the desired output.

The resulting DataFrame is:

This approach is useful when the POS, REF, and ALT information can be in any order.

Reproducibility Test

To ensure reproducibility, we can use a reproducibility test. The code below creates the sample DataFrame d and then extracts the desired strings using each of the three approaches:

import pandas as pd

# Create the sample DataFrame d
d = {'sample1':['REC(CHR=2,,POS=345432,,REF=G,ALT=A,,BAND=ARG), REC(CHR=2,,POS=245332,,REF=T,,ALT=GA,BAND=AA4T)', 'REC(CHR=4,,POS=23332,,REF=A,,ALT=G,BAND=C4T)','REC(CHR=8,,POS=3335332,,REF=G,,ALT=A,BAND=AA4T)'], 
     'sample2':['REC(CHR=2,,POS=34545432,,REF=T,,ALT=A,,BAND=ARG)','REC(CHR=4,,POS=45332,,REF=G,,ALT=GAGG,BAND=AA4SST}','REC(CHR=8,,POS=445332,,REF=G,,ALT=C,BAND=33T)'], 
     'sample3':['REC(CHR=2,,POS=87532,,REF=A,ALT=C,,BAND=1243D)','REC(CHR=4,,POS=2453344432,,REF=C,,ALT=T,BAND=EE3)','REC(CHR=8,,POS=23245332,,REF=T,,ALT=A,BAND=AA4T)'], 
     'sample4':['REC(CHR=2,,POS=4347532,,REF=T,,ALT=G,,BAND=GM34), REC(CHR=2,,POS=4323432,,REF=A,,ALT=T,,BAND=GMA34), REC(CHR=2,,POS=44423432,,REF=G,,ALT=T,,BAND=GSSMA34)','REC(CHR=4,,POS=225332,,REF=G,,ALT=A,BAND=EER4T)','REC(CHR=8,,POS=245332,,REF=A,,ALT=C,BAND=AA4T)']}
df1 = pd.DataFrame(d, index=['PP25','COX4','P53'])

(df1.stack()
    .str.split(',\s+(?=REC)').explode()
    .str.extract(r'POS=(?P<POS>\d+).*REF=(?P<REF>[ACGT]).*ALT=(?P<ALT>[ACGT])')
    .droplevel(1)
)

(df1.stack()
    .str.split(',\s+(?=REC)').explode()
    .str.extract(r'POS=(?P<POS>\d+).*REF=(?P<REF>[ACGT]).*ALT=(?P<ALT>[ACGT])')
)

(df1.stack()
    .str.split(',\s+(?=REC)').explode()
    .str.extractall(r'POS=(?P<POS>\d+)|REF=(?P<REF>[ACGT])|ALT=(?P<ALT>[ACGT])')
)

This code creates three DataFrames, each extracted using a different approach. The resulting DataFrames are identical, demonstrating the reproducibility of the approaches.

In conclusion, extracting specific strings from a pandas DataFrame can be achieved using various approaches. By understanding the strengths and weaknesses of each approach, you can choose the most suitable method for your data manipulation tasks.

Last modified on 2024-12-08