FMIndexで高速に塩基配列を検索してみた

今回は以下のリポジトリにあるCライブラリを使用してみた記録です。

https://github.com/TravisWheelerLab/AvxWindowFmIndex

ライブラリのビルド

まずは以下の要領でライブラリのビルド〜インストールを行いました。パス等は各環境で適切な値に変えます。

(MacOSの場合)

$ git clone https://github.com/TravisWheelerLab/AvxWindowFmIndex.git
$ cd AvxWindowFmIndex/
$ git submodule init
$ git submodule update
$ brew install gcc
$ # /path/to/gcc, /prefix/path は適宜変える
$ cmake -DCMAKE_C_COMPILER=/path/to/gcc -DCMAKE_INSTALL_PREFIX=/prefix/path .
$ make
$ make install

プログラムの作成

READMEを参考にFMIndexによる配列検索を行うプログラムをCで書きました。以下、コード全体です。main.cとして保存しています。

fastas/dnaSequence.fasta というファイルから配列を読み込み、コマンドライン引数に与えられた配列を検索し、ヒットしたポジションをプリントします。

#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include <AwFmIndex.h>

int main(int argc, char **argv) {
    char *indexFileSrc = "indexFiles/index.awfmi";
    char *fastaInputFileSrc = "fastas/dnaSequence.fasta";
    struct AwFmIndex *index;
    struct AwFmIndexConfiguration config = {
        .suffixArrayCompressionRatio = 8,
        .kmerLengthInSeedTable       = 12,
        .alphabetType                = AwFmAlphabetNucleotide,
        .keepSuffixArrayInMemory     = false,
        .storeOriginalSequence       = true
    };
    
    //enum AwFmReturnCode returnCode = awFmCreateIndexFromFasta(&index, &config, sequence, sequenceLength, indexFileSrc, true); origin
    enum AwFmReturnCode returnCode = awFmCreateIndexFromFasta(
        &index, &config, fastaInputFileSrc, indexFileSrc);
    if (awFmReturnCodeIsFailure(returnCode)) {
        printf("create index failed with return code %u\n", returnCode);
        exit(1);
    }
    uint64_t numQueries = argc-1; //first argument is the name of the executable, so skip it.
    struct AwFmKmerSearchList *searchList = awFmCreateKmerSearchList(numQueries);
    if (searchList == NULL) {
        printf("could not allocate memory for search list\n");
        exit(2);
    }
    
    searchList->count = numQueries;
    //initialize the queries in the search list
    for (size_t i = 0; i < numQueries; i++) {
        searchList->kmerSearchData[i].kmerLength = strlen(argv[i+1]);
        searchList->kmerSearchData[i].kmerString = argv[i+1];
    }
    
    //search for the queries
    int numThreads = 3;
    returnCode = awFmParallelSearchLocate(index, searchList, numThreads);
    if (awFmReturnCodeIsFailure(returnCode)) {
        printf("parallel search failed with return code %u\n", returnCode);
        exit(3);
    }
    
    //print all the hits.
    for (size_t kmerIndex = 0; kmerIndex < searchList->count; kmerIndex++) {
        const uint32_t numHitPositions = searchList->kmerSearchData[kmerIndex].count;
        uint64_t *positionList = searchList->kmerSearchData[kmerIndex].positionList;
      
        for (size_t i = 0; i < numHitPositions; i++) {
            printf("kmer at index %zu found at database position %zu.\n", kmerIndex, positionList[i]);
        }
    }
    
    //code cleanup
    awFmDeallocKmerSearchList(searchList);
    awFmDeallocIndex(index);
    
    exit(0);
}

結果の比較用にRustによる自前のFMIndex配列検索プログラムも書きました。以下にコード全体を示します。

#!/usr/bin/env rust-script
//!
//! ```cargo
//! [dependencies]
//! bio = "*"
//! 
//! ```
use std::path::Path;
use bio::io::fasta::{self, FastaRead};
use bio::alphabets;
use bio::data_structures::bwt::{bwt, less, Occ};
use bio::data_structures::fmindex::{BackwardSearchResult, FMIndex, FMIndexable};
use bio::data_structures::suffix_array::{suffix_array};
use std::collections::HashSet;

