starting to refactor data preprosess

gamestonk_terminal/prediction_techniques/knn_view.py

@@ -21,6 +21,7 @@ from gamestonk_terminal.prediction_techniques.pred_helper import (
     print_pretty_prediction,
     price_prediction_backtesting_color,
     print_prediction_kpis,
+    get_backtesting_data,
 )
 
 from gamestonk_terminal.config_plot import PLOT_DPI
@@ -108,34 +109,12 @@ def k_nearest_neighbors(other_args: List[str], s_ticker: str, df_stock: pd.DataF
 
         # BACKTESTING
         if ns_parser.s_end_date:
-            if ns_parser.s_end_date < df_stock.index[0]:
-                print(
-                    "Backtesting not allowed, since End Date is older than Start Date of historical data\n"
-                )
-                return
-
-            if ns_parser.s_end_date < get_next_stock_market_days(
-                last_stock_day=df_stock.index[0],
-                n_next_days=ns_parser.n_inputs + ns_parser.n_days,
-            )[-1]:
-                print(
-                    "Backtesting not allowed, since End Date is too close to Start Date to train model\n"
-                )
-                return
-
-            future_index = get_next_stock_market_days(
-                last_stock_day=ns_parser.s_end_date, n_next_days=ns_parser.n_days
+            df_future, df_stock, bt_flag = get_backtesting_data(
+                df_stock, ns_parser.s_end_date, ns_parser.n_days
             )
-
-            if future_index[-1] > datetime.datetime.now():
-                print(
-                    "Backtesting not allowed, since End Date + Prediction days is in the future\n"
-                )
+            if not bt_flag:
                 return
 
-            df_future = df_stock[future_index[0] : future_index[-1]]
-            df_stock = df_stock[: ns_parser.s_end_date]
-
         # Split training data
         stock_x, stock_y = splitTrain.split_train(
             df_stock["5. adjusted close"].values,

gamestonk_terminal/prediction_techniques/pred_helper.py

@@ -13,6 +13,8 @@ from sklearn.metrics import (
     r2_score,
     mean_squared_error,
 )
+from sklearn.model_selection import train_test_split
+
 from gamestonk_terminal.helper_funcs import (
     get_next_stock_market_days,
     patch_pandas_text_adjustment,
@@ -24,6 +26,51 @@ from gamestonk_terminal.config_plot import PLOT_DPI
 register_matplotlib_converters()
 
 
+def prepare_test_train_valid(
+    df_stock, parsed_end_date, n_input_days, n_predict_days, test_size
+):
+
+    dates = df_stock.index
+    prices = df_stock["5. adjusted close"].values
+
+    input_dates = []
+    input_prices = []
+    next_n_day_prices = []
+    next_n_day_dates = []
+
+    for idx in range(len(prices) - n_input_days - n_predict_days):
+        input_prices.append(prices[idx : idx + n_input_days])
+        input_dates.append(dates[idx : idx + n_input_days])
+        next_n_day_prices.append(
+            prices[idx + n_input_days : idx + n_input_days + n_predict_days]
+        )
+        next_n_day_dates.append(
+            dates[idx + n_input_days : idx + n_input_days + n_predict_days]
+        )
+
+    input_dates = np.asarray(input_dates)
+    input_prices = np.asarray(input_prices)
+    next_n_day_prices = np.asarray(next_n_day_prices)
+    next_n_day_dates = np.asarray(next_n_day_dates)
+
+    (
+        X_train,
+        X_test,
+        y_train,
+        y_test,
+        X_dates_train,
+        X_dates_test,
+        y_dates_train,
+        y_dates_test,
+    ) = train_test_split(
+        input_prices,
+        next_n_day_prices,
+        input_dates,
+        next_n_day_dates,
+        test_size=test_size,
+    )
+
+
 def get_backtesting_data(
     df_stock: pd.DataFrame, parsed_end_date: datetime, pred_days: int
 ) -> Tuple[pd.DataFrame, pd.DataFrame, bool]:
@@ -84,19 +131,12 @@ def plot_pred(
     df_stock: pd.DataFrame,
     df_future: pd.DataFrame,
     df_pred: pd.DataFrame,
+    title: str,
     bt_flag: bool,
 ):
     plt.figure(figsize=plot_autoscale(), dpi=PLOT_DPI)
     plt.plot(df_stock.index, df_stock["5. adjusted close"], lw=2)
-    # BACKTESTING
-    if ns_parser.s_end_date:
-        plt.title(
-            f"BACKTESTING: {ns_parser.n_length} Moving Average on {s_ticker} - {ns_parser.n_days} days prediction"
-        )
-    else:
-        plt.title(
-            f"{ns_parser.n_length} Moving Average on {s_ticker} - {ns_parser.n_days} days prediction"
-        )
+    plt.title(title)
     plt.xlim(df_stock.index[0], get_next_stock_market_days(df_pred.index[-1], 1)[-1])
     plt.xlabel("Time")
     plt.ylabel("Share Price ($)")

gamestonk_terminal/prediction_techniques/sma_view.py

@@ -139,8 +139,8 @@ def simple_moving_average(other_args: List[str], s_ticker: str, df_stock: pd.Dat
         plt.grid(b=True, which="major", color="#666666", linestyle="-")
         plt.minorticks_on()
         plt.grid(b=True, which="minor", color="#999999", linestyle="-", alpha=0.2)
-        df_ma = df_stock["5. adjusted close"].rolling(window=ns_parser.n_length).mean()
-        plt.plot(df_ma.index, df_ma, lw=2, linestyle="--", c="tab:orange")
+        # df_ma = df_stock["5. adjusted close"].rolling(window=ns_parser.n_length).mean()
+        # plt.plot(df_ma.index, df_ma, lw=2, linestyle="--", c="tab:orange")
         plt.plot(
             [df_stock.index[-1], df_pred.index[0]],
             [df_stock["5. adjusted close"].values[-1], df_pred.values[0]],