<commit_before>from numpy import histogram
import os, sys, numpy, argparse, logging, re, pandas
import edl.blastm8
from Bio import SeqIO
from edl.util import asciiHistogram
logger=logging.getLogger(__name__)

## 
# the calc_stat, plot_assembly, and getN50 methods originally come from the
# assemlbly_quality_stats.py script by:
# Author: <NAME>
# Date: 09 Feb. 2013
# http://travispoulsen.com/blog/2013/07/basic-assembly-statistics/
# https://gist.github.com/tpoulsen/422b1a19cbd8c0f514fe/raw/assembly_quality_stats.py
def calc_stats(file_in, return_data=False, txt_width=0, log=False, backend=None, format='fasta', minLength=0, **kwargs):
    """
    Given contigs in fastsa format:
     * calculate length stats (including N50)
     * plot histogram (use txt_width and backend to select format)
    """
    with open(file_in, 'r') as seq:
        sizes = [len(record) for record in SeqIO.parse(seq, format) if len(record)>=minLength]

    sizes = numpy.array(sizes)
    data = get_contig_length_stats(sizes)

    if not return_data:
        report = get_contig_length_report(data)

    if backend is not None:
        h = plot_assembly(sizes, file_in, data, log=log, backend=backend, **kwargs)
    if txt_width>0:
        if backend is None:
            h=histogram(sizes, **kwargs)
        histogramText = asciiHistogram(h, log=log, width=txt_width)
        if not return_data:
            if log:
                report += "\n\nContig length histogram (log):\n"
            else:
                report += "\n\nContig length histogram:\n"
            report += histogramText
        else:
            data['histogram']=histogramText

    if return_data:
        return data
    else:
        return report

def get_contig_length_stats(sizes):
    """
    return a dict of useful contig length stats
    """
    return {'min':numpy.min(sizes),
            'max':numpy.max(sizes),
            'mean':numpy.mean(sizes),
            'median':numpy.median(sizes),
            'N50':int(getN50(sizes)),
            'N75':int(getN50(sizes,N=75)),
            'N90':int(getN50(sizes,N=90)),
            'count':len(sizes),
            }

def get_contig_length_report(data):
    """
    return a formatted string summarizing contig length data
    """
    report ='Number of contigs:\t%i' % data['count']
    report += '\nN50:\t%i' % data['N50']
    report += '\nN75:\t%i' % data['N75']
    report += '\nN90:\t%i' % data['N90']
    report += '\nMean contig length:\t%.2f' % data['mean']
    report += '\nMedian contig length:\t%.2f' % data['median']
    report += '\nMinimum contig length:\t%i' % data['min']
    report += '\nMaximum contig length:\t%i' % data['max']
    return report

def mira_stats(contigStatsFile,minLength=0,bins=20,**kwargs):
    """
    Get length, coverage, and GC stats from mira info file
    Returns text with N50 and histograms
    """
    contigStats=pandas.read_csv(contigStatsFile,index_col=0,sep='\t')
    if minLength>0:
        contigStats=contigStats[contigStats.length>=minLength]
    sizes = contigStats['length']
    data = get_contig_length_stats(sizes)
    report = get_contig_length_report(data)

    # add histograms to report
    report += '\nHistograms:\n'
    for key in ['length','GC%','av.cov','mx.cov.','av.qual']:
        report += '\n'
        report += edl.util.asciiHistogram(histogram(contigStats[key],bins=bins),label=key,**kwargs)
    
    return report

def plot_assembly(sizes, file_in, length_data, backend=None,**kwargs):
    min_contig = length_data['min']
    max_contig = length_data['max']
    avg_contig = length_data['mean']
    num_contig = length_data['count']

    if backend:
        import matplotlib
        matplotlib.use(backend)
    from matplotlib import pyplot as plt
    #import pylab
    h=plt.hist(sizes, **kwargs)
    plt.title('%i %s sequences\nLengths %i to %i, Average contig length: %.2f' % (num_contig, file_in, min_contig, max_contig, avg_contig))
    plt.xlabel('Sequence length (bp)')
    plt.ylabel('Count')
    return h

# Get the N50 of the contigs. This is the sequence length at which point
# half of the bases in the entire assembly are contained in contigs of a
# smaller size.
def getN50_old(sizes):
    bases = []
    for read in sorted(sizes):
        for i in range(read):
            bases.append(read)
    return numpy.median(bases)
def getN50(sizes, N=50):
    if N==50:
        compare=getN50_old(sizes)

    totalLength = sum(sizes)
    targetLength = float(totalLength) * N/100.
    totalLength = 0
    for size in sorted(sizes,reverse=True):
        totalLength+=size
        if totalLength>=targetLength:
            if N==50:
                if size != compare:
                    raise Exception("N50s don't match: %s != %s" % (size,compare))
            return size
    else:
        raise Exception("Target length never reached!\nN=%d, target=%d, total=%d" % (N, targetLength, totalLength))

######
# A set of methods for plotting the quality of reseq hits to a set of contigs<commit_msg>removing comparison to old N50 method. I've seen no errors in the new one and teh old onee is failing on one really big assembly. Issue #8<commit_after>from numpy import histogram
import os, sys, numpy, argparse, logging, re, pandas
import edl.blastm8
from Bio import SeqIO
from edl.util import asciiHistogram
logger=logging.getLogger(__name__)

