samtools/bcftools/index.c.pysam.c

   1 #include "pysam.h"
   2
   3 #include <assert.h>
   4 #include <ctype.h>
   5 #include <sys/stat.h>
   6 #include "bam_endian.h"
   7 #include "kstring.h"
   8 #include "bcf.h"
   9 #ifdef _USE_KNETFILE
  10 #include "knetfile.h"
  11 #endif
  12
  13 #define TAD_LIDX_SHIFT 13
  14
  15 typedef struct {
  16         int32_t n, m;
  17         uint64_t *offset;
  18 } bcf_lidx_t;
  19
  20 struct __bcf_idx_t {
  21         int32_t n;
  22         bcf_lidx_t *index2;
  23 };
  24
  25 /************
  26  * indexing *
  27  ************/
  28
  29 static inline void insert_offset2(bcf_lidx_t *index2, int _beg, int _end, uint64_t offset)
  30 {
  31         int i, beg, end;
  32         beg = _beg >> TAD_LIDX_SHIFT;
  33         end = (_end - 1) >> TAD_LIDX_SHIFT;
  34         if (index2->m < end + 1) {
  35                 int old_m = index2->m;
  36                 index2->m = end + 1;
  37                 kroundup32(index2->m);
  38                 index2->offset = (uint64_t*)realloc(index2->offset, index2->m * 8);
  39                 memset(index2->offset + old_m, 0, 8 * (index2->m - old_m));
  40         }
  41         if (beg == end) {
  42                 if (index2->offset[beg] == 0) index2->offset[beg] = offset;
  43         } else {
  44                 for (i = beg; i <= end; ++i)
  45                         if (index2->offset[i] == 0) index2->offset[i] = offset;
  46         }
  47         if (index2->n < end + 1) index2->n = end + 1;
  48 }
  49
  50 bcf_idx_t *bcf_idx_core(bcf_t *bp, bcf_hdr_t *h)
  51 {
  52         bcf_idx_t *idx;
  53         int32_t last_coor, last_tid;
  54         uint64_t last_off;
  55         kstring_t *str;
  56         BGZF *fp = bp->fp;
  57         bcf1_t *b;
  58         int ret;
  59
  60         b = calloc(1, sizeof(bcf1_t));
  61         str = calloc(1, sizeof(kstring_t));
  62         idx = (bcf_idx_t*)calloc(1, sizeof(bcf_idx_t));
  63         idx->n = h->n_ref;
  64         idx->index2 = calloc(h->n_ref, sizeof(bcf_lidx_t));
  65
  66         last_tid = 0xffffffffu;
  67         last_off = bgzf_tell(fp); last_coor = 0xffffffffu;
  68         while ((ret = bcf_read(bp, h, b)) > 0) {
  69                 int end, tmp;
  70                 if (last_tid != b->tid) { // change of chromosomes
  71                         last_tid = b->tid;
  72                 } else if (last_coor > b->pos) {
  73                         fprintf(pysamerr, "[bcf_idx_core] the input is out of order\n");
  74                         free(str->s); free(str); free(idx); bcf_destroy(b);
  75                         return 0;
  76                 }
  77                 tmp = strlen(b->ref);
  78                 end = b->pos + (tmp > 0? tmp : 1);
  79                 insert_offset2(&idx->index2[b->tid], b->pos, end, last_off);
  80                 last_off = bgzf_tell(fp);
  81                 last_coor = b->pos;
  82         }
  83         free(str->s); free(str); bcf_destroy(b);
  84         return idx;
  85 }
  86
  87 void bcf_idx_destroy(bcf_idx_t *idx)
  88 {
  89         int i;
  90         if (idx == 0) return;
  91         for (i = 0; i < idx->n; ++i) free(idx->index2[i].offset);
  92         free(idx->index2);
  93         free(idx);
  94 }
  95
  96 /******************
  97  * index file I/O *
  98  ******************/
  99
 100 void bcf_idx_save(const bcf_idx_t *idx, BGZF *fp)
 101 {
 102         int32_t i, ti_is_be;
 103         ti_is_be = bam_is_big_endian();
 104         bgzf_write(fp, "BCI\4", 4);
 105         if (ti_is_be) {
 106                 uint32_t x = idx->n;
 107                 bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 108         } else bgzf_write(fp, &idx->n, 4);
 109         for (i = 0; i < idx->n; ++i) {
 110                 bcf_lidx_t *index2 = idx->index2 + i;
 111                 // write linear index (index2)
 112                 if (ti_is_be) {
 113                         int x = index2->n;
 114                         bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 115                 } else bgzf_write(fp, &index2->n, 4);
 116                 if (ti_is_be) { // big endian
 117                         int x;
 118                         for (x = 0; (int)x < index2->n; ++x)
 119                                 bam_swap_endian_8p(&index2->offset[x]);
 120                         bgzf_write(fp, index2->offset, 8 * index2->n);
 121                         for (x = 0; (int)x < index2->n; ++x)
 122                                 bam_swap_endian_8p(&index2->offset[x]);
 123                 } else bgzf_write(fp, index2->offset, 8 * index2->n);
 124         }
 125 }
 126
 127 static bcf_idx_t *bcf_idx_load_core(BGZF *fp)
 128 {
 129         int i, ti_is_be;
 130         char magic[4];
 131         bcf_idx_t *idx;
 132         ti_is_be = bam_is_big_endian();
 133         if (fp == 0) {
 134                 fprintf(pysamerr, "[%s] fail to load index.\n", __func__);
 135                 return 0;
 136         }
 137         bgzf_read(fp, magic, 4);
 138         if (strncmp(magic, "BCI\4", 4)) {
 139                 fprintf(pysamerr, "[%s] wrong magic number.\n", __func__);
 140                 return 0;
 141         }
 142         idx = (bcf_idx_t*)calloc(1, sizeof(bcf_idx_t));
 143         bgzf_read(fp, &idx->n, 4);
 144         if (ti_is_be) bam_swap_endian_4p(&idx->n);
 145         idx->index2 = (bcf_lidx_t*)calloc(idx->n, sizeof(bcf_lidx_t));
 146         for (i = 0; i < idx->n; ++i) {
 147                 bcf_lidx_t *index2 = idx->index2 + i;
 148                 int j;
 149                 bgzf_read(fp, &index2->n, 4);
 150                 if (ti_is_be) bam_swap_endian_4p(&index2->n);
 151                 index2->m = index2->n;
 152                 index2->offset = (uint64_t*)calloc(index2->m, 8);
 153                 bgzf_read(fp, index2->offset, index2->n * 8);
 154                 if (ti_is_be)
 155                         for (j = 0; j < index2->n; ++j) bam_swap_endian_8p(&index2->offset[j]);
 156         }
 157         return idx;
 158 }
 159
 160 bcf_idx_t *bcf_idx_load_local(const char *fnidx)
 161 {
 162         BGZF *fp;
 163         fp = bgzf_open(fnidx, "r");
 164         if (fp) {
 165                 bcf_idx_t *idx = bcf_idx_load_core(fp);
 166                 bgzf_close(fp);
 167                 return idx;
 168         } else return 0;
 169 }
 170
 171 #ifdef _USE_KNETFILE
 172 static void download_from_remote(const char *url)
 173 {
 174         const int buf_size = 1 * 1024 * 1024;
 175         char *fn;
 176         FILE *fp;
 177         uint8_t *buf;
 178         knetFile *fp_remote;
 179         int l;
 180         if (strstr(url, "ftp://") != url && strstr(url, "http://") != url) return;
 181         l = strlen(url);
 182         for (fn = (char*)url + l - 1; fn >= url; --fn)
 183                 if (*fn == '/') break;
 184         ++fn; // fn now points to the file name
 185         fp_remote = knet_open(url, "r");
 186         if (fp_remote == 0) {
 187                 fprintf(pysamerr, "[download_from_remote] fail to open remote file.\n");
 188                 return;
 189         }
 190         if ((fp = fopen(fn, "w")) == 0) {
 191                 fprintf(pysamerr, "[download_from_remote] fail to create file in the working directory.\n");
 192                 knet_close(fp_remote);
 193                 return;
 194         }
 195         buf = (uint8_t*)calloc(buf_size, 1);
 196         while ((l = knet_read(fp_remote, buf, buf_size)) != 0)
 197                 fwrite(buf, 1, l, fp);
 198         free(buf);
 199         fclose(fp);
 200         knet_close(fp_remote);
 201 }
 202 #else
 203 static void download_from_remote(const char *url)
 204 {
 205         return;
 206 }
 207 #endif
 208
 209 static char *get_local_version(const char *fn)
 210 {
 211     struct stat sbuf;
 212         char *fnidx = (char*)calloc(strlen(fn) + 5, 1);
 213         strcat(strcpy(fnidx, fn), ".bci");
 214         if ((strstr(fnidx, "ftp://") == fnidx || strstr(fnidx, "http://") == fnidx)) {
 215                 char *p, *url;
 216                 int l = strlen(fnidx);
 217                 for (p = fnidx + l - 1; p >= fnidx; --p)
 218                         if (*p == '/') break;
 219                 url = fnidx; fnidx = strdup(p + 1);
 220                 if (stat(fnidx, &sbuf) == 0) {
 221                         free(url);
 222                         return fnidx;
 223                 }
 224                 fprintf(pysamerr, "[%s] downloading the index file...\n", __func__);
 225                 download_from_remote(url);
 226                 free(url);
 227         }
 228     if (stat(fnidx, &sbuf) == 0) return fnidx;
 229         free(fnidx); return 0;
 230 }
 231
 232 bcf_idx_t *bcf_idx_load(const char *fn)
 233 {
 234         bcf_idx_t *idx;
 235     char *fname = get_local_version(fn);
 236         if (fname == 0) return 0;
 237         idx = bcf_idx_load_local(fname);
 238     free(fname);
 239         return idx;
 240 }
 241
 242 int bcf_idx_build2(const char *fn, const char *_fnidx)
 243 {
 244         char *fnidx;
 245         BGZF *fpidx;
 246         bcf_t *bp;
 247         bcf_idx_t *idx;
 248         bcf_hdr_t *h;
 249         if ((bp = bcf_open(fn, "r")) == 0) {
 250                 fprintf(pysamerr, "[bcf_idx_build2] fail to open the BAM file.\n");
 251                 return -1;
 252         }
 253         h = bcf_hdr_read(bp);
 254         idx = bcf_idx_core(bp, h);
 255         bcf_close(bp);
 256         if (_fnidx == 0) {
 257                 fnidx = (char*)calloc(strlen(fn) + 5, 1);
 258                 strcpy(fnidx, fn); strcat(fnidx, ".bci");
 259         } else fnidx = strdup(_fnidx);
 260         fpidx = bgzf_open(fnidx, "w");
 261         if (fpidx == 0) {
 262                 fprintf(pysamerr, "[bcf_idx_build2] fail to create the index file.\n");
 263                 free(fnidx);
 264                 return -1;
 265         }
 266         bcf_idx_save(idx, fpidx);
 267         bcf_idx_destroy(idx);
 268         bgzf_close(fpidx);
 269         free(fnidx);
 270         return 0;
 271 }
 272
 273 int bcf_idx_build(const char *fn)
 274 {
 275         return bcf_idx_build2(fn, 0);
 276 }
 277
 278 /********************************************
 279  * parse a region in the format chr:beg-end *
 280  ********************************************/
 281
 282 int bcf_parse_region(void *str2id, const char *str, int *tid, int *begin, int *end)
 283 {
 284         char *s, *p;
 285         int i, l, k;
 286         l = strlen(str);
 287         p = s = (char*)malloc(l+1);
 288         /* squeeze out "," */
 289         for (i = k = 0; i != l; ++i)
 290                 if (str[i] != ',' && !isspace(str[i])) s[k++] = str[i];
 291         s[k] = 0;
 292         for (i = 0; i != k; ++i) if (s[i] == ':') break;
 293         s[i] = 0;
 294         if ((*tid = bcf_str2id(str2id, s)) < 0) {
 295                 free(s);
 296                 return -1;
 297         }
 298         if (i == k) { /* dump the whole sequence */
 299                 *begin = 0; *end = 1<<29; free(s);
 300                 return 0;
 301         }
 302         for (p = s + i + 1; i != k; ++i) if (s[i] == '-') break;
 303         *begin = atoi(p);
 304         if (i < k) {
 305                 p = s + i + 1;
 306                 *end = atoi(p);
 307         } else *end = 1<<29;
 308         if (*begin > 0) --*begin;
 309         free(s);
 310         if (*begin > *end) return -1;
 311         return 0;
 312 }
 313
 314 /*******************************
 315  * retrieve a specified region *
 316  *******************************/
 317
 318 uint64_t bcf_idx_query(const bcf_idx_t *idx, int tid, int beg)
 319 {
 320         uint64_t min_off, *offset;
 321         int i;
 322         if (beg < 0) beg = 0;
 323         offset = idx->index2[tid].offset;
 324         for (i = beg>>TAD_LIDX_SHIFT; i < idx->index2[tid].n && offset[i] == 0; ++i);
 325         min_off = (i == idx->index2[tid].n)? offset[idx->index2[tid].n-1] : offset[i];
 326         return min_off;
 327 }
 328
 329 int bcf_main_index(int argc, char *argv[])
 330 {
 331         if (argc == 1) {
 332                 fprintf(pysamerr, "Usage: bcftools index <in.bcf>\n");
 333                 return 1;
 334         }
 335         bcf_idx_build(argv[1]);
 336         return 0;
 337 }