From c7526068405ecefdbbdc2ae97b132a9a5ef19c8a Mon Sep 17 00:00:00 2001
From: "wenwen.tang@thizgroup.com" <wenwen.tang@thizgroup.com>
Date: Fri, 26 Apr 2024 14:17:54 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E9=A2=84=E6=B5=8B=E6=A8=A1?=
 =?UTF-8?q?=E5=9E=8B=EF=BC=8C=E5=A2=9E=E5=8A=A08=E6=A1=A3=E9=A2=84?=
 =?UTF-8?q?=E6=B5=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ai/EstimateMarketTrendV20.py | 22 +++++++++++++---------
 ai/model_trainer.py          | 22 ++++++++++++++++++++--
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/ai/EstimateMarketTrendV20.py b/ai/EstimateMarketTrendV20.py
index 0c4a5ae..0d1f087 100644
--- a/ai/EstimateMarketTrendV20.py
+++ b/ai/EstimateMarketTrendV20.py
@@ -11,24 +11,23 @@ from ai.training_data_builder import TrainingDataBuilder
 from api import DataSync
 
 # 截止日期
-# max_date = None
-max_date = '2024-03-01'
+max_date = None
+# max_date = '2024-03-20'
 # max_date = '2024-01-11'
 
 toForecast = True  # False means test, True means forecast
 syncData = True  # 开启会同步数据库指数及基金数据
-uploadData = False  # 开启会上传预测结果
+uploadData = True  # 开启会上传预测结果
 doReport = True  # 开启会生成Excel报告
 
 # 待预测指数
 # PREDICT_LIST = [67, 121, 122, 123]
-PREDICT_LIST = [67, 121, 122, 123, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171,
-                174, 175, 177, 178]
+PREDICT_LIST = [67, 121, 122, 123, 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, 174, 175, 177, 178]
 eco = [65, 66, 74, 134]
 index = [67, 68, 69, 70, 71, 72, 73, 75, 76, 77, 105, 106, 116, 117, 138, 139, 142, 143, 140, 141, 144, 145, 146]
-fund = [121, 122, 123, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 174, 175,
-        177,
-        178]
+fund = [121, 122, 123, 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, 174, 175, 177, 178]
 
 
 @autowired
@@ -130,10 +129,15 @@ if __name__ == '__main__':
         rf_model = trainer.train_random_forest(X_train, y_train, X_test, y_test)
         gbt_model = trainer.train_GBT(X_train, y_train, X_test, y_test)
         svc_model = trainer.train_SVC(X_train, y_train, X_test, y_test)
-        ensemble_model = trainer.ensemble_model(rf_model, gbt_model, svc_model, X_train, y_train, X_test, y_test)
+        knn_model = trainer.train_nearest_neighbors(X_train, y_train, X_test, y_test)
+        ada_model = trainer.train_AdaBoost(X_train, y_train, X_test, y_test)
+        ensemble_model = trainer.ensemble_model(rf_model, gbt_model, svc_model,
+                                                knn_model, ada_model, X_train, y_train, X_test, y_test)
         model_predict = {'forest': rf_model.predict(scaledX_forecast),
                          'gbt': gbt_model.predict(scaledX_forecast),
                          'svc': svc_model.predict(scaledX_forecast),
+                         'knn': knn_model.predict(scaledX_forecast),
+                         'adaboost': ada_model.predict(scaledX_forecast),
                          'ensemble': ensemble_model.predict(scaledX_forecast)}
         print(f'预测结果：{model_predict}'.center(60, '+'))
         judgement(pid, infos_type[pid], model_predict)
diff --git a/ai/model_trainer.py b/ai/model_trainer.py
index 662fa01..939be9c 100644
--- a/ai/model_trainer.py
+++ b/ai/model_trainer.py
@@ -3,8 +3,10 @@ from abc import ABC
 import matplotlib.pyplot as plt
 from lightgbm import LGBMClassifier
 from sklearn import svm
+from sklearn.ensemble import AdaBoostClassifier
 from sklearn.ensemble import RandomForestClassifier, VotingClassifier
 from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay, accuracy_score
+from sklearn.neighbors import KNeighborsClassifier
 
 
 class ModelTrainer(ABC):
@@ -61,9 +63,25 @@ class ModelTrainer(ABC):
             self.test_model('Support Vector Machines', classifierSVC, X_test, y_test)
         return classifierSVC
 
-    def ensemble_model(self, rf_model, gbt_model, svc_model, X_train, y_train, X_test, y_test):
+    def train_nearest_neighbors(self, X_train, y_train, X_test, y_test):
+        classifier = KNeighborsClassifier()
+        classifier.fit(X_train, y_train)
+        if not self._toForecast:
+            self.test_model('K-Nearest Neighbors', classifier, X_test, y_test)
+        return classifier
+
+    def train_AdaBoost(self, X_train, y_train, X_test, y_test):
+        classifier = AdaBoostClassifier()
+        classifier.fit(X_train, y_train)
+        if not self._toForecast:
+            self.test_model('AdaBoost', classifier, X_test, y_test)
+        return classifier
+
+    def ensemble_model(self, rf_model, gbt_model, svc_model, knn_model,
+                       ada_model, X_train, y_train, X_test, y_test):
         # Create a dictionary of our models
-        estimators = [('rf', rf_model), ('gbt', gbt_model), ('svc', svc_model)]
+        estimators = [('rf', rf_model), ('gbt', gbt_model), ('svc', svc_model),
+                      ('knn', knn_model), ('AdaBoost', ada_model)]
         # Create our voting classifier, inputting our models
         ensemble = VotingClassifier(estimators, voting='hard')
         # fit model to training data
-- 
2.18.1