cnrand.py

from scipy.stats import chisquare
import sys

def cnrand(counts):
    """
    cnrand will return the chisquare p-value of the observed distribution
    of copy-number genomes **of the same CN with  0...n copies of H47R

    # test 4-copy genomes (0 ... 4)
    # 1 0 copy, 1 3-copy and 1 4-copy
    >>> cnrand([1, 0, 0, 1, 1])
    >>> cnrand([0, 10, 0, 0, 0])

    """
    total_alleles = sum(a * b for a, b in enumerate(counts))
    #sys.stderr.write("total_alleles: %d\n" % total_alleles)
    total_genomes = (len(counts) - 1) * sum(counts)
    #sys.stderr.write("total_genomes: %d\n" % total_genomes)

    af = total_alleles / float(total_genomes)
    m = sum(counts) #/ float(len(counts) - 1)

    expected = [1-af] + [af**i for i in range(1, len(counts))]
    #sys.stderr.write("%s\n" % expected)
    expected = [v * m for v in expected]

    sys.stderr.write("%s\n" % expected)
    assert len(counts) == len(expected)

    return af, expected, chisquare(counts, expected)

if __name__ == "__main__":
    import doctest
    doctest.testmod()