第一次修改

args.py

@@ -0,0 +1,51 @@
+# args.py
+import argparse
+
+def lstm_args_parser():
+    parser = argparse.ArgumentParser(description="LSTM模型训练参数")
+    
+    # 核心数据集参数
+    parser.add_argument('--train_start_date', type=str, default='2024-10-09', help='训练集开始日期')
+    parser.add_argument('--train_end_date', type=str, default='2025-03-24', help='训练集结束日期')
+    parser.add_argument('--val_start_date', type=str, default='2024-10-09', help='验证集开始日期')
+    parser.add_argument('--val_end_date', type=str, default='2025-03-24', help='验证集结束日期')
+    parser.add_argument('--test_start_date', type=str, default='2024-10-09', help='测试集开始日期')
+    parser.add_argument('--test_end_date', type=str, default='2025-03-24', help='测试集结束日期')
+
+    # 模型架构参数
+    parser.add_argument('--seq_len', type=int, default=10, help='输入序列长度')
+    parser.add_argument('--output_size', type=int, default=5, help='预测步长')
+    parser.add_argument('--step_size', type=int, default=5, help='采样步长')
+    parser.add_argument('--resolution', type=int, default=60, help='数据分辨率(分钟)')
+    parser.add_argument('--feature_num', type=int, default=42, help='输入特征维度')
+    parser.add_argument('--labels_num', type=int, default=4, help='预测标签数量(子模型数量)')
+    
+    # 训练超参数
+    parser.add_argument('--epochs', type=int, default=200, help='训练轮数')
+    parser.add_argument('--hidden_size', type=int, default=64, help='隐藏层大小')
+    parser.add_argument('--num_layers', type=int, default=1, help='LSTM层数')
+    parser.add_argument('--dropout', type=float, default=0, help='dropout概率')
+    parser.add_argument('--lr', type=float, default=0.01, help='学习率')
+    parser.add_argument('--batch_size', type=int, default=512, help='批次大小')
+    
+    parser.add_argument('--scheduler_step_size', type=int, default=100, help='学习率调整步长')
+    parser.add_argument('--scheduler_gamma', type=float, default=0.9, help='学习率衰减率')
+    
+    parser.add_argument('--patience', type=int, default=200, help='早停耐心值')
+    parser.add_argument('--min_delta', type=float, default=1e-10, help='最小改善阈值')
+    parser.add_argument('--device', type=int, default=1, help='GPU设备ID')
+
+    # 文件路径配置
+    parser.add_argument('--start_files', type=int, default=1, help='开始文件索引')
+    parser.add_argument('--end_files', type=int, default=17, help='结束文件索引')
+    parser.add_argument('--data_dir', type=str, default='datasets_anzhen', help='数据文件夹路径')
+    parser.add_argument('--file_pattern', type=str, default='data_process_{}.csv', help='数据文件命名模式')
+    
+    parser.add_argument('--model_path', type=str, default='model.pth', help='模型保存路径')
+    parser.add_argument('--scaler_path', type=str, default='scaler.pkl', help='归一化器路径')
+    parser.add_argument('--output_csv_path', type=str, default='predictions.csv', help='预测评估结果路径')
+    
+    parser.add_argument('--random_seed', type=int, default=1314, help='随机种子')
+
+    args = parser.parse_args()
+    return args

data_preprocessor.py

