Mapping DNA nucleotides into two-dimensional coordinates

Question

I am the author of this package that turns DNA sequences into two dimensional visualizations. DNA, for the unaware, consists of four letters (A, T, G, and C). To convert the sequence, I am using this code:

import numpy as np

def transform(sequence):
    running_value = 0
    x, y = np.linspace(0, len(sequence), 2 * len(sequence) + 1), [0]
    for character in sequence:
        if character == "A":
            y.extend([running_value + 0.5, running_value])
        elif character == "C":
            y.extend([running_value - 0.5, running_value])
        elif character == "T":
            y.extend([running_value - 0.5, running_value - 1])
            running_value -= 1
        elif character == "G":
            y.extend([running_value + 0.5, running_value + 1])
            running_value += 1
        else:
            y.extend([running_value] * 2)
    return list(x), y

Note that seq may be a very long sequence, and this transformation process may be done (hundreds of) thousands of time.

What are some things I could do to improve performance?

Edit:

• The else clause is the because some DNA sequences have ambiguity (i.e. non ATGC characters) which must be accounted for as horizontal lines.

• sequence may safely assumed to be a string. For example, "GATTACA" is a valid DNA sequence that should return ([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0], [0, 0.5, 1, 1.5, 1, 0.5, 0, -0.5, -1, -0.5, -1, -1.5, -1, -0.5, -1]).

• By "long", I mean anywhere from thousands to billions. That being said, the most common use case likely will involve fewer than ten thousand characters, although this process may be done to thousands of such sequences. For context, I am currently running this function over 400,000 sequences with a median length of 800 characters.

Answer 1

First of all I'd like to say that your code is fast. From studying it, the major bottlenecks that I've found come from the input being a string, and from converting to numpy arrays and back.

Since what you're doing is a kind of cumulative sum with some alterations, I'd go with numpy.cumsum. It's a bit tricky to get everything in place, but the calculations of the actual sequence is actually not the major bottleneck. From what I've seen, the solution I came up with is about twice as fast, and 75% of the time is used to convert the string input to something useful, which in this case is a numpy array representing the character index in the string "ATCG" for each character in the input string.

def transform_fast(seq, d):
    l = len(seq)
    x = np.linspace(0, l, 2 * l + 1, dtype = np.float32)
    y = np.zeros(2* l + 1, dtype = np.float32)

    atcg = np.array([d[x] for x in seq], dtype = np.int8)

    a = (atcg == 1).astype(np.int8)
    t = (atcg == 2).astype(np.int8)
    c = (atcg == 4).astype(np.int8)
    g = (atcg == 8).astype(np.int8)
    ac = a - c
    tg = -t + g

    cum_sum = np.concatenate((
        np.array([0]),
        np.cumsum(tg[:-1])
    ))

    y[1::2] = cum_sum + 0.5*(ac + tg)
    y[2::2] = cum_sum + tg

    return x, y

In the function above, the argument d is a dictionary which looks like:

d = {
    "A": 1, "T": 2, "C": 4, "G": 8,
    "U": 0, "W": 0, "S": 0, "M": 0,
    "K": 0, "R": 0, "Y": 0, "B": 0,
    "D": 0, "H": 0, "V": 0, "N": 0, "Z": 0
}

The only difference compared to your solution is that this returns numpy arrays rather than lists. If this doesn't suit your use case, you can convert them to lists. However, this will slow down the overall performance significantly.

The code above also does not handle the non-ATCG characters in the input string. However, I think that you'll be able to follow the code and implement that in a similar fashion.

Benchmark with random input string of length 10^3
transform(seq): 0.45ms
transform_fast(seq): 0.24ms
Speedup factor: 1.82

Benchmark with random input string of length 10^6
transform(seq): 388.42ms
transform_fast(seq): 133.40ms
Speedup factor: 2.91

EDIT: I found that the np.fromiter(map(...)) that was previously used was unnecessary, and replaced it. This increased performance, benchmarks have been updated. I also found that -t + g was used thrice, and saved the result to a variable. I also switched x and y to be np.float32 arrays, as that improved performance slightly.

It should be noted that the approach above uses more memory than your original approach, which could be a problem for large enough sequences. On my machine (16GB RAM) I couldn't get the result for a sequence of length \$10^8\$, as I ran out of memory.

For the sake of curiosity, here are the benchmarks for when the input string is parsed into a numpy array before transform_fast, which could be used if the input sequences are reused often (these benchmarks haven't been updated, but should remain roughly the same):

Benchmark with random input string of length 10^3
transform(seq): 0.42ms
transform_fast(seq): 0.13ms
Speedup factor: 3.15

Benchmark with random input string of length 10^6
transform(seq): 380.23ms
transform_fast(seq): 37.23ms
Speedup factor: 10.21