Adding vtkhdf option to write vtk data

openmc/mesh.py

@@ -5,6 +5,7 @@
 from functools import wraps
 from math import pi, sqrt, atan2
 from numbers import Integral, Real
+from pathlib import Path
 from typing import Protocol
 
 import h5py
@@ -2655,7 +2656,8 @@ def write_vtk_mesh(self, **kwargs):
         warnings.warn(
             "The 'UnstructuredMesh.write_vtk_mesh' method has been renamed "
             "to 'write_data_to_vtk' and will be removed in a future version "
-            " of OpenMC.", FutureWarning
+            " of OpenMC.",
+            FutureWarning,
         )
         self.write_data_to_vtk(**kwargs)
 
@@ -2673,26 +2675,54 @@ def write_data_to_vtk(
         Parameters
         ----------
         filename : str or pathlib.Path
-            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.
+            Name of the VTK file to write. If the filename ends in '.vtkhdf'
+            then a VTKHDF format file will be written. 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
         volume_normalization : bool
             Whether or not to normalize the data by the volume of the mesh
             elements
         """
+
+        if Path(filename).suffix == ".vtkhdf":
+
+            self._write_data_to_vtk_vtk_format(
+                filename=filename,
+                datasets=datasets,
+                volume_normalization=volume_normalization,
+            )
+
+        elif Path(filename).suffix == ".vtk" or Path(filename).suffix == ".vtu":
+
+            self._write_data_to_vtk_hdf5_format(
+                filename=filename,
+                datasets=datasets,
+                volume_normalization=volume_normalization,
+            )
+
+        else:
+            raise ValueError(
+                "Unsupported file extension, The filename must end with "
+                "'.vtkhdf', '.vtu' or '.vtk'"
+            )
+
+    def _write_data_to_vtk_vtk_format(
+        self,
+        filename: PathLike | None = None,
+        datasets: dict | None = None,
+        volume_normalization: bool = True,
+    ):
         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"
@@ -2774,29 +2804,129 @@ def write_data_to_vtk(
 
         writer.Write()
 
+    def _write_data_to_vtk_hdf_format(
+        self,
+        filename: PathLike | None = None,
+        datasets: dict | None = None,
+        volume_normalization: bool = True,
+    ):
+        def append_dataset(dset, array):
+            """Convenience function to append data to an HDF5 dataset"""
+            origLen = dset.shape[0]
+            dset.resize(origLen + array.shape[0], axis=0)
+            dset[origLen:] = array
+
+        if self.library != "moab":
+            raise NotImplemented("VTKHDF output is only supported for MOAB meshes")
+
+        # the self.connectivity contains an arrays of length 8, in the case of
+        # MOAB tetrahedra mesh elements, the last 4 values are -1 and can be removed
+        trimmed_connectivity = []
+        for cell in self.connectivity:
+            # Find the index of the first -1 value, if any
+            first_negative_index = np.where(cell == -1)[0]
+            if first_negative_index.size > 0:
+                # Slice the array up to the first -1 value
+                trimmed_connectivity.append(cell[: first_negative_index[0]])
+            else:
+                # No -1 values, append the whole cell
+                trimmed_connectivity.append(cell)
+        trimmed_connectivity = np.array(trimmed_connectivity, dtype="int32").flatten()
+
+        # MOAB meshes supports tet elements only so we know it has 4 points per cell
+        points_per_cell = 4
+
+        # offsets are the indices of the first point of each cell in the array of points
+        offsets = np.arange(0, self.n_elements * points_per_cell + 1, points_per_cell)
+
+        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}"
+                )
+
+        with h5py.File(filename, "w") as f:
+
+            root = f.create_group("VTKHDF")
+            vtk_file_format_version = (2, 1)
+            root.attrs["Version"] = vtk_file_format_version
+            ascii_type = "UnstructuredGrid".encode("ascii")
+            root.attrs.create(
+                "Type",
+                ascii_type,
+                dtype=h5py.string_dtype("ascii", len(ascii_type)),
+            )
+
+            # create hdf5 file structure
+            root.create_dataset("NumberOfPoints", (0,), maxshape=(None,), dtype="i8")
+            root.create_dataset("Types", (0,), maxshape=(None,), dtype="uint8")
+            root.create_dataset("Points", (0, 3), maxshape=(None, 3), dtype="f")
+            root.create_dataset(
+                "NumberOfConnectivityIds", (0,), maxshape=(None,), dtype="i8"
+            )
+            root.create_dataset("NumberOfCells", (0,), maxshape=(None,), dtype="i8")
+            root.create_dataset("Offsets", (0,), maxshape=(None,), dtype="i8")
+            root.create_dataset("Connectivity", (0,), maxshape=(None,), dtype="i8")
+
+            append_dataset(root["NumberOfPoints"], np.array([len(self.vertices)]))
+            append_dataset(root["Points"], self.vertices)
+            append_dataset(
+                root["NumberOfConnectivityIds"],
+                np.array([len(trimmed_connectivity)]),
+            )
+            append_dataset(root["Connectivity"], trimmed_connectivity)
+            append_dataset(root["NumberOfCells"], np.array([self.n_elements]))
+            append_dataset(root["Offsets"], offsets)
+
+            # VTK_TETRA type is known as DAGMC only supports tet meshes
+            element_type = 10  # VTK_TETRA
+            append_dataset(
+                root["Types"], np.full(self.n_elements, element_type, dtype="uint8")
+            )
+
+            cell_data_group = root.create_group("CellData")
+
+            for name, data in datasets.items():
+
+                cell_data_group.create_dataset(
+                    name, (0,), maxshape=(None,), dtype="float64", chunks=True
+                )
+
+                if volume_normalization:
+                    data /= self.volumes
+                append_dataset(cell_data_group[name], data)
+
     @classmethod
     def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
-        filename = group['filename'][()].decode()
-        library = group['library'][()].decode()
-        if 'options' in group.attrs:
-            options = group.attrs['options'].decode()
+        filename = group["filename"][()].decode()
+        library = group["library"][()].decode()
+        if "options" in group.attrs:
+            options = group.attrs["options"].decode()
         else:
             options = None
 
-        mesh = cls(filename=filename, library=library, mesh_id=mesh_id, name=name, options=options)
+        mesh = cls(
+            filename=filename,
+            library=library,
+            mesh_id=mesh_id,
+            name=name,
+            options=options,
+        )
         mesh._has_statepoint_data = True
-        vol_data = group['volumes'][()]
+        vol_data = group["volumes"][()]
         mesh.volumes = np.reshape(vol_data, (vol_data.shape[0],))
         mesh.n_elements = mesh.volumes.size
 
-        vertices = group['vertices'][()]
+        vertices = group["vertices"][()]
         mesh._vertices = vertices.reshape((-1, 3))
-        connectivity = group['connectivity'][()]
+        connectivity = group["connectivity"][()]
         mesh._connectivity = connectivity.reshape((-1, 8))
-        mesh._element_types = group['element_types'][()]
+        mesh._element_types = group["element_types"][()]
 
-        if 'length_multiplier' in group:
-            mesh.length_multiplier = group['length_multiplier'][()]
+        if "length_multiplier" in group:
+            mesh.length_multiplier = group["length_multiplier"][()]
 
         return mesh
 
@@ -2815,7 +2945,7 @@ def to_xml_element(self):
 
         element.set("library", self._library)
         if self.options is not None:
-            element.set('options', self.options)
+            element.set("options", self.options)
         subelement = ET.SubElement(element, "filename")
         subelement.text = str(self.filename)
 

tests/regression_tests/unstructured_mesh/test_mesh_dagmc_tets.vtk

@@ -0,0 +1,159 @@
+# vtk DataFile Version 2.0
+made_with_cad_to_dagmc_package, Created by Gmsh 4.12.1 
+ASCII
+DATASET UNSTRUCTURED_GRID
+POINTS 14 double
+-0.5 -0.5 0.5
+-0.5 -0.5 -0.5
+-0.5 0.5 0.5
+-0.5 0.5 -0.5
+0.5 -0.5 0.5
+0.5 -0.5 -0.5
+0.5 0.5 0.5
+0.5 0.5 -0.5
+-0.5 0 0
+0.5 0 0
+0 -0.5 0
+0 0.5 0
+0 0 -0.5
+0 0 0.5
+
+CELLS 68 268
+1 0
+1 1
+1 2
+1 3
+1 4
+1 5
+1 6
+1 7
+2 1 0
+2 0 2
+2 3 2
+2 1 3
+2 5 4
+2 4 6
+2 7 6
+2 5 7
+2 1 5
+2 0 4
+2 3 7
+2 2 6
+3 1 0 8
+3 0 2 8
+3 3 1 8
+3 2 3 8
+3 5 9 4
+3 4 9 6
+3 7 9 5
+3 6 9 7
+3 0 1 10
+3 4 0 10
+3 1 5 10
+3 5 4 10
+3 2 11 3
+3 6 11 2
+3 3 11 7
+3 7 11 6
+3 1 3 12
+3 5 1 12
+3 3 7 12
+3 7 5 12
+3 0 13 2
+3 4 13 0
+3 2 13 6
+3 6 13 4
+4 13 8 12 10
+4 11 8 12 13
+4 12 11 13 9
+4 10 12 13 9
+4 12 3 11 7
+4 13 2 8 0
+4 8 12 1 3
+4 11 3 8 2
+4 10 8 1 0
+4 0 10 13 4
+4 1 12 10 5
+4 11 2 13 6
+4 4 9 13 6
+4 6 9 11 7
+4 10 9 4 5
+4 7 9 12 5
+4 3 8 12 11
+4 1 12 8 10
+4 8 11 2 13
+4 10 13 8 0
+4 13 10 9 4
+4 11 13 9 6
+4 12 11 9 7
+4 9 10 12 5
+
+CELL_TYPES 68
+1
+1
+1
+1
+1
+1
+1
+1
+3
+3
+3
+3
+3
+3
+3
+3
+3
+3
+3
+3
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+5
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10
+10