src/fastq-sort.c

   1 /*
   2  * This file is part of fastq-tools.
   3  *
   4  * Copyright (c) 2012 by Daniel C. Jones <dcjones@cs.washington.edu>
   5  *
   6  * fastq-sample :
   7  * Sample reads with or without replacement from a FASTQ file.
   8  *
   9  */
  10
  11 #include <assert.h>
  12 #include <getopt.h>
  13 #include <string.h>
  14 #include <unistd.h>
  15
  16 #include "common.h"
  17 #include "parse.h"
  18
  19
  20 /* A collection of filenames of sorted chunks of fastq. */
  21 typedef struct seq_dumps_t_
  22 {
  23     char** fns;
  24     size_t n;
  25     size_t size;
  26 } seq_dumps_t;
  27
  28
  29 seq_dumps_t* seq_dumps_create()
  30 {
  31     seq_dumps_t* d = malloc_or_die(sizeof(seq_dumps_t));
  32     d->n = 0;
  33     d->size = 64;
  34     d->fns = malloc_or_die(d->size * sizeof(char*));
  35     return d;
  36 }
  37
  38
  39 void seq_dumps_free(seq_dumps_t* d)
  40 {
  41     size_t i;
  42     for (i = 0; i < d->n; ++i) {
  43         if (unlink(d->fns[i]) == -1) {
  44             fprintf(stderr, "Warning: unable to remove temporary file %s.\n",
  45                     d->fns[i]);
  46         }
  47         free(d->fns[i]);
  48     }
  49     free(d->fns);
  50     free(d);
  51 }
  52
  53
  54 static inline size_t heap_parent(size_t i) { return (i - 1) / 2; }
  55 static inline size_t heap_left  (size_t i) { return 2 * i + 1; }
  56 static inline size_t heap_right (size_t i) { return 2 * i + 2; }
  57
  58
  59 /* Enqueue an item in the heap.
  60  *
  61  * Args:
  62  *   heap: A binary heap in which stores indexs into seqs.
  63  *   n: Fixed maximum size of the heap.
  64  *   m: Current size of the heap.
  65  *   seqs: Sequences.
  66  *   idx: Index to enqueue.
  67  *   cmp: Comparison function.
  68  *
  69  */
  70 void heap_push(size_t* heap, size_t n, size_t* m, seq_t** seqs,
  71                int (*cmp)(const void*, const void*), size_t idx)
  72 {
  73     if (*m >= n) return;
  74     size_t tmp, j, i = (*m)++;
  75     heap[i] = idx;
  76
  77     /* percolate up */
  78     while (i > 0) {
  79         j = heap_parent(i);
  80         if (cmp(seqs[heap[i]], seqs[heap[j]]) < 0) {
  81             tmp = heap[i];
  82             heap[i] = heap[j];
  83             heap[j] = tmp;
  84             i = j;
  85         }
  86         else break;
  87     }
  88 }
  89
  90
  91 /* Dequeue an item from a heap. */
  92 size_t heap_pop(size_t* heap, size_t* m, seq_t** seqs,
  93                int (*cmp)(const void*, const void*))
  94 {
  95     assert(*m > 0);
  96     size_t ans = heap[0];
  97     heap[0] = heap[--(*m)];
  98     size_t tmp, l, r, j, i = 0;
  99     while (true) {
 100         l = heap_left(i);
 101         r = heap_right(i);
 102
 103         if (l >= (*m)) break;
 104
 105         if (r >= (*m)) {
 106             j = l;
 107         }
 108         else {
 109             j = cmp(seqs[heap[l]], seqs[heap[r]]) < 0 ? l : r;
 110         }
 111
 112         if (cmp(seqs[heap[i]], seqs[heap[j]]) > 0) {
 113             tmp = heap[i];
 114             heap[i] = heap[j];
 115             heap[j] = tmp;
 116             i = j;
 117         }
 118         else break;
 119     }
 120
 121     return ans;
 122 }
 123
 124
 125 /* n-way merge sort to stdout */
 126 void merge_sort(const seq_dumps_t* d, int (*cmp)(const void*, const void*))
 127 {
 128     FILE** files = malloc_or_die(d->n * sizeof(FILE*));
 129     size_t i;
 130     for (i = 0; i < d->n; ++i) {
 131         files[i] = fopen(d->fns[i], "rb");
 132         if (files[i] == NULL) {
 133             fprintf(stderr, "Cannot open temporary file %s for reading.\n",
 134                     d->fns[i]);
 135             exit(EXIT_FAILURE);
 136         }
 137     }
 138
 139     fastq_t** fs = malloc_or_die(d->n * sizeof(fastq_t*));
 140     seq_t** seqs = malloc_or_die(d->n * sizeof(seq_t*));
 141     for (i = 0; i < d->n; ++i) {
 142         fs[i] = fastq_create(files[i]);
 143         seqs[i] = seq_create();
 144     }
 145
 146     /* A binary heap of indexes to fs. We use this to repeatedly pop the
 147      * smallest fastq entry. */
 148     size_t* heap = malloc_or_die(d->n * sizeof(size_t));
 149
 150     /* heap size */
 151     size_t m = 0;
 152
 153     for (i = 0; i < d->n; ++i) {
 154         if (fastq_read(fs[i], seqs[i])) {
 155             heap_push(heap, d->n, &m, seqs, cmp, i);
 156         }
 157     }
 158
 159     while (m > 0) {
 160         i = heap_pop(heap, &m, seqs, cmp);
 161         fastq_print(stdout, seqs[i]);
 162         if (fastq_read(fs[i], seqs[i])) {
 163             heap_push(heap, d->n, &m, seqs, cmp, i);
 164         }
 165     }
 166
 167     for (i = 0; i < d->n; ++i) {
 168         seq_free(seqs[i]);
 169         fastq_free(fs[i]);
 170         fclose(files[i]);
 171     }
 172
 173     free(files);
 174     free(fs);
 175 }
 176
 177
 178 typedef struct seq_array_t_
 179 {
 180     seq_t* seqs;
 181
 182     /* Number of seq objects. */
 183     size_t n;
 184
 185     /* Size reserved in seqs. */
 186     size_t size;
 187
 188     /* Space reserved for strings. */
 189     char* data;
 190
 191     /* Data used. */
 192     size_t data_used;
 193
 194     /* Total data size. */
 195     size_t data_size;
 196 } seq_array_t;
 197
 198
 199 seq_array_t* seq_array_create(size_t data_size)
 200 {
 201     seq_array_t* a = malloc_or_die(sizeof(seq_array_t));
 202     a->size = 1024;
 203     a->n = 0;
 204     a->seqs = malloc_or_die(a->size * sizeof(seq_t));
 205
 206     a->data_size = data_size;
 207     a->data_used = 0;
 208     a->data = malloc_or_die(data_size);
 209
 210     return a;
 211 }
 212
 213
 214 void seq_array_free(seq_array_t* a)
 215 {
 216     free(a->seqs);
 217     free(a->data);
 218     free(a);
 219 }
 220
 221
 222 /* Push a fastq entry to back of the array. Return false if there was not enough
 223  * space. */
 224 bool seq_array_push(seq_array_t* a, const seq_t* seq)
 225 {
 226     size_t size_needed = (seq->id1.n + 1) + (seq->seq.n + 1) +
 227                          (seq->id2.n + 1) + (seq->qual.n + 1);
 228
 229     if (size_needed > a->data_size - a->data_used) return false;
 230
 231     if (a->n == a->size) {
 232         a->size *= 2;
 233         a->seqs = realloc_or_die(a->seqs, a->size * sizeof(seq_t));
 234     }
 235
 236     memcpy(&a->data[a->data_used], seq->id1.s, seq->id1.n + 1);
 237     a->seqs[a->n].id1.s = &a->data[a->data_used];
 238     a->seqs[a->n].id1.n = seq->id1.n + 1;
 239     a->data_used += seq->id1.n + 1;
 240
 241     memcpy(&a->data[a->data_used], seq->seq.s, seq->seq.n + 1);
 242     a->seqs[a->n].seq.s = &a->data[a->data_used];
 243     a->seqs[a->n].seq.n = seq->seq.n + 1;
 244     a->data_used += seq->seq.n + 1;
 245
 246     memcpy(&a->data[a->data_used], seq->id2.s, seq->id2.n + 1);
 247     a->seqs[a->n].id2.s = &a->data[a->data_used];
 248     a->seqs[a->n].id2.n = seq->id2.n + 1;
 249     a->data_used += seq->id2.n + 1;
 250
 251     memcpy(&a->data[a->data_used], seq->qual.s, seq->qual.n + 1);
 252     a->seqs[a->n].qual.s = &a->data[a->data_used];
 253     a->seqs[a->n].qual.n = seq->qual.n + 1;
 254     a->data_used += seq->qual.n + 1;
 255
 256     ++a->n;
 257
 258     return true;
 259 }
 260
 261
 262 void seq_array_clear(seq_array_t* a)
 263 {
 264     a->n = 0;
 265     a->data_used = 0;
 266 }
 267
 268
 269 void seq_array_sort(seq_array_t* a, int (*cmp)(const void*, const void*))
 270 {
 271     qsort(a->seqs, a->n, sizeof(seq_t), cmp);
 272 }
 273
 274
 275 int seq_cmp_hash(const void* a, const void* b)
 276 {
 277     uint32_t ha = seq_hash((seq_t*) a);
 278     uint32_t hb = seq_hash((seq_t*) b);
 279
 280     if      (ha < hb) return -1;
 281     else if (ha > hb) return  1;
 282     else              return  0;
 283     return 0;
 284 }
 285
 286
 287 int seq_cmp_id(const void* a, const void* b)
 288 {
 289     return strcmp(((seq_t*) a)->id1.s, ((seq_t*) b)->id1.s);
 290 }
 291
 292
 293 int seq_cmp_seq(const void* a, const void* b)
 294 {
 295     return strcmp(((seq_t*) a)->seq.s, ((seq_t*) b)->seq.s);
 296 }
 297
 298
 299 void seq_array_dump(seq_dumps_t* d, const seq_array_t* a)
 300 {
 301     const char* template = "/tmp/fastq_sort.XXXXXXXX";
 302     char* fn = malloc_or_die(strlen(template) + 1);
 303     memcpy(fn, template, strlen(template) + 1);
 304     if (mktemp(fn) == NULL) {
 305         fprintf(stderr, "Unable to create a temporary file.\n");
 306         exit(EXIT_FAILURE);
 307     }
 308
 309     FILE* f = fopen(fn, "wb");
 310     if (f == NULL) {
 311         fprintf(stderr, "Unable to open temporary file %s for writing.\n", fn);
 312         exit(EXIT_FAILURE);
 313     }
 314
 315     size_t i;
 316     for (i = 0; i < a->n; ++i) {
 317         fastq_print(f, &a->seqs[i]);
 318     }
 319
 320     if (d->n == d->size) {
 321         d->size *= 2;
 322         d->fns = realloc_or_die(d->fns, d->size * sizeof(char*));
 323     }
 324     d->fns[d->n++] = fn;
 325
 326     fclose(f);
 327 }
 328
 329
 330 static const char* prog_name = "fastq-sort";
 331
 332
 333 void print_help()
 334 {
 335     fprintf(stdout,
 336 "fastq-sort [OPTION]... [FILE]...\n"
 337 "Concatenate and sort FASTQ files and write to standard output.\n"
 338 "Options:\n"
 339 "  -I, --id           sort alphabetically by read identifier\n"
 340 "  -S, --seq          sort alphabetically by sequence\n"
 341 "  -R, --random       randomly shuffle the sequences\n"
 342 "  -G, --gc           sort by GC content\n" /* TODO */
 343 "  -M, --median-qual  sort by median quality score\n" /* TODO */
 344 "  -h, --help         print this message\n"
 345 "  -V, --version      output version information and exit\n"
 346    );
 347 }
 348
 349
 350 int main(int argc, char* argv[])
 351 {
 352     int opt, opt_idx;
 353     size_t buffer_size = 100000000;
 354     int (*cmp)(const void*, const void*) = seq_cmp_hash;;
 355
 356     static struct option long_options[] =
 357     {
 358         {"id",      no_argument, NULL, 'I'},
 359         {"seq",     no_argument, NULL, 'S'},
 360         {"random",  no_argument, NULL, 'R'},
 361         {"help",    no_argument, NULL, 'h'},
 362         {"version", no_argument, NULL, 'V'},
 363         {0, 0, 0, 0}
 364     };
 365
 366     while (true) {
 367         opt = getopt_long(argc, argv, "hV", long_options, &opt_idx);
 368         if (opt == -1) break;
 369
 370         switch (opt) {
 371             case 'I':
 372                 cmp = seq_cmp_id;
 373                 break;
 374
 375             case 'S':
 376                 cmp = seq_cmp_seq;
 377                 break;
 378
 379             case 'R':
 380                 cmp = seq_cmp_hash;
 381                 break;
 382
 383             case 'h':
 384                 print_help();
 385                 return 0;
 386
 387             case 'V':
 388                 print_version(stdout, prog_name);
 389                 return 0;
 390
 391             case '?':
 392                 return 1;
 393
 394             default:
 395                 abort();
 396         }
 397     }
 398
 399     seq_array_t* a = seq_array_create(buffer_size);
 400     seq_dumps_t* d = seq_dumps_create();
 401     seq_t* seq = seq_create();
 402
 403     fastq_t* f;
 404     if (optind >= argc) {
 405         f = fastq_create(stdin);
 406         while (fastq_read(f, seq)) {
 407             if (!seq_array_push(a, seq)) {
 408                 seq_array_sort(a, cmp);
 409                 seq_array_dump(d, a);
 410                 seq_array_clear(a);
 411                 if (!seq_array_push(a, seq)) {
 412                     fprintf(stderr, "The buffer size is to small.\n");
 413                     return EXIT_FAILURE;
 414                 }
 415             }
 416         }
 417         fastq_free(f);
 418     }
 419     else {
 420         FILE* file;
 421         for (; optind < argc; ++optind) {
 422             file = fopen(argv[optind], "rb");
 423             if (file == NULL) {
 424                 fprintf(stderr, "Cannot open %s for reading.\n", argv[optind]);
 425                 return EXIT_FAILURE;
 426             }
 427             f = fastq_create(file);
 428
 429             while (fastq_read(f, seq)) {
 430                 if (!seq_array_push(a, seq)) {
 431                     seq_array_sort(a, cmp);
 432                     seq_array_dump(d, a);
 433                     seq_array_clear(a);
 434                     if (!seq_array_push(a, seq)) {
 435                         fprintf(stderr, "The buffer size is to small.\n");
 436                         return EXIT_FAILURE;
 437                     }
 438                 }
 439             }
 440
 441             fastq_free(f);
 442             fclose(file);
 443         }
 444     }
 445
 446     if (a->n > 0) {
 447         seq_array_sort(a, cmp);
 448
 449         /* We were able to sort everything in memory. */
 450         if (d->n == 0) {
 451             size_t i;
 452             for (i = 0; i < a->n; ++i) {
 453                 fastq_print(stdout, &a->seqs[i]);
 454             }
 455         }
 456         else {
 457             seq_array_dump(d, a);
 458             merge_sort(d, cmp);
 459         }
 460     }
 461
 462     seq_dumps_free(d);
 463     seq_free(seq);
 464     seq_array_free(a);
 465
 466     return EXIT_SUCCESS;
 467 }
 468
 469