Finding roughly matching genome sequences in Python dictionary

Question

The purpose of my code here is to play a part in genome sequencing analysis, and while functional it takes days to run, so I am looking for any way I can improve speed. The input is up to 500 million lines long (making speed code efficiency important) and contains sequencing reads and corresponding info. Each read takes up 4 lines within the input file and looks something like this:

@A001             <-header
AAAAACCCCCCCCCCCC <-seq read (finalRead)
+
################# <-quality (trimmed_quality)

The portion of my code that is very slow takes a dictionary as input, which contains all of the data found within the input sequencing file and is in the form shown below:

duplexDict[umi] = {'header':header, 'seq':finalRead, 'qual':trimmed_quality}

In the first part of the code I am looking for pairs of sequences by checking for similar keys (termed umi in the code). The goal is to find keys that when converted to complement sequence are only different by a single letter. Then for each key if there is only one closely matching key, the associated dictionaries are retained. If there are no matches or more than one matching key, all of these keys should be ignored.

from Levenshtein import distance
deDuplexDict = {} # dict that will contain key pairs
finalList = [] # list to keep track of valid key pairs
for i in duplexDict: # dict with sequencing file info
    tempList = []
    for j in duplexDict:
        complement = str(Seq(j).complement()) # this is just finding complementary sequence
        if distance(i,complement) <= 1: # find similar umi/read seq pairs
            tempList.append(j) # make a list of all similar pairs
    # only keep a complementary pair if there are exactly two matching consensus reads
    if len(tempList) == 1:
        if i not in finalList and j not in finalList:
            finalList.append(i)
            finalList.append(j)

# only retain those dict values that are true pairs
for key in finalList:
    deDuplexDict[key] = duplexDict[key]

The second piece is designed to now collapse combine the sequences of two matching dictionary keys together and output to file. This is done by taking the complement of one of the sequences and then comparing each character position along the sequence strings. If anything doesn't match the character in a final string is just set to 'N' rather than the character found in the reads.

from itertools import combinations
prevScanned=[]
plus = '+'

# only pairs now exist, just search for them
for key1, key2 in combinations(deDuplexDict, 2):
    finalRead = ''
    complement = str(Seq(key2).complement()) # complement of second read sequence
    # if neither key has been analysed and they are a matching pair then use for consensus read
    if distance(key1, complement) <= 1 and key1 not in prevScanned and key2 not in prevScanned:
        prevScanned.extend([key1,key2]) # keep track of analyzed keys

        # convert to complementary matches
        refRead = deDuplexDict[key1]['seq']
        compRead = str(Seq(deDuplexDict[key2]['seq']).complement())

        # iterate through by locus and derive consensus
        for base in range(readLength):
            if refRead[base] == compRead[base]:
                finalRead += refRead[base]
            else:
                finalRead += 'N' # only perfect matches are permitted

        # output consensus and associated info
        target = open(final_output_file, 'a')
        target.write(deDuplexDict[umi]['header'] + '\n' + finalRead + '\n' + plus + '\n' + deDuplexDict[umi]['qual'] + '\n')
target.close()

Answer 1

You need a faster search method

You are comparing every entry in duplexDict directly with every entry in duplexDict. This means the number of operations will increase with the square of the number of entries in duplexDict. This stands out from the lines:

    for i in duplexDict:
        ...
        for j in duplexDict:

More formally, you algorithm runs in \$\mathcal{O}(n^2)\$, where n is the length of the input dictionary. So, for 500 million (5e8) reads of sequence data, you need to run about 250 thousand trillion (25e16) operations. This is why it takes days to run.

You will need to index your reads based on the sequences themselves. Find and implement an architecture, whether hash tables, binary trees, or something else, that allows fast searching of the input list of sequencing reads.

Of course, a hashing method is built-in with Python's dict. There is no hard limit on the length of the key strings, and the number of entries you can put in the dict is limited by available memory.

Using dictionary search

In your case, in order to use Python's built-in dictionary to make and search the hash table, you might first do something like this:

seq_dict = {}
for i in duplexDict:
    seq_dict[i['seq']] = {info: i['header']}  # add whatever info you need to find
                                              # the original entry again in `duplexDict`}

The resulting dictionary has the sequences themselves as the keys. Searching for a particular sequence takes \$\mathcal{O}(1)\$ operations.

