Awesome
MambaStock: Selective state space model for stock prediction
Mamba (Structured state space sequence models with selection mechanism and scan module, S6) has achieved remarkable success in sequence modeling tasks. This paper proposes a Mamba-based model to predict the stock price.
Requirements
The code has been tested running under Python 3.7.4, with the following packages and their dependencies installed:
numpy==1.16.5
matplotlib==3.1.0
sklearn==0.21.3
pandas==0.25.1
pytorch==1.7.1
The stock data used in this repository was downloaded from TuShare. The stock data on TuShare are with public availability. Some code of the Mamba model is from https://github.com/alxndrTL/mamba.py
Usage
python main.py
Options
We adopt an argument parser by package argparse
in Python, and the options for running code are defined as follow:
parser = argparse.ArgumentParser()
parser.add_argument('--use-cuda', default=False,
help='CUDA training.')
parser.add_argument('--seed', type=int, default=1, help='Random seed.')
parser.add_argument('--epochs', type=int, default=100,
help='Number of epochs to train.')
parser.add_argument('--lr', type=float, default=0.01,
help='Learning rate.')
parser.add_argument('--wd', type=float, default=1e-5,
help='Weight decay (L2 loss on parameters).')
parser.add_argument('--hidden', type=int, default=16,
help='Dimension of representations')
parser.add_argument('--layer', type=int, default=2,
help='Num of layers')
parser.add_argument('--n-test', type=int, default=300,
help='Size of test set')
parser.add_argument('--ts-code', type=str, default='601988',
help='Stock code')
args = parser.parse_args()
args.cuda = args.use_cuda and torch.cuda.is_available()
Citation
@article{shi2024mamba,
title={MambaStock: Selective state space model for stock prediction},
author={Zhuangwei Shi},
journal={arXiv preprint arXiv:2402.18959},
year={2024},
}