@@ -0,0 +1,193 @@
+# data_preprocessor.py
+import os
+import torch
+import joblib
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from sklearn.preprocessing import MinMaxScaler
+from torch.utils.data import DataLoader, TensorDataset
+from concurrent.futures import ThreadPoolExecutor
+
+class DataPreprocessor:
+    """数据预处理类"""
+    
+    # 定义必须保留的列
+    COLUMNS_TO_KEEP = [
+        'index',
+        "AR.1#UF_JSFLOW_O", "AR.2#UF_JSFLOW_O", "AR.1#RO_JSFLOW_O", "AR.2#RO_JSFLOW_O",
+        "AR.1#UF_JSPRESS_O", "AR.2#UF_JSPRESS_O", "AR.1#RO_JSPRESS_O", "AR.2#RO_JSPRESS_O",
+        "AR.1#RO_EDJSPRESS_O", "AR.1#RO_SDJSPRESS_O", "AR.2#RO_EDJSPRESS_O", "AR.2#RO_SDJSPRESS_O",
+        "AR.ZJS_TEMP_O", "AR.ZJS_ZD_O", "AR.RO_JSDD_O", "AR.RO_JSORP_O", "AR.RO_JSPH_O",
+        "AR.1#UF_V_FB_O", "AR.2#UF_V_FB_O", "AR.1#UFBWB_FRE_FB_O", "AR.2#UFBWB_FRE_FB_O",
+        "AR.1#RODJB_FRE_FB_O", "AR.1#ROGYB_FRE_FB_O", "AR.1#RODJB_CZ_O", "AR.1#ROGYB_CZ_O",
+        "AR.2#RODJB_CZ_O", "AR.2#ROGYB_CZ_O", "AR.ROGSB_FRE_FB_O", "AR.UFGSB_FRE_FB_O",
+        "AR.V_UF1_TJV_KD_FB", "AR.V_UF2_TJV_KD_FB", "AR.CS_LEVEL_O", "AR.UF_CSLEVEL_O",
+        "AR.UF1_SSD_KMYC", "AR.UF2_SSD_KMYC", "AR.RO1_2D_YC", "AR.PUBLIC_BY_REAL_1",
+        "1#RO_CSFLOW"
+    ]
+
+    @staticmethod
+    def load_and_process_data(args, data):
+        """加载并处理数据，划分训练/验证/测试集"""
+        # 处理日期
+        data['date'] = pd.to_datetime(data['date'])
+        time_interval = pd.Timedelta(minutes=(4 * args.resolution / 60))
+        window_time_span = time_interval * (args.seq_len + 1)
+
+        val_start_date = pd.to_datetime(args.val_start_date)
+        test_start_date = pd.to_datetime(args.test_start_date)
+        
+        # 调整时间窗口
+        adjusted_val_start = val_start_date - window_time_span
+        adjusted_test_start = test_start_date - window_time_span
+        
+        train_mask = (data['date'] >= pd.to_datetime(args.train_start_date)) & \
+                     (data['date'] <= pd.to_datetime(args.train_end_date))
+        val_mask = (data['date'] >= adjusted_val_start) & \
+                   (data['date'] <= pd.to_datetime(args.val_end_date))
+        test_mask = (data['date'] >= adjusted_test_start) & \
+                    (data['date'] <= pd.to_datetime(args.test_end_date))
+
+        train_data = data[train_mask].reset_index(drop=True)
+        val_data = data[val_mask].reset_index(drop=True)
+        test_data = data[test_mask].reset_index(drop=True)
+        
+        train_data = train_data.drop(columns=['date'])
+        val_data = val_data.drop(columns=['date'])
+        test_data = test_data.drop(columns=['date'])
+    
+        # 创建数据集
+        train_supervised = DataPreprocessor.create_supervised_dataset(args, train_data, 1)
+        val_supervised = DataPreprocessor.create_supervised_dataset(args, val_data, 1)
+        test_supervised = DataPreprocessor.create_supervised_dataset(args, test_data, args.step_size)
+        
+        # 转换为DataLoader
+        train_loader = DataPreprocessor.load_data(args, train_supervised, shuffle=True)
+        val_loader = DataPreprocessor.load_data(args, val_supervised, shuffle=False)
+        test_loader = DataPreprocessor.load_data(args, test_supervised, shuffle=False)
+        
+        return train_loader, val_loader, test_loader, data
+    
+    @staticmethod
+    def read_and_combine_csv_files(args):
+        """读取文件并进行特征筛选和预处理"""
+        current_dir = os.path.dirname(__file__)
+        parent_dir = os.path.dirname(current_dir)
+        args.data_dir = os.path.join(parent_dir, args.data_dir)
+        
+        def read_file(file_count):
+            file_name = args.file_pattern.format(file_count)
+            file_path = os.path.join(args.data_dir, file_name)
+            try:
+                df = pd.read_csv(file_path)
+                # 确保只读取需要的列，若列不存在则会报错提示
+                return df[DataPreprocessor.COLUMNS_TO_KEEP]
+            except KeyError as e:
+                print(f"文件 {file_name} 中缺少列: {e}")
+                raise
+        
+        file_indices = list(range(args.start_files, args.end_files + 1))
+        max_workers = os.cpu_count()
+        
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            results = list(tqdm(executor.map(read_file, file_indices),
+                                total=len(file_indices),
+                                desc="正在读取文件"))
+        
+        all_data = pd.concat(results, ignore_index=True)
+        
+        # 确保列顺序一致
+        all_data = all_data[DataPreprocessor.COLUMNS_TO_KEEP]
+        
+        # 下采样
+        chunk = all_data.iloc[::args.resolution, :].reset_index(drop=True)
+        
+        # 处理特征
+        chunk = DataPreprocessor.process_date(chunk, args)
+        chunk = DataPreprocessor.scaler_data(chunk, args)
+        
+        return chunk
+    
+    @staticmethod
+    def process_date(data, args):
+        data = data.rename(columns={'index': 'date'})
+        data['date'] = pd.to_datetime(data['date'])
+    
+        time_features = []
+        # 固定生成分钟级和日级特征，保持与Predictor一致
+        data['minute_of_day'] = data['date'].dt.hour * 60 + data['date'].dt.minute
+        data['minute_sin'] = np.sin(2 * np.pi * data['minute_of_day'] / 1440)
+        data['minute_cos'] = np.cos(2 * np.pi * data['minute_of_day'] / 1440)
+        
+        data['day_of_year'] = data['date'].dt.dayofyear
+        data['day_year_sin'] = np.sin(2 * np.pi * data['day_of_year'] / 366)
+        data['day_year_cos'] = np.cos(2 * np.pi * data['day_of_year'] / 366)
+        
+        time_features.extend(['minute_sin', 'minute_cos', 'day_year_sin', 'day_year_cos'])
+        data.drop(columns=['minute_of_day', 'day_of_year'], inplace=True)
+    
+        other_columns = [col for col in data.columns if col not in ['date'] and col not in time_features]
+        data = data[['date'] + time_features + other_columns]
+        return data
+    
+    @staticmethod
+    def scaler_data(data, args):
+        date_col = data[['date']]
+        data_to_scale = data.drop(columns=['date'])
+
+        scaler = MinMaxScaler(feature_range=(0, 1))
+        scaled_data = scaler.fit_transform(data_to_scale)
+        joblib.dump(scaler, args.scaler_path)
+
+        scaled_data = pd.DataFrame(scaled_data, columns=data_to_scale.columns)
+        scaled_data = pd.concat([date_col.reset_index(drop=True), scaled_data], axis=1)
+        return scaled_data
+    
+    @staticmethod
+    def create_supervised_dataset(args, data, step_size):
+        data = pd.DataFrame(data)
+        cols = []
+        col_names = []
+        feature_columns = data.columns.tolist()
+
+        # 输入序列
+        for col in feature_columns:
+            for i in range(args.seq_len - 1, -1, -1):
+                cols.append(data[[col]].shift(i))
+                col_names.append(f"{col}(t-{i})")
+        
+        # 目标序列 (取最后labels_num列)
+        target_columns = feature_columns[-args.labels_num:]
+        for i in range(1, args.output_size + 1):
+            for col in target_columns:
+                cols.append(data[[col]].shift(-i))
+                col_names.append(f"{col}(t+{i})")
+
+        dataset = pd.concat(cols, axis=1)
+        dataset.columns = col_names
+        dataset = dataset.iloc[::step_size, :]
+        dataset.dropna(inplace=True)
+        return dataset
+
+    @staticmethod
+    def load_data(args, dataset, shuffle):
+        input_length = args.seq_len
+        n_features = args.feature_num
+        labels_num = args.labels_num
+    
+        n_features_total = n_features * input_length
+        n_labels_total = args.output_size * labels_num
+
+        X = dataset.values[:, :n_features_total]
+        y = dataset.values[:, n_features_total:n_features_total + n_labels_total]
+    
+        X = X.reshape(X.shape[0], input_length, n_features)
+        X = torch.tensor(X, dtype=torch.float32).to(args.device)
+        y = torch.tensor(y, dtype=torch.float32).to(args.device)
+
+        dataset_tensor = TensorDataset(X, y)
+        generator = torch.Generator()
+        generator.manual_seed(args.random_seed)
+        
+        return DataLoader(dataset_tensor, batch_size=args.batch_size, shuffle=shuffle, generator=generator)

