<commit_before>#! /usr/bin/python

import screed
import sys
import argparse
import time
import khmer
from khmer.counting_args import build_construct_args, DEFAULT_MIN_HASHSIZE

# little class to store sequence information for the buffering class
class Seq:
    def __init__(self, name, color, seq):
        self.name = name
        self.color = color
        self.seq = seq
    def write(self, fp):
        fp.write('>{}\t{}\n{}\n'.format(self.name, self.color, self.seq))

# stores reads in memory and flushes them to their approriate files
# when certain criteria are met
# Basic idea is to buffer some number of reads in memory, then dump them all at once
# Hope that each file acrues, on average, BUFFER_SIZE / NUM_PARTS reads
# ie, if we buffer 1000000 reads, and we have 100000 partitions or colors,
# we should expect the mean buffer size to be 10 reads
class ReadBuffer:

    def __init__(self, max_buffers=10000, max_reads=1000000, est_files=100000, output_pref='reads_'):
        self.buffers = {}
        self.buffer_counts = {}
        self.max_buffers = max_buffers
        self.max_reads = max_reads

        self.est_files = est_files
        self.output_pref = output_pref
        self.buffer_flush = self.max_reads / self.est_files

        self.cur_reads = 0
        self.cur_files = 0

    def add_seq(self, seq):
        color = seq.color
        if color in self.buffers:
            count = self.buffer_counts[color]
            self.buffers[color].append(seq)
            self.buffer_counts[color] += 1
            if count > self.buffer_flush:
                self.flush_buffer(color)
                self.del_buffer(color)

        else:
            self.buffers[color] = [seq]
            self.buffer_counts[color] = 1
        self.cur_reads += 1
        if self.cur_reads > self.max_reads:
            self.flush_all()
        if len(self.buffers) > self.max_buffers:
            #self.clean_buffers(2)
            self.flush_all()
    
    def flush_buffer(self, color):
        with open('{}_{}.fa'.format(self.output_pref, color), 'a') as outfp:
            for read in self.buffers[color]:
                read.write(outfp)
                self.cur_reads -= 1
            
    def del_buffer(self, color):
        del self.buffer_counts[color]
        del self.buffers[color]

    def flush_all(self):
        print >>sys.stderr, '** reached max buffer size, flushing all to files...'
        for color in self.buffers:
            self.flush_buffer(color)
        colors = self.buffers.keys()
        for color in colors:
            self.del_buffer(color)
        del colors
        assert self.cur_reads == 0

    def clean_buffers(self, cutoff):
        print >>sys.stderr, '** flushing low-abundance buffers...'
        flushed = []
        for color in self.buffers:
            if self.buffer_counts[color] < cutoff:
                self.flush_buffer(color)
                flushed.append(color)
        for color in flushed:
            self.del_buffer(color)
        del flushed

def main():

    parser = build_construct_args()
    parser.add_argument('-i', '--input_fastp',dest='input_fastp')
    parser.add_argument('-r', '--traversal_range', type=int, dest='traversal_range')
    parser.add_argument('-b', '--buffer_size', dest='buffer_size', type=int)
    parser.add_argument('-e', '--files_estimate', dest='files_estimate', type=int)
    parser.add_argument('-o', '--output_prefix', dest='output_prefix')
    parser.add_argument('-m', '--max_buffers', dest='max_buffers', type=int)
    parser.add_argument('input_files', nargs='+')
    args = parser.parse_args()
    
    if not args.quiet:
        if args.min_hashsize == DEFAULT_MIN_HASHSIZE:
            print >>sys.stderr, \
                "** WARNING: hashsize is default!  " \
                "You absodefly want to increase this!\n** " \
                "Please read the docs!"

        print >>sys.stderr, '\nPARAMETERS:'
        print >>sys.stderr, \
            ' - kmer size =    {ksize:d} \t\t(-k)'.format(ksize=args.ksize)
        print >>sys.stderr, \
            ' - n hashes =     {nhash:d} \t\t(-N)'.format(nhash=args.n_hashes)
        print >>sys.stderr, \
            ' - min hashsize = {mh:-5.2g} \t(-x)'.format(mh=args.min_hashsize)
        print >>sys.stderr, ''<commit_msg>added default parameters to buffered sweep<commit_after>#! /usr/bin/python
#
# This file is part of khmer, http://github.com/ged-lab/khmer/, and is
# Copyright (C) Michigan State University, 2009-2013. It is licensed under
# the three-clause BSD license; see doc/LICENSE.txt. Contact: <EMAIL>
#

"""
Find all reads connected to the given contigs on a per-partition basis.

% python scripts/normalize-by-median.py -r <range> -i <contigs fastp> \
<reads1> <reads2> ... <readsN>
"""

