Merge pull request #469 from ValeevGroup/evaleev/feature/make-uniform…

…-range1 introduced `make_uniform(Range1,tilesize)`
ValeevGroup · Aug 27, 2024 · c10f8d5 · c10f8d5
2 parents a09717d + cd78961
commit c10f8d5
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 10 deletions.
diff --git a/src/TiledArray/tiled_range1.h b/src/TiledArray/tiled_range1.h
@@ -244,26 +244,27 @@ class TiledRange1 {
   // clang-format off
   /// @brief makes a uniform (or, as uniform as possible) TiledRange1
 
-  /// @param[in] range_size the range size
+  /// @param[in] range the Range to be tiled
   /// @param[in] target_tile_size the desired tile size
-  /// @return TiledRange1 obtained by tiling range `[0,range_size)` into
-  /// `ntiles = (range_size + target_tile_size - 1)/target_tile_size`
-  ///         tiles; if `x = range_size % ntiles` is not zero, first `x` tiles
+  /// @return TiledRange1 obtained by tiling \p range into
+  /// `ntiles = (range.extent() + target_tile_size - 1)/target_tile_size`
+  ///         tiles; if `x = range.extent() % ntiles` is not zero, first `x` tiles
   /// have size `target_tile_size` and last
   /// `ntiles - x` tiles have size `target_tile_size - 1`, else
   /// all tiles have size `target_tile_size` .
   // clang-format on
-  static TiledRange1 make_uniform(std::size_t range_size,
+  static TiledRange1 make_uniform(const Range1& range,
                                   std::size_t target_tile_size) {
-    if (range_size > 0) {
+    const auto range_extent = range.extent();
+    if (range_extent > 0) {
       TA_ASSERT(target_tile_size > 0);
       std::size_t ntiles =
-          (range_size + target_tile_size - 1) / target_tile_size;
-      auto dv = std::div((long)(range_size + ntiles - 1), (long)ntiles);
+          (range_extent + target_tile_size - 1) / target_tile_size;
+      auto dv = std::div((long)(range_extent + ntiles - 1), (long)ntiles);
       auto avg_tile_size = dv.quot - 1, num_avg_plus_one = dv.rem + 1;
       std::vector<std::size_t> hashmarks;
       hashmarks.reserve(ntiles + 1);
-      std::size_t element = 0;
+      std::size_t element = range.lobound();
       for (auto i = 0; i < num_avg_plus_one;
            ++i, element += avg_tile_size + 1) {
         hashmarks.push_back(element);
@@ -272,12 +273,19 @@ class TiledRange1 {
            ++i, element += avg_tile_size) {
         hashmarks.push_back(element);
       }
-      hashmarks.push_back(range_size);
+      hashmarks.push_back(range.upbound());
       return TiledRange1(hashmarks.begin(), hashmarks.end());
     } else
       return TiledRange1{};
   }
 
+  /// same as make_uniform(const Range1&, std::size_t) for a 0-based range
+  /// specified by its extent
+  static TiledRange1 make_uniform(std::size_t range_extent,
+                                  std::size_t target_tile_size) {
+    return make_uniform(Range1(0, range_extent), target_tile_size);
+  }
+
   /// swapper
 
   /// \param other the range with which the contents of this range will be

diff --git a/tests/tiled_range1.cpp b/tests/tiled_range1.cpp
@@ -294,6 +294,21 @@ BOOST_AUTO_TEST_CASE(concatenation) {
 }
 
 BOOST_AUTO_TEST_CASE(make_uniform) {
+  BOOST_REQUIRE_NO_THROW(TiledRange1::make_uniform(Range1{1, 1}, 0));
+  BOOST_CHECK(TiledRange1::make_uniform(Range1{1, 1}, 0) == TiledRange1{});
+  BOOST_REQUIRE_NO_THROW(TiledRange1::make_uniform(Range1{3, 6}, 10));
+  BOOST_CHECK(TiledRange1::make_uniform(Range1{3, 6}, 10) ==
+              (TiledRange1{3, 6}));
+  BOOST_REQUIRE_NO_THROW(TiledRange1::make_uniform(Range1{10, 60}, 10));
+  BOOST_CHECK(TiledRange1::make_uniform(Range1{10, 60}, 10) ==
+              (TiledRange1{10, 20, 30, 40, 50, 60}));
+  BOOST_REQUIRE_NO_THROW(TiledRange1::make_uniform(Range1{10, 65}, 10));
+  BOOST_CHECK(TiledRange1::make_uniform(Range1{10, 65}, 10) ==
+              (TiledRange1{10, 20, 29, 38, 47, 56, 65}));
+  BOOST_REQUIRE_NO_THROW(TiledRange1::make_uniform(Range1{10, 69}, 10));
+  BOOST_CHECK(TiledRange1::make_uniform(Range1{10, 69}, 10) ==
+              (TiledRange1{10, 20, 30, 40, 50, 60, 69}));
+
   BOOST_REQUIRE_NO_THROW(TiledRange1::make_uniform(0, 0));
   BOOST_CHECK(TiledRange1::make_uniform(0, 0) == TiledRange1{});
   BOOST_REQUIRE_NO_THROW(TiledRange1::make_uniform(0, 1));