From 8e3b8a69532df7f4fb8f4c4953271aa051ce1a3a Mon Sep 17 00:00:00 2001
From: Joana Niermann <joana.niermann@cern.ch>
Date: Thu, 9 Nov 2023 07:58:12 +0100
Subject: [PATCH] This implements the material maps writing. It splits the
 common writer functionality for the grids in a common helper class that is
 inherited by the surface_grid_writer and the new material_map_writer. The
 material_map_writer uses the grid_payload with a material_slab_payload as
 value type.

---
 .../detray/core/detail/multi_store.hpp        |   7 +
 .../io/common/{ => detail}/grid_writer.hpp    | 127 ++++++++----------
 .../detray/io/common/detail/type_traits.hpp   |  46 ++++++-
 .../detray/io/common/detector_writer.hpp      |  11 +-
 .../io/common/homogeneous_material_writer.hpp |   5 +
 .../detray/io/common/material_map_writer.hpp  |  92 +++++++++++++
 .../detray/io/common/surface_grid_writer.hpp  |  82 +++++++++++
 io/include/detray/io/json/json_writer.hpp     |  12 +-
 .../unit_tests/io/io_json_detector_writer.cpp |  22 ++-
 9 files changed, 320 insertions(+), 84 deletions(-)
 rename io/include/detray/io/common/{ => detail}/grid_writer.hpp (60%)
 create mode 100644 io/include/detray/io/common/material_map_writer.hpp
 create mode 100644 io/include/detray/io/common/surface_grid_writer.hpp

diff --git a/core/include/detray/core/detail/multi_store.hpp b/core/include/detray/core/detail/multi_store.hpp
index 83a9fcd62..6084dddeb 100644
--- a/core/include/detray/core/detail/multi_store.hpp
+++ b/core/include/detray/core/detail/multi_store.hpp
@@ -14,6 +14,7 @@
 #include "detray/core/detail/tuple_container.hpp"
 #include "detray/definitions/indexing.hpp"
 #include "detray/definitions/qualifiers.hpp"
+#include "detray/utils/type_list.hpp"
 #include "detray/utils/type_registry.hpp"
 
 // Vecmem include(s)
@@ -358,6 +359,12 @@ class multi_store {
             detail::get<1>(link), std::forward<Args>(args)...);
     }
 