import screed
import sys
import argparse
import time
import khmer
from khmer.counting_args import build_construct_args, DEFAULT_MIN_HASHSIZE

DEFAULT_NUM_BUFFERS=50000
DEFAULT_BUFFER_SIZE=1000000
DEFAULT_NUM_PARTITIONS=100000
DEFAULT_OUT_PREF='reads_'

# little class to store sequence information for the buffering class
class Seq:
    def __init__(self, name, color, seq):
        self.name = name
        self.color = color
        self.seq = seq
    def write(self, fp):
        fp.write('>{}\t{}\n{}\n'.format(self.name, self.color, self.seq))

# stores reads in memory and flushes them to their approriate files
# when certain criteria are met
# Basic idea is to buffer some number of reads in memory, then dump them all at once
# Hope that each file acrues, on average, BUFFER_SIZE / NUM_PARTS reads
# ie, if we buffer 1000000 reads, and we have 100000 partitions or colors,
# we should expect the mean buffer size to be 10 reads
class ReadBuffer:

    def __init__(self, max_buffers, max_reads, est_files, output_pref):
        self.buffers = {}
        self.buffer_counts = {}
        self.max_buffers = max_buffers
        self.max_reads = max_reads

        self.est_files = est_files
        self.output_pref = output_pref
        self.buffer_flush = self.max_reads / self.est_files

        self.cur_reads = 0
        self.cur_files = 0

    def add_seq(self, seq):
        color = seq.color
        if color in self.buffers:
            count = self.buffer_counts[color]
            self.buffers[color].append(seq)
            self.buffer_counts[color] += 1
            if count > self.buffer_flush:
                self.flush_buffer(color)
                self.del_buffer(color)

        else:
            self.buffers[color] = [seq]
            self.buffer_counts[color] = 1
        self.cur_reads += 1
        if self.cur_reads > self.max_reads:
            self.flush_all()
        if len(self.buffers) > self.max_buffers:
            #self.clean_buffers(2)
            self.flush_all()
    
    def flush_buffer(self, color):
        with open('{}_{}.fa'.format(self.output_pref, color), 'a') as outfp:
            for read in self.buffers[color]:
                read.write(outfp)
                self.cur_reads -= 1
            
    def del_buffer(self, color):
        del self.buffer_counts[color]
        del self.buffers[color]

    def flush_all(self):
        print >>sys.stderr, '** reached max buffer size, flushing all to files...'
        for color in self.buffers:
            self.flush_buffer(color)
        colors = self.buffers.keys()
        for color in colors:
            self.del_buffer(color)
        del colors
        assert self.cur_reads == 0

    def clean_buffers(self, cutoff):
        print >>sys.stderr, '** flushing low-abundance buffers...'
        flushed = []
        for color in self.buffers:
            if self.buffer_counts[color] < cutoff:
                self.flush_buffer(color)
                flushed.append(color)
        for color in flushed:
            self.del_buffer(color)
        del flushed

def main():

    parser = build_construct_args()
    parser.add_argument('-i', '--input_fastp',dest='input_fastp')
    parser.add_argument('-r', '--traversal_range', type=int, dest='traversal_range')
    parser.add_argument('-b', '--buffer_size', dest='buffer_size', type=int, \
                        default=DEFAULT_BUFFER_SIZE)
    parser.add_argument('-e', '--files_estimate', dest='files_estimate', type=int, \
                        default=DEFAULT_NUM_PARTITIONS)
    parser.add_argument('-o', '--output_prefix', dest='output_prefix',
                        default=DEFAULT_OUT_PREF)
    parser.add_argument('-m', '--max_buffers', dest='max_buffers', type=int, \
                        default=DEFAULT_NUM_BUFFERS)
    parser.add_argument('input_files', nargs='+')
    args = parser.parse_args()
    
    if not args.quiet:
        if args.min_hashsize == DEFAULT_MIN_HASHSIZE:
            print >>sys.stderr, \
                "** WARNING: hashsize is default!  " \
                "You absodefly want to increase this!\n** " \
                "Please read the docs!"

        print >>sys.stderr, '\nPARAMETERS:'
        print >>sys.stderr, \
            ' - kmer size =    {ksize:d} \t\t(-k)'.format(ksize=args.ksize)
        print >>sys.stderr, \
            ' - n hashes =     {nhash:d} \t\t(-N)'.format(nhash=args.n_hashes)
        print >>sys.stderr, \
            ' - min hashsize = {mh:-5.2g} \t(-x)'.format(mh=args.min_hashsize)
        print >>sys.stderr, ''
{'fraction_non_alphanumeric': 0.11589226311836034, 'fraction_numerical': 0.010301534499409808, 'mean_word_length': 2.3014523556500177, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1378', 'n_tokens_mistral': 3028, 'n_tokens_neox': 2853, 'n_words': 710}