296 lines
9.7 KiB
C
296 lines
9.7 KiB
C
/* The MIT License
|
||
|
||
Copyright (c) 2025- Zhonghai Zhang (ICT)
|
||
|
||
Permission is hereby granted, free of charge, to any person obtaining
|
||
a copy of this software and associated documentation files (the
|
||
"Software"), to deal in the Software without restriction, including
|
||
without limitation the rights to use, copy, modify, merge, publish,
|
||
distribute, sublicense, and/or sell copies of the Software, and to
|
||
permit persons to whom the Software is furnished to do so, subject to
|
||
the following conditions:
|
||
|
||
The above copyright notice and this permission notice shall be
|
||
included in all copies or substantial portions of the Software.
|
||
|
||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
|
||
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
|
||
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
|
||
BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
|
||
ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
||
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||
SOFTWARE.
|
||
*/
|
||
|
||
#include <stdio.h>
|
||
#include <stdlib.h>
|
||
#include <string.h>
|
||
#include <time.h>
|
||
#include <unistd.h>
|
||
#include <zlib.h>
|
||
|
||
#include "bntseq.h"
|
||
#include "bwa.h"
|
||
#include "bwt.h"
|
||
#include "rle.h"
|
||
#include "rope.h"
|
||
#include "utils.h"
|
||
#include "kvec.h"
|
||
#include "hyb_idx.h"
|
||
|
||
|
||
#ifdef _DIVBWT
|
||
#include "divsufsort.h"
|
||
#endif
|
||
|
||
#ifdef USE_MALLOC_WRAPPERS
|
||
#include "malloc_wrap.h"
|
||
#endif
|
||
|
||
int bwa_bwt2fullbytesa(int argc, char* argv[]) // the "bwt2bytesa" command
|
||
{
|
||
bwt_t* bwt;
|
||
if (optind + 2 > argc) {
|
||
fprintf(stderr, "Usage: %s bwt2bytesa <in.bwt> <out.bytesa>\n", PROG_NAME);
|
||
return 1;
|
||
}
|
||
bwt = bwt_restore_bwt(argv[optind]);
|
||
bwt_cal_byte_sa(bwt);
|
||
_store_data_to_file(argv[optind + 1], bwt->byte_sa, sizeof(bwtint_t), SA_BYTES(bwt->n_sa) >> 3);
|
||
return 0;
|
||
}
|
||
|
||
// extract suffix array from non-sampled suffix array
|
||
int bwa_extract_sa(int argc, char* argv[]) {
|
||
bwt_t* bwt;
|
||
int c, sa_intv = 32;
|
||
while ((c = getopt(argc, argv, "i:")) >= 0) {
|
||
switch (c) {
|
||
case 'i':
|
||
sa_intv = atoi(optarg);
|
||
break;
|
||
default:
|
||
return 1;
|
||
}
|
||
}
|
||
if (optind + 3 > argc) {
|
||
fprintf(stderr, "Usage: bwa extractsa [-i %d] <in.bwt> <in.full_bytesa> <out.sa>\n", sa_intv);
|
||
return 1;
|
||
}
|
||
bwt = bwt_restore_bwt(argv[optind]);
|
||
|
||
fprintf(stderr, "bwt seq_len %ld, intv: %d\n", bwt->seq_len, sa_intv);
|
||
//exit(0);
|
||
|
||
_load_file_to_data(argv[optind + 1], bwt->byte_sa);
|
||
|
||
if (bwt->sa) free(bwt->sa);
|
||
bwt->sa_intv = sa_intv;
|
||
bwt->n_sa = (bwt->seq_len + sa_intv) / sa_intv;
|
||
bwt->sa = (bwtint_t*)calloc(bwt->n_sa, sizeof(bwtint_t));
|
||
uint64_t i;
|
||
for (i = 0; i < bwt->seq_len; ++i) {
|
||
if (i % sa_intv == 0)
|
||
bwt->sa[i / sa_intv] = bwt_get_sa(bwt->byte_sa, i);
|
||
}
|
||
if (i % sa_intv == 0)
|
||
bwt->sa[i / sa_intv] = bwt_get_sa(bwt->byte_sa, i);
|
||
bwt->sa[0] = (bwtint_t)-1; // before this line, bwt->sa[0] = bwt->seq_len
|
||
// bwt_cal_sa(bwt, sa_intv);
|
||
bwt_dump_sa(argv[optind + 2], bwt);
|
||
// bwt_destroy(bwt);
|
||
return 0;
|
||
}
|
||
|
||
// extract byte suffix array from non-sampled suffix array
|
||
int bwa_extract_byte_sa(int argc, char* argv[]) {
|
||
bwt_t* bwt;
|
||
int c, sa_intv = 32;
|
||
while ((c = getopt(argc, argv, "i:")) >= 0) {
|
||
switch (c) {
|
||
case 'i':
|
||
sa_intv = atoi(optarg);
|
||
break;
|
||
default:
|
||
return 1;
|
||
}
|
||
}
|
||
if (optind + 3 > argc) {
|
||
fprintf(stderr, "Usage: bwa extractbytesa [-i %d] <in.bwt> <in.full_bytesa> <out.bytesa>\n\n", sa_intv);
|
||
return 1;
|
||
}
|
||
bwt = bwt_restore_bwt(argv[optind]);
|
||
|
||
// fprintf(stderr, "bwt seq_len %ld, intv: %d\n", bwt->seq_len, sa_intv);
|
||
// exit(0);
|
||
|
||
uint8_t *full_byte_sa = NULL;
|
||
_load_file_to_data(argv[optind + 1], full_byte_sa);
|
||
|
||
bwt->sa_intv = sa_intv;
|
||
bwt->n_sa = (bwt->seq_len + sa_intv) / sa_intv;
|
||
bwt->byte_sa = (uint8_t*)calloc(SA_BYTES(bwt->n_sa), 1);
|
||
uint64_t i;
|
||
for (i = 0; i < bwt->seq_len; ++i) {
|
||
if (i % sa_intv == 0)
|
||
bwt_set_sa(bwt->byte_sa, i / sa_intv, bwt_get_sa(full_byte_sa, i));
|
||
}
|
||
if (i % sa_intv == 0)
|
||
bwt_set_sa(bwt->byte_sa, i / sa_intv, bwt_get_sa(full_byte_sa, i));
|
||
bwt_set_sa(bwt->byte_sa, 0, 8589934591); // before this line, bwt->sa[0] = bwt->seq_len
|
||
bwt_dump_byte_sa(argv[optind + 2], bwt);
|
||
// bwt_destroy(bwt);
|
||
return 0;
|
||
}
|
||
|
||
// 生成所有KMER_LEN长度的序列，字节表示(0A,1C,2G,3T)
|
||
void gen_kmer_bits(uint8_t** seq_arr, uint64_t* kmer_arr_size, int kmer_len) {
|
||
uint64_t i;
|
||
uint64_t seq_up_val = (1L << (kmer_len << 1));
|
||
*kmer_arr_size = seq_up_val;
|
||
*seq_arr = (uint8_t*)realloc(*seq_arr, seq_up_val * (uint64_t)kmer_len);
|
||
for (i = 0; i < seq_up_val; ++i) {
|
||
const uint64_t base_idx = i * kmer_len;
|
||
for (int j = kmer_len - 1; j >= 0; --j) {
|
||
(*seq_arr)[base_idx + kmer_len - 1 - j] = (i >> (j << 1)) & 3;
|
||
}
|
||
}
|
||
}
|
||
|
||
void bwt_create_kmer_idx(bwt_t* bwt) {
|
||
uint64_t kmer_arr_size = 0;
|
||
uint8_t* seq_arr = 0;
|
||
bwtintv_t ik, ok[4];
|
||
uint64_t j;
|
||
int i, c1;
|
||
int qlen = 10;
|
||
KmerHash* kh = &bwt->kmer_hash;
|
||
|
||
kh->ke10 = (KmerEntryArr*)malloc((1 << (10 << 1)) * sizeof(KmerEntryArr));
|
||
kh->ke11 = (KmerEntry*)malloc((1 << (11 << 1)) * sizeof(KmerEntry));
|
||
kh->ke12 = (KmerEntry*)malloc((1 << (12 << 1)) * sizeof(KmerEntry));
|
||
kh->ke13 = (KmerEntry*)malloc((1 << (13 << 1)) * sizeof(KmerEntry));
|
||
kh->ke14 = (KmerEntry*)malloc((1 << (14 << 1)) * sizeof(KmerEntry));
|
||
|
||
// 长度为10的kmer
|
||
gen_kmer_bits(&seq_arr, &kmer_arr_size, 10);
|
||
for (j = 0; j < kmer_arr_size; ++j) {
|
||
uint8_t* q = &seq_arr[j * 10];
|
||
uint8_t* mem_addr = kh->ke10[j].intv_arr;
|
||
bwt_set_intv(bwt, q[0], ik);
|
||
kmer_setval_at(mem_addr, ik, 0);
|
||
// 每次扩展一个碱基
|
||
for (i = 1; i < qlen; ++i) {
|
||
c1 = 3 - q[i];
|
||
bwt_extend(bwt, &ik, ok, 0);
|
||
ik = ok[c1];
|
||
kmer_setval_at(mem_addr, ik, i);
|
||
}
|
||
}
|
||
fprintf(stderr, "Kmer idx for length 10 done.\n");
|
||
#define GEN_KMER_FOR_LEN(len) \
|
||
do { \
|
||
gen_kmer_bits(&seq_arr, &kmer_arr_size, (len)); \
|
||
for (j = 0; j < kmer_arr_size; ++j) { \
|
||
bwtintv_t p = bwt_forward_search(bwt, &seq_arr[j * (len)], (len)); \
|
||
kmer_setval_at(bwt->kmer_hash.ke##len[j].intv_arr, p, 0); \
|
||
} \
|
||
fprintf(stderr, "Kmer idx for length %d done.\n", len); \
|
||
} while (0)
|
||
|
||
GEN_KMER_FOR_LEN(11); // 长度11的kmer
|
||
GEN_KMER_FOR_LEN(12);
|
||
GEN_KMER_FOR_LEN(13);
|
||
GEN_KMER_FOR_LEN(14);
|
||
free(seq_arr);
|
||
}
|
||
|
||
// 从bwt中创建kmer index，并保存到文件
|
||
int bwa_bwt2kmer(int argc, char* argv[]) {
|
||
bwt_t* bwt;
|
||
if (optind + 2 > argc) {
|
||
fprintf(stderr, "Usage: bwa extractbytesa <in.bwt> <out.kmer>\n\n");
|
||
return 1;
|
||
}
|
||
bwt = bwt_restore_bwt(argv[optind]);
|
||
bwt_create_kmer_idx(bwt);
|
||
bwt_dump_kmer_idx(argv[optind + 1], &bwt->kmer_hash);
|
||
return 0;
|
||
}
|
||
|
||
// 创建hybrid index，并保存到文件
|
||
int bwa_bwt2hyb(int argc, char* argv[]) {
|
||
int hyb_idx_build_and_dump(int num_threads, bwt_t* bwt, const char* idx_prefix);
|
||
bwt_t* bwt;
|
||
int num_threads = 1;
|
||
char c;
|
||
int error = 0;
|
||
while ((c = getopt(argc, argv, "t:")) >= 0) {
|
||
switch (c) {
|
||
case 't':
|
||
num_threads = atoi(optarg);
|
||
assert(num_threads > 0 && num_threads < 512);
|
||
break;
|
||
default:
|
||
error = 1;
|
||
break;
|
||
}
|
||
}
|
||
if (optind + 1 > argc || error) {
|
||
fprintf(stderr, "Usage: bwa bwt2hyb [Options] <bwt-prefix>\n\n");
|
||
fprintf(stderr, "Options: -t INT number of threads for hybrid index building [%d]\n", num_threads);
|
||
fprintf(stderr, "\n");
|
||
return 1;
|
||
}
|
||
char fn[MAX_PATH];
|
||
snprintf(fn, MAX_PATH, "%s.bwt", argv[optind]);
|
||
bwt = bwt_restore_bwt(fn);
|
||
snprintf(fn, MAX_PATH, "%s.kmer", argv[optind]);
|
||
bwt->kmer_hash = bwt_restore_kmer_idx(fn);
|
||
snprintf(fn, MAX_PATH, "%s.hyb.bytesa", argv[optind]);
|
||
_load_file_to_data(fn, bwt->byte_sa);
|
||
hyb_idx_build_and_dump(num_threads, bwt, argv[optind]);
|
||
return 0;
|
||
}
|
||
|
||
// 创建正向的kmer
|
||
uint64_t build_forward_kmer(const uint8_t* q, int qlen, int kmer_len, int* base_consumed) {
|
||
uint64_t qbit = 0, i;
|
||
qlen = qlen < kmer_len ? qlen : kmer_len;
|
||
for (i = 0; i < qlen; ++i) {
|
||
if (q[i] > 3)
|
||
break; // 要考虑碱基是N
|
||
qbit |= (uint64_t)q[i] << ((kmer_len - 1 - i) << 1);
|
||
}
|
||
*base_consumed = i;
|
||
return qbit;
|
||
}
|
||
|
||
int hyb_test(int argc, char* argv[]) {
|
||
bwt_t* bwt;
|
||
if (optind + 1 > argc) {
|
||
fprintf(stderr, "Usage: bwa hybtest <idx_prefix>\n\n");
|
||
return 1;
|
||
}
|
||
|
||
char fn[MAX_PATH];
|
||
snprintf(fn, MAX_PATH, "%s.bwt", argv[optind]);
|
||
|
||
bwt = bwt_restore_bwt(fn);
|
||
snprintf(fn, MAX_PATH, "%s.hyb.kmer", argv[optind]);
|
||
bwt->kmer_hash = bwt_restore_kmer_idx(fn);
|
||
|
||
uint8_t test_seq[14] = {0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1}; // ACGTACGTACGTAC
|
||
bwtintv_t ok1, ok2;
|
||
ok1 = bwt_forward_search(bwt, test_seq, 14);
|
||
|
||
int base_consumed = 0;
|
||
uint64_t qbit = build_forward_kmer(test_seq, 14, 14, &base_consumed);
|
||
|
||
bwt_kmer_get(&bwt->kmer_hash, &ok2, qbit, 13);
|
||
|
||
fprintf(stderr, "forward search: [%lu, %lu, %lu), kmer idx: [%lu, %lu, %lu)\n", ok1.x[0], ok1.x[1], ok1.x[2], ok2.x[0], ok2.x[1], ok2.x[2]);
|
||
|
||
return 0;
|
||
} |