Flatten data structures in geometry

cpp/dolfinx/geometry/BoundingBoxTree.cpp

@@ -1,5 +1,5 @@
-// Copyright (C) 2013-2021 Chris N. Richardson, Anders Logg, Garth N. Wells and
-// Jørgen S. Dokken
+// Copyright (C) 2013-2022 Chris N. Richardson, Anders Logg, Garth N. Wells,
+// Jørgen S. Dokken, Sarah Roggendorf
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -34,8 +34,8 @@ std::vector<std::int32_t> range(const mesh::Mesh& mesh, int tdim)
 }
 //-----------------------------------------------------------------------------
 // Compute bounding box of mesh entity
-std::array<std::array<double, 3>, 2>
-compute_bbox_of_entity(const mesh::Mesh& mesh, int dim, std::int32_t index)
+std::array<double, 6> compute_bbox_of_entity(const mesh::Mesh& mesh, int dim,
+                                             std::int32_t index)
 {
   // Get the geometrical indices for the mesh entity
   std::span<const double> xg = mesh.geometry().x();
@@ -45,129 +45,114 @@ compute_bbox_of_entity(const mesh::Mesh& mesh, int dim, std::int32_t index)
   const std::vector<std::int32_t> vertex_indices
       = mesh::entities_to_geometry(mesh, dim, entity, false);
 
-  std::array<std::array<double, 3>, 2> b;
-  b[0] = {xg[3 * vertex_indices.front()], xg[3 * vertex_indices.front() + 1],
-          xg[3 * vertex_indices.front() + 2]};
-  b[1] = b[0];
+  std::array<double, 6> b;
+  auto b0 = std::span(b).subspan(0, 3);
+  auto b1 = std::span(b).subspan(3, 6);
+  common::impl::copy_N<3>(std::next(xg.begin(), 3 * vertex_indices.front()),
+                          b0.begin());
+  common::impl::copy_N<3>(std::next(xg.begin(), 3 * vertex_indices.front()),
+                          b1.begin());
 
   // Compute min and max over vertices
   for (int local_vertex : vertex_indices)
   {
     for (int j = 0; j < 3; ++j)
     {
-      b[0][j] = std::min(b[0][j], xg[3 * local_vertex + j]);
-      b[1][j] = std::max(b[1][j], xg[3 * local_vertex + j]);
+      b0[j] = std::min(b0[j], xg[3 * local_vertex + j]);
+      b1[j] = std::max(b1[j], xg[3 * local_vertex + j]);
     }
   }
 
   return b;
 }
 //-----------------------------------------------------------------------------
 // Compute bounding box of bounding boxes
