Addition of a collision tracking feature (#3417)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

include/openmc/bank.h

@@ -20,6 +20,8 @@ extern vector<SourceSite> source_bank;
 
 extern SharedArray<SourceSite> surf_source_bank;
 
+extern SharedArray<CollisionTrackSite> collision_track_bank;
+
 extern SharedArray<SourceSite> fission_bank;
 
 extern vector<vector<int>> ifp_source_delayed_group_bank;

include/openmc/bank_io.h

@@ -0,0 +1,103 @@
+#ifndef OPENMC_BANK_IO_H
+#define OPENMC_BANK_IO_H
+
+#include "hdf5.h"
+
+#include "openmc/message_passing.h"
+#include "openmc/span.h"
+#include "openmc/vector.h"
+
+#include <algorithm>
+
+#ifdef OPENMC_MPI
+#include <mpi.h>
+#endif
+
+namespace openmc {
+
+template<typename SiteType>
+void write_bank_dataset(const char* dataset_name, hid_t group_id,
+  span<SiteType> bank, const vector<int64_t>& bank_index, hid_t banktype
+#ifdef OPENMC_MPI
+  ,
+  MPI_Datatype mpi_dtype
+#endif
+)
+{
+  int64_t dims_size = bank_index.back();
+  int64_t count_size = bank_index[mpi::rank + 1] - bank_index[mpi::rank];
+
+#ifdef PHDF5
+  hsize_t dims[] {static_cast<hsize_t>(dims_size)};
+  hid_t dspace = H5Screate_simple(1, dims, nullptr);
+  hid_t dset = H5Dcreate(group_id, dataset_name, banktype, dspace, H5P_DEFAULT,
+    H5P_DEFAULT, H5P_DEFAULT);
+
+  hsize_t count[] {static_cast<hsize_t>(count_size)};
+  hid_t memspace = H5Screate_simple(1, count, nullptr);
+
+  hsize_t start[] {static_cast<hsize_t>(bank_index[mpi::rank])};
+  H5Sselect_hyperslab(dspace, H5S_SELECT_SET, start, nullptr, count, nullptr);
+
+  hid_t plist = H5Pcreate(H5P_DATASET_XFER);
+  H5Pset_dxpl_mpio(plist, H5FD_MPIO_COLLECTIVE);
+
+  H5Dwrite(dset, banktype, memspace, dspace, plist, bank.data());
+
+  H5Sclose(dspace);
+  H5Sclose(memspace);
+  H5Dclose(dset);
+  H5Pclose(plist);
+#else
+  if (mpi::master) {
+    hsize_t dims[] {static_cast<hsize_t>(dims_size)};
+    hid_t dspace = H5Screate_simple(1, dims, nullptr);
+    hid_t dset = H5Dcreate(group_id, dataset_name, banktype, dspace,
+      H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+
+#ifdef OPENMC_MPI
+    vector<SiteType> temp_bank {bank.begin(), bank.end()};
+#endif
+
+    for (int i = 0; i < mpi::n_procs; ++i) {
+      hsize_t count[] {static_cast<hsize_t>(bank_index[i + 1] - bank_index[i])};
+      hid_t memspace = H5Screate_simple(1, count, nullptr);
+
+#ifdef OPENMC_MPI
+      if (i > 0) {
+        MPI_Recv(bank.data(), count[0], mpi_dtype, i, i, mpi::intracomm,
+          MPI_STATUS_IGNORE);
+      }
+#endif
+
+      hid_t dspace_rank = H5Dget_space(dset);
+      hsize_t start[] {static_cast<hsize_t>(bank_index[i])};
+      H5Sselect_hyperslab(
+        dspace_rank, H5S_SELECT_SET, start, nullptr, count, nullptr);
+
+      H5Dwrite(dset, banktype, memspace, dspace_rank, H5P_DEFAULT, bank.data());
+
+      H5Sclose(memspace);
+      H5Sclose(dspace_rank);
+    }
+
+    H5Dclose(dset);
+
+#ifdef OPENMC_MPI
+    std::copy(temp_bank.begin(), temp_bank.end(), bank.begin());
+#endif
+  }
+#ifdef OPENMC_MPI
+  else {
+    if (!bank.empty()) {
+      MPI_Send(
+        bank.data(), bank.size(), mpi_dtype, 0, mpi::rank, mpi::intracomm);
+    }
+  }
+#endif
+#endif
+}
+
+} // namespace openmc
+
+#endif // OPENMC_BANK_IO_H

include/openmc/collision_track.h

@@ -0,0 +1,23 @@
+#ifndef OPENMC_COLLISION_TRACK_H
+#define OPENMC_COLLISION_TRACK_H
+
+#include <string>
+
+namespace openmc {
+
+class Particle;
+
+//! Reserve space in the collision track bank according to user settings.
+void collision_track_reserve_bank();
+
+//! Write collision track data to disk when the bank is full or the batch ends.
+void collision_track_flush_bank();
+
+//! Record the current particle as a collision-track entry when applicable.
+//!
+//! \param particle Particle whose collision should be recorded if eligible
+void collision_track_record(Particle& particle);
+
+} // namespace openmc
+
+#endif // OPENMC_COLLISION_TRACK_H

include/openmc/constants.h

@@ -34,6 +34,7 @@ constexpr array<int, 2> VERSION_VOXEL {2, 0};
 constexpr array<int, 2> VERSION_MGXS_LIBRARY {1, 0};
 constexpr array<int, 2> VERSION_PROPERTIES {1, 1};
 constexpr array<int, 2> VERSION_WEIGHT_WINDOWS {1, 0};
+constexpr array<int, 2> VERSION_COLLISION_TRACK {1, 0};
 
 // ============================================================================
 // ADJUSTABLE PARAMETERS

include/openmc/mcpl_interface.h

@@ -38,6 +38,21 @@ vector<SourceSite> mcpl_source_sites(std::string path);
 void write_mcpl_source_point(const char* filename, span<SourceSite> source_bank,
   const vector<int64_t>& bank_index);
 
+//! Write an MCPL collision track file
+//!
+//! This function writes collision track data to an MCPL file. Additional
+//! collision-specific metadata (such as energy deposition, material info, etc.)
+//! is stored in the file header as blob data.
+//!
+//! \param[in] filename     Path to MCPL file
+//! \param[in] collision_track_bank  Vector of CollisionTrackSites to write to
+//!                         file for this MPI rank.
+//! \param[in] bank_index   Pointer to vector of site index ranges over all
+//!                         MPI ranks.
+void write_mcpl_collision_track(const char* filename,
+  span<CollisionTrackSite> collision_track_bank,
+  const vector<int64_t>& bank_index);
+
 //! Check if MCPL functionality is available
 bool is_mcpl_interface_available();
 

include/openmc/message_passing.h

@@ -18,6 +18,7 @@ extern bool master;
 
 #ifdef OPENMC_MPI
 extern MPI_Datatype source_site;
+extern MPI_Datatype collision_track_site;
 extern MPI_Comm intracomm;
 #endif
 

include/openmc/particle_data.h

@@ -56,6 +56,25 @@ struct SourceSite {
   int64_t progeny_id;
 };
 
+struct CollisionTrackSite {
+  Position r;
+  Direction u;
+  double E;
+  double dE;
+  double time {0.0};
+  double wgt {1.0};
+  int event_mt {0};
+  int delayed_group {0};
+  int cell_id {0};
+  int nuclide_id;
+  int material_id {0};
+  int universe_id {0};
+  int n_collision {0};
+  ParticleType particle;
+  int64_t parent_id;
+  int64_t progeny_id;
+};
+
 //! State of a particle used for particle track files
 struct TrackState {
   Position r;           //!< Position in [cm]

include/openmc/settings.h

@@ -32,6 +32,24 @@ enum class IFPParameter {
   GenerationTime,
 };
 
+struct CollisionTrackConfig {
+  bool mcpl_write {false}; //!< Write collision tracks using MCPL?
+  std::unordered_set<int>
+    cell_ids; //!< Cell ids where collisions will be written
+  std::unordered_set<int>
+    mt_numbers; //!< MT Numbers where collisions will be written
+  std::unordered_set<int>
+    universe_ids; //!< Universe IDs where collisions will be written
+  std::unordered_set<int>
+    material_ids; //!< Material IDs where collisions will be written
+  std::unordered_set<std::string>
+    nuclides; //!< Nuclides where collisions will be written
+  double deposited_energy_threshold {0.0}; //!< Minimum deposited energy [eV]
+  int64_t max_collisions {
+    1000};               //!< Maximum events recorded per collision track file
+  int64_t max_files {1}; //!< Maximum number of collision track files
+};
+
 //==============================================================================
 // Global variable declarations
 //==============================================================================
@@ -41,6 +59,7 @@ namespace settings {
 // Boolean flags
 extern bool assume_separate;      //!< assume tallies are spatially separate?
 extern bool check_overlaps;       //!< check overlaps in geometry?
+extern bool collision_track;      //!< flag to use collision track feature?
 extern bool confidence_intervals; //!< use confidence intervals for results?
 extern bool
   create_fission_neutrons; //!< create fission neutrons (fixed source)?
@@ -145,6 +164,7 @@ extern std::unordered_set<int>
   statepoint_batch; //!< Batches when state should be written
 extern std::unordered_set<int>
   source_write_surf_id; //!< Surface ids where sources will be written
+extern CollisionTrackConfig collision_track_config;
 extern double source_rejection_fraction; //!< Minimum fraction of source sites
                                          //!< that must be accepted
 extern double free_gas_threshold;        //!< Threshold multiplier for free gas

include/openmc/simulation.h

@@ -22,6 +22,7 @@ constexpr int STATUS_EXIT_ON_TRIGGER {2};
 
 namespace simulation {
 
+extern int ct_current_file;   //!< current collision track file index
 extern "C" int current_batch; //!< current batch
 extern "C" int current_gen;   //!< current fission generation
 extern "C" bool initialized;  //!< has simulation been initialized?

include/openmc/urr.h

@@ -27,10 +27,10 @@ public:
     double heating;
   };
 
-  Interpolation interp_;          //!< interpolation type
-  int inelastic_flag_;            //!< inelastic competition flag
-  int absorption_flag_;           //!< other absorption flag
-  bool multiply_smooth_;          //!< multiply by smooth cross section?
+  Interpolation interp_; //!< interpolation type
+  int inelastic_flag_;   //!< inelastic competition flag
+  int absorption_flag_;  //!< other absorption flag
+  bool multiply_smooth_; //!< multiply by smooth cross section?
 
   vector<double> energy_; //!< incident energies
   auto n_energy() const { return energy_.size(); }