<commit_before>#! /usr/bin/env python2
#
# This file is part of khmer, https://github.com/dib-lab/khmer/, and is
# Copyright (C) Michigan State University, 2009-2015. It is licensed under
# the three-clause BSD license; see LICENSE.
# Contact: <EMAIL>
#
# pylint: disable=invalid-name,missing-docstring
"""
Eliminate surplus reads.

Eliminate reads with median k-mer abundance higher than
DESIRED_COVERAGE.  Output sequences will be placed in 'infile.keep', with the
option to output to STDOUT.

% python scripts/normalize-by-median.py [ -C <cutoff> ] <data1> <data2> ...

Use '-h' for parameter help.
"""

import sys
import screed
import os
import khmer
import textwrap
from itertools import izip
from contextlib import contextmanager

from khmer.khmer_args import (build_counting_args, add_loadhash_args,
                              report_on_config, info)
import argparse
from khmer.kfile import (check_space, check_space_for_hashtable,
                         check_valid_file_exists)
from khmer.utils import write_record, check_is_pair
DEFAULT_DESIRED_COVERAGE = 10

# Iterate a collection in arbitrary batches
# from: http://stackoverflow.com/questions/4628290/pairs-from-single-list


def batchwise(coll, size):
    iter_coll = iter(coll)
    return izip(*[iter_coll] * size)

# Returns true if the pair of records are properly pairs


def WithDiagnostics(ifile, batch_size, fp, paired):
    index = 0
    global total, discarded

    for batch in batchwise(screed.open(ifile, parse_description=False),
                           batch_size):

        if index > 0 and index % 100000 == 0:
            print >>sys.stderr, '... kept {kept} of {total} or'\
                ' {perc:2}%'.format(kept=total - discarded, total=total,
                                    perc=int(100. - discarded /
                                             float(total) * 100.))
            print >>sys.stderr, '... in file', input_filename

            if report_fp:
                print >> fp, total, total - discarded, \
                    1. - (discarded / float(total))
                report_fp.flush()

        total += batch_size

        # If in paired mode, check that the reads are properly interleaved

        if paired:
            if not check_is_pair(batch[0], batch[1]):
                raise IOError('Error: Improperly interleaved pairs \
                    {b0} {b1}'.format(b0=batch[0].name, b1=batch[1].name))

        yield batch

        index = index + 1


# pylint: disable=too-many-locals,too-many-branches
def normalize_by_median(input_filename, outfp, htable, paired, cutoff,
                        report_fp=None):

    desired_coverage = cutoff
    ksize = htable.ksize()

    # In paired mode we read two records at a time
    batch_size = 1
    if paired:
        batch_size = 2

    index = -1
    # global some things to work with our iterator
    global total, discarded
    total = 0
    discarded = 0
    for batch in WithDiagnostics(input_filename, batch_size, report_fp,
                                 paired):

        # Emit the batch of reads if any read passes the filter
        # and all reads are longer than K
        passed_filter = False
        passed_length = True
        for record in batch:
            if len(record.sequence) < ksize:
                passed_length = False
                continue

            seq = record.sequence.replace('N', 'A')
            med, _, _ = htable.get_median_count(seq)

            if med < desired_coverage:
                htable.consume(seq)
                passed_filter = True

        # Emit records if any passed
        if passed_length and passed_filter:
            for record in batch:
                write_record(record, outfp)
        else:
            discarded += batch_size

    if report_fp:
        print >> report_fp, total, total - discarded, \
            1. - (discarded / float(total))
        report_fp.flush()

    return total, discarded


def handle_error(error, output_name, input_name, fail_save, htable):
    print >> sys.stderr, '** ERROR:', error
    print >> sys.stderr, '** Failed on {name}: '.format(name=input_name)
    if fail_save:
        tablename = os.path.basename(input_name) + '.ct.failed'
        print >> sys.stderr, \
            '** ...dumping k-mer counting table to {tn}'.format(tn=tablename)
        htable.save(tablename)
    try:
        os.remove(output_name)
    except:  # pylint: disable=bare-except
        print >> sys.stderr, '** ERROR: problem removing corrupt filtered file'


@contextmanager
def FailSafe(ifile, ofile, save_on_fail, ht, total, discarded, force):
    global corrupt_files, total, discarded, total_acc, discarded_acc
    try:
        yield
    except IOError as err:
        handle_error(err, ofile, ifile, save_on_fail,
                     ht)
        if not force:
            print >> sys.stderr, '** Exiting!'

            sys.exit(1)
        else:
            print >> sys.stderr, '*** Skipping error file, moving on...'
            corrupt_files.append(ifile)


