Support track file writing for particle restart runs. (#2957)

Co-authored-by: Paul Romano <[email protected]>

src/particle_restart.cpp

@@ -87,8 +87,10 @@ void run_particle_restart()
   read_particle_restart(p, previous_run_mode);
 
   // write track if that was requested on command line
-  if (settings::write_all_tracks)
+  if (settings::write_all_tracks) {
+    open_track_file();
     p.write_track() = true;
+  }
 
   // Set all tallies to 0 for now (just tracking errors)
   model::tallies.clear();
@@ -123,6 +125,10 @@ void run_particle_restart()
 
   // Write output if particle made it
   print_particle(p);
+
+  if (settings::write_all_tracks) {
+    close_track_file();
+  }
 }
 
 } // namespace openmc

tests/unit_tests/test_tracks.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import h5py
 import numpy as np
 import openmc
 import pytest
@@ -157,3 +158,35 @@ def test_write_to_vtk(sphere_model):
 
     assert isinstance(polydata, vtk.vtkPolyData)
     assert Path('tracks.vtp').is_file()
+
+
+def test_restart_track(run_in_tmpdir, sphere_model):
+    # cut the sphere model in half with an improper boundary condition
+    plane = openmc.XPlane(x0=-1.0)
+    for cell in sphere_model.geometry.get_all_cells().values():
+        cell.region &= +plane
+
+    # generate lost particle files
+    with pytest.raises(RuntimeError, match='Maximum number of lost particles has been reached.'):
+        sphere_model.run(output=False)
+
+    lost_particle_files = list(Path.cwd().glob('particle_*.h5'))
+    assert len(lost_particle_files) > 0
+    particle_file = lost_particle_files[0]
+     # restart the lost particle with tracks enabled
+    sphere_model.run(tracks=True, restart_file=particle_file)
+    tracks_file = Path('tracks.h5')
+    assert tracks_file.is_file()
+
+    # check that the last track of the file matches the lost particle file
+    tracks = openmc.Tracks(tracks_file)
+    initial_state = tracks[0].particle_tracks[0].states[0]
+    restart_r = np.array(initial_state['r'])
+    restart_u = np.array(initial_state['u'])
+
+    with h5py.File(particle_file, 'r') as lost_particle_file:
+        lost_r = np.array(lost_particle_file['xyz'][()])
+        lost_u = np.array(lost_particle_file['uvw'][()])
+
+    pytest.approx(restart_r, lost_r)
+    pytest.approx(restart_u, lost_u)