Trouble to run a multiprocessing kmer count script

Viewed 60

Hi there I would like to look for words/kmers/ngrams in genomes. I have exactly 198000 files each containing one genome of length N. I have some code like:

def count_kmers(sequence, alphabet, k, base_filter=False):
    seq, alphabet = sequence.upper(), set(alphabet.upper())
    seq_len = len(seq)
    count = Counter(seq[i:i+k] for i in range(0, seq_len - k +1))
    if base_filter:
        return dict((key,value) for key, value in
                    count.items() if set(key).issubset(alphabet))
    return count

def kmer_positions(sequence, alphabet, k):
    """ returns the position of all k-mers in sequence as a dictionary"""
    mer_position = defaultdict(list)
    for i in range(1,len(sequence) - k + 1):
        kmer = sequence[i:i+k]
        if all(base in set(alphabet) for base in kmer):
            mer_position[kmer] = mer_position.get(kmer,[]) + [i]
    # combine kmers with their reverse complements
    pair_position = defaultdict(list)
    for kmer, pos in mer_position.items():
        krev = get_reverse_complement(kmer)
        if kmer < krev:
            pair_position[kmer] = sorted(pos + mer_position.get(krev, []))
        elif krev < kmer:
            pair_position[krev] = sorted(mer_position.get(krev, []) + pos)
        else:
            pair_position[kmer] = pos
    return pair_position

And others, but the kmer_positions(sequence, alphabet, k) function take at least 2min 42s to finish one genome of length 6407042. And the count_kmers I tried multiprocessing and it was from 1.3 s to 97.3 µs.

I am using multiprocessing like this:

args = [alphabet='ACGT', k=6]
f = partial(kmer_positions, genome)
pool =  multiprocessing.Pool(4)
r = pool.apply_async(f, args) # use apply_async that was the one who works to my arguments
r.get().keys()

I tried to find a good material to study, but still not satisfied. Any help to get a better use of this tool?

Thanks PS - I know there are many tools specific for this task like khmer, kat, etc... But I am learning python and I would like to learn the language best as I can

0 Answers
Related