python 计算股市技术分析PSY指标

PSY，心理线，百度百科链接：

https://baike.baidu.com/item/PSY%E6%8C%87%E6%A0%87/3083493?fr=aladdin

价格上涨可以通过两种方式判定：第一种是当日收盘价格高于前日收盘价格；

第二种是当日收盘价格高于当日开盘价格。这里用的是第一种方法。

import numpy as np
def getPSY(priceData, period):
    #rolling calculation of the PSY, psychological line
    #priceData should be the close price data, in np format
    #period is the length of days in which we look at how many days witness price increase
    difference = priceData[1:] - priceData[:-1]
    #price change
#    difference = np.append(np.nan, difference)
    #to make the length of the difference same as the priceData, lag of one day
    #to avoid the warning, use 0 instead of np.nan, the result should be the same
    difference = np.append(0, difference)    
    difference_dir = np.where(difference > 0, 1, 0)
    #get the direction of the price change, if increase, 1, else 0
    psy = np.zeros((len(priceData),))
    psy[:period] *= np.nan
    #there are two kind of lags here, the lag of the price change and the lag of the period
    for i in range(period, len(priceData)):
        psy[i] = (difference_dir[i-period+1:i+1].sum()) / period
        #definition of the psy: the number of the price increases to the total number of days
    return psy