def calccumsum(a):
    '''Calcaulates the cumulative-sum vectors from a 2d numpy array
    of [cov, num].  Note depends on upstream sort '''
    cn = a[:, 0]  # Coverage
    c1 = a[:, 1]  # number of distinct kmers.
    cp = cn * c1  # elementwise multiply     observed kmers by abundance
    # cumul number of distinct kmers (top to bottom)
    yd = np.flipud(np.flipud(c1).cumsum())
    # cumul number of observed kmers (top to bottom)
    yo = np.flipud(np.flipud(cp).cumsum())
    # cumul number of distinct kmers (bottom to top)
    zd = np.cumsum(c1)
    # cumul number of observed kmers (bottom to top)
    zo = np.cumsum(cp)
    if zo.max() == 0:
        raise ValueError  # There should be data here
    return(cn, c1, yd, yo, zd, zo) def loadfile(filename):
    '''Loads file, returns two-column ndarray or None on
    failure.  Uses filename to guess format.'''
    try:
        # parse velvet contig stats format
        if filename.find("stats.txt") >= 0:
            matrix = np.loadtxt(filename, usecols=(5, 1), skiprows=1)
        elif filename.find(".npo") == len(filename) - 4:
            matrix = np.loadtxt(filename, usecols=(0, 1), skiprows=6)
            L = getlength(filename)
            matrix[:, 0] = matrix[:, 0] * L
        else:  # default bare-bones spectrum format
            try:
                matrix = np.loadtxt(filename, comments="#")
            except ValueError:
                matrix = np.loadtxt(filename, skiprows=1,
                                    delimiter=",", usecols=(0, 1))
        # return None if the file is empty
        matrix[np.isinf(matrix)] = 0
        matrix = np.atleast_2d(matrix)
        if matrix.shape[1] == 0:
            return []
        else:
            return matrix
    except IOError:
        sys.stderr.write("ERROR: Can't find file %s\n" % filename)
        return [] def renyispectrum(x, spectrum):
    '''Given a two-dimensional spectrum spectrum and a one-dimensional
    vector of Renyi spectrum exponents, return a vector of corresponding
    log10 Renyi entropies.
    '''
    n = spectrum[:, 0]
    y = spectrum[:, 1]
    N = np.sum(n * y)
    p = n / N
    R = np.zeros(x.shape)
    for i, l in enumerate(x):
        G = np.log10(np.sum(y * np.power(p, l))) / (1 - l)
        R[i] = G
    return R def pad(xvalues, yvalues, fill=True):
    '''Adds missing integer values to x and corresponding zeros to y.'''
    yout = []
    xout = []
    xv = np.arange(1, int(1.5 * max(xvalues)))
    yv = np.zeros(shape=int(1.5 * max(xvalues) - 1), dtype=int)

    if fill is True:
        for i, x in enumerate(xvalues):
            yv[np.where(xv == x)[0][0]] = yvalues[i]
        xout = list(xv)
        yout = list(yv)
    else:
        for i, x in enumerate(xvalues):
            if (x - 1 not in xvalues) and (x - 1) not in xout:
                xout.append(x - 1)
                yout.append(0)
            if yvalues[xvalues.index(x)] > 0:
                xout.append(xvalues[xvalues.index(x)])
                yout.append(yvalues[xvalues.index(x)])
            if (x + 1 not in xvalues) and (x + 1) not in xout:
                xout.append(x + 1)
                yout.append(0)
    return(xout, yout) def smoothspectrum(data):
    bins = np.hstack((np.arange(1, 200), np.exp(
        np.arange(0, 4, .01) * np.log(10)) * 200))
    x = data[:, 0]
    y = data[:, 1]
    yo = np.zeros((len(bins)))
    for i in range(0, len(bins) - 1):
        yo[i] = y[np.where((x >= bins[i]) * (x < bins[i + 1]))].sum()
    xo = bins[:-1]
    yo = yo[:-1] / np.diff(bins)
    return xo, yo def calcmedian(yd, y, num):
    '''wrapper for np.interp; interpolates to return the value of yd corresponding to num on y
    sorts the data and prepends a 0,0 to get smooth behavior for entire range 0,1.'''
    ya = np.argsort(y)
    y2 = np.hstack(([0], y[ya]))
    y2d = np.hstack(([0], yd[ya]))
    r = np.interp(num, y2, y2d)
    return r def cleanlabel(label):
    '''Sanitizes graph labels of unintersting file extensions'''
    suffixes = [".histhist", ".fastq", "_info_contigstats.txt",
                ".stats.txt", ".txt", ".csv", ".037.kmerhistogram",
                "_001", ".21"]
    for suffix in suffixes:
        if label.find(suffix) > 0:
            label = label[0:(label.find(suffix))]
    if label[-3:] == ".21":
        label = label[:-3]
    return label def getmgrkmerspectrum(accessionnumber, mgrkey=None):
    '''Retrieve kmer spectrum from MG-RAST'''
    import json
    import time
    try:
        from urllib.request import urlopen
        from urllib.error import HTTPError
    except ImportError:
        from urllib2 import urlopen, HTTPError
    assert accessionnumber[0:3] == "mgm", sys.exit(
        "Data error: field %s not in mgm......... accession number format" % accessionnumber)
    some_url = "http://api.metagenomics.anl.gov/api.cgi/metagenome/%s?verbosity=full" % accessionnumber
    if mgrkey is not None:
        some_url = some_url + "&auth=%s" % mgrkey
    sys.stderr.write("Sending request for " + some_url + "\n")
    time.sleep(1)
