zhangzexuan
/
FemGIL
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

上传文件至 /

This commit is contained in:
zhangzexuan 2026-05-25 13:31:26 +08:00
commit dbcb20280a
5 changed files with 3165 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

217
build_graph.py Normal file
View File

@ -0,0 +1,217 @@
import torch
from torch_geometric.data import Data
import numpy as np
import scipy.sparse as sp
import os
import glob
import re
# ================= 配置区域 =================
# 从全局配置文件导入数据集配置
from config import DATASET_TYPE, DATA_ROOT_PATH, SCA_PREFIX, DATASET_DIRS_PATTERN
# 根据全局配置设置路径
root_data_path = os.path.join(os.path.dirname(__file__), DATA_ROOT_PATH)
sca_prefix = SCA_PREFIX
dataset_dirs_pattern = DATASET_DIRS_PATTERN
# ===========================================
def is_main_process():
"""检查是否是主进程（rank 0）"""
rank = int(os.environ.get("RANK", "0"))
return rank == 0
# 全局变量
data_mapping = {}
n_total = 0
# 移除模块级别的打印，避免DDP重复打印
# 配置信息将在实际使用时打印
# ===========================================
def scan_all_data(root_path):
"""
扫描A-4目录下的所有四组数据（A-TainDataset, A-TainDataset2, A-TainDataset3, A-TainDataset4）
依然以 edge 文件为锚点
"""
global data_mapping, n_total
# 如果已经扫描过，直接返回，避免重复扫描
if data_mapping and n_total > 0:
return data_mapping, n_total
data_mapping = {}
k_idx = 1
# 获取所有数据子目录
dataset_dirs = [d for d in os.listdir(root_path)
if os.path.isdir(os.path.join(root_path, d)) and d.startswith(dataset_dirs_pattern.replace('*', ''))]
if is_main_process():
print(f"发现 {len(dataset_dirs)} 个数据集目录: {dataset_dirs}")
total_folders = 0
for dataset_dir in sorted(dataset_dirs):
dataset_path = os.path.join(root_path, dataset_dir)
# 扫描每个数据集目录中的sca*_data文件夹
all_items = os.listdir(dataset_path)
subfolders = [f for f in all_items
if os.path.isdir(os.path.join(dataset_path, f))
and f.startswith('sca') and f.endswith('_data')]
def extract_folder_num(folder_name):
match = re.match(rf'{sca_prefix}(\d+)_data', folder_name)
return int(match.group(1)) if match else 9999
subfolders.sort(key=extract_folder_num)
# 移除此打印，避免DDP重复打印
for folder_name in subfolders:
folder_path = os.path.join(dataset_path, folder_name)
folder_num = extract_folder_num(folder_name)
# 查找 edge 文件
search_pattern = os.path.join(folder_path, f"edge_{sca_prefix}{folder_num}_*.txt")
edge_files = glob.glob(search_pattern)
file_ids = []
pattern = re.compile(rf'edge_{sca_prefix}{folder_num}_(\d+)\.txt')
for vf in edge_files:
match = pattern.match(os.path.basename(vf))
if match:
file_ids.append(int(match.group(1)))
file_ids.sort()
for data_id in file_ids:
data_mapping[k_idx] = (folder_path, folder_num, str(data_id))
k_idx += 1
total_folders += 1
n_total = k_idx - 1
if is_main_process():
print(f"初始化完成：总共扫描 {total_folders} 个文件夹，找到 {n_total} 组数据。")
return data_mapping, n_total
def load_file_data(folder_path, prefix, folder_num, data_id):
"""通用数据读取函数"""
filename = f"{prefix}_{sca_prefix}{folder_num}_{data_id}.txt"
filepath = os.path.join(folder_path, filename)
if not os.path.exists(filepath):
raise FileNotFoundError(f"文件缺失: {filepath}")
data = np.loadtxt(filepath, dtype=np.float32)
if data.ndim == 1:
data = data.reshape(1, -1)
return data
def build_graph_data(k):
"""
构建图数据，包含残差计算步骤: r = b - A * Ebz
"""
if not data_mapping:
scan_all_data(root_data_path)
if k not in data_mapping:
raise ValueError(f"索引 k={k} 超出范围")
folder_path, folder_num, data_id = data_mapping[k]
# -------------------------------------------------------
# 1. 基础数据读取
# -------------------------------------------------------
# 读取 edge (拓扑)
edge_data = load_file_data(folder_path, "edge", folder_num, data_id)
edge_index_np = edge_data.astype(np.int64) - 1 # MATLAB -> Python 索引
edge_index = torch.tensor(edge_index_np.T, dtype=torch.long)
num_nodes = int(edge_index.max()) + 1
# 读取 eps (材料), b (右端项), Ebz (当前解/初始解)
eps_data = load_file_data(folder_path, "eps", folder_num, data_id) # [N, 2]
b_data = load_file_data(folder_path, "b", folder_num, data_id) # [N, 2]
Ebz_data = load_file_data(folder_path, "Ebz", folder_num, data_id) # [N, 2]
# -------------------------------------------------------
# 2. 计算残差 r (Preprocessing)
# 公式: r = b - A * Ebz
# -------------------------------------------------------
# 2.1 读取稀疏矩阵 A 的信息 (Aij 和 Av)
try:
Aij_data = load_file_data(folder_path, "Aij", folder_num, data_id) # [NNZ, 2] 坐标
Av_data = load_file_data(folder_path, "Av", folder_num, data_id) # [NNZ, 2] 值
except FileNotFoundError:
raise FileNotFoundError(f"计算残差需要 Aij 和 Av 文件，但在数据组 {k} 中未找到。")
# 2.2 构建 Scipy 稀疏矩阵 (COO 格式)
# Aij 是 MATLAB 索引 (1-based)，需要减 1
rows = Aij_data[:, 0].astype(int) - 1
cols = Aij_data[:, 1].astype(int) - 1
# 构造复数数值: Real + 1j * Imag
values = Av_data[:, 0] + 1j * Av_data[:, 1]
# 创建稀疏矩阵 A (N x N)
A_mat = sp.coo_matrix((values, (rows, cols)), shape=(num_nodes, num_nodes))
# 2.3 准备向量 b 和 Ebz (复数形式)
b_vec = b_data[:, 0] + 1j * b_data[:, 1]
Ebz_vec = Ebz_data[:, 0] + 1j * Ebz_data[:, 1]
# 2.4 执行矩阵乘法和减法: r = b - A * Ebz
# A_mat.dot() 是高效的稀疏矩阵乘法
Ax = A_mat.dot(Ebz_vec)
r_vec = b_vec - Ax # 得到复数残差向量 [N,]
# 2.5 将残差拆分为实部和虚部 [N, 2]
r_real = r_vec.real
r_imag = r_vec.imag
# -------------------------------------------------------
# 3. 拼接节点特征 (Input Features)
# 目标输入: [eps, r, Ebz] (共 6 通道)
# -------------------------------------------------------
# eps: [N, 2]
# r: [N, 2] (由上一步计算得到)
# Ebz: [N, 2]
x_tensor = torch.cat([
torch.from_numpy(eps_data), # eps_real, eps_imag
torch.tensor(r_real, dtype=torch.float32).unsqueeze(1), # r_real
torch.tensor(r_imag, dtype=torch.float32).unsqueeze(1), # r_imag
torch.from_numpy(Ebz_data), # Ebz_real, Ebz_imag (当前解)
torch.from_numpy(Ebz_data), # bg_real, bg_imag (背景场，不随网络更新)
], dim=1)
# -------------------------------------------------------
# 4. 读取标签 (Ground Truth)
# -------------------------------------------------------
Esz_data = load_file_data(folder_path, "Esz", folder_num, data_id)
y_tensor = torch.from_numpy(Esz_data)
# -------------------------------------------------------
# 5. 返回 Data 对象
# -------------------------------------------------------
data = Data(x=x_tensor, edge_index=edge_index, y=y_tensor)
data.k_idx = torch.tensor([k])
return data
# 测试运行
if __name__ == "__main__":
scan_all_data(root_data_path)
if n_total > 0:
print("\n--- 读取并计算残差中 ---")
try:
data = build_graph_data(1)
print("数据构建成功!")
print(f"节点特征 x shape: {data.x.shape}")
print("特征顺序: [eps_re, eps_im, r_re, r_im, Ebz_re, Ebz_im, bg_re, bg_im]")
print(f"残差(r)部分均值: {data.x[:, 2:4].abs().mean().item():.4e}")
except Exception as e:
print(f"出错: {e}")

