Use correct energy when generating Cherenkov and scintillation photons

src/celeritas/optical/OpticalCollector.cc

@@ -56,7 +56,8 @@ OpticalCollector::OpticalCollector(CoreParams const& core, Input&& inp)
                                                            std::move(la_inp));
 
     // Create core action to gather pre-step data for populating distributions
-    gather_ = OffloadGatherAction::make_and_insert(core);
+    pre_gather_
+        = OffloadGatherAction<StepActionOrder::pre>::make_and_insert(core);
 
     // Create optical action to generate Cherenkov or scintillation photons
     generate_ = optical::GeneratorAction::make_and_insert(
@@ -66,7 +67,7 @@ OpticalCollector::OpticalCollector(CoreParams const& core, Input&& inp)
     {
         // Create core action to generate Cherenkov optical distributions
         OffloadAction<GT::cherenkov>::Input oa_inp;
-        oa_inp.step_id = gather_->aux_id();
+        oa_inp.pre_step_id = pre_gather_->aux_id();
         oa_inp.gen_id = generate_->aux_id();
         oa_inp.optical_id = launch_->aux_id();
         oa_inp.material = inp.optical_params->material();
@@ -76,9 +77,15 @@ OpticalCollector::OpticalCollector(CoreParams const& core, Input&& inp)
     }
     if (inp.optical_params->scintillation())
     {
+        // Create core action to gather post-along-step state data
+        pre_post_gather_
+            = OffloadGatherAction<StepActionOrder::pre_post>::make_and_insert(
+                core);
+
         // Create action to generate scintillation optical distributions
         OffloadAction<GT::scintillation>::Input oa_inp;
-        oa_inp.step_id = gather_->aux_id();
+        oa_inp.pre_step_id = pre_gather_->aux_id();
+        oa_inp.pre_post_step_id = pre_post_gather_->aux_id();
         oa_inp.gen_id = generate_->aux_id();
         oa_inp.optical_id = launch_->aux_id();
         oa_inp.material = inp.optical_params->material();

src/celeritas/optical/OpticalCollector.hh

@@ -30,6 +30,7 @@ class CoreStateInterface;
 class CoreParams;
 template<GeneratorType G>
 class OffloadAction;
+template<StepActionOrder S>
 class OffloadGatherAction;
 class ScintillationParams;
 
@@ -141,15 +142,19 @@ class OpticalCollector
     using GT = GeneratorType;
     using SPCherenkovOffload = std::shared_ptr<OffloadAction<GT::cherenkov>>;
     using SPScintOffload = std::shared_ptr<OffloadAction<GT::scintillation>>;
-    using SPGatherAction = std::shared_ptr<OffloadGatherAction>;
+    using SPPreGatherAction
+        = std::shared_ptr<OffloadGatherAction<StepActionOrder::pre>>;
+    using SPPrePostGatherAction
+        = std::shared_ptr<OffloadGatherAction<StepActionOrder::pre_post>>;
     using SPGenerator = std::shared_ptr<optical::GeneratorAction>;
     using SPLaunchAction = std::shared_ptr<detail::OpticalLaunchAction>;
     using SPActionTimes = std::shared_ptr<ActionTimes>;
 
     //// DATA ////
 
     SPConstOpticalParams optical_params_;
-    SPGatherAction gather_;
+    SPPreGatherAction pre_gather_;
+    SPPrePostGatherAction pre_post_gather_;
     SPCherenkovOffload cherenkov_offload_;
     SPScintOffload scint_offload_;
     SPGenerator generate_;

src/celeritas/optical/gen/OffloadAction.cc

@@ -55,6 +55,7 @@ OffloadAction<G>::OffloadAction(ActionId id, Input&& inp)
 {
     CELER_EXPECT(action_id_);
     CELER_EXPECT(data_);
+    CELER_EXPECT(data_.pre_post_step_id || G == GeneratorType::cherenkov);
 }
 
 //---------------------------------------------------------------------------//
@@ -122,16 +123,22 @@ template<GeneratorType G>
 void OffloadAction<G>::offload(CoreParams const& core_params,
                                CoreStateHost& core_state) const
 {
-    auto& step = core_state.aux_data<OffloadStepStateData>(data_.step_id);
+    auto& pre_step
+        = core_state.aux_data<PreTraitsT::template Data>(data_.pre_step_id);
     auto& gen_state = get<optical::GeneratorState<MemSpace::native>>(
         core_state.aux(), data_.gen_id);
-    TrackExecutor execute{core_params.ptr<MemSpace::native>(),
-                          core_state.ptr(),
-                          Executor{data_.material->host_ref(),
-                                   data_.shared->host_ref(),
-                                   gen_state.store.ref(),
-                                   step,
-                                   gen_state.counters.buffer_size}};
+    TrackExecutor execute{
+        core_params.ptr<MemSpace::native>(),
+        core_state.ptr(),
+        Executor{data_.material->host_ref(),
+                 data_.shared->host_ref(),
+                 gen_state.store.ref(),
+                 pre_step,
+                 (G == GeneratorType::scintillation)
+                     ? core_state.aux_data<PostTraitsT::template Data>(
+                           data_.pre_post_step_id)
+                     : NativeRef<PostTraitsT::template Data>{},
+                 gen_state.counters.buffer_size}};
     launch_action(*this, core_params, core_state, execute);
 }
 

src/celeritas/optical/gen/OffloadAction.cu

@@ -30,16 +30,22 @@ template<GeneratorType G>
 void OffloadAction<G>::offload(CoreParams const& core_params,
                                CoreStateDevice& core_state) const
 {
-    auto& step = core_state.aux_data<OffloadStepStateData>(data_.step_id);
+    auto& pre_step
+        = core_state.aux_data<PreTraitsT::template Data>(data_.pre_step_id);
     auto& gen_state = get<optical::GeneratorState<MemSpace::native>>(
         core_state.aux(), data_.gen_id);
-    TrackExecutor execute{core_params.ptr<MemSpace::native>(),
-                          core_state.ptr(),
-                          Executor{data_.material->device_ref(),
-                                   data_.shared->device_ref(),
-                                   gen_state.store.ref(),
-                                   step,
-                                   gen_state.counters.buffer_size}};
+    TrackExecutor execute{
+        core_params.ptr<MemSpace::native>(),
+        core_state.ptr(),
+        Executor{data_.material->device_ref(),
+                 data_.shared->device_ref(),
+                 gen_state.store.ref(),
+                 pre_step,
+                 (G == GeneratorType::scintillation)
+                     ? core_state.aux_data<PostTraitsT::template Data>(
+                           data_.pre_post_step_id)
+                     : NativeRef<PostTraitsT::template Data>{},
+                 gen_state.counters.buffer_size}};
     static ActionLauncher<decltype(execute)> const launch_kernel(*this);
     launch_kernel(core_state, execute);
 }

src/celeritas/optical/gen/OffloadAction.hh

@@ -15,6 +15,7 @@
 
 #include "GeneratorData.hh"
 
+#include "detail/OffloadGatherTraits.hh"
 #include "detail/OffloadTraits.hh"
 
 namespace celeritas
@@ -42,15 +43,16 @@ class OffloadAction final : public CoreStepActionInterface
     //! Offload input data
     struct Input
     {
-        AuxId step_id;
+        AuxId pre_step_id;
+        AuxId pre_post_step_id;
         AuxId gen_id;
         AuxId optical_id;
         SPConstMaterial material;
         SPConstParams shared;
 
         explicit operator bool() const
         {
-            return step_id && gen_id && optical_id && material && shared;
+            return pre_step_id && gen_id && optical_id && material && shared;
         }
     };
 
@@ -77,7 +79,7 @@ class OffloadAction final : public CoreStepActionInterface
     //! \name StepAction interface
 
     //! Dependency ordering of the action
-    StepActionOrder order() const final { return StepActionOrder::user_post; }
+    StepActionOrder order() const final { return TraitsT::order; }
     // Launch kernel with host data
     void step(CoreParams const&, CoreStateHost&) const final;
     // Launch kernel with device data
@@ -92,6 +94,8 @@ class OffloadAction final : public CoreStepActionInterface
   private:
     //// TYPES ////
 
+    using PreTraitsT = detail::OffloadGatherTraits<StepActionOrder::pre>;
+    using PostTraitsT = detail::OffloadGatherTraits<StepActionOrder::pre_post>;
     using Executor = typename TraitsT::Executor;
 
     //// DATA ////

src/celeritas/optical/gen/OffloadData.hh

@@ -56,22 +56,22 @@ struct OffloadPreStepData
 
 //---------------------------------------------------------------------------//
 /*!
- * Pre-step data that is cached and used to generate optical distributions.
+ * Particle speed after the continuous part of the step.
  */
-template<Ownership W, MemSpace M>
-struct OffloadStepStateData
+struct OffloadPrePostStepData
 {
-    //// TYPES ////
-
-    template<class T>
-    using StateItems = StateCollection<T, W, M>;
-
-    //// DATA ////
-
-    // Pre-step data for generating optical photon distributions
-    StateItems<OffloadPreStepData> step;
+    units::LightSpeed speed;
+};
 
-    //// METHODS ////
+//---------------------------------------------------------------------------//
+/*!
+ * State data that is cached and used to generate optical distributions.
+ */
+template<class StepDataT, Ownership W, MemSpace M>
+struct OffloadStateData
+{
+    // State data for generating optical photon distributions
+    StateCollection<StepDataT, W, M> step;
 
     //! Number of states
     CELER_FUNCTION size_type size() const { return step.size(); }
@@ -81,20 +81,26 @@ struct OffloadStepStateData
 
     //! Assign from another set of data
     template<Ownership W2, MemSpace M2>
-    OffloadStepStateData& operator=(OffloadStepStateData<W2, M2>& other)
+    OffloadStateData& operator=(OffloadStateData<StepDataT, W2, M2>& other)
     {
         CELER_EXPECT(other);
         step = other.step;
         return *this;
     }
 };
 
+template<Ownership W, MemSpace M>
+using OffloadPreStateData = OffloadStateData<OffloadPreStepData, W, M>;
+
+template<Ownership W, MemSpace M>
+using OffloadPrePostStateData = OffloadStateData<OffloadPrePostStepData, W, M>;
+
 //---------------------------------------------------------------------------//
 /*!
- * Resize optical step states.
+ * Resize optical offload step states.
  */
-template<MemSpace M>
-void resize(OffloadStepStateData<Ownership::value, M>* state,
+template<class StepDataT, MemSpace M>
+void resize(OffloadStateData<StepDataT, Ownership::value, M>* state,
             StreamId,
             size_type size)
 {

src/celeritas/optical/gen/OffloadGatherAction.cc

@@ -17,16 +17,18 @@
 #include "celeritas/global/CoreTrackData.hh"
 #include "celeritas/global/TrackExecutor.hh"
 
-#include "detail/OffloadGatherExecutor.hh"
+#include "detail/OffloadPreGatherExecutor.hh"
+#include "detail/OffloadPrePostGatherExecutor.hh"
 
 namespace celeritas
 {
 //---------------------------------------------------------------------------//
 /*!
  * Construct and add to core params.
  */
-std::shared_ptr<OffloadGatherAction>
-OffloadGatherAction::make_and_insert(CoreParams const& core)
+template<StepActionOrder S>
+std::shared_ptr<OffloadGatherAction<S>>
+OffloadGatherAction<S>::make_and_insert(CoreParams const& core)
 {
     ActionRegistry& actions = *core.action_reg();
     AuxParamsRegistry& aux = *core.aux_reg();
@@ -41,7 +43,8 @@ OffloadGatherAction::make_and_insert(CoreParams const& core)
 /*!
  * Construct with action ID and aux ID.
  */
-OffloadGatherAction::OffloadGatherAction(ActionId action_id, AuxId aux_id)
+template<StepActionOrder S>
+OffloadGatherAction<S>::OffloadGatherAction(ActionId action_id, AuxId aux_id)
     : action_id_(action_id), aux_id_(aux_id)
 {
     CELER_EXPECT(action_id_);
@@ -52,44 +55,43 @@ OffloadGatherAction::OffloadGatherAction(ActionId action_id, AuxId aux_id)
 /*!
  * Build state data for a stream.
  */
-auto OffloadGatherAction::create_state(MemSpace m,
-                                       StreamId id,
-                                       size_type size) const -> UPState
+template<StepActionOrder S>
+auto OffloadGatherAction<S>::create_state(MemSpace m,
+                                          StreamId id,
+                                          size_type size) const -> UPState
 {
-    return make_aux_state<OffloadStepStateData>(m, id, size);
-}
-
-//---------------------------------------------------------------------------//
-/*!
- * Descriptive name of the action.
- */
-std::string_view OffloadGatherAction::description() const
-{
-    return "gather pre-step data to generate optical distributions";
+    return make_aux_state<TraitsT::template Data>(m, id, size);
 }
 
 //---------------------------------------------------------------------------//
 /*!
  * Gather pre-step data.
  */
-void OffloadGatherAction::step(CoreParams const& params,
-                               CoreStateHost& state) const
+template<StepActionOrder S>
+void OffloadGatherAction<S>::step(CoreParams const& params,
+                                  CoreStateHost& state) const
 {
-    auto& step = state.aux_data<OffloadStepStateData>(aux_id_);
-    auto execute
-        = make_active_track_executor(params.ptr<MemSpace::native>(),
-                                     state.ptr(),
-                                     detail::OffloadGatherExecutor{step});
+    auto& step = state.aux_data<TraitsT::template Data>(aux_id_);
+    auto execute = make_active_track_executor(
+        params.ptr<MemSpace::native>(), state.ptr(), Executor{step});
     launch_action(*this, params, state, execute);
 }
 
 //---------------------------------------------------------------------------//
 #if !CELER_USE_DEVICE
-void OffloadGatherAction::step(CoreParams const&, CoreStateDevice&) const
+template<StepActionOrder S>
+void OffloadGatherAction<S>::step(CoreParams const&, CoreStateDevice&) const
 {
     CELER_NOT_CONFIGURED("CUDA OR HIP");
 }
 #endif
 
+//---------------------------------------------------------------------------//
+// EXPLICIT INSTANTIATION
+//---------------------------------------------------------------------------//
+
+template class OffloadGatherAction<StepActionOrder::pre>;
+template class OffloadGatherAction<StepActionOrder::pre_post>;
+
 //---------------------------------------------------------------------------//
 }  // namespace celeritas

src/celeritas/optical/gen/OffloadGatherAction.cu

@@ -12,25 +12,32 @@
 #include "celeritas/global/CoreState.hh"
 #include "celeritas/global/TrackExecutor.hh"
 
-#include "detail/OffloadGatherExecutor.hh"
+#include "detail/OffloadPreGatherExecutor.hh"
+#include "detail/OffloadPrePostGatherExecutor.hh"
 
 namespace celeritas
 {
 //---------------------------------------------------------------------------//
 /*!
  * Gather pre-step data.
  */
-void OffloadGatherAction::step(CoreParams const& params,
-                               CoreStateDevice& state) const
+template<StepActionOrder S>
+void OffloadGatherAction<S>::step(CoreParams const& params,
+                                  CoreStateDevice& state) const
 {
-    auto& step = state.aux_data<OffloadStepStateData>(aux_id_);
-    auto execute
-        = make_active_track_executor(params.ptr<MemSpace::native>(),
-                                     state.ptr(),
-                                     detail::OffloadGatherExecutor{step});
+    auto& step = state.aux_data<TraitsT::template Data>(aux_id_);
+    auto execute = make_active_track_executor(
+        params.ptr<MemSpace::native>(), state.ptr(), Executor{step});
     static ActionLauncher<decltype(execute)> const launch_kernel(*this);
     launch_kernel(state, execute);
 }
 
+//---------------------------------------------------------------------------//
+// EXPLICIT INSTANTIATION
+//---------------------------------------------------------------------------//
+
+template class OffloadGatherAction<StepActionOrder::pre>;
+template class OffloadGatherAction<StepActionOrder::pre_post>;
+
 //---------------------------------------------------------------------------//
 }  // namespace celeritas