mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-10-28 21:12:16 +00:00
195 lines
6.8 KiB
Python
195 lines
6.8 KiB
Python
# pylint: disable=too-many-arguments
|
||
"""Nhits Model"""
|
||
__docformat__ = "numpy"
|
||
|
||
import logging
|
||
import warnings
|
||
from typing import Union, Optional, List, Tuple
|
||
|
||
|
||
import pandas as pd
|
||
|
||
from darts import TimeSeries
|
||
from darts.models.forecasting.nhits import NHiTSModel
|
||
from darts.utils.likelihood_models import GaussianLikelihood
|
||
from openbb_terminal.decorators import log_start_end
|
||
from openbb_terminal.forecast import helpers
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@log_start_end(log=logger)
|
||
def get_nhits_data(
|
||
data: Union[pd.Series, pd.DataFrame],
|
||
target_column: str = "close",
|
||
n_predict: int = 5,
|
||
train_split: float = 0.85,
|
||
past_covariates: str = None,
|
||
forecast_horizon: int = 5,
|
||
input_chunk_length: int = 14,
|
||
output_chunk_length: int = 5,
|
||
num_stacks: int = 3,
|
||
num_blocks: int = 1,
|
||
num_layers: int = 2,
|
||
layer_widths: int = 512,
|
||
pooling_kernel_sizes: Optional[Tuple[Tuple[int]]] = None,
|
||
n_freq_downsample: Optional[Tuple[Tuple[int]]] = None,
|
||
dropout: float = 0.1,
|
||
activation: str = "ReLU",
|
||
max_pool_1d: bool = True,
|
||
batch_size: int = 32,
|
||
n_epochs: int = 100,
|
||
learning_rate: float = 1e-3,
|
||
model_save_name: str = "nhits_model",
|
||
force_reset: bool = True,
|
||
save_checkpoints: bool = True,
|
||
) -> Tuple[
|
||
Optional[List[TimeSeries]],
|
||
Optional[List[TimeSeries]],
|
||
Optional[List[TimeSeries]],
|
||
Optional[float],
|
||
Optional[type[NHiTSModel]],
|
||
]:
|
||
"""Performs Nhits forecasting
|
||
|
||
Parameters
|
||
----------
|
||
data: Union[pd.Series, pd.DataFrame]
|
||
Input Data
|
||
target_column: str
|
||
Target column to forecast. Defaults to "close".
|
||
n_predict: int
|
||
Days to predict. Defaults to 5.
|
||
train_split: float
|
||
Train/val split. Defaults to 0.85.
|
||
past_covariates: str
|
||
Multiple secondary columns to factor in when forecasting. Defaults to None.
|
||
forecast_horizon: int
|
||
Forecast horizon when performing historical forecasting. Defaults to 5.
|
||
input_chunk_length: int
|
||
Number of past time steps that are fed to the forecasting module at prediction time. Defaults to 14.
|
||
output_chunk_length: int
|
||
The length of the forecast of the model. Defaults to 5.
|
||
num_stacks: int
|
||
The number of stacks that make up the whole model.
|
||
num_blocks: int
|
||
The number of blocks making up every stack.
|
||
num_layers: int
|
||
The number of fully connected layers preceding the final forking layers in each block
|
||
of every stack.
|
||
layer_widths: int
|
||
Determines the number of neurons that make up each fully connected layer in each
|
||
block of every stack. If a list is passed, it must have a length equal to num_stacks
|
||
and every entry in that list corresponds to the layer width of the corresponding stack.
|
||
If an integer is passed, every stack will have blocks with FC layers of the same width.
|
||
pooling_kernel_size: Optional[Tuple[Tuple[int]]]:
|
||
If set, this parameter must be a tuple of tuples, of size (num_stacks x num_blocks),
|
||
specifying the kernel size for each block in each stack used for the input pooling
|
||
layer. If left to None, some default values will be used based on input_chunk_length.
|
||
n_freq_downsample: Optional[Tuple[Tuple[int]]]
|
||
If set, this parameter must be a tuple of tuples, of size (num_stacks x num_blocks),
|
||
specifying the downsampling factors before interpolation, for each block in each stack.
|
||
If left to None, some default values will be used based on output_chunk_length.
|
||
dropout: float
|
||
The dropout probability to be used in fully connected layers.
|
||
activation: str
|
||
Supported activations: [‘ReLU’,’RReLU’, ‘PReLU’, ‘Softplus’, ‘Tanh’, ‘SELU’, ‘LeakyReLU’, ‘Sigmoid’]
|
||
max_pool_1d: bool
|
||
Use max_pool_1d pooling. False uses AvgPool1d.
|
||
batch_size: int
|
||
Number of time series (input and output sequences) used in each training pass. Defaults to 32.
|
||
n_epochs: int
|
||
Number of epochs over which to train the model. Defaults to 100.
|
||
learning_rate: float
|
||
Defaults to 1e-3.
|
||
model_save_name: str
|
||
Name for model. Defaults to "brnn_model".
|
||
force_reset: bool
|
||
If set to True, any previously-existing model with the same name will be reset (all checkpoints will be
|
||
discarded). Defaults to True.
|
||
save_checkpoints: bool
|
||
Whether or not to automatically save the untrained model and checkpoints from training. Defaults to True.
|
||
|
||
Returns
|
||
-------
|
||
Tuple[List[TimeSeries], List[TimeSeries], List[TimeSeries], Optional[float], type[NHiTSModel]]
|
||
Adjusted Data series,
|
||
Historical forecast by best RNN model,
|
||
list of Predictions,
|
||
Mean average precision error,
|
||
Best BRNN Model.
|
||
"""
|
||
|
||
# TODO Check if torch GPU AVAILABLE
|
||
|
||
use_scalers = True
|
||
probabilistic = False
|
||
|
||
scaler, ticker_series = helpers.get_series(
|
||
data, target_column, is_scaler=use_scalers
|
||
)
|
||
train, val = ticker_series.split_before(train_split)
|
||
valid = helpers.check_data_length(
|
||
train, val, input_chunk_length, output_chunk_length
|
||
)
|
||
if not valid:
|
||
return [], [], [], None, None
|
||
|
||
(
|
||
past_covariate_whole,
|
||
past_covariate_train,
|
||
past_covariate_val,
|
||
) = helpers.past_covs(past_covariates, data, train_split, use_scalers)
|
||
|
||
# Early Stopping
|
||
|
||
nhits_model = NHiTSModel(
|
||
input_chunk_length=input_chunk_length,
|
||
output_chunk_length=output_chunk_length,
|
||
num_stacks=num_stacks,
|
||
num_blocks=num_blocks,
|
||
num_layers=num_layers,
|
||
layer_widths=layer_widths,
|
||
pooling_kernel_sizes=pooling_kernel_sizes,
|
||
n_freq_downsample=n_freq_downsample,
|
||
dropout=dropout,
|
||
activation=activation,
|
||
MaxPool1d=max_pool_1d,
|
||
n_epochs=n_epochs,
|
||
batch_size=batch_size,
|
||
optimizer_kwargs={"lr": learning_rate},
|
||
model_name=model_save_name,
|
||
random_state=42,
|
||
pl_trainer_kwargs=helpers.get_pl_kwargs(accelerator="cpu"),
|
||
force_reset=force_reset,
|
||
save_checkpoints=save_checkpoints,
|
||
likelihood=GaussianLikelihood(),
|
||
)
|
||
|
||
# fit model on train series for historical forecasting
|
||
with warnings.catch_warnings():
|
||
warnings.simplefilter("ignore")
|
||
helpers.fit_model(
|
||
nhits_model,
|
||
train,
|
||
val,
|
||
past_covariate_train,
|
||
past_covariate_val,
|
||
)
|
||
best_model = NHiTSModel.load_from_checkpoint(model_name=model_save_name, best=True)
|
||
|
||
# Showing historical backtesting without retraining model (too slow)
|
||
return helpers.get_prediction(
|
||
"NHITS",
|
||
probabilistic,
|
||
use_scalers,
|
||
scaler,
|
||
past_covariates,
|
||
best_model,
|
||
ticker_series,
|
||
past_covariate_whole,
|
||
train_split,
|
||
forecast_horizon,
|
||
n_predict,
|
||
)
|