Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-09-16 08:52:28

0001 //------------------------------- -*- C++ -*- -------------------------------//
0002 // Copyright Celeritas contributors: see top-level COPYRIGHT file for details
0003 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0004 //---------------------------------------------------------------------------//
0005 //! \file celeritas/user/StepData.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008 
0009 #include "corecel/Macros.hh"
0010 #include "corecel/cont/EnumArray.hh"
0011 #include "corecel/cont/Range.hh"
0012 #include "corecel/data/Collection.hh"
0013 #include "corecel/data/CollectionBuilder.hh"
0014 #include "celeritas/Quantities.hh"
0015 #include "celeritas/Types.hh"
0016 #include "celeritas/Units.hh"
0017 
0018 namespace celeritas
0019 {
0020 //---------------------------------------------------------------------------//
0021 // TYPES
0022 //---------------------------------------------------------------------------//
0023 /*!
0024  * Which track properties to gather at the beginning and end of a step.
0025  *
0026  * These should all default to *false* so that this list can be extended
0027  * without adversely affecting existing interfaces.
0028  */
0029 struct StepPointSelection
0030 {
0031     bool time{false};
0032     bool pos{false};
0033     bool dir{false};
0034     bool energy{false};
0035 
0036     bool volume_id{false};
0037     bool volume_instance_ids{false};
0038 
0039     //! Create StepPointSelection with all options set to true
0040     static constexpr StepPointSelection all()
0041     {
0042         return StepPointSelection{true, true, true, true, true, true};
0043     }
0044 
0045     //! Whether any selection is requested
0046     explicit CELER_FUNCTION operator bool() const
0047     {
0048         return time || pos || dir || energy || volume_id || volume_instance_ids;
0049     }
0050 
0051     //! Combine the selection with another
0052     StepPointSelection& operator|=(StepPointSelection const& other)
0053     {
0054         this->time |= other.time;
0055         this->pos |= other.pos;
0056         this->dir |= other.dir;
0057         this->energy |= other.energy;
0058         this->volume_id |= other.volume_id;
0059         this->volume_instance_ids |= other.volume_instance_ids;
0060         return *this;
0061     }
0062 };
0063 
0064 //---------------------------------------------------------------------------//
0065 /*!
0066  * Which track properties to gather at every step.
0067  *
0068  * These should correspond to the data items in StepStateDataImpl.
0069  *
0070  * TODO: particle -> particle_id for consistency
0071  */
0072 struct StepSelection
0073 {
0074     EnumArray<StepPoint, StepPointSelection> points;
0075 
0076     bool event_id{false};
0077     bool parent_id{false};
0078     bool track_step_count{false};
0079     bool action_id{false};
0080     bool step_length{false};
0081     bool particle{false};
0082     bool energy_deposition{false};
0083 
0084     //! Create StepSelection with all options set to true
0085     static constexpr StepSelection all()
0086     {
0087         return StepSelection{
0088             {StepPointSelection::all(), StepPointSelection::all()},
0089             true,
0090             true,
0091             true,
0092             true,
0093             true,
0094             true,
0095             true};
0096     }
0097 
0098     //! Whether any selection is requested
0099     explicit CELER_FUNCTION operator bool() const
0100     {
0101         return points[StepPoint::pre] || points[StepPoint::post] || event_id
0102                || parent_id || track_step_count || action_id || step_length
0103                || particle || energy_deposition;
0104     }
0105 
0106     //! Combine the selection with another
0107     StepSelection& operator|=(StepSelection const& other)
0108     {
0109         for (auto sp : range(StepPoint::size_))
0110         {
0111             points[sp] |= other.points[sp];
0112         }
0113 
0114         this->event_id |= other.event_id;
0115         this->parent_id |= other.parent_id;
0116         this->track_step_count |= other.track_step_count;
0117         this->action_id |= other.action_id;
0118         this->step_length |= other.step_length;
0119         this->particle |= other.particle;
0120         this->energy_deposition |= other.energy_deposition;
0121         return *this;
0122     }
0123 };
0124 
0125 //---------------------------------------------------------------------------//
0126 /*!
0127  * Shared attributes about the hits being collected.
0128  *
0129  * This will be expanded to include filters for particle type, region, etc.
0130  */
0131 template<Ownership W, MemSpace M>
0132 struct StepParamsData
0133 {
0134     //// DATA ////
0135 
0136     //! Options for gathering data at each step
0137     StepSelection selection;
0138 
0139     //! Optional mapping for volume -> sensitive detector
0140     Collection<DetectorId, W, M, VolumeId> detector;
0141 
0142     //! Filter out steps that have not deposited energy (for sensitive det)
0143     bool nonzero_energy_deposition{false};
0144 
0145     //! Per-state volume instance size if volume_instance_ids selected
0146     size_type volume_instance_depth{0};
0147 
0148     //// METHODS ////
0149 
0150     //! Whether the data is assigned
0151     explicit CELER_FUNCTION operator bool() const
0152     {
0153         return static_cast<bool>(selection);
0154     }
0155 
0156     //! Assign from another set of data
0157     template<Ownership W2, MemSpace M2>
0158     StepParamsData& operator=(StepParamsData<W2, M2> const& other)
0159     {
0160         CELER_EXPECT(other);
0161         selection = other.selection;
0162         detector = other.detector;
0163         nonzero_energy_deposition = other.nonzero_energy_deposition;
0164         volume_instance_depth = other.volume_instance_depth;
0165         return *this;
0166     }
0167 };
0168 
0169 //---------------------------------------------------------------------------//
0170 /*!
0171  * Gathered state data for beginning/end of step data for tracks in parallel.
0172  *
0173  * - Each data member corresponds exactly to a flag in \c StepPointSelection
0174  * - If the flag is disabled (no step interfaces require the data), then the
0175  *   corresponding member data will be empty.
0176  * - If a track is outside the volume (which can only happen at the end-of-step
0177  *   evaluation) the VolumeId will be "false".
0178  */
0179 template<Ownership W, MemSpace M>
0180 struct StepPointStateData
0181 {
0182     //// TYPES ////
0183 
0184     template<class T>
0185     using StateItems = celeritas::StateCollection<T, W, M>;
0186     template<class T>
0187     using Items = celeritas::Collection<T, W, M>;
0188     using Energy = units::MevEnergy;
0189 
0190     // Sim
0191     StateItems<real_type> time;
0192 
0193     // Geo
0194     StateItems<Real3> pos;
0195     StateItems<Real3> dir;
0196     StateItems<VolumeId> volume_id;
0197     Items<VolumeInstanceId> volume_instance_ids;
0198 
0199     // Physics
0200     StateItems<Energy> energy;
0201 
0202     //// METHODS ////
0203 
0204     //! Always true since all step-point data could be disabled
0205     explicit CELER_FUNCTION operator bool() const { return true; }
0206 
0207     //! Assign from another set of states
0208     template<Ownership W2, MemSpace M2>
0209     StepPointStateData& operator=(StepPointStateData<W2, M2>& other)
0210     {
0211         CELER_EXPECT(other);
0212         time = other.time;
0213         pos = other.pos;
0214         dir = other.dir;
0215         volume_id = other.volume_id;
0216         volume_instance_ids = other.volume_instance_ids;
0217         energy = other.energy;
0218         return *this;
0219     }
0220 };
0221 
0222 //---------------------------------------------------------------------------//
0223 /*!
0224  * Gathered data for a single step for many tracks in parallel.
0225  *
0226  * - Each data member corresponds exactly to a flag in \c StepSelection .
0227  * - If the flag is disabled (no step interfaces require the data), then the
0228  *   corresponding member data will be empty.
0229  * - The track ID will be set to "false" if the track is inactive.
0230  * - If sensitive detector are specified, the \c detector field is set based
0231  *   on the pre-step geometric volume. Data members will have \b unspecified
0232  *   values if the detector ID is "false" (i.e. no information is being
0233  *   collected). The detector ID for inactive threads is always "false".
0234  */
0235 template<Ownership W, MemSpace M>
0236 struct StepStateDataImpl
0237 {
0238     //// TYPES ////
0239 
0240     using StepPointData = StepPointStateData<W, M>;
0241     template<class T>
0242     using StateItems = celeritas::StateCollection<T, W, M>;
0243     using Energy = units::MevEnergy;
0244 
0245     //// DATA ////
0246 
0247     // Pre- and post-step data
0248     EnumArray<StepPoint, StepPointData> points;
0249 
0250     //! Track ID is always assigned (but will be false for inactive tracks)
0251     StateItems<TrackId> track_id;
0252 
0253     //! Detector ID is non-empty if params.detector is nonempty
0254     StateItems<DetectorId> detector;
0255 
0256     // Sim
0257     StateItems<EventId> event_id;
0258     StateItems<TrackId> parent_id;
0259     StateItems<ActionId> action_id;
0260     StateItems<size_type> track_step_count;
0261     StateItems<real_type> step_length;
0262 
0263     // Physics
0264     StateItems<ParticleId> particle;
0265     StateItems<Energy> energy_deposition;
0266 
0267     //// METHODS ////
0268 
0269     //! True if constructed and correctly sized
0270     explicit CELER_FUNCTION operator bool() const
0271     {
0272         auto right_sized = [this](auto const& t) {
0273             return (t.size() == this->size()) || t.empty();
0274         };
0275 
0276         return !track_id.empty() && right_sized(detector)
0277                && right_sized(event_id) && right_sized(parent_id)
0278                && right_sized(track_step_count) && right_sized(action_id)
0279                && right_sized(step_length) && right_sized(particle)
0280                && right_sized(energy_deposition);
0281     }
0282 
0283     //! State size
0284     CELER_FUNCTION TrackSlotId::size_type size() const
0285     {
0286         return track_id.size();
0287     }
0288 
0289     //! Assign from another set of states
0290     template<Ownership W2, MemSpace M2>
0291     StepStateDataImpl& operator=(StepStateDataImpl<W2, M2>& other)
0292     {
0293         // The extra storage used to gather the step data is only required on
0294         // the device
0295         CELER_EXPECT(other || (M == MemSpace::host && other.size() == 0));
0296 
0297         for (auto sp : range(StepPoint::size_))
0298         {
0299             points[sp] = other.points[sp];
0300         }
0301 
0302         track_id = other.track_id;
0303         parent_id = other.parent_id;
0304         detector = other.detector;
0305         event_id = other.event_id;
0306         track_step_count = other.track_step_count;
0307         action_id = other.action_id;
0308         step_length = other.step_length;
0309         particle = other.particle;
0310         energy_deposition = other.energy_deposition;
0311         return *this;
0312     }
0313 };
0314 
0315 //---------------------------------------------------------------------------//
0316 /*!
0317  * Gathered data and persistent scratch space for gathering and copying data.
0318  *
0319  * Extra storage \c scratch and \c valid_id is needed to efficiently gather and
0320  * copy the step data on the device but will not be allocated on the host.
0321  *
0322  * \todo Refactor the step interface stuff so that we passing params alongside
0323  * state data to the step interfaces, so that we don't have to keep a copy of
0324  * the volume instance depth (and for better consistency with the rest of the
0325  * code).
0326  */
0327 template<Ownership W, MemSpace M>
0328 struct StepStateData
0329 {
0330     //// TYPES ////
0331 
0332     using StepDataImpl = StepStateDataImpl<W, M>;
0333     template<class T>
0334     using StateItems = celeritas::StateCollection<T, W, M>;
0335 
0336     //// DATA ////
0337 
0338     //! Gathered data for a single step
0339     StepDataImpl data;
0340 
0341     //! Scratch space for gathering the data on device based on track validity
0342     StepDataImpl scratch;
0343 
0344     //! Thread IDs of active tracks that are in a detector
0345     StateItems<size_type> valid_id;
0346 
0347     // Copy of params max depth for dimensioning volume_instance_ids
0348     size_type volume_instance_depth{0};
0349 
0350     //! Unique identifier for "thread-local" data.
0351     StreamId stream_id;
0352 
0353     //// METHODS ////
0354 
0355     //! True if constructed and correctly sized
0356     explicit CELER_FUNCTION operator bool() const
0357     {
0358         auto right_sized = [this](auto const& t) {
0359             return (t.size() == this->size())
0360                    || (t.size() == 0 && M == MemSpace::host);
0361         };
0362 
0363         return data.size() > 0 && right_sized(scratch) && right_sized(valid_id)
0364                && stream_id;
0365     }
0366 
0367     //! State size
0368     CELER_FUNCTION TrackSlotId::size_type size() const { return data.size(); }
0369 
0370     //! Assign from another set of states
0371     template<Ownership W2, MemSpace M2>
0372     StepStateData& operator=(StepStateData<W2, M2>& other)
0373     {
0374         CELER_EXPECT(other);
0375 
0376         data = other.data;
0377         scratch = other.scratch;
0378         valid_id = other.valid_id;
0379         volume_instance_depth = other.volume_instance_depth;
0380         stream_id = other.stream_id;
0381         return *this;
0382     }
0383 };
0384 
0385 //---------------------------------------------------------------------------//
0386 // HELPER FUNCTIONS
0387 //---------------------------------------------------------------------------//
0388 /*!
0389  * Resize a state point.
0390  */
0391 template<MemSpace M>
0392 inline void resize(StepPointStateData<Ownership::value, M>* state,
0393                    StepPointSelection selection,
0394                    HostCRef<StepParamsData> const& params,
0395                    size_type size)
0396 {
0397     CELER_EXPECT(size > 0);
0398 #define SD_RESIZE_IF_SELECTED(ATTR)     \
0399     do                                  \
0400     {                                   \
0401         if (selection.ATTR)             \
0402         {                               \
0403             resize(&state->ATTR, size); \
0404         }                               \
0405     } while (0)
0406 
0407     SD_RESIZE_IF_SELECTED(time);
0408     SD_RESIZE_IF_SELECTED(pos);
0409     SD_RESIZE_IF_SELECTED(dir);
0410     SD_RESIZE_IF_SELECTED(volume_id);
0411     if (selection.volume_instance_ids)
0412     {
0413         size_type const vi_depth = params.volume_instance_depth;
0414         CELER_ASSERT(vi_depth > 0);
0415         resize(&state->volume_instance_ids, size * vi_depth);
0416     }
0417     SD_RESIZE_IF_SELECTED(energy);
0418 
0419 #undef SD_RESIZE_IF_SELECTED
0420 }
0421 
0422 //---------------------------------------------------------------------------//
0423 /*!
0424  * Resize the step data.
0425  */
0426 template<MemSpace M>
0427 inline void resize(StepStateDataImpl<Ownership::value, M>* state,
0428                    HostCRef<StepParamsData> const& params,
0429                    size_type size)
0430 {
0431     CELER_EXPECT(state->size() == 0);
0432     CELER_EXPECT(size > 0);
0433 
0434     for (auto sp : range(StepPoint::size_))
0435     {
0436         resize(&state->points[sp], params.selection.points[sp], params, size);
0437     }
0438 
0439 #define SD_RESIZE_IF_SELECTED(ATTR)     \
0440     do                                  \
0441     {                                   \
0442         if (params.selection.ATTR)      \
0443         {                               \
0444             resize(&state->ATTR, size); \
0445         }                               \
0446     } while (0)
0447 
0448     resize(&state->track_id, size);
0449     if (!params.detector.empty())
0450     {
0451         resize(&state->detector, size);
0452     }
0453 
0454     SD_RESIZE_IF_SELECTED(event_id);
0455     SD_RESIZE_IF_SELECTED(parent_id);
0456     SD_RESIZE_IF_SELECTED(track_step_count);
0457     SD_RESIZE_IF_SELECTED(step_length);
0458     SD_RESIZE_IF_SELECTED(action_id);
0459     SD_RESIZE_IF_SELECTED(particle);
0460     SD_RESIZE_IF_SELECTED(energy_deposition);
0461 }
0462 
0463 //---------------------------------------------------------------------------//
0464 /*!
0465  * Resize the state.
0466  */
0467 template<MemSpace M>
0468 inline void resize(StepStateData<Ownership::value, M>* state,
0469                    HostCRef<StepParamsData> const& params,
0470                    StreamId stream_id,
0471                    size_type size)
0472 {
0473     CELER_EXPECT(state->size() == 0);
0474     CELER_EXPECT(size > 0);
0475 
0476     state->volume_instance_depth = params.volume_instance_depth;
0477     state->stream_id = stream_id;
0478 
0479     resize(&state->data, params, size);
0480 
0481     if constexpr (M == MemSpace::device)
0482     {
0483         // Allocate extra space on device for gathering step data
0484         resize(&state->scratch, params, size);
0485         resize(&state->valid_id, size);
0486     }
0487 }
0488 
0489 //---------------------------------------------------------------------------//
0490 }  // namespace celeritas