Time Series stock price prediction
from enum import Enum


class PredictionType(str, Enum):
    NEXT_POINT = 'Next point prediction'
    MULTIPLE_SEQUENCES = 'Multiple sequence prediction'
    FULL_SEQUENCE = 'Full sequence prediction'


# select prediction type
PREDICTION_TYPE = PredictionType.MULTIPLE_SEQUENCES

COMPANY_TICKER = "^FTSE"
TRAIN_SPLIT = 0.85

STEPS = 6
LEARNING_RATE = 0.8
HIDDEN_LAYERS = 64
WINDOW_SIZE = 50

config.py
RoryCoulson/LSTM-Stock-Price-Prediction-PyTorch
import torch
import numpy as np
import yfinance as yf

from sklearn.preprocessing import MinMaxScaler
from plot import plot_normalized_data, plot_stock_price_history, plot_train_test_split
from config import COMPANY_TICKER, PREDICTION_TYPE, TRAIN_SPLIT, PredictionType, WINDOW_SIZE


def normalize(train_data, test_data, plot=False):
    # normalizes the data in separate windows to allow lower price periods to still have a significance
    normalize_window = 2000
    scaler = MinMaxScaler()
    for i in range(0, len(train_data), normalize_window):
        train_data[i:i+normalize_window,
                   :] = scaler.fit_transform(train_data[i:i+normalize_window, :])
    # apply the normalized transformation to the test data afterwards
    test_data = scaler.transform(test_data)

    if plot:
        plot_normalized_data(train_data, test_data)
    return train_data.reshape(-1), test_data.reshape(-1)


def train_test_split(prices, currency, plot=False):
    train_size = int(np.round(TRAIN_SPLIT * len(prices)))
    train_data = prices[:train_size].reshape(-1, 1)
    test_data = prices[train_size:].reshape(-1, 1)
    if plot:
        plot_train_test_split(train_data, test_data, currency)
    return train_data, test_data


def get_train_data_split(train_data):
    # splits out the train data into train_input and train_target
    train_windows = np.empty((len(train_data) - WINDOW_SIZE + 1, WINDOW_SIZE), dtype=np.float32)
    # fill train_windows using a sliding window across the train data
    for i in range(len(train_data) - WINDOW_SIZE + 1):
        train_windows[i] = train_data[i:i+WINDOW_SIZE]

    # all train windows but excludes the final price in each
    train_input = torch.from_numpy(train_windows[:, :-1]).reshape(-1, WINDOW_SIZE-1, 1)
    # all train windows but excludes the first price in each, corresponding to the next timestep in the sequence
    train_target = torch.from_numpy(train_windows[:, 1:]).reshape(-1, WINDOW_SIZE-1, 1)

    return train_input, train_target


def get_test_data_split(test_data):
    # splits out the test data into test_input and test_target
    test_prices = test_data.astype(np.float32)
    match PREDICTION_TYPE:
        case PredictionType.NEXT_POINT:
            # makes predictions using the true test data windows
            test_windows = np.empty(
                (len(test_data) - WINDOW_SIZE + 1, WINDOW_SIZE), dtype=np.float32)
            for i in range(len(test_data) - WINDOW_SIZE + 1):
                test_windows[i] = test_data[i:i+WINDOW_SIZE]
            # all test window excluding the last point
            test_input = torch.from_numpy(test_windows[:, :-1]).reshape(-1, WINDOW_SIZE-1, 1)
            # all values after the first input window - equivalent to the last point of every test_input window
            test_target = torch.from_numpy(test_prices[WINDOW_SIZE-1:])
            return test_input, test_target

        case PredictionType.MULTIPLE_SEQUENCES:
            # full test data since this will be split into sequences when predicting
            test_input = torch.from_numpy(test_prices)
            # starts at future prediction starting point, excluding predictions made with test data as input
            test_target = torch.from_numpy(test_prices[WINDOW_SIZE:])
            return test_input, test_target

        case PredictionType.FULL_SEQUENCE:
            # only use the first input window for full sequence prediction
            test_input = torch.from_numpy(test_prices[:WINDOW_SIZE-1]).reshape(1, -1, 1)
            # starts at future prediction starting point, excluding predictions made with test data as input
            test_target = torch.from_numpy(test_prices[WINDOW_SIZE:])
            return test_input, test_target


