Update app.py

app.py

@@ -2,165 +2,113 @@ from flask import Flask, request, jsonify
 import pandas as pd
 import numpy as np
 import baostock as bs
-from sklearn.preprocessing import MinMaxScaler, StandardScaler
+from sklearn.preprocessing import MinMaxScaler
 from sklearn.model_selection import train_test_split
-from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, mean_absolute_error
+from sklearn.metrics import mean_absolute_error
 from neuralprophet import NeuralProphet, set_log_level
-from torch.optim import Adam
-from torch.nn import LSTM
 import torch
 import torch.nn as nn
+from torch.optim import Adam
 import os
 
-
+# Initialize Flask app
 app = Flask(__name__)
 
-# Set log level to ERROR to suppress unnecessary warnings
+# Set log level to suppress unnecessary warnings
 set_log_level("ERROR")
 
-# Baostock API settings
-bs.login()
+# Baostock API login
+lg = bs.login()
+if lg.error_code != '0':
+    raise ConnectionError(f"Baostock login failed. Error code: {lg.error_code}, Error message: {lg.error_msg}")
 
 # Collect historical data
-data = bs.query_history_k_data_plus(
-    "sz.000001",  # Shanghai Composite Index
-    "date,open,high,low,close,volume",
-    start_date="2005-05-30",
-    end_date="2024-01-31",
-    frequency="d"
-)
-
-# Convert ResultData object to pandas DataFrame
-data_list = []
-while (data.error_code == '0') & data.next():
-    # 获取一条记录，将记录合并在一起
-    data_list.append(data.get_row_data())
-data_df = pd.DataFrame(data_list, columns=data.fields)
-
-# Convert 'open' and 'close' columns to numeric type
-data_df['open'] = pd.to_numeric(data_df['open'])
-data_df['close'] = pd.to_numeric(data_df['close'])
-
-# Filter out stocks that meet the conditions
-  data_df = data_df[(data_df["open"] >= 0.98 * data_df["close"].shift(1).fillna(0)) & (data_df["open"] <= 1.02 * data_df["close"].shift(1).fillna(0))] 
-  data_df = data_df[(data_df["high"] == data_df["close"]) & (data_df["low"] == data_df["close"])]  # limit-up condition
-  data_df = data_df[(data_df["open"]!= 0) & (data_df["close"]!= 0)]  # exclude zero prices
-
-# Check if data_df is empty before proceeding
-if data_df.empty:
-    print("Warning: data_df is empty after filtering. Check your filtering conditions.")
-else:
-    # Now use data_df (the DataFrame) in train_test_split
-    train_data, val_data = train_test_split(data_df, test_size=0.2, random_state=42)
-
-# Define the custom model
+def get_historical_data(start_date, end_date):
+    data = bs.query_history_k_data_plus(
+        "sz.000001",  # Shanghai Composite Index
+        "date,open,high,low,close,volume",
+        start_date=start_date,
+        end_date=end_date,
+        frequency="d"
+    )
+    if data.error_code != '0':
+        raise ValueError(f"Error in fetching data: {data.error_msg}")
+
+    data_list = []
+    while data.next():
+        data_list.append(data.get_row_data())
+    data_df = pd.DataFrame(data_list, columns=data.fields)
+
+    # Convert relevant columns to numeric type
+    data_df[['open', 'close', 'high', 'low', 'volume']] = data_df[['open', 'close', 'high', 'low', 'volume']].apply(pd.to_numeric, errors='coerce')
+    return data_df.dropna()
+
+# Filter stocks based on conditions
+def filter_stocks(data_df):
+    data_df = data_df[(data_df["open"] >= 0.98 * data_df["close"].shift(1).fillna(0)) & (data_df["open"] <= 1.02 * data_df["close"].shift(1).fillna(0))]
+    data_df = data_df[(data_df["high"] == data_df["close"]) & (data_df["low"] == data_df["close"]) & (data_df["open"] != 0) & (data_df["close"] != 0)]
+    return data_df
+
+# Prepare the training and validation data
+data_df = get_historical_data("2005-05-30", "2024-01-31")
+filtered_df = filter_stocks(data_df)
+
+if filtered_df.empty:
+    raise ValueError("Filtered dataset is empty. Please adjust the filtering conditions.")
+
+train_data, val_data = train_test_split(filtered_df, test_size=0.2, random_state=42)
+
+# Define custom model
 class CustomModel(nn.Module):
     def __init__(self):
         super(CustomModel, self).__init__()
         self.neural_prophet = NeuralProphet(
             n_forecasts=1,
-            n_lags=30, 
-            n_changepoints=10,  
-            changepoints_range=0.8,  
+            n_lags=30,
+            n_changepoints=10,
+            changepoints_range=0.8,
             learning_rate=1e-3,
             optimizer=Adam,
         )
