From ecd0455b0647fb9a872616fbb43195cee203a957 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Sat, 13 Aug 2022 12:32:14 +0100
Subject: [PATCH] added vtk writing to tracks object

---
 openmc/tracks.py                | 49 +++++++++++++++++++++++++++++++++
 tests/unit_tests/test_tracks.py | 14 ++++++++++
 2 files changed, 63 insertions(+)

diff --git a/openmc/tracks.py b/openmc/tracks.py
index e9097e8eeb..4b4ecc4f2b 100644
--- a/openmc/tracks.py
+++ b/openmc/tracks.py
@@ -266,6 +266,55 @@ class Tracks(list):
             track.plot(ax)
         return ax
 
+    def write_tracks_to_vtk(self, filename):
+        """Creates a VTK file of the tracks
+
+        Parameters
+        ----------
+        filename : str
+            Name of the VTK file to write.
+
+        Returns
+        -------
+        vtk.vtkStructuredGrid
+            the VTK object
+        """
+
+        import vtk
+
+        # Initialize data arrays and offset.
+        points = vtk.vtkPoints()
+        cells = vtk.vtkCellArray()
+    
+        for particle in self:
+            for state in particle.states:
+                points.InsertNextPoint(state['r'])
+
+            # Create VTK line and assign points to line.
+            n = particle.states.size
+            line = vtk.vtkPolyLine()
+            line.GetPointIds().SetNumberOfIds(n)
+            for i in range(n):
+                line.GetPointIds().SetId(i, point_offset + i)
+            point_offset += n
+
+            # Add line to cell array
+            cells.InsertNextCell(line)
+
+        data = vtk.vtkPolyData()
+        data.SetPoints(points)
+        data.SetLines(cells)
+
+        writer = vtk.vtkXMLPPolyDataWriter()
+        if vtk.vtkVersion.GetVTKMajorVersion() > 5:
+            writer.SetInputData(data)
+        else:
+            writer.SetInput(data)
+        writer.SetFileName(filename)
+        writer.Write()
+
+        return filename
+
     @staticmethod
     def combine(track_files, path='tracks.h5'):
         """Combine multiple track files into a single track file
diff --git a/tests/unit_tests/test_tracks.py b/tests/unit_tests/test_tracks.py
index e18015d81a..bb1210d2ea 100644
--- a/tests/unit_tests/test_tracks.py
+++ b/tests/unit_tests/test_tracks.py
@@ -141,3 +141,17 @@ def test_filter(sphere_model, run_in_tmpdir):
     assert matches == tracks
     matches = tracks.filter(particle='bunnytron')
     assert matches == []
+
+def test_write_tracks_to_vtk():
+    # Set maximum number of tracks per process to write
+    sphere_model.settings.max_tracks = 25
+    sphere_model.settings.photon_transport = True
+
+    # Run OpenMC to generate tracks.h5 file
+    generate_track_file(sphere_model, tracks=True)
+
+    tracks = openmc.Tracks('tracks.h5')
+    filename = tracks.write_tracks_to_vtk('tracks.vtk')
+
+    assert filename == 'tracks.vtk'
+    assert Path('tracks.vtk').is_file()