Pandas dataframe to nested counter dictionary

Viewed 159

I've seen a lot of questions on how to convert pandas dataframes to nested dictionaries, but none of them deal with aggregating the information. I may even be able to do what I need within pandas, but I'm stuck.

Input

I have a dataframe that looks like this:

  FeatureID    gene  Target  pos  bc_count
0     1_1_1  NRAS_3  TAGCAC    0      0.42
1     1_1_1  NRAS_3  TGCACA    1      1.00
2     1_1_1  NRAS_3  GCACAA    2      0.50
3     1_1_1  NRAS_3  CACAAA    3      2.00
4     1_1_1  NRAS_3  CAGAAA    3      0.42

# create df as below
import pandas as pd
df = pd.DataFrame([{"FeatureID":"1_1_1", "gene":"NRAS_3", "Target":"TAGCAC", 
   "pos":0, "bc_count":.42},
   {"FeatureID":"1_1_1", "gene":"NRAS_3", "Target":"TGCACA", "pos":1, 
   "bc_count":1.00},
   {"FeatureID":"1_1_1", "gene":"NRAS_3", "Target":"GCACAA", "pos":2, 
   "bc_count":0.50},
   {"FeatureID":"1_1_1", "gene":"NRAS_3", "Target":"CACAAA", "pos":3, 
   "bc_count":2.00},
   {"FeatureID":"1_1_1", "gene":"NRAS_3", "Target":"CAGAAA", "pos":4, 
   "bc_count":0.42}])

The problem

I need to break apart the Target column for each row to return a tuple of (position, letter, count), where the starting position is given in the "pos" column, and then enumerating the string for each position following, and the count is the value found for that row in the "bc_count" column.

For example, in the first row, the desired list of tuples would be:

[(0, "T", 0.42), (1,"A", 0.42), (2,"G", 0.42), (3,"C", 0.42), (4,"A", 0.42), (5,"C", 0.42)]

What I've tried

I've created code that breaks up the target column into the position found, returning a tuple of position, nucleotide (letter) and count for that letter, and adds them as a column to the dataframe:

def index_target(row):
    count_list = [((row.pos + x),y, 
        row.bc_count) for x,y in 
        enumerate(row.Target)]

df['pos_count'] = df.apply(self.index_target, axis=1)

Which returns a list of tuples for each row based on that row's target column.

I need to take every row in df, for each target, and sum the counts. Which is why I thought of using a dictionary as a counter:

position[letter] += bc_count

I've tried creating a defaultdict, but it is appending each list of tuples separately instead of summing the counts for each position:

from collections import defaultdict

d = defaultdict(dict) # also tried defaultdict(list) here
for x,y,z in row.pos_count:
    d[x][y] += z

Desired Output

For each feature in the dataframe, where the numbers below represent a sum of the individual counts found in the bc_count column for each position and x representing positions where ties were found and no one letter can be returned as the max:

pos A   T   G   C
0   25  80  25  57
1   32  19  100 32
2   27  18  16  27
3   90  90  90  90
4   10  42  37  18

consensus= TGXXT

2 Answers

Unsure how to get your desired output, but I created the list d which contains the tuples you desired for a dataframe. Hopefully, it provides some direction in what you want to create:

d = []

for t,c,p in zip(df.Target,df.bc_count,df.pos):
    d.extend([(p,c,i) for i in list(t)])

df_new = pd.DataFrame(d, columns = ['pos','count','val'])
df_new = df_new.groupby(['pos','val']).agg({'count':'sum'}).reset_index()

df_new.pivot(index = 'pos', columns = 'val', values = 'count')

This may not be the most elegant solution, but I think it might accomplish what you need:

new_df = pd.DataFrame(
    df.apply(
        # this lambda is basically the same thing you're doing,
        # but we create a pd.Series with it
        lambda row: pd.Series(
            [(row.pos + i, c, row.bc_count) for i, c in enumerate(row.Target)]
        ),
        axis=1)
        .stack().tolist(),
    columns=["pos", "nucl", "count"]

)

Where new_df looks like this:

  pos nucl count
0   0    T  0.42
1   1    A  0.42
2   2    G  0.42
3   3    C  0.42
4   4    A  0.42
5   5    C  0.42
6   1    T  1.00
7   2    G  1.00
8   3    C  1.00
9   4    A  1.00

Then I would pivot this to get the aggregated counts:

nucleotide_count_by_pos = new_df.pivot_table(
    index="pos",
    columns="nucl",
    values="count",
    aggfunc="sum",
    fill_value=0
)

Where nucleotide_count_by_pos looks like:

nucl     A     C     G     T
 pos
   0  0.00  0.00  0.00  0.42
   1  0.42  0.00  0.00  1.00
   2  0.00  0.00  1.92  0.00
   3  0.00  4.34  0.00  0.00
   4  4.34  0.00  0.00  0.00

And then to get the consensus:

def get_consensus(row):
    max_value = row.max()
    nuc = row.idxmax()
    if (row == max_value).sum() == 1:
        return nuc
   else:
        return "X"

consensus = ''.join(nucleotide_count_by_pos.apply(get_consensus, axis=1).tolist())

Which in the case of your example data would be:

'TTGCACAAA'
Related