ADX.py

import numpy as np
import pandas as pd

# ADX
def adx_di(high, low, close, period_di, period_adx):
  last_close = shift(close, 1, cval=0)
  last_high = shift(high, 1, cval=0)
  last_low = shift(low, 1, cval=0)

  dmp = np.where((high - last_high) < 0, 0, (high - last_high))
  dmp = np.where((last_low - low) < 0, 0, (last_low - low))
  dmp = np.where(dmp > dmm, dmp, 0)
  dmp = np.where(dmp > dmm, 0, dmm)
  tr = np.maximum(high - low, high - last_close, last_close - low)

  for i in range(1, len(close)):
    dmp[i] = dmp[i-1] - (dmp[i-1] / period_di) + dmp[i]
    dmm[i] = dmm[i-1] - (dmm[i-1] / period_di) + dmm[i]
    tr[i] = tr[i-1] - (tr[i-1] / period_di) + tr[i]

  plusdi = (dmp / tr) * 100
  minusdi = (dmm / tr) * 100

  dx = (abs(plusdi - minusdi) / (plusdi + minusdi)) * 100
  adx = pd.Series(dx).rolling(period_adx).mean()
  ADX_DI = np.c_[adx, plusdi, minusdi]
  return ADX_DI