From 7be28ef156a96fb4c525130d9c52e6966cf43539 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 29 Jun 2022 11:08:58 +0200
Subject: [PATCH] added vtk_grid to meshes

---
 openmc/mesh.py | 99 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/openmc/mesh.py b/openmc/mesh.py
index 7f7b07879..dbdfb10b0 100644
--- a/openmc/mesh.py
+++ b/openmc/mesh.py
@@ -621,7 +621,45 @@ class RegularMesh(StructuredMesh):
         root_cell.fill = lattice
 
         return root_cell, cells
+    
+    def vtk_grid(self, filename=None):
+        import vtk
+        from vtk.util import numpy_support as nps
 
+        x_vals = np.linspace(
+            self.lower_left[0],
+            self.upper_right[0],
+            num=self.dimension[0] + 1,
+        )
+        y_vals = np.linspace(
+            self.lower_left[1],
+            self.upper_right[1],
+            num=self.dimension[1] + 1,
+        )
+        z_vals = np.linspace(
+            self.lower_left[2],
+            self.upper_right[2],
+            num=self.dimension[2] + 1,
+        )
+        vtk_grid = vtk.vtkStructuredGrid()
+
+        vtk_grid.SetDimensions(len(x_vals), len(y_vals), len(z_vals))
+
+        # create points
+        pts_cartesian = np.array([[x, y, z] for z in z_vals for y in y_vals for x in x_vals])
+
+        vtkPts = vtk.vtkPoints()
+        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()
+
+        return vtk_grid
 
 def Mesh(*args, **kwargs):
     warnings.warn("Mesh has been renamed RegularMesh. Future versions of "
@@ -821,6 +859,33 @@ class RectilinearMesh(StructuredMesh):
 
         return element
 
+    def vtk_grid(self, filename=None):
+        import vtk
+        from vtk.util import numpy_support as nps
+
+        x_vals = self.x_grid
+        y_vals = self.y_grid
+        z_vals = self.z_grid
+        vtk_grid = vtk.vtkStructuredGrid()
+
+        vtk_grid.SetDimensions(len(x_vals), len(y_vals), len(z_vals))
+
+        # create points
+        pts_cartesian = np.array([[x, y, z] for z in z_vals for y in y_vals for x in x_vals])
+
+        vtkPts = vtk.vtkPoints()
+        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()
+
+        return vtk_grid
+
 
 class CylindricalMesh(StructuredMesh):
     """A 3D cylindrical mesh
@@ -1012,6 +1077,37 @@ class CylindricalMesh(StructuredMesh):
 
         return np.multiply.outer(np.outer(V_r, V_p), V_z)
 
+    def vtk_grid(self, filename=None):
+        import vtk
+        from vtk.util import numpy_support as nps
+
+        r_vals = self.r_grid
+        phi_vals = self.phi_grid
+        z_vals = self.z_grid
+        vtk_grid = vtk.vtkStructuredGrid()
+
+        vtk_grid.SetDimensions(len(r_vals), len(phi_vals), len(z_vals))
+
+        # create points
+        pts_cylindrical = np.array([[r, phi, z] for z in z_vals for phi in phi_vals for r in r_vals])
+        pts_cartesian = np.copy(pts_cylindrical)
+        r, phi, z = pts_cylindrical[:, 0], pts_cylindrical[:, 1], pts_cylindrical[:, 2]
+        pts_cartesian[:, 0] = r * np.cos(phi)
+        pts_cartesian[:, 1] = r * np.sin(phi)
+        pts_cartesian[:, 2] = z
+
+        vtkPts = vtk.vtkPoints()
+        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()
+
+        return vtk_grid
 
 class SphericalMesh(StructuredMesh):
     """A 3D spherical mesh
@@ -1204,6 +1300,9 @@ class SphericalMesh(StructuredMesh):
 
         return np.multiply.outer(np.outer(V_r, V_t), V_p)
 
+    def vtk_grid(self, filename=None):
+        raise NotImplementedError("vtk_grid not implemented for SphericalMesh")
+
 
 class UnstructuredMesh(MeshBase):
     """A 3D unstructured mesh