有没有一种方法可以向量化可依赖于前一行的行的函数?

时间:2019-10-21 06:43:52

标签: python pandas numpy vectorization yahoo-finance

我正在向股市数据添加指标,但是我找到ATR的解决方案太慢。有没有办法向量化这个问题?

我正在使用yfinance库获取股票价格。 最初,我做了一个简单的for循环,但是后来我意识到这会花很长时间,现在我正在使用Index来获取先前的行数据。我尝试使用.shift()函数,但是我需要获取以前的数据,因为如果我使用平均值174,将很难编写。 例如,我有此表:

日期开盘高低盘关闭
0 2017-10-23 156.839996 157.039993 155.949997 156.509995
1 2017-10-23 156.529999 156.550095 155.929901 156.339996
2 2017-10-23 156.320007 157.600006 156.119995 157.593994
3 2017-10-23 157.599899 157.679993 157.279999 157.279999
4 2017-10-23 157.297501 157.369995 156.559998 156.619995

我想添加一个由以下公式计算的ATR列:
当前ATR =(先前ATR +当前TR)/ 2,其中
当前TR = max [(高-低),(高-上一个收盘价),(低-上一个收盘价)

import yfinance as yf
import numpy as np
import pandas as pd

def get_prices(stock):
    ul = yf.Ticker(stock)
    pdate = (datetime.now() - timedelta(days=729)).date()
    df = ul.history(start=pdate, interval="1H").drop(["Volume","Dividends", "Stock Splits"], axis=1)
    df = df.reset_index()
    return df

def addATR_(row, df, av=13):
    if(row.Index>av):
        lrid = row.Index-1
        high = row.High
        low = row.Low
        prev_close = df.loc[df.index==(lrid),"Close"]
        matr = df[row.Index-av:row.Index]["ATR"].sum()
        tr=[(high - low),np.abs(high - prev_close).values,np.abs(low - prev_close).values] 
        return (matr + np.max(tr))/(av+1)

def add_ATR(df):
    df.insert(5,"ATR", 0.2)
    df['Index'] = df.index
    df["ATR"] = df.apply(lambda row: addATR_(row, df), axis=1)
    df = df.drop(range(14))
    df = df.reset_index().drop(["index","Index"],axis=1)
    print("Added ATR")
    return df

print(add_ATR(get_prices("msft")))

我真的想加快速度,因为仅添加一个指标需要9秒钟。

0 个答案:

没有答案