245 lines
10 KiB
Python
245 lines
10 KiB
Python
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)}") |