110
config.py Normal file
View File

@ -0,0 +1,110 @@
"""
PhiSAGE 项目全局配置文件
集中管理所有可配置的参数，避免循环导入问题
"""
# ==========================================
# 全局迭代次数配置（统一管理所有相关文件的迭代次数）
# ==========================================
N_ITER = 5 # 全局迭代次数，修改此值将影响所有相关文件
# ==========================================
# 全局模型保存路径配置
# ==========================================
SAVE_DIR = "/public/home/zzx/gnn/PhiSAGE/PhiSAGE-test/saved_models/A-0.6" # 模型保存目录
OUTPUT_DIR = "/public/home/zzx/gnn/PhiSAGE/PhiSAGE-test/training_outputs/A-0.6" # 训练曲线输出目录
# ==========================================
# 数据集配置
# ==========================================
# 数据集选择：设置为 'A' 或 'B'
# - 'A': 读取 data/A-4 目录（包含A-TainDataset等4个子目录），使用 scaA 命名格式
# - 'B': 读取 data/B-4 目录（包含B-TainDataset等4个子目录），使用 scaB 命名格式
# - 'C': 读取 data/C-4 目录（包含C-TainDataset等4个子目录），使用 scaC 命名格式
# 使用方法：
# 1. 将数据放在对应的目录结构中
# 2. 修改 DATASET_TYPE 的值
# 3. 程序会自动扫描所有子目录中的数据
#
DATASET_TYPE = 'A' # 可选: 'A' 或 'B' 或 'C'
DATA_ROOT_PATH = "data/A-1"
# 根据选择动态设置相关参数（由 build_graph.py 使用）
if DATASET_TYPE == 'A':
SCA_PREFIX = 'scaA'
DATASET_DIRS_PATTERN = 'A-TainDataset*'
elif DATASET_TYPE == 'B':
SCA_PREFIX = 'scaB'
DATASET_DIRS_PATTERN = 'B-TainDataset*'
elif DATASET_TYPE == 'C':
SCA_PREFIX = 'scaC'
DATASET_DIRS_PATTERN = 'C-TainDataset*'
else:
raise ValueError("DATASET_TYPE 必须设置为 'A' 或 'B' 或 'C'")
# ==========================================
# ==========================================
# 训练启动器配置 (run.py 使用)
# ==========================================
# 默认 GPU 配置
DEFAULT_TARGET_GPUS = "4,5" # 默认使用的 GPU 列表
# 训练模式配置
DEFAULT_TRAIN_MODE = "ddp" # 默认训练模式: "single", "multi", "ddp"
# 单卡训练配置
DEFAULT_SINGLE_GPU_ID = 0 # 默认单卡 GPU ID
# ==========================================
# DDP分布式训练配置
# ==========================================
# DDP通信端口配置
MASTER_PORT = "20870" # DDP主进程通信端口
# ==========================================
# 网络维度配置
# ==========================================
# 网络架构配置：支持多种维度设置方式
#
# 方式1：使用base_dim自动计算（推荐）
# NETWORK_BASE_DIM = 56 # 基础维度，会自动生成 [56, 112, 192]
#
# 方式2：自定义维度配置
# NETWORK_CUSTOM_DIMS = [56, 112, 192] # 手动指定每层维度
#
# 方式3：使用预设配置
# NETWORK_CONFIG = "default" # 可选: "small", "medium", "large", "xlarge"
# 当前使用的配置方式
NETWORK_USE_CUSTOM_DIMS = True # True=使用自定义维度，False=使用base_dim自动计算
NETWORK_BASE_DIM = 64 # 基础维度（当不使用自定义维度时）
NETWORK_CUSTOM_DIMS = [64, 128, 256] # 自定义维度配置
# 池化配置
NETWORK_POOL_RATIOS = [0.8, 0.6] # 池化比例
# ==========================================
# Loss函数配置
# ==========================================
# LOSS_TYPE: 选择使用的损失函数类型
# - "mse": 传统的MSE损失 (||pred - true||^2)
# - "phi": Phi损失 (||A*x - b||^2) - 直接计算物理残差
# - "asinh": Asinh损失 sqrt(asinh(norm(x-x_ref)^2/N)) - 平滑的损失函数
# - "hybrid": 混合损失 - 前100epoch纯MSE，后续MSE + λ*phi (λ从0.001开始每50epoch×10倍，至0.1)
LOSS_TYPE = "hybrid" # 可选: "mse", "phi", "asinh" 或 "hybrid"
# ==========================================
# 预训练模型加载配置（迁移学习）
# ==========================================
# LOAD_PRETRAINED_MODEL: 是否加载预训练模型权重（仅模型权重，不加载优化器等状态）
# - True: 从 PRETRAINED_MODEL_DIR 目录加载模型权重，优化器等重新初始化
# - False: 不加载预训练模型，从头开始训练
LOAD_PRETRAINED_MODEL = False # 是否加载预训练模型权重
# PRETRAINED_MODEL_DIR: 预训练模型目录
# 如果 LOAD_PRETRAINED_MODEL=True，将从该目录加载模型权重文件（real_iter_*.pth 和 imag_iter_*.pth）
# None表示使用SAVE_DIR，也可以指定其他目录
PRETRAINED_MODEL_DIR = None # None表示使用SAVE_DIR，也可以指定其他目录
# ==========================================
# 其他全局配置参数
# ==========================================
# 可以在这里添加其他需要全局共享的配置参数

