Extend Coulomb scattering to support additional particle types (#1574)

* Update MSC params to support additional particle types

* Update WentzelOKVI to support additional particle types

Changes to the test results are from using the imported electron mass to
calculate the kin factor in the Wentzel helper instead of the constant

* Update Urban MSC to support additional particle types

* Import MSC lateral displacement flag

* Clean up a bit

* Support Urban model for muons in CelerEmStandardPhysics

* Update lar-sphere geant setup options

* Address review feedback

* Store number of particles and add documentation

app/celer-sim/Runner.cc

@@ -336,7 +336,8 @@ void Runner::build_core_params(RunnerInput const& inp,
         imported, params.particle, params.material);
 
     // Construct shared data for Coulomb scattering
-    params.wentzel = WentzelOKVIParams::from_import(imported, params.material);
+    params.wentzel = WentzelOKVIParams::from_import(
+        imported, params.material, params.particle);
 
     // Load physics: create individual processes with make_shared
     params.physics = [&params, &inp, &imported] {

src/accel/SharedParams.cc

@@ -513,7 +513,8 @@ void SharedParams::initialize_core(SetupOptions const& options)
         *imported, params.particle, params.material);
 
     // Construct shared data for Coulomb scattering
-    params.wentzel = WentzelOKVIParams::from_import(*imported, params.material);
+    params.wentzel = WentzelOKVIParams::from_import(
+        *imported, params.material, params.particle);
 
     // Load physics: create individual processes with make_shared
     params.physics = [&params, &options, &imported] {

src/celeritas/em/data/CommonCoulombData.hh

@@ -12,14 +12,12 @@
 
 namespace celeritas
 {
+//! Opaque index to MSC-applicable particles
+using MscParticleId = OpaqueId<struct MscParticle_>;
+
 //---------------------------------------------------------------------------//
 /*!
  * Physics IDs for MSC.
- *
- * \todo If we want to extend this *generally*, we should have an array (length
- * \c ParticleParams::size() ) that maps IDs to "model parameters". For
- * example, electrons and positrons probably map to the same ID. Light ions and
- * protons probably do as well.
  */
 struct CoulombIds
 {

src/celeritas/em/data/UrbanMscData.hh

@@ -17,6 +17,15 @@
 
 namespace celeritas
 {
+//! Particle categories for Urban MSC particle and material-dependent data
+enum class UrbanParMatType
+{
+    electron = 0,
+    positron,
+    muhad,
+    size_
+};
+
 //---------------------------------------------------------------------------//
 /*!
  * Settable parameters and default values for Urban multiple scattering.
@@ -101,13 +110,18 @@ struct UrbanMscMaterialData
  *
  * The scaled Zeff parameters are:
  *
- *   Particle | a    | b
- *   -------- | ---- | ----
- *   electron | 0.87 | 2/3
- *   positron | 0.7  | 1/2
+ *   Particle             | a    | b
+ *   -------------------- | ---- | ----
+ *   electron/muon/hadron | 0.87 | 2/3
+ *   positron             | 0.7  | 1/2
+ *
+ * Two different \c d_over_r values are used: one for electrons and positrons,
+ * and another for muons and hadrons.
  */
 struct UrbanMscParMatData
 {
+    using UrbanParMatId = OpaqueId<UrbanMscParMatData>;
+
     real_type scaled_zeff{};  //!< a * Z^b
     real_type d_over_r{};  //!< Maximum distance/range heuristic
 
@@ -118,10 +132,6 @@ struct UrbanMscParMatData
 //---------------------------------------------------------------------------//
 /*!
  * Device data for Urban MSC.
- *
- * Since the model currently applies only to electrons and positrons, the
- * particles are hardcoded to be length 2. TODO: extend to other charged
- * particles when further physics is implemented.
  */
 template<Ownership W, MemSpace M>
 struct UrbanMscData
@@ -131,7 +141,9 @@ struct UrbanMscData
     template<class T>
     using Items = Collection<T, W, M>;
     template<class T>
-    using MaterialItems = celeritas::Collection<T, W, M, MaterialId>;
+    using MaterialItems = Collection<T, W, M, MaterialId>;
+    template<class T>
+    using ParticleItems = Collection<T, W, M, ParticleId>;
 
     //// DATA ////
 
@@ -143,6 +155,14 @@ struct UrbanMscData
     UrbanMscParameters params;
     //! Material-dependent data
     MaterialItems<UrbanMscMaterialData> material_data;
+    //! Number of particles this model applies to
+    ParticleId::size_type num_particles;
+    //! Number of particle categories for particle and material-dependent data
+    ParticleId::size_type num_par_mat;
+    //! Map from particle ID to index in particle and material-dependent data
+    ParticleItems<UrbanMscParMatData::UrbanParMatId> pid_to_pmdata;
+    //! Map from particle ID to index in cross sections
+    ParticleItems<MscParticleId> pid_to_xs;
     //! Particle and material-dependent data
     Items<UrbanMscParMatData> par_mat_data;  //!< [mat][particle]
     //! Scaled xs data
@@ -157,6 +177,8 @@ struct UrbanMscData
     explicit CELER_FUNCTION operator bool() const
     {
         return ids && electron_mass > zero_quantity() && !material_data.empty()
+               && num_particles >= 2 && num_par_mat >= 2
+               && !pid_to_pmdata.empty() && !pid_to_xs.empty()
                && !par_mat_data.empty() && !xs.empty() && !reals.empty();
     }
 
@@ -169,22 +191,15 @@ struct UrbanMscData
         electron_mass = other.electron_mass;
         params = other.params;
         material_data = other.material_data;
+        num_particles = other.num_particles;
+        num_par_mat = other.num_par_mat;
+        pid_to_pmdata = other.pid_to_pmdata;
+        pid_to_xs = other.pid_to_xs;
         par_mat_data = other.par_mat_data;
         xs = other.xs;
         reals = other.reals;
         return *this;
     }
-
-    //! Get the data location for a material + particle
-    template<class T>
-    CELER_FUNCTION ItemId<T> at(MaterialId mat, ParticleId par) const
-    {
-        CELER_EXPECT(mat && par);
-        size_type result = mat.unchecked_get() * 2;
-        result += (par == this->ids.electron ? 0 : 1);
-        CELER_ENSURE(result < this->par_mat_data.size());
-        return ItemId<T>{result};
-    }
 };
 
 }  // namespace celeritas

src/celeritas/em/data/WentzelOKVIData.hh

@@ -82,37 +82,44 @@ struct MottElementData
 template<Ownership W, MemSpace M>
 struct WentzelOKVIData
 {
+    //// TYPES ////
+
     template<class T>
     using ElementItems = celeritas::Collection<T, W, M, ElementId>;
     template<class T>
     using IsotopeItems = celeritas::Collection<T, W, M, IsotopeId>;
     template<class T>
     using MaterialItems = Collection<T, W, M, MaterialId>;
 
-    // User-assignable parameters
-    CoulombParameters params;
+    //// DATA ////
 
-    // Constant prefactor for the squared momentum transfer [(MeV/c)^{-2}]
+    //! User-assignable parameters
+    CoulombParameters params;
+    //! Mass of electron in MeV
+    units::MevMass electron_mass;
+    //! Constant prefactor for the squared momentum transfer [(MeV/c)^{-2}]
     IsotopeItems<real_type> nuclear_form_prefactor;
-
-    // Per element form factors
+    //! Per element form factors
     ElementItems<MottElementData> mott_coeffs;
-
-    // Inverse effective A^2/3 [1/mass^2/3]
+    //! Inverse effective A^2/3 [1/mass^2/3]
     MaterialItems<real_type> inv_mass_cbrt_sq;
 
-    // Check if the data is initialized
+    //// METHODS ////
+
+    //! Check if the data is initialized
     explicit CELER_FUNCTION operator bool() const
     {
-        return params && !mott_coeffs.empty()
+        return params && electron_mass > zero_quantity() && !mott_coeffs.empty()
                && params.is_combined == !inv_mass_cbrt_sq.empty();
     }
 
+    //! Assign from another set of data
     template<Ownership W2, MemSpace M2>
     WentzelOKVIData& operator=(WentzelOKVIData<W2, M2> const& other)
     {
         CELER_EXPECT(other);
         params = other.params;
+        electron_mass = other.electron_mass;
         nuclear_form_prefactor = other.nuclear_form_prefactor;
         mott_coeffs = other.mott_coeffs;
         inv_mass_cbrt_sq = other.inv_mass_cbrt_sq;

src/celeritas/em/data/WentzelVIMscData.hh

@@ -47,15 +47,21 @@ struct WentzelVIMscData
 
     template<class T>
     using Items = Collection<T, W, M>;
+    template<class T>
+    using ParticleItems = Collection<T, W, M, ParticleId>;
 
     //// DATA ////
 
     //! Particle IDs
     CoulombIds ids;
-    //! Mass of of electron in MeV
+    //! Mass of electron in MeV
     units::MevMass electron_mass;
     //! User-assignable options
     WentzelVIMscParameters params;
+    //! Number of particles this model applies to
+    ParticleId::size_type num_particles;
+    //! Map from particle ID to index in cross sections
+    ParticleItems<MscParticleId> pid_to_xs;
     //! Scaled xs data
     Items<XsGridData> xs;  //!< [mat][particle]
 
@@ -67,8 +73,8 @@ struct WentzelVIMscData
     //! Check whether the data is assigned
     explicit CELER_FUNCTION operator bool() const
     {
-        return ids && electron_mass > zero_quantity() && !xs.empty()
-               && !reals.empty();
+        return ids && electron_mass > zero_quantity() && num_particles >= 2
+               && !pid_to_xs.empty() && !xs.empty() && !reals.empty();
     }
 
     //! Assign from another set of data
@@ -79,20 +85,12 @@ struct WentzelVIMscData
         ids = other.ids;
         electron_mass = other.electron_mass;
         params = other.params;
+        num_particles = other.num_particles;
+        pid_to_xs = other.pid_to_xs;
         xs = other.xs;
         reals = other.reals;
         return *this;
     }
-
-    //! Get the data location for a material + particle
-    CELER_FUNCTION ItemId<XsGridData> at(MaterialId mat, ParticleId par) const
-    {
-        CELER_EXPECT(mat && par);
-        size_type result = mat.unchecked_get() * 2;
-        result += (par == this->ids.electron ? 0 : 1);
-        CELER_ENSURE(result < this->xs.size());
-        return ItemId<XsGridData>{result};
-    }
 };
 
 }  // namespace celeritas

src/celeritas/em/msc/UrbanMsc.hh

@@ -81,8 +81,7 @@ UrbanMsc::is_applicable(CoreTrackView const& track, real_type step) const
         return false;
 
     auto par = track.make_particle_view();
-    if (par.particle_id() != shared_.ids.electron
-        && par.particle_id() != shared_.ids.positron)
+    if (!shared_.pid_to_xs[par.particle_id()])
         return false;
 
     return par.energy() > shared_.params.low_energy_limit
@@ -148,7 +147,7 @@ CELER_FUNCTION void UrbanMsc::limit_step(CoreTrackView const& track)
         // Calculate step limit using "safety" or "safety plus" algorithm
         detail::UrbanMscSafetyStepLimit calc_limit(shared_,
                                                    msc_helper,
-                                                   par.energy(),
+                                                   par,
                                                    &phys,
                                                    phys.material_id(),
                                                    geo.is_on_boundary(),

src/celeritas/em/msc/detail/UrbanMscHelper.hh

@@ -74,6 +74,12 @@ class UrbanMscHelper
     // The kinetic energy at the end of a given step length corrected by dedx
     inline CELER_FUNCTION Energy calc_end_energy(real_type step) const;
 
+    // Data for this particle+material
+    inline CELER_FUNCTION UrbanMscParMatData const& pmdata() const;
+
+    // Scaled cross section data for this particle+material
+    inline CELER_FUNCTION XsGridData const& xs() const;
+
   private:
     //// DATA ////
 
@@ -84,20 +90,6 @@ class UrbanMscHelper
 
     // Precalculated mean free path (TODO: move to physics step view)
     real_type lambda_;  // [len]
-
-    // Data for this particle+material
-    CELER_FUNCTION UrbanMscParMatData const& pmdata() const
-    {
-        return shared_.par_mat_data[shared_.at<UrbanMscParMatData>(
-            physics_.material_id(), particle_.particle_id())];
-    }
-
-    // Scaled cross section data for this particle+material
-    CELER_FUNCTION XsGridData const& xs() const
-    {
-        return shared_.xs[shared_.at<XsGridData>(physics_.material_id(),
-                                                 particle_.particle_id())];
-    }
 };
 
 //---------------------------------------------------------------------------//
@@ -115,8 +107,6 @@ UrbanMscHelper::UrbanMscHelper(UrbanMscRef const& shared,
     , physics_(physics)
     , lambda_(this->calc_msc_mfp(particle_.energy()))
 {
-    CELER_EXPECT(particle.particle_id() == shared_.ids.electron
-                 || particle.particle_id() == shared_.ids.positron);
 }
 
 //---------------------------------------------------------------------------//
@@ -181,6 +171,38 @@ UrbanMscHelper::calc_end_energy(real_type step) const -> Energy
     }
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Scaled cross section data for this particle+material.
+ */
+CELER_FUNCTION XsGridData const& UrbanMscHelper::xs() const
+{
+    auto par_id = shared_.pid_to_xs[particle_.particle_id()];
+    CELER_ASSERT(par_id < shared_.num_particles);
+
+    size_type idx = physics_.material_id().get() * shared_.num_particles
+                    + par_id.unchecked_get();
+    CELER_ASSERT(idx < shared_.xs.size());
+
+    return shared_.xs[ItemId<XsGridData>(idx)];
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Data for this particle+material.
+ */
+CELER_FUNCTION UrbanMscParMatData const& UrbanMscHelper::pmdata() const
+{
+    auto par_id = shared_.pid_to_pmdata[particle_.particle_id()];
+    CELER_ASSERT(par_id < shared_.num_par_mat);
+
+    size_type idx = physics_.material_id().get() * shared_.num_par_mat
+                    + par_id.unchecked_get();
+    CELER_ASSERT(idx < shared_.par_mat_data.size());
+
+    return shared_.par_mat_data[ItemId<UrbanMscParMatData>(idx)];
+}
+
 //---------------------------------------------------------------------------//
 }  // namespace detail
 }  // namespace celeritas