performance - 从核苷酸序列生成 kmer 计数向量的最快方法 (Julia)

Question

给定一个核苷酸序列，我正在编写一些 Julia 代码来生成(屏蔽的)kmer 计数的稀疏向量，我希望它尽可能快地运行。

这是我当前的实现，

using Distributions
using SparseArrays

function kmer_profile(seq, k, mask)
    basis = [4^i for i in (k - 1):-1:0]

    d = Dict('A'=>0, 'C'=>1, 'G'=>2, 'T'=>3)

    kmer_dict = Dict{Int, Int32}(4^k=>0)

    for n in 1:(length(seq) - length(mask) + 1)
        kmer_hash = 1
        j = 1
        for i in 1:length(mask)
            if mask[i]
                kmer_hash += d[seq[n+i-1]] * basis[j]
                j += 1
            end
        end
        haskey(kmer_dict, kmer_hash) ? kmer_dict[kmer_hash] += 1 : kmer_dict[kmer_hash] = 1
    end

    return sparsevec(kmer_dict)
end

seq = join(sample(['A','C','G','T'], 1000000))

mask_str = "111111011111001111111111111110"

mask = BitArray([parse(Bool, string(m)) for m in split(mask_str, "")])
k = sum(mask)

@time kmer_profile(seq, k, mask)

这段代码在我的 M1 MacBook Pro 上运行大约需要 0.3 秒，有什么方法可以让它运行得更快吗？

函数kmer_profile使用大小为length(mask)的滑动窗口来计算每个掩码kmer出现在核苷酸序列中的次数。掩码是二进制序列，掩码kmer是在掩码为零的位置处掉落核苷酸的kmer。例如。 kmer ACGT 和掩码 1001 将生成掩码 kmer AT。

为了生成 kmer 哈希，该函数将每个 kmer 视为基数 4 的数字，然后将其转换为(基数 10)64 位整数，用于索引到 kmer 向量。

k 的大小等于掩码字符串中 1 的数量，并且隐式限制为 31，以便 kmer 哈希可以适合 64 位整数类型。

Answer 1

有几种可能的优化可以使此代码更快。

首先，可以将 Dict 转换为数组，因为基于数组的索引比基于字典的索引更快，而且这里可以实现这一点，因为键是 ASCII 字符。 p>

此外，通过预先计算代码并将结果放入临时数组，可以提取一次序列代码，而不是length(mask)次。

此外，基于mask的条件和循环携带的依赖使事情变得很慢。事实上，处理器无法(轻松)预测该情况，导致其停顿几个周期。循环携带的依赖性使事情变得更糟，因为处理器在此停顿期间几乎无法执行其他指令。这个问题可以通过基于mask和basis预先计算因子来解决。结果是更快的无分支循环。

完成上述优化后，最大的瓶颈是sparsevec。事实上，它也花费了最初实现的近一半时间!优化这一步很困难，但并非不可能。由于 Julia 实现中的随机访问，速度很慢。首先可以通过对键值对进行排序来加快速度。由于更适合缓存的执行，它的速度更快，并且还可以帮助处理器的预测单元。这是一个复杂的话题。有关其工作原理的更多详细信息，请阅读 Why is processing a sorted array faster than processing an unsorted array? .

这是最终优化的代码:

function kmer_profile_opt(seq, k, mask)
    basis = [4^i for i in (k - 1):-1:0]

    d = zeros(Int8, 128)
    d[Int64('A')] = 0
    d[Int64('C')] = 1
    d[Int64('G')] = 2
    d[Int64('T')] = 3

    seq_codes = [d[Int8(e)] for e in seq]

    j = 1
    premult = zeros(Int64, length(mask))
    for i in 1:length(mask)
        if mask[i]
            premult[i] = basis[j]
            j += 1
        end
    end

    kmer_dict = Dict{Int, Int32}(4^k=>0)

    for n in 1:(length(seq) - length(mask) + 1)
        kmer_hash = 1
        j = 1
        for i in 1:length(mask)
            kmer_hash += seq_codes[n+i-1] * premult[i]
        end
        haskey(kmer_dict, kmer_hash) ? kmer_dict[kmer_hash] += 1 : kmer_dict[kmer_hash] = 1
    end

    sorted_kmer_pairs = sort(collect(kmer_dict))
    sorted_kmer_keys = [e[1] for e in sorted_kmer_pairs]
    sorted_kmer_values = [e[2] for e in sorted_kmer_pairs]
    return sparsevec(sorted_kmer_keys, sorted_kmer_values)
end

此代码比我机器上的初始实现快两倍。很大一部分时间仍然花在排序算法上。

代码还可以进一步优化。一种方法是使用并行排序算法。另一种方法是用移位替换 premult[i] 乘法，假设 premult[i] 已修改为包含指数，则移位速度更快。我预计该代码比原始代码快大约 4 倍。主要瓶颈应该是大字典的创建。进一步提高其性能非常困难(尽管仍然有可能)。