Add VTU export for Unstructured meshes (#3290)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
This commit is contained in:
rherrero-pf 2025-02-13 22:26:36 +01:00 committed by GitHub
parent 18c3112415
commit 78bf4cf145
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 136 additions and 31 deletions

View file

@ -9,6 +9,7 @@ from numbers import Integral, Real
import h5py
import lxml.etree as ET
import numpy as np
from pathlib import Path
import openmc
import openmc.checkvalue as cv
@ -2358,55 +2359,71 @@ class UnstructuredMesh(MeshBase):
self.write_data_to_vtk(**kwargs)
def write_data_to_vtk(
self,
filename: PathLike | None = None,
datasets: dict | None = None,
volume_normalization: bool = True
self,
filename: PathLike | None = None,
datasets: dict | None = None,
volume_normalization: bool = True,
):
"""Map data to unstructured VTK mesh elements.
If filename is None, then a filename will be generated based on the mesh
ID, and exported to VTK format.
Parameters
----------
filename : str or pathlib.Path
Name of the VTK file to write
Name of the VTK file to write. If the filename ends in '.vtu' then a
binary VTU format file will be written, if the filename ends in
'.vtk' then a legacy VTK file will be written.
datasets : dict
Dictionary whose keys are the data labels
and values are numpy appropriately sized arrays
of the data
Dictionary whose keys are the data labels and values are numpy
appropriately sized arrays of the data
volume_normalization : bool
Whether or not to normalize the data by the
volume of the mesh elements
Whether or not to normalize the data by the volume of the mesh
elements
"""
import vtk
from vtk.util import numpy_support as nps
from vtkmodules.util import numpy_support
from vtkmodules import vtkCommonCore
from vtkmodules import vtkCommonDataModel
from vtkmodules import vtkIOLegacy
from vtkmodules import vtkIOXML
if self.connectivity is None or self.vertices is None:
raise RuntimeError('This mesh has not been '
'loaded from a statepoint file.')
raise RuntimeError(
"This mesh has not been loaded from a statepoint file."
)
if filename is None:
filename = f'mesh_{self.id}.vtk'
filename = f"mesh_{self.id}.vtk"
writer = vtk.vtkUnstructuredGridWriter()
if Path(filename).suffix == ".vtk":
writer = vtkIOLegacy.vtkUnstructuredGridWriter()
elif Path(filename).suffix == ".vtu":
writer = vtkIOXML.vtkXMLUnstructuredGridWriter()
writer.SetCompressorTypeToZLib()
writer.SetDataModeToBinary()
writer.SetFileName(str(filename))
grid = vtk.vtkUnstructuredGrid()
grid = vtkCommonDataModel.vtkUnstructuredGrid()
vtk_pnts = vtk.vtkPoints()
vtk_pnts.SetData(nps.numpy_to_vtk(self.vertices))
grid.SetPoints(vtk_pnts)
points = vtkCommonCore.vtkPoints()
points.SetData(numpy_support.numpy_to_vtk(self.vertices))
grid.SetPoints(points)
n_skipped = 0
for elem_type, conn in zip(self.element_types, self.connectivity):
if elem_type == self._LINEAR_TET:
elem = vtk.vtkTetra()
elem = vtkCommonDataModel.vtkTetra()
elif elem_type == self._LINEAR_HEX:
elem = vtk.vtkHexahedron()
elem = vtkCommonDataModel.vtkHexahedron()
elif elem_type == self._UNSUPPORTED_ELEM:
n_skipped += 1
continue
else:
raise RuntimeError(f'Invalid element type {elem_type} found')
raise RuntimeError(f"Invalid element type {elem_type} found")
for i, c in enumerate(conn):
if c == -1:
break
@ -2415,30 +2432,36 @@ class UnstructuredMesh(MeshBase):
grid.InsertNextCell(elem.GetCellType(), elem.GetPointIds())
if n_skipped > 0:
warnings.warn(f'{n_skipped} elements were not written because '
'they are not of type linear tet/hex')
warnings.warn(
f"{n_skipped} elements were not written because "
"they are not of type linear tet/hex"
)
# check that datasets are the correct size
datasets_out = []
if datasets is not None:
for name, data in datasets.items():
if data.shape != self.dimension:
raise ValueError(f'Cannot apply dataset "{name}" with '
f'shape {data.shape} to mesh {self.id} '
f'with dimensions {self.dimension}')
raise ValueError(
f'Cannot apply dataset "{name}" with '
f"shape {data.shape} to mesh {self.id} "
f"with dimensions {self.dimension}"
)
if volume_normalization:
for name, data in datasets.items():
if np.issubdtype(data.dtype, np.integer):
warnings.warn(f'Integer data set "{name}" will '
'not be volume-normalized.')
warnings.warn(
f'Integer data set "{name}" will '
"not be volume-normalized."
)
continue
data /= self.volumes
# add data to the mesh
for name, data in datasets.items():
datasets_out.append(data)
arr = vtk.vtkDoubleArray()
arr = vtkCommonCore.vtkDoubleArray()
arr.SetName(name)
arr.SetNumberOfTuples(data.size)

View file

@ -1,8 +1,11 @@
from math import pi
from pathlib import Path
import numpy as np
import pytest
import openmc
import openmc.lib
from openmc.utility_funcs import change_directory
@pytest.mark.parametrize("val_left,val_right", [(0, 0), (-1., -1.), (2.0, 2)])
@ -398,6 +401,85 @@ def test_umesh_roundtrip(run_in_tmpdir, request):
assert umesh.id == xml_mesh.id
@pytest.fixture(scope='module')
def simple_umesh(request):
"""Fixture returning UnstructuredMesh with all attributes"""
surf1 = openmc.Sphere(r=20.0, boundary_type="vacuum")
material1 = openmc.Material()
material1.add_element("H", 1.0)
material1.set_density('g/cm3', 1.0)
materials = openmc.Materials([material1])
cell1 = openmc.Cell(region=-surf1, fill=material1)
geometry = openmc.Geometry([cell1])
umesh = openmc.UnstructuredMesh(
filename=request.path.parent.parent
/ "regression_tests/external_moab/test_mesh_tets.h5m",
library="moab",
mesh_id=1
)
# setting ID to make it easier to get the mesh from the statepoint later
mesh_filter = openmc.MeshFilter(umesh)
# Create flux mesh tally to score alpha production
mesh_tally = openmc.Tally(name="test_tally")
mesh_tally.filters = [mesh_filter]
mesh_tally.scores = ["total"]
tallies = openmc.Tallies([mesh_tally])
settings = openmc.Settings()
settings.run_mode = "fixed source"
settings.batches = 2
settings.particles = 100
settings.source = openmc.IndependentSource(
space=openmc.stats.Point((0.1, 0.1, 0.1))
)
model = openmc.Model(
materials=materials, geometry=geometry, settings=settings, tallies=tallies
)
with change_directory(tmpdir=True):
statepoint_file = model.run()
with openmc.StatePoint(statepoint_file) as sp:
return sp.meshes[1]
@pytest.mark.skipif(not openmc.lib._dagmc_enabled(), reason="DAGMC not enabled.")
@pytest.mark.parametrize('export_type', ('.vtk', '.vtu'))
def test_umesh(run_in_tmpdir, simple_umesh, export_type):
"""Performs a minimal UnstructuredMesh simulation, reads in the resulting
statepoint file and writes the mesh data to vtk and vtkhdf files. It is
necessary to read in the unstructured mesh from a statepoint file to ensure
it has all the required attributes
"""
# Get VTK modules
vtkIOLegacy = pytest.importorskip("vtkmodules.vtkIOLegacy")
vtkIOXML = pytest.importorskip("vtkmodules.vtkIOXML")
# Sample some random data and write to VTK
rng = np.random.default_rng()
ref_data = rng.random(simple_umesh.dimension)
filename = f"test_mesh{export_type}"
simple_umesh.write_data_to_vtk(datasets={"mean": ref_data}, filename=filename)
assert Path(filename).exists()
if export_type == ".vtk":
reader = vtkIOLegacy.vtkGenericDataObjectReader()
elif export_type == ".vtu":
reader = vtkIOXML.vtkXMLGenericDataObjectReader()
reader.SetFileName(str(filename))
reader.Update()
# Get mean from file and make sure it matches original data
arr = reader.GetOutput().GetCellData().GetArray("mean")
mean = np.array([arr.GetTuple1(i) for i in range(ref_data.size)])
np.testing.assert_almost_equal(mean, ref_data)
def test_mesh_get_homogenized_materials():
"""Test the get_homogenized_materials method"""
# Simple model with 1 cm of Fe56 next to 1 cm of H1