from sklearn.preprocessing import StandardScaler import numpy as np import logging from src.feature_analysis import FeatureAnalysis from src.database import get_db_connection from .logger import setup_logger logger = setup_logger(__name__) # PyTorch 为可选依赖 try: import torch from torch.utils.data import Dataset, DataLoader _HAS_TORCH = True except ImportError: torch = None Dataset = object DataLoader = None _HAS_TORCH = False class EquipmentDataset(Dataset if _HAS_TORCH else object): """装备数据集类""" def __init__(self, features, targets=None): if not _HAS_TORCH: raise ImportError("PyTorch is not installed. Install with: pip install torch") self.features = torch.FloatTensor(features) self.targets = torch.FloatTensor(targets) if targets is not None else None def __len__(self): return len(self.features) def __getitem__(self, idx): if self.targets is not None: return self.features[idx], self.targets[idx] return self.features[idx] class DataPreparation: def __init__(self): self.feature_analyzer = FeatureAnalysis() self.feature_scaler = StandardScaler() self.target_scaler = StandardScaler() def prepare_training_data(self, equipment_data, equipment_type, batch_size=32, validation_split=None): """准备训练数据，可选自动划分验证集""" try: logger.info(f"Preparing training data for {equipment_type}") # 获取特征名称（包含生产商特征） feature_names = self.feature_analyzer.get_equipment_specific_features(equipment_type) features = [] targets = [] # 获取数据库连接 with get_db_connection() as conn: cursor = conn.cursor() # 获取所有生产商数据，用于计算特征 cursor.execute(""" SELECT * FROM manufacturers """) manufacturers = {row['id']: row for row in cursor.fetchall()} for item in equipment_data: # 获取生产商数据 manufacturer = manufacturers.get(item['manufacturer_id'], {}) # 计算生产商特征 manufacturer_features = self.feature_analyzer.calculate_manufacturer_features(manufacturer) # 合并装备特征和生产商特征 feature_values = [] for name in feature_names: if name.startswith('manufacturer_'): value = manufacturer_features.get(name, 0.0) else: value = item.get(name) feature_values.append(float(value) if value is not None else 0.0) features.append(feature_values) targets.append(float(item['actual_cost'])) # 转换为numpy数组 X = np.array(features, dtype=float) y = np.array(targets, dtype=float) # 记录数据范围 logger.info(f"Raw X range: min={X.min()}, max={X.max()}") logger.info(f"Raw y range: min={y.min()}, max={y.max()}") # 标准化特征和目标值 X_scaled = self.feature_scaler.fit_transform(X) y_scaled = self.target_scaler.fit_transform(y.reshape(-1, 1)).ravel() result = { 'feature_names': feature_names, 'feature_scaler': self.feature_scaler, 'target_scaler': self.target_scaler, 'raw_shape': X.shape, 'X': X_scaled, 'y': y_scaled } # 如果指定了验证集比例，自动划分 if validation_split is not None and 0 < validation_split < 1: from sklearn.model_selection import train_test_split X_train, X_val, y_train, y_val = train_test_split( X_scaled, y_scaled, test_size=validation_split, random_state=42 ) result['X'] = X_train result['y'] = y_train result['X_val'] = X_val result['y_val'] = y_val logger.info(f"Auto-split: train={X_train.shape[0]}, val={X_val.shape[0]}") # 创建数据集和数据加载器（仅用于训练集） dataset = EquipmentDataset(result['X'], result['y']) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) result['dataloader'] = dataloader return result except Exception as e: logger.error(f"Error in data preparation: {str(e)}") raise Exception(f"Training error: {str(e)}") def prepare_validation_data(self, validation_data, equipment_type, feature_names=None, scalers=None): """ 准备验证数据 """ try: logger.info(f"Preparing validation data for {equipment_type}") logger.info(f"Raw validation data size: {len(validation_data)}") # 如果输入已经是 numpy 数组，直接使用 if isinstance(validation_data, np.ndarray): X = validation_data logger.info(f"Input is already numpy array with shape: {X.shape}") # 处理无效值 X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0) # 使用训练数据的标准化器 if scalers and 'feature_scaler' in scalers: X_scaled = scalers['feature_scaler'].transform(X) else: X_scaled = X logger.info(f"Preprocessed data shape: {X_scaled.shape}") logger.info(f"Validation features shape: {X_scaled.shape}") logger.info(f"Validation features type: {X_scaled.dtype}") return { 'X': X_scaled, 'y': None # 验证数据可能没有标签 } # 从原始数据中提取特征和目标值 if not feature_names: feature_names = self.feature_analyzer.get_equipment_specific_features(equipment_type) features = [] targets = [] for item in validation_data: # 提取特征值 feature_values = [] for name in feature_names: value = item.get(name) try: feature_values.append(float(value) if value is not None else 0.0) except (ValueError, TypeError): feature_values.append(0.0) features.append(feature_values) # 提取目标值（成本）并验证范围 try: cost = float(item['actual_cost']) if cost > 0: # 只使用正数成本值 targets.append(cost) else: logger.warning(f"Skipping non-positive cost value: {cost}") except (ValueError, TypeError): logger.error(f"Invalid cost value: {item.get('actual_cost')}") continue # 转换为numpy数组 X = np.array(features, dtype=float) y = np.array(targets, dtype=float) # 记录数据范围 logger.info(f"Features range: min={X.min()}, max={X.max()}") logger.info(f"Targets range: min={y.min()}, max={y.max()}") # 处理无效值 X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0) # 使用训练数据的标准化器 if scalers and 'feature_scaler' in scalers: X_scaled = scalers['feature_scaler'].transform(X) if 'target_scaler' in scalers: y_scaled = scalers['target_scaler'].transform(y.reshape(-1, 1)).ravel() else: y_scaled = y else: X_scaled = X y_scaled = y logger.info(f"Preprocessed data shape: {X_scaled.shape}") logger.info(f"Validation features shape: {X_scaled.shape}") logger.info(f"Validation features type: {X_scaled.dtype}") # 记录标准化后的数据范围 logger.info(f"Scaled validation X range: min={X_scaled.min()}, max={X_scaled.max()}") logger.info(f"Scaled validation y range: min={y_scaled.min()}, max={y_scaled.max()}") # 确保特征维度一致 if not feature_names: feature_names = self.feature_analyzer.get_equipment_specific_features(equipment_type) logger.info(f"Expected features: {len(feature_names)}") logger.info(f"Actual features: {X_scaled.shape[1]}") if X_scaled.shape[1] != len(feature_names): raise ValueError(f"Feature dimension mismatch: expected {len(feature_names)}, got {X_scaled.shape[1]}") return { 'X': X_scaled, 'y': y_scaled # 返回标准化后的成本值 } except Exception as e: logger.error(f"Error in validation data preparation: {str(e)}") logger.error(f"Feature names: {feature_names}") logger.error(f"Equipment type: {equipment_type}") raise Exception(f"Validation error: {str(e)}") def calculate_derived_features(self, data, equipment_type): """ 计算衍生特征 """ try: return self.feature_analyzer.calculate_derived_features(data, equipment_type) except Exception as e: logger.error(f"Error calculating derived features: {str(e)}") raise Exception(f"Feature calculation error: {str(e)}")