金融指标工具

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from dateutil.parser import parse

# 显示dataframe中的每一行和每一列
pd.set_option('display.unicode.ambiguous_as_wide', True)
pd.set_option('display.unicode.east_asian_width', True)
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.set_option('display.width', 5000)
pd.set_option('mode.chained_assignment', None)


def timeprocess(df, starttime, endtime):  # 去除不同时间间隔的影响
    extra = pd.date_range(start=starttime + ' ' + '00:01:00.000', end=endtime + ' ' + '23:59:00.000',
                          freq='T')  # 生成一个从starttime到endtime，每分钟间隔的数据
    extra_df = pd.DataFrame({'receiveTime': extra})
    df['receiveTime'] = df['receiveTime'].apply(parse)
    df = pd.merge(extra_df, df, how='outer', on='receiveTime')
    df = df.fillna(method='pad')
    df['receiveTime'] = df['receiveTime'].apply(lambda x: x.strftime("%Y-%m-%d %H:%M:%S"))
    return df

def get_EMA(df, N, colname='Close'):
    ema = []
    for i in range(len(df)):
        if i == 0:
            ema.append(df.loc[i, colname])
        else:
            ema.append((2 * df.loc[i, colname] + (N - 1) * ema[i - 1]) / (N + 1))
    return ema


def get_MACD(df, short=12, long=26, M=9):
    a = get_EMA(df, short)
    b = get_EMA(df, long)
    df['diff'] = pd.Series(a) - pd.Series(b)
    dea = get_EMA(df, M, 'diff')
    macd = [2 * (df['diff'][i] - dea[i]) for i in range(len(dea))]
    return macd


def cal_macd_system(data, short_, long_, m):
    '''
    short_,long_,m分别是macd的三个参数
    返回值是包含原始数据和diff,dea,macd三个列的dataframe
    '''
    data['diff'] = data['Close'].ewm(adjust=False, alpha=2 / (short_ + 1), ignore_na=True).mean() - \
                   data['Close'].ewm(adjust=False, alpha=2 / (long_ + 1), ignore_na=True).mean()
    data['dea'] = data['diff'].ewm(adjust=False, alpha=2 / (m + 1), ignore_na=True).mean()
    data['macd'] = 2 * (data['diff'] - data['dea'])
    return data


def double_ma(df, sma=5, lma=120):  # ma双均线
    # df['receiveTime'] = pd.to_datetime(df['receiveTime'])  # 更改date列的数据类型
    # df.set_index('receiveTime', inplace=True)  # 把行索引设置为date列，且直接在源数据上进行改动
    '''绘制5日均线图与120日均线图来示范'''
    small_ma = df['lastPrice'].rolling(sma).mean()  # 求5日均线
    big_ma = df['lastPrice'].rolling(lma).mean()  # 120日均线`
    '''因为我们长均线是120日均线所以我们要先对df,small_ma,big_ma进行切片操作'''
    small_ma = small_ma[lma:]
    big_ma = big_ma[lma:]
    df = df[lma:]
    s1 = small_ma < big_ma
    s2 = small_ma > big_ma
    death = s1 & s2.shift(1)  # 判定死叉的条件
    death_date = df.loc[death].index  # 死叉的所有时间
    golden = ~(s1 | s2.shift())  # 判定金叉的条件
    golden_date = df.loc[golden].index  # 金叉的所有时间
    return death_date, golden_date


def timeadd(times):  # 为times时间轴加上一分钟生成新的时间轴
    b = times  # 类型为Timestamp
    if (b.minute + 1) > 59:  ##超出60分钟界限，小时加1
        b = b.replace(hour=b.hour + 1)
        b = b.replace(minute=(b.minute + 1) % 60)
    else:
        b = b.replace(minute=b.minute + 1)
    return b


def macd_strategy(df):  # macd 线交叉策略
    data = cal_macd_system(df, 12, 26, 9)
    diff = data['diff']
    dea = data['dea']
    s1 = diff < dea
    s2 = diff > dea
    death = s1 & s2.shift(1)  # 判定死叉的条件
    death_date = df.loc[death].index  # 死叉的所有时间
    golden = ~(s1 | s2.shift())  # 判定金叉的条件
    golden_date = df.loc[golden].index  # 金叉的所有时间
    return death_date, golden_date  # 金叉和死叉的时间