data_trainer.py

@@ -0,0 +1,165 @@
+# data_trainer.py
+import torch
+import joblib
+import numpy as np
+import pandas as pd
+from sklearn.metrics import r2_score
+from datetime import datetime, timedelta
+from sklearn.preprocessing import MinMaxScaler
+
+class Trainer:
+    def __init__(self, model, args, data):
+        self.args = args
+        self.model = model
+        self.data = data
+        self.patience = args.patience
+        self.min_delta = args.min_delta
+        self.counter = 0
+        self.early_stop = False
+        self.best_val_loss = float('inf')
+        self.best_model_state = None
+        self.best_epoch = 0
+
+    def train_full_model(self, train_loader, val_loader, optimizer, criterion, scheduler):
+        self.counter = 0
+        self.best_val_loss = float('inf')
+        self.early_stop = False
+        self.best_model_state = None
+        self.best_epoch = 0
+        max_epochs = self.args.epochs
+
+        for epoch in range(max_epochs):
+            self.model.train()
+            running_loss = 0.0
+            
+            for inputs, targets in train_loader:
+                inputs = inputs.to(self.args.device)
+                targets = targets.to(self.args.device)
+                optimizer.zero_grad()
+                outputs = self.model(inputs)
+                loss = criterion(outputs, targets)
+                loss.backward()
+                optimizer.step()
+                running_loss += loss.item()
+            
+            train_loss = running_loss / len(train_loader)
+            val_loss = self.validate_full(val_loader, criterion) if val_loader else 0.0
+
+            print(f'Epoch {epoch+1}/{max_epochs}, Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}')
+
+            if val_loader:
+                if val_loss < (self.best_val_loss - self.min_delta):
+                    self.best_val_loss = val_loss
+                    self.counter = 0
+                    self.best_model_state = self.model.state_dict()
+                    self.best_epoch = epoch
+                else:
+                    self.counter += 1
+                    if self.counter >= self.patience:
+                        self.early_stop = True
+                        print(f"早停触发")
+                        
+            scheduler.step()
+            torch.cuda.empty_cache()
+            if self.early_stop:
+                break
+
+        if self.best_model_state is not None:
+            self.model.load_state_dict(self.best_model_state)
+        print(f"最佳迭代: {self.best_epoch+1}, 最佳验证损失: {self.best_val_loss:.6f}")
+        return self.model
+
+    def validate_full(self, val_loader, criterion):
+        self.model.eval()
+        total_loss = 0.0
+        with torch.no_grad():
+            for inputs, targets in val_loader:
+                inputs = inputs.to(self.args.device)
+                targets = targets.to(self.args.device)
+                outputs = self.model(inputs)
+                loss = criterion(outputs, targets)
+                total_loss += loss.item()
+        return total_loss / len(val_loader)
+
+    def save_model(self):
+        torch.save(self.model.state_dict(), self.args.model_path)
+        print(f"模型已保存到：{self.args.model_path}")
+            
+    def evaluate_model(self, test_loader, criterion):
+        self.model.eval()
+        scaler = joblib.load(self.args.scaler_path)
+        predictions = []
+        true_values = []
+        
+        with torch.no_grad():
+            for inputs, targets in test_loader:
+                inputs = inputs.to(self.args.device)
+                targets = targets.to(self.args.device)
+                outputs = self.model(inputs)
+                predictions.append(outputs.cpu().numpy())
+                true_values.append(targets.cpu().numpy())
+    
+        predictions = np.concatenate(predictions, axis=0)
+        true_values = np.concatenate(true_values, axis=0)
+    
+        # 重塑
+        reshaped_predictions = predictions.reshape(predictions.shape[0], self.args.output_size, self.args.labels_num)
+        predictions = reshaped_predictions.reshape(-1, self.args.labels_num)
+        
+        reshaped_true_values = true_values.reshape(true_values.shape[0], self.args.output_size, self.args.labels_num)
+        true_values = reshaped_true_values.reshape(-1, self.args.labels_num)
+    
+        # 反归一化 (仅标签列)
+        column_scaler = MinMaxScaler(feature_range=(0, 1))
+        column_scaler.min_ = scaler.min_[-self.args.labels_num:] 
+        column_scaler.scale_ = scaler.scale_[-self.args.labels_num:] 
+        
+        true_values = column_scaler.inverse_transform(true_values)
+        predictions = column_scaler.inverse_transform(predictions)
+    
+        # 定义4个核心变量
+        column_names = [
+            "AR.UF1_SSD_KMYC",          # UF1跨膜压差
+            "AR.RO1_2D_YC",             # RO1二段压差
+            "AR.PUBLIC_BY_REAL_1",      # RO1三段压差
+            "1#RO_CSFLOW",              # 1#RO产水流量
+        ]
+    
+        # 生成时间
+        start_datetime = datetime.strptime(self.args.test_start_date, "%Y-%m-%d")
+        time_interval = timedelta(minutes=(4 * self.args.resolution / 60))
+        total_points = len(predictions)
+        date_times = [start_datetime + i * time_interval for i in range(total_points)]
+        
+        results = pd.DataFrame({'date': date_times})
+        metrics_details = []
+        
+        for i, col_name in enumerate(column_names):
+            if i >= self.args.labels_num: break # 防止越界
+            
+            results[f'{col_name}_True'] = true_values[:, i]
+            results[f'{col_name}_Predicted'] = predictions[:, i]
+            
+            var_true = true_values[:, i]
+            var_pred = predictions[:, i]
+            
+            # 指标计算
+            non_zero_mask = var_true != 0
+            var_true_nonzero = var_true[non_zero_mask]
+            var_pred_nonzero = var_pred[non_zero_mask]
+            
+            if len(var_true_nonzero) > 0:
+                r2 = r2_score(var_true_nonzero, var_pred_nonzero)
+                rmse = np.sqrt(np.mean((var_true_nonzero - var_pred_nonzero) ** 2))
+                mape = np.mean(np.abs((var_true_nonzero - var_pred_nonzero) / np.abs(var_true_nonzero))) * 100
+                metrics_details.append(f"{col_name}: R2={r2:.4f}, RMSE={rmse:.4f}, MAPE={mape:.4f}%")
+            else:
+                metrics_details.append(f"{col_name}: 无效数据")
+
+        results.to_csv(self.args.output_csv_path, index=False)
+        
+        txt_path = self.args.output_csv_path.replace('.csv', '_metrics.txt')
+        with open(txt_path, 'w') as f:
+            f.write('\n'.join(metrics_details))
+            
+        return metrics_details

