Merge pull request #1115 from aprokop/custom_distance

Add an example using custom distance for nearest search

Author: aprokop

examples/CMakeLists.txt

@@ -8,6 +8,7 @@ endif()
 add_subdirectory(access_traits)
 add_subdirectory(brute_force)
 add_subdirectory(callback)
+add_subdirectory(custom_distance)
 add_subdirectory(dbscan)
 add_subdirectory(molecular_dynamics)
 add_subdirectory(simple_intersection)

examples/README.md

@@ -5,6 +5,7 @@ Directory | Description
 `access_traits`         | provide non-View user data to ArborX
 `brute_force`           | use brute force to solve a problem
 `callback`              | execute search which calls user function on match
+`custom_distance`       | execute nearest search using a custom distance metric
 `dbscan`                | cluster a set of points using DBSCAN
 `molecular_dynamics`    | do one step of moving particles under forces
 `moving_least_squares`  | interpolate data between point clouds using MLS

examples/custom_distance/CMakeLists.txt

@@ -0,0 +1,3 @@
+add_executable(ArborX_Example_CustomDistance.exe example_custom_distance.cpp)
+target_link_libraries(ArborX_Example_CustomDistance.exe ArborX::ArborX)
+add_test(NAME ArborX_Example_CustomDistance COMMAND ArborX_Example_CustomDistance.exe)

examples/custom_distance/example_custom_distance.cpp

@@ -0,0 +1,138 @@
+/****************************************************************************
+ * Copyright (c) 2025, ArborX authors                                       *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the ArborX library. ArborX is                       *
+ * distributed under a BSD 3-clause license. For the licensing terms see    *
+ * the LICENSE file in the top-level directory.                             *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+#include <ArborX.hpp>
+
+#include <Kokkos_Core.hpp>
+
+#include <algorithm> // copy
+#include <iostream>  // cout
+#include <iterator>  // ostream_iterator
+
+// Define custom type that will be used for distance calculations
+struct CustomPoint
+{
+  float x;
+  float y;
+
+  KOKKOS_FUNCTION auto operator[](int i) const
+  {
+    KOKKOS_ASSERT(i == 0 || i == 1);
+    return (i == 0 ? x : y);
+  }
+  KOKKOS_FUNCTION auto &operator[](int i)
+  {
+    KOKKOS_ASSERT(i == 0 || i == 1);
+    return (i == 0 ? x : y);
+  }
+};
+
+template <>
+struct ArborX::GeometryTraits::dimension<CustomPoint>
+{
+  static constexpr int value = 2;
+};
+template <>
+struct ArborX::GeometryTraits::coordinate_type<CustomPoint>
+{
+  using type = float;
+};
+template <>
+struct ArborX::GeometryTraits::tag<CustomPoint>
+{
+  using type = ArborX::GeometryTraits::PointTag;
+};
+
+// Provide the distance function between indexable getter geometry and the
+// custom data type
+KOKKOS_FUNCTION auto distance(CustomPoint const &p, CustomPoint const &q)
+{
+  return Kokkos::abs(p.x - q.x) * 2 + Kokkos::abs(p.y - q.y) / 2;
+}
+
+// Provide the distance function between the bounding volume geometry and
+// the custom data type
+using BoundingVolume = ArborX::Box<2>;
+KOKKOS_FUNCTION auto distance(CustomPoint const &point,
+                              BoundingVolume const &box)
+{
+  CustomPoint projected_point{
+      Kokkos::clamp(point.x, box.minCorner()[0], box.maxCorner()[0]),
+      Kokkos::clamp(point.y, box.minCorner()[0], box.maxCorner()[0])};
+  return distance(point, projected_point);
+}
+
+// Callback to store the result indices
+struct ExtractIndex
+{
+  template <typename Query, typename Value, typename Output>
+  KOKKOS_FUNCTION void operator()(Query const &, Value const &value,
+                                  Output const &out) const
+  {
+    out(value.index);
+  }
+};
+
+int main(int argc, char *argv[])
+{
+  Kokkos::ScopeGuard guard(argc, argv);
+
+  using ExecutionSpace = Kokkos::DefaultExecutionSpace;
+  using MemorySpace = ExecutionSpace::memory_space;
+
+  ExecutionSpace space;
+
+  // In L2 metric, two bottom points are the closest.
+  // In custom metric, left two points are the closest.
+  //
+  // | 1
+  // |
+  // | x
+  // | 2 0
+  // +-----------
+  Kokkos::View<CustomPoint *, MemorySpace> points("Example::points", 3);
+  auto points_host = Kokkos::create_mirror_view(points);
+  points_host[0] = {1, 0};
+  points_host[1] = {0, 3};
+  points_host[2] = {0, 0};
+  Kokkos::deep_copy(points, points_host);
+
+  Kokkos::View<CustomPoint *, MemorySpace> query_points("Example::query_points",
+                                                        1);
+  auto query_points_host = Kokkos::create_mirror_view(query_points);
+  query_points_host[0] = {0, 1};
+  Kokkos::deep_copy(query_points, query_points_host);
+
+  ArborX::BoundingVolumeHierarchy bvh(
+      space, ArborX::Experimental::attach_indices(points));
+
+  Kokkos::View<int *, MemorySpace> offsets("Example::offsets", 0);
+  Kokkos::View<int *, MemorySpace> indices("Example::indices", 0);
+  bvh.query(space, ArborX::Experimental::make_nearest(query_points, 2),
+            ExtractIndex{}, indices, offsets);
+
+  auto offsets_host =
+      Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, offsets);
+  auto indices_host =
+      Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, indices);
+
+  // Expected output:
+  // offsets: 0 2
+  // indices: 2 1
+  std::cout << "offsets: ";
+  std::copy(offsets_host.data(), offsets_host.data() + offsets.size(),
+            std::ostream_iterator<int>(std::cout, " "));
+  std::cout << "\nindices: ";
+  std::copy(indices_host.data(), indices_host.data() + indices.size(),
+            std::ostream_iterator<int>(std::cout, " "));
+  std::cout << "\n";
+
+  return 0;
+}

src/spatial/detail/ArborX_Callbacks.hpp

@@ -12,7 +12,6 @@
 #define ARBORX_CALLBACKS_HPP
 
 #include <detail/ArborX_AccessTraits.hpp>
-#include <detail/ArborX_PairValueIndex.hpp>
 #include <detail/ArborX_Predicates.hpp> // is_valid_predicate_tag
 
 #include <Kokkos_DetectionIdiom.hpp>

src/spatial/detail/ArborX_PairValueIndex.hpp

@@ -12,10 +12,10 @@
 #ifndef ARBORX_PAIR_VALUE_INDEX_HPP
 #define ARBORX_PAIR_VALUE_INDEX_HPP
 
-#include <detail/ArborX_AccessTraits.hpp>
-
 #include <Kokkos_Macros.hpp>
 
+#include <type_traits>
+
 namespace ArborX
 {