Добавление столбцов в Pandas DataFrame из класса

Question

Я стремлюсь улучшить свой OOP набор навыков и написал скрипт для извлечения биржевых данных и запуска простой статистики. Я могу запускать и вызывать каждую функцию в классе Evaluation по отдельности (см. Ниже), но у меня возникают проблемы при попытке пролистать список тикеров и добавить статистику в исходный фрейм данных.

import datetime as d
import pandas as pd
import pandas_datareader.data as web
import numpy as np

start = d.datetime(2019, 1, 1)
end = d.datetime(2020, 4, 17)

class Security(object):
    def _init__(self, ticker, data_platform, start_date, end_date):
        self.ticker = ticker
        self.data_platform = data_platform
        self.start_date = start_date
        self.end_date = end_date

    def fetch_stock_data(self, ticker, data_platform, start_date, end_date):
        df = web.DataReader(ticker, data_platform, start_date, end_date)
        return df


class Evaluation(Security):

    def __init__(self, ticker, data_platform, start_date, end_date):
        self.df = Security.fetch_stock_data(
            self, ticker, data_platform, start_date, end_date)

    def simple_moving_average(self, period):
        df = self.df
        df['SMA-{}'.format(period)] = df['Adj Close'].rolling(period).mean()
        return df['SMA-{}'.format(period)]

    def exp_moving_average(self, period):
        df = self.df
        df['EMA_{}'.format(period)] = df['Adj Close'].ewm(span=period).mean()
        return df['EMA_{}'.format(period)]

    def rsi(self, period):
        df = self.df
        delta = df['Adj Close'].diff()

        up = delta * 0
        down = up.copy()

        up[delta > 0] = delta[delta > 0]
        down[delta < 0] = -delta[delta < 0]

        up[up.index[period - 1]] = np.mean(up[:period])
        up = up.drop(up.index[:(period - 1)])

        down[down.index[period - 1]] = np.mean(down[:period])
        down = down.drop(down.index[:(period - 1)])

        rs = up.ewm(span=period - 1).mean() / down.ewm(span=period - 1).mean()

        rsi_calc = 100 - 100 / (1 + rs)
        df['rsi'] = rsi_calc
        return df['rsi']


# pypl = Evaluation('PYPL', 'yahoo', start, end)
# print(csgs.df)
# print(csgs.simple_moving_average(50))
# print(csgs.exp_moving_average(26))
# print(csgs.rsi(14))


tickers = ['PYPL', 'TSLA']

for i in tickers:
    df = Evaluation(i, 'yahoo', start, end)
    df['SMA'] = df.simple_moving_average(50)
    df['EMA'] = df.exp_moving_average(26)
    df['rsi'] = df.rsi(14)
    print(df)

Я получаю ошибку TypeError, которая, по моему мнению, связана со ссылкой на класс Evaluation.

TypeError: 'Evaluation' object does not support item assignment

Answer 1

Вы путаете методы объекта с методами dataframe. В вашем примере df - это объект Evaluation, а не датафрейм.

>>>e = Evaluation()                                                                                                  
>>>type(e)                                                                                                  
__main__.Evaluation

>>>type(e.df)                                                                                         
pandas.core.frame.DataFrame

Строка df['SMA'] = df.simple_moving_average(50) не работает, поскольку вы не можете добавить столбец к объекту. Вам нужно использовать df.df['SMA'] = df.simple_moving_average(50).

Как указывал NomadMonad, использование df в качестве имени переменной для объекта Evaluation сбивает с толку, поэтому было бы лучше дать ему другое имя. Однако eval является встроенной функцией в python, поэтому было бы лучше использовать e.

Кроме того, вам следует изменить дизайн класса по нескольким причинам

• В python 3 нет необходимости наследовать от Object

• Метод __init__ для Security имеет только одно начальное подчеркивание вместо двух .

• Вы не хотите, чтобы Evaluation наследовал от Security. Вместо этого передайте объект Security в методе __init__ Evaluation.

• Вы не хотите вызывать метод, который очищает веб-сайт при создании экземпляра объекта. Вызов pandas_datareader должен быть отдельным методом.

• Вам не нужно передавать параметры в методы, если эти значения установлены в методе __init__. Вы можете получить к ним доступ с помощью self.

• Вам не нужно изменять базовый фрейм данных в ваших методах оценки. Вместо этого верните значение, полученное методом.

import datetime
import pandas as pd
import numpy as np
import pandas_datareader.data as web 


class Security:
    def _init__(self, ticker, data_platform, start_date, end_date):
        self.ticker = ticker
        self.data_platform = data_platform
        self.start_date = start_date
        self.end_date = end_date
        self.df = None

    def fetch_stock_data(self):
        self.df = web.DataReader(self.ticker, self.data_platform, self.start_date, self.end_date)


class Evaluation:

    def __init__(self, security):
        self.security = security

    def simple_moving_average(self, period):
        df = self.security.df
        return df['Adj Close'].rolling(period).mean()

    def exp_moving_average(self, period):
        df = self.security.df
        return df['Adj Close'].ewm(span=period).mean()

    def rsi(self, period):
        df = self.security.df
        delta = df['Adj Close'].diff()
        up = delta * 0
        down = up.copy()
        up[delta > 0] = delta[delta > 0]
        down[delta < 0] = -delta[delta < 0]
        up[up.index[period - 1]] = np.mean(up[:period])
        up = up.drop(up.index[:(period - 1)])
        down[down.index[period - 1]] = np.mean(down[:period])
        down = down.drop(down.index[:(period - 1)])
        rs = up.ewm(span=period - 1).mean() / down.ewm(span=period - 1).mean()
        return 100 - 100 / (1 + rs)


start = datetime.datetime(2019, 1, 1)
end = datetime.datetime(2019, 4, 17)

s = Security(ticker, 'yahoo', start, end)
e = Evaluation(security=s)