refactoring

Author: JordiCorbilla

stock_prediction_class.py

@@ -0,0 +1,46 @@
+# Copyright 2020 Jordi Corbilla. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+class StockPrediction:
+    def __init__(self, ticker, start_date, validation_date, project_folder):
+        self._ticker = ticker
+        self._start_date = start_date
+        self._validation_date = validation_date
+        self._project_folder = project_folder
+
+    def get_ticker(self):
+        return self._ticker
+
+    def set_ticker(self, value):
+        self._ticker = value
+
+    def get_start_date(self):
+        return self._start_date
+
+    def set_start_date(self, value):
+        self._start_date = value
+
+    def get_validation_date(self):
+        return self._validation_date
+
+    def set_validation_date(self, value):
+        self._validation_date = value
+
+    def get_project_folder(self):
+        return self._project_folder
+
+    def set_project_folder(self, value):
+        self._project_folder = value

stock_prediction_deep_learning.py

@@ -16,41 +16,26 @@
 import secrets
 import pandas as pd
 import tensorflow as tf
-from sklearn.preprocessing import MinMaxScaler
 from datetime import datetime
-import yfinance as yf
+
+from stock_prediction_class import StockPrediction
 from stock_prediction_lstm import LongShortTermMemory
 from stock_prediction_numpy import StockData
 from stock_prediction_plotter import Plotter
 
 os.environ["PATH"] += os.pathsep + 'C:/Program Files (x86)/Graphviz2.38/bin/'
 
 
-def train_LSTM_network(start_date, ticker, validation_date):
-    min_max = MinMaxScaler(feature_range=(0, 1))
-    sec = yf.Ticker(ticker)
-    end_date = datetime.today()
-    print('End Date: ' + end_date.strftime("%Y-%m-%d"))
-    data = yf.download([ticker], start=start_date, end=end_date)[['Close']]
-    data = data.reset_index()
-    print(data)
-
-    plotter = Plotter(True, project_folder, sec.info['shortName'], sec.info['currency'], STOCK_TICKER)
+def train_LSTM_network(stock):
+    data = StockData(stock)
 
-    training_data = data[data['Date'] < validation_date].copy()
-    test_data = data[data['Date'] >= validation_date].copy()
-    training_data = training_data.set_index('Date')
-    # Set the data frame index using column Date
-    test_data = test_data.set_index('Date')
-    print(test_data)
-    plotter.plot_histogram_data_split(training_data, test_data, validation_date)
+    plotter = Plotter(True, stock.get_project_folder(), data.get_stock_short_name(), data.get_stock_currency(), stock.get_ticker())
 
-    data = StockData()
-    (x_train, y_train), (x_test, y_test) = data.to_numpy(TIME_STEPS, min_max, training_data, test_data)
+    (x_train, y_train), (x_test, y_test), (min_max, test_data) = data.download_transform_to_numpy(TIME_STEPS)
 
     print(x_test)
 
