def set_input_data(self, df): super().set_input_data(df) # Set dataFrame keys self.volume_key = Constants.get_volume_key() self.indicator_key = Constants.get_key("OBV") prices_temp = pd.DataFrame() df_list = [] for ticker in self.tickers: if ticker in df: df_list.append( pd.concat([ df[ticker].loc[:, [self.volume_key, self.prices_key]], prices_temp ], axis=1, keys=[ticker])) df_indicator = pd.concat(df_list, axis=1) self.df = df_indicator.copy()
def plot_stock(self, stock, tickers=None, collapse_indicators=False): Plotter.legend_id = 0 Plotter.current_color_indicator = 0 if stock is None: print("There is no ticker Information, nothing to be plot") return if tickers is None: tickers = stock.tickers elif isinstance(tickers, list) is True: tickers = tickers else: tickers = [tickers] if self.fig is None or self.axes_main is None or self.axes_indicators is None: if stock.price_info is None or stock.price_info[tickers].empty: raise ValueError( "There is no price information for this stock") adj_close_key = Constants.get_adj_close_key() volume_key = Constants.get_volume_key() self.price_series = {} self.volume_series = {} self.x_series = {} for ticker in tickers: if (adj_close_key in stock.price_info[ticker]) == False: adj_close_key = Constants.get_close_key() self.x_series[ticker] = stock.price_info[ticker].iloc[ -self.period:, :].index self.price_series[ticker] = stock.price_info[ticker].iloc[ -self.period:, :][adj_close_key] self.volume_series[ticker] = stock.price_info[ticker].iloc[ -self.period:, :][volume_key] self.axes_main = dict() self.axes_indicators = dict() if len(stock.indicators) == 0: subplots = 1 elif len(stock.indicators) > 0 and collapse_indicators is True: extra = len( list( filter( lambda x: x.collapse is False or x.in_main_plot is False, stock.indicators))) no_collapse = len( list( filter( lambda x: x.collapse is False and x. in_main_plot is False, stock.indicators))) if no_collapse > 0: no_collapse = no_collapse - 1 fixed = 2 if extra == 0: fixed = 1 subplots = fixed + no_collapse else: subplots = len( list( filter(lambda x: x.in_main_plot is False, stock.indicators))) + 1 heights_list = [2 for i in range(subplots - 1)] if subplots == 1: heights_list.insert(0, 2) else: heights_list.insert(0, 3) self.fig = plt.figure(figsize=(8, 6), dpi=80) gridspec = self.fig.add_gridspec(ncols=1, nrows=subplots, height_ratios=heights_list) # gridspec_kw = {'height_ratios': heights_list} self.axes_main[Constants.volume_axis] = self.fig.add_subplot( gridspec[0, 0]) self.set_volume(ticker=ticker) self.set_stock_price(ticker=ticker, color=self.stock_color) i = 1 # Indicator axis begins in 2 indicator_axis = None Plotter.legend_id = 0 for indicator in stock.indicators: if indicator.in_main_plot is False: if collapse_indicators == True: if i > 1: if indicator.collapse is False: self.axes_indicators = self.fig.add_subplot( gridspec[i, 0]) indicator_axis = self.axes_indicators else: indicator_axis = indicator_axis.twinx() else: # Executed first self.axes_indicators = \ self.fig.add_subplot( gridspec[i, 0], sharex=self.axes_main[Constants.prices_axis]) indicator_axis = self.axes_indicators else: Plotter.legend_id = 0 sharex = None if indicator.collapse is True: sharex = self.axes_main[Constants.prices_axis] self.axes_indicators = \ self.fig.add_subplot( gridspec[i, 0], sharex=sharex) indicator_axis = self.axes_indicators if indicator_axis is None: indicator_axis = self.axes_indicators i += 1 else: indicator_axis = self.axes_main[Constants.prices_axis] self.set_plot_indicator(indicator=indicator, ticker=ticker, axis=indicator_axis) plt.tight_layout()