467
model.py Normal file
View File

@ -0,0 +1,467 @@
import torch
import torch.nn as nn
from torch_geometric.nn import MessagePassing, TopKPooling
# 从全局配置文件导入迭代次数配置
from config import N_ITER
# ==========================================
# 6. 网络维度配置导入
# ==========================================
# 从全局配置导入网络维度设置
from config import (
NETWORK_USE_CUSTOM_DIMS,
NETWORK_BASE_DIM,
NETWORK_CUSTOM_DIMS,
NETWORK_POOL_RATIOS
)
# ==========================================
# 0. BatchNorm 辅助函数
# ==========================================
def get_batch_norm(num_features):
return nn.BatchNorm1d(num_features)
# ==========================================
# 1. 物理感知图卷积层
# ==========================================
class EMFullComplexLayer(MessagePassing):
def __init__(self, in_feats, out_feats):
super(EMFullComplexLayer, self).__init__()
self.aggr = 'mean'
total_in_dim = in_feats + in_feats
self.lin_fusion = nn.Linear(total_in_dim, out_feats)
def forward(self, x, edge_index):
return self.propagate(edge_index, x=x)
def message(self, x_j):
return x_j
def update(self, aggs, x):
combined = torch.cat([x, aggs], dim=1)
fused_msg = self.lin_fusion(combined)
return fused_msg
# ==========================================
# 2. 修改后的 GCN 模块
# ==========================================
class GCN(nn.Module):
def __init__(self, in_feats, out_feats):
super(GCN, self).__init__()
self.conv = EMFullComplexLayer(in_feats, out_feats)
self.bn = get_batch_norm(out_feats)
self.gelu = nn.GELU()
def forward(self, x, edge_index):
x = self.conv(x, edge_index)
x = self.bn(x)
x = self.gelu(x)
return x
class FFTFeatureLayer(nn.Module):
"""
对特征维度进行 FFT
"""
def __init__(self, in_feats, out_feats):
super(FFTFeatureLayer, self).__init__()
# 线性投影调整维度
self.lin_in = nn.Linear(in_feats, out_feats)
# 可学习的频域滤波器 (复数权重)
# RFFT 后，频域维度为 (out_feats // 2) + 1
self.freq_dim = out_feats // 2 + 1
# 初始化复数权重
# shape: [freq_dim] - 对每个频率分量进行缩放
self.complex_weight = nn.Parameter(
torch.randn(self.freq_dim, 2, dtype=torch.float32) * 0.02
)
def forward(self, x):
# x: [N, in_feats] -> [N, out_feats]
x = self.lin_in(x)
# 修复：cuFFT在float16模式下只支持维度大小为2的幂次
# 如果输入是float16且维度不是2的幂次，需要转换为float32
original_dtype = x.dtype
converted_to_float32 = False
if x.dtype == torch.float16:
# 检查最后一个维度是否为2的幂次
last_dim = x.size(-1)
is_power_of_two = (last_dim & (last_dim - 1)) == 0 and last_dim != 0
if not is_power_of_two:
# 转换为float32以避免cuFFT限制
x = x.to(torch.float32)
converted_to_float32 = True
x_fft = torch.fft.rfft(x, dim=-1)
# 2. 频域滤波 / 混合
weight = torch.view_as_complex(self.complex_weight)
# 广播乘法: [N, freq_dim] * [freq_dim]
x_fft = x_fft * weight
# 3. 傅里叶逆变换 (Complex -> Real)
x_out = torch.fft.irfft(x_fft, n=x.size(-1), dim=-1)
# 如果原始是float16且被转换过，转换回float16（保持精度一致性）
if original_dtype == torch.float16 and converted_to_float32:
x_out = x_out.to(original_dtype)
return x_out
# ==========================================
# 2.5 [核心修改] FFM 增强的输入层
# ==========================================
class FFMEncodingLayer(nn.Module):
"""
使用傅里叶特征映射 (FFM) 替换传统的线性输入层。
逻辑：
1. 输入 features (如坐标或物理量)
2. 映射 -> sin(2*pi*B*x), cos(2*pi*B*x)
3. 线性投影融合
4. 物理图卷积 (EMFullComplexLayer)
"""
def __init__(self, in_feats, out_feats, sigma=1.0, mapping_size=None):
super(FFMEncodingLayer, self).__init__()
# 1. 配置 FFM 参数
self.input_dim = in_feats
# 如果未指定 mapping_size，默认设为 out_feats 的一半（因为sin/cos会翻倍）
# 或者设为一个固定的高维空间，例如 128 或 256
self.mapping_size = mapping_size if mapping_size is not None else max(out_feats, 64)
self.sigma = sigma
# 2. 初始化高斯随机矩阵 B (不可学习，类似位置编码)
# shape: [in_feats, mapping_size]
self.B = nn.Parameter(
torch.randn(self.input_dim, self.mapping_size) * self.sigma,
requires_grad=False
)
# FFM 后的维度 = mapping_size * 2 (sin + cos)
ffm_out_dim = self.mapping_size * 2
# 3. 特征融合层：将高维 FFM 特征投影回目标维度
self.feature_projection = nn.Linear(ffm_out_dim, out_feats)
# 4. 空间混合：保留图卷积能力
self.conv = EMFullComplexLayer(out_feats, out_feats)
# 5. 激活与归一化
self.bn = get_batch_norm(out_feats)
self.gelu = nn.GELU()
def input_mapping(self, x):
# x: [N, input_dim]
# 投影: (2 * pi * x) @ B
# 结果 shape: [N, mapping_size]
x_proj = (2.0 * torch.pi * x) @ self.B
# 拼接 sin 和 cos -> [N, mapping_size * 2]
return torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1)
def forward(self, x, edge_index):
# 1. FFM 映射 (提升频率感知能力)
x_ffm = self.input_mapping(x)
# 2. 线性投影 (融合特征)
x_embed = self.feature_projection(x_ffm)
# 3. 图卷积 (聚合邻居信息)
x_out = self.conv(x_embed, edge_index)
# 4. 后处理
x_out = self.bn(x_out)
x_out = self.gelu(x_out)
return x_out
# ==========================================
# 2.6 新增：混合谱图卷积层 (替代原 GCN)
# ==========================================
class SpectralGCN(nn.Module):
"""
混合层：结合 局部图卷积(GCN) 和 频域特征变换(FFT)
Out = GCN(x, edge) + FFT(x)
"""
def __init__(self, in_feats, out_feats):
super(SpectralGCN, self).__init__()
# 分支1：物理感知图卷积 (处理局部连接)
self.spatial_conv = EMFullComplexLayer(in_feats, out_feats)
# 分支2：FFT 频域层 (处理全局/频域特征)
self.spectral_conv = FFTFeatureLayer(in_feats, out_feats)
# 融合门控 (可学习的加权系数)
self.alpha = nn.Parameter(torch.tensor(0.5))
self.bn = get_batch_norm(out_feats)
self.gelu = nn.GELU()
def forward(self, x, edge_index):
# 1. 空间路径：消息传递
x_spatial = self.spatial_conv(x, edge_index)
# 2. 频谱路径：FFT 变换 (不依赖 edge_index，纯节点特征变换)
x_spectral = self.spectral_conv(x)
# 3. 残差融合
# 使用 sigmoid 确保融合比例在 0-1 之间，或者直接相加
# 这里使用加权求和，兼顾空间和频域信息
x_out = x_spatial + self.alpha * x_spectral
# 4. 激活与归一化
x_out = self.bn(x_out)
x_out = self.gelu(x_out)
return x_out
# ==========================================
# 2.7 FFM处理
# ==========================================
class ffm_process(nn.Module):
def __init__(self, in_feats, out_feats, ffm_sigma=1.0):
super(ffm_process, self).__init__()
self.B = nn.Parameter(torch.randn(in_feats, out_feats) * ffm_sigma, requires_grad=False)
def forward(self, x):
x_proj = (2.0 * torch.pi * x) @ self.B
return torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1)
# ==========================================
# 3. Graph U-Nets 组件 (Checkpoint稳定版)
# ==========================================
class gPool(nn.Module):
"""Graph Pooling layer - 简化版，不特殊处理背景场"""
def __init__(self, in_dim, ratio):
super(gPool, self).__init__()
self.ratio = ratio
self.proj = nn.Linear(in_dim, 1) # 评分投影
self.sigmoid = nn.Sigmoid()
def forward(self, x, edge_index):
N, C = x.size()
# 计算节点重要性得分（使用所有特征）
scores = self.proj(x) # [N, 1]
scores = self.sigmoid(scores).squeeze() # [N]
# 确定保留的节点数量
k = max(1, int(self.ratio * N))
# 选择top-k节点
values, idx = torch.topk(scores, k)
# 门控机制：保留的节点特征按重要性加权
x_pool = x[idx] * values.view(-1, 1) # [k, C]
# 筛选保留节点之间的边
row, col = edge_index
row_mask = torch.zeros(N, dtype=torch.bool, device=x.device)
col_mask = torch.zeros(N, dtype=torch.bool, device=x.device)
row_mask[idx] = True
col_mask[idx] = True
edge_mask = row_mask[row] & col_mask[col]
edge_index_pool = edge_index[:, edge_mask]
# 重新映射节点索引到0-k范围
idx_map = -torch.ones(N, dtype=torch.long, device=x.device)
idx_map[idx] = torch.arange(k, device=x.device)
edge_index_pool[0] = idx_map[edge_index_pool[0]]
edge_index_pool[1] = idx_map[edge_index_pool[1]]
return x_pool, edge_index_pool, idx
class gUnpool(nn.Module):
"""Graph Unpooling layer (inverse of Pool)"""
def __init__(self):
super(gUnpool, self).__init__()
def forward(self, x_pool, x_skip, idx):
# idx 对应 TopKPooling 返回的 perm (保留节点的索引)
N = x_skip.size(0)
C = x_pool.size(1)
# 恢复原始图结构大小
x_unpool = torch.zeros(N, C, device=x_pool.device, dtype=x_skip.dtype)
if idx is not None and len(idx) > 0:
# 填充保留节点的特征
# PyG TopKPooling 的 perm 保证了 idx < N
# 确保 x_pool 的行数与 idx 长度一致
if x_pool.size(0) == len(idx):
x_unpool[idx] = x_pool
else:
# 理论上不应发生，除非维度对不上
valid_len = min(x_pool.size(0), len(idx))
x_unpool[idx[:valid_len]] = x_pool[:valid_len]
# 跳跃连接（残差连接）
x_unpool = x_unpool + x_skip
return x_unpool
# ==========================================
# 4. 辅助模块
# ==========================================
class StackedLinearBlock(nn.Module):
def __init__(self, in_feats, out_feats, dropout):
super(StackedLinearBlock, self).__init__()
self.fc = nn.Linear(in_feats, out_feats)
self.bn = get_batch_norm(out_feats)
self.gelu = nn.GELU()
self.dropout = nn.Dropout(dropout)
self.reset_parameters()
def reset_parameters(self):
nn.init.kaiming_uniform_(self.fc.weight, a=1.0)
if self.fc.bias is not None:
nn.init.zeros_(self.fc.bias)
def forward(self, x):
x = self.fc(x)
x = self.bn(x)
x = self.gelu(x)
x = self.dropout(x)
return x
class FinalLinear(nn.Module):
def __init__(self, in_feats, out_feats):
super(FinalLinear, self).__init__()
self.fc = nn.Linear(in_feats, out_feats)
self.reset_parameters()
def reset_parameters(self):
nn.init.kaiming_uniform_(self.fc.weight, a=1.0)
if self.fc.bias is not None:
nn.init.zeros_(self.fc.bias)
def forward(self, x):
x = self.fc(x)
return x
# ==========================================
# 5. BuildGCN 网络构建 (最接近原始版本)
# ==========================================
class BuildGCN(nn.Module):
def __init__(self, input_feats, output_feats, base_dim=None, custom_dims=None):
super(BuildGCN, self).__init__()
# 优先级：函数参数 > 全局配置 > 默认值
if custom_dims is not None:
# 函数参数指定的自定义维度（最高优先级）
dims = custom_dims
elif base_dim is not None:
# 函数参数指定的基础维度
dims = [base_dim, base_dim * 2, base_dim * 4]
elif NETWORK_USE_CUSTOM_DIMS:
# 全局配置的自定义维度
dims = NETWORK_CUSTOM_DIMS
else:
# 全局配置的基础维度自动计算
dims = [NETWORK_BASE_DIM, NETWORK_BASE_DIM * 2, NETWORK_BASE_DIM * 4]
# 使用全局配置的池化比例
pool_ratios = NETWORK_POOL_RATIOS
# ========== 编码器路径 ==========
self.gcn1 = SpectralGCN(input_feats, dims[0])
self.pool1 = TopKPooling(dims[0], pool_ratios[0])
self.gcn2 = GCN(dims[0], dims[1])
self.pool2 = TopKPooling(dims[1], pool_ratios[1])
self.gcn3 = GCN(dims[1], dims[2])
# ========== 瓶颈层 ==========
self.bottom_gcn = GCN(dims[2], dims[2])
# ========== 解码器路径 ==========
self.dec_gcn1 = GCN(dims[2], dims[1])
self.unpool1 = gUnpool()
self.dec_gcn2 = GCN(dims[1], dims[0])
self.unpool2 = gUnpool()
self.dec_gcn3 = GCN(dims[0], dims[0])
# ========== 输出层 ==========
self.lin_stacked = StackedLinearBlock(dims[0], dims[0], dropout=0.5)
self.final_linear = FinalLinear(dims[0], output_feats)
def forward(self, ndata, edata, batch=None):
x, edge_index = ndata, edata
if batch is None:
batch = torch.zeros(x.size(0), dtype=torch.long, device=x.device)
# === 编码器路径 ===
x1 = self.gcn1(x, edge_index)
# 第1次池化 (TopKPooling)
# 返回值: x, edge_index, edge_attr, batch, perm, score
x2, edge_index2, _, batch2, idx1, _ = self.pool1(x1, edge_index, batch=batch)
x2 = self.gcn2(x2, edge_index2)
# 第2次池化 (TopKPooling)
x3, edge_index3, _, batch3, idx2, _ = self.pool2(x2, edge_index2, batch=batch2)
x3 = self.gcn3(x3, edge_index3)
# === 瓶颈层 ===
x_bottom = self.bottom_gcn(x3, edge_index3)
# === 解码器路径 ===
# 注意：这里继续使用 edge_index3，和原逻辑一致
x_up1 = self.dec_gcn1(x_bottom, edge_index3)
# 反池化：需要传入 idx (即 perm) 来恢复到上一层的大小 (x2 的大小)
x_up1 = self.unpool1(x_up1, x2, idx2)
# 注意：这里使用 edge_index2
x_up2 = self.dec_gcn2(x_up1, edge_index2)
# 反池化：恢复到输入层大小 (x1 的大小)
x_up2 = self.unpool2(x_up2, x1, idx1)
x_out = self.dec_gcn3(x_up2, edge_index)
# === 输出层 ===
x_out = self.lin_stacked(x_out)
x_out = self.final_linear(x_out)
return x_out
# ==========================================
# 7. BuildGCNList (保持不变)
# ==========================================
class BuildGCNList(nn.Module):
def __init__(self, input_feats, output_feats, n_iter=N_ITER, base_dim=None, custom_dims=None):
super(BuildGCNList, self).__init__()
self.n_iter = n_iter
self.input_feats = input_feats
self.output_feats = output_feats
self.networks = nn.ModuleList()
for i in range(n_iter):
network = BuildGCN(input_feats, output_feats, base_dim=base_dim, custom_dims=custom_dims)
self.networks.append(network)
setattr(self, f'iter_{i}', network)
print(f"已创建 {n_iter} 个集成Graph U-Nets的BuildGCN网络")
def forward(self, ndata, edata, batch=None, iter_idx=None):
if iter_idx is not None:
if iter_idx < 0 or iter_idx >= self.n_iter:
raise ValueError(f"iter_idx必须在[0, {self.n_iter-1}]范围内")
return self.networks[iter_idx](ndata, edata, batch)
else:
outputs = []
for i in range(self.n_iter):
output = self.networks[i](ndata, edata, batch)
outputs.append(output)
return outputs
def get_network(self, iter_idx):
return self.networks[iter_idx]
def __len__(self):
return self.n_iter