## 
# the calc_stat, plot_assembly, and getN50 methods originally come from the
# assemlbly_quality_stats.py script by:
# Author: <NAME>
# Date: 09 Feb. 2013
# http://travispoulsen.com/blog/2013/07/basic-assembly-statistics/
# https://gist.github.com/tpoulsen/422b1a19cbd8c0f514fe/raw/assembly_quality_stats.py
def calc_stats(file_in, return_data=False, txt_width=0, log=False, backend=None, format='fasta', minLength=0, **kwargs):
    """
    Given contigs in fastsa format:
     * calculate length stats (including N50)
     * plot histogram (use txt_width and backend to select format)
    """
    with open(file_in, 'r') as seq:
        sizes = [len(record) for record in SeqIO.parse(seq, format) if len(record)>=minLength]

    sizes = numpy.array(sizes)
    data = get_contig_length_stats(sizes)

    if not return_data:
        report = get_contig_length_report(data)

    if backend is not None:
        h = plot_assembly(sizes, file_in, data, log=log, backend=backend, **kwargs)
    if txt_width>0:
        if backend is None:
            h=histogram(sizes, **kwargs)
        histogramText = asciiHistogram(h, log=log, width=txt_width)
        if not return_data:
            if log:
                report += "\n\nContig length histogram (log):\n"
            else:
                report += "\n\nContig length histogram:\n"
            report += histogramText
        else:
            data['histogram']=histogramText

    if return_data:
        return data
    else:
        return report

def get_contig_length_stats(sizes):
    """
    return a dict of useful contig length stats
    """
    return {'min':numpy.min(sizes),
            'max':numpy.max(sizes),
            'mean':numpy.mean(sizes),
            'median':numpy.median(sizes),
            'N50':int(getN50(sizes)),
            'N75':int(getN50(sizes,N=75)),
            'N90':int(getN50(sizes,N=90)),
            'count':len(sizes),
            }

def get_contig_length_report(data):
    """
    return a formatted string summarizing contig length data
    """
    report ='Number of contigs:\t%i' % data['count']
    report += '\nN50:\t%i' % data['N50']
    report += '\nN75:\t%i' % data['N75']
    report += '\nN90:\t%i' % data['N90']
    report += '\nMean contig length:\t%.2f' % data['mean']
    report += '\nMedian contig length:\t%.2f' % data['median']
    report += '\nMinimum contig length:\t%i' % data['min']
    report += '\nMaximum contig length:\t%i' % data['max']
    return report

def mira_stats(contigStatsFile,minLength=0,bins=20,**kwargs):
    """
    Get length, coverage, and GC stats from mira info file
    Returns text with N50 and histograms
    """
    contigStats=pandas.read_csv(contigStatsFile,index_col=0,sep='\t')
    if minLength>0:
        contigStats=contigStats[contigStats.length>=minLength]
    sizes = contigStats['length']
    data = get_contig_length_stats(sizes)
    report = get_contig_length_report(data)

    # add histograms to report
    report += '\nHistograms:\n'
    for key in ['length','GC%','av.cov','mx.cov.','av.qual']:
        report += '\n'
        report += edl.util.asciiHistogram(histogram(contigStats[key],bins=bins),label=key,**kwargs)
    
    return report

def plot_assembly(sizes, file_in, length_data, backend=None,**kwargs):
    min_contig = length_data['min']
    max_contig = length_data['max']
    avg_contig = length_data['mean']
    num_contig = length_data['count']

    if backend:
        import matplotlib
        matplotlib.use(backend)
    from matplotlib import pyplot as plt
    #import pylab
    h=plt.hist(sizes, **kwargs)
    plt.title('%i %s sequences\nLengths %i to %i, Average contig length: %.2f' % (num_contig, file_in, min_contig, max_contig, avg_contig))
    plt.xlabel('Sequence length (bp)')
    plt.ylabel('Count')
    return h

# Get the N50 of the contigs. This is the sequence length at which point
# half of the bases in the entire assembly are contained in contigs of a
# smaller size.
def getN50_old(sizes):
    bases = []
    for read in sorted(sizes):
        for i in range(read):
            bases.append(read)
    return numpy.median(bases)

def getN50(sizes, N=50):

    #if N==50:
    #    # This adds a significant chunk of time and has yet to find an error...should we drop it?
    #    try:
    #        compare=getN50_old(sizes)
    #    except:
    #        compare=None

    totalLength = sum(sizes)
    targetLength = float(totalLength) * N/100.
    totalLength = 0
    for size in sorted(sizes,reverse=True):
        totalLength+=size
        if totalLength>=targetLength:
            #if N==50 and compare is not None:
            #    if size != compare:
            #        raise Exception("N50s don't match: %s != %s" % (size,compare))
            return size
    else:
        raise Exception("Target length never reached!\nN=%d, target=%d, total=%d" % (N, targetLength, totalLength))

######
# A set of methods for plotting the quality of reseq hits to a set of contigs
{'fraction_non_alphanumeric': 0.12580019398642095, 'fraction_numerical': 0.017555771096023277, 'mean_word_length': 2.993415956622773, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 15, 'https://': 2, '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': '1016285', 'n_tokens_mistral': 3532, 'n_tokens_neox': 3346, 'n_words': 866}