From b32d4b5ed3b82b25d5f86cb22b0aaf079ae35f10 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Fri, 13 May 2022 07:18:34 -0500
Subject: [PATCH] Write particle types separately for track files

---
 include/openmc/particle_data.h |  7 ++++++-
 src/particle_data.cpp          |  1 -
 src/track_output.cpp           | 12 +++++++-----
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/include/openmc/particle_data.h b/include/openmc/particle_data.h
index 8c1f4deaba..3e6c8515cb 100644
--- a/include/openmc/particle_data.h
+++ b/include/openmc/particle_data.h
@@ -65,7 +65,12 @@ struct TrackState {
   double wgt {1.0};
   int cell_id;
   int material_id {-1};
+};
+
+//! Full history of a single particle's track states
+struct TrackStateHistory {
   ParticleType particle;
+  std::vector<TrackState> states;
 };
 
 //! Saved ("banked") state of a particle, for nu-fission tallying
@@ -302,7 +307,7 @@ private:
 
   vector<FilterMatch> filter_matches_; // tally filter matches
 
-  vector<vector<TrackState>> tracks_; // tracks for outputting to file
+  vector<TrackStateHistory> tracks_; // tracks for outputting to file
 
   vector<NuBank> nu_bank_; // bank of most recently fissioned particles
 
diff --git a/src/particle_data.cpp b/src/particle_data.cpp
index be34fbc9b5..4206d9cd94 100644
--- a/src/particle_data.cpp
+++ b/src/particle_data.cpp
@@ -67,7 +67,6 @@ TrackState ParticleData::get_track_state() const
   if (this->material() != MATERIAL_VOID) {
     state.material_id = model::materials[material()]->id_;
   }
-  state.particle = this->type();
   return state;
 }
 
diff --git a/src/track_output.cpp b/src/track_output.cpp
index 3d04a73bab..5a7aff834f 100644
--- a/src/track_output.cpp
+++ b/src/track_output.cpp
@@ -26,11 +26,12 @@ namespace openmc {
 void add_particle_track(Particle& p)
 {
   p.tracks().emplace_back();
+  p.tracks().back().particle = p.type();
 }
 
 void write_particle_track(Particle& p)
 {
-  p.tracks().back().push_back(p.get_track_state());
+  p.tracks().back().states.push_back(p.get_track_state());
 }
 
 hid_t trackstate_type()
@@ -51,8 +52,6 @@ hid_t trackstate_type()
   H5Tinsert(tracktype, "cell_id", HOFFSET(TrackState, cell_id), H5T_NATIVE_INT);
   H5Tinsert(
     tracktype, "material_id", HOFFSET(TrackState, material_id), H5T_NATIVE_INT);
-  H5Tinsert(
-    tracktype, "particle", HOFFSET(TrackState, particle), H5T_NATIVE_INT);
 
   H5Tclose(postype);
   return tracktype;
@@ -66,12 +65,14 @@ void finalize_particle_track(Particle& p)
 
   // Determine number of coordinates for each particle
   vector<int> offsets;
+  vector<int> particles;
   vector<TrackState> tracks;
   int offset = 0;
   for (auto& track_i : p.tracks()) {
     offsets.push_back(offset);
-    offset += track_i.size();
-    tracks.insert(tracks.end(), track_i.begin(), track_i.end());
+    particles.push_back(static_cast<int>(track_i.particle));
+    offset += track_i.states.size();
+    tracks.insert(tracks.end(), track_i.states.begin(), track_i.states.end());
   }
   offsets.push_back(offset);
 
@@ -82,6 +83,7 @@ void finalize_particle_track(Particle& p)
     write_attribute(file_id, "version", VERSION_TRACK);
     write_attribute(file_id, "n_particles", p.tracks().size());
     write_attribute(file_id, "offsets", offsets);
+    write_attribute(file_id, "particles", particles);
 
     // Create HDF5 datatype for TrackState
     hid_t track_type = trackstate_type();