diff --git a/tests/unit_tests/test_tallies.py b/tests/unit_tests/test_tallies.py index aeeb0612d5..b6fc6cc478 100644 --- a/tests/unit_tests/test_tallies.py +++ b/tests/unit_tests/test_tallies.py @@ -1,4 +1,8 @@ import numpy as np +import pytest +import vtk +from os.path import exists + import openmc @@ -38,3 +42,61 @@ def test_xml_roundtrip(run_in_tmpdir): assert new_tally.triggers[0].trigger_type == tally.triggers[0].trigger_type assert new_tally.triggers[0].threshold == tally.triggers[0].threshold assert new_tally.triggers[0].scores == tally.triggers[0].scores + + +cylinder_mesh = openmc.CylindricalMesh() +cylinder_mesh.r_grid = np.linspace(1, 2, num=30) +cylinder_mesh.phi_grid = np.linspace(0, np.pi / 2, num=50) +cylinder_mesh.z_grid = np.linspace(0, 1, num=30) + +regular_mesh = openmc.RegularMesh() +regular_mesh.lower_left = [0, 0, 0] +regular_mesh.upper_right = [1, 1, 1] +regular_mesh.dimension = [10, 5, 6] + +rectilinear_mesh = openmc.RectilinearMesh() +rectilinear_mesh.x_grid = np.linspace(0, 1) +rectilinear_mesh.y_grid = np.linspace(0, 1) +rectilinear_mesh.z_grid = np.linspace(0, 1) + + +@pytest.mark.parametrize("mesh", [cylinder_mesh, regular_mesh, rectilinear_mesh]) +def test_voxels_to_vtk(mesh): + vtk_grid = openmc.voxels_to_vtk(mesh, mean=None, std_dev=None) + assert isinstance(vtk_grid, vtk.vtkStructuredGrid) + + +@pytest.mark.parametrize("mesh", [cylinder_mesh, regular_mesh, rectilinear_mesh]) +def test_write_to_vtk(mesh, tmpdir): + # build + tally = openmc.Tally() + tally.filters = [openmc.MeshFilter(mesh)] + filename = tmpdir / "out.vtk" + # run + tally.write_to_vtk(filename) + # test + assert exists(filename) + + +def test_write_to_vtk_raises_error_when_no_meshfilter(): + # build + tally = openmc.Tally() + + # test + expected_err_msg = "write_to_vtk requires a MeshFilter in the tally filters" + with pytest.raises(ValueError, match=expected_err_msg): + tally.write_to_vtk("out.vtk") + + +def test_voxel_to_vtk_raises_error_with_wrong_mesh(): + # build + tally = openmc.Tally() + spherical_mesh = openmc.SphericalMesh() + spherical_mesh.r_grid = np.linspace(1, 2) + spherical_mesh.phi_grid = np.linspace(1, 2) + spherical_mesh.theta_grid = np.linspace(1, 2) + tally.filters = [openmc.MeshFilter(spherical_mesh)] + # test + expected_err_msg = "voxels_to_vtk only works with openmc.RegularMesh, openmc.RectilinearMesh, openmc.CylindricalMesh" + with pytest.raises(ValueError, match=expected_err_msg): + tally.write_to_vtk("out.vtk") \ No newline at end of file