-std::array<std::array<double, 3>, 2> compute_bbox_of_bboxes(
-    const std::span<const std::pair<std::array<std::array<double, 3>, 2>,
-                                    std::int32_t>>& leaf_bboxes)
+std::array<double, 6> compute_bbox_of_bboxes(
+    const std::span<const std::pair<std::array<double, 6>, std::int32_t>>&
+        leaf_bboxes)
 {
   // Compute min and max over remaining boxes
-  std::array<std::array<double, 3>, 2> b;
-  b[0] = leaf_bboxes[0].first[0];
-  b[1] = leaf_bboxes[0].first[1];
+  std::array<double, 6> b = leaf_bboxes.front().first;
+
   for (auto& box : leaf_bboxes)
   {
-    std::transform(box.first[0].begin(), box.first[0].end(), b[0].begin(),
-                   b[0].begin(),
+    std::transform(box.first.cbegin(), std::next(box.first.cbegin(), 3),
+                   b.cbegin(), b.begin(),
                    [](double a, double b) { return std::min(a, b); });
-    std::transform(box.first[1].begin(), box.first[1].end(), b[1].begin(),
-                   b[1].begin(),
+    std::transform(std::next(box.first.cbegin(), 3), box.first.cend(),
+                   std::next(b.cbegin(), 3), std::next(b.begin(), 3),
                    [](double a, double b) { return std::max(a, b); });
   }
 
   return b;
 }
 //------------------------------------------------------------------------------
 int _build_from_leaf(
-    std::span<std::pair<std::array<std::array<double, 3>, 2>, std::int32_t>>
-        leaf_bboxes,
-    std::vector<std::array<int, 2>>& bboxes,
-    std::vector<double>& bbox_coordinates)
+    std::span<std::pair<std::array<double, 6>, std::int32_t>> leaf_bboxes,
+    std::vector<int>& bboxes, std::vector<double>& bbox_coordinates)
 {
   if (leaf_bboxes.size() == 1)
   {
     // Reached leaf
 
     // Get bounding box coordinates for leaf
-    const std::int32_t entity_index = leaf_bboxes[0].second;
-    const std::array<double, 3> b0 = leaf_bboxes[0].first[0];
-    const std::array<double, 3> b1 = leaf_bboxes[0].first[1];
+    const auto [b, entity_index] = leaf_bboxes.front();
 
     // Store bounding box data
-    bboxes.push_back({entity_index, entity_index});
-    bbox_coordinates.insert(bbox_coordinates.end(), b0.begin(), b0.end());
-    bbox_coordinates.insert(bbox_coordinates.end(), b1.begin(), b1.end());
-    return bboxes.size() - 1;
+    bboxes.push_back(entity_index);
+    bboxes.push_back(entity_index);
+    common::impl::copy_N<6>(b.begin(), std::back_inserter(bbox_coordinates));
+    return bboxes.size() / 2 - 1;
   }
   else
   {
     // Compute bounding box of all bounding boxes
-    std::array<std::array<double, 3>, 2> b
-        = compute_bbox_of_bboxes(leaf_bboxes);
-    std::array<double, 3> b0 = b[0];
-    std::array<double, 3> b1 = b[1];
+    std::array<double, 6> b = compute_bbox_of_bboxes(leaf_bboxes);
 
     // Sort bounding boxes along longest axis
     std::array<double, 3> b_diff;
-    std::transform(b1.begin(), b1.end(), b0.begin(), b_diff.begin(),
-                   std::minus<double>());
+    std::transform(std::next(b.cbegin(), 3), b.cend(), b.cbegin(),
+                   b_diff.begin(), std::minus<double>());
     const std::size_t axis = std::distance(
         b_diff.begin(), std::max_element(b_diff.begin(), b_diff.end()));
 
     auto middle = std::next(leaf_bboxes.begin(), leaf_bboxes.size() / 2);
     std::nth_element(leaf_bboxes.begin(), middle, leaf_bboxes.end(),
                      [axis](auto& p0, auto& p1) -> bool
                      {
-                       double x0 = p0.first[0][axis] + p0.first[1][axis];
-                       double x1 = p1.first[0][axis] + p1.first[1][axis];
+                       double x0 = p0.first[axis] + p0.first[3 + axis];
+                       double x1 = p1.first[axis] + p1.first[3 + axis];
                        return x0 < x1;
                      });
 
     // Split bounding boxes into two groups and call recursively
     assert(!leaf_bboxes.empty());
     std::size_t part = leaf_bboxes.size() / 2;
-    std::array bbox{
-        _build_from_leaf(leaf_bboxes.first(part), bboxes, bbox_coordinates),
-        _build_from_leaf(leaf_bboxes.last(leaf_bboxes.size() - part), bboxes,
-                         bbox_coordinates)};
+    int bbox0
+        = _build_from_leaf(leaf_bboxes.first(part), bboxes, bbox_coordinates);
+    int bbox1 = _build_from_leaf(leaf_bboxes.last(leaf_bboxes.size() - part),
+                                 bboxes, bbox_coordinates);
 
     // Store bounding box data. Note that root box will be added last.
-    bboxes.push_back(bbox);
-    bbox_coordinates.insert(bbox_coordinates.end(), b0.begin(), b0.end());
-    bbox_coordinates.insert(bbox_coordinates.end(), b1.begin(), b1.end());
-    return bboxes.size() - 1;
+    bboxes.push_back(bbox0);
+    bboxes.push_back(bbox1);
+    common::impl::copy_N<6>(b.begin(), std::back_inserter(bbox_coordinates));
+    return bboxes.size() / 2 - 1;
   }
 }
 //-----------------------------------------------------------------------------
 std::pair<std::vector<std::int32_t>, std::vector<double>> build_from_leaf(
-    std::vector<std::pair<std::array<std::array<double, 3>, 2>, std::int32_t>>
-        leaf_bboxes)
+    std::vector<std::pair<std::array<double, 6>, std::int32_t>> leaf_bboxes)
 {
-  std::vector<std::array<std::int32_t, 2>> bboxes;
+  std::vector<std::int32_t> bboxes;
   std::vector<double> bbox_coordinates;
   _build_from_leaf(leaf_bboxes, bboxes, bbox_coordinates);
-
-  std::vector<std::int32_t> bbox_array(2 * bboxes.size());
-  for (std::size_t i = 0; i < bboxes.size(); ++i)
-  {
-    bbox_array[2 * i] = bboxes[i][0];
-    bbox_array[2 * i + 1] = bboxes[i][1];
-  }
-
-  return {std::move(bbox_array), std::move(bbox_coordinates)};
+  return {std::move(bboxes), std::move(bbox_coordinates)};
 }
 //-----------------------------------------------------------------------------
 int _build_from_point(
     std::span<std::pair<std::array<double, 3>, std::int32_t>> points,
-    std::vector<std::array<std::int32_t, 2>>& bboxes,
-    std::vector<double>& bbox_coordinates)
+    std::vector<std::int32_t>& bboxes, std::vector<double>& bbox_coordinates)
 {
   // Reached leaf
   if (points.size() == 1)
@@ -176,12 +161,13 @@ int _build_from_point(
 
     // Index of entity contained in leaf
     const std::int32_t c1 = points[0].second;
-    bboxes.push_back({c1, c1});
+    bboxes.push_back(c1);
+    bboxes.push_back(c1);
     bbox_coordinates.insert(bbox_coordinates.end(), points[0].first.begin(),
                             points[0].first.end());
     bbox_coordinates.insert(bbox_coordinates.end(), points[0].first.begin(),
                             points[0].first.end());
-    return bboxes.size() - 1;
+    return bboxes.size() / 2 - 1;
   }
 
   // Compute bounding box of all points
@@ -204,16 +190,17 @@ int _build_from_point(
   // Split bounding boxes into two groups and call recursively
   assert(!points.empty());
   std::size_t part = points.size() / 2;
-  std::array bbox{
-      _build_from_point(points.first(part), bboxes, bbox_coordinates),
-      _build_from_point(points.last(points.size() - part), bboxes,
-                        bbox_coordinates)};
+  std::int32_t bbox0
+      = _build_from_point(points.first(part), bboxes, bbox_coordinates);
+  std::int32_t bbox1 = _build_from_point(points.last(points.size() - part),
+                                         bboxes, bbox_coordinates);
 
   // Store bounding box data. Note that root box will be added last.
-  bboxes.push_back(bbox);
+  bboxes.push_back(bbox0);
+  bboxes.push_back(bbox1);
   bbox_coordinates.insert(bbox_coordinates.end(), b0.begin(), b0.end());
   bbox_coordinates.insert(bbox_coordinates.end(), b1.begin(), b1.end());
-  return bboxes.size() - 1;
+  return bboxes.size() / 2 - 1;
 }
 //-----------------------------------------------------------------------------
 } // namespace
