Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-11-01 10:03:48

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