-    lstm = LongShortTermMemory(project_folder)
+    lstm = LongShortTermMemory(stock.get_project_folder())
     model = lstm.create_model(x_train)
 
     defined_metrics = [
@@ -63,7 +48,7 @@ def train_LSTM_network(start_date, ticker, validation_date):
     history = model.fit(x_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_data=(x_test, y_test),
                         callbacks=[callback])
     print("saving weights")
-    model.save(os.path.join(project_folder, 'model_weights.h5'))
+    model.save(os.path.join(stock.get_project_folder(), 'model_weights.h5'))
     plotter.plot_loss(history)
     plotter.plot_mse(history)
 
@@ -77,7 +62,7 @@ def train_LSTM_network(start_date, ticker, validation_date):
     test_predictions_baseline = model.predict(x_test)
     test_predictions_baseline = min_max.inverse_transform(test_predictions_baseline)
     test_predictions_baseline = pd.DataFrame(test_predictions_baseline)
-    test_predictions_baseline.to_csv(os.path.join(project_folder, 'predictions.csv'))
+    test_predictions_baseline.to_csv(os.path.join(stock.get_project_folder(), 'predictions.csv'))
 
     test_predictions_baseline.rename(columns={0: STOCK_TICKER + '_predicted'}, inplace=True)
     test_predictions_baseline = test_predictions_baseline.round(decimals=0)
@@ -105,9 +90,10 @@ def train_LSTM_network(start_date, ticker, validation_date):
     print('Validation Date: ' + STOCK_START_DATE.strftime("%Y-%m-%d"))
     print('Generating folder: ' + TOKEN)
     # create project run folder
-    project_folder = os.path.join(os.getcwd(), TOKEN)
-    if not os.path.exists(project_folder):
-        os.makedirs(project_folder)
+    PROJECT_FOLDER = os.path.join(os.getcwd(), TOKEN)
+    if not os.path.exists(PROJECT_FOLDER):
+        os.makedirs(PROJECT_FOLDER)
 
+    stock_prediction = StockPrediction(STOCK_TICKER, STOCK_START_DATE, STOCK_VALIDATION_DATE, PROJECT_FOLDER)
     # Execute Deep Learning model
-    train_LSTM_network(STOCK_START_DATE, STOCK_TICKER, STOCK_VALIDATION_DATE)
+    train_LSTM_network(stock_prediction)

stock_prediction_numpy.py

@@ -13,19 +13,50 @@
 # limitations under the License.
 # ==============================================================================
 import numpy as np
-import pandas as pd
 from datetime import timedelta
 import random
+import pandas as pd
+from sklearn.preprocessing import MinMaxScaler
+from datetime import datetime
+import yfinance as yf
+from stock_prediction_plotter import Plotter
 
 
 class StockData:
+    def __init__(self, stock):
+        self._stock = stock
+        self._sec = yf.Ticker(self._stock.get_ticker())
+
     def __data_verification(self, train):
         print('mean:', train.mean(axis=0))
         print('max', train.max())
         print('min', train.min())
         print('Std dev:', train.std(axis=0))
 
-    def to_numpy(self, time_steps, min_max, training_data, test_data):
+    def get_stock_short_name(self):
+        return self._sec.info['shortName']
+
+    def get_stock_currency(self):
+        return self._sec.info['currency']
+
+    def download_transform_to_numpy(self, time_steps):
+        min_max = MinMaxScaler(feature_range=(0, 1))
+        end_date = datetime.today()
+        print('End Date: ' + end_date.strftime("%Y-%m-%d"))
+        data = yf.download([self._stock.get_ticker()], start=self._stock.get_start_date(), end=end_date)[['Close']]
+        data = data.reset_index()
+        print(data)
+
+        plotter = Plotter(True, self._stock.get_project_folder(), self._sec.info['shortName'], self._sec.info['currency'], self._stock.get_ticker())
+
+        training_data = data[data['Date'] < self._stock.get_validation_date()].copy()
+        test_data = data[data['Date'] >= self._stock.get_validation_date()].copy()
+        training_data = training_data.set_index('Date')
+        # Set the data frame index using column Date
+        test_data = test_data.set_index('Date')
+        print(test_data)
+        plotter.plot_histogram_data_split(training_data, test_data, self._stock.get_validation_date())
+
         train_scaled = min_max.fit_transform(training_data)
         self.__data_verification(train_scaled)
 
@@ -52,7 +83,7 @@ def to_numpy(self, time_steps, min_max, training_data, test_data):
 
         x_test, y_test = np.array(x_test), np.array(y_test)
         x_test = np.reshape(x_test, (x_test.shape[0], x_test.shape[1], 1))
-        return (x_train, y_train), (x_test, y_test)
+        return (x_train, y_train), (x_test, y_test), (min_max, test_data)
 
     def __daterange(self, start_date, end_date):
         for n in range(int((end_date - start_date).days)):