def load_data(start='2000-01-01', end='2023-09-01', plot=False):
    stock_df = yf.download(COMPANY_TICKER, start=start, end=end)
    currency = yf.Ticker(COMPANY_TICKER).info['currency']

    if plot:
        plot_stock_price_history(currency=currency, stock_df=stock_df)

    prices = stock_df['Close'].values

    train_data, test_data = train_test_split(prices, currency, plot=True)
    train_data, test_data = normalize(train_data, test_data, plot=True)

    train_input, train_target = get_train_data_split(train_data)
    test_input, test_target = get_test_data_split(test_data)

    return train_input, train_target, test_input, test_target

data_loader.py
RoryCoulson/LSTM-Stock-Price-Prediction-PyTorch
import torch
import torch.nn as nn

from config import WINDOW_SIZE


class LSTM(nn.Module):
    def __init__(self, hidden_layers=64, num_layers=2):
        super(LSTM, self).__init__()
        self.hidden_layers = hidden_layers
        self.num_layers = num_layers

        # using dropout as a regularization technique to prevent overfitting
        self.lstm = nn.LSTM(1, self.hidden_layers, self.num_layers, batch_first=True, dropout=0.2)
        # self.gru = nn.GRU(1, self.hidden_layers, self.num_layers, batch_first=True, dropout=0.2) # gru

        self.linear = nn.Linear(self.hidden_layers, 1)

    def forward(self, input, future_preds=0):
        outputs = []
        h_t = torch.zeros(self.num_layers, input.size(0), self.hidden_layers, dtype=torch.float32)
        c_t = torch.zeros(self.num_layers, input.size(0), self.hidden_layers, dtype=torch.float32)

        for input_t in input.split(1, dim=1):
            out, _ = self.lstm(input_t, (h_t, c_t))
            # out, _ = self.gru(input_t, h_t) # gru
            output = self.linear(out)
            outputs.append(output)

        # if future_preds then the model will make predictions using output as input
        for _ in range(future_preds):
            out, _ = self.lstm(output, (h_t, c_t))
            # out, _ = self.gru(output, h_t) # gru
            output = self.linear(out)
            outputs.append(output)

        outputs = torch.cat(outputs, dim=1)
        return outputs


def predict_next_point(test_input, test_target, model, criterion):
    # predicts the next point for each test data sliding window, only using test data as input
    prediction = model(test_input)
    # the prediction is all the last values for every predicted sliding window
    prediction = prediction[:, -1].reshape(-1)
    test_loss = criterion(prediction, test_target)
    print('test loss:', test_loss.item())
    prediction = prediction.detach().numpy()
    return prediction


def predict_sequence(test_input, test_target, i, model, criterion):
    # predicts a sequence of future steps given the ith test window
    prediction = model(
        test_input[i * WINDOW_SIZE: (i + 1) * WINDOW_SIZE - 1].reshape(1, -1, 1),
        future_preds=WINDOW_SIZE)
    # filter to only include future predictions
    prediction = prediction[:, WINDOW_SIZE-1:].reshape(-1)
    test_loss = criterion(prediction, test_target[i*WINDOW_SIZE:(i+1)*WINDOW_SIZE])
    print('test loss:', test_loss.item())
    prediction = prediction.detach().numpy()
    return prediction


def predict_full_sequence(test_input, test_target, model, criterion):
    # predicts all future steps given the first test window
    prediction = model(test_input, future_preds=len(test_target))
    # filter to only include future predictions
    prediction = prediction[:, WINDOW_SIZE-1:].reshape(-1)
    test_loss = criterion(prediction, test_target)
    print('test loss:', test_loss.item())
    prediction = prediction.detach().numpy()
    return prediction

model.py
RoryCoulson/LSTM-Stock-Price-Prediction-PyTorch
import matplotlib.pyplot as plt
from matplotlib.pyplot import cm
import numpy as np
from config import COMPANY_TICKER, PREDICTION_TYPE

FIGURE_SIZE = (13, 8)


def plot_stock_price_history(currency, stock_df):
    plt.figure(figsize=FIGURE_SIZE)
    stock_df['Close'].plot()
    plt.title(f'{COMPANY_TICKER} Stock Price History')
    plt.xlabel('Date')
    plt.ylabel(f'Price ({currency})')
    plt.show()