# Ok, exception handling here is a important.  HTTP errors and
# malformed JSON are likely failure modes.
    try:
        opener = urlopen(some_url)
    except HTTPError as e:
        sys.stderr.write("Error retrieving %s\n" % some_url)
        sys.stderr.write("Error with HTTP request: %d %s\n%s\n" %
                         (e.code, e.reason, e.read()))
        return np.atleast_2d(np.array([1, 0]))
    try:
        j = json.loads(opener.read().decode("ISO-8859-1"))
    except ValueError:
        sys.stderr.write("Error parsing result from %s\n" % some_url)
        j = {}
    try:
        sys.stderr.write("Error with %s\nERROR: %s\n" % (some_url, j["ERROR"]))
        dataarray = None
    except KeyError:
        try:
            #        This is the data object containing the 15mer spectrum
            spectrum = j["statistics"]["qc"]["kmer"]["15_mer"]["data"]
            dataarray = np.array(spectrum, dtype="float")
            try:
                dataarray = dataarray[:, 0:2]
            except IndexError:
                dataarray = np.atleast_2d(np.array([1, 0]))
        except KeyError:
            dataarray = np.atleast_2d(np.array([1, 0]))
    return dataarray def printstats(a, filename, filehandle=None, n=0):
    '''Prints summary statistics to filename'''
    cn, c1, yd, yo, zd, zo = calccumsum(a)
    T = zo.max()
    j = cn / T
    # avoids zero-log warning, since products of j are used
    j[np.where(j == 0)] = 1
    intermediate = - c1 * cn / T * np.log(j)
    # allow calculation with 0 counts in some rows
    intermediate[np.isnan(intermediate)] = 0
    H = np.exp(sum(intermediate))   # antilog Shannon entropy
    if T == 0:
        H = np.NaN
    H2 = 1 / sum(c1 * j * j)        # antilog Reyni entropy / Simpson index
    w = np.array(yo, dtype="float") / yo.max()
    wd = yd
    M90 = calcmedian(wd, w, .9)      # 90th percentile by observations
    M50 = calcmedian(wd, w, .5)      # 50th percentile by observations
    M10 = calcmedian(wd, w, .1)      # 10th percentile by observations
    M100 = calcmedian(wd, w, 1.0)    # should be the same as wd.max()
    F100 = calcmedian(w, wd, 100)    # fraction of data in top 100 kmers
    F10K = calcmedian(w, wd, 10000)  # in 10K kmers
    F1M = calcmedian(w, wd, 1000000)  # in 1M kmers
    C50 = calcmedian(cn, w, .5)      # 50th percentile by observations
    AVONE = np.sum(cn * c1) / T
    assert (AVONE - 1.) < 1E-7
    AVCOV = np.sum(cn * c1 * cn) / T
    AVGEN = np.sum(cn * c1 * c1) / T
    if filehandle is None:
        consensusfh = open(filename, "w")
    else:
        consensusfh = filehandle
    if filehandle is None or n == 0:
        consensusfh.write(
            "#filename\tM10\tM50\tM90\tM100\tF100\tF10K" +
            "\tF1M\tH\tH2\tAVC\tAVG\tC50\n")
    consensusfh.write(
        "%s\t%.1f\t%.1f\t%.1f\t%d\t%f\t%f\t%f\t%.1f\t%.1f\t%.1f\t%.1f\t%.2f\n" %
        (filename, M10, M50, M90, M100, F100, F10K, F1M, H, H2, AVCOV, AVGEN, C50))
    if filehandle is None:
        consensusfh.close() def printstratify(spectrum, bands=None, flat=False, label=""):
    '''prints table of sizes and data fracitons for separate
    bands, inclusive of lower bound, excluding upper band
    boundary'''
    bands, frac, size = stratify(spectrum, bands)
    if flat is False:
        for i in range(len(bands)):
            if i != len(bands) - 1:
                print("{:.04f}\t{: 13.0f}\t{:.0f}-{:.0f}\n".format((frac[i] - frac[i + 1]), (
                    size[i] - size[i + 1]), bands[i], bands[i + 1]))
    else:
        print("#name\t" + "\t".join(map(str, list(bands[:-1] + bands[1:]))))
        print(label + "\t" + "\t".join(map(str, list(size)[:-1] +
                                           list(frac)[1:]))) def makegraphs(*args, **kwargs):
    spectra, filenames = args
    if type(spectra) is list:
        assert type(filenames) is list
        assert len(filenames) == len(spectra)
        for i, (spect, filen) in enumerate(zip(spectra, filenames)):
            makegraph(spect, filen, n=i, **kwargs)
    else:
        makegraph(*args, **kwargs) def run_indir(cmd, directory=None):
    cwd = os.getcwd()
    if directory:
        os.chdir(directory)
    else:
        os.chdir(cwd)
    call(cmd, shell=True)
    os.chdir(cwd)
    return() def show_pretty_graphs(spectra, filenames, labels=None):
    import matplotlib.pyplot as plt
    plt.subplot(131)
    makegraphs(spectra, filenames, option=6)
    plt.subplot(132)
    makegraphs(spectra, filenames, option=5)
    plt.subplot(133)
    makegraphs(spectra, filenames, option=3) TESTDIR = "tests/data/"
