What is an efficient method using pandas for retrieving data with a given start/end point when the datapoints also have start/end points?

Viewed 35

I am currently developing an application which leverages genomic count data. This data is loaded into a pandas dataframe. An example of this data would look like this:

chrom chromStart chromEnd height
chr3 93504854 93504940 0
chr3 93504940 93504977 1
chr3 93504977 93505016 2
chr3 93505016 93505018 1
chr3 93505018 93505053 2
chr3 93505053 93505093 1
chr3 93505093 93505151 0
chr3 93505151 93505227 1
chr3 93505227 93505252 0
chr3 93505252 93505328 1

I currently have an implementation of the title which accounts for quite a bit of my time taken. Here is the implementation:

def checkInBounds(row, chrom, chromStart, chromEnd):
    try:
        if not chrom == row['chrom']:
            return False
    except KeyError:
        print('row count', row.count)
        print("CheckInBoundsKeyError\nRow\n", row, '\n chrom', chrom, 'start', chromStart, 'end', chromEnd)
        return False

    if chromStart <= row['chromStart'] <= chromEnd:
        return True
    elif chromStart <= row['chromEnd'] <= chromEnd:
        return True
    else:
        return (row['chromStart'] < chromEnd) and (row['chromEnd'] > chromEnd)

This is ran using an apply statement which looks like this:

df.apply(checkInBounds, axis=1, args=(chrom, start, end))

Is there a faster way to accomplish this same task?

1 Answers

Here is a solution using the searchedsorted function in pandas Series. In a sorted Series, it tells you where you should put the value you pass to keep the data sorted.

Using your function on your data I get:

chrom = 'chr3'
start = 93504854
end = 93505053
df = df.sort_values(['chrom','chromStart','chromEnd'])
df.apply(checkInBounds, axis=1, args=(chrom, start, end))

    0     True
    1     True
    2     True
    3     True
    4     True
    5     True
    6    False
    7    False
    8    False
    9    False

Here is the function with search sorted.

def checkInBounds_new(df, start, end, chrom):
    bound1 = df.chromStart.searchsorted(start)
    bound2 = df.chromStart.searchsorted(end)
    df =  df.loc[bound1:bound2]
    return df[df['chrom'] == chrom]
checkInBounds_new(df, start, end, chrom):

which gives you:

chrom   chromStart  chromEnd    height
0   chr3    93504854    93504940    0
1   chr3    93504940    93504977    1
2   chr3    93504977    93505016    2
3   chr3    93505016    93505018    1
4   chr3    93505018    93505053    2
5   chr3    93505053    93505093    1

Notice how both functions basically returns the first 5 rows for these inputs. Also, I didn't implement any checks here like making sure that end > start and so on.

Related