Fix celer-g4 primary generator MT reproducibility (#1057)

celeritas-project · Dec 10, 2023 · 238b17c · 238b17c
1 parent ef8ec19
commit 238b17c
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 14 deletions.
diff --git a/app/celer-g4/PGPrimaryGeneratorAction.cc b/app/celer-g4/PGPrimaryGeneratorAction.cc
@@ -32,9 +32,7 @@ PGPrimaryGeneratorAction::PGPrimaryGeneratorAction(
     // Generate one particle at each call to \c GeneratePrimaryVertex()
     gun_.SetNumberOfParticles(1);
 
-    // Seed with an independent value for each thread
-    rng_.seed(options.seed + get_geant_thread_id());
-
+    seed_ = options.seed;
     num_events_ = options.num_events;
     primaries_per_event_ = options.primaries_per_event;
     sample_energy_ = make_energy_sampler(options.energy);
@@ -57,32 +55,38 @@ PGPrimaryGeneratorAction::PGPrimaryGeneratorAction(
 void PGPrimaryGeneratorAction::GeneratePrimaries(G4Event* event)
 {
     CELER_EXPECT(event);
+    CELER_EXPECT(event->GetEventID() >= 0);
 
-    if (event_count_ == num_events_)
+    size_type event_id = event->GetEventID();
+    if (event_id >= num_events_)
     {
         return;
     }
 
+    // Seed with an independent value for each event. Since Geant4 schedules
+    // events dynamically, the same event ID may not be mapped to the same
+    // thread across multiple runs. For reproducibility, Geant4 reseeds each
+    // worker thread's RNG at the start of each event with a seed calculated
+    // from the event ID.
+    rng_.seed(seed_ + event_id);
+
     for (size_type i = 0; i < primaries_per_event_; ++i)
     {
-        gun_.SetParticleDefinition(
-            particle_def_[primary_count_ % particle_def_.size()]);
+        gun_.SetParticleDefinition(particle_def_[i % particle_def_.size()]);
         gun_.SetParticlePosition(
             convert_to_geant(sample_pos_(rng_), CLHEP::cm));
         gun_.SetParticleMomentumDirection(
             convert_to_geant(sample_dir_(rng_), 1));
         gun_.SetParticleEnergy(
             convert_to_geant(sample_energy_(rng_), CLHEP::MeV));
         gun_.GeneratePrimaryVertex(event);
-        ++primary_count_;
 
         if (CELERITAS_DEBUG)
         {
             CELER_ASSERT(G4VPrimaryGenerator::CheckVertexInsideWorld(
                 gun_.GetParticlePosition()));
         }
     }
-    ++event_count_;
 
     CELER_ENSURE(event->GetNumberOfPrimaryVertex()
                  == static_cast<int>(primaries_per_event_));

diff --git a/app/celer-g4/PGPrimaryGeneratorAction.hh b/app/celer-g4/PGPrimaryGeneratorAction.hh
@@ -58,8 +58,7 @@ class PGPrimaryGeneratorAction final : public G4VUserPrimaryGeneratorAction
     EnergySampler sample_energy_;
     PositionSampler sample_pos_;
     DirectionSampler sample_dir_;
-    size_type primary_count_{0};
-    size_type event_count_{0};
+    size_type seed_{0};
 };
 
 //---------------------------------------------------------------------------//

diff --git a/src/celeritas/phys/PrimaryGenerator.cc b/src/celeritas/phys/PrimaryGenerator.cc
@@ -78,14 +78,13 @@ auto PrimaryGenerator::operator()() -> result_type
     for (auto i : range(primaries_per_event_))
     {
         Primary& p = result[i];
-        p.particle_id = particle_id_[primary_count_ % particle_id_.size()];
+        p.particle_id = particle_id_[i % particle_id_.size()];
         p.energy = units::MevEnergy{sample_energy_(rng_)};
         p.position = sample_pos_(rng_);
         p.direction = sample_dir_(rng_);
         p.time = 0;
         p.event_id = EventId{event_count_};
         p.track_id = TrackId{i};
-        ++primary_count_;
     }
     ++event_count_;
     return result;

diff --git a/src/celeritas/phys/PrimaryGenerator.hh b/src/celeritas/phys/PrimaryGenerator.hh
@@ -77,7 +77,6 @@ class PrimaryGenerator : public EventReaderInterface
     PositionSampler sample_pos_;
     DirectionSampler sample_dir_;
     std::vector<ParticleId> particle_id_;
-    size_type primary_count_{0};
     size_type event_count_{0};
     PrimaryGeneratorEngine rng_;
 };

diff --git a/test/celeritas/phys/PrimaryGenerator.test.cc b/test/celeritas/phys/PrimaryGenerator.test.cc
@@ -89,7 +89,7 @@ TEST_F(PrimaryGeneratorTest, basic)
     auto primaries = generate_primaries();
     EXPECT_TRUE(primaries.empty());
 
-    static int const expected_particle_id[] = {0, 1, 0, 1, 0, 1};
+    static int const expected_particle_id[] = {0, 1, 0, 0, 1, 0};
     static int const expected_event_id[] = {0, 0, 0, 1, 1, 1};
     static int const expected_track_id[] = {0, 1, 2, 0, 1, 2};