rvar

Author: microprediction

momentum/functions.py

@@ -44,12 +44,14 @@ def kurtosis_update(m: dict, x: float) -> dict:
     return m
 
 
-def rvar(m: dict, x: float = None, alpha=0.01, n=10):
-    """ Mimics the skater style where we pass {} to initialize """
+def rvar(m: dict, x: float = None, rho=0.01, n=10):
+    """ One function that performs either initialization or an update.
+        Pass m={} to initialize
+    """
     if m:
         return rvar_update(m=m, x=x)
     else:
-        return rvar_init(rho=alpha, n=n)
+        return rvar_init(rho=rho, n=n)
 
 
 def rvar_init(rho: float, n=10) -> dict:
@@ -59,21 +61,21 @@ def rvar_init(rho: float, n=10) -> dict:
     """
     assert 0 <= rho <= 1
     state = var_init()
-    state.update({'alpha': rho, 'burnin': n})
+    state.update({'rho': rho, 'n': n})
     return state
 
 
 def rvar_update(m: dict, x: float) -> dict:
-    if m['count'] < m['burnin']:
-        alpha = m['alpha']
+    if m['count'] < m['n']:
+        rho = m['rho']
         m = var_update(m, x)
-        m['alpha'] = alpha
+        m['rho'] = rho
         return m
     else:
         m['count'] += 1
-        alpha = m['alpha']
-        m['var'] = (1 - alpha) * (m['var'] + alpha * ((x - m['mean']) ** 2))
-        m['mean'] = (1 - alpha) * m['mean'] + alpha * x
+        rho = m['rho']
+        m['var'] = (1 - rho) * (m['var'] + rho * ((x - m['mean']) ** 2))
+        m['mean'] = (1 - rho) * m['mean'] + rho * x
         m['pvar'] = ((m['count'] - 1) / m['count']) * m['var']  # Not sure this really makes sense :)
         if m.get('M2'):
             del m['M2']

setup.py

@@ -7,7 +7,7 @@
 
 setup(
     name="momentum",
-    version="0.1.1",
+    version="0.1.2",
     description="Running estimates of moments",
     long_description=README,
     long_description_content_type="text/markdown",

tests/test_forgetting.py

@@ -1,4 +1,4 @@
-from momentum.functions import rvar_init, rvar_update
+from momentum.functions import rvar_init, rvar_update, rvar
 from pprint import pprint
 
 
@@ -12,6 +12,17 @@ def test_forget():
     assert -0.5<m['mean']<0.5
 
 
+def test_forget_other_style():
+    import numpy as np
+    xs = list(np.random.randn(5000))+list(2*np.random.randn(5000))
+    m = rvar(m={},rho=0.01)
+    for x in xs:
+        m = rvar_update(m, x)
+    assert 1.6<m['std']<2.4
+    assert -0.5<m['mean']<0.5
+
+
+
 if __name__=='__main__':
     test_forget()