mirror of
https://github.com/openmc-dev/openmc.git
synced 2026-07-28 14:15:42 -04:00
Merge pull request #2096 from RemDelaporteMathurin/cylindermesh_to_vtk
`.write_data_to_vtk()` method in meshes
This commit is contained in:
Jonathan Shimwell
GitHub
commit
f081b28b1d
3 changed files with 307 additions and 3 deletions
233
openmc/mesh.py
233
openmc/mesh.py
|
|
@ -193,6 +193,77 @@ class StructuredMesh(MeshBase):
|
|||
s1 = (slice(1, None),)*ndim + (slice(None),)
|
||||
return (vertices[s0] + vertices[s1]) / 2
|
||||
|
||||
def write_data_to_vtk(self, points, filename, datasets, volume_normalization=True):
|
||||
"""Creates a VTK object of the mesh
|
||||
|
||||
Parameters
|
||||
----------
|
||||
points : list or np.array
|
||||
List of (X,Y,Z) tuples.
|
||||
filename : str
|
||||
Name of the VTK file to write.
|
||||
datasets : dict
|
||||
Dictionary whose keys are the data labels
|
||||
and values are the data sets.
|
||||
volume_normalization : bool, optional
|
||||
Whether or not to normalize the data by
|
||||
the volume of the mesh elements.
|
||||
|
||||
Raises
|
||||
------
|
||||
RuntimeError
|
||||
When the size of a dataset doesn't match the number of cells
|
||||
|
||||
Returns
|
||||
-------
|
||||
vtk.vtkStructuredGrid
|
||||
the VTK object
|
||||
"""
|
||||
|
||||
import vtk
|
||||
from vtk.util import numpy_support as nps
|
||||
|
||||
# check that the data sets are appropriately sized
|
||||
errmsg = "The size of the dataset {} should be equal to the number of cells"
|
||||
for label, dataset in datasets.items():
|
||||
if isinstance(dataset, np.ndarray):
|
||||
if not dataset.size == self.dimension[0] * self.dimension[1]* self.dimension[2]:
|
||||
raise RuntimeError(errmsg.format(label))
|
||||
else:
|
||||
if len(dataset) == self.dimension[0] * self.dimension[1]* self.dimension[2]:
|
||||
raise RuntimeError(errmsg.format(label))
|
||||
cv.check_type('label', label, str)
|
||||
|
||||
vtk_grid = vtk.vtkStructuredGrid()
|
||||
|
||||
vtk_grid.SetDimensions(*self.dimension)
|
||||
|
||||
vtkPts = vtk.vtkPoints()
|
||||
vtkPts.SetData(nps.numpy_to_vtk(points, deep=True))
|
||||
vtk_grid.SetPoints(vtkPts)
|
||||
|
||||
# create VTK arrays for each of
|
||||
# the data sets
|
||||
for label, dataset in datasets.items():
|
||||
dataset = np.asarray(dataset).flatten()
|
||||
|
||||
if volume_normalization:
|
||||
dataset /= self.volumes.flatten()
|
||||
|
||||
dataset_array = vtk.vtkDoubleArray()
|
||||
dataset_array.SetName(label)
|
||||
dataset_array.SetArray(nps.numpy_to_vtk(dataset),
|
||||
dataset.size,
|
||||
True)
|
||||
vtk_grid.GetCellData().AddArray(dataset_array)
|
||||
|
||||
# write the .vtk file
|
||||
writer = vtk.vtkStructuredGridWriter()
|
||||
writer.SetFileName(str(filename))
|
||||
writer.SetInputData(vtk_grid)
|
||||
writer.Write()
|
||||
|
||||
return vtk_grid
|
||||
|
||||
class RegularMesh(StructuredMesh):
|
||||
"""A regular Cartesian mesh in one, two, or three dimensions
|
||||
|
|
@ -621,7 +692,42 @@ class RegularMesh(StructuredMesh):
|
|||
root_cell.fill = lattice
|
||||
|
||||
return root_cell, cells
|
||||
|
||||
def write_data_to_vtk(self, filename, datasets, volume_normalization=True):
|
||||
"""Creates a VTK object of the mesh
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename : str or pathlib.Path
|
||||
Name of the VTK file to write.
|
||||
datasets : dict
|
||||
Dictionary whose keys are the data labels
|
||||
and values are the data sets.
|
||||
volume_normalization : bool, optional
|
||||
Whether or not to normalize the data by
|
||||
the volume of the mesh elements.
|
||||
Defaults to True.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vtk.vtkStructuredGrid
|
||||
the VTK object
|
||||
"""
|
||||
|
||||
ll, ur = self.lower_left, self.upper_right
|
||||
x_vals = np.linspace(ll[0], ur[0], num=self.dimension[0] + 1)
|
||||
y_vals = np.linspace(ll[1], ur[1], num=self.dimension[1] + 1)
|
||||
z_vals = np.linspace(ll[2], ur[2], num=self.dimension[2] + 1)
|
||||
|
||||
# create points
|
||||
pts_cartesian = np.array([[x, y, z] for z in z_vals for y in y_vals for x in x_vals])
|
||||
|
||||
return super().write_data_to_vtk(
|
||||
points=pts_cartesian,
|
||||
filename=filename,
|
||||
datasets=datasets,
|
||||
volume_normalization=volume_normalization
|
||||
)
|
||||
|
||||
def Mesh(*args, **kwargs):
|
||||
warnings.warn("Mesh has been renamed RegularMesh. Future versions of "
|
||||
|
|
@ -821,6 +927,36 @@ class RectilinearMesh(StructuredMesh):
|
|||
|
||||
return element
|
||||
|
||||
def write_data_to_vtk(self, filename, datasets, volume_normalization=True):
|
||||
"""Creates a VTK object of the mesh
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename : str or pathlib.Path
|
||||
Name of the VTK file to write.
|
||||
datasets : dict
|
||||
Dictionary whose keys are the data labels
|
||||
and values are the data sets.
|
||||
volume_normalization : bool, optional
|
||||
Whether or not to normalize the data by
|
||||
the volume of the mesh elements.
|
||||
Defaults to True.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vtk.vtkStructuredGrid
|
||||
the VTK object
|
||||
"""
|
||||
# create points
|
||||
pts_cartesian = np.array([[x, y, z] for z in self.z_grid for y in self.y_grid for x in self.x_grid])
|
||||
|
||||
return super().write_data_to_vtk(
|
||||
points=pts_cartesian,
|
||||
filename=filename,
|
||||
datasets=datasets,
|
||||
volume_normalization=volume_normalization
|
||||
)
|
||||
|
||||
|
||||
class CylindricalMesh(StructuredMesh):
|
||||
"""A 3D cylindrical mesh
|
||||
|
|
@ -1012,6 +1148,46 @@ class CylindricalMesh(StructuredMesh):
|
|||
|
||||
return np.multiply.outer(np.outer(V_r, V_p), V_z)
|
||||
|
||||
def write_data_to_vtk(self, filename, datasets, volume_normalization=True):
|
||||
"""Creates a VTK object of the mesh
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename : str or pathlib.Path
|
||||
Name of the VTK file to write.
|
||||
datasets : dict
|
||||
Dictionary whose keys are the data labels
|
||||
and values are the data sets.
|
||||
volume_normalization : bool, optional
|
||||
Whether or not to normalize the data by
|
||||
the volume of the mesh elements.
|
||||
Defaults to True.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vtk.vtkStructuredGrid
|
||||
the VTK object
|
||||
"""
|
||||
# create points
|
||||
pts_cylindrical = np.array(
|
||||
[
|
||||
[r, phi, z]
|
||||
for z in self.z_grid
|
||||
for phi in self.phi_grid
|
||||
for r in self.r_grid
|
||||
]
|
||||
)
|
||||
pts_cartesian = np.copy(pts_cylindrical)
|
||||
r, phi = pts_cylindrical[:, 0], pts_cylindrical[:, 1]
|
||||
pts_cartesian[:, 0] = r * np.cos(phi)
|
||||
pts_cartesian[:, 1] = r * np.sin(phi)
|
||||
|
||||
return super().write_data_to_vtk(
|
||||
points=pts_cartesian,
|
||||
filename=filename,
|
||||
datasets=datasets,
|
||||
volume_normalization=volume_normalization
|
||||
)
|
||||
|
||||
class SphericalMesh(StructuredMesh):
|
||||
"""A 3D spherical mesh
|
||||
|
|
@ -1204,6 +1380,49 @@ class SphericalMesh(StructuredMesh):
|
|||
|
||||
return np.multiply.outer(np.outer(V_r, V_t), V_p)
|
||||
|
||||
def write_data_to_vtk(self, filename, datasets, volume_normalization=True):
|
||||
"""Creates a VTK object of the mesh
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename : str or pathlib.Path
|
||||
Name of the VTK file to write.
|
||||
datasets : dict
|
||||
Dictionary whose keys are the data labels
|
||||
and values are the data sets.
|
||||
volume_normalization : bool, optional
|
||||
Whether or not to normalize the data by
|
||||
the volume of the mesh elements.
|
||||
Defaults to True.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vtk.vtkStructuredGrid
|
||||
the VTK object
|
||||
"""
|
||||
|
||||
# create points
|
||||
pts_spherical = np.array(
|
||||
[
|
||||
[r, theta, phi]
|
||||
for phi in self.phi_grid
|
||||
for theta in self.theta_grid
|
||||
for r in self.r_grid
|
||||
]
|
||||
)
|
||||
pts_cartesian = np.copy(pts_spherical)
|
||||
r, theta, phi = pts_spherical[:, 0], pts_spherical[:, 1], pts_spherical[:, 2]
|
||||
pts_cartesian[:, 0] = r * np.sin(phi) * np.cos(theta)
|
||||
pts_cartesian[:, 1] = r * np.sin(phi) * np.sin(theta)
|
||||
pts_cartesian[:, 2] = r * np.cos(phi)
|
||||
|
||||
return super().write_data_to_vtk(
|
||||
points=pts_cartesian,
|
||||
filename=filename,
|
||||
datasets=datasets,
|
||||
volume_normalization=volume_normalization
|
||||
)
|
||||
|
||||
|
||||
class UnstructuredMesh(MeshBase):
|
||||
"""A 3D unstructured mesh
|
||||
|
|
@ -1373,7 +1592,7 @@ class UnstructuredMesh(MeshBase):
|
|||
|
||||
Parameters
|
||||
----------
|
||||
filename : str
|
||||
filename : str or pathlib.Path
|
||||
Name of the VTK file to write.
|
||||
datasets : dict
|
||||
Dictionary whose keys are the data labels
|
||||
|
|
@ -1381,6 +1600,11 @@ class UnstructuredMesh(MeshBase):
|
|||
volume_normalization : bool
|
||||
Whether or not to normalize the data by the
|
||||
volume of the mesh elements
|
||||
|
||||
Raises
|
||||
------
|
||||
RuntimeError
|
||||
when the size of a dataset doesn't match the number of cells
|
||||
"""
|
||||
|
||||
import vtk
|
||||
|
|
@ -1397,11 +1621,14 @@ class UnstructuredMesh(MeshBase):
|
|||
" mesh information from a statepoint file.")
|
||||
|
||||
# check that the data sets are appropriately sized
|
||||
errmsg = "The size of the dataset {} should be equal to the number of cells"
|
||||
for label, dataset in datasets.items():
|
||||
if isinstance(dataset, np.ndarray):
|
||||
assert dataset.size == self.n_elements
|
||||
if not dataset.size == self.dimension[0] * self.dimension[1]* self.dimension[2]:
|
||||
raise RuntimeError(errmsg.format(label))
|
||||
else:
|
||||
assert len(dataset) == self.n_elements
|
||||
if len(dataset) == self.dimension[0] * self.dimension[1]* self.dimension[2]:
|
||||
raise RuntimeError(errmsg.format(label))
|
||||
cv.check_type('label', label, str)
|
||||
|
||||
# create data arrays for the cells/points
|
||||
|
|
|
|||
76
tests/unit_tests/test_mesh_to_vtk.py
Normal file
76
tests/unit_tests/test_mesh_to_vtk.py
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
import numpy as np
|
||||
from pathlib import Path
|
||||
import pytest
|
||||
|
||||
vtk = pytest.importorskip("vtk")
|
||||
from vtk.util import numpy_support as nps
|
||||
|
||||
import openmc
|
||||
|
||||
|
||||
regular_mesh = openmc.RegularMesh()
|
||||
regular_mesh.lower_left = (0, 0, 0)
|
||||
regular_mesh.upper_right = (1, 1, 1)
|
||||
regular_mesh.dimension = [30, 20, 10]
|
||||
|
||||
rectilinear_mesh = openmc.RectilinearMesh()
|
||||
rectilinear_mesh.x_grid = np.linspace(1, 2, num=30)
|
||||
rectilinear_mesh.y_grid = np.linspace(1, 2, num=30)
|
||||
rectilinear_mesh.z_grid = np.linspace(1, 2, num=30)
|
||||
|
||||
cylinder_mesh = openmc.CylindricalMesh()
|
||||
cylinder_mesh.r_grid = np.linspace(1, 2, num=30)
|
||||
cylinder_mesh.phi_grid = np.linspace(0, np.pi, num=50)
|
||||
cylinder_mesh.z_grid = np.linspace(0, 1, num=30)
|
||||
|
||||
spherical_mesh = openmc.SphericalMesh()
|
||||
spherical_mesh.r_grid = np.linspace(1, 2, num=30)
|
||||
spherical_mesh.phi_grid = np.linspace(0, np.pi, num=50)
|
||||
spherical_mesh.theta_grid = np.linspace(0, np.pi / 2, num=30)
|
||||
|
||||
@pytest.mark.parametrize("mesh", [cylinder_mesh, regular_mesh, rectilinear_mesh, spherical_mesh])
|
||||
def test_write_data_to_vtk(mesh, tmpdir):
|
||||
# BUILD
|
||||
filename = Path(tmpdir) / "out.vtk"
|
||||
|
||||
data = np.random.random(mesh.dimension[0]*mesh.dimension[1]*mesh.dimension[2])
|
||||
|
||||
# RUN
|
||||
mesh.write_data_to_vtk(filename=filename, datasets={"label1": data, "label2": data})
|
||||
|
||||
# TEST
|
||||
assert filename.is_file()
|
||||
|
||||
# read file
|
||||
reader = vtk.vtkStructuredGridReader()
|
||||
reader.SetFileName(filename)
|
||||
reader.Update()
|
||||
|
||||
# check name of datasets
|
||||
vtk_grid = reader.GetOutput()
|
||||
array1 = vtk_grid.GetCellData().GetArray(0)
|
||||
array2 = vtk_grid.GetCellData().GetArray(1)
|
||||
|
||||
assert array1.GetName() == "label1"
|
||||
assert array2.GetName() == "label2"
|
||||
|
||||
# check size of datasets
|
||||
assert nps.vtk_to_numpy(array1).size == data.size
|
||||
assert nps.vtk_to_numpy(array2).size == data.size
|
||||
|
||||
@pytest.mark.parametrize("mesh", [cylinder_mesh, regular_mesh, rectilinear_mesh, spherical_mesh])
|
||||
def test_write_data_to_vtk_size_mismatch(mesh):
|
||||
"""Checks that an error is raised when the size of the dataset
|
||||
doesn't match the mesh number of cells
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mesh : openmc.StructuredMesh
|
||||
The mesh to test
|
||||
"""
|
||||
right_size = mesh.dimension[0]*mesh.dimension[1]*mesh.dimension[2]
|
||||
data = np.random.random(right_size + 1)
|
||||
|
||||
expected_error_msg = "The size of the dataset label should be equal to the number of cells"
|
||||
with pytest.raises(RuntimeError, match=expected_error_msg):
|
||||
mesh.write_data_to_vtk(filename="out.vtk", datasets={"label": data})
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import numpy as np
|
||||
|
||||
import openmc
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue