utils.py

def _init_():
  global \
    Singleton, OneOfKind, \
    UnionFind, DynamicLRU, \
    max_, select, prod, split_on_none, powerset, \
    matrix_to_formula, \
    stabilize_period, matrix_formula_period

  from itertools import chain, combinations
  from functools import reduce
  from collections import defaultdict
  import operator
  import weakref
  import inspect


  def push_tests(cls):
    mod = inspect.getmodule(cls)
    if not hasattr(mod, "__test__"): setattr(mod, "__test__", {})
    mod_tests = getattr(mod, "__test__")
    name = cls.__name__

    if name in mod_tests:
      i = 1
      while ("%s_%i" % (name, i)) in mod_tests: i += 1
      name = "%s_%i" % (name, i)

    mod_tests[name] = cls


  def Singleton(cls):
    """Decorator which replaces the class with object of it."""
    push_tests(cls)
    return cls()

  def OneOfKind(cls):
    """Decorator which ensures that there exisits at most one
    object of given class for every tuple of constuctor
    arguments."""

    push_tests(cls)

    cache = DynamicLRU()

    def call(*args, **kwargs):
      key = (tuple(args), tuple(sorted(kwargs.items())))
      return cache.get(key, lambda: cls(*args, **kwargs))

    return call

  class UnionFind:
    """Simple union-find datastructure.

    It is implemented with pointers. No path compression
    or ranks are implemented.

    EXAMPLES:
    >>> uf = UnionFind(5)
    >>> any( uf.same(i, j) for i in range(5) for j in range(i) )
    False
    >>> uf.union(3, 2)
    >>> any( uf.same(i, j) for i in range(5) for j in range(i) )
    True
    >>> uf.same(2, 3)
    True
    >>> uf.same(1, 2)
    False
    >>> uf.union(2, 0)
    >>> uf.same(0, 3)
    True
    """

    def __init__(self, size):
      self.x = [ None ] * size

    def union(self, a, b):
      a = self.find(a)
      b = self.find(b)
      if a != b: self.x[b] = a

    def find(self, a):
      while self.x[a] is not None:
        a = self.x[a]

      return a

    def same(self, a, b):
      return self.find(a) == self.find(b)


  def max_(l):
    """Version of `max` which ignores `None` in the list of arguments.

    EXPAMPLES:
    >>> max_([1, 5, 2, 3])
    5
    >>> max_([0, -14, 2, 7, None, 9])
    9
    """
    return max( x for x in l if x is not None )

  def select(pat, l):
    """Shuffle `l` according to pattern `pat`.

    Returns list-like object of the same type as `pat`
    and of the same length as `pat` where on `i`-th
    position is element `l[pat[i]]`.

    EXAMPLES:
    >>> select([5, 3, 1], [1, 2, 3, 4, 5, 6])
    [6, 4, 2]
    >>> select((5, 3, 1), [1, 2, 3, 4, 5, 6])
    (6, 4, 2)
    """

    return type(pat)( l[i] for i in pat )


  def prod(iterable):
    """Calculate product of elements of `iterable`."""
    return reduce(operator.mul, iterable, 1)


  class DynamicLRU:
    def __init__(self, capacity = 100):
      self._capacity = capacity
      self._data = dict()
      self._time = 0

    def get(self, key, make_value):
      self._time += 1

      if key not in self._data:
        self._data[key] = (make_value(), self._time)

      return self._data[key][0]


  def split_on_none(t):
    """Split list (or tuple) at first `None`.

    EXAMPLE:
    >>> split_on_none([1, 2, 3, None, 4, 5])
    ([1, 2, 3], [4, 5])
    >>> split_on_none([1, 2])
    ([1, 2], [])
    """

    i = 0
    while i < len(t) and t[i] is not None: i += 1
    return (t[:i], t[i+1:])


  def powerset(iterable):
      """Enumerate all subsets of `iterable`.

      Vraci tuply!!! TODO

      EXAMPLE:
      >>> list(powerset([1,2,3]))
      [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
      """

      s = list(iterable)
      return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))


  def stabilize_period(X, vectors=False, shift=0):
    return matrix_formula_period(
      X['finalize'], X['step_matrix'], X['initial_vector'], vectors, shift)


  def matrix_formula_period(l, A, r, vectors=False, shift=0):

    from sage.all import matrix

    assert all( x >= 0 for x in l.dict().values() )
    assert all( x >= 0 for x in A.dict().values() )
    assert all( x >= 0 for x in r.dict().values() )

    def to_zero_one(M):
      for (r, c), v in M.dict().items():
        assert v > 0
        M[r, c] = 1
      return M

    l = to_zero_one(matrix(l))
    A = to_zero_one(matrix(A))
    r = to_zero_one(matrix(r))

    seen = []
    while True:
      for i in range(len(seen)):
        if seen[i] == r:
          return (i + shift,
            [None]*shift + (seen if vectors else [ int(l*v != 0) for v in seen ]))

      seen.append(r)
      r = to_zero_one(A*r)

  def matrix_to_formula(l, A, r, k=None, base_ring=None):
    """Tranform matrix formula l A^k r to simple formula.

    First rewrite A as P J P^-1 using Jordan normal form.
    Then convert to formula using sage symbolic expressions.

    Sage (backed by Maxima) simplifies 0^x to 0 which is
    not correct because we need 0^0 == 1. So we check correctness
    of resulting formula for values of k less than size of
    the largest block of J.

    EXAMPLE:
    We derive formula for Fibonacci numbers:
    >>> from sage.all import *
    >>> m = matrix(QQ, 2, 2, [1, 1, 1, 0])
    >>> r = matrix(QQ, 2, 1, [1, 0])
    >>> l = matrix(QQ, 1, 2, [0, 1])
    >>> [ l * m**i * r for i in range(10) ]
    [[0], [1], [1], [2], [3], [5], [8], [13], [21], [34]]
    >>> f = matrix_to_formula(l, m, r, base_ring=QQbar)['formula']; f
    0.4472135954999580?*1.618033988749895?^k - 0.4472135954999580?*(-0.618033988749895?)^k
    >>> [ f(k=i) for i in range(10) ]
    [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
    """

    from sage.all import binomial, SR

    J, P = A.jordan_form(transformation=True, sparse=True, base_ring=base_ring)
    assert P * J * P**(-1) == A, "Error %s != %s" % (P * J * P**(-1), A)
    l = l * P
    r = P**(-1) * r

    if k is None: k = SR.symbol('k')
    out_k = k
    k = SR.symbol('k')
    f = 0

    subdivisions = J.subdivisions()
    assert subdivisions[0] == subdivisions[1]
    subdivisions = [ 0 ] + subdivisions[0] + [ J.nrows() ]
    to_delete = []

    max_s = 0
    for i in range(len(subdivisions) - 1):
      s = subdivisions[i+1] - subdivisions[i]
      max_s = max(max_s, s)
      for j in range(s):
        ind = subdivisions[i] + j
        f += l[0, ind] * sum(
          binomial(k, ll) * J[ind, ind]**(k - ll) * r[ind + ll, 0] for ll in range(s-j)
        )

    try:
      f = f.simplify_full()
    except: pass

    while max_s > 0:
      val_r = (l * J**(max_s - 1) * r)[0,0]
      val_f = f(k=max_s - 1)
      if val_r != val_f: break
      max_s -= 1

    try:
      f = f.subs(k=out_k).simplify_full()
    except: pass

    return {
      'simplified': (l, J, r),
      'formula': f,
      'formula_if': out_k >= max_s,
    }


_init_()