Merge pull request #753 from asalzburger/feat-allow-material-maps-custom-boundaries

feat: allow to read material map boundaries

Author: asalzburger

core/include/detray/builders/grid_factory.hpp

@@ -36,6 +36,11 @@ namespace detray {
 /// @tparam serialzier_t  type of the serializer to the storage represenations
 /// @tparam algebra_t the matrix/vector/point types to use
 /// @tparam container_t the container types to use
+///
+/// throughout:
+/// @note that bin_edges is only used for variable binning
+/// @note that if non-zero axis_spans are provided the values of the
+/// mask is overwritten
 template <typename bin_t, template <std::size_t> class serializer_t,
           typename algebra_t = ALGEBRA_PLUGIN<detray::scalar>>
 class grid_factory {
@@ -76,12 +81,14 @@ class grid_factory {
         typename phi_binning = axis::regular<host_container_types, scalar_type>,
         std::enable_if_t<std::is_enum_v<decltype(r_bounds::label)>, bool> =
             true>
-    auto new_grid(const mask<annulus2D> &grid_bounds,
-                  const std::array<std::size_t, 2UL> n_bins,
-                  const std::vector<std::pair<loc_bin_index<annulus2D>, dindex>>
-                      &bin_capacities = {},
-                  const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {
-                      {}}) const {
+    auto new_grid(
+        const mask<annulus2D> &grid_bounds,
+        const std::array<std::size_t, 2UL> n_bins,
+        const std::vector<std::pair<loc_bin_index<annulus2D>, dindex>>
+            &bin_capacities = {},
+        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 2UL> &axis_spans = {
+            {}}) const {
 
         static_assert(
             std::is_same_v<phi_bounds, axis::circular<>>,
@@ -96,13 +103,25 @@ class grid_factory {
         constexpr auto e_phi_axis = static_cast<dindex>(axes_t::label1);
 
         auto b_values = grid_bounds.values();
+        // Overwrite the mask values if axis spans are provided
+        if (!axis_spans[0UL].empty()) {
+            assert(axis_spans[0UL].size() == 2UL);
+            b_values[boundary::e_min_r] = axis_spans[0UL].at(0UL);
+            b_values[boundary::e_max_r] = axis_spans[0UL].at(1UL);
+        }
+        scalar_type min_phi = b_values[boundary::e_average_phi] -
+                              b_values[boundary::e_min_phi_rel];
+        scalar_type max_phi = b_values[boundary::e_average_phi] +
+                              b_values[boundary::e_max_phi_rel];
+        if (!axis_spans[1UL].empty()) {
+            assert(axis_spans[1UL].size() == 2UL);
+            min_phi = axis_spans[1UL].at(0UL);
+            max_phi = axis_spans[1UL].at(1UL);
+        }
 
         return new_grid<local_frame>(
-            {b_values[boundary::e_min_r], b_values[boundary::e_max_r],
-             b_values[boundary::e_average_phi] -
-                 b_values[boundary::e_min_phi_rel],
-             b_values[boundary::e_average_phi] +
-                 b_values[boundary::e_max_phi_rel]},
+            {b_values[boundary::e_min_r], b_values[boundary::e_max_r], min_phi,
+             max_phi},
             {n_bins[e_r_axis], n_bins[e_phi_axis]}, bin_capacities,
             {bin_edges[e_r_axis], bin_edges[e_phi_axis]},
             types::list<r_bounds, phi_bounds>{},
@@ -121,12 +140,14 @@ class grid_factory {
         typename z_binning = axis::regular<host_container_types, scalar_type>,
         std::enable_if_t<std::is_enum_v<decltype(x_bounds::label)>, bool> =
             true>
-    auto new_grid(const mask<cuboid3D> &grid_bounds,
-                  const std::array<std::size_t, 3UL> n_bins,
-                  const std::vector<std::pair<loc_bin_index<cuboid3D>, dindex>>
-                      &bin_capacities = {},
-                  const std::array<std::vector<scalar_type>, 3UL> &bin_edges = {
-                      {}}) const {
+    auto new_grid(
+        const mask<cuboid3D> &grid_bounds,
+        const std::array<std::size_t, 3UL> n_bins,
+        const std::vector<std::pair<loc_bin_index<cuboid3D>, dindex>>
+            &bin_capacities = {},
+        const std::array<std::vector<scalar_type>, 3UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 3UL> &axis_spans = {
+            {}}) const {
         // Axes boundaries and local indices
         using boundary = cuboid3D::boundaries;
         using axes_t = axes<cuboid3D>::template type<algebra_t>;
@@ -136,7 +157,23 @@ class grid_factory {
         constexpr auto e_y_axis = static_cast<dindex>(axes_t::label1);
         constexpr auto e_z_axis = static_cast<dindex>(axes_t::label2);
 
+        // Overwrite the mask values if axis spans are provided
         auto b_values = grid_bounds.values();
+        if (!axis_spans[0UL].empty()) {
+            assert(axis_spans[0UL].size() == 2UL);
+            b_values[boundary::e_min_x] = axis_spans[0UL].at(0UL);
+            b_values[boundary::e_max_x] = axis_spans[0UL].at(1UL);
+        }
+        if (!axis_spans[1UL].empty()) {
+            assert(axis_spans[1UL].size() == 2UL);
+            b_values[boundary::e_min_y] = axis_spans[1UL].at(0UL);
+            b_values[boundary::e_max_y] = axis_spans[1UL].at(1UL);
+        }
+        if (!axis_spans[2UL].empty()) {
+            assert(axis_spans[2UL].size() == 2UL);
+            b_values[boundary::e_min_z] = axis_spans[2UL].at(0UL);
+            b_values[boundary::e_max_z] = axis_spans[2UL].at(1UL);
+        }
 
         return new_grid<local_frame>(
             {b_values[boundary::e_min_x], b_values[boundary::e_max_x],
@@ -165,7 +202,8 @@ class grid_factory {
         const std::array<std::size_t, 2UL> n_bins,
         const std::vector<std::pair<loc_bin_index<cylinder2D>, dindex>>
             &bin_capacities = {},
-        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {
+        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 2UL> &axis_spans = {
             {}}) const {
 
         static_assert(
@@ -181,7 +219,11 @@ class grid_factory {
         constexpr auto e_z_axis = static_cast<dindex>(axes_t::label1);
 
         auto b_values = grid_bounds.values();
-
+        if (!axis_spans[1UL].empty()) {
+            assert(axis_spans[1UL].size() == 2UL);
+            b_values[boundary::e_lower_z] = axis_spans[1UL].at(0UL);
+            b_values[boundary::e_upper_z] = axis_spans[1UL].at(1UL);
+        }
         return new_grid<local_frame>(
             {-constant<scalar_type>::pi, constant<scalar_type>::pi,
              b_values[boundary::e_lower_z], b_values[boundary::e_upper_z]},
@@ -202,13 +244,14 @@ class grid_factory {
         typename z_binning = axis::regular<host_container_types, scalar_type>,
         std::enable_if_t<std::is_enum_v<decltype(rphi_bounds::label)>, bool> =
             true>
-    auto new_grid(const mask<concentric_cylinder2D> &grid_bounds,
-                  const std::array<std::size_t, 2UL> n_bins,
-                  const std::vector<
-                      std::pair<loc_bin_index<concentric_cylinder2D>, dindex>>
-                      &bin_capacities = {},
-                  const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {
-                      {}}) const {
+    auto new_grid(
+        const mask<concentric_cylinder2D> &grid_bounds,
+        const std::array<std::size_t, 2UL> n_bins,
+        const std::vector<std::pair<loc_bin_index<concentric_cylinder2D>,
+                                    dindex>> &bin_capacities = {},
+        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 2UL> &axis_spans = {
+            {}}) const {
 
         static_assert(
             std::is_same_v<rphi_bounds, axis::circular<axis::label::e_rphi>>,
@@ -223,6 +266,11 @@ class grid_factory {
         constexpr auto e_z_axis = static_cast<dindex>(axes_t::label1);
 
         auto b_values = grid_bounds.values();
+        if (!axis_spans[1UL].empty()) {
+            assert(axis_spans[1UL].size() == 2UL);
+            b_values[boundary::e_lower_z] = axis_spans[1UL].at(0UL);
+            b_values[boundary::e_upper_z] = axis_spans[1UL].at(1UL);
+        }
 
         return new_grid<local_frame>(
             {-constant<scalar_type>::pi, constant<scalar_type>::pi,
@@ -250,7 +298,8 @@ class grid_factory {
         const std::array<std::size_t, 3UL> n_bins,
         const std::vector<std::pair<loc_bin_index<cylinder3D>, dindex>>
             &bin_capacities = {},
-        const std::array<std::vector<scalar_type>, 3UL> &bin_edges = {
+        const std::array<std::vector<scalar_type>, 3UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 3UL> &axis_spans = {
             {}}) const {
 
         static_assert(
@@ -267,11 +316,29 @@ class grid_factory {
         constexpr auto e_z_axis = static_cast<dindex>(axes_t::label2);
 
         auto b_values = grid_bounds.values();
+        // Overwrite the mask values if axis spans are provided
+        if (!axis_spans[0UL].empty()) {
+            assert(axis_spans[0UL].size() == 2UL);
+            b_values[boundary::e_min_r] = axis_spans[0UL].at(0UL);
+            b_values[boundary::e_max_r] = axis_spans[0UL].at(1UL);
+        }
+        scalar_type min_phi = -constant<scalar_type>::pi;
+        scalar_type max_phi = constant<scalar_type>::pi;
+        if (!axis_spans[1UL].empty()) {
+            assert(axis_spans[1UL].size() == 2UL);
+            min_phi = axis_spans[1UL].at(0UL);
+            max_phi = axis_spans[1UL].at(1UL);
+        }
+        if (!axis_spans[2UL].empty()) {
+            assert(axis_spans[2UL].size() == 2UL);
+            b_values[boundary::e_min_z] = axis_spans[2UL].at(0UL);
+            b_values[boundary::e_max_z] = axis_spans[2UL].at(1UL);
+        }
 
         return new_grid<local_frame>(
-            {b_values[boundary::e_min_r], b_values[boundary::e_max_r],
-             b_values[boundary::e_min_phi], b_values[boundary::e_max_phi],
-             -b_values[boundary::e_min_z], b_values[boundary::e_max_z]},
+            {b_values[boundary::e_min_r], b_values[boundary::e_max_r], min_phi,
+             max_phi, -b_values[boundary::e_min_z],
+             b_values[boundary::e_max_z]},
             {n_bins[e_r_axis], n_bins[e_phi_axis], n_bins[e_z_axis]},
             bin_capacities,
             {bin_edges[e_r_axis], bin_edges[e_phi_axis], bin_edges[e_z_axis]},
@@ -289,12 +356,14 @@ class grid_factory {
         typename phi_binning = axis::regular<host_container_types, scalar_type>,
         std::enable_if_t<std::is_enum_v<decltype(r_bounds::label)>, bool> =
             true>
-    auto new_grid(const mask<ring2D> &grid_bounds,
-                  const std::array<std::size_t, 2UL> n_bins,
-                  const std::vector<std::pair<loc_bin_index<ring2D>, dindex>>
-                      &bin_capacities = {},
-                  const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {
-                      {}}) const {
+    auto new_grid(
+        const mask<ring2D> &grid_bounds,
+        const std::array<std::size_t, 2UL> n_bins,
+        const std::vector<std::pair<loc_bin_index<ring2D>, dindex>>
+            &bin_capacities = {},
+        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 2UL> &axis_spans = {
+            {}}) const {
 
         static_assert(std::is_same_v<phi_bounds, axis::circular<>>,
                       "Phi axis bounds need to be circular for ring shape");
@@ -308,6 +377,12 @@ class grid_factory {
         constexpr auto e_phi_axis = static_cast<dindex>(axes_t::label1);
 
         auto b_values = grid_bounds.values();
+        // Overwrite the mask values if axis spans are provided
+        if (!axis_spans[0UL].empty()) {
+            assert(axis_spans[0UL].size() == 2UL);
+            b_values[boundary::e_inner_r] = axis_spans[0UL].at(0UL);
+            b_values[boundary::e_outer_r] = axis_spans[0UL].at(1UL);
+        }
 
         return new_grid<local_frame>(
             {b_values[boundary::e_inner_r], b_values[boundary::e_outer_r],
@@ -333,7 +408,8 @@ class grid_factory {
         const std::array<std::size_t, 2UL> n_bins,
         const std::vector<std::pair<loc_bin_index<rectangle2D>, dindex>>
             &bin_capacities = {},
-        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {
+        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 2UL> &axis_spans = {
             {}}) const {
         // Axes boundaries and local indices
         using boundary = rectangle2D::boundaries;
@@ -344,6 +420,15 @@ class grid_factory {
         constexpr auto e_y_axis = static_cast<dindex>(axes_t::label1);
 
         auto b_values = grid_bounds.values();
+        // Overwrite the mask values if axis spans are provided
+        if (!axis_spans[0UL].empty()) {
+            assert(axis_spans[0UL].size() == 2UL);
+            b_values[boundary::e_half_x] = axis_spans[0UL].at(1UL);
+        }
+        if (!axis_spans[1UL].empty()) {
+            assert(axis_spans[1UL].size() == 2UL);
+            b_values[boundary::e_half_y] = axis_spans[1UL].at(1UL);
+        }
 
         return new_grid<local_frame>(
             {-b_values[boundary::e_half_x], b_values[boundary::e_half_x],
@@ -369,7 +454,8 @@ class grid_factory {
         const std::array<std::size_t, 2UL> n_bins,
         const std::vector<std::pair<loc_bin_index<trapezoid2D>, dindex>>
             &bin_capacities = {},
-        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {
+        const std::array<std::vector<scalar_type>, 2UL> &bin_edges = {{}},
+        const std::array<std::vector<scalar_type>, 2UL> &axis_spans = {
             {}}) const {
         // Axes boundaries and local indices
         using boundary = trapezoid2D::boundaries;
@@ -380,6 +466,15 @@ class grid_factory {
         constexpr auto e_y_axis = static_cast<dindex>(axes_t::label1);
 
         auto b_values = grid_bounds.values();
+        // Overwrite the mask values if axis spans are provided
+        if (!axis_spans[0UL].empty()) {
+            assert(axis_spans[0UL].size() == 2UL);
+            b_values[boundary::e_half_length_1] = axis_spans[0UL].at(1UL);
+        }
+        if (!axis_spans[1UL].empty()) {
+            assert(axis_spans[1UL].size() == 2UL);
+            b_values[boundary::e_half_length_2] = axis_spans[1UL].at(1UL);
+        }
 
         return new_grid<local_frame>({-b_values[boundary::e_half_length_1],
                                       b_values[boundary::e_half_length_1],

core/include/detray/builders/material_map_builder.hpp

@@ -87,7 +87,8 @@ class material_map_builder : public volume_decorator<detector_t> {
         auto mat_factory = std::dynamic_pointer_cast<
             material_map_factory<detector_t, axis::multi_bin<DIM>>>(sf_factory);
         if (mat_factory) {
-            (*mat_factory)(this->surfaces(), m_bin_data, m_n_bins);
+            (*mat_factory)(this->surfaces(), m_bin_data, m_n_bins,
+                           m_axis_spans);
         }
         auto mat_generator =
             std::dynamic_pointer_cast<material_map_generator<detector_t>>(
@@ -122,12 +123,19 @@ class material_map_builder : public volume_decorator<detector_t> {
                 continue;
             }
 
+            // Construct and append the material map for a given surface shape
+            std::array<std::vector<scalar_type>, DIM> axis_spans{};
+            auto axis_spans_itr = m_axis_spans.find(sf_idx - 1u);
+            if (axis_spans_itr != m_axis_spans.end()) {
+                axis_spans = axis_spans_itr->second;
+            }
+
             // Construct and append the material map for a given surface shape
             auto sf = surface{det, sf_desc};
             [[maybe_unused]] auto [mat_id, mat_idx] = sf.template visit_mask<
                 detail::add_sf_material_map<materials_t>>(
                 m_factory, m_bin_data.at(sf_idx - 1u), m_n_bins.at(sf_idx - 1u),
-                det.material_store());
+                axis_spans, det.material_store());
 
             // Make sure the linking was precomputed correctly
             assert(mat_id == sf_desc.material().id());
@@ -189,6 +197,8 @@ class material_map_builder : public volume_decorator<detector_t> {
     std::map<dindex, std::vector<bin_data_type>> m_bin_data;
     /// Number of bins for the material grid axes
     std::map<dindex, std::array<std::size_t, DIM>> m_n_bins{};
+    /// The Axis spans for the material grid axes
+    std::map<dindex, std::array<std::vector<scalar_type>, DIM>> m_axis_spans{};
     /// Helper to generate empty grids
     mat_map_factory_t m_factory{};
 };
@@ -211,13 +221,16 @@ template <typename materials_t>
 struct add_sf_material_map {
 
     template <typename coll_t, typename index_t, typename mat_factory_t,
-              typename bin_data_t, std::size_t DIM, typename material_store_t>
+              typename bin_data_t, std::size_t DIM, typename material_store_t,
+              typename scalar_t>
     DETRAY_HOST inline std::pair<typename materials_t::id, dindex> operator()(
         [[maybe_unused]] const coll_t& coll,
         [[maybe_unused]] const index_t& index,
         [[maybe_unused]] const mat_factory_t& mat_factory,
         [[maybe_unused]] std::vector<bin_data_t>& bin_data,
         [[maybe_unused]] const std::array<std::size_t, DIM>& n_bins,
+        [[maybe_unused]] const std::array<std::vector<scalar_t>, DIM>&
+            axis_spans,
         [[maybe_unused]] material_store_t& mat_store) const {
 
         using mask_shape_t = typename coll_t::value_type::shape;
@@ -230,7 +243,8 @@ struct add_sf_material_map {
         if constexpr (!is_line && mask_shape_t::dim == DIM) {
             // Map a grid onto the surface mask
             const auto& sf_mask = coll[index];
-            auto mat_grid = mat_factory.new_grid(sf_mask, n_bins);
+            auto mat_grid =
+                mat_factory.new_grid(sf_mask, n_bins, {}, {}, axis_spans);
 
             // The detector only knows the non-owning grid types
             using non_owning_t =