Merge pull request #538 from ValeevGroup/evaleev/fix/dist-array-fill-zero-is-sparse

DistArray::fill(0) produces an empty array if sparse

Author: evaleev

doc/dox/dev/Basic-Programming.md

@@ -533,7 +533,7 @@ for(TArrayD::iterator it = array.begin(); it != array.end(); ++it) {
 ```
 The outer loop in this example iterates over the local tiles of `array`. Within the loop body we first create an tile (an object of `TArrayD::value_type` type, which in this case is `TA::Tensor<double>`). Then we loop over its elements and assign each to zero.
 
-N.B. Of course, filling a DistArray with a constant is such a common use case that there's already a method for exactly that: `array.fill(0.0)`.
+N.B. Of course, filling a DistArray with a constant is such a common use case that there's already a method for exactly that: `array.fill(0.0)`. Note that for sparse arrays `fill(0.0)` will produce an empty array with no tiles, rather than an array full of tiles filled with zeroes.
 
 There are more serious issues with the last example. First, it is too verbose. Second, it's not generic enough (i.e. trying to reuse it for a sparse DistArray would require changing a few lines). Both issues can be solved by using modern C++ features:
 ```.cpp

src/TiledArray/dist_array.h

@@ -918,6 +918,9 @@ class DistArray : public madness::archive::ParallelSerializableObject {
 
   /// Fill all local tiles with the specified value
 
+  /// For sparse arrays, if \p value is zero (i.e., equal to
+  /// `element_type()`), the array is empty (no nonzero tiles)
+  /// and 0 is returned.
   /// \tparam fence If Fence::No, the operation will return early,
   ///         before the tasks have completed
   /// \param[in] value What each local tile should be filled with.
@@ -930,6 +933,15 @@ class DistArray : public madness::archive::ParallelSerializableObject {
   template <Fence fence = Fence::No>
   std::int64_t fill(const element_type& value = numeric_type(),
                     bool skip_set = false) {
+    // for sparse arrays filled with zero, replace with an empty array
+    if constexpr (!is_dense_v<Policy>) {
+      if (value == element_type()) {
+        *this = DistArray(
+            world(), trange(),
+            shape_type(typename shape_type::value_type(0), trange()), pmap());
+        return 0;
+      }
+    }
     return fill_local<fence>(value, skip_set);
   }
 

tests/dist_array.cpp

@@ -1007,4 +1007,16 @@ BOOST_AUTO_TEST_CASE(size_of) {
   BOOST_REQUIRE(sz0 == sz0_expected);
 }
 
+BOOST_FIXTURE_TEST_CASE(fill_zero_sparse, ArrayFixture) {
+  // construct a sparse array with some non-zero tiles and fill it
+  SpArrayN as(world, tr, TiledArray::SparseShape<float>(shape_tensor, tr));
+  as.fill(1);
+  BOOST_CHECK_GT(as.shape().nnz(), 0);
+
+  // fill with zero should result in an empty array (no tiles)
+  as.fill(0);
+  BOOST_CHECK_EQUAL(as.shape().nnz(), 0);
+  BOOST_CHECK(as.begin() == as.end());
+}
+
 BOOST_AUTO_TEST_SUITE_END()