exponentially weighted moving average

#!/usr/bin/python

def approx_eq(A, B, eps):
  """ [real] [real] epsilon => Boolean

  Checks if A and B are approximately equal, within epsilon

  >>> approx_eq([1, 2], [1.0, 2.0], 1e-6)
  True
  >>> approx_eq([1.5, 2.75], [1.5, 2.7], 0.01)
  [1.5, 2.75]
  [1.5, 2.7]
  False
  """
  for a, b in zip(A, B):
    if abs(a-b) > eps:
      print A
      print B
      return False
  return True

def algo0(l, queries, a=0.01):
  """
  [positive nums], [positive nums] => [positive nums]

  Compute exponentially weighted moving average

  Assumes l and queries are sorted

  >>> algo0([3.0,7.0,10.0,20.0,200.0], [23,68,103])
  [0.03552711850193005, 0.022601837152599107, 0.015899210247751473]
  """
  assert l == sorted(l)
  assert queries == sorted(queries)

  y = [0]*1000
  for item in l:
    y[int(item)]=1
  s = [0]*1000
  for i in xrange(1000):
    s[i] = a*y[i-1]+(1-a)*s[i-1]

  return [s[q] for q in queries]

def decay(t0, x, t1, a):
  """
  (real, real, real, real) => real

  Exponentially decays value x at time t0 to a value x' at time t1
  according to decay rate a.

  >>> decay(0, 0, 4, 0.01)
  0.0
  >>> decay(4, 0.01, 5, 0.01)
  0.0099
  >>> decay(4, 0.01, 6, 0.01)
  0.009801
  >>> decay(4, 0.01, 7, 0.01)
  0.00970299
  >>> decay(4, 0.01, 8, 0.01)
  0.0096059601
  """
  from math import pow
  return pow((1-a), (t1-t0))*x

def algo1(l, queries, a=0.01):
  """
  Compute exponentially weighted moving average

  Assumes l and queries are sorted

  >>> l,queries = [3.0, 7.0, 10.0], range(11)
  >>> approx_eq(algo0(l, queries), algo1(l, queries), 1e-7)
  True
  >>> l, queries = [3.0,7.0,10.0,20.0,200.0], [23,68,103]
  >>> approx_eq(algo0(l, queries), algo1(l, queries), 1e-7)
  True
  """
  assert l == sorted(l)
  assert queries == sorted(queries)

  samples = []
  try:
    t0, x0 = (0.0, 0.0)
    it = iter(queries)
    q = it.next()-1.0

    for t1 in l:
      # new value is decayed previous value, plus a
      x1 = decay(t0, x0, t1, a) + a
      # take care of all queries between t0 and t1
      while q < t1:
        samples.append(decay(t0, x0, q, a))
        q = it.next()-1.0
      # take care of all queries equal to t1
      while q == t1:
        samples.append(x1)
        q = it.next()-1.0
      # update t0, x0
      t0, x0 = t1, x1

    # take care of any remaining queries
    while True:
      samples.append(decay(t0, x0, q, a))
      q = it.next()-1.0
  except StopIteration:
    return samples

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