-        self.lstm = LSTM(input_size=1, hidden_size=128, num_layers=1, batch_first=True)
-
-    def forward(self, x):
-        x = self.neural_prophet(x)
-        x = self.lstm(x)
-        return x
 
     def predict(self, df):
-        """
-        Custom predict method for CustomModel. Utilizes NeuralProphet's prediction.
-
-        Args:
-            df: The input DataFrame for prediction.
-
-        Returns:
-            Predictions from the NeuralProphet model.
-        """
-        # Assuming your NeuralProphet model expects a DataFrame in a specific format
-        # You might need to adjust this based on your data and model setup
-        future = self.neural_prophet.make_future_dataframe(df, periods=1)  # Adjust periods as needed
+        future = self.neural_prophet.make_future_dataframe(df, periods=1)
         forecast = self.neural_prophet.predict(future)
-        return forecast['yhat1'].values  # Or access the relevant prediction column
+        return forecast['yhat1'].values
 
-# Instantiate the model
+# Instantiate model
 model = CustomModel()
 
-# Define loss function and optimizer
-criterion = nn.BCELoss()  
-optimizer = Adam(model.parameters(), lr=1e-3)
-
-# Training loop
-def fit(model, train_data, epochs, batch_size, validation_data):
-    """
-    Custom training loop for the CustomModel.
-
-    Args:
-        model: The CustomModel instance.
-        train_data: Training data.
-        epochs: Number of training epochs.
-        batch_size: Batch size for training.
-        validation_data: Validation data.
-    """
-    for epoch in range(epochs):
-        model.train()  # Set model to training mode
-        for batch_idx, (data, target) in enumerate(train_data):  # Assuming train_data is a DataLoader
-            optimizer.zero_grad()  # Zero the gradients
-            output = model(data)  # Forward pass
-            loss = criterion(output, target)  # Calculate loss
-            loss.backward()  # Backpropagate gradients
-            optimizer.step()  # Update model parameters
-
-            # Print training progress
-            if batch_idx % 100 == 0:
-                pass Load the prediction model
-model = CustomModel()
-
-# Define a function to prepare the data for prediction
+# Prepare data for prediction
 def prepare_data(date):
-    # Get the historical data for the given date
-    data = bs.query_history_k_data_plus(
-        "sz.000001",  # Shanghai Composite Index
-        "date,open,high,low,close,volume",
-        start_date="2005-05-30",
-        end_date=date,
-        frequency="d"
-    )
-    data_list = []
-    while (data.error_code == '0') & data.next():
-        data_list.append(data.get_row_data())
-    data_df = pd.DataFrame(data_list, columns=data.fields)
-
-    # Convert 'open' and 'close' columns to numeric type
-    data_df['open'] = pd.to_numeric(data_df['open'])
-    data_df['close'] = pd.to_numeric(data_df['close'])
-
-    # Filter out stocks that meet the conditions
-    data_df = data_df[(data_df["open"] >= 0.98 * data_df["close"].shift(1).fillna(0)) & (data_df["open"] <= 1.02 * data_df["close"].shift(1).fillna(0))] 
-    data_df = data_df[(data_df["high"] == data_df["close"]) & (data_df["low"] == data_df["close"])]  # limit-up condition
-    data_df = data_df[(data_df["open"]!= 0) & (data_df["close"]!= 0)]  # exclude zero prices
+    data_df = get_historical_data("2005-05-30", date)
+    filtered_df = filter_stocks(data_df)
+    if filtered_df.empty:
+        return pd.DataFrame()  # Return empty DataFrame if no data matches the filter
 
     # Scale the data using MinMaxScaler
     scaler = MinMaxScaler()
-    data_df[['open', 'high', 'low', 'close', 'volume']] = scaler.fit_transform(data_df[['open', 'high', 'low', 'close', 'volume']])
-
-    return data_df
+    filtered_df[['open', 'high', 'low', 'close', 'volume']] = scaler.fit_transform(filtered_df[['open', 'high', 'low', 'close', 'volume']])
+    return filtered_df
 
 # Define a route to predict the top 5 stock codes
 @app.route('/predict', methods=['POST'])
 def predict():
-    date = request.json['date']
-    data_df = prepare_data(date)
-    if data_df.empty:
-        return jsonify({'error': 'No data available for the given date'}), 400
-    y_pred = model.predict(data_df)
-    top_5_stocks = predict_stock_codes(y_pred, data_df)
-    return jsonify({'top_5_stocks': top_5_stocks})
-
+    try:
+        date = request.json['date']
+        data_df = prepare_data(date)
+        if data_df.empty:
+            return jsonify({'error': 'No data available for the given date'}), 400
+
+        y_pred = model.predict(data_df)
+        top_5_stocks = y_pred[:5]  # Assuming y_pred contains the predicted values for stocks
+        return jsonify({'top_5_stocks': top_5_stocks.tolist()})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+# Run the Flask app
 if __name__ == '__main__':
-    app.run(debug=True)
+    app.run(debug=True)
+
+# Logout from Baostock API
+bs.logout()