""" 药店销售预测系统 - 核心预测器类 (已重构) 支持多店铺销售预测功能，并完全集成新的ModelManager """ import os import pandas as pd import time from datetime import datetime from trainers import ( train_product_model_with_mlstm, train_product_model_with_kan, train_product_model_with_tcn, train_product_model_with_transformer ) from predictors.model_predictor import load_model_and_predict from utils.multi_store_data_utils import ( load_multi_store_data, get_store_product_sales_data, aggregate_multi_store_data ) from core.config import DEVICE, DEFAULT_MODEL_DIR, DEFAULT_DATA_PATH from utils.model_manager import model_manager class PharmacyPredictor: """ 药店销售预测系统核心类，用于训练模型和进行预测 """ def __init__(self, data_path=None, model_dir=DEFAULT_MODEL_DIR): """ 初始化预测器 """ self.data_path = data_path if data_path else DEFAULT_DATA_PATH self.model_dir = model_dir self.device = DEVICE if not os.path.exists(model_dir): os.makedirs(model_dir) print(f"使用设备: {self.device}") try: self.data = load_multi_store_data(self.data_path) print(f"已加载多店铺数据，来源: {self.data_path}") except Exception as e: print(f"加载数据失败: {e}") self.data = None def train_model(self, product_id, model_type='transformer', epochs=100, learning_rate=0.001, use_optimized=False, store_id=None, training_mode='product', aggregation_method='sum', socketio=None, task_id=None, progress_callback=None, patience=10): """ 训练预测模型 - 完全适配新的训练器接口 """ def log_message(message, log_type='info'): print(f"[{log_type.upper()}] {message}", flush=True) if progress_callback: try: progress_callback({'log_type': log_type, 'message': message}) except Exception as e: print(f"[ERROR] 进度回调失败: {e}", flush=True) if self.data is None: log_message("没有可用的数据，请先加载或生成数据", 'error') return None # --- 数据准备 --- try: if training_mode == 'store': product_data = get_store_product_sales_data(store_id, product_id, self.data_path) log_message(f"按店铺-产品训练: 店铺 {store_id}, 产品 {product_id}, 数据量: {len(product_data)}") elif training_mode == 'global': product_data = aggregate_multi_store_data(product_id, aggregation_method, self.data_path) log_message(f"全局训练模式: 产品 {product_id}, 聚合方法 {aggregation_method}, 数据���: {len(product_data)}") else: # 'product' product_data = self.data[self.data['product_id'] == product_id].copy() log_message(f"按产品训练模式: 产品 {product_id}, 数据量: {len(product_data)}") except Exception as e: log_message(f"数据准备失败: {e}", 'error') return None if product_data.empty: log_message(f"找不到产品 {product_id} 的数据", 'error') return None # --- 训练器选择与参数准备 --- trainers = { 'transformer': train_product_model_with_transformer, 'mlstm': train_product_model_with_mlstm, 'tcn': train_product_model_with_tcn, 'kan': train_product_model_with_kan, 'optimized_kan': train_product_model_with_kan, } if model_type not in trainers: log_message(f"不支持的模型类型: {model_type}", 'error') return None trainer_func = trainers[model_type] # 统一所有训练器的参数 trainer_args = { "product_id": product_id, "product_df": product_data, "store_id": store_id, "training_mode": training_mode, "aggregation_method": aggregation_method, "epochs": epochs, "socketio": socketio, "task_id": task_id, "progress_callback": progress_callback, "patience": patience, "learning_rate": learning_rate } # 为 KAN 模型添加特殊参数 if 'kan' in model_type: trainer_args['use_optimized'] = (model_type == 'optimized_kan') # --- 调用训练器 --- try: log_message(f"🤖 开始调用 {model_type} 训练器") model, metrics, version, model_version_path = trainer_func(**trainer_args) log_message(f"✅ {model_type} 训练器成功返回", 'success') if metrics: metrics.update({ 'model_type': model_type, 'version': version, 'model_path': model_version_path, 'training_mode': training_mode, 'store_id': store_id, 'product_id': product_id, 'aggregation_method': aggregation_method if training_mode == 'global' else None }) log_message(f"📈 最终返回的metrics: {metrics}", 'success') return metrics else: log_message("⚠️ 训练器返回的metrics为空", 'warning') return None except Exception as e: import traceback log_message(f"模型训练过程中发生严重错误: {e}\n{traceback.format_exc()}", 'error') return None def predict(self, model_version_path, future_days=7, start_date=None, analyze_result=False): """ 使用已训练的模型进行预测 - 直接使用模型版本路径 """ if not os.path.exists(model_version_path): raise FileNotFoundError(f"指定的模型路径不存在: {model_version_path}") return load_model_and_predict( model_version_path=model_version_path, future_days=future_days, start_date=start_date, analyze_result=analyze_result ) def list_models(self, **kwargs): """ 列出所有可用的模型版本。 直接调用 ModelManager 的 list_models 方法。 支持的过滤参数: model_type, training_mode, scope, version """ return model_manager.list_models(**kwargs) def delete_model(self, model_version_path): """ 删除一个指定的模型版本目录。 """ return model_manager.delete_model_version(model_version_path) def compare_models(self, product_id, epochs=50, **kwargs): """ 在相同数据上训练并比较多个模型的性能。 """ results = {} model_types_to_compare = ['tcn', 'mlstm', 'transformer', 'kan', 'optimized_kan'] for model_type in model_types_to_compare: print(f"\n{'='*20} 训练模型: {model_type.upper()} {'='*20}") try: metrics = self.train_model( product_id=product_id, model_type=model_type, epochs=epochs, **kwargs ) results[model_type] = metrics if metrics else {} except Exception as e: print(f"训练 {model_type} 模型失败: {e}") results[model_type] = {'error': str(e)} # 打印比较结果 print(f"\n{'='*25} 模型性能比较 {'='*25}") # 准备数据 df_data = [] for model, metrics in results.items(): if metrics and 'rmse' in metrics: df_data.append({ 'Model': model.upper(), 'RMSE': metrics.get('rmse'), 'R²': metrics.get('r2'), 'MAPE (%)': metrics.get('mape'), 'Time (s)': metrics.get('training_time') }) if not df_data: print("没有可供比较的模型结果。") return results comparison_df = pd.DataFrame(df_data).set_index('Model') print(comparison_df.to_string(float_format="%.4f")) return results