NDESolver/plot_diffusion.py at main · SphenHe/NDESolver · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
"""
中子扩散方程求解器 - 绘图模块

该模块负责从HDF5文件读取计算结果并生成可视化图像。

主要功能：
1. 从HDF5文件读取中子通量和有效增殖因子数据
2. 绘制中子通量分布图（phi1和phi2）
3. 绘制有效增殖因子收敛曲线和误差曲线
4. 保存图像到指定目录

输入：HDF5格式的计算结果文件
输出：PDF格式的图像文件
"""

import argparse
import numpy as np
import matplotlib.pyplot as plt
import h5py
import os


def plot_phi(h5_file, output_dir=None):
    '''
    绘制中子通量分布图

    参数:
        h5_file: HDF5文件路径
        output_dir: 输出目录，如果为None则使用HDF5文件中的predir
    '''
    with h5py.File(h5_file, 'r') as hf:
        phi1 = hf['phi1'][:]
        phi2 = hf['phi2'][:]
        length_x = hf['length_x'][()]
        length_y = hf['length_y'][()]
        geo_length_x = hf['geo_length_x'][()]
        geo_length_y = hf['geo_length_y'][()]
        geo_num_x = hf['geo_num_x'][()]
        geo_num_y = hf['geo_num_y'][()]

        if output_dir is None:
            output_dir = hf.attrs['figures_dir']

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 准备绘图数据
    xarr = [i*geo_length_x for i in range(geo_num_x+1)]
    yarr = [i*geo_length_y for i in range(geo_num_y+1)]

    fig, ax = plt.subplots(1, 2, figsize=(14, 7))

    # 绘制 phi1
    phi1max = np.max(np.abs(phi1))
    im = ax[0].imshow(phi1, cmap='coolwarm', vmin=0, vmax=phi1max, origin='lower',
                      extent=(0, length_x, 0, length_y))
    fig.colorbar(im, ax=ax[0], orientation='vertical', fraction=0.046, pad=0.04)
    ax[0].set_title("$\\phi_1$ Neutron Flux")
    ax[0].grid(True, which='both', linestyle='--', linewidth=0.5, color='white')
    ax[0].set_xticks(xarr)
    ax[0].set_yticks(yarr)

    # 绘制 phi2
    phi2max = np.max(np.abs(phi2))
    im = ax[1].imshow(phi2, cmap='coolwarm', vmin=0, vmax=phi2max,
                      extent=(0, length_x, 0, length_y), origin='lower')
    fig.colorbar(im, ax=ax[1], orientation='vertical', fraction=0.046, pad=0.04)
    ax[1].set_title("$\\phi_2$ Neutron Flux")
    ax[1].grid(True, which='both', linestyle='--', linewidth=0.5, color='white')
    ax[1].set_xticks(xarr)
    ax[1].set_yticks(yarr)

    plt.tight_layout()
    output_path = f'{output_dir}/phi.pdf'
    plt.savefig(output_path)
    plt.close()

    print(f"中子通量图已保存到: {output_path}")


def plot_keff(h5_file, output_dir=None):
    '''
    绘制有效增殖因子收敛曲线和误差曲线

    参数:
        h5_file: HDF5文件路径
        output_dir: 输出目录，如果为None则使用HDF5文件中的predir
    '''
    with h5py.File(h5_file, 'r') as hf:
        keff = hf['keff'][()]
        keff_data = hf['keff_data'][:]
        keff_ref = hf['keff_ref'][()]

        if output_dir is None:
            output_dir = hf.attrs['figures_dir']

    os.makedirs(output_dir, exist_ok=True)

    fig, ax = plt.subplots(1, 2, figsize=(10, 5))

    # 绘制收敛曲线
    ax[0].plot(keff_data, label='$k_{eff}$')
    ax[0].plot([0, len(keff_data)], [keff_ref, keff_ref], 'r--',
               label=f"$k_{{eff,ref}}={keff_ref}$")
    ax[0].plot([0, len(keff_data)], [keff, keff], 'g--',
               label=f"$k_{{eff,final}}={keff:.6f}$")
    ax[0].set_xlabel('Iteration')
    ax[0].set_ylabel('keff')
    ax[0].set_title('Effective Multiplication Factor')
    ax[0].grid(True, which='both', linestyle='--', linewidth=0.5)
    ax[0].legend()

    # 绘制相对误差曲线
    relative_error = np.abs(keff_data - keff_ref) / keff_ref
    final_error = np.abs(keff - keff_ref) / keff_ref

    ax[1].plot(relative_error, label='Relative error of $k_{eff}$')
    ax[1].plot([0, len(keff_data)], [final_error, final_error], 'g--',
               label=f'Relative error of $k_{{eff, final}}={final_error:.2e}$')
    ax[1].set_xlabel('Iteration')
    ax[1].set_ylabel('Relative error of keff')
    ax[1].set_yscale('log')
    ax[1].set_title('Relative Error of Effective Multiplication Factor')
    ax[1].grid(True, which='both', linestyle='--', linewidth=0.5)
    ax[1].legend()

    plt.tight_layout()
    output_path = f'{output_dir}/keff.pdf'
    plt.savefig(output_path)
    plt.close()

    print(f"有效增殖因子图已保存到: {output_path}")


def plot_all(h5_file, output_dir=None):
    '''
    绘制所有图像

    参数:
        h5_file: HDF5文件路径
        output_dir: 输出目录，如果为None则使用HDF5文件中的predir
    '''
    plot_phi(h5_file, output_dir)
    plot_keff(h5_file, output_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='中子扩散方程求解器 - 绘图模块')
    parser.add_argument("--data", type=str, required=True, help='HDF5数据文件路径')
    parser.add_argument("--output", type=str, default=None,
                       help='输出目录（可选，默认使用数据文件中的predir）')
    parser.add_argument("--type", type=str, choices=['phi', 'keff', 'all'], default='all',
                       help='绘图类型：phi（中子通量）、keff（有效增殖因子）或all（全部）')
    args = parser.parse_args()

    if not os.path.exists(args.data):
        raise FileNotFoundError(f"数据文件未找到: {args.data}")

    if args.type == 'phi':
        plot_phi(args.data, args.output)
    elif args.type == 'keff':
        plot_keff(args.data, args.output)
    else:
        plot_all(args.data, args.output)