OpenMC/tests/regression_tests/surface_source_write/_visualize.py
Joffrey Dorville ddc9526966
Storing surface source points using a cell ID (#2888)
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
2024-06-19 10:11:58 -05:00

66 lines
2 KiB
Python

"""Helper script to visualize the surface_source.h5 files created with this test.
"""
import h5py
import matplotlib.pyplot as plt
if __name__ == "__main__":
# Select an option
# "show": 3D visualization using matplotlib
# "savefig": 2D representation using matplotlib and storing the fig under plot_2d.png
option = "show"
# option = "savefig"
# Select the case from its folder name
folder = "case-20"
# Reading the surface source file
with h5py.File(f"{folder}/surface_source_true.h5", "r") as fp:
source_bank = fp["source_bank"][()]
r_xs = source_bank['r']['x']
r_ys = source_bank['r']['y']
r_zs = source_bank['r']['z']
print("Size of the source bank: ", len(source_bank))
# Select data range to visualize
idx_1 = 0
idx_2 = -1
# Show 3D representation
if option == "show":
fig = plt.figure(figsize=(10, 10))
ax1 = fig.add_subplot(projection="3d", proj_type="ortho")
ax1.scatter(r_xs[idx_1:idx_2], r_ys[idx_1:idx_2], r_zs[idx_1:idx_2], marker=".")
ax1.view_init(0, 0)
ax1.xaxis.set_ticklabels([])
ax1.set_ylabel("y-axis [cm]")
ax1.set_zlabel("z-axis [cm]")
ax1.set_aspect("equal", "box")
plt.show()
# Save 2D representations
elif option == "savefig":
fig = plt.figure(figsize=(14, 5))
ax1 = fig.add_subplot(121, projection="3d", proj_type="ortho")
ax1.scatter(r_xs[idx_1:idx_2], r_ys[idx_1:idx_2], r_zs[idx_1:idx_2], marker=".")
ax1.view_init(0, 0)
ax1.xaxis.set_ticklabels([])
ax1.set_ylabel("y-axis [cm]")
ax1.set_zlabel("z-axis [cm]")
ax1.set_aspect("equal", "box")
ax2 = fig.add_subplot(122, projection="3d", proj_type="ortho")
ax2.scatter(r_xs[idx_1:idx_2], r_ys[idx_1:idx_2], r_zs[idx_1:idx_2], marker=".")
ax2.view_init(90, -90)
ax2.zaxis.set_ticklabels([])
ax2.set_xlabel("x-axis [cm]")
ax2.set_ylabel("y-axis [cm]")
ax2.set_aspect("equal", "box")
plt.savefig("plot_2d.png")