Python rc stuff

rc.py

#!/usr/bin/env python
from __future__ import division, generators, nested_scopes, with_statement
import atexit, os, readline, rlcompleter

readline.parse_and_bind("tab: complete")
readline.set_history_length(10000)

histfile = os.path.expanduser("~/.python/history")
try:
   readline.read_history_file(histfile)
except IOError:
   pass

def isint(n):
    return type(n) in (int, long)

try:
    # python2 only
    from compiler.ast import flatten

except ImportError:
    # python3 stuff
    from functools import reduce

    long = int     # py3 has no long type

    def flatten(*iterable):
        """hopefully equivalent to python2's compiler.ast.flatten"""
        # thanks http://stackoverflow.com/a/16176779
        import collections
        from itertools import chain
        for el in chain(*iterable):
            if isinstance(el, collections.Iterable) and not isinstance(el, str):
                for subel in flatten(el):
                    yield subel
            else:
                yield el

def gamma(z, x=1.000000000000000174663, p=[5716.400188274341379136,
    -14815.30426768413909044, 14291.49277657478554025, -6348.160217641458813289,
    1301.608286058321874105, -108.1767053514369634679, 2.605696505611755827729,
    -0.7423452510201416151527e-2, 0.5384136432509564062961e-7,
    -0.4023533141268236372067e-8]):
    """Lanczos approximation good to approximately 15 digits of precision.
    Supports complex input and output."""
    epsilon = 1e-10
    def withinepsilon(x):
        return abs(x - abs(x)) <= epsilon

    from cmath import sin, sqrt, pi, exp

    z = complex(z)

    # Reflection formula
    if z.real < 0.5:
        result = pi / (sin(pi * z) * gamma(1 - z, x=x, p=p))
    else:
        z -= 1

        for (i, pval) in enumerate(p):
            x += pval / (z + i + 1)

        t = z + len(p) - 0.5
        result = sqrt(2 * pi) * t**(z + 0.5) * exp(-t) * x

    if withinepsilon(result.imag):
        return result.real
    return result


def binary(n, prefix='0b', length=0, spacer=0):
    "drop-in replacement for bin() builtin but accepts additional parameters"
    return '{0}{{:{1}>{2}}}'.format(prefix, spacer, length).format(bin(n)[2:])


def gcd(*numbers):
    try:
        from math import gcd                # py3
    except ImportError:
        try:
            from fractions import gcd       # requires 2.6+
        except ImportError:                 # for older Python courtesy <http://stackoverflow.com/questions/7337807/help-me-find-mistake-in-greatest-common-divisor-algorithm-in-python>
            def gcd(a, b):
                if b != 0:
                    return gcd(b, a % b)
                return a

    return reduce(gcd, flatten(numbers))


def lcm(*numbers):
    def _lcm(a,b):
        return (a * b) // gcd(a, b)
    return reduce(_lcm, flatten(numbers), 1)


def factorial(n):
    if isint(n):
        return long(round(gamma(n + 1)))  # if we're given an int, return something that looks like an int.

    return gamma(n + 1)


def rfactorial(n):
    if n >= 0:
        return factorial(n)

    try:
        return ((-1) ** (-n - 1)) / factorial(-n - 1)
    except ValueError:
        return ((-1) ** (-n - 1 + 0j)) / factorial(-n - 1)


def permutation(n, k, fact=factorial):
    if k > n:
        return 0

    import operator
    op = operator.truediv

    # if we have integer operators, return an int
    if isint(k) and isint(n):
        op = operator.ifloordiv

    return op(fact(n), fact(n - k))

def de_bruijn(k, n, includetail=True):
    '''De Bruijn sequence for alphabet k and subsequences of length n.
    You may also pass an integer for k, in which case the alphabet will be range(0,k).
    Set the includetail parameter to false if you wish to have a "toroid".'''
    # see FKM algorithm
    try:
        _ = int(k)
        alphabet = list(map(str, range(k)))

    except (ValueError, TypeError):
        alphabet = k
        k = len(k)

    a = [0] * k * n
    sequence = []

    def db(t, p):
        if t > n:
            if n % p == 0:
                sequence.extend(a[1:p + 1])
        else:
            a[t] = a[t - p]
            db(t + 1, p)
            for j in range(a[t - p] + 1, k):
                a[t] = j
                db(t + 1, t)
    db(1, 1)

    if includetail:
        sequence += sequence[0:n-1]

    return "".join(str(alphabet[i]) for i in sequence)


def binomial(z, w, fact=factorial):
   """return binominal coefficients for (z, w)"""
   # TODO: some kind of up-front bounds checking?  Not sure what the rules are exactly

   if 0 >= w and isint(w):        # w is a negative integer
      return 0
   elif 0 >= z and isint(z):      # z is a negative integer
      if isint(w):                # w is an integer
         return ((-1) ** w) * fact(w - z - 1) / (fact(w) * fact(-z - 1))
      else:
         return float("inf")
   else:                                # z is not a negative integer
      if isint(z) and isint(w):
         return int(round(gamma(z + 1)/(gamma(w + 1) * gamma(z - w + 1))))

      return gamma(z + 1)/(gamma(w + 1) * gamma(z - w + 1))