gat_lstm.py

@@ -0,0 +1,54 @@
+# gat_lstm.py
+import torch
+import torch.nn as nn
+
+class SingleGATLSTM(nn.Module):
+    """单个子模型：预测1个目标指标"""
+    def __init__(self, args):
+        super(SingleGATLSTM, self).__init__()
+        self.args = args
+        
+        self.lstm = nn.LSTM(
+            input_size=args.feature_num,
+            hidden_size=args.hidden_size,
+            num_layers=args.num_layers,
+            batch_first=True
+        )
+        
+        self.final_linear = nn.Sequential(
+            nn.Linear(args.hidden_size, args.hidden_size),
+            nn.LeakyReLU(0.01),
+            nn.Dropout(args.dropout * 0.4),
+            nn.Linear(args.hidden_size, args.output_size)
+        )
+        self._init_weights()
+        
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None: nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        batch_size, seq_len, feature_num = x.size()
+        lstm_out, _ = self.lstm(x)
+        last_out = lstm_out[:, -1, :]
+        output = self.final_linear(last_out)
+        return output
+
+class GAT_LSTM(nn.Module):
+    """总模型：包含多个SingleGATLSTM子模型"""
+    def __init__(self, args):
+        super(GAT_LSTM, self).__init__()
+        self.args = args
+        # 创建4个独立模型（对应labels_num=4）
+        self.models = nn.ModuleList([SingleGATLSTM(args) for _ in range(args.labels_num)])
+    
+    def set_edge_index(self, edge_index):
+        self.edge_index = edge_index
+        
+    def forward(self, x):
+        outputs = []
+        for model in self.models:
+            outputs.append(model(x))
+        return torch.cat(outputs, dim=1)

