282 lines
11 KiB
Python
282 lines
11 KiB
Python
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)
|
||
|
||
|
||
|