started sketching index algorithm, TODO: make a generic function whic…

…h adds two tuples together and returns a new one, currently the offset shift is hard coded for 2D distributed arrays

include/bits/communication/distributed_array.hpp

@@ -82,6 +82,19 @@ static inline size_t local_element_count(const DistributedArray<N, T>& array) {
     return size;
 }
 
+///
+///@brief Returns the global element count (without padding) of the input
+/// distributed array.
+///
+///@param array The input array to query the global element count from.
+///@return size_t The global element count without padding.
+///
+template <size_t N, class T>
+static inline size_t global_element_count(const DistributedArray<N, T>& array) {
+
+    return flat_size(array.topology().get_domain().get_extent());
+}
+
 ///
 ///@brief Returns the local capacity (with padding) held by the input array
 /// distributed array.
@@ -313,4 +326,35 @@ array) {
     // return global;
 }
 */
+
+template <class ExecutionPolicy, size_t N, class T, class UnaryFunction>
+static void for_each(ExecutionPolicy&&       policy,
+                     DistributedArray<N, T>& arr,
+                     UnaryFunction           f) {
+
+    for (auto span : make_subspans(arr)) { for_each(policy, span, f); }
+}
+
+template <class ExecutionPolicy, size_t N, class T, class BinaryIndexFunction>
+static void for_each_indexed(ExecutionPolicy&&       policy,
+                             DistributedArray<N, T>& arr,
+                             BinaryIndexFunction     f) {
+
+    auto boxes = arr.local_boxes();
+    auto subspans = make_subspans(arr);
+    for (size_t i = 0; i < subspans.size(); ++i){
+        auto offset = boxes[i].box.m_begin;
+        auto span = subspans[i];
+
+        auto F = [=](auto md_idx){
+            auto copy = md_idx;
+            std::get<0>(copy) += std::get<0>(offset);
+            std::get<1>(copy) += std::get<1>(offset);
+            f(copy, span(md_idx));
+        };
+        detail::md_for_each(policy, all_indices(span), F);
+    }
+
+}
+
 } // namespace jada

include/bits/communication/topology.hpp

@@ -327,4 +327,7 @@ auto make_subspans(const Data& data, const Topology<N>& topo, int rank) {
     return ret;
 }
 
+
+
+
 } // namespace jada

test/test_communication.cpp

@@ -865,6 +865,57 @@ TEST_CASE("Test DistributedArray")
     }
     */
 
+
+    SECTION("algorithms"){
+
+        index_type nj = 2;
+        index_type ni = 3;
+
+        std::array<index_type, 2> bpad{1,0};
+        std::array<index_type, 2> epad{4,2};
+
+        Box<2> domain({0,0}, {nj, ni});
+        auto topo = decompose(domain, mpi::world_size(), {false, false});
+
+        const std::vector<int> org_data(size_t(nj * ni), 0);
+
+
+        SECTION("for_each"){
+            auto arr = distribute(org_data, topo, mpi::get_world_rank(), bpad, epad);
+            size_t n = global_element_count(arr);
+            auto op = [](auto& e){e = 43;};
+            for_each(std::execution::par_unseq, arr, op);
+
+            CHECK
+            (
+                to_vector(arr) == std::vector<int>(n, 43)
+            );
+
+        }
+        SECTION("for_each_indexed"){
+
+            auto arr = distribute(org_data, topo, mpi::get_world_rank(), bpad, epad);
+
+            auto op = [](auto idx, int& v){
+                auto j = std::get<0>(idx);
+                auto i = std::get<1>(idx);
+                v = int(i + j);
+            };
+
+            for_each_indexed(std::execution::par_unseq, arr, op);
+
+            std::vector<int> correct =
+            {
+                0,1,2,
+                1,2,3
+            };
+
+            CHECK(to_vector(arr) == correct);
+
+        }
+
+   }
+
 
 
 }