What's the fastest way to find all occurences of the substring in large string in python

Question

Edited to clarify input/output. I think this will be somewhat slow no matter what, but up until now I haven't really considered speed in my python scripts, and I'm trying to figure out ways to speed up operations like these.

My input is pickled dictionaries of the genome sequences. I'm currently working with two genomes, the budding yeast genome (11.5 MB on disk), and the human genome (2.8 GB on disk). These dictionaries have the form:

seq_d = { 'chr1' : 'ATCGCTCGCTGCTCGCT', 'chr2' : 'CGATCAGTCATGCATGCAT', 
        'chr3' : 'ACTCATCATCATCATACTGGC' }

I want to find all single-base instances of a nucleotide(s) in both strands of the genome. Where the + strand refers to the sequence in the above dictionary, and the - strand is the reverse complement of the sequences. My output is a nested dictionary, where the top level keys are + or -, the nested keys are chromosome names, and the values are lists of 0-indexed positions:

nts = 'T'
test_d = {'+': {'chr3': [2, 5, 8, 11, 14, 17], 'chr2': [3, 7, 10, 14, 18], 
           'chr1': [1, 5, 9, 12, 16]}, '-': {'chr3': [0, 4, 7, 10, 13, 15], 
           'chr2': [2, 5, 9, 13, 17], 'chr1': [0]}}

test_d defines a set of positions to examine in a large Illumina sequencing dataset later in the script.

My first attempt uses enumerate, and iteration.

import time
import numpy as np

rev_comps = { 'A' : 'T', 'T' : 'A', 'G' : 'C', 'C' : 'G', 'N' : 'N'}
test_d = { '+' : {}, '-' : {}}
nts = 'T'

s = time.time()
for chrom in seq_d:
    plus_pos, minus_pos = [], []
    chrom_seq = seq_d[chrom]
    for pos, nt in enumerate(chrom_seq):
        if nt in nts:
            plus_pos.append(pos)
        if rev_comps[nt] in nts:
            minus_pos.append(pos)

    test_d['+'][chrom] = plus_pos
    test_d['-'][chrom] = minus_pos
e = time.time()
print 'The serial version took {} minutes...'.format((e-s)/60)

The output for yeast:

The serial version took 0.0455190300941 minutes...

The output for human:

The serial version took 10.1694028815 minutes...

I tried using numpy arrays rather than enumerate() and iteration:

s = time.time()
for chrom in seq_d:
    chrom_seq = np.array(list(seq_d[chrom]))
    nts = list(nts)
    rev_nts = [rev_comps[nt] for nt in nts]
    plus_pos = list(np.where(np.in1d(chrom_seq, nts) == True)[0])
    minus_pos = list(np.where(np.in1d(chrom_seq, rev_nts) == True)[0])

    test_d['+'][chrom] = plus_pos
    test_d['-'][chrom] = minus_pos
e = time.time()
print 'The numpy version took {} minutes...'.format((e-s)/60)

The output for yeast:

The numpy version took 0.0309354345004 minutes...

The output for human:

The numpy version took 9.86174853643 minutes...

Why does the numpy version lose its performance advantage for the larger human genome? Is there a faster way to do this? I tried implementing a version using multiprocessing.Pool, but that's slower than either version:

def getTestHelper(chrom_seq, nts, rev_comp):

    rev_comps = { 'A' : 'T', 'T' : 'A', 'G' : 'C', 'C' : 'G', 'N' : 'N'}
    if rev_comp:
        nts = [rev_comps[nt] for nt in nts]
    else:
        nts = list(nts)
    chrom_seq = np.array(list(chrom_seq))
    mask = np.in1d(chrom_seq, nts)
    pos_l = list(np.where(mask == True)[0])
    return pos_l

s = time.time()
pool = Pool(4)
plus_pos = pool.map(functools.partial(getTestHelper, nts=nts, rev_comp=False), seq_d.values())
minus_pos = pool.map(functools.partial(getTestHelper, nts=nts, rev_comp=True), seq_d.values())
e = time.time()
print 'The parallel version took {} minutes...'.format((e-s)/60)

I haven't run this on the human genome, but the yeast version is slower:

The parallel version took 0.652778700987 minutes...

Answer 1

Use built-ins

Instead of manually iterating through your long string, try str.find or str.index. Don't slice the string yourself, use these methods' built-in slicing.

This also ditches enumerate-ing, although that shouldn't be costly anyway.

Also, you could use set to store indices, not list - additions could be faster.

You would have to do it twice, though, to find both your nucleotide and its complement. Of course, look up the complement outside of the loop.

Try regular expressions

You could also try regular expressions to do the same thing (if you are going to try this, try both 2 regex (for "T" and "A") and one for "T|A").

Also, instead of doing

for chrom in seq_d:

You could do

for chromosome_number, chomosome in seq_d.items():

Which has little to do with performance, but makes the code more readable.

Answer 2

To be honest, I think you've been doing the right things.

There are a few more tweaks you can make to your code, though. In general, when performance is key, only do the bare minimum in your innermost loops. Looking through your code, there are still some quick optimizations left on this front:

1. Use if...elif instead of if...if.
2. Don't use lists where you don't have to - e.g. just a single string is sufficient for nts and the reverse.
3. Don't evaluate the same result multiple times - e.g. the reverse lookup.

I'm guessing that your problem with multi-processing is down to the serialization of these very large strings, offsetting any performance gain you might have from running in parallel. However, there may be another way to do this - see Parallelizing a Numpy vector operation. I can't verify as I am having difficulty installing numexpr.

Putting them together and trying out some of the other suggestions in this trail, I get the following results:

