File size: 9,343 Bytes
f828b45
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
import os
import pickle
import joblib
import pandas as pd
import numpy as np
from statsmodels.tsa.arima.model import ARIMA
from sklearn.preprocessing import MinMaxScaler
import warnings
warnings.filterwarnings('ignore')

class StockForecaster:
    """
    A comprehensive stock forecasting class that combines ARIMA and LSTM models.
    """
    
    def __init__(self):
        self.arima_model = None
        self.lstm_model = None
        self.scaler = MinMaxScaler()
        self.optimal_arima_params = (1, 1, 1)  # Default parameters
        
    def preprocess_data(self, data):
        """
        Preprocess stock data for forecasting.
        
        Args:
            data (pd.DataFrame): Stock data with 'date' and 'close' columns
            
        Returns:
            pd.Series: Processed time series data
        """
        if isinstance(data, pd.DataFrame):
            if 'date' in data.columns:
                data['date'] = pd.to_datetime(data['date'])
                data = data.set_index('date').sort_index()
            
            if 'close' in data.columns:
                return data['close'].dropna()
        
        return data.dropna()
    
    def find_optimal_arima_params(self, ts_data, max_p=3, max_d=2, max_q=3):
        """
        Find optimal ARIMA parameters using AIC criterion.
        
        Args:
            ts_data (pd.Series): Time series data
            max_p, max_d, max_q (int): Maximum values for ARIMA parameters
            
        Returns:
            tuple: Optimal (p, d, q) parameters
        """
        best_aic = np.inf
        best_params = (1, 1, 1)
        
        for p in range(max_p + 1):
            for d in range(max_d + 1):
                for q in range(max_q + 1):
                    try:
                        model = ARIMA(ts_data, order=(p, d, q))
                        fitted_model = model.fit()
                        aic = fitted_model.aic
                        
                        if aic < best_aic:
                            best_aic = aic
                            best_params = (p, d, q)
                    except:
                        continue
        
        self.optimal_arima_params = best_params
        return best_params
    
    def train_arima(self, ts_data):
        """
        Train ARIMA model on time series data.
        
        Args:
            ts_data (pd.Series): Time series data
            
        Returns:
            statsmodels.ARIMAResults: Fitted ARIMA model
        """
        try:
            # Find optimal parameters if not set
            if self.optimal_arima_params == (1, 1, 1):
                self.find_optimal_arima_params(ts_data)
            
            # Fit ARIMA model
            arima_model = ARIMA(ts_data, order=self.optimal_arima_params)
            self.arima_model = arima_model.fit()
            
            return self.arima_model
            
        except Exception as e:
            print(f"ARIMA training error: {e}")
            return None
    
    def create_lstm_sequences(self, data, sequence_length=60):
        """
        Create sequences for LSTM training.
        
        Args:
            data (np.array): Scaled time series data
            sequence_length (int): Length of input sequences
            
        Returns:
            tuple: (X, y) arrays for LSTM training
        """
        X, y = [], []
        for i in range(sequence_length, len(data)):
            X.append(data[i-sequence_length:i])
            y.append(data[i])
        return np.array(X), np.array(y)
    
    def train_simple_lstm(self, ts_data, sequence_length=60):
        """
        Train a simplified LSTM model or use trend-based prediction.
        
        Args:
            ts_data (pd.Series): Time series data
            sequence_length (int): Sequence length for LSTM
            
        Returns:
            dict: Model information and scaler
        """
        try:
            # Scale the data
            scaled_data = self.scaler.fit_transform(ts_data.values.reshape(-1, 1))
            
            # For demo purposes, we'll use a trend-based approach
            # In production, you'd train an actual LSTM here
            
            self.lstm_model = {
                'type': 'trend_based',
                'data': scaled_data,
                'sequence_length': sequence_length
            }
            
            return self.lstm_model
            
        except Exception as e:
            print(f"LSTM training error: {e}")
            return None
    
    def forecast_arima(self, steps=30):
        """
        Generate ARIMA forecast.
        
        Args:
            steps (int): Number of steps to forecast
            
        Returns:
            np.array: Forecasted values
        """
        if self.arima_model is None:
            raise ValueError("ARIMA model not trained. Call train_arima first.")
        
        try:
            forecast = self.arima_model.forecast(steps=steps)
            return forecast.values if hasattr(forecast, 'values') else forecast
        except Exception as e:
            print(f"ARIMA forecast error: {e}")
            return None
    
    def forecast_lstm(self, steps=30):
        """
        Generate LSTM forecast (simplified trend-based approach).
        
        Args:
            steps (int): Number of steps to forecast
            
        Returns:
            np.array: Forecasted values
        """
        if self.lstm_model is None:
            raise ValueError("LSTM model not trained. Call train_simple_lstm first.")
        
        try:
            # Simple trend-based forecast for demo
            scaled_data = self.lstm_model['data']
            sequence_length = self.lstm_model['sequence_length']
            
            # Calculate recent trend
            recent_data = scaled_data[-min(10, len(scaled_data)):]
            recent_trend = np.mean(np.diff(recent_data.flatten()))
            
            # Generate forecast
            last_value = scaled_data[-1][0]
            forecast_scaled = []
            
            for i in range(steps):
                next_val = last_value + recent_trend * (i + 1) * 0.1
                forecast_scaled.append([next_val])
            
            # Inverse transform
            forecast = self.scaler.inverse_transform(forecast_scaled).flatten()
            return forecast
            
        except Exception as e:
            print(f"LSTM forecast error: {e}")
            return None
    
    def save_models(self, filepath_prefix):
        """
        Save trained models to disk.
        
        Args:
            filepath_prefix (str): Prefix for saved model files
        """
        try:
            if self.arima_model is not None:
                with open(f"{filepath_prefix}_arima.pkl", 'wb') as f:
                    pickle.dump(self.arima_model, f)
            
            if self.lstm_model is not None:
                joblib.dump(self.lstm_model, f"{filepath_prefix}_lstm.pkl")
            
            joblib.dump(self.scaler, f"{filepath_prefix}_scaler.pkl")
            
            print(f"Models saved with prefix: {filepath_prefix}")
            
        except Exception as e:
            print(f"Error saving models: {e}")
    
    def load_models(self, filepath_prefix):
        """
        Load trained models from disk.
        
        Args:
            filepath_prefix (str): Prefix for saved model files
        """
        try:
            # Load ARIMA model
            arima_path = f"{filepath_prefix}_arima.pkl"
            if os.path.exists(arima_path):
                with open(arima_path, 'rb') as f:
                    self.arima_model = pickle.load(f)
            
            # Load LSTM model
            lstm_path = f"{filepath_prefix}_lstm.pkl"
            if os.path.exists(lstm_path):
                self.lstm_model = joblib.load(lstm_path)
            
            # Load scaler
            scaler_path = f"{filepath_prefix}_scaler.pkl"
            if os.path.exists(scaler_path):
                self.scaler = joblib.load(scaler_path)
            
            print(f"Models loaded from prefix: {filepath_prefix}")
            
        except Exception as e:
            print(f"Error loading models: {e}")

# Example usage and model training script
if __name__ == "__main__":
    # This would be run to pre-train models
    forecaster = StockForecaster()
    
    # Example with sample data
    dates = pd.date_range('2020-01-01', periods=1000, freq='D')
    prices = 100 + np.cumsum(np.random.randn(1000) * 0.5)
    
    sample_data = pd.DataFrame({
        'date': dates,
        'close': prices
    })
    
    # Preprocess and train
    ts_data = forecaster.preprocess_data(sample_data)
    
    print("Training ARIMA model...")
    forecaster.train_arima(ts_data)
    
    print("Training LSTM model...")
    forecaster.train_simple_lstm(ts_data)
    
    # Generate forecasts
    arima_forecast = forecaster.forecast_arima(30)
    lstm_forecast = forecaster.forecast_lstm(30)
    
    print(f"ARIMA forecast shape: {arima_forecast.shape if arima_forecast is not None else 'None'}")
    print(f"LSTM forecast shape: {lstm_forecast.shape if lstm_forecast is not None else 'None'}")
    
    # Save models
    forecaster.save_models("models/stock_forecaster")