+    /// Print the types that are in the store
+    DETRAY_HOST
+    static constexpr void print() {
+        types::print<types::type_list<detail::get_value_t<Ts>...>>();
+    }
+
     /// @return the view on this tuple container - non-const
     DETRAY_HOST auto get_data() -> view_type {
         return m_tuple_container.get_data();
diff --git a/io/include/detray/io/common/grid_writer.hpp b/io/include/detray/io/common/detail/grid_writer.hpp
similarity index 60%
rename from io/include/detray/io/common/grid_writer.hpp
rename to io/include/detray/io/common/detail/grid_writer.hpp
index 497aad44b..2fe6a96b6 100644
--- a/io/include/detray/io/common/grid_writer.hpp
+++ b/io/include/detray/io/common/detail/grid_writer.hpp
@@ -14,6 +14,7 @@
 #include "detray/io/common/io_interface.hpp"
 #include "detray/io/common/payloads.hpp"
 #include "detray/surface_finders/accelerator_grid.hpp"
+#include "detray/utils/type_list.hpp"
 
 // System include(s)
 #include <algorithm>
@@ -23,83 +24,43 @@
 #include <type_traits>
 #include <vector>
 
-namespace detray {
+namespace detray::detail {
 
 /// @brief Abstract base class for accelerator grid writers
-template <class detector_t,
-          typename value_t = typename detector_t::surface_type>
+template <typename detector_t, typename value_t>
 class grid_writer : public writer_interface<detector_t> {
 
     using base_type = writer_interface<detector_t>;
 
-    // Bin content type to be written into the grid bin payloads
-    // For detector surface descriptors, write only the surface index
-    using content_t = std::conditional_t<
-        std::is_same_v<value_t, typename detector_t::surface_type>, std::size_t,
-        value_t>;
-
-    protected:
-    /// Tag the writer as "surface_grids"
-    inline static const std::string tag = "surface_grids";
-
     public:
     /// Same constructors for this class as for base_type
     using base_type::base_type;
 
-    protected:
     /// Serialize the header information into its payload
-    static grid_header_payload write_header(const detector_t& det,
+    template <typename grid_store_t>
+    static grid_header_payload write_header(const std::string_view writer_tag,
+                                            const grid_store_t& store,
                                             const std::string_view det_name) {
 
         grid_header_payload header_data;
 
-        header_data.common = base_type::serialize(det_name, tag);
+        header_data.common = base_type::serialize(det_name, writer_tag);
 
         header_data.sub_header.emplace();
         auto& grid_sub_header = header_data.sub_header.value();
-        grid_sub_header.n_grids = get_n_grids(det.accelerator_store());
+        grid_sub_header.n_grids = get_n_grids(store);
 
         return header_data;
     }
 
-    /// Serialize the grid collections of a detector @param det into their io
-    /// payload
-    template <
-        typename V = value_t,
-        std::enable_if_t<std::is_same_v<V, typename detector_t::surface_type>,
-                         bool> = true>
-    static detector_grids_payload<std::size_t> serialize(
-        const detector_t& det, const typename detector_t::name_map&) {
-
-        detector_grids_payload<std::size_t> grids_data;
-
-        for (const auto& vol_desc : det.volumes()) {
-            // Links to all acceleration data structures in the volume
-            const auto& multi_link = vol_desc.full_link();
-
-            for (dindex i = 0u; i < multi_link.size(); ++i) {
-                const auto& acc_link = multi_link[i];
-                // Don't look at empty links
-                if (acc_link.is_invalid()) {
-                    continue;
-                }
-
-                // If the accelerator is a grid, insert the payload
-                det.accelerator_store().template visit<get_grid_payload>(
-                    acc_link, vol_desc.index(), grids_data);
-            }
-        }
-
-        return grids_data;
-    }
-
+    protected:
     /// Serialize a grid @param gr of type @param type and index @param idx
-    /// into its io payload
-    template <class grid_t>
-    static grid_payload<content_t> serialize(std::size_t volume_index,
-                                             io::detail::acc_type type,
-                                             const std::size_t idx,
-                                             const grid_t& gr) {
+    /// into its io payload, using @param serializer for the bin content
+    template <typename content_t, class grid_t>
+    static grid_payload<content_t> serialize(
+        std::size_t volume_index, io::detail::acc_type type,
+        const std::size_t idx, const grid_t& gr,
+        std::function<content_t(const value_t&)> serializer) {
 
         grid_payload<content_t> grid_data;
 
@@ -117,7 +78,8 @@ class grid_writer : public writer_interface<detector_t> {
         // Write the surface indices
         for (unsigned int gid = 0u; gid < gr.nbins(); ++gid) {
             // Get the local bin indices and serialize the bin into its payload
-            grid_bin_payload binp = serialize(gr.serialize(gid), gr.bin(gid));
+            grid_bin_payload binp =
+                serialize(gr.serialize(gid), gr.bin(gid), serializer);
             grid_data.bins.push_back(std::move(binp));
         }
 
@@ -160,9 +122,10 @@ class grid_writer : public writer_interface<detector_t> {
     }
 
     /// Serialize a multi-bin @param mbin into its io payload
-    template <std::size_t DIM, typename content_range_t>
+    template <typename content_t, std::size_t DIM, typename content_range_t>
     static grid_bin_payload<content_t> serialize(
-        const n_axis::multi_bin<DIM> mbin, const content_range_t& content) {
+        const n_axis::multi_bin<DIM> mbin, const content_range_t& content,
+        std::function<content_t(const value_t&)> serializer) {
 
         grid_bin_payload<content_t> bin_data;
 
@@ -174,46 +137,64 @@ class grid_writer : public writer_interface<detector_t> {
         // Put all entries of the bin into the payload
         bin_data.content.reserve(content.size());
         for (const auto& entry : content) {
-            bin_data.content.push_back(entry.index());
+            bin_data.content.push_back(serializer(entry));
         }
 
         return bin_data;
     }
 
+    /// Serialize a grid from a collection into its payload
+    ///
+    /// @param store the data store of grids (tuple of grid collections)
+    /// @param grid_link type and index of the grid
+    /// @param owner_idx inder of the owner of the grid (e.g. volume index)
+    /// @param grid_data the grid payload to be filled
+    /// @param serializer callable that can serialize a grid bin entry into its
+    /// respective IO payload (of type @tparam content_t)
+    template <typename store_t, typename content_t, typename serializer_t>
+    static void serialize(const store_t& store,
+                          typename store_t::single_link grid_link,
+                          dindex owner_idx,
+                          detector_grids_payload<content_t>& grids_data,
+                          serializer_t serializer) {
+
+        // If the accelerator is a grid, insert the payload
+        store.template visit<get_grid_payload>(grid_link, owner_idx, grids_data,
+                                               serializer);
+    }
+
     private:
     /// Retrieve a @c grid_payload from grid collection elements
     struct get_grid_payload {
 
-        template <typename grid_group_t, typename index_t>
+        template <typename grid_group_t, typename index_t, typename content_t,
+                  typename serializer_t>
         inline void operator()(
             [[maybe_unused]] const grid_group_t& coll,
             [[maybe_unused]] const index_t& index,
-            [[maybe_unused]] std::size_t volume_index,
-            [[maybe_unused]] detector_grids_payload<content_t>& grids_data)
-            const {
+            [[maybe_unused]] std::size_t link,
+            [[maybe_unused]] detector_grids_payload<content_t>& grids_data,
+            [[maybe_unused]] serializer_t& serializer) const {
 
             using accel_t = typename grid_group_t::value_type;
 
             if constexpr (detail::is_grid_v<accel_t>) {
 
-                grids_data.grids.push_back(
-                    serialize(volume_index, io::detail::get_grid_id<accel_t>(),
-                              index, coll[index]));
+                grids_data.grids.push_back(serialize<content_t>(
+                    link, io::detail::get_grid_id<accel_t>(), index,
+                    coll[index], serializer));
             }
         }
     };
 
     /// Retrieve number of overall grids in detector
-    template <std::size_t I = 0u>
-    static std::size_t get_n_grids(
-        const typename detector_t::accelerator_container& store,
-        std::size_t n = 0u) {
+    template <std::size_t I = 0u, typename store_t>
+    static std::size_t get_n_grids(const store_t& store, std::size_t n = 0u) {
 
-        using store_t = typename detector_t::accelerator_container;
         constexpr auto coll_id{store_t::value_types::to_id(I)};
-        using accel_t = typename store_t::template get_type<coll_id>;
+        using value_type = typename store_t::template get_type<coll_id>;
 
-        if constexpr (detail::is_grid_v<accel_t>) {
+        if constexpr (detail::is_grid_v<value_type>) {
             n += store.template size<coll_id>();
         }
 
@@ -224,4 +205,4 @@ class grid_writer : public writer_interface<detector_t> {
     }
 };
 
-}  // namespace detray
+}  // namespace detray::detail
diff --git a/io/include/detray/io/common/detail/type_traits.hpp b/io/include/detray/io/common/detail/type_traits.hpp
index 4be458f9a..2adb406a6 100644
--- a/io/include/detray/io/common/detail/type_traits.hpp
+++ b/io/include/detray/io/common/detail/type_traits.hpp
@@ -156,6 +156,37 @@ struct mask_info<io::detail::mask_shape::trapezoid2, detector_t,
 };
 /// @}
 
+/// Check for grid types in a data store
+/// @{
+template <typename>
+struct constains_grids;
+
+template <template <typename, typename...> class R, class ID, typename... Ts>
+struct constains_grids<R<ID, Ts...>> {
+    static constexpr bool value{std::disjunction_v<detail::is_grid<Ts>...>};
+};
+
+template <typename T>
+inline constexpr bool constains_grids_v = constains_grids<T>::value;
+/// @}
+
+/// Check for the various types of material
+/// @{
+template <class detector_t, typename = void>
+struct has_surface_grids : public std::false_type {};
+
+template <class detector_t>
+struct has_surface_grids<
+    detector_t,
+    std::enable_if_t<constains_grids_v<typename detector_t::accel>, void>>
+    : public std::true_type {};
+
+template <typename T>
+inline constexpr bool has_surface_grids_v = has_surface_grids<T>::value;
+/// @}
+
+/// Check for the various types of material
+/// @{
 template <class detector_t, typename = void>
 struct has_material_slabs : public std::false_type {};
 
@@ -185,8 +216,6 @@ inline constexpr bool has_material_rods_v = has_material_rods<T>::value;
 template <class detector_t, typename = void>
 struct has_homogeneous_material : public std::false_type {};
 
-/// Is the value type in the detector material store a simple material or is it
-/// wrapped in another class (e.g. grids for material maps)
 template <class detector_t>
 struct has_homogeneous_material<
     detector_t, std::enable_if_t<has_material_slabs_v<detector_t> or
@@ -197,4 +226,17 @@ template <typename T>
 inline constexpr bool has_homogeneous_material_v =
     has_homogeneous_material<T>::value;
 
+template <class detector_t, typename = void>
+struct has_material_grids : public std::false_type {};
+
+template <class detector_t>
+struct has_material_grids<
+    detector_t,
+    std::enable_if_t<constains_grids_v<typename detector_t::materials>, void>>
+    : public std::true_type {};
+
+template <typename T>
+inline constexpr bool has_material_grids_v = has_material_grids<T>::value;
+/// @}
+
 }  // namespace detray::detail
diff --git a/io/include/detray/io/common/detector_writer.hpp b/io/include/detray/io/common/detector_writer.hpp
index 14d3a9115..3ddf79207 100644
--- a/io/include/detray/io/common/detector_writer.hpp
+++ b/io/include/detray/io/common/detector_writer.hpp
@@ -97,10 +97,15 @@ detail::detector_component_writers<detector_t> assemble_writer(
                 writers.template add<json_homogeneous_material_writer>();
             }
             // Material maps
-            // ...
+            if constexpr (detail::has_material_grids_v<detector_t>) {
+                writers.template add<json_material_map_writer>();
+            }
         }
-        if (cfg.write_grids()) {
-            writers.template add<json_grid_writer>();
+        // Navigation acceleration structures
+        if constexpr (detail::has_surface_grids_v<detector_t>) {
+            if (cfg.write_grids()) {
+                writers.template add<json_surface_grid_writer>();
+            }
         }
     }
 
diff --git a/io/include/detray/io/common/homogeneous_material_writer.hpp b/io/include/detray/io/common/homogeneous_material_writer.hpp
index d247ebbd2..189f0e452 100644
--- a/io/include/detray/io/common/homogeneous_material_writer.hpp
+++ b/io/include/detray/io/common/homogeneous_material_writer.hpp
@@ -22,6 +22,9 @@
 
 namespace detray {
 
+template <typename detector_t>
+class material_map_writer;
+
 /// @brief Abstract base class for simple material description writers
 template <class detector_t>
 class homogeneous_material_writer : public writer_interface<detector_t> {
@@ -29,6 +32,8 @@ class homogeneous_material_writer : public writer_interface<detector_t> {
     using base_type = writer_interface<detector_t>;
     using scalar_t = typename detector_t::scalar_type;
 
+    friend class material_map_writer<detector_t>;
+
     protected:
     /// Tag the writer as "homogeneous_material"
     inline static const std::string tag = "homogeneous_material";
diff --git a/io/include/detray/io/common/material_map_writer.hpp b/io/include/detray/io/common/material_map_writer.hpp
new file mode 100644
index 000000000..02e3071b6
--- /dev/null
+++ b/io/include/detray/io/common/material_map_writer.hpp
@@ -0,0 +1,92 @@
+/** Detray library, part of the ACTS project (R&D line)
+ *
+ * (c) 2023 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s)
+#include "detray/io/common/detail/grid_writer.hpp"
+#include "detray/io/common/detail/type_traits.hpp"
+#include "detray/io/common/homogeneous_material_writer.hpp"
+#include "detray/io/common/io_interface.hpp"
+#include "detray/io/common/payloads.hpp"
+#include "detray/materials/material_slab.hpp"
+
+// System include(s)
+#include <string>
+
+namespace detray {
+
+/// @brief Abstract base class for material maps writers
+template <class detector_t>
+class material_map_writer
+    : public detail::grid_writer<
+          detector_t, material_slab<typename detector_t::scalar_type>> {
+
+    using scalar_t = typename detector_t::scalar_type;
+    using material_t = material_slab<scalar_t>;
+
+    using base_type =
+        detail::grid_writer<detector_t,
+                            material_slab<typename detector_t::scalar_type>>;
+    using grid_writer_t = base_type;
+    using mat_writer_t = homogeneous_material_writer<detector_t>;
+
+    protected:
+    /// Tag the writer as "material_map"
+    inline static const std::string tag = "material_maps";
+
+    public:
+    /// Same constructors for this class as for base_type
+    using base_type::base_type;
+
+    protected:
+    /// Serialize the header information into its payload
+    static auto write_header(const detector_t& det,
+                             const std::string_view det_name) {
+
+        return grid_writer_t::write_header(tag, det.material_store(), det_name);
+    }
+
+    /// Serialize the material description of a detector @param det into its io
+    /// payload
+    static detector_grids_payload<material_slab_payload> serialize(
+        const detector_t& det, const typename detector_t::name_map&) {
+
+        detector_grids_payload<material_slab_payload> grids_data;
+
+        // How to serialize a material slab in the grid
+        auto mat_serializer = [](const material_t& mat) {
+            return mat_writer_t::serialize(mat, dindex_invalid);
+        };
+
+        for (const auto& vol_desc : det.volumes()) {
+
+            /// Check if a surface has a metrial map
+            for (const auto& sf_desc : det.surface_lookup()) {
+                if (sf_desc.volume() != vol_desc.index()) {
+                    continue;
+                }
+
+                const auto& mat_link = sf_desc.material();
+                // Don't look at empty links
+                if (mat_link.is_invalid() or
+                    mat_link.id() == detector_t::materials::id::e_none) {
+                    continue;
+                }
+
+                // Generate the payload
+                grid_writer_t::serialize(det.material_store(), mat_link,
+                                         sf_desc.index(), grids_data,
+                                         mat_serializer);
+            }
+        }
+
+        return grids_data;
+    }
+};
+
+}  // namespace detray
diff --git a/io/include/detray/io/common/surface_grid_writer.hpp b/io/include/detray/io/common/surface_grid_writer.hpp
new file mode 100644
index 000000000..48446b921
--- /dev/null
+++ b/io/include/detray/io/common/surface_grid_writer.hpp
@@ -0,0 +1,82 @@
+/** Detray library, part of the ACTS project (R&D line)
+ *
+ * (c) 2023 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s)
+#include "detray/definitions/indexing.hpp"
+#include "detray/io/common/detail/grid_writer.hpp"
+#include "detray/io/common/io_interface.hpp"
+#include "detray/io/common/payloads.hpp"
+
+// System include(s)
+#include <string>
+
+namespace detray {
+
+/// @brief Abstract base class for accelerator grid writers
+template <class detector_t>
+class surface_grid_writer
+    : public detail::grid_writer<detector_t,
+                                 typename detector_t::surface_type> {
+
+    using surface_t = typename detector_t::surface_type;
+    using base_type = detail::grid_writer<detector_t, surface_t>;
+    using grid_writer_t = base_type;
+
+    protected:
+    /// Tag the writer as "surface_grids"
+    inline static const std::string tag = "surface_grids";
+
+    public:
+    /// Same constructors for this class as for base_type
+    using base_type::base_type;
+
+    protected:
+    /// Serialize the header information into its payload
+    static auto write_header(const detector_t& det,
+                             const std::string_view det_name) {
+
+        return grid_writer_t::write_header(tag, det.accelerator_store(),
+                                           det_name);
+    }
+
+    /// Serialize the grid collections of a detector @param det into their io
+    /// payload
+    static detector_grids_payload<std::size_t> serialize(
+        const detector_t& det, const typename detector_t::name_map&) {
+
+        detector_grids_payload<std::size_t> grids_data;
+
+        // How to serialize the surface descriptors in the grid
+        auto sf_serializer = [](const surface_t& sf_desc) {
+            return sf_desc.index();
+        };
+
+        for (const auto& vol_desc : det.volumes()) {
+            // Links to all acceleration data structures in the volume
+            const auto& multi_link = vol_desc.full_link();
+
+            for (dindex i = 0u; i < multi_link.size(); ++i) {
+                const auto& acc_link = multi_link[i];
+                // Don't look at empty links
+                if (acc_link.is_invalid()) {
+                    continue;
+                }
+
+                // Generate the payload
+                grid_writer_t::serialize(det.accelerator_store(), acc_link,
+                                         vol_desc.index(), grids_data,
+                                         sf_serializer);
+            }
+        }
+
+        return grids_data;
+    }
+};
+
+}  // namespace detray
diff --git a/io/include/detray/io/json/json_writer.hpp b/io/include/detray/io/json/json_writer.hpp
index 538640afd..a73785de4 100644
--- a/io/include/detray/io/json/json_writer.hpp
+++ b/io/include/detray/io/json/json_writer.hpp
@@ -10,8 +10,9 @@
 // Project include(s)
 #include "detray/io/common/detail/file_handle.hpp"
 #include "detray/io/common/geometry_writer.hpp"
-#include "detray/io/common/grid_writer.hpp"
 #include "detray/io/common/homogeneous_material_writer.hpp"
+#include "detray/io/common/material_map_writer.hpp"
+#include "detray/io/common/surface_grid_writer.hpp"
 #include "detray/io/json/json.hpp"
 #include "detray/io/json/json_serializers.hpp"
 
@@ -91,9 +92,12 @@ template <typename detector_t>
 using json_homogeneous_material_writer =
     json_writer<detector_t, homogeneous_material_writer>;
 
+/// Write a material map description to file in json format
+template <typename detector_t>
+using json_material_map_writer = json_writer<detector_t, material_map_writer>;
+
 /// Write the detector grid collections to file in json format
-template <typename detector_t,
-          typename value_t = typename detector_t::surface_type>
-using json_grid_writer = json_writer<detector_t, grid_writer, value_t>;
+template <typename detector_t>
+using json_surface_grid_writer = json_writer<detector_t, surface_grid_writer>;
 
 }  // namespace detray
diff --git a/tests/unit_tests/io/io_json_detector_writer.cpp b/tests/unit_tests/io/io_json_detector_writer.cpp
index 61a775a26..40cff2d47 100644
--- a/tests/unit_tests/io/io_json_detector_writer.cpp
+++ b/tests/unit_tests/io/io_json_detector_writer.cpp
@@ -73,6 +73,22 @@ TEST(io, json_telescope_material_writer) {
     mat_writer.write(det, names);
 }
 
+/// Test the writing of the toy detector grids to json
+TEST(io, json_toy_material_maps_writer) {
+
+    using detector_t = detector<toy_metadata>;
+
+    // Toy detector
+    vecmem::host_memory_resource host_mr;
+    toy_det_config toy_cfg{};
+    toy_cfg.use_material_maps(true);
+    auto [det, names] = create_toy_geometry(host_mr, toy_cfg);
+
+    json_material_map_writer<detector_t> map_writer;
+    map_writer.write(det, names,
+                     std::ios::out | std::ios::binary | std::ios::trunc);
+}
+
 /// Test the writing of the toy detector grids to json
 TEST(io, json_toy_grid_writer) {
 
@@ -82,7 +98,7 @@ TEST(io, json_toy_grid_writer) {
     vecmem::host_memory_resource host_mr;
     auto [det, names] = create_toy_geometry(host_mr);
 
-    json_grid_writer<detector_t> grid_writer;
+    json_surface_grid_writer<detector_t> grid_writer;
     grid_writer.write(det, names,
                       std::ios::out | std::ios::binary | std::ios::trunc);
 }
@@ -92,7 +108,9 @@ TEST(io, json_toy_detector_writer) {
 
     // Toy detector
     vecmem::host_memory_resource host_mr;
-    auto [det, names] = create_toy_geometry(host_mr);
+    toy_det_config toy_cfg{};
+    toy_cfg.use_material_maps(true);
+    const auto [det, names] = create_toy_geometry(host_mr, toy_cfg);
 
     auto writer_cfg = io::detector_writer_config{}
                           .format(io::format::json)