From 48b561f14ef05ab6fd52af5855f4c77f38aec440 Mon Sep 17 00:00:00 2001 From: Jamie Date: Wed, 5 May 2021 11:33:45 -0400 Subject: [PATCH] Docstring-ed and Typing-ed all functions in nn_view --- gamestonk_terminal/feature_flags.py | 4 +- .../neural_networks_helper.py | 0 .../neural_networks_view.py | 129 ++++++++++++++++-- 3 files changed, 118 insertions(+), 15 deletions(-) create mode 100644 gamestonk_terminal/prediction_techniques/neural_networks_helper.py diff --git a/gamestonk_terminal/feature_flags.py b/gamestonk_terminal/feature_flags.py index 5b32fb003fa..061b6688f82 100644 --- a/gamestonk_terminal/feature_flags.py +++ b/gamestonk_terminal/feature_flags.py @@ -14,10 +14,10 @@ USE_FLAIR = os.getenv("GTFF_USE_FLAIR") or "stars" USE_ION = strtobool(os.getenv("GTFF_USE_ION", "True")) # Enable Prompt Toolkit -USE_PROMPT_TOOLKIT = strtobool(os.getenv("GTFF_USE_PROMPT_TOOLKIT", "False")) +USE_PROMPT_TOOLKIT = strtobool(os.getenv("GTFF_USE_PROMPT_TOOLKIT", "True")) # Enable Prediction features -ENABLE_PREDICT = strtobool(os.getenv("GTFF_ENABLE_PREDICT", "False")) +ENABLE_PREDICT = strtobool(os.getenv("GTFF_ENABLE_PREDICT", "True")) # Enable FB Prophet ENABLE_FBPROPHET = strtobool(os.getenv("GTFF_ENABLE_FBPROPHET", "False")) diff --git a/gamestonk_terminal/prediction_techniques/neural_networks_helper.py b/gamestonk_terminal/prediction_techniques/neural_networks_helper.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/gamestonk_terminal/prediction_techniques/neural_networks_view.py b/gamestonk_terminal/prediction_techniques/neural_networks_view.py index 498606321d5..9c7bacbc5eb 100644 --- a/gamestonk_terminal/prediction_techniques/neural_networks_view.py +++ b/gamestonk_terminal/prediction_techniques/neural_networks_view.py @@ -2,7 +2,7 @@ __docformat__ = "numpy" import argparse -from typing import List, Any +from typing import List, Any, Union import datetime import os import traceback @@ -111,7 +111,6 @@ def _parse_args(prog: str, description: str, other_args: List[str]): Returns ------- ns_parser: argparse.Namespace - Parsed argument parser """ @@ -347,10 +346,32 @@ def _setup_backtesting(df_stock: pd.DataFrame, ns_parser: argparse.Namespace): return df_stock, df_future -def _preprocess_split(df_stock, ns_parser): - """Preprocess and split training data. - :return: (scaler, stock_train_data, stock_x, stock_y) - :raises Exception: if more training data is needed.""" +def _preprocess_split(df_stock: pd.DataFrame, ns_parser: argparse.Namespace): + """ + Preprocess and split dstock data + Parameters + ---------- + df_stock: pd.DataFrame + Dataframe of stock prices + ns_parser: argparse.Namespace + Parsed argparse arguments + + Returns + ------- + scaler: Union[MinMaxScaler, StandardScaler, Normalizer] + sklearn scaler used + stock_train_data: np.ndarray + All training data + stock_x: np.ndarray + Inputs for model + stock_y: np.ndarray + Outputs for model + Raises + ------- + Exception + Not enough training data is provided + """ + # Pre-process data if ns_parser.s_preprocessing == "standardization": scaler = StandardScaler() @@ -379,13 +400,40 @@ def _preprocess_split(df_stock, ns_parser): ) if not stock_x: raise Exception("Given the model parameters more training data is needed.") + stock_x = np.array(stock_x) stock_y = np.array(stock_y) + return scaler, stock_train_data, stock_x, stock_y -def _rescale_data(df_stock, ns_parser, scaler, yhat, idx_loop): - """Re-scale the data back and return the prediction dataframe.""" +def _rescale_data( + df_stock: pd.DataFrame, + ns_parser: argparse.Namespace, + scaler: Union[MinMaxScaler, StandardScaler, Normalizer], + yhat: np.ndarray, + idx_loop: int, +): + """ + Re-scale the data back and return the prediction dataframe. + Parameters + ---------- + df_stock: pd.DataFrame + Dataframe of stock prices + ns_parser: argparse.Namespace + Parsed argument parser + scaler: + One of the predefined sklearn scalers + yhat: np.ndarry + Predicted values + idx_loop: int + Loop index + + Returns + ------- + df_pred: pd.Series + Series of resclaled predictions + """ if ( (ns_parser.s_preprocessing == "standardization") or (ns_parser.s_preprocessing == "normalization") @@ -403,11 +451,36 @@ def _rescale_data(df_stock, ns_parser, scaler, yhat, idx_loop): df_pred = pd.Series(y_pred_test_t[0].tolist(), index=l_pred_days, name=column_name) return df_pred - def _plot_and_print_results( - df_stock, ns_parser, df_future, df_pred, model_name, s_ticker + df_stock: pd.DataFrame, + ns_parser: argparse.Namespace, + df_future: pd.DataFrame, + df_pred: pd.DataFrame, + model_name: str, + s_ticker: str, ): - """Plot and print the results.""" + """ + Plot and print results + Parameters + ---------- + df_stock: pd.DataFrame: + Dataframe of stock prices + ns_parser: argparse.Namespace + Parsed argument parser + df_future: pd.DataFrame + Dataframe of future prices + df_pred: pd.DataFrame + Dataframe of predicted prices + model_name: str + Model Name + s_ticker: str + Ticker + + Returns + ------- + + """ + # Plotting plt.figure(figsize=plot_autoscale(), dpi=PLOT_DPI) plt.plot(df_stock.index, df_stock["5. adjusted close"], lw=3) @@ -828,7 +901,22 @@ def mlp(other_args: List[str], s_ticker: str, df_stock: pd.DataFrame): _restore_env() -def rnn(other_args: List[str], s_ticker, df_stock): +def rnn(other_args: List[str], s_ticker: str, df_stock: pd.DataFrame): + """ + Train a Recurrent Neural Network (rnn) + Parameters + ---------- + other_args:List[str] + Argparse arguments + s_ticker: str + Stock ticker + df_stock: pd.DataFrame + Dataframe of stock prices + + Returns + ------- + + """ try: ns_parser = _parse_args( prog="rnn", @@ -901,7 +989,22 @@ def rnn(other_args: List[str], s_ticker, df_stock): _restore_env() -def lstm(other_args: List[str], s_ticker, df_stock): +def lstm(other_args: List[str], s_ticker: str, df_stock: pd.DataFrame): + """ + Train a Long-Short-Term-Memory Neural Net (lstm) + Parameters + ---------- + other_args:List[str] + Argparse arguments + s_ticker: str + Stock ticker + df_stock: pd.DataFrame + Dataframe of stock prices + + Returns + ------- + + """ try: ns_parser = _parse_args( prog="lstm",