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)