geocel/rasterize/Raytracer.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2024 UT-Battelle, LLC, and other Celeritas developers.
0003 // See the top-level COPYRIGHT file for details.
0004 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0005 //---------------------------------------------------------------------------//
0006 //! \file geocel/rasterize/Raytracer.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include "corecel/math/Algorithms.hh"
0011
0012 #include "ImageLineView.hh"
0013
0014 namespace celeritas
0015 {
0016 //---------------------------------------------------------------------------//
0017 /*!
0018  * Trace each pixel along a line.
0019  *
0020  * \tparam GTV Geometry Track View
0021  * \tparam F   Calculate the pixel value (result_type) given GTV
0022  */
0023 template<class GTV, class F>
0024 class Raytracer
0025 {
0026   public:
0027     //!@{
0028     //! \name Type aliases
0029     using result_type = int;
0030     //!@}
0031
0032   public:
0033     // Construct from geo track view, function to calculate ID, and image
0034     inline CELER_FUNCTION
0035     Raytracer(GTV&& geo, F&& calc_id, ImageLineView const& image);
0036
0037     // Calculate the value for the next pixel
0038     inline CELER_FUNCTION result_type operator()(size_type pix);
0039
0040   private:
0041     //// DATA ////
0042
0043     GTV geo_;
0044     F calc_id_;
0045     ImageLineView const& image_;
0046
0047     size_type pixel_{invalid_pixel()};  //!< Current pixel
0048     real_type distance_{0};  //!< Distance to next boundary
0049     int cur_id_{-1};  //!< Current ID
0050
0051     //// HELPER FUNCTIONS ////
0052
0053     // Initialize the geometry at this pixel index
0054     inline CELER_FUNCTION void initialize_at_pixel(size_type pix);
0055
0056     // Find the next step
0057     inline CELER_FUNCTION void find_next_step();
0058
0059     //! Sentinel value for needing the pixel to be reset
0060     static CELER_CONSTEXPR_FUNCTION size_type invalid_pixel()
0061     {
0062         return static_cast<size_type>(-1);
0063     }
0064
0065     //! Stop tracking the pixel if this many crossings occur
0066     static CELER_CONSTEXPR_FUNCTION short int max_crossings_per_pixel()
0067     {
0068         return 32;
0069     }
0070 };
0071
0072 //---------------------------------------------------------------------------//
0073 // DEDUCTION GUIDES
0074 //---------------------------------------------------------------------------//
0075
0076 template<class GTV, class F>
0077 CELER_FUNCTION Raytracer(GTV&&, F&&, ImageLineView const&) -> Raytracer<GTV, F>;
0078
0079 //---------------------------------------------------------------------------//
0080 // INLINE DEFINITIONS
0081 //---------------------------------------------------------------------------//
0082 /*!
0083  * Construct from geo track view, function to calculate ID, and image.
0084  */
0085 template<class GTV, class F>
0086 CELER_FUNCTION
0087 Raytracer<GTV, F>::Raytracer(GTV&& geo, F&& calc_id, ImageLineView const& image)
0088     : geo_{celeritas::forward<GTV>(geo)}
0089     , calc_id_{celeritas::forward<F>(calc_id)}
0090     , image_{image}
0091 {
0092 }
0093
0094 //---------------------------------------------------------------------------//
0095 /*!
0096  * Calculate the value for a single pixel.
0097  */
0098 template<class GTV, class F>
0099 CELER_FUNCTION auto Raytracer<GTV, F>::operator()(size_type pix) -> result_type
0100 {
0101     if (pix != pixel_)
0102     {
0103         this->initialize_at_pixel(pix);
0104         if (pixel_ == invalid_pixel())
0105         {
0106             // Pixel starts outside the geometry
0107             return cur_id_;
0108         }
0109     }
0110
0111     // Remaining distance to the end of the pixel
0112     real_type pix_dist = image_.pixel_width();
0113     // Distance associated with the largest segment so far
0114     real_type max_dist = 0;
0115     int max_id = cur_id_;
0116     // Countdown to limit stuck tracks or overdetailed pixels
0117     auto abort_counter = this->max_crossings_per_pixel();
0118
0119     while (cur_id_ >= 0 && distance_ < pix_dist)
0120     {
0121         // Move to geometry boundary
0122         if (max_id == cur_id_)
0123         {
0124             max_dist += distance_;
0125         }
0126         else if (distance_ > max_dist)
0127         {
0128             max_dist = distance_;
0129             max_id = cur_id_;
0130         }
0131
0132         // Update pixel and boundary distance
0133         pix_dist -= distance_;
0134         distance_ = 0;
0135
0136         // Cross surface and update post-crossing ID
0137         geo_.move_to_boundary();
0138         geo_.cross_boundary();
0139
0140         if (--abort_counter == 0)
0141         {
0142             // Give up and move to end of pixel, reinitialize on the next call
0143             pix_dist = 0;
0144         }
0145         if (pix_dist > 0)
0146         {
0147             // Update next distance and current ID for the remaining pixel
0148             this->find_next_step();
0149         }
0150     }
0151
0152     if (pix_dist > 0)
0153     {
0154         // Move to pixel boundary
0155         distance_ -= pix_dist;
0156         if (pix_dist > max_dist)
0157         {
0158             max_dist = pix_dist;
0159             max_id = cur_id_;
0160         }
0161     }
0162
0163     if (abort_counter != 0)
0164     {
0165         ++pixel_;
0166     }
0167     else
0168     {
0169         pixel_ = invalid_pixel();
0170     }
0171
0172     return max_id;
0173 }
0174
0175 //---------------------------------------------------------------------------//
0176 /*!
0177  * Initialize the geometry at the given pixel position.
0178  */
0179 template<class GTV, class F>
0180 CELER_FUNCTION void Raytracer<GTV, F>::initialize_at_pixel(size_type pix)
0181 {
0182     CELER_EXPECT(pix < image_.max_index());
0183
0184     GeoTrackInitializer init{image_.start_pos(), image_.start_dir()};
0185     axpy(pix * image_.pixel_width(), init.dir, &init.pos);
0186     geo_ = init;
0187     pixel_ = pix;
0188     distance_ = 0;
0189     this->find_next_step();
0190 }
0191
0192 //---------------------------------------------------------------------------//
0193 /*!
0194  * Find the distance to the next step.
0195  */
0196 template<class GTV, class F>
0197 CELER_FUNCTION void Raytracer<GTV, F>::find_next_step()
0198 {
0199     CELER_EXPECT(pixel_ != invalid_pixel());
0200     CELER_EXPECT(distance_ <= 0);
0201     if (geo_.is_outside())
0202     {
0203         // Skip this pixel since not all navigation engines can trace from
0204         // outside the geometry
0205         pixel_ = invalid_pixel();
0206         distance_ = numeric_limits<real_type>::infinity();
0207     }
0208     else
0209     {
0210         // Search to distances just past the end of the saved pixels (so the
0211         // last point isn't coincident with a boundary)
0212         distance_ = geo_.find_next_step((image_.max_index() + 1 - pixel_)
0213                                         * image_.pixel_width())
0214                         .distance;
0215     }
0216     cur_id_ = this->calc_id_(geo_);
0217
0218     CELER_ENSURE(distance_ > 0);
0219 }
0220
0221 //---------------------------------------------------------------------------//
0222 }  // namespace celeritas