improve WMA algorithm

algorithm/wma.py

@@ -1,22 +1,33 @@
 class WMA:
-  def __init__(self, unit, period, width):
-    self._unit   = unit
-    self._period = period
-    self._width  = width
+  def __init__(self, unit, width):
+    self._unit         = unit
+    self._width        = width
+
+    self._period_short = 7
+    self._period_mid   = 15
+    self._period_long  = 60
+
+    self._posi_amp = 0.1
 
   def judge(self, status):
     cands = status.candles[self._unit]
+    short = self._wma(cands, self._period_short)
+    mid   = self._wma(cands, self._period_mid)
+    long  = self._wma(cands, self._period_long)
+
+    width = mid * self._width
+    posi  = max(-1, min((short-long)/width, 1)) * self._posi_amp
+
+    diff  = short - mid
+    if diff < 0:
+      return max(-1, diff/width+posi)  # sell
+    else:
+      return min(1, diff/width+posi)  # buy
+
+  def _wma(self, cands, period):
     sum = 0
     den = 0
-    for i in range(min(self._period-1, len(cands)-1)):
-      sum += cands[i+1][3] * (i+1)
+    for i in range(min(period, len(cands))):
+      sum += cands[i][3] * (i+1)
       den += i+1
-    avg   = sum / den
-    price = cands[0][3]
-
-    width = price * self._width
-    diff  = price - avg
-    if diff < 0:
-      return max(-1, diff/width)  # sell
-    else:
-      return min(1, diff/width)  # buy
+    return sum / den