DualFlow/predict.py

# predict.py
import os
import torch
import joblib
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from gat_lstm import GAT_LSTM

class RealTimePredictor:
    def __init__(self, model_path='model.pth', scaler_path='scaler.pkl', device=None):
        """
        初始化预测器
        """
        # 1. 参数配置 (与训练 args.py 保持一致)
        self.seq_len = 10         # 输入序列长度
        self.feature_num = 42     # 输入特征数 (4时间编码 + 38业务特征)
        self.labels_num = 4       # 输出标签数
        self.hidden_size = 64
        self.num_layers = 1
        self.output_size = 5      # 预测未来 5 步
        self.dropout = 0
        
        # 2. 设备与资源加载
        self.device = device if device else torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model_path = model_path
        self.scaler_path = scaler_path
        
        # 加载归一化器
        if not os.path.exists(self.scaler_path):
             raise FileNotFoundError(f"未找到归一化文件: {self.scaler_path}，请确保已完成训练。")
        self.scaler = joblib.load(self.scaler_path)

        # 加载模型
        self._load_model()

        # 定义必须存在的列名 (39个，包含index，顺序必须固定)
        self.required_columns = [
            'index', 
            "AR.1#UF_JSFLOW_O",         # 1#UF进水流量
            "AR.2#UF_JSFLOW_O",         # 2#UF进水流量
            "AR.1#RO_JSFLOW_O",         # 1#RO进水流量
            "AR.2#RO_JSFLOW_O",         # 2#RO进水流量
            "AR.1#UF_JSPRESS_O",        # 1#UF进水压力
            "AR.2#UF_JSPRESS_O",        # 2#UF进水压力
            "AR.1#RO_JSPRESS_O",        # 1#RO进水压力
            "AR.2#RO_JSPRESS_O",        # 2#RO进水压力
            "AR.1#RO_EDJSPRESS_O",      # 1#RO二段进水压力
            "AR.1#RO_SDJSPRESS_O",      # 1#RO三段进水压力
            "AR.2#RO_EDJSPRESS_O",      # 2#RO二段进水压力
            "AR.2#RO_SDJSPRESS_O",      # 2#RO三段进水压力
            "AR.ZJS_TEMP_O",            # 进水温度
            "AR.ZJS_ZD_O",              # UF进水浊度
            "AR.RO_JSDD_O",             # RO进水电导
            "AR.RO_JSORP_O",            # RO进水ORP
            "AR.RO_JSPH_O",             # RO进水PH
            "AR.1#UF_V_FB_O",           # 1#UF调节阀开度反馈
            "AR.2#UF_V_FB_O",           # 2#UF调节阀开度反馈
            "AR.1#UFBWB_FRE_FB_O",      # 1#UF反洗泵频率反馈
            "AR.2#UFBWB_FRE_FB_O",      # 2#UF反洗泵频率反馈
            "AR.1#RODJB_FRE_FB_O",      # 1#RO段间泵频率反馈
            "AR.1#ROGYB_FRE_FB_O",      # 1#RO高压泵频率反馈
            "AR.1#RODJB_CZ_O",          # 1#RO段间泵测振反馈
            "AR.1#ROGYB_CZ_O",          # 1#RO高压泵测振反馈
            "AR.2#RODJB_CZ_O",          # 2#RO段间泵测振反馈
            "AR.2#ROGYB_CZ_O",          # 2#RO高压泵测振反馈
            "AR.ROGSB_FRE_FB_O",        # RO供水泵频率反馈
            "AR.UFGSB_FRE_FB_O",        # UF供水泵频率反馈
            "AR.V_UF1_TJV_KD_FB",       # UF1调节阀开度反馈
            "AR.V_UF2_TJV_KD_FB",       # UF2调节阀开度反馈
            "AR.CS_LEVEL_O",            # RO产水箱液位
            "AR.UF_CSLEVEL_O",          # UF产水箱液位
            "AR.UF1_SSD_KMYC",          # UF1跨膜压差
            "AR.UF2_SSD_KMYC",          # UF2跨膜压差
            "AR.RO1_2D_YC",             # RO1二段压差
            "AR.PUBLIC_BY_REAL_1",      # RO1三段压差
            "1#RO_CSFLOW",              # 1#RO产水流量
        ]

    def _load_model(self):
        """内部方法：加载模型权重"""
        class ModelArgs: pass
        args = ModelArgs()
        args.feature_num = self.feature_num
        args.hidden_size = self.hidden_size
        args.num_layers = self.num_layers
        args.output_size = self.output_size
        args.labels_num = self.labels_num
        args.dropout = self.dropout

        self.model = GAT_LSTM(args).to(self.device)
        
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(f"未找到模型权重文件: {self.model_path}")
            
        state_dict = torch.load(self.model_path, map_location=self.device, weights_only=True)
        self.model.load_state_dict(state_dict)
        self.model.eval()

    def _preprocess(self, df):
        """数据预处理：补全、排序、生成时间特征、整体归一化"""
        data = df.copy()
        
        # 1. 统一时间列名
        if 'datetime' in data.columns:
            data = data.rename(columns={'datetime': 'index'})
        if 'index' not in data.columns:
             data['index'] = pd.date_range(end=datetime.now(), periods=len(data), freq='min')
        data['index'] = pd.to_datetime(data['index'])
        
        # 2. 补全长度 (Padding)
        if len(data) < self.seq_len:
            pad_len = self.seq_len - len(data)
            first_row = data.iloc[0:1]
            pads = pd.concat([first_row] * pad_len, ignore_index=True)
            start_time = data['index'].iloc[0]
            for i in range(pad_len):
                pads.at[i, 'index'] = start_time - timedelta(minutes=(pad_len-i))
            data = pd.concat([pads, data], ignore_index=True)

        # 3. 列筛选排序 (提取业务数据，不含index)
        try:
            # required_columns[0] 是 'index'，我们取后面的业务列
            business_cols = self.required_columns[1:]
            data_business = data[business_cols]
        except KeyError:
            missing = list(set(self.required_columns) - set(data.columns))
            raise ValueError(f"缺少列: {missing}")

        # 4. 生成时间特征
        date_col = data['index']
        minute_of_day = date_col.dt.hour * 60 + date_col.dt.minute
        day_of_year = date_col.dt.dayofyear
        
        time_features = pd.DataFrame({
            'minute_sin': np.sin(2 * np.pi * minute_of_day / 1440),
            'minute_cos': np.cos(2 * np.pi * minute_of_day / 1440),
            'day_year_sin': np.sin(2 * np.pi * day_of_year / 366),
            'day_year_cos': np.cos(2 * np.pi * day_of_year / 366)
        })
        
        # 5. 拼接：[时间特征 + 业务特征]
        # 注意：训练时的顺序是 time_features + other_columns
        # 必须重置索引以避免拼接错位
        data_to_scale = pd.concat([
            time_features.reset_index(drop=True), 
            data_business.reset_index(drop=True)
        ], axis=1)
        
        # 6. 整体归一化
        # 此时 columns 应该包含: minute_sin, minute_cos..., AR.1#UF_JSFLOW_O...
        # 顺序和名字必须与 fit 时一致
        scaled_array = self.scaler.transform(data_to_scale)
        
        return scaled_array

    def predict(self, df):
        """
        返回: List[List[float]]
        格式: [[t+1时刻的4个值], [t+2时刻的4个值], ..., [t+5时刻的4个值]]
        """
        # 1. 预处理 (返回的是归一化后的 numpy 数组)
        processed_data = self._preprocess(df)
        
        # 2. 取最后 seq_len 个时间步构建 Tensor
        input_seq = processed_data[-self.seq_len:] 
        input_tensor = torch.tensor(input_seq, dtype=torch.float32).unsqueeze(0).to(self.device)
        
        # 3. 推理
        with torch.no_grad():
            output = self.model(input_tensor)
        
        # 4. 反归一化
        # 输出形状调整为 (5, 4) -> 5个步长, 4个变量
        preds = output.cpu().numpy().reshape(self.output_size, self.labels_num)
        
        # 获取最后4列的归一化参数 (目标变量)
        target_min = self.scaler.min_[-self.labels_num:]
        target_scale = self.scaler.scale_[-self.labels_num:]
        
        real_preds = (preds - target_min) / target_scale
        real_preds = np.abs(real_preds)
        
        # 5. 返回纯数值列表
        return real_preds.tolist()

if __name__ == "__main__":
    # 测试代码
    try:
        # 初始化
        predictor = RealTimePredictor()
        
        # 生成模拟数据
        mock_data = pd.DataFrame()
        mock_data['index'] = pd.date_range(end=datetime.now(), periods=15, freq='min')
        for col in predictor.required_columns[1:]:
            mock_data[col] = np.random.rand(15) * 10
            
        # 预测
        result = predictor.predict(mock_data)
        
        print("预测结果 (5x4 数组):")
        print(result)
        
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()