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

@@ -52,6 +52,7 @@ struct SourceSite {
   int parent_nuclide {-1};
   int64_t parent_id {0};
   int64_t progeny_id {0};
+  int64_t root_index {-1};
   double wgt_born {1.0};
   double wgt_ww_born {-1.0};
   int64_t n_split {0};
@@ -556,6 +557,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 +739,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

@@ -49,6 +49,8 @@ extern const RegularMesh* ufs_mesh;
 extern vector<double> k_generation;
 extern vector<int64_t> work_index;
 
+extern vector<int64_t> phase1_work_index;
+
 extern int64_t
   simulation_tracks_completed; //!< Number of tracks completed on this rank
 

include/openmc/tallies/pulse_height.h

@@ -0,0 +1,71 @@
+//! \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
+  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 pht per-cell energy deposited by this particle alone
+void stage_pulse_height(int64_t root_index, const vector<double>& pht);
+
+//! Aggregate staged contributions by history and score them.
+//
+//! Sends each contribution to the rank that owns its root according to
+//! simulation::phase1_work_index, 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.
+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

@@ -32,6 +32,7 @@ class _SourceSite(Structure):
                 ('parent_nuclide', c_int),
                 ('parent_id', c_int64),
                 ('progeny_id', c_int64),
+                ('root_index', c_int64),
                 ('wgt_born', c_double),
                 ('wgt_ww_born', c_double),
                 ('n_split', c_int64),

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

@@ -161,7 +161,7 @@ void initialize_mpi(MPI_Comm intracomm)
 
   // Create bank datatype
   SourceSite b;
-  MPI_Aint disp[15];
+  MPI_Aint disp[16];
   MPI_Get_address(&b.r, &disp[0]);
   MPI_Get_address(&b.u, &disp[1]);
   MPI_Get_address(&b.E, &disp[2]);
@@ -173,16 +173,17 @@ void initialize_mpi(MPI_Comm intracomm)
   MPI_Get_address(&b.parent_nuclide, &disp[8]);
   MPI_Get_address(&b.parent_id, &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]);