def normalize_by_median_and_check(input_filename, htable, single_output_file,
                                  fail_save, paired, cutoff, force,
                                  report_fp=None):
    global corrupt_files, total, discarded, total_acc, discarded_acc
    total = 0
    discarded = 0

    total_acc = None
    discarded_acc = None

    if single_output_file:
        if single_output_file is sys.stdout:
            output_name = '/dev/stdout'
        else:
            output_name = single_output_file.name
        outfp = single_output_file

    else:
        output_name = os.path.basename(input_filename) + '.keep'
        outfp = open(output_name, 'w')

    with FailSafe(input_filename, outfp, fail_save, htable, total, discarded,
                  force):

        total_acc, discarded_acc = normalize_by_median(
            input_filename, outfp, htable, paired, cutoff, report_fp)

        if total_acc == 0 and discarded_acc == 0:
            print >> sys.stderr, 'SKIPPED empty file', input_filename
        elif total_acc is not None and discarded_acc is not None:
            # will be none if the context manager caught an error
            total += total_acc
            discarded += discarded_acc
            print >> sys.stderr, \
                'DONE with {inp}; kept {kept} of {total} or {perc:2}%'\
                .format(inp=input_filename, kept=total - discarded,
                        total=total, perc=int(100. - discarded /
                                              float(total) * 100.))
            print >> sys.stderr, 'output in', output_name

    return total_acc, discarded_acc, corrupt_files