input_format.txt

@@ -0,0 +1,39 @@
+index
+AR.1#UF_JSFLOW_O
+AR.2#UF_JSFLOW_O
+AR.1#RO_JSFLOW_O
+AR.2#RO_JSFLOW_O
+AR.1#UF_JSPRESS_O
+AR.2#UF_JSPRESS_O
+AR.1#RO_JSPRESS_O
+AR.2#RO_JSPRESS_O
+AR.1#RO_EDJSPRESS_O
+AR.1#RO_SDJSPRESS_O
+AR.2#RO_EDJSPRESS_O
+AR.2#RO_SDJSPRESS_O
+AR.ZJS_TEMP_O
+AR.ZJS_ZD_O
+AR.RO_JSDD_O
+AR.RO_JSORP_O
+AR.RO_JSPH_O
+AR.1#UF_V_FB_O
+AR.2#UF_V_FB_O
+AR.1#UFBWB_FRE_FB_O
+AR.2#UFBWB_FRE_FB_O
+AR.1#RODJB_FRE_FB_O
+AR.1#ROGYB_FRE_FB_O
+AR.1#RODJB_CZ_O
+AR.1#ROGYB_CZ_O
+AR.2#RODJB_CZ_O
+AR.2#ROGYB_CZ_O
+AR.ROGSB_FRE_FB_O
+AR.UFGSB_FRE_FB_O
+AR.V_UF1_TJV_KD_FB
+AR.V_UF2_TJV_KD_FB
+AR.CS_LEVEL_O
+AR.UF_CSLEVEL_O
+AR.UF1_SSD_KMYC
+AR.UF2_SSD_KMYC
+AR.RO1_2D_YC
+AR.PUBLIC_BY_REAL_1
+1#RO_CSFLOW

