Support pulse-height tallies in shared secondary mode

CMakeLists.txt

@@ -478,6 +478,7 @@ list(APPEND libopenmc_SOURCES
   src/tallies/filter_universe.cpp
   src/tallies/filter_weight.cpp
   src/tallies/filter_zernike.cpp
+  src/tallies/pulse_height.cpp
   src/tallies/tally.cpp
   src/tallies/tally_scoring.cpp
   src/tallies/trigger.cpp

include/openmc/particle_data.h

@@ -50,12 +50,33 @@ struct SourceSite {
 
   // Extra attributes that don't show up in source written to file
   int parent_nuclide {-1};
-  int64_t parent_id {0};
+
+  //! Index of an ancestor of this site, in one of two senses that never
+  //! overlap. Before the site is placed by collection it is the immediate
+  //! parent's slot in the current generation's work, which with progeny_id
+  //! gives the site its position in the collected bank and is used for nothing
+  //! else. Placement consumes that key, and resolve_root_indices() then
+  //! overwrites the field with the index of the primary at the root of the
+  //! site's history, which is what survives the sort and the MPI migration the
+  //! shared secondary bank performs between generations. Read it through
+  //! parent_slot() or root_index() so the sense is explicit at every use.
+  int64_t ancestor_index {0};
+
   int64_t progeny_id {0};
   double wgt_born {1.0};
   double wgt_ww_born {-1.0};
   int64_t n_split {0};
   int n_collision {0};
+
+  //! Slot of the immediate parent within the current generation's work. Valid
+  //! from creation until the site is placed by collection.
+  int64_t& parent_slot() { return ancestor_index; }
+  int64_t parent_slot() const { return ancestor_index; }
+
+  //! Index of the primary at the root of this site's history. Valid once the
+  //! site has been placed by collection.
+  int64_t& root_index() { return ancestor_index; }
+  int64_t root_index() const { return ancestor_index; }
 };
 
 struct CollisionTrackSite {
@@ -556,6 +577,8 @@ class ParticleData : public GeometryState {
 
   vector<double> pht_storage_;
 
+  int64_t root_index_ {-1};
+
   double keff_tally_absorption_ {0.0};
   double keff_tally_collision_ {0.0};
   double keff_tally_tracklength_ {0.0};
@@ -736,6 +759,11 @@ class ParticleData : public GeometryState {
 
   // Interim pulse height tally storage
   vector<double>& pht_storage() { return pht_storage_; }
+  const vector<double>& pht_storage() const { return pht_storage_; }
+
+  // Index of the primary particle at the root of this particle's tree
+  int64_t& root_index() { return root_index_; }
+  int64_t root_index() const { return root_index_; }
 
   // Global tally accumulators
   double& keff_tally_absorption() { return keff_tally_absorption_; }

include/openmc/simulation.h

@@ -6,6 +6,7 @@
 
 #include "openmc/mesh.h"
 #include "openmc/particle.h"
+#include "openmc/shared_array.h"
 #include "openmc/vector.h"
 
 #include <cstdint>
@@ -64,6 +65,36 @@ void allocate_banks();
 //! Determine number of particles to transport per process
 void calculate_work(int64_t n_particles);
 
+//! First primary index owned by a rank under the phase-1 partition
+//!
+//! Recomputes what calculate_work(settings::n_particles) would produce, so
+//! that the primary partition is available after work_index has been
+//! overwritten for a secondary generation. Valid for rank in [0, n_procs],
+//! where n_procs returns the total primary count.
+//!
+//! \param rank MPI rank
+//! \return Index of that rank's first primary
+int64_t phase1_first_root(int rank);
+
+//! Rank owning a given root index under the phase-1 partition
+//!
+//! \param root Root index in [0, n_particles)
+//! \return Rank that transported that primary
+int phase1_owner_of_root(int64_t root);
+
+//! Replace the placement key of every site in a collected bank with the root
+//! of its history
+//!
+//! Must be called once per generation, after the sites have been placed and
+//! before any MPI migration, since the placement key indexes a bank that is
+//! local to the rank that produced the sites.
+//!
+//! \param sites Bank whose sites were just collected
+//! \param parents Bank the parents were transported from, or nullptr when the
+//!   parents were the primaries
+void resolve_root_indices(
+  SharedArray<SourceSite>& sites, const SharedArray<SourceSite>* parents);
+
 //! Initialize nuclear data before a simulation
 void initialize_data();
 

include/openmc/tallies/pulse_height.h

@@ -0,0 +1,79 @@
+//! \file pulse_height.h
+//! \brief Deferred, per-history aggregation of pulse-height results
+//!
+//! A pulse-height tally scores one count per source history, in the bin
+//! containing the total energy that history's entire particle tree deposited in
+//! a given cell. In the default transport modes a whole tree is carried by a
+//! single Particle object, so Particle::pht_storage() already holds the
+//! per-history total by the time event_death() runs and can be scored directly.
+//!
+//! Under the shared secondary bank each secondary generation is transported as
+//! a fresh set of Particle objects, redistributed across MPI ranks between
+//! generations. A history's deposition is therefore spread over many Particle
+//! objects on potentially many ranks. This module collects those fragments,
+//! keyed by the root index carried on every SourceSite, and scores them once
+//! per history after the generation loop has drained.
+
+#ifndef OPENMC_TALLIES_PULSE_HEIGHT_H
+#define OPENMC_TALLIES_PULSE_HEIGHT_H
+
+#include <cstdint>
+
+#include "openmc/vector.h"
+
+namespace openmc {
+
+//==============================================================================
+//! One flushed pulse-height fragment, tagged with the history it belongs to.
+//==============================================================================
+
+struct PulseHeightContribution {
+  int64_t root_index;    //!< index of the primary at the root of the tree
+  int64_t track_id;      //!< id of the track that deposited this fragment
+  vector<double> energy; //!< per-cell energy, indexed as pulse_height_cells
+};
+
+namespace simulation {
+
+//! Per-thread staging buffers, merged in finalize_pulse_height_tallies().
+extern vector<vector<PulseHeightContribution>> pht_thread_buffers;
+
+} // namespace simulation
+
+//! Allocate the per-thread staging buffers. Called from initialize_simulation()
+//! when pulse-height tallies and the shared secondary bank are both active.
+void init_pulse_height_buffers();
+
+//! Release the staging buffers and the phase-1 partition snapshot.
+void free_memory_pulse_height();
+
+//! Stage one Particle's contribution to its history's pulse height.
+//
+//! Thread-safe by construction: each thread appends only to its own buffer.
+//! Contributions that are identically zero in every cell are dropped; histories
+//! that deposit nothing are recovered in finalize_pulse_height_tallies() by
+//! iterating over the full root range rather than over staged entries.
+//
+//! \param root_index index of the primary at the root of this particle's tree
+//! \param track_id this particle's id, used to give the fragments of a history
+//!   a canonical summation order in finalize_pulse_height_tallies()
+//! \param pht per-cell energy deposited by this particle alone
+void stage_pulse_height(
+  int64_t root_index, int64_t track_id, const vector<double>& pht);
+
+//! Aggregate staged contributions by history and score them.
+//
+//! Sends each contribution to the rank that owns its root according to
+//! the phase-1 primary partition, sums per (history, cell), and scores every
+//! owned history including those with no deposition. Must be called after the
+//! last secondary generation has been transported and before tally results are
+//! accumulated for the batch.
+//
+//! Fragments are summed in order of track id rather than in arrival order, so
+//! a history's total is bit-for-bit independent of thread scheduling and of the
+//! rank a descendant happened to land on.
+void finalize_pulse_height_tallies();
+
+} // namespace openmc
+
+#endif // OPENMC_TALLIES_PULSE_HEIGHT_H

include/openmc/tallies/tally_scoring.h

@@ -120,7 +120,14 @@ void score_surface_tally(
 //
 //! \param p The particle being tracked
 //! \param tallies A vector of the indices of the tallies to score to
-void score_pulse_height_tally(Particle& p, const vector<int>& tallies);
+//! Score a completed per-history pulse-height result.
+//
+//! \param p particle used to drive filter matching; its cell and E_last are
+//!   temporarily overwritten and restored
+//! \param pht per-cell deposited energy, indexed as model::pulse_height_cells
+//! \param tallies indices of the pulse-height tallies to score into
+void score_pulse_height_tally(
+  Particle& p, const vector<double>& pht, const vector<int>& tallies);
 
 } // namespace openmc
 

openmc/lib/core.py

@@ -30,7 +30,7 @@ class _SourceSite(Structure):
                 ('surf_id', c_int),
                 ('particle', c_int32),
                 ('parent_nuclide', c_int),
-                ('parent_id', c_int64),
+                ('ancestor_index', c_int64),
                 ('progeny_id', c_int64),
                 ('wgt_born', c_double),
                 ('wgt_ww_born', c_double),

src/bank.cpp

@@ -79,8 +79,8 @@ void init_fission_bank(int64_t max)
 }
 
 // Performs an O(n) sort on a fission or secondary bank, by leveraging
-// the parent_id and progeny_id fields of banked particles. See the following
-// paper for more details:
+// the ancestor_index and progeny_id fields of banked particles. See the
+// following paper for more details:
 // "Reproducibility and Monte Carlo Eigenvalue Calculations," F.B. Brown and
 // T.M. Sutton, 1992 ANS Annual Meeting, Transactions of the American Nuclear
 // Society, Volume 65, Page 235.
@@ -121,15 +121,15 @@ void sort_bank(SharedArray<SourceSite>& bank, bool is_fission_bank)
   // Use parent and progeny indices to sort bank
   for (int64_t i = 0; i < bank.size(); i++) {
     const auto& site = bank[i];
-    if (site.parent_id < 0 ||
-        site.parent_id >=
+    if (site.parent_slot() < 0 ||
+        site.parent_slot() >=
           static_cast<int64_t>(simulation::progeny_per_particle.size())) {
       fatal_error(fmt::format("Invalid parent_id {} for banked site (expected "
                               "range [0, {})).",
-        site.parent_id, simulation::progeny_per_particle.size()));
+        site.parent_slot(), simulation::progeny_per_particle.size()));
     }
     int64_t idx =
-      simulation::progeny_per_particle[site.parent_id] + site.progeny_id;
+      simulation::progeny_per_particle[site.parent_slot()] + site.progeny_id;
     if (idx < 0 || idx >= bank.size()) {
       fatal_error("Mismatch detected between sum of all particle progeny and "
                   "bank size during sorting.");

src/finalize.cpp

@@ -25,6 +25,7 @@
 #include "openmc/simulation.h"
 #include "openmc/source.h"
 #include "openmc/surface.h"
+#include "openmc/tallies/pulse_height.h"
 #include "openmc/tallies/tally.h"
 #include "openmc/thermal.h"
 #include "openmc/timer.h"
@@ -51,6 +52,7 @@ void free_memory()
   free_memory_source();
   free_memory_mesh();
   free_memory_tally();
+  free_memory_pulse_height();
   free_memory_bank();
   free_memory_plot();
   free_memory_weight_windows();

src/initialize.cpp

@@ -171,7 +171,7 @@ void initialize_mpi(MPI_Comm intracomm)
   MPI_Get_address(&b.surf_id, &disp[6]);
   MPI_Get_address(&b.particle, &disp[7]);
   MPI_Get_address(&b.parent_nuclide, &disp[8]);
-  MPI_Get_address(&b.parent_id, &disp[9]);
+  MPI_Get_address(&b.ancestor_index, &disp[9]);
   MPI_Get_address(&b.progeny_id, &disp[10]);
   MPI_Get_address(&b.wgt_born, &disp[11]);
   MPI_Get_address(&b.wgt_ww_born, &disp[12]);
@@ -195,7 +195,7 @@ void initialize_mpi(MPI_Comm intracomm)
     MPI_INT,     // surf_id
     MPI_INT,     // particle (enum)
     MPI_INT,     // parent_nuclide
-    MPI_INT64_T, // parent_id
+    MPI_INT64_T, // ancestor_index
     MPI_INT64_T, // progeny_id
     MPI_DOUBLE,  // wgt_born
     MPI_DOUBLE,  // wgt_ww_born
@@ -389,28 +389,6 @@ int parse_command_line(int argc, char* argv[])
   return 0;
 }
 
-// TODO: Pulse-height tallies require per-history scoring across the full
-// particle tree (parent + all descendants). The shared secondary bank
-// transports each secondary as an independent Particle, breaking this
-// assumption. A proper fix would defer pulse-height scoring: save
-// (root_source_id, cell, pht_storage) per particle, then aggregate by
-// root_source_id after all secondary generations complete before scoring
-// into the histogram. For now, disable shared secondary when pulse-height
-// tallies are present.
-static void check_pulse_height_compatibility()
-{
-  if (settings::use_shared_secondary_bank) {
-    for (const auto& t : model::tallies) {
-      if (t->type_ == TallyType::PULSE_HEIGHT) {
-        settings::use_shared_secondary_bank = false;
-        warning("Pulse-height tallies are not yet compatible with the shared "
-                "secondary bank. Disabling shared secondary bank.");
-        break;
-      }
-    }
-  }
-}
-
 bool read_model_xml()
 {
   std::string model_filename = settings::path_input;
@@ -505,8 +483,6 @@ bool read_model_xml()
   if (check_for_node(root, "tallies"))
     read_tallies_xml(root.child("tallies"));
 
-  check_pulse_height_compatibility();
-
   // Initialize distribcell_filters
   prepare_distribcell();
 
@@ -552,8 +528,6 @@ void read_separate_xml_files()
 
   read_tallies_xml();
 
-  check_pulse_height_compatibility();
-
   // Initialize distribcell_filters
   prepare_distribcell();