def get_parser():
    epilog = ("""
    Discard sequences based on whether or not their median k-mer abundance lies
    above a specified cutoff. Kept sequences will be placed in <fileN>.keep.

    Paired end reads will be considered together if :option:`-p` is set. If
    either read will be kept, then both will be kept. This should result in
    keeping (or discarding) each sequencing fragment. This helps with retention
    of repeats, especially. With :option: `-u`/:option:`--unpaired-reads`, 
    unpaired reads from the specified file will be read after the paired data
    is read. 

    With :option:`-s`/:option:`--savetable`, the k-mer counting table
    will be saved to the specified file after all sequences have been
    processed. With :option:`-d`, the k-mer counting table will be
    saved every d files for multifile runs; if :option:`-s` is set,
    the specified name will be used, and if not, the name `backup.ct`
    will be used.  :option:`-l`/:option:`--loadtable` will load the
    specified k-mer counting table before processing the specified
    files.  Note that these tables are are in the same format as those<commit_msg>the generators strike back!<commit_after>#! /usr/bin/env python2
#
# This file is part of khmer, https://github.com/dib-lab/khmer/, and is
# Copyright (C) Michigan State University, 2009-2015. It is licensed under
# the three-clause BSD license; see LICENSE.
# Contact: <EMAIL>
#
# pylint: disable=invalid-name,missing-docstring
"""
Eliminate surplus reads.

Eliminate reads with median k-mer abundance higher than
DESIRED_COVERAGE.  Output sequences will be placed in 'infile.keep', with the
option to output to STDOUT.

% python scripts/normalize-by-median.py [ -C <cutoff> ] <data1> <data2> ...

Use '-h' for parameter help.
"""

import sys
import screed
import os
import khmer
import textwrap
from itertools import izip
from contextlib import contextmanager

from khmer.khmer_args import (build_counting_args, add_loadhash_args,
                              report_on_config, info)
import argparse
from khmer.kfile import (check_space, check_space_for_hashtable,
                         check_valid_file_exists)
from khmer.utils import write_record, check_is_pair
DEFAULT_DESIRED_COVERAGE = 10

# Iterate a collection in arbitrary batches
# from: http://stackoverflow.com/questions/4628290/pairs-from-single-list


def batchwise(coll, size):
    iter_coll = iter(coll)
    return izip(*[iter_coll] * size)

# Returns true if the pair of records are properly pairs


def WithDiagnostics(ifile, batch_size, fp, paired):
    index = 0
    global total, discarded

    for batch in batchwise(screed.open(ifile, parse_description=False),
                           batch_size):

        if index > 0 and index % 100000 == 0:
            print >>sys.stderr, '... kept {kept} of {total} or'\
                ' {perc:2}%'.format(kept=total - discarded, total=total,
                                    perc=int(100. - discarded /
                                             float(total) * 100.))
            print >>sys.stderr, '... in file', input_filename

            if report_fp:
                print >> fp, total, total - discarded, \
                    1. - (discarded / float(total))
                report_fp.flush()

        total += batch_size

        # If in paired mode, check that the reads are properly interleaved

        if paired:
            if not check_is_pair(batch[0], batch[1]):
                raise IOError('Error: Improperly interleaved pairs \
                    {b0} {b1}'.format(b0=batch[0].name, b1=batch[1].name))

        yield batch

        index = index + 1


# pylint: disable=too-many-locals,too-many-branches
def normalize_by_median(input_filename, outfp, htable, paired, cutoff,
                        report_fp=None):

    desired_coverage = cutoff
    ksize = htable.ksize()

    # In paired mode we read two records at a time
    batch_size = 1
    if paired:
        batch_size = 2

    index = -1
    # global some things to work with our iterator
    global total, discarded
    total = 0
    discarded = 0
    for batch in WithDiagnostics(input_filename, batch_size, report_fp,
                                 paired):

        # Emit the batch of reads if any read passes the filter
        # and all reads are longer than K
        passed_filter = False
        passed_length = True
        for record in batch:
            if len(record.sequence) < ksize:
                passed_length = False
                continue

            seq = record.sequence.replace('N', 'A')
            med, _, _ = htable.get_median_count(seq)

            if med < desired_coverage:
                htable.consume(seq)
                passed_filter = True

        # Emit records if any passed
        if passed_length and passed_filter:
            for record in batch:
                write_record(record, outfp)
        else:
            discarded += batch_size

    if report_fp:
        print >> report_fp, total, total - discarded, \
            1. - (discarded / float(total))
        report_fp.flush()

    return total, discarded


def handle_error(error, output_name, input_name, fail_save, htable):
    print >> sys.stderr, '** ERROR:', error
    print >> sys.stderr, '** Failed on {name}: '.format(name=input_name)
    if fail_save:
        tablename = os.path.basename(input_name) + '.ct.failed'
        print >> sys.stderr, \
            '** ...dumping k-mer counting table to {tn}'.format(tn=tablename)
        htable.save(tablename)
    try:
        os.remove(output_name)
    except:  # pylint: disable=bare-except
        print >> sys.stderr, '** ERROR: problem removing corrupt filtered file'


@contextmanager
def FailSafe(ifile, ofile, save_on_fail, ht, force):
    global corrupt_files, total, discarded, total_acc, discarded_acc
    try:
        yield
    except IOError as err:
        handle_error(err, ofile, ifile, save_on_fail,
                     ht)
        if not force:
            print >> sys.stderr, '** Exiting!'

            sys.exit(1)
        else:
            print >> sys.stderr, '*** Skipping error file, moving on...'
            corrupt_files.append(ifile)

    if total_acc == 0 and discarded_acc == 0:
        print >> sys.stderr, 'SKIPPED empty file', ifile
    elif total_acc is not None and discarded_acc is not None:
        # will be none if the context manager caught an error
        total += total_acc
        discarded += discarded_acc
        print >> sys.stderr, \
            'DONE with {inp}; kept {kept} of {total} or {perc:2}%'\
            .format(inp=ifile, kept=total - discarded,
                    total=total, perc=int(100. - discarded /
                                          float(total) * 100.))
        print >> sys.stderr, 'output in', ofile


def normalize_by_median_and_check(input_filename, htable, single_output_file,
                                  fail_save, paired, cutoff, force,
                                  report_fp=None):
    global corrupt_files, total, discarded, total_acc, discarded_acc
    total = 0
    discarded = 0

    total_acc = None
    discarded_acc = None

    if single_output_file:
        if single_output_file is sys.stdout:
            output_name = '/dev/stdout'
        else:
            output_name = single_output_file.name
        outfp = single_output_file

    else:
        output_name = os.path.basename(input_filename) + '.keep'
        outfp = open(output_name, 'w')

    with FailSafe(input_filename, outfp, fail_save, htable, force):

        total_acc, discarded_acc = normalize_by_median(
            input_filename, outfp, htable, paired, cutoff, report_fp)

    return total_acc, discarded_acc, corrupt_files


def get_parser():
    epilog = ("""
    Discard sequences based on whether or not their median k-mer abundance lies
    above a specified cutoff. Kept sequences will be placed in <fileN>.keep.

    Paired end reads will be considered together if :option:`-p` is set. If
    either read will be kept, then both will be kept. This should result in
    keeping (or discarding) each sequencing fragment. This helps with retention
    of repeats, especially. With :option: `-u`/:option:`--unpaired-reads`, 
    unpaired reads from the specified file will be read after the paired data
    is read. 

    With :option:`-s`/:option:`--savetable`, the k-mer counting table
    will be saved to the specified file after all sequences have been
    processed. With :option:`-d`, the k-mer counting table will be
    saved every d files for multifile runs; if :option:`-s` is set,
    the specified name will be used, and if not, the name `backup.ct`
    will be used.  :option:`-l`/:option:`--loadtable` will load the
    specified k-mer counting table before processing the specified
    files.  Note that these tables are are in the same format as those
{'fraction_non_alphanumeric': 0.08950141351837573, 'fraction_numerical': 0.008224106913389874, 'mean_word_length': 2.2386600083229298, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 67, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2338149', 'n_tokens_mistral': 4682, 'n_tokens_neox': 4355, 'n_words': 1434}