def plot_train_test_split(train_data, test_data, currency):
    plt.figure(figsize=FIGURE_SIZE)
    plt.title('Train test split')
    plt.plot(np.arange(len(train_data)), train_data, label='Train data')
    plt.plot(np.arange(len(train_data), len(train_data) + len(test_data)), test_data,
             label='Test data')
    plt.ylabel(f'Price ({currency})')
    plt.legend()
    plt.show()


def plot_normalized_data(train_data, test_data):
    plt.figure(figsize=FIGURE_SIZE)
    plt.title('Normalized Train and Test data')
    plt.plot(np.arange(len(train_data)), train_data, label='Normalized Train data')
    plt.plot(np.arange(len(train_data), len(
        train_data) + len(test_data)), test_data, label='Normalized Test data')
    plt.legend()
    plt.show()


def plot_multiple_sequence_predictions(
        multiple_sequence_predictions, test_target, step, show=False):
    colours = cm.tab20(np.linspace(0, 1, len(multiple_sequence_predictions)+2))
    plt.figure(figsize=FIGURE_SIZE)
    plt.title(f'{PREDICTION_TYPE} {COMPANY_TICKER} results')
    plt.plot(test_target.detach().numpy(), label='Test data')
    for c, (x, y_pred) in zip(colours[2:], multiple_sequence_predictions):
        plt.plot(x, y_pred, label=f'Prediction', c=c)
    plt.legend()
    save_plot(step, show)


def plot_predictions(prediction, test_target, step, show=False):
    plt.figure(figsize=FIGURE_SIZE)
    plt.title(f'{PREDICTION_TYPE} {COMPANY_TICKER} results')
    plt.plot(test_target.detach().numpy(), label='Test data')
    plt.plot(prediction, label='Prediction')
    plt.legend()
    save_plot(step, show)


def save_plot(step, show):
    path = f'prediction_plots/step_{step}_predictions.png'
    plt.savefig(path)
    if show:
        plt.show()
    plt.close()

plot.py
RoryCoulson/LSTM-Stock-Price-Prediction-PyTorch
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

from config import HIDDEN_LAYERS, LEARNING_RATE, PREDICTION_TYPE, STEPS, WINDOW_SIZE, PredictionType
from data_loader import load_data
from model import LSTM, predict_full_sequence, predict_next_point, predict_sequence
from plot import plot_multiple_sequence_predictions, plot_predictions

# create directory for prediction plots
Path("prediction_plots/").mkdir(parents=True, exist_ok=True)

np.random.seed(0)
torch.manual_seed(0)

train_input, train_target, test_input, test_target = load_data(plot=True)

model = LSTM(hidden_layers=HIDDEN_LAYERS)
criterion = nn.MSELoss()
# LBFGS optimizer since it can load the whole data for training
optimizer = optim.LBFGS(model.parameters(), lr=LEARNING_RATE)

for step in range(STEPS):

    def closure():
        optimizer.zero_grad()
        out = model(train_input)
        train_loss = criterion(out, train_target)
        print(f'train loss: {train_loss.item()}')
        train_loss.backward()
        return train_loss
    optimizer.step(closure)

    with torch.no_grad():
        match PREDICTION_TYPE:
            case PredictionType.NEXT_POINT:
                prediction = predict_next_point(
                    test_input, test_target, model, criterion)
                plot_predictions(prediction, test_target, step)
            case PredictionType.MULTIPLE_SEQUENCES:
                multiple_sequence_predictions = []
                # non overlapping windows of test data to use for each sequence prediction
                for i in range(int(len(test_target)/WINDOW_SIZE)):
                    prediction = predict_sequence(
                        test_input, test_target, i, model, criterion)
                    multiple_sequence_predictions.append(
                        (np.arange(i*WINDOW_SIZE, (i+1)*WINDOW_SIZE), prediction))
                plot_multiple_sequence_predictions(
                    multiple_sequence_predictions, test_target, step)
            case PredictionType.FULL_SEQUENCE:
                prediction = predict_full_sequence(
                    test_input, test_target, model, criterion)
                plot_predictions(prediction, test_target, step)

run.py
RoryCoulson/LSTM-Stock-Price-Prediction-PyTorch