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 #include <getopt.h>
  12 #include <math.h>
  13 #include <stdio.h>
  14 #include <stdlib.h>
  15 #include <string.h>
  16
  17 #include "common.h"
  18 #include "parse.h"
  19 #include "rng.h"
  20
  21
  22 static const char* prog_name = "fastq-sample";
  23
  24 static int replacement_flag;
  25
  26
  27 void print_help()
  28 {
  29     fprintf(stdout,
  30 "fastq-sample [OPTION]... FILE [FILE2]\n"
  31 "Sample random reads from a FASTQ file."
  32 "Options:\n"
  33 "  -n N                    the number of reads to sample (default: 10000)\n"
  34 "  -p N                    the proportion of the total reads to sample\n"
  35 "  -o, --output=PREFIX     output file prefix\n (Default: \"sample\")"
  36 "  -c, --complement-output=PREFIX\n"
  37 "                          output reads not included in the random sample to\n"
  38 "                          a file (or files) with the given prefix (by default,\n"
  39 "                          they are not output).\n"
  40 "  -r, --with-replacement  sample with relpacement\n"
  41 "  -s, --seed=SEED         a manual seed to the random number generator\n"
  42 "  -h, --help              print this message\n"
  43 "  -V, --version           output version information and exit\n"
  44 );
  45 }
  46
  47
  48 /* count the number of entries in a fastq file */
  49 unsigned long count_entries(fastq_t* fqf)
  50 {
  51     seq_t* seq = seq_create();
  52     unsigned long n = 0;
  53     while (fastq_read(fqf, seq)) ++n;
  54     seq_free(seq);
  55
  56     return n;
  57 }
  58
  59
  60 /* compare two unsigned integers (for qsort) */
  61 int cmpul( const void* a, const void* b ) {
  62     if(      *(unsigned long*) a < *(unsigned long*) b ) return -1;
  63     else if( *(unsigned long*) a > *(unsigned long*) b ) return 1;
  64     else                                                 return 0;
  65 }
  66
  67
  68 /* randomly shuffle an array of unsigned longs */
  69 void shuffle(rng_t* rng, unsigned long* xs, unsigned long n)
  70 {
  71     unsigned long i, j, k;
  72     for (i = n - 1; i > 0; --i) {
  73         j = fastq_rng_uniform_int(rng, i + 1);
  74         k = xs[j]; xs[j] = xs[i]; xs[i] = k; // swap
  75     }
  76 }
  77
  78
  79 unsigned long* index_with_replacement(rng_t* rng, unsigned long n, unsigned long k)
  80 {
  81     unsigned long* xs = malloc_or_die(k * sizeof(unsigned long));
  82     size_t i;
  83     for (i = 0; i < k; ++i) xs[i] = fastq_rng_uniform_int(rng, n);
  84     return xs;
  85 }
  86
  87
  88 unsigned long* index_without_replacement(rng_t* rng, unsigned long n)
  89 {
  90     unsigned long* xs = malloc_or_die(n * sizeof(unsigned long));
  91     size_t i;
  92     for (i = 0; i < n; ++i) xs[i] = i;
  93     shuffle(rng, xs, n);
  94     return xs;
  95 }
  96
  97
  98 void fastq_sample(unsigned long rng_seed,
  99         const char* prefix, const char* cprefix,
 100         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_create(file1);
 123     fastq_t* f2 = file2 == NULL ? NULL : fastq_create(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) round(p * (double) n);
 139         if (!replacement_flag && k > n) k = n;
 140     }
 141
 142     rng_t* rng = fastq_rng_alloc();
 143     fastq_rng_seed(rng, rng_seed);
 144
 145     unsigned long* xs;
 146     if (replacement_flag) xs = index_with_replacement(rng, n, k);
 147     else                  xs = index_without_replacement(rng, n);
 148
 149     qsort(xs, k, sizeof(unsigned long), cmpul);
 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 = open_without_clobber(output_name);
 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 = open_without_clobber(output_name);
 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         fout1 = open_without_clobber(output_name);
 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     /* open complement output */
 194     FILE* cfout1 = NULL;
 195     FILE* cfout2 = NULL;
 196
 197     if (cprefix != NULL && file2 == NULL) {
 198         output_len = strlen(cprefix) + 7;
 199         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 200
 201         snprintf(output_name, output_len, "%s.fastq", cprefix);
 202         cfout1 = fopen(output_name, "wb");
 203         if (cfout1 == NULL) {
 204             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 205             exit(1);
 206         }
 207
 208         cfout2 = NULL;
 209
 210         free(output_name);
 211     }
 212     else if (cprefix != NULL) {
 213         output_len = strlen(cprefix) + 9;
 214         output_name = malloc_or_die((output_len + 1) * sizeof(char));
 215
 216         snprintf(output_name, output_len, "%s.1.fastq", cprefix);
 217         cfout1 = fopen(output_name, "wb");
 218         if (cfout1 == NULL) {
 219             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 220             exit(1);
 221         }
 222
 223         snprintf(output_name, output_len, "%s.2.fastq", cprefix);
 224         cfout2 = fopen(output_name, "wb");
 225         if (cfout1 == NULL) {
 226             fprintf(stderr, "Cannot open file %s for writing.\n", output_name);
 227             exit(1);
 228         }
 229
 230         free(output_name);
 231     }
 232
 233     unsigned long i = 0; // read number
 234     unsigned long j = 0; // index into xs
 235
 236     int ret;
 237     seq_t* seq1 = seq_create();
 238     seq_t* seq2 = seq_create();
 239
 240     while (j < k && fastq_read(f1, seq1)) {
 241         if (f2 != NULL){
 242             ret = fastq_read(f2, seq2);
 243             if (ret == 0) {
 244                 fputs("Input files have differing numbers of entries.\n", stderr);
 245                 exit(1);
 246             }
 247         }
 248
 249         if (xs[j] == i) {
 250             while (j < k && xs[j] == i) {
 251                 fastq_print(fout1, seq1);
 252                 if (f2 != NULL) fastq_print(fout2, seq2);
 253                 ++j;
 254             }
 255         }
 256         else if (cfout1 != NULL) {
 257             fastq_print(cfout1, seq1);
 258             if (f2 != NULL) fastq_print(cfout2, seq2);
 259         }
 260
 261         ++i;
 262     }
 263
 264     seq_free(seq1);
 265     seq_free(seq2);
 266     fastq_free(f1);
 267     if (f2 != NULL) fastq_free(f2);
 268
 269     fclose(fout1);
 270     if (fout2 != NULL) fclose(fout2);
 271
 272     if (cfout1 != NULL) fclose(cfout1);
 273     if (cfout2 != NULL) fclose(cfout2);
 274
 275     fastq_rng_free(rng);
 276     free(xs);
 277 }
 278
 279
 280 int main(int argc, char* argv[])
 281 {
 282     int opt;
 283     int opt_idx;
 284
 285     const char* prefix = "sample";
 286     const char* cprefix = NULL;
 287     unsigned long rng_seed = 4357;
 288     unsigned long k = 10000; // number of reads to sample
 289     double        p = -1;    // proportion of reads to sample
 290
 291     static struct option long_options[] =
 292         {
 293           {"with-replacement",  no_argument,       NULL, 'r'},
 294           {"complement-output", required_argument, NULL, 'c'},
 295           {"seed",              required_argument, NULL, 's'},
 296           {"output",            no_argument,       NULL, 'o'},
 297           {"help",              no_argument,       NULL, 'h'},
 298           {"version",           no_argument,       NULL, 'V'},
 299           {0, 0, 0, 0}
 300         };
 301
 302     replacement_flag = 0;
 303
 304     while (1) {
 305         opt = getopt_long(argc, argv, "n:p:o:rs:hV", long_options, &opt_idx);
 306
 307         if( opt == -1 ) break;
 308
 309         switch (opt) {
 310             case 0:
 311                 if (long_options[opt_idx].flag != 0) break;
 312                 if (optarg) {
 313                 }
 314                 break;
 315
 316             case 'n':
 317                 k = strtoul(optarg, NULL, 10);
 318                 break;
 319
 320             case 'p':
 321                 p = atof(optarg);
 322                 if (p < 0.0) {
 323                     fputs("Sample proportion ('-p') is less than zero.\n", stderr);
 324                     return 1;
 325                 }
 326                 break;
 327
 328             case 'r':
 329                 replacement_flag = 1;
 330                 break;
 331
 332             case 's':
 333                 rng_seed = strtoul(optarg, NULL, 10);
 334                 break;
 335
 336             case 'o':
 337                 prefix = optarg;
 338                 break;
 339
 340             case 'c':
 341                 cprefix = optarg;
 342                 break;
 343
 344             case 'h':
 345                 print_help();
 346                 return 0;
 347
 348             case 'V':
 349                 print_version(stdout, prog_name);
 350                 return 0;
 351
 352             case '?':
 353                 return 1;
 354
 355             default:
 356                 abort();
 357         }
 358     }
 359
 360     FILE* file1 = NULL;
 361     FILE* file2 = NULL;
 362
 363     if (optind >= argc) {
 364         fputs("An input file must be given.\n", stderr);
 365         print_help();
 366         return EXIT_FAILURE;
 367     }
 368
 369     file1 = fopen(argv[optind], "rb");
 370     if (file1 == NULL) {
 371         fprintf(stderr, "Cannot open '%s' for reading.\n", argv[optind]);
 372         return 1;
 373     }
 374
 375     if (++optind < argc) {
 376         file2 = fopen(argv[optind], "rb");
 377         if (file2 == NULL) {
 378             fprintf(stderr, "Cannot open '%s' for reading.\n", argv[optind]);
 379             return 1;
 380         }
 381     }
 382
 383     fastq_sample(rng_seed, prefix, cprefix, file1, file2, k, p);
 384
 385     return EXIT_SUCCESS;
 386 }
 387