import os
import numpy as np
import matplotlib as mpl
from numpy import array
from nose.tools import assert_equal, assert_almost_equal, assert_true, \
    assert_false, assert_raises, assert_is_instance
from ksatools.ksatools import calccumsum, loadfile, renyispectrum, pad, \
    smoothspectrum, calcmedian, cleanlabel, getmgrkmerspectrum, printstats, \
    printstratify, makegraphs, run_indir, show_pretty_graphs, TESTDIR


def test_calccumsum_float():
    data = FIXTURE1
    cn, c1, yd, yo, zd, zo = calccumsum(data)


def test_calccumsum_int():
    data = FIXTURE3
    cn, c1, yd, yo, zd, zo = calccumsum(data)


def test_calccumsum_bad1():
    data = FIXTUREBAD1
    with assert_raises(ValueError):
        cn, c1, yd, yo, zd, zo = calccumsum(data)


def test_calccumsum_bad2():
    data = FIXTUREBAD2
    with assert_raises(ValueError):
        cn, c1, yd, yo, zd, zo = calccumsum(data)


def test_calccumsum_bad3():
    data = FIXTUREBAD3
    print(type(data), data)
    with assert_raises(ValueError):
        cn, c1, yd, yo, zd, zo = calccumsum(data)


def test_loadfile_SRR():
    FIXTURE2 = loadfile(
        "repeatresolutionpaper/counts-validationgenomedata/SRR000333.fastq.21")


def test_loadfile_empty():
    fixture = loadfile(TESTDIR + "emptyfile")
    print(type(fixture))
    assert fixture == []


def test_loadfile_npo():
    fixture = loadfile(TESTDIR + "test_1.npo")
    assert type(fixture) == np.ndarray


def test_loadfile_stats():
    fixture = loadfile(TESTDIR + "stats.txt")
    assert type(fixture) == np.ndarray


def test_loadfile_khmer():
    fixture = loadfile(TESTDIR + "test_khmer.21")
    assert type(fixture) == np.ndarray


def test_missingfile():
    fixture = loadfile("nonexistentfile")
    assert fixture == []


def test_calccumsum_00_file():
    data = loadfile(TESTDIR + "test00.21")
    cn, c1, yd, yo, zd, zo = calccumsum(data)
    assert_equal(cn, 1)
    assert_equal(c1, 1)
    assert_equal(yd, 1)
    assert_equal(yo, 1)
    assert_equal(zd, 1)
    assert_equal(zo, 1)


def test_calccumsum_01_array_int():
    data = np.array([[1, 910], [2, 45]], dtype="int")
    cn, c1, yd, yo, zd, zo = calccumsum(data)
    assert_true(np.all(cn == array([1, 2])))
    assert_true(np.all(c1 == array([910, 45])))
    assert_true(np.all(yd == array([955, 45])))
    assert_true(np.all(yo == array([1000, 90])))
    assert_true(np.all(zd == array([910, 955])))
    assert_true(np.all(zo == array([910, 1000])))


def test_calccumsum_01_array_float():
    data = np.array([[1, 910], [2, 45]], dtype="float")
    cn, c1, yd, yo, zd, zo = calccumsum(data)
    assert_true(np.all(cn == array([1, 2])))
    assert_true(np.all(c1 == array([910, 45])))
    assert_true(np.all(yd == array([955, 45])))
    assert_true(np.all(yo == array([1000, 90])))
    assert_true(np.all(zd == array([910, 955])))
    assert_true(np.all(zo == array([910, 1000])))


def test_calccumsum_int():
    data = FIXTURE1
    cn, c1, yd, yo, zd, zo = calccumsum(data)
    pass


def test_calccumsum_float():
    data = FIXTURE1
    cn, c1, yd, yo, zd, zo = calccumsum(data)
    pass


def test_reyni_returns():
    exponents = np.arange(.01, 5, .01)
    renyi_sepctrum = renyispectrum(exponents, FIXTURE1)


def test_pad():
    xpadded, ypadded = pad(FIXTURE1[:, 0], FIXTURE1[:, 1])
    assert_true(len(xpadded) > max(FIXTURE1[:, 0]) * 1.2)
    assert_true(len(ypadded) > max(FIXTURE1[:, 0]) * 1.2)
    assert_true(sum(ypadded) == sum(FIXTURE1[:, 1]))


def test_smooth():

    s = smoothspectrum(FIXTURE1)