(Note: I am assuming you need to be able to refer back to the original duplexDict once you have your hits. Then, you don't need all of the extra information to accompany the sequencing reads in the new dictionary. If each entry in duplexDict has an identifier like umi, just put that alone as the value in the new dictionary.)

Generate mismatches and search

As you are only looking for sequences different by one base, you can just generate all possible mismatches. If you eventually want to include other types of sequence similarity, you will need to use more specialized sequence analysis tools (such as BLAST).

So, you will need a simple function to generate all possible one-base-pair mismatches from each input sequence. Assuming you are using only genomic bases A,T,C,G, for each sequence, there will be 3 x n possible one-base-pair mismatches, where n is the length of the sequencing reads.

def get_one_bp_mismatches(seq):
    mismatches = []
    bases = ['A','T','G','C']
    for e,i in enumerate(seq):
        for b in bases:
            mismatches.append(seq[:e] + b + seq[e+1:])
    return mismatches

Then, search the dictionary like this:

for seq_list,info in [get_one_bp_mismatches(i['seq']),i['info'] for i in duplexDict.items()]:
    for seq in seq_list:
        if seq in seq_dict:
            finalList.append({'search_seq': info,
                              'found_seq': seq_dict[seq]['info']})

Your finalList will contain all matching pairs, identified by whatever information you use to look them up in the original duplexDict.

The whole search process will take on the order of \$\mathcal{O}(n)\$ operations, and should likely finish within minutes for 500 million sequencing reads. You can then use the last lines of your existing code to generate the output file.

Answer 2

You have quite a few problems with your code:

• You should use a set for \$O(1)\$ containment checks. Where checking if something is in a list takes \$O(n)\$ time.
• You should make some functions
• You should make your first code block return an array of tuples, (i, j).
• You don't need deDuplexDict, it makes your code require itertools.combinations, which is unneeded.
• You shouldn't need prevScanned if you follow the above. But if you did want to keep it make it a set, for the reason stated above.
• Your range(readLength) can instead be a comprehension, and should use ''.join rather than \$O(n^2)\$ string concatenation.
• You should look into str.format to simplify the creation of complex strings.
• Rather than manually calling file.close() instead use with. It guarantees it'll be called, even if there's an error when writing to the file.

This along with changing your \$O(n^2)\$ algorithm to use a trie can get:

import itertools


def de_duplex(duplex):
    trie = Trie()
    for key in duplex:
        trie[str(Seq(key).complement())] = key

    islice = itertools.islice
    matched = set()
    for key in duplex:
        matches = trie.levenshtein(key, limit=1)
        matches = list(islice(matches, 2))
        if len(matches) != 1:
            continue
        match = matches[0]
        if not {key, match} & matched:
            matched |= {key, matched}
            yield key, match


def write_de_duplex(final_output_file, duplex, read_length, umi):
    for key1, key2 in de_duplex(duplex):
        ref_read = duplex[key1]['seq']
        comp_read = str(Seq(duplex[key2]['seq']).complement())
        final_read = ''.join(
            ref_read[base] if ref_read[base] == comp_read[base] else 'N'
            for base in range(read_length)
        )
        with open(final_output_file, 'a') as target:
            target.write('{umi[header]}\n{}\n+\n{umi[qual]}\n'
                         .format(final_read, umi=duplex[umi]))

---

I agree with @RyanMills on that you want to generate all the mismatches purely based around the key, rather than looping through the entire dictionary, duplex. However, I think their algorithm is quite static, and so if your code uses more than just the letters; A, T, G, and C, then it doesn't work. You can increase the size of this array so that it contains all the characters you need, but there are currently more than 128,000 characters defined in the Unicode spec. And so including that all is a waste of time, and makes their function \$O(ck^2)\$, where \$c\$ is the amount of characters in your list, which isn't the best.

And so instead I'd create a simple Trie, that contains a function levenshtein. To make this code efficient, we need to know the levenshtein limit - This is how many characters the node can be different from the input. It should also take the key that we want to match against. And finally an optional argument that deletes output that isn't the same as the levenshtein limit. To note, this only outputs keys that are the same length as the input key.

To do this, you don't want to use recursion, as that has a limit of roughly a 1000 character wide Trie. You want to also use generators and iterators to reduce wasted time, as we only need to check if there is a second value, and all others are thrown away. Fortunately using generators and iterators are as simple as using lists in Python.

To do the above we want to loop through each possible node. Then we want to loop though their children, if the child's key is not the same as the current inputs letter then we add one to it's levenshtein limit. If the levenshtein limit exceeds the limit we passed the function, then we don't add that node back to the possible nodes.

This has a much worse algorithmic performance than @RyanMills' at \$O(\frac{c^{k+1}}{c-1})\$, where \$c\$ is the amount of unique characters, and \$k\$ is the size of the input. However when used with a limit of 1, as we are above, it's \$O(ck)\$.

And so I'd use:

class Trie:
    class Node(dict):
        _DEFAULT = object()
        def __init__(self, children=None, data=_DEFAULT):
            if children is not None:
                for key, value in children.items():
                    children[key] = type(self)(value)
                self.update(children)
            self.data = data

        def __repr__(self):
            if self.data == self._DEFAULT:
                return f"{super().__repr__()}"
            return f"Node({super().__repr__()}, {self.data!r})"

        def __missing__(self, key):
            ret = self[key] = type(self)()
            return ret

    def __init__(self, children=None):
        self.root = self.Node(children)

    def __repr__(self):
        return f"Trie({self.root})"

    def _get_node(self, key):
        node = self.root
        for key_item in key:
            node = node[key_item]
        return node

    def __setitem__(self, key, value):
        self._get_node(key).data = value

    def __getitem__(self, key):
        return self._get_node(key).data

    def _levenshtein(self, key_item, possible):
        for (p, amount) in possible:
            for k, v in p.items():
                amt = amount + (k != key_item)
                if amt > limit:
                    continue
                yield v, amt

    def levenshtein(self, key, limit, delete=True):
        _levenshtein = self._levenshtein
        possible = [(self.root, 0)]
        for key_item in key:
            possible = _levenshtein(key_item, possible)
        if delete:
            possible = ((n, a) for n, a in possible if a == limit)
        yield from (n.data for n, _ in possible)