utils/grid/grid.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 // Project include(s).
0012 #include "detray/core/detail/container_views.hpp"
0013 #include "detray/definitions/detail/qualifiers.hpp"
0014 #include "detray/utils/grid/axis.hpp"
0015 #include "detray/utils/grid/detail/axis_helpers.hpp"
0016 #include "detray/utils/grid/detail/bin_storage.hpp"
0017 #include "detray/utils/grid/populators.hpp"
0018 #include "detray/utils/grid/serializers.hpp"
0019 #include "detray/utils/ranges.hpp"
0020
0021 // VecMem include(s).
0022 #include <vecmem/memory/memory_resource.hpp>
0023
0024 // System include(s).
0025 #include <cstddef>
0026 #include <type_traits>
0027
0028 namespace detray {
0029
0030 /// @brief An N-dimensional grid for object storage.
0031 ///
0032 /// @tparam axes_t the types of grid axes
0033 /// @tparam bin_t type of bin in the (global) bin storage.
0034 /// @tparam serializer_t how to serialize axis-local bin indices into global bin
0035 ///                      indices in the grid backend storage and vice versa.
0036 template <typename axes_t, typename bin_t,
0037           template <std::size_t> class serializer_t = simple_serializer>
0038 class grid_impl {
0039  public:
0040   /// Grid dimension
0041   static constexpr unsigned int dim = axes_t::dim;
0042
0043   /// Single value in a bin entry
0044   using bin_type = bin_t;
0045   using value_type = typename bin_type::entry_type;
0046
0047   template <std::size_t DIM>
0048   using serializer_type = serializer_t<DIM>;
0049
0050   /// The type of the multi-axis is tied to the type of the grid: a non-
0051   /// owning grid holds a non-owning multi-axis member.
0052   using axes_type = axes_t;
0053   using glob_bin_index = dindex;
0054   using loc_bin_index = typename axes_type::loc_bin_index;
0055   using local_frame_type = typename axes_type::local_frame_type;
0056   using algebra_type = typename axes_type::algebra_type;
0057   using scalar_type = dscalar<algebra_type>;
0058   using point_type = typename axes_type::point_type;
0059
0060   static constexpr bool is_owning{axes_type::is_owning};
0061
0062   /// Backend storage type for the grid
0063   using bin_storage =
0064       detray::detail::bin_storage<is_owning, bin_type,
0065                                   typename axes_type::container_types>;
0066   using bin_container_type = typename bin_storage::bin_container_type;
0067
0068   /// Vecmem based grid view type
0069   using view_type = dmulti_view<typename bin_storage::view_type,
0070                                 typename axes_type::view_type>;
0071   /// Vecmem based grid view type - const
0072   using const_view_type = dmulti_view<typename bin_storage::const_view_type,
0073                                       typename axes_type::const_view_type>;
0074   /// Vecmem based buffer type
0075   using buffer_type = dmulti_buffer<typename bin_storage::buffer_type,
0076                                     typename axes_type::buffer_type>;
0077
0078   /// Find the corresponding (non-)owning grid type
0079   template <bool owning>
0080   using type =
0081       grid_impl<typename axes_t::template type<owning>, bin_t, serializer_t>;
0082
0083   /// Make grid default constructible: Empty grid with empty axis
0084   grid_impl() = default;
0085
0086   /// Create empty grid with empty axes from specific vecmem memory resource
0087   DETRAY_HOST
0088   explicit grid_impl(vecmem::memory_resource &resource)
0089       : m_bins(resource), m_axes(resource) {}
0090
0091   /// Create grid with well defined @param axes and @param bins_data - move
0092   DETRAY_HOST_DEVICE
0093   grid_impl(bin_container_type &&bin_data, axes_type &&axes)
0094       : m_bins(std::move(bin_data)), m_axes(std::move(axes)) {}
0095
0096   /// Create grid from container pointers - non-owning (both grid and axes)
0097   DETRAY_HOST_DEVICE
0098   grid_impl(const bin_container_type *bin_data_ptr, const axes_type &axes,
0099             const dindex offset = 0)
0100       : m_bins(*bin_data_ptr, offset, axes.nbins()), m_axes(axes) {}
0101
0102   /// Create grid from container pointers - non-owning (both grid and axes)
0103   DETRAY_HOST_DEVICE
0104   grid_impl(bin_container_type *bin_data_ptr, axes_type &axes,
0105             const dindex offset = 0u)
0106       : m_bins(*bin_data_ptr, offset, axes.nbins()), m_axes(axes) {}
0107
0108   /// Create grid from container pointers - non-owning (both grid and axes)
0109   // TODO: const correctness
0110   DETRAY_HOST_DEVICE
0111   grid_impl(const bin_container_type *bin_data_ptr, axes_type &&axes,
0112             const dindex offset = 0)
0113       : m_bins(*(const_cast<bin_container_type *>(bin_data_ptr)), offset,
0114                axes.nbins()),
0115         m_axes(std::move(axes)) {}
0116
0117   /// Create grid from container pointers - non-owning (both grid and axes)
0118   DETRAY_HOST_DEVICE
0119   grid_impl(bin_container_type *bin_data_ptr, axes_type &&axes,
0120             const dindex offset = 0u)
0121       : m_bins(*bin_data_ptr, offset, axes.nbins()), m_axes(std::move(axes)) {}
0122
0123   /// Device-side construction from a vecmem based view type
0124   template <concepts::device_view grid_view_t>
0125   DETRAY_HOST_DEVICE explicit grid_impl(grid_view_t &view)
0126       : m_bins(detray::detail::get<0>(view.m_view)),
0127         m_axes(detray::detail::get<1>(view.m_view)) {}
0128
0129   /// @returns the multi-axis used by the grid - const
0130   DETRAY_HOST_DEVICE
0131   auto axes() const -> const axes_type & { return m_axes; }
0132
0133   /// @returns the grid local coordinate system
0134   DETRAY_HOST_DEVICE
0135   static constexpr auto get_local_frame() -> local_frame_type { return {}; }
0136
0137   /// Transform a point in global cartesian coordinates to bound coordinates
0138   ///
0139   /// @param trf the placement transform of the grid (e.g. from a volume or
0140   ///            a surface).
0141   /// @param p   the point in global coordinates
0142   /// @param d   direction of a track at position p
0143   ///
0144   /// @returns a point in the coordinate system that is spanned by the grid's
0145   /// axes.
0146   template <concepts::transform3D transform3_t, concepts::point3D point3_t,
0147             concepts::vector3D vector3_t>
0148   DETRAY_HOST_DEVICE point_type project(const transform3_t &trf,
0149                                         const point3_t &p,
0150                                         const vector3_t &d) const {
0151     return grid_impl::get_local_frame().global_to_local(trf, p, d);
0152   }
0153
0154   /// @returns an axis object, corresponding to the index.
0155   template <std::size_t index>
0156   DETRAY_HOST_DEVICE constexpr auto get_axis() const {
0157     return m_axes.template get_axis<index>();
0158   }
0159
0160   /// @returns an axis object, corresponding to the label.
0161   template <axis::label L>
0162   DETRAY_HOST_DEVICE constexpr auto get_axis() const {
0163     return m_axes.template get_axis<L>();
0164   }
0165
0166   /// @returns an axis object of the given type.
0167   template <typename axis_t>
0168   DETRAY_HOST_DEVICE constexpr axis_t get_axis() const {
0169     return m_axes.template get_axis<axis_t>();
0170   }
0171
0172   /// @returns the total number of bins in the grid
0173   DETRAY_HOST_DEVICE constexpr auto nbins() const -> dindex {
0174     return m_axes.nbins();
0175   }
0176
0177   /// @returns the total number of values in the grid
0178   /// @note this has to query every bin for the number of elements
0179   DETRAY_HOST_DEVICE constexpr auto size() const -> dindex {
0180     return static_cast<dindex>(all().size());
0181   }
0182
0183   /// @returns an instance of the grid serializer
0184   static constexpr auto serializer() -> serializer_t<dim> { return {}; }
0185
0186   /// @returns a local multi-bin index from a global bin index @param gid
0187   constexpr auto deserialize(const glob_bin_index gid) const -> loc_bin_index {
0188     return serializer()(axes(), gid);
0189   }
0190
0191   /// @returns a global bin index from a local bin index @param mbin
0192   constexpr auto serialize(const loc_bin_index &mbin) const -> glob_bin_index {
0193     return serializer()(axes(), mbin);
0194   }
0195
0196   /// @returns the full range of bins - const
0197   DETRAY_HOST_DEVICE
0198   auto bins() const -> const bin_storage & { return m_bins; }
0199
0200   /// @returns the full range of bins
0201   DETRAY_HOST_DEVICE
0202   auto bins() -> bin_storage & { return m_bins; }
0203
0204   /// @returns the iterable view of the bin content
0205   /// @{
0206   /// @param gbin the global bin index - const
0207   DETRAY_HOST_DEVICE
0208   decltype(auto) bin(const glob_bin_index gbin) const { return m_bins[gbin]; }
0209
0210   /// @param gbin the global bin index
0211   DETRAY_HOST_DEVICE
0212   decltype(auto) bin(const glob_bin_index gbin) { return m_bins[gbin]; }
0213
0214   /// @param mbin the multi-index of bins over all axes - const
0215   DETRAY_HOST_DEVICE
0216   decltype(auto) bin(const loc_bin_index &mbin) const {
0217     return bin(serialize(mbin));
0218   }
0219
0220   /// @param mbin the multi-index of bins over all axes
0221   DETRAY_HOST_DEVICE
0222   decltype(auto) bin(const loc_bin_index &mbin) { return bin(serialize(mbin)); }
0223
0224   /// Find the value of a single bin - const
0225   ///
0226   /// @param p is point in the local (bound) frame
0227   ///
0228   /// @return the iterable view of the bin content
0229   DETRAY_HOST_DEVICE decltype(auto) bin(const point_type &p) const {
0230     return bin(m_axes.bins(p));
0231   }
0232
0233   /// Find the value of a single bin
0234   ///
0235   /// @param p is point in the local (bound) frame
0236   ///
0237   /// @return the iterable view of the bin content
0238   DETRAY_HOST_DEVICE decltype(auto) bin(const point_type &p) {
0239     return bin(m_axes.bins(p));
0240   }
0241
0242   /// @param indices the single indices corresponding to a multi_bin
0243   template <concepts::index... I>
0244     requires(sizeof...(I) == dim)
0245   DETRAY_HOST_DEVICE decltype(auto) bin(I... indices) const {
0246     return bin(loc_bin_index{indices...});
0247   }
0248   /// @}
0249
0250   /// Access a single entry in a bin from the global bin index, as well as
0251   /// the index of the entry in the bin.
0252   ///
0253   /// @param idx the index of a specific grid entry
0254   ///
0255   /// @returns a single bin entry.
0256   /// @{
0257   DETRAY_HOST_DEVICE
0258   decltype(auto) at(const glob_bin_index gbin, const dindex idx) {
0259     return bin(gbin)[idx];
0260   }
0261   DETRAY_HOST_DEVICE
0262   decltype(auto) at(const glob_bin_index gbin, const dindex idx) const {
0263     return bin(gbin)[idx];
0264   }
0265   DETRAY_HOST_DEVICE
0266   decltype(auto) at(const loc_bin_index &mbin, const dindex idx) {
0267     return bin(mbin)[idx];
0268   }
0269   DETRAY_HOST_DEVICE
0270   decltype(auto) at(const loc_bin_index &mbin, const dindex idx) const {
0271     return bin(mbin)[idx];
0272   }
0273   /// @}
0274
0275   /// @returns a view over the flattened bin content by joining the bin ranges
0276   DETRAY_HOST_DEVICE auto all() { return detray::views::join(bins()); }
0277
0278   /// @returns a view over the flattened bin content by joining the bin ranges
0279   DETRAY_HOST_DEVICE auto all() const { return detray::views::join(bins()); }
0280
0281   /// Poupulate a bin with a single one of its corresponding values @param v
0282   /// @{
0283   /// @param mbin the multi bin index to be populated
0284   template <typename populator_t, typename V = value_type>
0285   DETRAY_HOST_DEVICE void populate(const loc_bin_index &mbin, V &&v) {
0286     populator_t{}(bin(mbin), std::forward<V>(v));
0287   }
0288
0289   /// @param gbin the global bin index to be populated
0290   template <typename populator_t, typename V = value_type>
0291   DETRAY_HOST_DEVICE void populate(const glob_bin_index gbin, V &&v) {
0292     populator_t{}(bin(gbin), std::forward<V>(v));
0293   }
0294
0295   /// @param p the point in local coordinates that defines the bin to be
0296   ///          populated
0297   template <typename populator_t, typename V = value_type>
0298   DETRAY_HOST_DEVICE void populate(const point_type &p, V &&v) {
0299     populator_t{}(bin(m_axes.bins(p)), std::forward<V>(v));
0300   }
0301   /// @}
0302
0303   /// @returns view of a grid, including the grids multi_axis. Also valid if
0304   /// the value type of the grid is cv qualified (then value_t propagates
0305   /// quialifiers) - non-const
0306   template <bool owning = is_owning>
0307     requires owning
0308   DETRAY_HOST auto get_data() -> view_type {
0309     return view_type{detray::get_data(m_bins), detray::get_data(m_axes)};
0310   }
0311
0312   /// @returns view of a grid, including the grids multi_axis. Also valid if
0313   /// the value type of the grid is cv qualified (then value_t propagates
0314   /// quialifiers) - const
0315   template <bool owning = is_owning>
0316     requires owning
0317   DETRAY_HOST auto get_data() const -> const_view_type {
0318     return const_view_type{detray::get_data(m_bins), detray::get_data(m_axes)};
0319   }
0320
0321   /// Equality comparison
0322   ///
0323   /// @param rhs the right-hand side of the comparison
0324   ///
0325   /// @note grids could have different bin storage ranges, but could still be
0326   /// identical, hence compare the actual grid content
0327   ///
0328   /// @returns whether the two grids are equal
0329   DETRAY_HOST_DEVICE constexpr auto operator==(const grid_impl &rhs) const
0330       -> bool {
0331     // Check axes: they need to be identical
0332     if (m_axes != rhs.m_axes) {
0333       return false;
0334     }
0335     // Loop over global bin index and compare the two
0336     for (glob_bin_index i = 0; i < nbins(); ++i) {
0337       if (bin(i) != rhs.bin(i)) {
0338         return false;
0339       }
0340     }
0341     return true;
0342   }
0343
0344   /// @returns a string stream that prints the grid details
0345   DETRAY_HOST
0346   friend std::ostream &operator<<(std::ostream &os, const grid_impl &gr) {
0347     os << gr.axes() << std::endl;
0348
0349     for (glob_bin_index gbin = 0; gbin < gr.nbins(); ++gbin) {
0350       os << "bin " << gbin << ": ";
0351       for (const auto &bin_content : gr.bin(gbin)) {
0352         os << bin_content << ", ";
0353       }
0354       os << std::endl;
0355     }
0356
0357     return os;
0358   }
0359
0360  private:
0361   /// Struct that contains the grid's data state
0362   bin_storage m_bins{};
0363   /// The axes of the grid
0364   axes_type m_axes{};
0365 };
0366
0367 /// Type alias for easier construction
0368 template <concepts::algebra algebra_t, typename axes_t, typename bin_t,
0369           template <std::size_t> class serializer_t = simple_serializer,
0370           typename containers = host_container_types, bool ownership = true>
0371 using grid =
0372     grid_impl<coordinate_axes<axes_t, algebra_t, ownership, containers>, bin_t,
0373               simple_serializer>;
0374
0375 }  // namespace detray