$ python test.py
Original serial took 0.08330821593602498 minutes...
Using sets took 0.09072601397832235 minutes...
Using built-ins took 0.061421032746632895 minutes...
Using regex took 0.11649663050969442 minutes...
Optimized serial took 0.05909080108006795 minutes...
Original numpy took 0.04050511916478475 minutes...
Optimized numpy took 0.03438538312911987 minutes...

My code is as follows.

import time
import numpy as np
from random import choice
import re

# Create single large chromosome for the test...
seq = ""
for i in range(10000000):
    seq += choice("ATGCN")
seq_d = {"Chromosome1": seq}


rev_comps = { 'A' : 'T', 'T' : 'A', 'G' : 'C', 'C' : 'G', 'N' : 'N'}
nts = 'T'

# Original serial implementation
def serial():
    test_d = { '+' : {}, '-' : {}}
    for chrom in seq_d:
        plus_pos, minus_pos = [], []
        chrom_seq = seq_d[chrom]
        for pos, nt in enumerate(chrom_seq):
            if nt in nts:
                plus_pos.append(pos)
            if rev_comps[nt] in nts:
                minus_pos.append(pos)

        test_d['+'][chrom] = plus_pos
        test_d['-'][chrom] = minus_pos

# Optimized for single character tests
def serial2():
    test_d = {'+': {}, '-': {}}
    rev_nts = rev_comps[nts]
    for chrom in seq_d:
        plus_pos, minus_pos = [], []
        chrom_seq = seq_d[chrom]
        for pos, nt in enumerate(chrom_seq):
            if nt == nts:
                plus_pos.append(pos)
            elif nt == rev_nts:
                minus_pos.append(pos)

        test_d['+'][chrom] = plus_pos
        test_d['-'][chrom] = minus_pos

# Use sets instead of lists
def set_style():
    test_d = { '+' : {}, '-' : {}}
    for chrom in seq_d:
        plus_pos, minus_pos = set(), set()
        chrom_seq = seq_d[chrom]
        for pos, nt in enumerate(chrom_seq):
            if nt in nts:
                plus_pos.add(pos)
            if rev_comps[nt] in nts:
                minus_pos.add(pos)

        test_d['+'][chrom] = plus_pos
        test_d['-'][chrom] = minus_pos

# Use regex to find either nucleotide...
def regex_it():
    test_d = { '+' : {}, '-' : {}}
    search = re.compile("(T|A)")
    for chrom in seq_d:
        pos = 0
        plus_pos, minus_pos = [], []
        chrom_seq = seq_d[chrom]
        for sub_seq in search.split(chrom_seq):
            if len(sub_seq) == 0:
                continue
            if sub_seq[0] == 'T':
                plus_pos.append(pos)
            elif sub_seq[0] == 'A':
                minus_pos.append(pos)
            pos += len(sub_seq)

        test_d['+'][chrom] = plus_pos
        test_d['-'][chrom] = minus_pos

# Use str.find instead of iteration
def use_builtins():
    test_d = { '+' : {}, '-' : {}}
    for chrom in seq_d:
        plus_pos, minus_pos = [], []
        chrom_seq = seq_d[chrom]
        start = 0
        while True:
            pos = chrom_seq.find("T", start)
            if pos == -1:
                break
            plus_pos.append(pos)
            start = pos + 1

        start = 0
        while True:
            pos = chrom_seq.find("A", start)
            if pos == -1:
                break
            minus_pos.append(pos)
            start = pos + 1

        test_d['+'][chrom] = plus_pos
        test_d['-'][chrom] = minus_pos

# Original numpy implementation
def numpy1():
    test_d = { '+' : {}, '-' : {}}
    for chrom in seq_d:
        chrom_seq = np.array(list(seq_d[chrom]))
        for_nts = list(nts)
        rev_nts = [rev_comps[nt] for nt in nts]
        plus_pos = list(np.where(np.in1d(chrom_seq, for_nts) == True)[0])
        minus_pos = list(np.where(np.in1d(chrom_seq, rev_nts) == True)[0])

        test_d['+'][chrom] = plus_pos
        test_d['-'][chrom] = minus_pos

# Optimized for single character look-ups
def numpy2():
    test_d = {'+': {}, '-': {}}
    rev_nts = rev_comps[nts]
    for chrom in seq_d:
        chrom_seq = np.array(list(seq_d[chrom]))
        plus_pos = np.where(chrom_seq == nts)
        minus_pos = np.where(chrom_seq == rev_nts)

        test_d['+'][chrom] = plus_pos
        test_d['-'][chrom] = minus_pos

for fn, name in [
        (serial, "Original serial"),
        (set_style, "Using sets"),
        (use_builtins, "Using built-ins"),
        (regex_it, "Using regex"),
        (serial2, "Optimized serial"),
        (numpy1, "Original numpy"),
        (numpy2, "Optimized numpy")]:
    s = time.time()
    fn()
    e = time.time()
    print('{} took {} minutes...'.format(name, (e-s)/60))

Answer 3

Finds the longest chromosome string, and creates an empty array with one row per chromosome, and columns up to the longest one in the dictionary. Then it puts each chromosome into its own row, where you can call np.where on the whole array

import numpy as np

longest_chrom = max([len(x) for x in seq_d.values()])
genome_data = np.zeros((len(seq_d), dtype=str, longest_chrom)
for index, entry in enumerate(seq_d):
    gen_string = list(seq_d[entry]
    genome_data[index, :len(gen_string) + 1] =  gen_string
nuc_search = "T"
nuc_locs = np.where(genome_data == nuc_search)

Using this method, it's slightly different for strings of > 1 nucleotide, but I'm sure it's achievable in only a couple more steps.