From 5ecac76b6960f635f375fd48f3cb567733073d3f Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 29 Jun 2022 12:34:05 +0200
Subject: [PATCH] write_data_to_vtk rectilinear mesh

---
 openmc/mesh.py | 41 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 7 deletions(-)

diff --git a/openmc/mesh.py b/openmc/mesh.py
index 179e9a1713..db2e1593c5 100644
--- a/openmc/mesh.py
+++ b/openmc/mesh.py
@@ -895,7 +895,7 @@ class RectilinearMesh(StructuredMesh):
 
         return element
 
-    def vtk_grid(self, filename=None):
+    def write_data_to_vtk(self, filename, datasets, volume_normalization=True):
         """Creates a VTK object of the mesh
 
         Args:
@@ -908,6 +908,19 @@ class RectilinearMesh(StructuredMesh):
         import vtk
         from vtk.util import numpy_support as nps
 
+        if self.volumes is None and volume_normalization:
+            raise RuntimeError("No volume data is present on this "
+                               "unstructured mesh. Please load the "
+                               " mesh information from a statepoint file.")
+
+        # check that the data sets are appropriately sized
+        for label, dataset in datasets.items():
+            if isinstance(dataset, np.ndarray):
+                assert dataset.size == self.dimension[0] * self.dimension[1]* self.dimension[2]
+            else:
+                assert len(dataset) == self.dimension[0] * self.dimension[1]* self.dimension[2]
+            cv.check_type('label', label, str)
+
         x_vals = self.x_grid
         y_vals = self.y_grid
         z_vals = self.z_grid
@@ -922,12 +935,26 @@ class RectilinearMesh(StructuredMesh):
         vtkPts.SetData(nps.numpy_to_vtk(pts_cartesian, deep=True))
         vtk_grid.SetPoints(vtkPts)
 
-        if filename:
-            # write the .vtk file
-            writer = vtk.vtkStructuredGridWriter()
-            writer.SetFileName(filename)
-            writer.SetInputData(vtk_grid)
-            writer.Write()
+        # 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(filename)
+        writer.SetInputData(vtk_grid)
+        writer.Write()
 
         return vtk_grid