BWT 算法和序列比对的基本实现

昨天晚上和今天抽空实现了Burrows Wheleer Tansform,并且尝试利用BWT，将短序列比对到长序列中。BWT的核心我觉得是要理解两个原则：

1. F序列的每个元素是下标对应的L元素的后一位。

2. 排序后，F中第一个A和L中第一个A是同一个A。（排序不改变相对位置），公共前缀不改变排序位置。 

mapping 过程实现的非常基础，只能全序列不对，不能有gap。

#!/usr/bin/env python3
'''
burrows wheleer transform
核心是 排序后的F,L数组， L数组的某个元素是F数组对应的元素的前一位。
其次是，F数组的某个元素A在所有A中的相对位置不变。也就是F中的第一个A对应着L中的第一个A.
'''

import os
import sys
from collections import Counter


def BWT(s, end):
    '''
    输出1：L string
    输出2：L string 每个字符的原来的位置
    '''
    # 记录原序列中字符出现位置
    C = {i: [] for i in set(s)}  # 初始化L列的元素相对位置
    for i in range(len(s)):  # 分别给每个元素标记出现的位置
        C[s[i]].append(i)

    s = s + end  # 添加最后的符号
    M = []  # 初始化输出
    count = len(s)  # 计算字符串长度

    # 设置偏移矩阵
    while True:
        M.append(s)
        s = s[len(s)-1]+s[0:len(s)-1]
        count -= 1
        if count <= 0:
            break
    N = [[v, k] for k, v in enumerate(M[::-1])] # 标记位置信息
    N_sort = sorted(N, key=lambda x: x[0])  #字典排序
    C = {i: [] for i in set(s)}  # 初始化L列的元素相对位置
    L = []
    for k, v in N_sort: # 以A为例，Lstring中 从头开始每一个A在原来字符串的位置 
        L.append(k[-1])
        C[k[-1]].append(v)

    return(["".join(L), C])


def reverseBWT(s, end):
    '''
    '''

    B = Counter(s)  # 压缩的F列
    C = {i: [] for i in set(B)}  # 初始化L列的元素相对位置
    for i in range(len(s)):  # 分别给每个元素标记出现的位置
        C[s[i]].append(i)

    out = ""  # 初始化输出
    arrow = 1
    while True:
        Lbase = s[arrow-1]  # 取出对应的字符
        out = out + Lbase  # 添加到输出

        if Lbase == end:  # 如果遇到终止符，则退出循环，并输出
            break
        # 从数组C中查找对应字符出现的位置，并且枚举，这样可以得到字符所在位置对应的相同字符的偏移量。
        for i, j in enumerate(C[Lbase]):
            if j == arrow-1:
                pianyi = i+1  # 这里应该是对应的L中的序号
                break

        arrow = CheckF(B, Lbase, pianyi)
    return(out[::-1])


def simple_mapping(ref, seed):

    s, rank = BWT(ref, "#")

    B = Counter(s)  # 压缩的F列
    C = {i: [] for i in set(s)}  # 初始化L列的元素相对位置
    for i in range(len(s)):  # 分别给每个元素标记出现的位置
        C[s[i]].append(i)

    seed = seed[::-1]
    print(ref)

    for mark in range(len(C[seed[0]])):
        arrow = C[seed[0]][mark]  # mark:0,1,2 编号，0开始
        # arrow: L 中的 12,34,65... 实际位置

        out = ""
        count = 0
        while True:
            Lbase = s[arrow]
            if Lbase != seed[count] or Lbase == '#':
                break
            out = out + Lbase
            for i, v in enumerate(C[Lbase]):
                if v == arrow:
                    pianyi = i + 1
                    break
            arrow = CheckF_0base(B, Lbase, pianyi)
            count += 1
            if count > len(seed)-1:
                left = rank[seed[0]][mark]
                out = "*"*(left-len(out)+1) + out[::-1] + "*"*(len(s)-left-2)
                print(out)
                break


def CheckF(mtx, base, order):
    out_order = 0
    for k in sorted(mtx.keys()):  # 从F数组中根据字符，以及偏移量，计算出下一个坐标 P = 字符前所有的字符的数目+偏移量
        if base != k:
            out_order += mtx[k]
        else:
            out_order += order
            break
    return(out_order)


def CheckF_0base(mtx, base, order):
    out_order = 0
    for k in sorted(mtx.keys()):  # 从F数组中根据字符，以及偏移量，计算出下一个坐标 P = 字符前所有的字符的数目+偏移量
        if base != k:
            out_order += mtx[k]
        else:
            out_order += order
            break
    return(out_order-1)


if __name__ == '__main__':

    x = "actgagctttagcgtagctttaggagagcttcctagctacgtatcgagcgggcatctatc"
    seed = "ga"
    print("origin string is :" + x)
    print("L string      is :" + BWT(x, '#')[0])
    print("seed string   is :" + seed)
    print("mapped seq    is :")
    simple_mapping(x, seed)

 

 

结果：