Implement BIH tree

src/orange/BoundingBox.hh

@@ -32,7 +32,7 @@ class BoundingBox
     //!@{
     //! \name Type aliases
     using value_type = T;
-    using Array3 = Array<value_type, 3>;
+    using Real3 = Array<value_type, 3>;
     //!@}
 
     // Construct from infinite extents
@@ -42,22 +42,22 @@ class BoundingBox
     inline CELER_FUNCTION BoundingBox();
 
     // Construct from upper and lower points
-    inline CELER_FUNCTION BoundingBox(Array3 const& lower, Array3 const& upper);
+    inline CELER_FUNCTION BoundingBox(Real3 const& lower, Real3 const& upper);
 
     //// ACCESSORS ////
 
     // Lower bbox coordinate
-    CELER_FORCEINLINE_FUNCTION Array3 const& lower() const;
+    CELER_FORCEINLINE_FUNCTION Real3 const& lower() const;
 
     // Upper bbox coordinate
-    CELER_FORCEINLINE_FUNCTION Array3 const& upper() const;
+    CELER_FORCEINLINE_FUNCTION Real3 const& upper() const;
 
     // Whether the bbox is assigned
     CELER_FORCEINLINE_FUNCTION explicit operator bool() const;
 
   private:
-    Array3 lower_;
-    Array3 upper_;
+    Real3 lower_;
+    Real3 upper_;
 };
 
 //---------------------------------------------------------------------------//
@@ -98,7 +98,7 @@ CELER_FUNCTION BoundingBox<T>::BoundingBox()
  * at a single point) but upper must not be less than lower.
  */
 template<class T>
