Updates to VTK data checks (#3371)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

docs/source/pythonapi/base.rst

@@ -147,17 +147,26 @@ Constructing Tallies
    openmc.ZernikeRadialFilter
    openmc.ParentNuclideFilter
    openmc.ParticleFilter
-   openmc.RegularMesh
-   openmc.RectilinearMesh
-   openmc.CylindricalMesh
-   openmc.SphericalMesh
-   openmc.UnstructuredMesh
    openmc.MeshMaterialVolumes
    openmc.Trigger
    openmc.TallyDerivative
    openmc.Tally
    openmc.Tallies
 
+Meshes
+------
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: myclassinherit.rst
+
+   openmc.RegularMesh
+   openmc.RectilinearMesh
+   openmc.CylindricalMesh
+   openmc.SphericalMesh
+   openmc.UnstructuredMesh
+
 Geometry Plotting
 -----------------
 

docs/source/usersguide/processing.rst

@@ -91,3 +91,82 @@ from a statepoint file, the ``openmc.statepoint`` module can be used. An
 `example notebook`_ demontrates how to analyze and plot source information.
 
 .. _example notebook: https://nbviewer.jupyter.org/github/openmc-dev/openmc-notebooks/blob/main/post-processing.ipynb
+
+------------------------
+VTK Mesh File Generation
+------------------------
+
+VTK files of OpenMC meshes can be created using the
+:meth:`openmc.Mesh.write_data_to_vtk`  method. Data can be applied to the
+elements of the resulting mesh from mesh filter objects. This data can be
+provided either as a flat array or, in the case of structured meshes
+(:class:`~openmc.RegularMesh`, :class:`~openmc.RectilinearMesh`,
+:class:`~openmc.CylindricalMesh`, or :class:`SphericalMesh`), the data can be
+shaped with dimensions that match the dimensions of the mesh itself.
+
+
+.. image:: ../_images/sphere-mesh-vtk.png
+   :width: 400px
+   :align: center
+   :alt: OpenMC spherical mesh exported to VTK
+
+
+For all mesh types, if a flat data array is provided to the mesh, it is expected
+that the data is ordered in the same ordering as the :attr:`openmc.Mesh.indices`
+for that mesh object. When providing data directly from a tally, as shown below,
+a flat array for a given dataset can be passed directly to this method.
+
+::
+
+    # create model above
+
+    # create a mesh tally
+    mesh = openmc.RegularMesh()
+    mesh.dimension = [10, 20, 30]
+    mesh.lower_left = [-5, -10, -15]
+    mesh.upper_right = [5, 10, 15]
+    mesh_filter = openmc.MeshFilter(mesh)
+    tally = openmc.Tally()
+    tally.filters = [mesh_filter]
+    tally.scores = ['flux']
+
+    model.tallies = [tally]
+    model.run(apply_tally_results=True)
+
+    # provide the data as-is to the method
+    mesh.write_data_to_vtk('flux.vtk', {'flux-mean': tally.mean})
+
+The :class:`~openmc.Tally` object also provides a way to expand the dimensions
+of the mesh filter into a meaningful form where indexing the mesh filter
+dimensions results in intuitive slicing of structured meshes by setting
+``expand_dims=True`` when using :meth:`openmc.Tally.get_reshaped_data`. This
+reshaping does cause flat indexing of the data to change, however. As noted
+above, provided datasets are allowed to be shaped so long as such datasets have
+shapes that match the mesh dimensions. The ability to pass datasets in this way
+is useful when additional filters are applied to a tally. The example below
+demonstrates such a case for tally with both a :class:`~openmc.MeshFilter` and
+:class:`~openmc.EnergyFilter` applied.
+
+::
+
+    # create model above
+
+    # create a mesh tally with energy filter
+    mesh = openmc.RegularMesh()
+    mesh.dimension = [10, 20, 30]
+    mesh.lower_left = [-5, -10, -15]
+    mesh.upper_right = [5, 10, 15]
+    mesh_filter = openmc.MeshFilter(mesh)
+    energy_filter = openmc.EnergyFilter([0.0, 1.0, 20.0e6])
+    tally = openmc.Tally()
+    tally.filters = [mesh_filter, energy_filter]
+    tally.scores = ['flux']
+
+    model.tallies = [tally]
+    model.run(apply_tally_results=True)
+
+    # get the data with mesh dimensions expanded, squeeze out length-one dimensions (nuclides, scores)
+    flux = tally.get_reshaped_data(expand_dims=True).squeeze() # shape: (10, 20, 30, 2)
+
+    # write the lowest energy group to a VTK file
+    mesh.write_data_to_vtk('flux-group1.vtk', datasets={'flux-mean': flux[..., 0]})

openmc/mesh.py

@@ -576,11 +576,16 @@ class StructuredMesh(MeshBase):
         filename : str
             Name of the VTK file to write.
         datasets : dict
-            Dictionary whose keys are the data labels
-            and values are the data sets.
+            Dictionary whose keys are the data labels and values are the data
+            sets. 1D datasets are expected to be extracted directly from
+            statepoint data without reordering/reshaping. Multidimensional
+            datasets are expected to have the same dimensions as the mesh itself
+            with structured indexing in "C" ordering. See the "expand_dims" flag
+            of :meth:`~openmc.Tally.get_reshaped_data` on reshaping tally data when using
+            :class:`~openmc.MeshFilter`'s.
         volume_normalization : bool, optional
-            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.
         curvilinear : bool
             Whether or not to write curvilinear elements. Only applies to
             ``SphericalMesh`` and ``CylindricalMesh``.
@@ -594,14 +599,27 @@ class StructuredMesh(MeshBase):
         -------
         vtk.StructuredGrid or vtk.UnstructuredGrid
             a VTK grid object representing the mesh
+
+        Examples
+        --------
+        1D data from a tally with only a mesh filter and heating score:
+
+            # pass the tally mean property of shape (N, 1, 1) directly to this
+            # method; dimensions of size 1 will automatically removed
+            >>> heating = tally.mean
+            >>> mesh.write_data_to_vtk({'heating': heating})
+
+        Multidimensional data from a tally with only a mesh
+
+           # retrieve a data array with the mesh filter expanded into three
+           # dimensions, ijk; additional dimensions of size one will
+           # automatically be removed
+           >>> heating = tally.get_reshaped_data(expand_dims=True)
+           >>> mesh.write_data_to_vtk({'heating': heating})
         """
         import vtk
         from vtk.util import numpy_support as nps
 
-        # check that the data sets are appropriately sized
-        if datasets is not None:
-            self._check_vtk_datasets(datasets)
-
         # write linear elements using a structured grid
         if not curvilinear or isinstance(self, (RegularMesh, RectilinearMesh)):
             vtk_grid = self._create_vtk_structured_grid()
@@ -612,22 +630,27 @@ class StructuredMesh(MeshBase):
             writer = vtk.vtkUnstructuredGridWriter()
 
         if datasets is not None:
-            # maintain a list of the datasets as added
-            # to the VTK arrays to ensure they persist
-            # in memory until the file is written
+            # maintain a list of the datasets as added to the VTK arrays to
+            # ensure they persist in memory until the file is written
             datasets_out = []
             for label, dataset in datasets.items():
-                dataset = np.asarray(dataset).flatten()
+                dataset = self._reshape_vtk_dataset(dataset)
+                self._check_vtk_dataset(label, dataset)
+                # If the array data is 3D, assume is in C ordering and transpose
+                # before flattening to match the ordering expected by the VTK
+                # array based on the way mesh indices are ordered in the Python
+                # API
+                # TODO: update to "C" ordering throughout
+                if dataset.ndim == 3:
+                    dataset = dataset.T.ravel()
                 datasets_out.append(dataset)
 
                 if volume_normalization:
-                    dataset /= self.volumes.T.flatten()
+                    dataset /= self.volumes.T.ravel()
 
                 dataset_array = vtk.vtkDoubleArray()
                 dataset_array.SetName(label)
-                dataset_array.SetArray(nps.numpy_to_vtk(dataset),
-                                    dataset.size,
-                                    True)
+                dataset_array.SetArray(nps.numpy_to_vtk(dataset), dataset.size, True)
                 vtk_grid.GetCellData().AddArray(dataset_array)
 
         writer.SetFileName(str(filename))
@@ -754,28 +777,69 @@ class StructuredMesh(MeshBase):
 
         return vtk_grid
 
-    def _check_vtk_datasets(self, datasets: dict):
-        """Perform some basic checks that the datasets are valid for this mesh
+    @staticmethod
+    def _reshape_vtk_dataset(dataset):
+        """Reshape a dataset to be compatible with VTK output
+
+        This method performs the following operations on a dataset:
+        1. Convert to numpy array if not already
+        2. Remove any trailing dimensions of size 1
+        3. Squeeze out any extra dimensions of size 1 beyond the first 3
 
         Parameters
         ----------
-        datasets : dict
-            Dictionary whose keys are the data labels
-            and values are the data sets.
+        dataset : array-like
+            The dataset to reshape
+
+        Returns
+        -------
+        numpy.ndarray
+            The reshaped dataset
+        """
+        reshaped_data = np.asarray(dataset)
+
+        # detect flat array with extra dims
+        if all(d == 1 for d in reshaped_data.shape[1:]):
+            reshaped_data = reshaped_data.squeeze()
+
+        # remove any higher dimensions with size 1
+        if reshaped_data.ndim > 3 and all(d == 1 for d in reshaped_data.shape[3:]):
+            reshaped_data = reshaped_data.reshape(reshaped_data.shape[:3])
+
+        if np.shares_memory(reshaped_data, dataset):
+            return np.copy(reshaped_data)
+        else:
+            return reshaped_data
+
+    def _check_vtk_dataset(self, label: str, dataset: np.ndarray):
+        """Perform some basic checks that a dataset is valid for this Mesh
+
+        Parameters
+        ----------
+        label : str
+            The label for the dataset being checked
+        dataset : numpy.ndarray
+            The dataset array to check against this mesh's dimensions
 
         """
-        for label, dataset in datasets.items():
-            errmsg = (
+        cv.check_type('data label', label, str)
+
+        if dataset.size != self.num_mesh_cells:
+            raise ValueError(
                 f"The size of the dataset '{label}' ({dataset.size}) should be"
                 f" equal to the number of mesh cells ({self.num_mesh_cells})"
             )
-            if isinstance(dataset, np.ndarray):
-                if not dataset.size == self.num_mesh_cells:
-                    raise ValueError(errmsg)
-            else:
-                if len(dataset) == self.num_mesh_cells:
-                    raise ValueError(errmsg)
-            cv.check_type('data label', label, str)
+
+        # accept a flat array as-is, assuming it is in the correct order
+        if dataset.ndim == 1:
+            return
+
+        if dataset.shape != self.dimension:
+            raise ValueError(
+                f'Cannot apply multidimensional dataset "{label}" with '
+                f"shape {dataset.shape} to mesh {self.id} "
+                f"with dimensions {self.dimension}"
+            )
 
 
 class RegularMesh(StructuredMesh):

tests/unit_tests/mesh_to_vtk/test.py

@@ -49,7 +49,10 @@ sphere_mesh = openmc.SphericalMesh(
 
 def mesh_data(mesh_dims):
     data = 100 * np.arange(np.prod(mesh_dims), dtype=float)
-    return data.reshape(*mesh_dims)
+    # data is returned reshaped with order 'F' to ensure that
+    # the resulting data is interpreted correctly by the
+    # write_data_to_vtk method
+    return data.reshape(*mesh_dims, order='F')
 
 test_data = ((reg_mesh, False, 'regular'),
              (rect_mesh, False, 'rectilinear'),
@@ -83,7 +86,6 @@ def test_mesh_write_vtk(mesh_params, run_in_tmpdir):
 
 # check data writing
 def test_mesh_write_vtk_data(run_in_tmpdir):
-
     data = {'ascending_data': mesh_data(cyl_mesh.dimension)}
     filename_expected = full_path('cyl-data.vtk')
     filename_actual = full_path('cyl-data-actual.vtk')

tests/unit_tests/mesh_to_vtk/test_vtk_dims.py

@@ -90,7 +90,7 @@ def test_write_data_to_vtk(mesh, tmpdir):
     # kji (i changing fastest) orering is expected for input data
     # by using the volumes transposed as the data here, we can ensure the
     # normalization is happening correctly
-    data = mesh.volumes.T
+    data = mesh.volumes
 
     # RUN
     mesh.write_data_to_vtk(filename=filename, datasets={"label1": data, "label2": data})

tests/unit_tests/test_mesh.py

@@ -480,6 +480,9 @@ def test_umesh(run_in_tmpdir, simple_umesh, export_type):
     mean = np.array([arr.GetTuple1(i) for i in range(ref_data.size)])
     np.testing.assert_almost_equal(mean, ref_data)
 
+    # attempt to apply a dataset with an improper size to a VTK write
+    with pytest.raises(ValueError, match='Cannot apply dataset "mean"') as e:
+        simple_umesh.write_data_to_vtk(datasets={'mean': ref_data[:-2]}, filename=filename)
 
 def test_mesh_get_homogenized_materials():
     """Test the get_homogenized_materials method"""