Merge pull request #459 from ValeevGroup/gaudel/feature/tot_array_norm

Generalize `TA::squared_norm`

src/TiledArray/dist_array.h

@@ -1224,6 +1224,32 @@ class DistArray : public madness::archive::ParallelSerializableObject {
     return TiledArray::expressions::TsrExpr<DistArray>(*this, vars);
   }
 
+  /// Create a tensor expression from an annotation (possibly free of
+  /// inner-tensor sub-annotation).
+
+  /// \brief This method creates a tensor expression but does not insist the
+  ///        annotation to be bipartite (outer and inner tensor annotations).
+  /// \param vars A string with a comma-separated list of variables.
+  /// \note Only use for unary evaluations when the indexing of the inner
+  ///       tensors is not significant, eg. norm computation.
+  ///
+  auto make_tsrexpr(const std::string& vars) {
+    return TiledArray::expressions::TsrExpr<DistArray>(*this, vars);
+  }
+
+  /// Create a tensor expression from an annotation (possibly free of
+  /// inner-tensor sub-annotation).
+
+  /// \brief This method creates a tensor expression but does not insist the
+  ///        annotation to be bipartite (outer and inner tensor annotations).
+  /// \param vars A string with a comma-separated list of variables.
+  /// \note Only use for unary evaluations when the indexing of the inner
+  ///       tensors is not significant, eg. norm computation.
+  ///
+  auto make_tsrexpr(const std::string& vars) const {
+    return TiledArray::expressions::TsrExpr<const DistArray>(*this, vars);
+  }
+
   /// \deprecated use DistArray::world()
   [[deprecated]] World& get_world() const { return world(); }
 
@@ -1892,32 +1918,32 @@ size_t volume(const DistArray<Tile, Policy>& array) {
 
 template <typename Tile, typename Policy>
 auto abs_min(const DistArray<Tile, Policy>& a) {
-  return a(detail::dummy_annotation(rank(a))).abs_min();
+  return a.make_tsrexpr(detail::dummy_annotation(rank(a))).abs_min();
 }
 
 template <typename Tile, typename Policy>
 auto abs_max(const DistArray<Tile, Policy>& a) {
-  return a(detail::dummy_annotation(rank(a))).abs_max();
+  return a.make_tsrexpr(detail::dummy_annotation(rank(a))).abs_max();
 }
 
 template <typename Tile, typename Policy>
 auto dot(const DistArray<Tile, Policy>& a, const DistArray<Tile, Policy>& b) {
-  return (a(detail::dummy_annotation(rank(a)))
-              .dot(b(detail::dummy_annotation(rank(b)))))
-      .get();
+  auto&& expr_a = a.make_tsrexpr(detail::dummy_annotation(rank(a)));
+  auto&& expr_b = b.make_tsrexpr(detail::dummy_annotation(rank(b)));
+  return expr_a.dot(expr_b).get();
 }
 
 template <typename Tile, typename Policy>
 auto inner_product(const DistArray<Tile, Policy>& a,
                    const DistArray<Tile, Policy>& b) {
-  return (a(detail::dummy_annotation(rank(a)))
-              .inner_product(b(detail::dummy_annotation(rank(b)))))
-      .get();
+  auto&& expr_a = a.make_tsrexpr(detail::dummy_annotation(rank(a)));
+  auto&& expr_b = b.make_tsrexpr(detail::dummy_annotation(rank(b)));
+  return expr_a.inner_product(expr_b).get();
 }
 
 template <typename Tile, typename Policy>
 auto squared_norm(const DistArray<Tile, Policy>& a) {
-  return a(detail::dummy_annotation(rank(a))).squared_norm();
+  return a.make_tsrexpr(detail::dummy_annotation(rank(a))).squared_norm();
 }
 
 template <typename Tile, typename Policy>

tests/dist_array.cpp

@@ -883,4 +883,65 @@ BOOST_AUTO_TEST_CASE(volume) {
   BOOST_REQUIRE(vol == TA::volume(array));
 }
 
+BOOST_AUTO_TEST_CASE(reduction) {
+  using Numeric = double;
+  using T = Tensor<Numeric>;
+  using ToT = Tensor<T>;
+  using Policy = SparsePolicy;
+  using ArrayToT = DistArray<ToT, Policy>;
+
+  auto unit_T = [](Range const& rng) { return T(rng, Numeric{1}); };
+
+  auto unit_ToT = [unit_T](Range const& rngo, Range const& rngi) {
+    return ToT(rngo, unit_T(rngi));
+  };
+
+  size_t constexpr nrows = 3;
+  size_t constexpr ncols = 4;
+  TiledRange const trange({{0, 2, 5, 7}, {0, 5, 7, 10, 12}});
+  TA_ASSERT(trange.tiles_range().extent().at(0) == nrows &&
+                trange.tiles_range().extent().at(1) == ncols,
+            "Following code depends on this condition.");
+
+  // this Range is used to construct all inner tensors of the tile with
+  // tile index @c tix.
+  auto inner_dims = [nrows, ncols](Range::index_type const& tix) -> Range {
+    static std::array<size_t, nrows> const rows{7, 8, 9};
+    static std::array<size_t, ncols> const cols{7, 8, 9, 10};
+
+    TA_ASSERT(tix.size() == 2, "Only rank-2 tensor expected.");
+    return Range({rows[tix.at(0) % nrows], cols[tix.at(1) % ncols]});
+  };
+
+  // let's make all 'diagonal' tiles zero
+  auto zero_tile = [](Range::index_type const& tix) -> bool {
+    return tix.at(0) == tix.at(1);
+  };
+
+  auto make_tile = [inner_dims,  //
+                    zero_tile,   //
+                    &trange,     //
+                    unit_ToT](auto& tile, auto const& rng) {
+    auto&& tix = trange.element_to_tile(rng.lobound());
+    if (zero_tile(tix))
+      return 0.;
+    else {
+      tile = unit_ToT(rng, inner_dims(tix));
+      return tile.norm();
+    }
+  };
+
+  auto& world = get_default_world();
+
+  // all non-zero inner tensors of this ToT array are unit (ie all
+  // inner tensors' elements are 1.)
+  auto array = make_array<ArrayToT>(world, trange, make_tile);
+
+  // since all inner tensors are filled with 1.
+  double array_norm = std::sqrt(TA::volume(array));
+
+  BOOST_REQUIRE(array_norm == TA::norm2(array));
+  BOOST_REQUIRE(array_norm = std::sqrt(TA::dot(array, array)));
+}
+
 BOOST_AUTO_TEST_SUITE_END()