-  MPI_Get_address(&b.n_split, &disp[13]);
-  MPI_Get_address(&b.n_collision, &disp[14]);
-  for (int i = 14; i >= 0; --i) {
+  MPI_Get_address(&b.root_index, &disp[11]);
+  MPI_Get_address(&b.wgt_born, &disp[12]);
+  MPI_Get_address(&b.wgt_ww_born, &disp[13]);
+  MPI_Get_address(&b.n_split, &disp[14]);
+  MPI_Get_address(&b.n_collision, &disp[15]);
+  for (int i = 15; i >= 0; --i) {
     disp[i] -= disp[0];
   }
 
   // Block counts for each field
-  int blocks[] = {3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
+  int blocks[] = {3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
 
   // Types for each field
   MPI_Datatype types[] = {
@@ -197,13 +198,14 @@ void initialize_mpi(MPI_Comm intracomm)
     MPI_INT,     // parent_nuclide
     MPI_INT64_T, // parent_id
     MPI_INT64_T, // progeny_id
+    MPI_INT64_T, // root_index
     MPI_DOUBLE,  // wgt_born
     MPI_DOUBLE,  // wgt_ww_born
     MPI_INT64_T, // n_split
     MPI_INT      // n_collision
   };
 
-  MPI_Type_create_struct(15, blocks, disp, types, &mpi::source_site);
+  MPI_Type_create_struct(16, blocks, disp, types, &mpi::source_site);
   MPI_Type_commit(&mpi::source_site);
 
   CollisionTrackSite bc;
@@ -389,28 +391,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 +485,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 +530,6 @@ void read_separate_xml_files()
 
   read_tallies_xml();
 
-  check_pulse_height_compatibility();
-
   // Initialize distribcell_filters
   prepare_distribcell();
 

src/particle.cpp

@@ -29,6 +29,7 @@
 #include "openmc/source.h"
 #include "openmc/surface.h"
 #include "openmc/tallies/derivative.h"
+#include "openmc/tallies/pulse_height.h"
 #include "openmc/tallies/tally.h"
 #include "openmc/tallies/tally_scoring.h"
 #include "openmc/track_output.h"
@@ -109,10 +110,29 @@ bool Particle::create_secondary(
   if (settings::use_shared_secondary_bank) {
     bank.progeny_id = n_progeny()++;
   }
+  bank.root_index = root_index();
   bank.wgt_born = wgt_born();
   bank.wgt_ww_born = wgt_ww_born();
   bank.n_split = n_split();
 
+  // Remove the energy carried off by this secondary from the parent's interim
+  // pulse-height result for the cell the parent is currently in. In non-shared
+  // mode the equivalent subtraction is performed at revival by
+  // pht_secondary_particles(); doing it here instead is equivalent, because the
+  // secondary is born at the parent's position and therefore in the parent's
+  // current cell, and it avoids the exhaustive_find_cell() call needed there.
+  // Placing this after the energy-cutoff early return above means a secondary
+  // that is never created is never subtracted, matching the non-shared path.
+  if (settings::use_shared_secondary_bank &&
+      !model::active_pulse_height_tallies.empty() && type.is_photon()) {
+    auto it = std::find(model::pulse_height_cells.begin(),
+      model::pulse_height_cells.end(), lowest_coord().cell());
+    if (it != model::pulse_height_cells.end()) {
+      int index = std::distance(model::pulse_height_cells.begin(), it);
+      pht_storage()[index] -= bank.E;
+    }
+  }
+
   local_secondary_bank().emplace_back(bank);
   return true;
 }
@@ -143,6 +163,10 @@ void Particle::split(double wgt)
   if (settings::use_shared_secondary_bank) {
     bank.progeny_id = n_progeny()++;
   }
+  // A split clone belongs to the same history as its parent. No pulse-height
+  // subtraction is applied here: a split is a weight artifact, not a physical
+  // secondary, and its energy is not carried away from the parent.
+  bank.root_index = root_index();
 
   local_secondary_bank().emplace_back(bank);
 }
@@ -502,15 +526,20 @@ void Particle::event_revive_from_secondary(const SourceSite& site)
 
   from_source(&site);
 
+  // Inherit the root of the tree this secondary belongs to. from_source() does
+  // not copy this, because it is also used for primaries read from the source
+  // bank, whose root index is assigned in initialize_particle_track().
+  root_index() = site.root_index;
+
   n_event() = 0;
   if (!settings::use_shared_secondary_bank) {
     n_tracks()++;
   }
   bank_second_E() = 0.0;
 
   // Subtract secondary particle energy from interim pulse-height results.
-  // In shared secondary mode, this subtraction was already done on the parent
-  // particle during create_secondary(), so skip it here.
+  // In shared secondary mode this subtraction is performed on the parent in
+  // create_secondary(), so skip it here.
   if (!settings::use_shared_secondary_bank &&
       !model::active_pulse_height_tallies.empty() && this->type().is_photon()) {
     // Since the birth cell of the particle has not been set we
@@ -604,7 +633,15 @@ void Particle::event_death()
   keff_tally_leakage() = 0.0;
 
   if (!model::active_pulse_height_tallies.empty()) {
-    score_pulse_height_tally(*this, model::active_pulse_height_tallies);
+    if (settings::use_shared_secondary_bank) {
+      // This Particle carries only one fragment of its history's pulse. Stage
+      // it for aggregation by root index; scoring happens once per history in
+      // finalize_pulse_height_tallies() after all generations have drained.
+      stage_pulse_height(root_index(), pht_storage());
+    } else {
+      score_pulse_height_tally(
+        *this, pht_storage(), model::active_pulse_height_tallies);
+    }
   }
 
   // Accumulate track count for this particle history