bam_index.c

   1 #include <ctype.h>
   2 #include "bam.h"
   3 #include "khash.h"
   4 #include "ksort.h"
   5 #include "bam_endian.h"
   6
   7 /*!
   8   @header
   9
  10   Alignment indexing. Before indexing, BAM must be sorted based on the
  11   leftmost coordinate of alignments. In indexing, BAM uses two indices:
  12   a UCSC binning index and a simple linear index. The binning index is
  13   efficient for alignments spanning long distance, while the auxiliary
  14   linear index helps to reduce unnecessary seek calls especially for
  15   short alignments.
  16
  17   The UCSC binning scheme was suggested by Richard Durbin and Lincoln
  18   Stein and is explained by Kent et al. (2002). In this scheme, each bin
  19   represents a contiguous genomic region which can be fully contained in
  20   another bin; each alignment is associated with a bin which represents
  21   the smallest region containing the entire alignment. The binning
  22   scheme is essentially another representation of R-tree. A distinct bin
  23   uniquely corresponds to a distinct internal node in a R-tree. Bin A is
  24   a child of Bin B if region A is contained in B.
  25
  26   In BAM, each bin may span 2^29, 2^26, 2^23, 2^20, 2^17 or 2^14 bp. Bin
  27   0 spans a 512Mbp region, bins 1-8 span 64Mbp, 9-72 8Mbp, 73-584 1Mbp,
  28   585-4680 128Kbp and bins 4681-37449 span 16Kbp regions. If we want to
  29   find the alignments overlapped with a region [rbeg,rend), we need to
  30   calculate the list of bins that may be overlapped the region and test
  31   the alignments in the bins to confirm the overlaps. If the specified
  32   region is short, typically only a few alignments in six bins need to
  33   be retrieved. The overlapping alignments can be quickly fetched.
  34
  35  */
  36
  37 #define BAM_MIN_CHUNK_GAP 32768
  38 // 1<<14 is the size of minimum bin.
  39 #define BAM_LIDX_SHIFT    14
  40
  41 typedef struct {
  42         uint64_t u, v;
  43 } pair64_t;
  44
  45 #define pair64_lt(a,b) ((a).u < (b).u)
  46 KSORT_INIT(off, pair64_t, pair64_lt)
  47
  48 typedef struct {
  49         uint32_t m, n;
  50         pair64_t *list;
  51 } bam_binlist_t;
  52
  53 typedef struct {
  54         int32_t n, m;
  55         uint64_t *offset;
  56 } bam_lidx_t;
  57
  58 KHASH_MAP_INIT_INT(i, bam_binlist_t)
  59
  60 struct __bam_index_t {
  61         int32_t n;
  62         khash_t(i) **index;
  63         bam_lidx_t *index2;
  64 };
  65
  66 // requirement: len <= LEN_MASK
  67 static inline void insert_offset(khash_t(i) *h, int bin, uint64_t beg, uint64_t end)
  68 {
  69         khint_t k;
  70         bam_binlist_t *l;
  71         int ret;
  72         k = kh_put(i, h, bin, &ret);
  73         l = &kh_value(h, k);
  74         if (ret) { // not present
  75                 l->m = 1; l->n = 0;
  76                 l->list = (pair64_t*)calloc(l->m, 16);
  77         }
  78         if (l->n == l->m) {
  79                 l->m <<= 1;
  80                 l->list = (pair64_t*)realloc(l->list, l->m * 16);
  81         }
  82         l->list[l->n].u = beg; l->list[l->n++].v = end;
  83 }
  84
  85 static inline void insert_offset2(bam_lidx_t *index2, bam1_t *b, uint64_t offset)
  86 {
  87         int i, beg, end;
  88         beg = b->core.pos >> BAM_LIDX_SHIFT;
  89         end = (bam_calend(&b->core, bam1_cigar(b)) - 1) >> BAM_LIDX_SHIFT;
  90         if (index2->m < end + 1) {
  91                 int old_m = index2->m;
  92                 index2->m = end + 1;
  93                 kroundup32(index2->m);
  94                 index2->offset = (uint64_t*)realloc(index2->offset, index2->m * 8);
  95                 memset(index2->offset + old_m, 0, 8 * (index2->m - old_m));
  96         }
  97         for (i = beg + 1; i <= end; ++i)
  98                 if (index2->offset[i] == 0) index2->offset[i] = offset;
  99         index2->n = end + 1;
 100 }
 101
 102 static void merge_chunks(bam_index_t *idx)
 103 {
 104 #if defined(BAM_TRUE_OFFSET) || defined(BAM_VIRTUAL_OFFSET16)
 105         khash_t(i) *index;
 106         int i, l, m;
 107         khint_t k;
 108         for (i = 0; i < idx->n; ++i) {
 109                 index = idx->index[i];
 110                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 111                         bam_binlist_t *p;
 112                         if (!kh_exist(index, k)) continue;
 113                         p = &kh_value(index, k);
 114                         m = 0;
 115                         for (l = 1; l < p->n; ++l) {
 116 #ifdef BAM_TRUE_OFFSET
 117                                 if (p->list[m].v + BAM_MIN_CHUNK_GAP > p->list[l].u) p->list[m].v = p->list[l].v;
 118 #else
 119                                 if (p->list[m].v>>16 == p->list[l].u>>16) p->list[m].v = p->list[l].v;
 120 #endif
 121                                 else p->list[++m] = p->list[l];
 122                         } // ~for(l)
 123                         p->n = m + 1;
 124                 } // ~for(k)
 125         } // ~for(i)
 126 #endif // defined(BAM_TRUE_OFFSET) || defined(BAM_BGZF)
 127 }
 128
 129 bam_index_t *bam_index_core(bamFile fp)
 130 {
 131         bam1_t *b;
 132         bam_header_t *h;
 133         int i, ret;
 134         bam_index_t *idx;
 135         uint32_t last_bin, save_bin;
 136         int32_t last_coor, last_tid, save_tid;
 137         bam1_core_t *c;
 138         uint64_t save_off, last_off;
 139
 140         idx = (bam_index_t*)calloc(1, sizeof(bam_index_t));
 141         b = (bam1_t*)calloc(1, sizeof(bam1_t));
 142         h = bam_header_read(fp);
 143         c = &b->core;
 144
 145         idx->n = h->n_targets;
 146         bam_header_destroy(h);
 147         idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*));
 148         for (i = 0; i < idx->n; ++i) idx->index[i] = kh_init(i);
 149         idx->index2 = (bam_lidx_t*)calloc(idx->n, sizeof(bam_lidx_t));
 150
 151         save_bin = save_tid = last_tid = last_bin = 0xffffffffu;
 152         save_off = last_off = bam_tell(fp); last_coor = 0xffffffffu;
 153         while ((ret = bam_read1(fp, b)) >= 0) {
 154                 if (last_tid != c->tid) { // change of chromosomes
 155                         last_tid = c->tid;
 156                         last_bin = 0xffffffffu;
 157                 } else if (last_coor > c->pos) {
 158                         fprintf(stderr, "[bam_index_core] the alignment is not sorted. Abort!\n");
 159                         exit(1);
 160                 }
 161                 if (b->core.tid >= 0 && b->core.bin < 4681) insert_offset2(&idx->index2[b->core.tid], b, last_off);
 162                 if (c->bin != last_bin) { // then possibly write the binning index
 163                         if (save_bin != 0xffffffffu) // save_bin==0xffffffffu only happens to the first record
 164                                 insert_offset(idx->index[save_tid], save_bin, save_off, last_off);
 165                         save_off = last_off;
 166                         save_bin = last_bin = c->bin;
 167                         save_tid = c->tid;
 168                         if (save_tid < 0) break;
 169                 }
 170                 if (bam_tell(fp) <= last_off) {
 171                         fprintf(stderr, "[bam_index_core] bug in BGZF/RAZF: %llx < %llx\n",
 172                                         (unsigned long long)bam_tell(fp), (unsigned long long)last_off);
 173                         exit(1);
 174                 }
 175                 last_off = bam_tell(fp);
 176                 last_coor = b->core.pos;
 177         }
 178         if (save_tid >= 0) insert_offset(idx->index[save_tid], save_bin, save_off, bam_tell(fp));
 179         merge_chunks(idx);
 180         if (ret < -1) fprintf(stderr, "[bam_index_core] truncated file? Continue anyway. (%d)\n", ret);
 181         free(b->data); free(b);
 182         return idx;
 183 }
 184
 185 void bam_index_destroy(bam_index_t *idx)
 186 {
 187         khint_t k;
 188         int i;
 189         if (idx == 0) return;
 190         for (i = 0; i < idx->n; ++i) {
 191                 khash_t(i) *index = idx->index[i];
 192                 bam_lidx_t *index2 = idx->index2 + i;
 193                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 194                         if (kh_exist(index, k))
 195                                 free(kh_value(index, k).list);
 196                 }
 197                 kh_destroy(i, index);
 198                 free(index2->offset);
 199         }
 200         free(idx->index); free(idx->index2);
 201         free(idx);
 202 }
 203
 204 void bam_index_save(const bam_index_t *idx, FILE *fp)
 205 {
 206         int32_t i, size;
 207         khint_t k;
 208         fwrite("BAI\1", 1, 4, fp);
 209         if (bam_is_be) {
 210                 uint32_t x = idx->n;
 211                 fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 212         } else fwrite(&idx->n, 4, 1, fp);
 213         for (i = 0; i < idx->n; ++i) {
 214                 khash_t(i) *index = idx->index[i];
 215                 bam_lidx_t *index2 = idx->index2 + i;
 216                 // write binning index
 217                 size = kh_size(index);
 218                 if (bam_is_be) { // big endian
 219                         uint32_t x = size;
 220                         fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 221                 } else fwrite(&size, 4, 1, fp);
 222                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 223                         if (kh_exist(index, k)) {
 224                                 bam_binlist_t *p = &kh_value(index, k);
 225                                 if (bam_is_be) { // big endian
 226                                         uint32_t x;
 227                                         x = kh_key(index, k); fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 228                                         x = p->n; fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 229                                         for (x = 0; (int)x < p->n; ++x) {
 230                                                 bam_swap_endian_8p(&p->list[x].u);
 231                                                 bam_swap_endian_8p(&p->list[x].v);
 232                                         }
 233                                         fwrite(p->list, 16, p->n, fp);
 234                                         for (x = 0; (int)x < p->n; ++x) {
 235                                                 bam_swap_endian_8p(&p->list[x].u);
 236                                                 bam_swap_endian_8p(&p->list[x].v);
 237                                         }
 238                                 } else {
 239                                         fwrite(&kh_key(index, k), 4, 1, fp);
 240                                         fwrite(&p->n, 4, 1, fp);
 241                                         fwrite(p->list, 16, p->n, fp);
 242                                 }
 243                         }
 244                 }
 245                 // write linear index (index2)
 246                 if (bam_is_be) {
 247                         int x = index2->n;
 248                         fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 249                 } else fwrite(&index2->n, 4, 1, fp);
 250                 if (bam_is_be) { // big endian
 251                         int x;
 252                         for (x = 0; (int)x < index2->n; ++x)
 253                                 bam_swap_endian_8p(&index2->offset[x]);
 254                         fwrite(index2->offset, 8, index2->n, fp);
 255                         for (x = 0; (int)x < index2->n; ++x)
 256                                 bam_swap_endian_8p(&index2->offset[x]);
 257                 } else fwrite(index2->offset, 8, index2->n, fp);
 258         }
 259         fflush(fp);
 260 }
 261
 262 bam_index_t *bam_index_load(const char *fn)
 263 {
 264         bam_index_t *idx;
 265         FILE *fp;
 266         int i;
 267         char *fnidx, magic[4];
 268
 269         fnidx = (char*)calloc(strlen(fn) + 5, 1);
 270         strcpy(fnidx, fn); strcat(fnidx, ".bai");
 271         fp = fopen(fnidx, "r");
 272         if (fp == 0) { // try "{base}.bai"
 273                 char *s = strstr(fn, "bam");
 274                 if (s == fn + strlen(fn) - 3) {
 275                         strcpy(fnidx, fn);
 276                         fnidx[strlen(fn)-1] = 'i';
 277                         fp = fopen(fnidx, "r");
 278                 }
 279         }
 280         if (fp == 0) {
 281                 fprintf(stderr, "[bam_index_load] the alignment is not indexed. Please run `index' command first. Abort!\n");
 282                 exit(1);
 283         }
 284         free(fnidx);
 285
 286         fread(magic, 1, 4, fp);
 287         if (strncmp(magic, "BAI\1", 4)) {
 288                 fprintf(stderr, "[bam_index_load] wrong magic number.\n");
 289                 fclose(fp);
 290                 return 0;
 291         }
 292         idx = (bam_index_t*)calloc(1, sizeof(bam_index_t));
 293         fread(&idx->n, 4, 1, fp);
 294         if (bam_is_be) bam_swap_endian_4p(&idx->n);
 295         idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*));
 296         idx->index2 = (bam_lidx_t*)calloc(idx->n, sizeof(bam_lidx_t));
 297         for (i = 0; i < idx->n; ++i) {
 298                 khash_t(i) *index;
 299                 bam_lidx_t *index2 = idx->index2 + i;
 300                 uint32_t key, size;
 301                 khint_t k;
 302                 int j, ret;
 303                 bam_binlist_t *p;
 304                 index = idx->index[i] = kh_init(i);
 305                 // load binning index
 306                 fread(&size, 4, 1, fp);
 307                 if (bam_is_be) bam_swap_endian_4p(&size);
 308                 for (j = 0; j < (int)size; ++j) {
 309                         fread(&key, 4, 1, fp);
 310                         if (bam_is_be) bam_swap_endian_4p(&key);
 311                         k = kh_put(i, index, key, &ret);
 312                         p = &kh_value(index, k);
 313                         fread(&p->n, 4, 1, fp);
 314                         if (bam_is_be) bam_swap_endian_4p(&p->n);
 315                         p->m = p->n;
 316                         p->list = (pair64_t*)malloc(p->m * 16);
 317                         fread(p->list, 16, p->n, fp);
 318                         if (bam_is_be) {
 319                                 int x;
 320                                 for (x = 0; x < p->n; ++x) {
 321                                         bam_swap_endian_8p(&p->list[x].u);
 322                                         bam_swap_endian_8p(&p->list[x].v);
 323                                 }
 324                         }
 325                 }
 326                 // load linear index
 327                 fread(&index2->n, 4, 1, fp);
 328                 if (bam_is_be) bam_swap_endian_4p(&index2->n);
 329                 index2->m = index2->n;
 330                 index2->offset = (uint64_t*)calloc(index2->m, 8);
 331                 fread(index2->offset, index2->n, 8, fp);
 332                 if (bam_is_be)
 333                         for (j = 0; j < index2->n; ++j) bam_swap_endian_8p(&index2->offset[j]);
 334         }
 335         fclose(fp);
 336         return idx;
 337 }
 338
 339 int bam_index_build2(const char *fn, const char *_fnidx)
 340 {
 341         char *fnidx;
 342         FILE *fpidx;
 343         bamFile fp;
 344         bam_index_t *idx;
 345         assert(fp = bam_open(fn, "r"));
 346         idx = bam_index_core(fp);
 347         bam_close(fp);
 348         if (_fnidx == 0) {
 349                 fnidx = (char*)calloc(strlen(fn) + 5, 1);
 350                 strcpy(fnidx, fn); strcat(fnidx, ".bai");
 351         } else fnidx = strdup(_fnidx);
 352         fpidx = fopen(fnidx, "w");
 353         if (fpidx == 0) {
 354                 fprintf(stderr, "[bam_index_build2] fail to create the index file.\n");
 355                 free(fnidx);
 356                 return 1;
 357         }
 358         bam_index_save(idx, fpidx);
 359         bam_index_destroy(idx);
 360         fclose(fpidx);
 361         free(fnidx);
 362         return 0;
 363 }
 364
 365 int bam_index_build(const char *fn)
 366 {
 367         return bam_index_build2(fn, 0);
 368 }
 369
 370 int bam_index(int argc, char *argv[])
 371 {
 372         if (argc < 2) {
 373                 fprintf(stderr, "Usage: samtools index <in.bam> [<out.index>]\n");
 374                 return 1;
 375         }
 376         if (argc >= 3) bam_index_build2(argv[1], argv[2]);
 377         else bam_index_build(argv[1]);
 378         return 0;
 379 }
 380
 381 #define MAX_BIN 37450 // =(8^6-1)/7+1
 382
 383 static inline int reg2bins(uint32_t beg, uint32_t end, uint16_t list[MAX_BIN])
 384 {
 385         int i = 0, k;
 386         --end;
 387         list[i++] = 0;
 388         for (k =    1 + (beg>>26); k <=    1 + (end>>26); ++k) list[i++] = k;
 389         for (k =    9 + (beg>>23); k <=    9 + (end>>23); ++k) list[i++] = k;
 390         for (k =   73 + (beg>>20); k <=   73 + (end>>20); ++k) list[i++] = k;
 391         for (k =  585 + (beg>>17); k <=  585 + (end>>17); ++k) list[i++] = k;
 392         for (k = 4681 + (beg>>14); k <= 4681 + (end>>14); ++k) list[i++] = k;
 393         return i;
 394 }
 395
 396 static inline int is_overlap(uint32_t beg, uint32_t end, const bam1_t *b)
 397 {
 398         uint32_t rbeg = b->core.pos;
 399         uint32_t rend = b->core.n_cigar? bam_calend(&b->core, bam1_cigar(b)) : b->core.pos + 1;
 400         return (rend > beg && rbeg < end);
 401 }
 402
 403 int bam_fetch(bamFile fp, const bam_index_t *idx, int tid, int beg, int end, void *data, bam_fetch_f func)
 404 {
 405         uint16_t *bins;
 406         int i, n_bins, n_off;
 407         pair64_t *off;
 408         khint_t k;
 409         khash_t(i) *index;
 410         uint64_t min_off;
 411
 412         bins = (uint16_t*)calloc(MAX_BIN, 2);
 413         n_bins = reg2bins(beg, end, bins);
 414         index = idx->index[tid];
 415         min_off = (beg>>BAM_LIDX_SHIFT >= idx->index2[tid].n)? 0 : idx->index2[tid].offset[beg>>BAM_LIDX_SHIFT];
 416         for (i = n_off = 0; i < n_bins; ++i) {
 417                 if ((k = kh_get(i, index, bins[i])) != kh_end(index))
 418                         n_off += kh_value(index, k).n;
 419         }
 420         if (n_off == 0) {
 421                 free(bins); return 0;
 422         }
 423         off = (pair64_t*)calloc(n_off, 16);
 424         for (i = n_off = 0; i < n_bins; ++i) {
 425                 if ((k = kh_get(i, index, bins[i])) != kh_end(index)) {
 426                         int j;
 427                         bam_binlist_t *p = &kh_value(index, k);
 428                         for (j = 0; j < p->n; ++j)
 429                                 if (p->list[j].v > min_off) off[n_off++] = p->list[j];
 430                 }
 431         }
 432         free(bins);
 433         {
 434                 bam1_t *b;
 435                 int l, ret, n_seeks;
 436                 uint64_t curr_off;
 437                 b = (bam1_t*)calloc(1, sizeof(bam1_t));
 438                 ks_introsort(off, n_off, off);
 439                 // resolve completely contained adjacent blocks
 440                 for (i = 1, l = 0; i < n_off; ++i)
 441                         if (off[l].v < off[i].v)
 442                                 off[++l] = off[i];
 443                 n_off = l + 1;
 444                 // resolve overlaps between adjacent blocks; this may happen due to the merge in indexing
 445                 for (i = 1; i < n_off; ++i)
 446                         if (off[i-1].v >= off[i].u) off[i-1].v = off[i].u;
 447                 { // merge adjacent blocks
 448 #if defined(BAM_TRUE_OFFSET) || defined(BAM_VIRTUAL_OFFSET16)
 449                         for (i = 1, l = 0; i < n_off; ++i) {
 450 #ifdef BAM_TRUE_OFFSET
 451                                 if (off[l].v + BAM_MIN_CHUNK_GAP > off[i].u) off[l].v = off[i].v;
 452 #else
 453                                 if (off[l].v>>16 == off[i].u>>16) off[l].v = off[i].v;
 454 #endif
 455                                 else off[++l] = off[i];
 456                         }
 457                         n_off = l + 1;
 458 #endif
 459                 }
 460                 // retrive alignments
 461                 n_seeks = 0; i = -1; curr_off = 0;
 462                 for (;;) {
 463                         if (curr_off == 0 || curr_off >= off[i].v) { // then jump to the next chunk
 464                                 if (i == n_off - 1) break; // no more chunks
 465                                 if (i >= 0) assert(curr_off == off[i].v); // otherwise bug
 466                                 if (i < 0 || off[i].v != off[i+1].u) { // not adjacent chunks; then seek
 467                                         bam_seek(fp, off[i+1].u, SEEK_SET);
 468                                         curr_off = bam_tell(fp);
 469                                         ++n_seeks;
 470                                 }
 471                                 ++i;
 472                         }
 473                         if ((ret = bam_read1(fp, b)) > 0) {
 474                                 curr_off = bam_tell(fp);
 475                                 if (b->core.tid != tid || b->core.pos >= end) break; // no need to proceed
 476                                 else if (is_overlap(beg, end, b)) func(b, data);
 477                         } else break; // end of file
 478                 }
 479 //              fprintf(stderr, "[bam_fetch] # seek calls: %d\n", n_seeks);
 480                 bam_destroy1(b);
 481         }
 482         free(off);
 483         return 0;
 484 }