Add Linear Regression model to prediction techniques menu

prediction_techniques/lr.py

@@ -0,0 +1,66 @@
+import argparse
+import numpy as np
+from stock_market_helper_funcs import *
+import pandas as pd
+import matplotlib.pyplot as plt
+from pandas.plotting import register_matplotlib_converters
+register_matplotlib_converters()
+from TimeSeriesCrossValidation import splitTrain
+from sklearn import linear_model
+
+# ----------------------------------------------------- LR -----------------------------------------------------
+def linear_regression(l_args, s_ticker, s_interval, df_stock):
+    parser = argparse.ArgumentParser(prog='linear_regression',
+                                     description="""Linear regression attempts to model the relationship between 
+                                     two variables by fitting a linear equation to observed data. One variable is 
+                                     considered to be an explanatory variable, and the other is considered to be a 
+                                     dependent variable. """)
+
+    parser.add_argument('-i', "--input", action="store", dest="n_inputs", type=check_positive, default=40, help='number of days to use for prediction')
+    parser.add_argument('-d', "--days", action="store", dest="n_days", type=check_positive, default=5, help='prediction days')
+    parser.add_argument('-j', "--jumps", action="store", dest="n_jumps", type=check_non_negative, default=0, help='number of jumps in training data')
+
+    try:
+        (ns_parser, l_unknown_args) = parser.parse_known_args(l_args)
+
+        if l_unknown_args:
+            print(f"The following args couldn't be interpreted: {l_unknown_args}\n")
+            return
+
+        # Split training data
+        stock_x, stock_y = splitTrain.split_train(df_stock['4. close'].values, ns_parser.n_inputs, ns_parser.n_days, ns_parser.n_jumps)
+
+        # Machine Learning model
+        lr = linear_model.LinearRegression(n_jobs=-1)
+        lr.fit(stock_x, stock_y)
+        
+        # Prediction data
+        l_predictions = lr.predict(df_stock['4. close'].values[-ns_parser.n_inputs:].reshape(1, -1))[0]
+        l_pred_days = get_next_stock_market_days(last_stock_day=df_stock['4. close'].index[-1], n_next_days=ns_parser.n_days)
+        df_pred = pd.Series(l_predictions, index=l_pred_days, name='Price') 
+
+        # Plotting
+        plt.plot(df_stock.index, df_stock['4. close'], lw=2)
+        plt.title(f"Linear Regression on {s_ticker} - {ns_parser.n_days} days prediction")
+        plt.xlim(df_stock.index[0], get_next_stock_market_days(df_pred.index[-1], 1)[-1] )
+        plt.xlabel('Time')
+        plt.ylabel('Share Price ($)')
+        plt.grid(b=True, which='major', color='#666666', linestyle='-')
+        plt.minorticks_on()
+        plt.grid(b=True, which='minor', color='#999999', linestyle='-', alpha=0.2)
+        plt.plot([df_stock.index[-1], df_pred.index[0]], [df_stock['4. close'].values[-1], df_pred.values[0]], lw=1, c='tab:green', linestyle='--')
+        plt.plot(df_pred.index, df_pred, lw=2, c='tab:green')
+        cmap = plt.get_cmap("tab10")
+        plt.axvspan(df_stock.index[-1], df_pred.index[-1], facecolor='tab:orange', alpha=0.2)
+        xmin, xmax, ymin, ymax = plt.axis()
+        plt.vlines(df_stock.index[-1], ymin, ymax, linewidth=1, linestyle='--', color='k')
+        plt.show()
+
+        # Print prediction data
+        df_pred = df_pred.apply(lambda x: f"{x:.2f} $")
+        print(df_pred.to_string())
+        print("")
+
+    except:
+        print("")
+

prediction_techniques/pred_menu.py

@@ -11,6 +11,7 @@ register_matplotlib_converters()
 
 from prediction_techniques import sma
 from prediction_techniques import knn
+from prediction_techniques import lr
 
 # -----------------------------------------------------------------------------------------------------------------------
 def print_prediction(s_ticker, s_start, s_interval):
@@ -30,6 +31,7 @@ def print_prediction(s_ticker, s_start, s_interval):
     print("")
     print("   sma         simple moving average")
     print("   knn         k-Nearest Neighbors")
+    print("   lr          linear regression")
     print("")
  
 
@@ -39,7 +41,7 @@ def pred_menu(df_stock, s_ticker, s_start, s_interval):
     # Add list of arguments that the prediction techniques parser accepts
     pred_parser = argparse.ArgumentParser(prog='technical_analysis', add_help=False)
     pred_parser.add_argument('cmd', choices=['help', 'q', 'quit',
-                                             'sma', 'knn'])
+                                             'sma', 'knn', 'lr'])
 
     print_prediction(s_ticker, s_start, s_interval)
 
@@ -75,6 +77,10 @@ def pred_menu(df_stock, s_ticker, s_start, s_interval):
         elif ns_known_args.cmd == 'knn':
             knn.k_nearest_neighbors(l_args, s_ticker, s_interval, df_stock)
 
+        # --------------------------------------------- LINEAR REGRESSION ---------------------------------------------
+        elif ns_known_args.cmd == 'lr':
+            lr.linear_regression(l_args, s_ticker, s_interval, df_stock)
+
         # ------------------------------------------------------------------------------------------------------------
         else:
             print("Command not recognized!")