src/fastq-sample.c

   1
   2 /*
   3  * This file is part of fastq-tools.
   4  *
   5  * Copyright (c) 2011 by Daniel C. Jones <dcjones@cs.washington.edu>
   6  *
   7  * fastq-sample :
   8  * Sample reads with or without replacement from a FASTQ file.
   9  *
  10  */
  11
  12 #include "common.h"
  13 #include "parse.h"
  14 #include "rng.h"
  15 #include <stdlib.h>
  16 #include <stdio.h>
  17 #include <string.h>
  18 #include <getopt.h>
  19 #include <zlib.h>
  20
  21 #if defined(MSDOS) || defined(OS2) || defined(WIN32) || defined(__CYGWIN__)
  22 #  include <fcntl.h>
  23 #  include <io.h>
  24 #  define SET_BINARY_MODE(file) setmode(fileno(file), O_BINARY)
  25 #else
  26 #  define SET_BINARY_MODE(file)
  27 #endif
  28
  29 static const char* prog_name = "fastq-sample";
  30
  31 static int replacement_flag;
  32
  33
  34 void print_help()
  35 {
  36     fprintf(stdout,
  37 "fastq-sample [OPTION]... FILE [FILE2]\n"
  38 "Sample random reads from a FASTQ file."
  39 "Options:\n"
  40 "  -n N                    the number of reads to sample (default: 10000)\n"
  41 "  -p N                    the proportion of the total reads to sample\n"
  42 "  -o, --output=PREFIX     output file prefix\n"
  43 "  -c, --complement-output=PREFIX\n"
  44 "                          output reads not included in the random sample to\n"
  45 "                          a file (or files) with the given prefix (by default,\n"
  46 "                          they are not output).\n"
  47 "  -r, --with-replacement  sample with relpacement\n"
  48 "  -s, --seed=SEED         a manual seed to the random number generator\n"
  49 "  -h, --help              print this message\n"
  50 "  -V, --version           output version information and exit\n"
  51     );
  52 }
  53
  54
  55 /* count the number of entries in a fastq file */
  56 unsigned long count_entries(fastq_t* fqf)
  57 {
  58     seq_t* seq = seq_create();
  59     unsigned long n = 0;
  60     while (fastq_read(fqf, seq)) ++n;
  61     seq_free(seq);
  62
  63     return n;
  64 }
  65
  66
  67 /* compare two unsigned integers (for qsort) */
  68 int cmpul( const void* a, const void* b ) {
  69     if(      *(unsigned long*) a < *(unsigned long*) b ) return -1;
  70     else if( *(unsigned long*) a > *(unsigned long*) b ) return 1;
  71     else                                                 return 0;
  72 }
  73
  74
  75 /* randomly shuffle an array of unsigned longs */
  76 void shuffle(rng_t* rng, unsigned long* xs, unsigned long n)
  77 {
  78     unsigned long i, j, k;
  79     for (i = n - 1; i > 0; --i) {
  80         j = fastq_rng_uniform_int(rng, i + 1);
  81         k = xs[j]; xs[j] = xs[i]; xs[i] = k; // swap
  82     }
  83 }
  84
  85
  86 unsigned long* index_with_replacement(rng_t* rng, unsigned long n, unsigned long k)
  87 {
  88     unsigned long* xs = malloc_or_die(k * sizeof(unsigned long));
  89     size_t i;
  90     for (i = 0; i < k; ++i) xs[i] = fastq_rng_uniform_int(rng, n);
  91     return xs;
  92 }
  93
  94
  95 unsigned long* index_without_replacement(rng_t* rng, unsigned long n)
  96 {
  97     unsigned long* xs = malloc_or_die(n * sizeof(unsigned long));
  98     size_t i;
  99     for (i = 0; i < n; ++i) xs[i] = i;
 100     shuffle(rng, xs, n);
 101     return xs;
 102 }
 103
 104
 105 void fastq_sample(unsigned long rng_seed,
 106                   const char* prefix, const char* cprefix,
 107                   FILE* file1, FILE* file2, unsigned long k, double p)
 108 {
 109      /*
 110       * The basic idea is this:
 111       *
 112       * 1. Count the number of lines in the file, n.
 113       *
 114       * 2a. If sampling with replacement, generate k random integers in [0, n-1].
 115       *
 116       * 2b. If sampling without replacement, generate a list of integers 0..(n-1),
 117       *     shuffle with fisher-yates, then consider the first k.
 118       *
 119       * 3. Sort the integer list.
 120       *
 121       * 3. Read through the file again, when the number at the front of the integer
 122       *    list matches the index of the fastq etry, print the entry, and pop the
 123       *    number.
 124       */
 125
 126
 127     unsigned long n, n2;
 128
 129     fastq_t* f1 = fastq_create(file1);
 130     fastq_t* f2 = file2 == NULL ? NULL : fastq_create(file2);
 131
 132     n = count_entries(f1);
 133     if (f2 != NULL) {
 134         n2 = count_entries(f2);
 135         if (n != n2) {
 136             fprintf(stderr, "Input files have differing numbers of entries (%lu != %lu).\n", n, n2);
 137             exit(1);
 138         }
 139     }
 140
 141     fastq_rewind(f1);
 142     if (f2 != NULL) fastq_rewind(f2);
 143
 144     if (p > 0.0) {
 145         k = (unsigned long) (p * (double) n);
 146         if (!replacement_flag && k > n) k = n;
 147     }
 148
 149     rng_t* rng = fastq_rng_alloc();
 150     fastq_rng_seed(rng, rng_seed);
 151
 152     unsigned long* xs;
 153     if (replacement_flag) xs = index_with_replacement(rng, n, k);
 154     else                  xs = index_without_replacement(rng, n);
 155
 156     qsort(xs, k, sizeof(unsigned long), cmpul);
 157
 158
 159     /* open output */
 160     FILE* fout1;
 161     FILE* fout2;
 162
 163     char* output_name;
 164     size_t output_len;
 165     if (file2 == NULL) {
 166         output_len = strlen(prefix) + 7;
 167         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 168
 169         snprintf(output_name, output_len, "%s.fastq", prefix);
 170         fout1 = fopen(output_name, "wb");
 171         if (fout1 == NULL) {
 172             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 173             exit(1);
 174         }
 175
 176         fout2 = NULL;
 177
 178         free(output_name);
 179     }
 180     else {
 181         output_len = strlen(prefix) + 9;
 182         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 183
 184         snprintf(output_name, output_len, "%s.1.fastq", prefix);
 185         fout1 = fopen(output_name, "wb");
 186         if (fout1 == NULL) {
 187             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 188             exit(1);
 189         }
 190
 191         snprintf(output_name, output_len, "%s.2.fastq", prefix);
 192         fout2 = fopen(output_name, "wb");
 193         if (fout1 == NULL) {
 194             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 195             exit(1);
 196         }
 197
 198         free(output_name);
 199     }
 200
 201
 202     /* open complement output */
 203     FILE* cfout1 = NULL;
 204     FILE* cfout2 = NULL;
 205
 206     if (cprefix != NULL && file2 == NULL) {
 207         output_len = strlen(cprefix) + 7;
 208         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 209
 210         snprintf(output_name, output_len, "%s.fastq", cprefix);
 211         cfout1 = fopen(output_name, "wb");
 212         if (cfout1 == NULL) {
 213             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 214             exit(1);
 215         }
 216
 217         cfout2 = NULL;
 218
 219         free(output_name);
 220     }
 221     else if (cprefix != NULL) {
 222         output_len = strlen(cprefix) + 9;
 223         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 224
 225         snprintf(output_name, output_len, "%s.1.fastq", cprefix);
 226         cfout1 = fopen(output_name, "wb");
 227         if (cfout1 == NULL) {
 228             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 229             exit(1);
 230         }
 231
 232         snprintf(output_name, output_len, "%s.2.fastq", cprefix);
 233         cfout2 = fopen(output_name, "wb");
 234         if (cfout1 == NULL) {
 235             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 236             exit(1);
 237         }
 238
 239         free(output_name);
 240     }
 241
 242
 243
 244     unsigned long i = 0; // read number
 245     unsigned long j = 0; // index into xs
 246
 247     int ret;
 248     seq_t* seq1 = seq_create();
 249     seq_t* seq2 = seq_create();
 250
 251     while (j < k && fastq_read(f1, seq1)) {
 252         if (f2 != NULL){
 253             ret = fastq_read(f2, seq2);
 254             if (ret == 0) {
 255                 fputs("Input files have differing numbers of entries.\n", stderr);
 256                 exit(1);
 257             }
 258         }
 259
 260         if (xs[j] == i) {
 261             while (j < k && xs[j] == i) {
 262                 fastq_print(fout1, seq1);
 263                 if (f2 != NULL) fastq_print(fout2, seq2);
 264                 ++j;
 265             }
 266         }
 267         else if (cfout1 != NULL) {
 268             fastq_print(cfout1, seq1);
 269             if (f2 != NULL) fastq_print(cfout2, seq2);
 270         }
 271
 272         ++i;
 273     }
 274
 275     seq_free(seq1);
 276     seq_free(seq2);
 277     fastq_free(f1);
 278     if (f2 != NULL) fastq_free(f2);
 279
 280     fclose(fout1);
 281     if (fout2 != NULL) fclose(fout2);
 282
 283     if (cfout1 != NULL) fclose(cfout1);
 284     if (cfout2 != NULL) fclose(cfout2);
 285
 286     fastq_rng_free(rng);
 287     free(xs);
 288 }
 289
 290
 291 int main(int argc, char* argv[])
 292 {
 293     SET_BINARY_MODE(stdin);
 294     SET_BINARY_MODE(stdout);
 295
 296     int opt;
 297     int opt_idx;
 298
 299     const char* prefix = NULL;
 300     const char* cprefix = NULL;
 301     unsigned long rng_seed = 4357;
 302     unsigned long k = 10000; // number of reads to sample
 303     double        p = -1;    // proportion of reads to sample
 304
 305
 306     static struct option long_options[] =
 307         {
 308           {"with-replacement",  no_argument,       NULL, 'r'},
 309           {"complement-output", required_argument, NULL, 'c'},
 310           {"seed",              required_argument, NULL, 's'},
 311           {"output",            no_argument,       NULL, 'o'},
 312           {"help",              no_argument,       NULL, 'h'},
 313           {"version",           no_argument,       NULL, 'V'},
 314           {0, 0, 0, 0}
 315         };
 316
 317     replacement_flag = 0;
 318
 319     while (1) {
 320         opt = getopt_long(argc, argv, "n:p:o:rs:hV", long_options, &opt_idx);
 321
 322         if( opt == -1 ) break;
 323
 324         switch (opt) {
 325             case 0:
 326                 if (long_options[opt_idx].flag != 0) break;
 327                 if (optarg) {
 328                 }
 329                 break;
 330
 331             case 'n':
 332                 k = strtoul(optarg, NULL, 10);
 333                 break;
 334
 335             case 'p':
 336                 p = atof(optarg);
 337                 if (p < 0.0) {
 338                     fputs("Sample proportion ('-p') is less than zero.\n", stderr);
 339                     return 1;
 340                 }
 341                 break;
 342
 343             case 'r':
 344                 replacement_flag = 1;
 345                 break;
 346
 347             case 's':
 348                 rng_seed = strtoul(optarg, NULL, 10);
 349                 break;
 350
 351             case 'o':
 352                 prefix = optarg;
 353                 break;
 354
 355             case 'c':
 356                 cprefix = optarg;
 357                 break;
 358
 359             case 'h':
 360                 print_help();
 361                 return 0;
 362
 363             case 'V':
 364                 print_version(stdout, prog_name);
 365                 return 0;
 366
 367             case '?':
 368                 return 1;
 369
 370             default:
 371                 abort();
 372         }
 373     }
 374
 375     FILE* file1 = NULL;
 376     FILE* file2 = NULL;
 377
 378     char* prefix_alloc = NULL;
 379
 380     if (optind >= argc) {
 381         fputs("An input file must be given.\n", stderr);
 382         print_help();
 383         exit(1);
 384     }
 385     else {
 386         file1 = fopen(argv[optind], "rb");
 387         if (file1 == NULL) {
 388             fprintf(stderr, "Cannot open '%s' for reading.\n", argv[optind]);
 389             return 1;
 390         }
 391
 392         if (prefix == NULL) {
 393             /* guess at a reasonable output refix by trimming the
 394              * trailing file extension, if any.  */
 395             char* tmp;
 396
 397             /* base name */
 398             tmp = strrchr(argv[optind], '/');
 399             if (tmp != NULL) argv[optind] = tmp + 1;
 400
 401             /* exclude file suffixes */
 402             tmp = strchr(argv[optind], '.');
 403             if (tmp == NULL) prefix = argv[optind];
 404             else {
 405                 prefix_alloc = malloc_or_die((tmp - argv[optind] + 1) * sizeof(char));
 406                 memcpy(prefix_alloc, argv[optind], (tmp - argv[optind]) * sizeof(char));
 407                 prefix_alloc[tmp - argv[optind]] = '\0';
 408                 prefix = prefix_alloc;
 409             }
 410         }
 411         ++optind;
 412
 413         if (optind < argc) {
 414             file2 = fopen(argv[optind], "rb");
 415             if (file2 == NULL) {
 416                 fprintf(stderr, "Cannot open '%s' for reading.\n", argv[optind]);
 417                 return 1;
 418             }
 419         }
 420     }
 421
 422     fastq_sample(rng_seed, prefix, cprefix, file1, file2, k, p);
 423
 424     free(prefix_alloc);
 425     return 0;
 426 }
 427
 428
 429
 430
 431
 432
 433
 434