From ecd0455b0647fb9a872616fbb43195cee203a957 Mon Sep 17 00:00:00 2001 From: shimwell Date: Sat, 13 Aug 2022 12:32:14 +0100 Subject: [PATCH 1/6] 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 e9097e8ee..4b4ecc4f2 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 e18015d81..bb1210d2e 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() From d15c051263510e6c81042228c35786921a1ddba0 Mon Sep 17 00:00:00 2001 From: shimwell Date: Sat, 13 Aug 2022 12:34:17 +0100 Subject: [PATCH 2/6] added missing arg --- tests/unit_tests/test_tracks.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_tracks.py b/tests/unit_tests/test_tracks.py index bb1210d2e..d2548be64 100644 --- a/tests/unit_tests/test_tracks.py +++ b/tests/unit_tests/test_tracks.py @@ -142,7 +142,8 @@ def test_filter(sphere_model, run_in_tmpdir): matches = tracks.filter(particle='bunnytron') assert matches == [] -def test_write_tracks_to_vtk(): + +def test_write_tracks_to_vtk(sphere_model): # Set maximum number of tracks per process to write sphere_model.settings.max_tracks = 25 sphere_model.settings.photon_transport = True From 2263549b9b59f9b49929c417c3e8ae1111d8bef4 Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Mon, 15 Aug 2022 13:22:44 +0100 Subject: [PATCH 3/6] producing vtp vtk files --- openmc/tracks.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/openmc/tracks.py b/openmc/tracks.py index 4b4ecc4f2..076ffc78d 100644 --- a/openmc/tracks.py +++ b/openmc/tracks.py @@ -266,7 +266,7 @@ class Tracks(list): track.plot(ax) return ax - def write_tracks_to_vtk(self, filename): + def write_tracks_to_vtk(self, filename='tracks.vtk'): """Creates a VTK file of the tracks Parameters @@ -286,12 +286,14 @@ class Tracks(list): points = vtk.vtkPoints() cells = vtk.vtkCellArray() + point_offset = 0 for particle in self: - for state in particle.states: - points.InsertNextPoint(state['r']) + for pt in particle.particle_tracks: + for state in pt.states: + points.InsertNextPoint(state['r']) # Create VTK line and assign points to line. - n = particle.states.size + n = pt.states.size line = vtk.vtkPolyLine() line.GetPointIds().SetNumberOfIds(n) for i in range(n): From 21b874f0be55d61f434397c7ff745e9fb2f97999 Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Mon, 15 Aug 2022 13:39:47 +0100 Subject: [PATCH 4/6] review suggestions from @paulromano --- openmc/tracks.py | 16 +++++++------- scripts/openmc-track-to-vtk | 42 +++++-------------------------------- 2 files changed, 13 insertions(+), 45 deletions(-) diff --git a/openmc/tracks.py b/openmc/tracks.py index 076ffc78d..55a5fb060 100644 --- a/openmc/tracks.py +++ b/openmc/tracks.py @@ -6,7 +6,7 @@ import h5py from .checkvalue import check_filetype_version from .source import SourceParticle, ParticleType - +from pathlib import Path ParticleTrack = namedtuple('ParticleTrack', ['particle', 'states']) ParticleTrack.__doc__ = """\ Particle track information @@ -266,18 +266,18 @@ class Tracks(list): track.plot(ax) return ax - def write_tracks_to_vtk(self, filename='tracks.vtk'): - """Creates a VTK file of the tracks + def write_tracks_to_vtk(self, filename=Path('tracks.vtk')): + """Creates a VTP file of the tracks Parameters ---------- - filename : str + filename : path-like Name of the VTK file to write. Returns ------- - vtk.vtkStructuredGrid - the VTK object + vtk.vtkPolyData + the VTK vtkPolyData object produced """ import vtk @@ -312,10 +312,10 @@ class Tracks(list): writer.SetInputData(data) else: writer.SetInput(data) - writer.SetFileName(filename) + writer.SetFileName(str(filename)) # SetFileName requires a string writer.Write() - return filename + return data @staticmethod def combine(track_files, path='tracks.h5'): diff --git a/scripts/openmc-track-to-vtk b/scripts/openmc-track-to-vtk index 3b3507974..118e65827 100755 --- a/scripts/openmc-track-to-vtk +++ b/scripts/openmc-track-to-vtk @@ -14,8 +14,8 @@ def _parse_args(): # Create argument parser. parser = argparse.ArgumentParser( description='Convert particle track file(s) to a .pvtp file.') - parser.add_argument('input', metavar='IN', type=str, nargs='+', - help='Input particle track data filename(s).') + parser.add_argument('input', metavar='IN', type=str, + help='Input particle track data filename.') parser.add_argument('-o', '--out', metavar='OUT', type=str, dest='out', help='Output VTK poly data filename.') @@ -33,41 +33,9 @@ def main(): elif not args.out.endswith('.pvtp'): args.out += '.pvtp' - # Initialize data arrays and offset. - points = vtk.vtkPoints() - cells = vtk.vtkCellArray() - point_offset = 0 - for fname in args.input: - # Write coordinate values to points array. - track_file = openmc.Tracks(fname) - for track in track_file: - for particle in track: - 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(args.out) - writer.Write() - + # Write coordinate values to points array. + track_file = openmc.Tracks(args.input) + track_file.write_tracks_to_vtk(args.out) if __name__ == '__main__': main() From 9847d053f35c62c8c5e73d702e938c6a2f9f8db0 Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Mon, 15 Aug 2022 13:40:56 +0100 Subject: [PATCH 5/6] changed vtk file to vtp --- openmc/tracks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openmc/tracks.py b/openmc/tracks.py index 55a5fb060..213a0867e 100644 --- a/openmc/tracks.py +++ b/openmc/tracks.py @@ -272,7 +272,7 @@ class Tracks(list): Parameters ---------- filename : path-like - Name of the VTK file to write. + Name of the VTP file to write. Returns ------- From 3a2d4a9138f95ec8c47f705de7d782c50137459e Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Mon, 15 Aug 2022 16:05:02 +0100 Subject: [PATCH 6/6] suggestions from code review by @paulromano Co-authored-by: Paul Romano --- openmc/tracks.py | 3 ++- tests/unit_tests/test_tracks.py | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/openmc/tracks.py b/openmc/tracks.py index 213a0867e..671aa65cf 100644 --- a/openmc/tracks.py +++ b/openmc/tracks.py @@ -7,6 +7,7 @@ from .checkvalue import check_filetype_version from .source import SourceParticle, ParticleType from pathlib import Path + ParticleTrack = namedtuple('ParticleTrack', ['particle', 'states']) ParticleTrack.__doc__ = """\ Particle track information @@ -266,7 +267,7 @@ class Tracks(list): track.plot(ax) return ax - def write_tracks_to_vtk(self, filename=Path('tracks.vtk')): + def write_tracks_to_vtk(self, filename=Path('tracks.vtp')): """Creates a VTP file of the tracks Parameters diff --git a/tests/unit_tests/test_tracks.py b/tests/unit_tests/test_tracks.py index d2548be64..5f3940fad 100644 --- a/tests/unit_tests/test_tracks.py +++ b/tests/unit_tests/test_tracks.py @@ -144,6 +144,7 @@ def test_filter(sphere_model, run_in_tmpdir): def test_write_tracks_to_vtk(sphere_model): + vtk = pytest.importorskip('vtk') # Set maximum number of tracks per process to write sphere_model.settings.max_tracks = 25 sphere_model.settings.photon_transport = True @@ -152,7 +153,7 @@ def test_write_tracks_to_vtk(sphere_model): generate_track_file(sphere_model, tracks=True) tracks = openmc.Tracks('tracks.h5') - filename = tracks.write_tracks_to_vtk('tracks.vtk') + polydata = tracks.write_tracks_to_vtk('tracks.vtp') - assert filename == 'tracks.vtk' - assert Path('tracks.vtk').is_file() + assert isinstance(polydata, vtk.vtkPolyData) + assert Path('tracks.vtp').is_file()