def multinomial(*arr, **kwargs):
    "computer multinomial coefficients; use fact keyword argument to specify factorial function"
    fact = kwargs.setdefault('fact', factorial)

    accum  = 0
    result = 1

    for k in flatten(arr):
        accum  += k
        result *= binomial(accum, k, fact)

    return result


def factor(n, naive=False):
    "given an integer, return an array of its prime factors"

    if not isint(n):
        raise TypeError("factorization is only defined for integers")
    if n < 0:
        raise ValueError("can't factorize negative numbers")

    # cheat and bust out GNU factor, if available
    from distutils.spawn import find_executable
    factor_path = find_executable('factor')
    try:
        from subprocess import check_output
    except ImportError:
        naive = True

    if factor_path and not naive:
        from re import split
        return [int(factor) for factor in split(":| ", check_output([factor_path, str(n)]).decode(encoding="ascii").strip())[2:]]

#   # cool idea from  http://stackoverflow.com/a/6800214 which yields all factors
#   return reduce(list.__add__, ([i, n//i] for i in range(1, int(n**0.5) + 1) if n % i == 0))

    # http://stackoverflow.com/a/14520938
    from math import floor
    if n == 0 or n == 1:
        return []
    res = []
    for i in range(2, int(floor(sqrt(n) + 1))):
        while n % i == 0:
            n //= i
            res.append(i)
    if n != 1:   # unusual numbers
        res.append(n)
    return res


def log(n, base=10):
   "return log, complex if necessary; base 10 default"
   from math import log
   if complex == type(n) or n < 0:
      from cmath import log
      return log(n)/log(base)
   else:
      from math import log
      return log(n)/log(base)


def lg(n):
   "return log base 2, complex if necessary"
   return log(n, base=2)


def ln(n):
   "return natural log, complex if necessary"
   from math import e
   return log(n, base=e)


def sqrt(n):
   "not quite a drop-in for math.sqrt; handles complex numbers by default"
   if complex == type(n) or ((type(n) == float or isint(n)) and n < 0):
      from cmath import sqrt
      return sqrt(n)
   elif type(n) in (float, int, long):
      from math import sqrt
      return sqrt(n)
   else:
      raise ValueError

def geomean(*arr):
    from operator import truediv, mul
    arr = list(flatten(arr))
    return reduce(mul, arr) ** truediv(1, len(arr))

# credit Beazley
def gen_open(filenames):
    for name in filenames:
        if name.endswith(".gz"):
            import gzip
            yield gzip.open(name)
        elif name.endswith(".bz2"):
            import bz2
            yield bz2.BZ2File(name)
        elif name.endswith(".xz") or name.endswith(".lz"):
            import lzip
            yield lzip.open(name)
        else:
            yield open(name)

atexit.register(readline.write_history_file, histfile)
del histfile, os

if __name__ == "__main__":
    assert(binomial(8,4) == 70)
    assert(de_bruijn(range(1,9,2), 3) == '111311511713313513715315515717317517733353373553573753775557577711')
    assert(de_bruijn(2, 2) == '00110')
    assert(de_bruijn(2, 3) == '0001011100')
    assert(de_bruijn(2, 4) == '0000100110101111000')
    assert(de_bruijn(5, 3) == '0001002003004011012013014021022023024031032033034041042043044111211311412212312413213313414214314422232242332342432443334344400')
    assert(de_bruijn(6, 2) == '0010203040511213141522324253343544550')
    assert(de_bruijn("abcd", 2) == 'aabacadbbcbdccdda')
    assert(binomial(20,4) == 4845)
    assert(multinomial(4,4,4,4,4) == 305540235000)
    assert(multinomial(4,[4,4],(4,4)) == 305540235000)
    assert(sqrt(-1) == 0+1j)
    assert(round(binomial(-10.2,3),3) == -232.288)
    assert(round(binomial(36,10.2)) == 305061085)
    assert(factor(6189814107) == [3, 3, 347, 457, 4337])
    assert(factor(6189814107, naive=True) == [3, 3, 347, 457, 4337])
    assert(permutation(8,4) == 1680)
    assert(round(permutation(8.5,4.2),7) == 3132.8466772)
    assert(sqrt(1+1j) == 1.09868411346781+0.45508986056222733j)
    assert(lg(1024) == 10.0)
    assert(ln(10) == 2.3025850929940459)
    assert(log(10) == 1.0)
    assert(round(log(1+1j).real, 13) == 0.150514997832)
    assert(lg(129510957-62897158194442687j).imag ==-2.266180067942957)