fn exact_search<P: AsRef<Path> + std::fmt::Debug>(ref_: P, que: P) {
    
    // read reference sequences and concatenate them
    let mut reader = fasta::Reader::from_file(ref_)
        .expect("Failed to open the reference fasta file.");
    let mut record = fasta::Record::new();

    let mut ref_seq = Vec::new();
    let mut ref_len = Vec::new();
    let mut ref_ids = HashSet::new();
    reader.read(&mut record).expect("Failed to parse the reference fasta record.");
    while !record.is_empty() {
        if record.check().is_err() {
            panic!("Rubbish record in the reference fasta");
        }
        ref_seq.extend_from_slice(record.seq());
        ref_seq.push(b'$');
        if ref_ids.contains(record.id()) {
            panic!("Record IDs are duplicated");
        }
        ref_ids.insert(record.id().to_string());
        ref_len.push((record.id().to_string(), record.seq().len()));
        reader.read(&mut record).expect("Failed to parse the reference fasta record.");
    }
    ref_seq.make_ascii_uppercase();
    //ref_seq.push(b'$');

    //  fmindex construction
    let alphabet = alphabets::Alphabet::new(&ref_seq);
    let sa = suffix_array(&ref_seq);
    let bwt = bwt(&ref_seq, &sa);
    let less = less(&bwt, &alphabet);
    let occ = Occ::new(&bwt, 128, &alphabet);
    let fm = FMIndex::new(&bwt, &less, &occ);

    // query fasta read
    let mut reader = fasta::Reader::from_file(que)
        .expect("Failed to open the query fasta file.");
    let mut record = fasta::Record::new();
    reader.read(&mut record).expect("Failed to parse the query fasta record.");
    while !record.is_empty() {
        if record.check().is_err() {
            panic!("Rubbish record in the query fasta");
        }
        let mut seq = record.seq().to_vec();
        seq.make_ascii_uppercase();
        if alphabet.is_word(&seq) {
            let interval = fm.backward_search(seq.iter());
            let mut positions = match interval {
                BackwardSearchResult::Complete(saint) => saint.occ(&sa),
                _ => Vec::new()
            };
            positions.sort();
            for i in positions {
                println!("{} found ad database position {}", record.id(), i);
            }
        }
        reader.read(&mut record)
            .expect("Failed to parse the query fasta record.");
    }

}

fn main() {
    exact_search("fastas/dnaSequence.fasta", "fastas/query.fasta");
}

実行結果

結局コンパイルオプションの適切な与え方はわかりませんでしたがコンパイルできたので適当なFASTAに対して 適当なクエリを検索してみました。以下、実行結果になります。ポジションは正確に出力できているかと思います。

ひとまず動いたので今回はここで終了です。SIMDあたりはもう少し勉強してみたいと思います。

$ # 入力ファイルのサイズ
$ seqkit stat fastas/dnaSequence.fasta
file                      format  type  num_seqs  sum_len  min_len  avg_len  max_len
fastas/dnaSequence.fasta  FASTA   DNA          5   50,000   10,000   10,000   10,000
$ # コンパイル(SIMDの有効化がよくわからなかった)
$ gcc main.c -I prefix/path/include -I AvxWindowFmIndex/lib/FastaVector/src -L prefix/path/lib/ -lawfmindex  -o main
$ export DYLD_LIBRARY_PATH=prefix/path/lib
$ # 実行結果
$ time ./main GCTCGCTGTTCTCACACATGAAACCTGACTGAATACGAT TATTGGAACATTGCACAACCTATAGACTGCGCGAAGGGCA GTGTCAACTTTAAGTTAGCCCATC
kmer at index 0 found at database position 18829.
kmer at index 1 found at database position 28108.
kmer at index 2 found at database position 49980.
./main GCTCGCTGTTCTCACACATGAAACCTGACTGAATACGAT  GTGTCAACTTTAAGTTAGCCCATC  2.90s user 0.19s system 99% cpu 3.110 total
$ # 実行結果(Rust版比較用プログラム。Indexファイル出力がないので早いのは当たり前。検索結果はちゃんと一致している。)
$ time ./main.rs
query1 found ad database position 18829
query2 found ad database position 28108
query3 found ad database position 49980
./main.rs  0.08s user 0.03s system 46% cpu 0.230 total