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 "  -r, --with-replacement  sample with relpacement\n"
  44 "  -h, --help              print this message\n"
  45 "  -V, --version           output version information and exit\n"
  46     );
  47 }
  48
  49
  50 /* count the number of entries in a fastq file */
  51 unsigned long count_entries(fastq_t* fqf)
  52 {
  53     seq_t* seq = fastq_alloc_seq();
  54     unsigned long n = 0;
  55     while (fastq_next(fqf, seq)) ++n;
  56     fastq_free_seq(seq);
  57
  58     return n;
  59 }
  60
  61
  62 /* compare two unsigned integers (for qsort) */
  63 int cmpul( const void* a, const void* b ) {
  64     if(      *(unsigned long*) a < *(unsigned long*) b ) return -1;
  65     else if( *(unsigned long*) a > *(unsigned long*) b ) return 1;
  66     else                                                 return 0;
  67 }
  68
  69
  70 /* randomly shuffle an array of unsigned longs */
  71 void shuffle(rng_t* rng, unsigned long* xs, unsigned long n)
  72 {
  73     unsigned long i, j, k;
  74     for (i = n - 1; i > 0; --i) {
  75         j = fastq_rng_uniform_int(rng, i + 1);
  76         k = xs[j]; xs[j] = xs[i]; xs[i] = k; // swap
  77     }
  78 }
  79
  80
  81 unsigned long* index_with_replacement(rng_t* rng, unsigned long n, unsigned long k)
  82 {
  83     unsigned long* xs = malloc_or_die(k * sizeof(unsigned long));
  84     size_t i;
  85     for (i = 0; i < k; ++i) xs[i] = fastq_rng_uniform_int(rng, n);
  86     return xs;
  87 }
  88
  89
  90 unsigned long* index_without_replacement(rng_t* rng, unsigned long n)
  91 {
  92     unsigned long* xs = malloc_or_die(n * sizeof(unsigned long));
  93     size_t i;
  94     for (i = 0; i < n; ++i) xs[i] = i;
  95     shuffle(rng, xs, n);
  96     return xs;
  97 }
  98
  99
 100 void fastq_sample(const char* prefix, FILE* file1, FILE* file2, unsigned long k, double p)
 101 {
 102      /*
 103       * The basic idea is this:
 104       *
 105       * 1. Count the number of lines in the file, n.
 106       *
 107       * 2a. If sampling with replacement, generate k random integers in [0, n-1].
 108       *
 109       * 2b. If sampling without replacement, generate a list of integers 0..(n-1),
 110       *     shuffle with fisher-yates, then consider the first k.
 111       *
 112       * 3. Sort the integer list.
 113       *
 114       * 3. Read through the file again, when the number at the front of the integer
 115       *    list matches the index of the fastq etry, print the entry, and pop the
 116       *    number.
 117       */
 118
 119
 120     unsigned long n, n2;
 121
 122     fastq_t* f1 = fastq_open(file1);
 123     fastq_t* f2 = file2 == NULL ? NULL : fastq_open(file2);
 124
 125     n = count_entries(f1);
 126     if (f2 != NULL) {
 127         n2 = count_entries(f2);
 128         if (n != n2) {
 129             fprintf(stderr, "Input files have differing numbers of entries (%lu != %lu).\n", n, n2);
 130             exit(1);
 131         }
 132     }
 133
 134     fastq_rewind(f1);
 135     if (f2 != NULL) fastq_rewind(f2);
 136
 137     if (p > 0.0) {
 138         k = (unsigned long) (p * (double) n);
 139         if (!replacement_flag && k > n) k = n;
 140     }
 141
 142     rng_t* rng = fastq_rng_alloc();
 143
 144     unsigned long* xs;
 145     if (replacement_flag) xs = index_with_replacement(rng, n, k);
 146     else                  xs = index_without_replacement(rng, n);
 147
 148     qsort(xs, k, sizeof(unsigned long), cmpul);
 149
 150
 151     /* open output */
 152     FILE* fout1;
 153     FILE* fout2;
 154
 155     char* output_name;
 156     size_t output_len;
 157     if (file2 == NULL) {
 158         output_len = strlen(prefix) + 7;
 159         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 160
 161         snprintf(output_name, output_len, "%s.fastq", prefix);
 162         fout1 = fopen(output_name, "wb");
 163         if (fout1 == NULL) {
 164             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 165             exit(1);
 166         }
 167
 168         fout2 = NULL;
 169
 170         free(output_name);
 171     }
 172     else {
 173         output_len = strlen(prefix) + 9;
 174         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 175
 176         snprintf(output_name, output_len, "%s.1.fastq", prefix);
 177         fout1 = fopen(output_name, "wb");
 178         if (fout1 == NULL) {
 179             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 180             exit(1);
 181         }
 182
 183         snprintf(output_name, output_len, "%s.2.fastq", prefix);
 184         fout2 = fopen(output_name, "wb");
 185         if (fout1 == NULL) {
 186             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 187             exit(1);
 188         }
 189
 190         free(output_name);
 191     }
 192
 193
 194
 195     unsigned long i = 0; // read number
 196     unsigned long j = 0; // index into xs
 197
 198     int ret;
 199     seq_t* seq1 = fastq_alloc_seq();
 200     seq_t* seq2 = fastq_alloc_seq();
 201
 202     while (j < k && fastq_next(f1, seq1)) {
 203         if (f2 != NULL){
 204             ret = fastq_next(f2, seq2);
 205             if (ret == 0) {
 206                 fputs("Input files have differing numbers of entries.\n", stderr);
 207                 exit(1);
 208             }
 209         }
 210
 211         while (j < k && xs[j] == i) {
 212             fastq_print(fout1, seq1);
 213             if (f2 != NULL) fastq_print(fout2, seq2);
 214             ++j;
 215         }
 216
 217         ++i;
 218     }
 219
 220     fastq_free_seq(seq1);
 221     fastq_free_seq(seq2);
 222     fastq_close(f1);
 223     if (f2 != NULL) fastq_close(f2);
 224
 225     fclose(fout1);
 226     if (fout2 != NULL) fclose(fout2);
 227
 228     fastq_rng_free(rng);
 229     free(xs);
 230 }
 231
 232
 233 int main(int argc, char* argv[])
 234 {
 235     SET_BINARY_MODE(stdin);
 236     SET_BINARY_MODE(stdout);
 237
 238     int opt;
 239     int opt_idx;
 240
 241     const char* prefix = NULL;
 242     unsigned long k = 10000; // number of reads to sample
 243     double        p = -1;    // proportion of reads to sample
 244
 245
 246     static struct option long_options[] =
 247         {
 248           {"with-replacement", no_argument, NULL, 'r'},
 249           {"output",           no_argument, NULL, 'o'},
 250           {"help",             no_argument, NULL, 'h'},
 251           {"version",          no_argument, NULL, 'V'},
 252           {0, 0, 0, 0}
 253         };
 254
 255     replacement_flag = 0;
 256
 257     while (1) {
 258         opt = getopt_long(argc, argv, "n:p:o:rhV", long_options, &opt_idx);
 259
 260         if( opt == -1 ) break;
 261
 262         switch (opt) {
 263             case 0:
 264                 if (long_options[opt_idx].flag != 0) break;
 265                 if (optarg) {
 266                 }
 267                 break;
 268
 269             case 'n':
 270                 k = strtoul(optarg, NULL, 10);
 271                 break;
 272
 273             case 'p':
 274                 p = atof(optarg);
 275                 if (p < 0.0) {
 276                     fputs("Sample proportion ('-p') is less than zero.\n", stderr);
 277                     return 1;
 278                 }
 279                 break;
 280
 281             case 'r':
 282                 replacement_flag = 1;
 283                 break;
 284
 285             case 'o':
 286                 prefix = optarg;
 287                 break;
 288
 289             case 'h':
 290                 print_help();
 291                 return 0;
 292
 293             case 'V':
 294                 print_version(stdout, prog_name);
 295                 return 0;
 296
 297             case '?':
 298                 return 1;
 299
 300             default:
 301                 abort();
 302         }
 303     }
 304
 305     FILE* file1 = NULL;
 306     FILE* file2 = NULL;
 307
 308     char* prefix_alloc = NULL;
 309
 310     if (optind >= argc) {
 311         file1 = stdin;
 312     }
 313     else {
 314         if (strcmp(argv[optind], "-") == 0) {
 315             file1 = stdin;
 316             if (prefix == NULL) prefix = "stdin.sample";
 317         }
 318         else {
 319             file1 = fopen(argv[optind], "rb");
 320             if (file1 == NULL) {
 321                 fprintf(stderr, "Cannot open '%s' for reading.\n", argv[optind]);
 322                 return 1;
 323             }
 324
 325             if (prefix == NULL) {
 326                 /* guess at a reasonable output refix by trimming the
 327                  * trailing file extension, if any.  */
 328                 char* tmp;
 329
 330                 /* base name */
 331                 tmp = strrchr(argv[optind], '/');
 332                 if (tmp != NULL) argv[optind] = tmp + 1;
 333
 334                 /* exclude file suffixes */
 335                 tmp = strchr(argv[optind], '.');
 336                 if (tmp == NULL) prefix = argv[optind];
 337                 else {
 338                     prefix_alloc = malloc_or_die((tmp - argv[optind] + 1) * sizeof(char));
 339                     memcpy(prefix_alloc, argv[optind], (tmp - argv[optind]) * sizeof(char));
 340                     prefix_alloc[tmp - argv[optind]] = '\0';
 341                     prefix = prefix_alloc;
 342                 }
 343             }
 344         }
 345         ++optind;
 346
 347         if (optind < argc) {
 348             if (strcmp(argv[optind], "-") == 0) {
 349                 if (file1 == stdin) {
 350                     fprintf(stderr, "Both input files cannot be standard input.\n");
 351                     return 1;
 352                 }
 353                 file2 = stdin;
 354             }
 355             else {
 356                 file2 = fopen(argv[optind], "rb");
 357                 if (file2 == NULL) {
 358                     fprintf(stderr, "Cannot open '%s' for reading.\n", argv[optind]);
 359                     return 1;
 360                 }
 361             }
 362         }
 363     }
 364
 365     fastq_sample(prefix, file1, file2, k, p);
 366
 367     free(prefix_alloc);
 368     return 0;
 369 }
 370
 371
 372
 373
 374
 375
 376
 377