-CELER_FUNCTION BoundingBox<T>::BoundingBox(Array3 const& lo, Array3 const& hi)
+CELER_FUNCTION BoundingBox<T>::BoundingBox(Real3 const& lo, Real3 const& hi)
     : lower_(lo), upper_(hi)
 {
     CELER_EXPECT(lower_[0] <= upper_[0] && lower_[1] <= upper_[1]
@@ -110,8 +110,7 @@ CELER_FUNCTION BoundingBox<T>::BoundingBox(Array3 const& lo, Array3 const& hi)
  * Lower bbox coordinate (must be valid).
  */
 template<class T>
-CELER_FUNCTION typename BoundingBox<T>::Array3 const&
-BoundingBox<T>::lower() const
+CELER_FUNCTION auto BoundingBox<T>::lower() const -> Real3 const&
 {
     CELER_EXPECT(*this);
     return lower_;
@@ -122,8 +121,7 @@ BoundingBox<T>::lower() const
  * Upper bbox coordinate (must be valid).
  */
 template<class T>
-CELER_FUNCTION typename BoundingBox<T>::Array3 const&
-BoundingBox<T>::upper() const
+CELER_FUNCTION auto BoundingBox<T>::upper() const -> Real3 const&
 {
     CELER_EXPECT(*this);
     return upper_;

src/orange/CMakeLists.txt

@@ -16,6 +16,8 @@ list(APPEND SOURCES
   OrangeParams.cc
   OrangeTypes.cc
   construct/SurfaceInputBuilder.cc
+  detail/BIHBuilder.cc
+  detail/BIHPartitioner.cc
   detail/RectArrayInserter.cc
   detail/UnitInserter.cc
   surf/SurfaceIO.cc

src/orange/OrangeData.hh

@@ -15,7 +15,9 @@
 #include "corecel/data/CollectionBuilder.hh"
 #include "corecel/sys/ThreadId.hh"
 
+#include "BoundingBox.hh"
 #include "OrangeTypes.hh"
+#include "detail/BIHData.hh"
 #include "univ/detail/Types.hh"
 
 namespace celeritas
@@ -61,6 +63,7 @@ struct VolumeRecord
     logic_int max_intersections{0};
     logic_int flags{0};
     DaughterId daughter_id;
+
     // TODO (KENO geometry): zorder
 
     //! Flag values (bit field)
@@ -161,8 +164,10 @@ struct SimpleUnitRecord
     // Volume data [index by LocalVolumeId]
     ItemMap<LocalVolumeId, VolumeRecordId> volumes;
 
+    // Bounding Interval Hierachy tree parameters
+    detail::BIHTree bih_params;
+
     // TODO: transforms
-    // TODO: acceleration structure (bvh/kdtree/grid)
     LocalVolumeId background{};  //!< Default if not in any other volume
     bool simple_safety{};
 
@@ -273,6 +278,9 @@ struct OrangeParamsData
     Items<SurfaceType> surface_types;
     Items<Connectivity> connectivities;
     Items<VolumeRecord> volume_records;
+    Items<FastBBox> bboxes;
+    Items<detail::BIHInnerNode> bih_inner_nodes;
+    Items<detail::BIHLeafNode> bih_leaf_nodes;
 
     Items<Daughter> daughters;
     Items<Translation> translations;

src/orange/OrangeTypes.hh

@@ -17,6 +17,7 @@
 #include "corecel/Types.hh"
 #include "corecel/cont/Array.hh"
 #include "corecel/math/NumericLimits.hh"
+#include "orange/BoundingBox.hh"
 #include "orange/Types.hh"
 
 #include "Types.hh"  // IWYU pragma: export
@@ -27,15 +28,27 @@ namespace celeritas
 // TYPE ALIASES
 //---------------------------------------------------------------------------//
 
+//! Real type used for acceleration
+using fast_real_type = float;
+
 //! Integer type for volume CSG tree representation
 using logic_int = unsigned short int;
 
+//! Identifier for a BIHNode objects
+using BIHNodeId = OpaqueId<struct BIHNode>;
+
 //! Identifier for a daughter universe
 using DaughterId = OpaqueId<struct Daughter>;
 
 //! Identifier for a face within a volume
 using FaceId = OpaqueId<struct Face>;
 
+//! Bounding box used for acceleration
+using FastBBox = BoundingBox<fast_real_type>;
+
+//! Identifier for a bounding box used for acceleration
+using FastBBoxId = OpaqueId<FastBBox>;
+
 //! Identifier for the current "level", i.e., depth of embedded universe
 using LevelId = OpaqueId<struct Level>;
 

src/orange/detail/BIHBuilder.cc

@@ -0,0 +1,222 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2022-2023 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file orange/detail/BIHBuilder.cc
+//---------------------------------------------------------------------------//
+#include "BIHBuilder.hh"
+
+#include "orange/BoundingBoxUtils.hh"
+#include "orange/detail/BIHUtils.hh"
+
+namespace celeritas
+{
+namespace detail
+{
+namespace
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Variadic-templated struct for overloaded operator() calls.
+ */
+template<typename... Ts>
+struct Overload : Ts...
+{
+    using Ts::operator()...;
+};
+
+//---------------------------------------------------------------------------//
+/*!
+ * "Deduction guide" for instantiating Overload objects w/o specifying types.
+ */
+template<class... Ts>
+Overload(Ts&&...) -> Overload<Ts...>;
+}  // namespace
+//---------------------------------------------------------------------------//
+/*!
+ * Construct from a Storage object
+ */
+BIHBuilder::BIHBuilder(BIHStorage storage) : storage_(storage)
+{
+    CELER_EXPECT(storage_);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Create BIH Nodes.
+ */
+BIHTree BIHBuilder::operator()(VecBBox bboxes)
+{
+    // Store bounding boxes and their corresponding centers
+    CELER_EXPECT(!bboxes.empty());
+    bboxes_ = std::move(bboxes);
+    centers_.resize(bboxes_.size());
+    std::transform(bboxes_.begin(),
+                   bboxes_.end(),
+                   centers_.begin(),
+                   &celeritas::center<fast_real_type>);
+
+    VecIndices indices;
+    VecIndices inf_volids;
+
+    for (auto i : range(bboxes_.size()))
+    {
+        LocalVolumeId id(i);
+
+        if (!is_infinite(bboxes_[i]))
+        {
+            indices.push_back(id);
+        }
+        else
+        {
+            inf_volids.push_back(id);
+        }
+    }
+
+    VecNodes nodes;
+    this->construct_tree(indices, &nodes, BIHNodeId{});
+
+    auto [inner_nodes, leaf_nodes] = this->arrange_nodes(std::move(nodes));
+
+    BIHTree params;
+
+    params.bboxes = ItemMap<LocalVolumeId, FastBBoxId>(
+        make_builder(storage_.bboxes)
+            .insert_back(bboxes_.begin(), bboxes_.end()));
+
+    params.inner_nodes
+        = make_builder(storage_.inner_nodes)
+              .insert_back(inner_nodes.begin(), inner_nodes.end());
+    params.leaf_nodes = make_builder(storage_.leaf_nodes)
+                            .insert_back(leaf_nodes.begin(), leaf_nodes.end());
+
+    params.inf_volids = make_builder(storage_.local_volume_ids)
+                            .insert_back(inf_volids.begin(), inf_volids.end());
+
+    return params;
+}
+
+//---------------------------------------------------------------------------//
+// HELPER FUNCTIONS
+//---------------------------------------------------------------------------//
+/*!
+ * Recursively construct BIH nodes for a vector of bbox indices.
+ */
+void BIHBuilder::construct_tree(VecIndices const& indices,
+                                VecNodes* nodes,
+                                BIHNodeId parent) const
+{
+    using Edge = BIHInnerNode::Edge;
+
+    auto current_index = nodes->size();
+    nodes->resize(nodes->size() + 1);
+
+    BIHPartitioner partition(&bboxes_, &centers_);
+
+    if (auto p = partition(indices))
+    {
+        BIHInnerNode node;
+        node.parent = parent;
+        node.axis = p.axis;
+
+        auto ax = to_int(p.axis);
+
+        node.bounding_planes[Edge::left].position
+            = p.bboxes[Edge::left].upper()[ax];
+        node.bounding_planes[Edge::right].position
+            = p.bboxes[Edge::right].lower()[ax];
+
+        // Recursively construct the left and right branches
+
+        for (auto edge : range(Edge::size_))
+        {
+            node.bounding_planes[edge].child = BIHNodeId(nodes->size());
+            this->construct_tree(
+                p.indices[edge], nodes, BIHNodeId(current_index));
+        }
+
+        CELER_EXPECT(node);
+        (*nodes)[current_index] = node;
+    }
+    else
+    {
+        BIHLeafNode node;
+        node.parent = parent;
+        node.vol_ids = make_builder(storage_.local_volume_ids)
+                           .insert_back(indices.begin(), indices.end());
+
+        CELER_EXPECT(node);
+        (*nodes)[current_index] = node;
+    }
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Seperate nodes into inner and leaf vectors and renumber accordingly.
+ */
+BIHBuilder::ArrangedNodes BIHBuilder::arrange_nodes(VecNodes nodes) const
+{
+    VecInnerNodes inner_nodes;
+    VecLeafNodes leaf_nodes;
+
+    std::vector<bool> is_leaf;
+    std::vector<BIHNodeId> new_ids;
+
+    auto visit_node
+        = Overload{[&](BIHInnerNode const& node) {
+                       new_ids.push_back(BIHNodeId(inner_nodes.size()));
+                       inner_nodes.push_back(node);
+                       is_leaf.push_back(false);
+                   },
+                   [&](BIHLeafNode const& node) {
+                       new_ids.push_back(BIHNodeId(leaf_nodes.size()));
+                       leaf_nodes.push_back(node);
+                       is_leaf.push_back(true);
+                   }};
+
+    for (auto const& node : nodes)
+    {
+        std::visit(visit_node, node);
+    }
+
+    size_type offset = inner_nodes.size();
+
+    for (auto i : range(nodes.size()))
+    {
+        if (is_leaf[i])
+        {
+            new_ids[i] = new_ids[i] + offset;
+        }
+    }
+
+    for (auto& inner_node : inner_nodes)
+    {
+        for (auto edge : range(BIHInnerNode::Edge::size_))
+        {
+            inner_node.bounding_planes[edge].child
+                = new_ids[inner_node.bounding_planes[edge].child.unchecked_get()];
+        }
+
+        // Handle root node
+        if (inner_node.parent)
+        {
+            inner_node.parent = new_ids[inner_node.parent.unchecked_get()];
+        }
+    }
+
+    for (auto& leaf_node : leaf_nodes)
+    {
+        // Handle where the entire tree is a single leaf node
+        if (leaf_node.parent)
+        {
+            leaf_node.parent = new_ids[leaf_node.parent.unchecked_get()];
+        }
+    }
+
+    return {std::move(inner_nodes), std::move(leaf_nodes)};
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace detail
+}  // namespace celeritas