2090
trainddp.py Normal file
View File

File diff suppressed because it is too large Load Diff

281
visualize.py Normal file
View File

@ -0,0 +1,281 @@
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as mtri
# 导入 build_graph 模块的函数和变量
from build_graph import (
scan_all_data,
root_data_path,
load_file_data,
data_mapping
)
# 从全局配置导入数据集配置，确保文件命名格式同步
from config import DATASET_TYPE, SCA_PREFIX
# 注意：移除对save_test模块的导入以避免循环导入问题
# k值现在通过函数参数传递，而不是模块级变量
# 默认算例编号
DEFAULT_K = 8777
os.makedirs(root_data_path, exist_ok=True)
def load_data(k, use_prediction=True):
"""
加载数据用于可视化（使用与build_graph.py一致的文件读取方式）
Args:
k: 算例编号（全局索引）
use_prediction: 如果True，加载预测结果，否则加载真实值(Esz)
Returns:
vertices: 节点坐标 [N, 2]
triangles: 三角形索引 [M, 3]
x_real: 实部 [N]
x_imag: 虚部 [N]
"""
# 扫描数据并获取data_mapping
data_mapping, n_total = scan_all_data(root_data_path)
if k not in data_mapping:
raise ValueError(f"索引 k={k} 超出范围 (总数据量: {n_total})")
folder_path, folder_num, data_id = data_mapping[k]
# 读取节点坐标（vertex文件）
vertex_file = os.path.join(folder_path, f"vertex_{SCA_PREFIX}{folder_num}_{data_id}.txt")
if not os.path.exists(vertex_file):
raise FileNotFoundError(f"找不到节点坐标文件: {vertex_file}")
vertices = np.loadtxt(vertex_file)
# 读取三角形索引（tri文件）
tri_file = os.path.join(folder_path, f"tri_{SCA_PREFIX}{folder_num}_{data_id}.txt")
if not os.path.exists(tri_file):
raise FileNotFoundError(f"找不到三角形索引文件: {tri_file}")
triangles = np.loadtxt(tri_file, dtype=int) - 1 # MATLAB索引转Python索引
# 根据use_prediction选择加载预测结果或真实值
if use_prediction:
# 加载预测结果
x_real_file = os.path.join(folder_path, f"Esz_pred_real_{SCA_PREFIX}{folder_num}_{data_id}.txt")
x_imag_file = os.path.join(folder_path, f"Esz_pred_imag_{SCA_PREFIX}{folder_num}_{data_id}.txt")
if not os.path.exists(x_real_file) or not os.path.exists(x_imag_file):
raise FileNotFoundError(f"找不到预测结果文件: {x_real_file} 或 {x_imag_file}")
x_real = np.loadtxt(x_real_file)
x_imag = np.loadtxt(x_imag_file)
else:
# 加载真实值（使用build_graph的load_file_data函数）
Esz_data = load_file_data(folder_path, "Esz", folder_num, data_id)
x_real = Esz_data[:, 0]
x_imag = Esz_data[:, 1]
return vertices, triangles, x_real, x_imag
def load_mesh_data(k):
"""
只加载网格数据（节点坐标和三角形索引），不加载场数据
Args:
k: 算例编号（全局索引）
Returns:
vertices: 节点坐标 [N, 2]
triangles: 三角形索引 [M, 3]
"""
# 扫描数据并获取data_mapping
data_mapping, n_total = scan_all_data(root_data_path)
if k not in data_mapping:
raise ValueError(f"索引 k={k} 超出范围 (总数据量: {n_total})")
folder_path, folder_num, data_id = data_mapping[k]
# 读取节点坐标（vertex文件）
vertex_file = os.path.join(folder_path, f"vertex_{SCA_PREFIX}{folder_num}_{data_id}.txt")
if not os.path.exists(vertex_file):
raise FileNotFoundError(f"找不到节点坐标文件: {vertex_file}")
vertices = np.loadtxt(vertex_file)
# 读取三角形索引（tri文件）
tri_file = os.path.join(folder_path, f"tri_{SCA_PREFIX}{folder_num}_{data_id}.txt")
if not os.path.exists(tri_file):
raise FileNotFoundError(f"找不到三角形索引文件: {tri_file}")
triangles = np.loadtxt(tri_file, dtype=int) - 1 # MATLAB索引转Python索引
return vertices, triangles
def visualize_solution(k, output_dir="/public/home/zzx/gnn/PhiSAGE/PhiSAGE/visualizations",
use_prediction=True, save_combined=True, save_separate=False,
custom_data=None, custom_filename=None):
"""
可视化解的实部、虚部和模值场图
Args:
k: 算例编号（全局索引）
output_dir: 输出目录
use_prediction: 如果True，可视化预测结果，否则可视化真实值
save_combined: 是否保存包含三个子图的组合图
save_separate: 是否分别保存三个单独的图
custom_data: 自定义数据字典，包含'vertices', 'triangles', 'x_real', 'x_imag'键
custom_filename: 自定义文件名后缀，用于区分不同的迭代
"""
if custom_data is not None:
# 使用提供的自定义数据
vertices = custom_data['vertices']
triangles = custom_data['triangles']
x_real = custom_data['x_real']
x_imag = custom_data['x_imag']
else:
# 从文件加载数据
vertices, triangles, x_real, x_imag = load_data(k, use_prediction=use_prediction)
# 计算模值
x_complex = x_real + 1j * x_imag
x_magnitude = np.abs(x_complex)
# 创建输出目录
os.makedirs(output_dir, exist_ok=True)
x_coord = vertices[:, 0]
y_coord = vertices[:, 1]
triang = mtri.Triangulation(x_coord, y_coord, triangles)
# 确定数据标签（用于文件名和标题）
if custom_filename is not None:
data_label = custom_filename
title_suffix = custom_filename
else:
data_label = "prediction" if use_prediction else "true"
title_suffix = f"k={k}"
# 方案1: 保存组合图（三个子图）
if save_combined:
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
# 实部
tcf1 = axes[0].tricontourf(triang, x_real, levels=100, cmap='RdBu_r')
cbar1 = fig.colorbar(tcf1, ax=axes[0])
cbar1.set_label('Re(u)', fontsize=12)
axes[0].set_title(f'Real Part ({title_suffix})', fontsize=14, fontweight='bold')
axes[0].set_aspect('equal')
axes[0].set_xlim(x_coord.min(), x_coord.max())
axes[0].set_ylim(y_coord.min(), y_coord.max())
axes[0].set_xticks([])
axes[0].set_yticks([])
# 虚部
tcf2 = axes[1].tricontourf(triang, x_imag, levels=100, cmap='RdBu_r')
cbar2 = fig.colorbar(tcf2, ax=axes[1])
cbar2.set_label('Im(u)', fontsize=12)
axes[1].set_title(f'Imaginary Part ({title_suffix})', fontsize=14, fontweight='bold')
axes[1].set_aspect('equal')
axes[1].set_xlim(x_coord.min(), x_coord.max())
axes[1].set_ylim(y_coord.min(), y_coord.max())
axes[1].set_xticks([])
axes[1].set_yticks([])
# 模值
tcf3 = axes[2].tricontourf(triang, x_magnitude, levels=100, cmap='jet')
cbar3 = fig.colorbar(tcf3, ax=axes[2])
cbar3.set_label('|u|', fontsize=12)
axes[2].set_title(f'Magnitude ({title_suffix})', fontsize=14, fontweight='bold')
axes[2].set_aspect('equal')
axes[2].set_xlim(x_coord.min(), x_coord.max())
axes[2].set_ylim(y_coord.min(), y_coord.max())
axes[2].set_xticks([])
axes[2].set_yticks([])
plt.tight_layout()
out_file_combined = os.path.join(output_dir, f"{data_label}_combined.png")
plt.savefig(out_file_combined, dpi=300, bbox_inches='tight')
plt.close()
print(f"已保存组合图到: {out_file_combined}")
# 方案2: 分别保存三个单独的图
if save_separate:
# 实部
fig, ax = plt.subplots(figsize=(8, 8))
tcf = ax.tricontourf(triang, x_real, levels=100, cmap='RdBu_r')
cbar = fig.colorbar(tcf, ax=ax)
cbar.set_label('Re(u)', fontsize=14)
ax.set_title(f'Real Part ({title_suffix})', fontsize=16, fontweight='bold')
ax.set_aspect('equal')
ax.set_xlim(x_coord.min(), x_coord.max())
ax.set_ylim(y_coord.min(), y_coord.max())
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
out_file_real = os.path.join(output_dir, f"{data_label}_real.png")
plt.savefig(out_file_real, dpi=300, bbox_inches='tight')
plt.close()
print(f"已保存实部图到: {out_file_real}")
# 虚部
fig, ax = plt.subplots(figsize=(8, 8))
tcf = ax.tricontourf(triang, x_imag, levels=100, cmap='RdBu_r')
cbar = fig.colorbar(tcf, ax=ax)
cbar.set_label('Im(u)', fontsize=14)
ax.set_title(f'Imaginary Part ({title_suffix})', fontsize=16, fontweight='bold')
ax.set_aspect('equal')
ax.set_xlim(x_coord.min(), x_coord.max())
ax.set_ylim(y_coord.min(), y_coord.max())
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
out_file_imag = os.path.join(output_dir, f"{data_label}_imag.png")
plt.savefig(out_file_imag, dpi=300, bbox_inches='tight')
plt.close()
print(f"已保存虚部图到: {out_file_imag}")
# 模值
fig, ax = plt.subplots(figsize=(8, 8))
tcf = ax.tricontourf(triang, x_magnitude, levels=100, cmap='jet')
cbar = fig.colorbar(tcf, ax=ax)
cbar.set_label('|u|', fontsize=14)
ax.set_title(f'Magnitude ({title_suffix})', fontsize=16, fontweight='bold')
ax.set_aspect('equal')
ax.set_xlim(x_coord.min(), x_coord.max())
ax.set_ylim(y_coord.min(), y_coord.max())
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
out_file_mag = os.path.join(output_dir, f"{data_label}_magnitude.png")
plt.savefig(out_file_mag, dpi=300, bbox_inches='tight')
plt.close()
print(f"已保存模值图到: {out_file_mag}")
if __name__ == "__main__":
# 可视化指定算例
# 尝试从 save_test 模块读取 k 值，避免循环导入
try:
# 使用动态导入来避免启动时的循环导入问题
import sys
if 'save_test' in sys.modules:
save_test_module = sys.modules['save_test']
current_k = getattr(save_test_module, 'k', DEFAULT_K)
print(f"从已加载的 save_test 模块读取到 k={current_k}")
else:
current_k = DEFAULT_K
print(f"save_test 模块未加载，使用默认 k={current_k}")
except Exception as e:
current_k = DEFAULT_K
print(f"读取 k 值失败: {e}，使用默认 k={current_k}")
# 可视化预测结果（默认）
print(f"正在可视化算例 k={current_k} 的预测结果...")
visualize_solution(k=current_k, use_prediction=True, save_combined=True, save_separate=False)
# 可选：同时可视化真实值进行对比
print(f"正在可视化算例 k={current_k} 的真实值...")
visualize_solution(k=current_k, use_prediction=False, save_combined=True, save_separate=False)