Expand Particle::tracks_ to include full state

include/openmc/particle_data.h

@@ -56,6 +56,18 @@ struct SourceSite {
   int64_t progeny_id;
 };
 
+//! State of a particle used for particle track files
+struct TrackState {
+  Position r;
+  Direction u;
+  double E;
+  double time {0.0};
+  double wgt {1.0};
+  int cell_id;
+  int material_id {-1};
+  ParticleType particle;
+};
+
 //! Saved ("banked") state of a particle, for nu-fission tallying
 struct NuBank {
   double E;          //!< particle energy
@@ -290,7 +302,7 @@ private:
 
   vector<FilterMatch> filter_matches_; // tally filter matches
 
-  vector<vector<Position>> tracks_; // tracks for outputting to file
+  vector<vector<TrackState>> tracks_; // tracks for outputting to file
 
   vector<NuBank> nu_bank_; // bank of most recently fissioned particles
 
@@ -342,6 +354,9 @@ public:
   const LocalCoord& coord(int i) const { return coord_[i]; }
   const vector<LocalCoord>& coord() const { return coord_; }
 
+  LocalCoord& lowest_coord() { return coord_[n_coord_ - 1]; }
+  const LocalCoord& lowest_coord() const { return coord_[n_coord_ - 1]; }
+
   int& n_coord_last() { return n_coord_last_; }
   const int& n_coord_last() const { return n_coord_last_; }
   int& cell_last(int i) { return cell_last_[i]; }
@@ -357,6 +372,7 @@ public:
   const int& g_last() const { return g_last_; }
 
   double& wgt() { return wgt_; }
+  double wgt() const { return wgt_; }
   double& mu() { return mu_; }
   const double& mu() const { return mu_; }
   double& time() { return time_; }
@@ -479,6 +495,9 @@ public:
     n_coord_ = 1;
   }
 
+  //! Get track information based on particle's current state
+  TrackState get_track_state() const;
+
   void zero_delayed_bank()
   {
     for (int& n : n_delayed_bank_) {

src/particle_data.cpp

@@ -1,6 +1,8 @@
 #include "openmc/particle_data.h"
 
+#include "openmc/cell.h"
 #include "openmc/geometry.h"
+#include "openmc/material.h"
 #include "openmc/nuclide.h"
 #include "openmc/photon.h"
 #include "openmc/settings.h"
@@ -53,4 +55,20 @@ ParticleData::ParticleData()
   photon_xs_.resize(data::elements.size());
 }
 
+TrackState ParticleData::get_track_state() const
+{
+  TrackState state;
+  state.r = this->r();
+  state.u = this->u();
+  state.E = this->E();
+  state.time = this->time();
+  state.wgt = this->wgt();
+  state.cell_id = model::cells[this->lowest_coord().cell]->id_;
+  if (this->material() != MATERIAL_VOID) {
+    state.material_id = model::materials[material()]->id_;
+  }
+  state.particle = this->type();
+  return state;
+}
+
 } // namespace openmc

src/track_output.cpp

@@ -30,7 +30,7 @@ void add_particle_track(Particle& p)
 
 void write_particle_track(Particle& p)
 {
-  p.tracks().back().push_back(p.r());
+  p.tracks().back().push_back(p.get_track_state());
 }
 
 void finalize_particle_track(Particle& p)
@@ -57,9 +57,9 @@ void finalize_particle_track(Particle& p)
       size_t n = t.size();
       xt::xtensor<double, 2> data({n, 3});
       for (int j = 0; j < n; ++j) {
-        data(j, 0) = t[j].x;
-        data(j, 1) = t[j].y;
-        data(j, 2) = t[j].z;
+        data(j, 0) = t[j].r.x;
+        data(j, 1) = t[j].r.y;
+        data(j, 2) = t[j].r.z;
       }
       std::string name = fmt::format("coordinates_{}", i);
       write_dataset(file_id, name.c_str(), data);