main.py

@@ -0,0 +1,51 @@
+# main.py
+import os
+import torch
+import numpy as np
+import random
+from gat_lstm import GAT_LSTM
+from data_trainer import Trainer
+from args import lstm_args_parser
+from torch.nn import MSELoss
+from data_preprocessor import DataPreprocessor
+
+def set_seed(seed):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def main():
+    args = lstm_args_parser()
+    set_seed(args.random_seed)
+    
+    device = torch.device(f"cuda:{args.device}" if torch.cuda.is_available() else "cpu")
+    args.device = device
+
+    print(f"当前配置: 序列长度={args.seq_len}, 特征数={args.feature_num}, 目标数={args.labels_num}")
+
+    # 数据预处理
+    data = DataPreprocessor.read_and_combine_csv_files(args)
+    train_loader, val_loader, test_loader, _ = DataPreprocessor.load_and_process_data(args, data)
+    
+    # 初始化模型
+    model = GAT_LSTM(args).to(device)
+
+    # 训练器
+    trainer = Trainer(model, args, data)
+    criterion = MSELoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.scheduler_step_size, gamma=args.scheduler_gamma)
+
+    print("=== 开始训练 ===")
+    trainer.train_full_model(train_loader, val_loader, optimizer, criterion, scheduler)
+    trainer.save_model()
+    
+    print("=== 开始评估 ===")
+    trainer.evaluate_model(test_loader, MSELoss())
+
+if __name__ == "__main__":
+    main()

output_format.txt

@@ -0,0 +1,2 @@
+预测结果 (5x4 数组):
+[[0.022336346255615352, 0.0721074445027709, 0.09917135512113572, 122.41020530160138], [0.050365259237587445, 0.14928617544579506, 0.029558163076043132, 444.2995635541028], [0.10231768742203712, 0.20150794134879113, 0.15267262662887573, 110.25981214446792], [0.1623324816673994, 0.06816830243599414, 0.049503796252608305, 106.16647537057565], [0.039512041788548224, 0.044865222564578054, 0.0686768901526928, 145.52421178978847]]

predict.py

@@ -0,0 +1,208 @@
+# 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()