Merge pull request #992 from mrlag31/mls-extra-fixes

Moving least squares (extra fixes)

examples/moving_least_squares/CMakeLists.txt

@@ -1,3 +1,3 @@
 add_executable(ArborX_Example_MovingLeastSquares.exe moving_least_squares.cpp)
 target_link_libraries(ArborX_Example_MovingLeastSquares.exe ArborX::ArborX)
-add_test(NAME ArborX_Example_MovingLeastSquares COMMAND ArborX_Example_MovingLeastSquares.exe)
+add_test(NAME ArborX_Example_MovingLeastSquares COMMAND ArborX_Example_MovingLeastSquares.exe)

examples/moving_least_squares/moving_least_squares.cpp

@@ -110,4 +110,4 @@ int main(int argc, char *argv[])
             << diff(2) << '\n';
 
   return 0;
-}
+}

src/geometry/ArborX_GeometryTraits.hpp

@@ -90,7 +90,7 @@ struct is_triangle : std::is_same<typename tag<Geometry>::type, TriangleTag>
 {};
 
 template <typename Geometry>
-KOKKOS_FUNCTION constexpr void check_valid_geometry_traits(Geometry const &)
+void check_valid_geometry_traits(Geometry const &)
 {
   static_assert(
       !Kokkos::is_detected<DimensionNotSpecializedArchetypeAlias, Geometry>{},

src/interpolation/ArborX_InterpMovingLeastSquares.hpp

@@ -31,40 +31,58 @@ namespace ArborX::Interpolation::Details
 {
 
 // This is done to avoid a clash with another predicates access trait
-template <typename Points>
-struct MLSTargetPointsPredicateWrapper
+template <typename TargetAccess>
+struct MLSPredicateWrapper
 {
-  Points target_points;
+  TargetAccess target_access;
   int num_neighbors;
 };
 
+// Functor used in the tree query to create the 2D source view and indices
+template <typename SourceView, typename IndicesView, typename CounterView>
+struct MLSSearchNeighborsCallback
+{
+  SourceView source_view;
+  IndicesView indices;
+  CounterView counter;
+
+  using SourcePoint = typename SourceView::non_const_value_type;
+
+  template <typename Predicate>
+  KOKKOS_FUNCTION void
+  operator()(Predicate const &predicate,
+             ArborX::PairValueIndex<SourcePoint> const &primitive) const
+  {
+    int const target = getData(predicate);
+    int const source = primitive.index;
+    auto count = Kokkos::atomic_fetch_add(&counter(target), 1);
+    indices(target, count) = source;
+    source_view(target, count) = primitive.value;
+  }
+};
+
 } // namespace ArborX::Interpolation::Details
 
 namespace ArborX
 {
 
-template <typename Points>
-struct AccessTraits<
-    Interpolation::Details::MLSTargetPointsPredicateWrapper<Points>,
-    PredicatesTag>
+template <typename TargetAccess>
+struct AccessTraits<Interpolation::Details::MLSPredicateWrapper<TargetAccess>,
+                    PredicatesTag>
 {
-  KOKKOS_INLINE_FUNCTION static auto size(
-      Interpolation::Details::MLSTargetPointsPredicateWrapper<Points> const &tp)
+  using Self = Interpolation::Details::MLSPredicateWrapper<TargetAccess>;
+
+  KOKKOS_FUNCTION static auto size(Self const &tp)
   {
-    return AccessTraits<Points, PrimitivesTag>::size(tp.target_points);
+    return tp.target_access.size();
   }
 
-  KOKKOS_INLINE_FUNCTION static auto
-  get(Interpolation::Details::MLSTargetPointsPredicateWrapper<Points> const &tp,
-      int const i)
+  KOKKOS_FUNCTION static auto get(Self const &tp, int const i)
   {
-    return nearest(
-        AccessTraits<Points, PrimitivesTag>::get(tp.target_points, i),
-        tp.num_neighbors);
+    return attach(nearest(tp.target_access(i), tp.num_neighbors), i);
   }
 
-  using memory_space =
-      typename AccessTraits<Points, PrimitivesTag>::memory_space;
+  using memory_space = typename TargetAccess::memory_space;
 };
 
 } // namespace ArborX
@@ -77,11 +95,11 @@ class MovingLeastSquares
 {
 public:
   template <typename ExecutionSpace, typename SourcePoints,
-            typename TargetPoints, typename CRBF = CRBF::Wendland<0>,
+            typename TargetPoints, typename CRBFunc = CRBF::Wendland<0>,
             typename PolynomialDegree = PolynomialDegree<2>>
   MovingLeastSquares(ExecutionSpace const &space,
                      SourcePoints const &source_points,
-                     TargetPoints const &target_points, CRBF = {},
+                     TargetPoints const &target_points, CRBFunc = {},
                      PolynomialDegree = {},
                      std::optional<int> num_neighbors = std::nullopt)
   {
@@ -95,105 +113,53 @@ class MovingLeastSquares
 
     // SourcePoints is an access trait of points
     ArborX::Details::check_valid_access_traits(PrimitivesTag{}, source_points);
-    using src_acc = AccessTraits<SourcePoints, PrimitivesTag>;
-    static_assert(KokkosExt::is_accessible_from<typename src_acc::memory_space,
-                                                ExecutionSpace>::value,
-                  "Source points must be accessible from the execution space");
-    using src_point =
-        typename ArborX::Details::AccessTraitsHelper<src_acc>::type;
-    GeometryTraits::check_valid_geometry_traits(src_point{});
-    static_assert(GeometryTraits::is_point<src_point>::value,
+    using SourceAccess =
+        ArborX::Details::AccessValues<SourcePoints, PrimitivesTag>;
+    static_assert(
+        KokkosExt::is_accessible_from<typename SourceAccess::memory_space,
+                                      ExecutionSpace>::value,
+        "Source points must be accessible from the execution space");
+    using SourcePoint = typename SourceAccess::value_type;
+    GeometryTraits::check_valid_geometry_traits(SourcePoint{});
+    static_assert(GeometryTraits::is_point<SourcePoint>::value,
                   "Source points elements must be points");
-    static constexpr int dimension = GeometryTraits::dimension_v<src_point>;
+    static constexpr int dimension = GeometryTraits::dimension_v<SourcePoint>;
 
     // TargetPoints is an access trait of points
     ArborX::Details::check_valid_access_traits(PrimitivesTag{}, target_points);
-    using tgt_acc = AccessTraits<TargetPoints, PrimitivesTag>;
-    static_assert(KokkosExt::is_accessible_from<typename tgt_acc::memory_space,
-                                                ExecutionSpace>::value,
-                  "Target points must be accessible from the execution space");
-    using tgt_point =
-        typename ArborX::Details::AccessTraitsHelper<tgt_acc>::type;
-    GeometryTraits::check_valid_geometry_traits(tgt_point{});
-    static_assert(GeometryTraits::is_point<tgt_point>::value,
+    using TargetAccess =
+        ArborX::Details::AccessValues<TargetPoints, PrimitivesTag>;
+    static_assert(
+        KokkosExt::is_accessible_from<typename TargetAccess::memory_space,
+                                      ExecutionSpace>::value,
+        "Target points must be accessible from the execution space");
+    using TargetPoint = typename TargetAccess::value_type;
+    GeometryTraits::check_valid_geometry_traits(TargetPoint{});
+    static_assert(GeometryTraits::is_point<TargetPoint>::value,
                   "Target points elements must be points");
-    static_assert(dimension == GeometryTraits::dimension_v<tgt_point>,
+    static_assert(dimension == GeometryTraits::dimension_v<TargetPoint>,
                   "Target and source points must have the same dimension");
 
-    int num_neighbors_val =
+    _num_neighbors =
         num_neighbors ? *num_neighbors
                       : Details::polynomialBasisSize<dimension,
                                                      PolynomialDegree::value>();
 
-    int const num_targets = tgt_acc::size(target_points);
-    _source_size = src_acc::size(source_points);
+    TargetAccess target_access{target_points};
+    SourceAccess source_access{source_points};
+
+    _num_targets = target_access.size();
+    _source_size = source_access.size();
     // There must be enough source points
-    KOKKOS_ASSERT(0 < num_neighbors_val && num_neighbors_val <= _source_size);
+    KOKKOS_ASSERT(0 < _num_neighbors && _num_neighbors <= _source_size);
 
-    // Organize the source points as a tree
-    BoundingVolumeHierarchy<MemorySpace, ArborX::PairValueIndex<src_point>>
-        source_tree(space, ArborX::Experimental::attach_indices(source_points));
-
-    // Create the predicates
-    Details::MLSTargetPointsPredicateWrapper<TargetPoints> predicates{
-        target_points, num_neighbors_val};
-
-    // Query the source
-    Kokkos::View<int *, MemorySpace> indices(
-        "ArborX::MovingLeastSquares::indices", 0);
-    Kokkos::View<int *, MemorySpace> offsets(
-        "ArborX::MovingLeastSquares::offsets", 0);
-    source_tree.query(space, predicates,
-                      ArborX::Details::LegacyDefaultCallback{}, indices,
-                      offsets);
-
-    // Fill in the value indices object so values can be transferred from a 1D
-    // source data to a properly distributed 2D array for each target.
-    auto const source_view = fillValuesIndicesAndGetSourceView(
-        space, indices, offsets, num_targets, num_neighbors_val, source_points);
+    // Search for neighbors and get the arranged source points
+    auto source_view = searchNeighbors(space, source_access, target_access);
 
     // Compute the moving least squares coefficients
-    _coeffs = Details::movingLeastSquaresCoefficients<
-        CRBF, PolynomialDegree, FloatingCalculationType, MemorySpace>(
-        space, source_view, target_points);
-  }
-
-  template <typename ExecutionSpace, typename SourcePoints>
-  Kokkos::View<typename ArborX::Details::AccessTraitsHelper<
-                   AccessTraits<SourcePoints, PrimitivesTag>>::type **,
-               MemorySpace>
-  fillValuesIndicesAndGetSourceView(
-      ExecutionSpace const &space,
-      Kokkos::View<int *, MemorySpace> const &indices,
-      Kokkos::View<int *, MemorySpace> const &offsets, int const num_targets,
-      int const num_neighbors, SourcePoints const &source_points)
-  {
-    auto guard = Kokkos::Profiling::ScopedRegion(
-        "ArborX::MovingLeastSquares::fillValuesIndicesAndGetSourceView");
-
-    using src_acc = AccessTraits<SourcePoints, PrimitivesTag>;
-    using src_point =
-        typename ArborX::Details::AccessTraitsHelper<src_acc>::type;
-
-    _values_indices = Kokkos::View<int **, MemorySpace>(
-        Kokkos::view_alloc(Kokkos::WithoutInitializing,
-                           "ArborX::MovingLeastSquares::values_indices"),
-        num_targets, num_neighbors);
-    Kokkos::View<src_point **, MemorySpace> source_view(
-        Kokkos::view_alloc(Kokkos::WithoutInitializing,
-                           "ArborX::MovingLeastSquares::source_view"),
-        num_targets, num_neighbors);
-    Kokkos::parallel_for(
-        "ArborX::MovingLeastSquares::values_indices_and_source_view_fill",
-        Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<2>>(
-            space, {0, 0}, {num_targets, num_neighbors}),
-        KOKKOS_CLASS_LAMBDA(int const i, int const j) {
-          auto index = indices(offsets(i) + j);
-          _values_indices(i, j) = index;
-          source_view(i, j) = src_acc::get(source_points, index);
-        });
-
-    return source_view;
+    _coeffs = Details::movingLeastSquaresCoefficients<CRBFunc, PolynomialDegree,
+                                                      FloatingCalculationType>(
+        space, source_view, target_access._values);
   }
 
   template <typename ExecutionSpace, typename SourceValues,
@@ -233,27 +199,66 @@ class MovingLeastSquares
     // original input
     KOKKOS_ASSERT(_source_size == source_values.extent_int(0));
 
-    using value_t = typename ApproxValues::non_const_value_type;
-
-    int const num_targets = _values_indices.extent(0);
-    int const num_neighbors = _values_indices.extent(1);
+    // Approx values must have the correct size
+    KOKKOS_ASSERT(approx_values.extent_int(0) == _num_targets);
 
-    KokkosExt::reallocWithoutInitializing(space, approx_values, num_targets);
+    using Value = typename ApproxValues::non_const_value_type;
 
     Kokkos::parallel_for(
         "ArborX::MovingLeastSquares::target_interpolation",
-        Kokkos::RangePolicy<ExecutionSpace>(space, 0, num_targets),
+        Kokkos::RangePolicy<ExecutionSpace>(space, 0, _num_targets),
         KOKKOS_CLASS_LAMBDA(int const i) {
-          value_t tmp = 0;
-          for (int j = 0; j < num_neighbors; j++)
-            tmp += _coeffs(i, j) * source_values(_values_indices(i, j));
+          Value tmp = 0;
+          for (int j = 0; j < _num_neighbors; j++)
+            tmp += _coeffs(i, j) * source_values(_indices(i, j));
           approx_values(i) = tmp;
         });
   }
 
 private:
+  template <typename ExecutionSpace, typename SourceAccess,
+            typename TargetAccess>
+  auto searchNeighbors(ExecutionSpace const &space,
+                       SourceAccess const &source_access,
+                       TargetAccess const &target_access)
+  {
+    auto guard = Kokkos::Profiling::ScopedRegion(
+        "ArborX::MovingLeastSquares::searchNeighbors");
+
+    // Organize the source points as a tree
+    using SourcePoint = typename SourceAccess::value_type;
+    BoundingVolumeHierarchy<MemorySpace, ArborX::PairValueIndex<SourcePoint>>
+        source_tree(space, ArborX::Experimental::attach_indices(source_access));
+
+    // Create the predicates
+    Details::MLSPredicateWrapper<TargetAccess> predicates{target_access,
+                                                          _num_neighbors};
+
+    // Create the callback
+    Kokkos::View<SourcePoint **, MemorySpace> source_view(
+        Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                           "ArborX::MovingLeastSquares::source_view"),
+        _num_targets, _num_neighbors);
+    _indices = Kokkos::View<int **, MemorySpace>(
+        Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                           "ArborX::MovingLeastSquares::indices"),
+        _num_targets, _num_neighbors);
+    Kokkos::View<int *, MemorySpace> counter(
+        "ArborX::MovingLeastSquares::counter", _num_targets);
+    Details::MLSSearchNeighborsCallback<decltype(source_view),
+                                        decltype(_indices), decltype(counter)>
+        callback{source_view, _indices, counter};
+
+    // Query the source tree
+    source_tree.query(space, predicates, callback);
+
+    return source_view;
+  }
+
   Kokkos::View<FloatingCalculationType **, MemorySpace> _coeffs;
-  Kokkos::View<int **, MemorySpace> _values_indices;
+  Kokkos::View<int **, MemorySpace> _indices;
+  int _num_targets;
+  int _num_neighbors;
   int _source_size;
 };
 

src/interpolation/details/ArborX_InterpDetailsCompactRadialBasisFunction.hpp

@@ -12,6 +12,9 @@
 #ifndef ARBORX_INTERP_DETAILS_COMPACT_RADIAL_BASIS_FUNCTION_HPP
 #define ARBORX_INTERP_DETAILS_COMPACT_RADIAL_BASIS_FUNCTION_HPP
 
+#include <ArborX_DetailsAlgorithms.hpp>
+#include <ArborX_GeometryTraits.hpp>
+
 #include <Kokkos_Core.hpp>
 
 #include <initializer_list>
@@ -40,20 +43,22 @@ namespace CRBF
 {
 
 #define CRBF_DECL(NAME)                                                        \
-  template <std::size_t>                                                       \
+  template <std::size_t N>                                                     \
   struct NAME;
 
 #define CRBF_DEF(NAME, N, FUNC)                                                \
   template <>                                                                  \
   struct NAME<N>                                                               \
   {                                                                            \
-    template <typename T>                                                      \
+    template <std::size_t DIM, typename T>                                     \
     KOKKOS_INLINE_FUNCTION static constexpr T evaluate(T const y)              \
     {                                                                          \
       /* We force the input to be between 0 and 1.                             \
          Because CRBF(-a) = CRBF(a) = CRBF(|a|), we take the absolute value    \
          and clamp the range to [0, 1] before entering in the definition of    \
-         the CRBF. */                                                          \
+         the CRBF.                                                             \
+         We also template the internal function on the dimension as CRBFs      \
+         depend on the point's dimensionality. */                              \
       T const x = Kokkos::min(Kokkos::abs(y), T(1));                           \
       return Kokkos::abs(FUNC);                                                \
     }                                                                          \
@@ -90,8 +95,20 @@ CRBF_DEF(Buhmann, 4,
 #undef CRBF_DEF
 #undef CRBF_DECL
 
+template <typename CRBFunc, typename Point>
+KOKKOS_INLINE_FUNCTION constexpr auto
+evaluate(Point const &point,
+         typename GeometryTraits::coordinate_type<Point>::type const radius)
+{
+  static_assert(GeometryTraits::is_point<Point>::value,
+                "point must be a point");
+  constexpr std::size_t dim = GeometryTraits::dimension_v<Point>;
+  return CRBFunc::template evaluate<dim>(
+      ArborX::Details::distance(point, Point{}) / radius);
+}
+
 } // namespace CRBF
 
 } // namespace ArborX::Interpolation
 
-#endif
+#endif