@@ -242,16 +229,14 @@ BoundingBoxTree::BoundingBoxTree(const mesh::Mesh& mesh, int tdim,
   mesh.topology_mutable().create_connectivity(tdim, mesh.topology().dim());
 
   // Create bounding boxes for all mesh entities (leaves)
-  std::vector<std::pair<std::array<std::array<double, 3>, 2>, std::int32_t>>
-      leaf_bboxes;
+  std::vector<std::pair<std::array<double, 6>, std::int32_t>> leaf_bboxes;
   leaf_bboxes.reserve(entities.size());
   for (std::int32_t e : entities)
   {
-    std::array<std::array<double, 3>, 2> b
-        = compute_bbox_of_entity(mesh, tdim, e);
-    std::transform(b[0].begin(), b[0].end(), b[0].begin(),
+    std::array<double, 6> b = compute_bbox_of_entity(mesh, tdim, e);
+    std::transform(b.cbegin(), std::next(b.cbegin(), 3), b.begin(),
                    [padding](double x) { return x - padding; });
-    std::transform(b[1].begin(), b[1].end(), b[1].begin(),
+    std::transform(std::next(b.begin(), 3), b.end(), std::next(b.begin(), 3),
                    [padding](double& x) { return x + padding; });
     leaf_bboxes.emplace_back(b, e);
   }
@@ -269,16 +254,10 @@ BoundingBoxTree::BoundingBoxTree(
     : _tdim(0)
 {
   // Recursively build the bounding box tree from the leaves
-  std::vector<std::array<int, 2>> bboxes;
   if (!points.empty())
   {
-    _build_from_point(std::span(points), bboxes, _bbox_coordinates);
-    _bboxes.resize(2 * bboxes.size());
-    for (std::size_t i = 0; i < bboxes.size(); ++i)
-    {
-      _bboxes[2 * i] = bboxes[i][0];
-      _bboxes[2 * i + 1] = bboxes[i][1];
-    }
+    _bboxes.clear();
+    _build_from_point(std::span(points), _bboxes, _bbox_coordinates);
   }
 
   LOG(INFO) << "Computed bounding box tree with " << num_bboxes()
@@ -305,14 +284,13 @@ BoundingBoxTree BoundingBoxTree::create_global_tree(MPI_Comm comm) const
   MPI_Allgather(send_bbox.data(), 6, MPI_DOUBLE, recv_bbox.data(), 6,
                 MPI_DOUBLE, comm);
 
-  std::vector<std::pair<std::array<std::array<double, 3>, 2>, std::int32_t>>
-      _recv_bbox(mpi_size);
+  std::vector<std::pair<std::array<double, 6>, std::int32_t>> _recv_bbox(
+      mpi_size);
   for (std::size_t i = 0; i < _recv_bbox.size(); ++i)
   {
-    common::impl::copy_N<3>(std::next(recv_bbox.begin(), 6 * i),
-                            _recv_bbox[i].first[0].begin());
-    common::impl::copy_N<3>(std::next(recv_bbox.begin(), 6 * i + 3),
-                            _recv_bbox[i].first[1].begin());
+    common::impl::copy_N<6>(std::next(recv_bbox.begin(), 6 * i),
+                            _recv_bbox[i].first.begin());
+
     _recv_bbox[i].second = i;
   }
 
@@ -362,14 +340,18 @@ void BoundingBoxTree::tree_print(std::stringstream& s, int i) const
   }
 }
 //-----------------------------------------------------------------------------
-std::array<std::array<double, 3>, 2>
-BoundingBoxTree::get_bbox(std::size_t node) const
+std::span<const double, 6> BoundingBoxTree::get_bbox(std::size_t node) const
 {
-  std::array<std::array<double, 3>, 2> x;
-  common::impl::copy_N<3>(std::next(_bbox_coordinates.begin(), 6 * node),
-                          x[0].begin());
-  common::impl::copy_N<3>(std::next(_bbox_coordinates.begin(), 6 * node + 3),
-                          x[1].begin());
+
+  return std::span<const double, 6>(_bbox_coordinates.data() + 6 * node, 6);
+}
+//-----------------------------------------------------------------------------
+std::array<double, 6> BoundingBoxTree::copy_bbox(std::size_t node) const
+{
+  std::array<double, 6> x;
+  common::impl::copy_N<6>(std::next(_bbox_coordinates.begin(), 6 * node),
+                          x.begin());
+
   return x;
 }
 //-----------------------------------------------------------------------------

cpp/dolfinx/geometry/BoundingBoxTree.h

@@ -70,9 +70,15 @@ class BoundingBoxTree
 
   /// Return bounding box coordinates for a given node in the tree
   /// @param[in] node The bounding box node index
-  /// @return The bounding box where [0] is the lower corner and [1] is
-  /// the upper corner
-  std::array<std::array<double, 3>, 2> get_bbox(std::size_t node) const;
+  /// @return A copy of bounding box (lower_corner, upper_corner). Shape (2,3).
+  /// Flattened row-major.
+  std::array<double, 6> copy_bbox(std::size_t node) const;
+
+  /// Return bounding box coordinates for a given node in the tree
+  /// @param[in] node The bounding box node index
+  /// @return The bounding box (lower_corner, upper_corner) as a subspan of the
+  /// bounding box coordinates. Shape (2,3). Flattened row-major.
+  std::span<const double, 6> get_bbox(std::size_t node) const;
 
   /// Compute a global bounding tree (collective on comm)
   /// This can be used to find which process a point might have a
@@ -103,6 +109,19 @@ class BoundingBoxTree
     return {_bboxes[2 * node], _bboxes[2 * node + 1]};
   }
 
+  /// Get bounding box child nodes
+  ///
+  /// @param[in] node The bounding box node index
+  /// @return The indices of the two child nodes. If @p node is a leaf
+  /// nodes, then the values in the returned array are equal and
+  /// correspond to the index of the entity that the leaf node bounds,
+  /// e.g. the index of the cell that it bounds.
+  std::span<const int, 2> bbox_span(std::size_t node) const
+  {
+    assert(2 * node + 1 < _bboxes.size());
+    return std::span<const int, 2>(_bboxes.data() + 2 * node, 2);
+  }
+
 private:
   // Constructor
   BoundingBoxTree(std::vector<std::int32_t>&& bboxes,