Repository: computationalgeography/lue (language: python)

Issue:
`lue.pcraster`'s `ycoordinate` outputs wrong values
### Describe the bug

The operation outputs coordinates per partition, not per array. Within the partition the output seems OK, but the overall results are wrong.

### Expected behaviour

Same results as PCRaster's `ycoordinate`.

### Steps to reproduce

Run `ycoordinate` on a partitioned raster.

### Additional information

The problem likely is in LUE's `cell_index` algorithm.


Discussion / hints:

Fix patch:
diff --git a/source/framework/algorithm/include/lue/framework/algorithm/definition/cell_index.hpp b/source/framework/algorithm/include/lue/framework/algorithm/definition/cell_index.hpp
index 157c66a90..878d96562 100644
--- a/source/framework/algorithm/include/lue/framework/algorithm/definition/cell_index.hpp
+++ b/source/framework/algorithm/include/lue/framework/algorithm/definition/cell_index.hpp
@@ -11,7 +11,9 @@ namespace lue {
         auto cell_index_partition(
             Policies const& policies,
             ConditionPartition const& condition_partition,
-            Index const dimension_idx) -> PartitionT<ConditionPartition, IndexElement>
+            Index const dimension_idx,
+            Index const idx_offset)
+            -> std::tuple<PartitionT<ConditionPartition, IndexElement>, hpx::future<Count>>
         {
             // TODO Generalize for all ranks and all dimension_idx values. Use mdspan?
             static_assert(rank<ConditionPartition> == 2);
@@ -24,15 +26,17 @@ namespace lue {
 
             lue_hpx_assert(condition_partition.is_ready());
 
-            return hpx::dataflow(
+            return hpx::split_future(hpx::dataflow(
                 hpx::launch::async,
                 hpx::unwrapping(
-                    [policies, condition_partition, dimension_idx](
-                        Offset const& offset, ConditionData const& condition_data)
+                    [policies, condition_partition, dimension_idx, idx_offset](
+                        Offset const& offset,
+                        ConditionData const& condition_data) -> std::tuple<IndexPartition, Count>
                     {
                         HPX_UNUSED(condition_partition);
 
                         auto const partition_shape = condition_data.shape();
+                        Count count{};
                         IndexData index_data{partition_shape};
 
                         auto const& indp = std::get<0>(policies.inputs_policies()).input_no_data_policy();
@@ -55,13 +59,15 @@ namespace lue {
                                     {
                                         if (condition_data[idx])
                                         {
-                                            index_data[idx] = idx0;
+                                            index_data[idx] = idx_offset + idx0;
                                         }
                                     }
 
                                     ++idx;
                                 }
                             }
+
+                            count = std::get<1>(partition_shape);
                         }
                         else if (dimension_idx == 1)
                         {
@@ -78,18 +84,22 @@ namespace lue {
                                     }
                                     else
                                     {
-                                        index_data[idx] = idx1;
+                                        index_data[idx] = idx_offset + idx1;
                                     }
 
                                     ++idx;
                                 }
                             }
+
+                            count = std::get<0>(partition_shape);
                         }
 
-                        return IndexPartition{hpx::find_here(), offset, std::move(index_data)};
+                        IndexPartition partition{hpx::find_here(), offset, std::move(index_data)};
+
+                        return {std::move(partition), count};
                     }),
                 condition_partition.offset(),
-                condition_partition.data());
+                condition_partition.data()));
         }
 
 
@@ -105,6 +115,16 @@ namespace lue {
     }  // namespace detail::cell_index
 
 
+    /*!
+        @brief      Return an array with for each valid true value in @a condition, the cell index along the
+                    dimension @a dimension_idx
+        @ingroup    local_operation
+
+        All elements with a non-valid or false value in @a condition will become non-valid in the result.
+
+        Cell indices don't vary along the dimension @a dimension_idx. They range from [0,
+        array_shape[dimension_idx] - 1].
+    */
     template<typename Policies, Rank rank>
     auto cell_index(
         Policies const& policies,
@@ -117,6 +137,7 @@ namespace lue {
 
         using IndexElement = policy::OutputElementT<Policies, 0>;
         using IndexArray = PartitionedArray<IndexElement, rank>;
+        using IndexPartition = PartitionT<IndexArray>;
         using IndexPartitions = PartitionsT<IndexArray>;
 
         static_assert(std::is_integral_v<ConditionElement>);
@@ -129,16 +150,137 @@ namespace lue {
         Localities<rank> const& localities{condition.localities()};
         auto const& condition_partitions{condition.partitions()};
         IndexPartitions index_partitions{condition_partitions.shape()};
-        Count const nr_partitions{nr_elements(condition_partitions.shape())};
 
-        for (Index partition_idx = 0; partition_idx < nr_partitions; ++partition_idx)
+        // TODO Generalize for all ranks and all dimension_idx values. Use mdspan?
+        static_assert(rank == 2);
+
+        auto const shape_in_partitions = condition_partitions.shape();
+
+        if (dimension_idx == 0)
         {
-            index_partitions[partition_idx] = hpx::dataflow(
-                hpx::launch::async,
-                [locality_id = localities[partition_idx], action, policies, dimension_idx](
-                    ConditionPartition const& condition_partition)
-                { return action(locality_id, policies, condition_partition, dimension_idx); },
-                condition_partitions[partition_idx]);
+            // Spawn tasks for all partitions with idx 0 in the dimension at idx 0
+            Index const partition_idx0{0};
+            hpx::shared_future<Count> nr_elements0_f{};
+
+            IndexPartition partition{};
+            hpx::future<Count> nr_elements0_f_{};
+
+            for (Index partition_idx1 = 0; partition_idx1 < std::get<1>(shape_in_partitions);
+                 ++partition_idx1)
+            {
+                std::tie(partition, nr_elements0_f_) = hpx::split_future(hpx::dataflow(
+                    hpx::launch::async,
+                    [locality_id = localities(partition_idx0, partition_idx1),
+                     action,
+                     policies,
+                     dimension_idx](ConditionPartition const& condition_partition)
+                    {
+                        Index const offset{0};
+
+                        return action(locality_id, policies, condition_partition, dimension_idx, offset);
+                    },
+                    condition_partitions(partition_idx0, partition_idx1)));
+
+                index_partitions(partition_idx0, partition_idx1) = std::move(partition);
+
+                if (partition_idx1 == 0)
+                {
+                    nr_elements0_f = nr_elements0_f_.share();
+                }
+            }
+
+            // Given that we now know how many cells are in a partition, we can spawn tasks for all other
+            // partitions
+
+            for (Index partition_idx0 = 1; partition_idx0 < std::get<0>(shape_in_partitions);
+                 ++partition_idx0)
+            {
+                for (Index partition_idx1 = 0; partition_idx1 < std::get<1>(shape_in_partitions);
+                     ++partition_idx1)
+                {
+                    std::tie(partition, std::ignore) = hpx::split_future(hpx::dataflow(
+                        hpx::launch::async,
+                        [locality_id = localities(partition_idx0, partition_idx1),
+                         action,
+                         policies,
+                         dimension_idx,
+                         partition_idx0](
+                            ConditionPartition const& condition_partition,
+                            hpx::shared_future<Count> const& nr_elements)
+                        {
+                            Index const offset{partition_idx0 * nr_elements.get()};
+
+                            return action(locality_id, policies, condition_partition, dimension_idx, offset);
+                        },
+                        condition_partitions(partition_idx0, partition_idx1),
+                        nr_elements0_f));
+
+                    index_partitions(partition_idx0, partition_idx1) = std::move(partition);
+                }
+            }
+        }
+        else if (dimension_idx == 1)
+        {
+            // Spawn tasks for all partitions with idx 0 in the dimension at idx 1
+            Index const partition_idx1{0};
+            hpx::shared_future<Count> nr_elements1_f{};
+
+            IndexPartition partition{};
+            hpx::future<Count> nr_elements1_f_{};
+
+            for (Index partition_idx0 = 0; partition_idx0 < std::get<0>(shape_in_partitions);
+                 ++partition_idx0)
+            {
+                std::tie(partition, nr_elements1_f_) = hpx::split_future(hpx::dataflow(
+                    hpx::launch::async,
+                    [locality_id = localities(partition_idx0, partition_idx1),
+                     action,
+                     policies,
+                     dimension_idx](ConditionPartition const& condition_partition)
+                    {
+                        Index const offset{0};
+
+                        return action(locality_id, policies, condition_partition, dimension_idx, offset);
+                    },
+                    condition_partitions(partition_idx0, partition_idx1)));
+
+                index_partitions(partition_idx0, partition_idx1) = std::move(partition);
+
+                if (partition_idx0 == 0)
+                {
+                    nr_elements1_f = nr_elements1_f_.share();
+                }
+            }
+
+            // Given that we now know how many cells are in a partition, we can spawn tasks for all other
+            // partitions
+
+            for (Index partition_idx0 = 0; partition_idx0 < std::get<0>(shape_in_partitions);
+                 ++partition_idx0)
+            {
+                for (Index partition_idx1 = 1; partition_idx1 < std::get<1>(shape_in_partitions);
+                     ++partition_idx1)
+                {
+                    std::tie(partition, std::ignore) = hpx::split_future(hpx::dataflow(
+                        hpx::launch::async,
+                        [locality_id = localities(partition_idx0, partition_idx1),
+                         action,
+                         policies,
+                         dimension_idx,
+                         partition_idx1](
+                            ConditionPartition const& condition_partition,
+                            hpx::shared_future<Count> const& nr_elements)
+                        {
+                            Index const offset{partition_idx1 * nr_elements.get()};
+
+                            return action(locality_id, policies, condition_partition, dimension_idx, offset);
+                        },
+                        condition_partitions(partition_idx0, partition_idx1),
+                        nr_elements1_f));
+
+                    index_partitions(partition_idx0, partition_idx1) = std::move(partition);
+                }
+            }
         }
 
         return {shape(condition), localities, std::move(index_partitions)};


Test patch:
diff --git a/source/framework/algorithm/test/cell_index_test.cpp b/source/framework/algorithm/test/cell_index_test.cpp
index 7ffad8b51..a23edf5d7 100644
--- a/source/framework/algorithm/test/cell_index_test.cpp
+++ b/source/framework/algorithm/test/cell_index_test.cpp
@@ -3,75 +3,170 @@
 #include "lue/framework/test/hpx_unit_test.hpp"
 
 
-BOOST_AUTO_TEST_CASE(cell_index)
+template<typename ConditionElement, typename IndexElement>
+void test_cell_index(
+    lue::Shape<lue::Count, 2> const& array_shape,
+    lue::Shape<lue::Count, 2> const& partition_shape,
+    std::initializer_list<std::initializer_list<ConditionElement>> const& condition_elements,
+    std::initializer_list<std::initializer_list<IndexElement>> const& cell_index_elements,
+    IndexElement const index)
 {
-    using namespace lue::value_policies;
+    auto const condition = lue::test::create_partitioned_array<lue::PartitionedArray<ConditionElement, 2>>(
+        array_shape, partition_shape, condition_elements);
+    auto const result_we_want = lue::test::create_partitioned_array<lue::PartitionedArray<IndexElement, 2>>(
+        array_shape, partition_shape, cell_index_elements);
+    auto const result_we_got = lue::value_policies::cell_index<IndexElement>(condition, index);
+    lue::test::check_arrays_are_equal(result_we_got, result_we_want);
+}
 
-    lue::Shape<lue::Count, 2> const array_shape{3, 3};
-    lue::Shape<lue::Count, 2> const partition_shape{3, 3};
 
+BOOST_AUTO_TEST_CASE(single_partition)
+{
     using ConditionElement = std::uint8_t;
     using IndexElement = std::uint64_t;
 
-    auto const condition_nd{lue::no_data<ConditionElement>};
-    auto const index_nd{lue::no_data<IndexElement>};
+    auto const cx{lue::no_data<ConditionElement>};
+    auto const ix{lue::no_data<IndexElement>};
 
-    auto const condition = lue::test::create_partitioned_array<lue::PartitionedArray<ConditionElement, 2>>(
-        array_shape,
-        partition_shape,
+    test_cell_index<ConditionElement, IndexElement>(
+        {3, 3},
+        {3, 3},
+        // clang-format off
+        {{
+            0, cx, 1,
+            1, 0, 1,
+            1, cx, 0,
+        }},
         {{
-            0,
-            condition_nd,
-            1,
-            1,
-            0,
-            1,
-            1,
-            condition_nd,
-            0,
-        }});
+            ix, ix, 0,
+            1, ix, 1,
+            2, ix, ix,
+        }},  // clang-format on
+        0);
 
-    {
-        auto const result_we_want =
-            lue::test::create_partitioned_array<lue::PartitionedArray<IndexElement, 2>>(
-                array_shape,
-                partition_shape,
-                {{
-                    index_nd,
-                    index_nd,
-                    0,
-                    1,
-                    index_nd,
-                    1,
-                    2,
-                    index_nd,
-                    index_nd,
-                }});
+    test_cell_index<ConditionElement, IndexElement>(
+        {3, 3},
+        {3, 3},
+        // clang-format off
+        {{
+            0, cx, 1,
+            1, 0, 1,
+            1, cx, 0,
+        }},
+        {{
+            ix, ix, 2,
+            0, ix, 2,
+            0, ix, ix,
+        }},  // clang-format on
+        1);
+}
 
-        auto const result_we_got = lue::value_policies::cell_index<IndexElement>(condition, 0);
 
-        lue::test::check_arrays_are_equal(result_we_got, result_we_want);
-    }
+BOOST_AUTO_TEST_CASE(multiple_partitions)
+{
+    using ConditionElement = std::uint8_t;
+    using IndexElement = std::uint64_t;
 
-    {
-        auto const result_we_want =
-            lue::test::create_partitioned_array<lue::PartitionedArray<IndexElement, 2>>(
-                array_shape,
-                partition_shape,
-                {{
-                    index_nd,
-                    index_nd,
-                    2,
-                    0,
-                    index_nd,
-                    2,
-                    0,
-                    index_nd,
-                    index_nd,
-                }});
+    auto const cx{lue::no_data<ConditionElement>};
+    auto const ix{lue::no_data<IndexElement>};
 
-        auto const result_we_got = lue::value_policies::cell_index<IndexElement>(condition, 1);
+    test_cell_index<ConditionElement, IndexElement>(
+        {6, 6},
+        {3, 3},
+        // clang-format off
+        {
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+        },
+        {
+            {
+                ix, ix, 0,
+                1, ix, 1,
+                2, ix, ix,
+            },
+            {
+                ix, ix, 0,
+                1, ix, 1,
+                2, ix, ix,
+            },
+            {
+                ix, ix, 3,
+                4, ix, 4,
+                5, ix, ix,
+            },
+            {
+                ix, ix, 3,
+                4, ix, 4,
+                5, ix, ix,
+            },
+        },  // clang-format on
+        0);
 
-        lue::test::check_arrays_are_equal(result_we_got, result_we_want);
-    }
+    test_cell_index<ConditionElement, IndexElement>(
+        {6, 6},
+        {3, 3},
+        // clang-format off
+        {
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+            {
+                0, cx, 1,
+                1, 0, 1,
+                1, cx, 0,
+            },
+        },
+        {
+            {
+                ix, ix, 2,
+                0, ix, 2,
+                0, ix, ix,
+            },
+            {
+                ix, ix, 5,
+                3, ix, 5,
+                3, ix, ix,
+            },
+            {
+                ix, ix, 2,
+                0, ix, 2,
+                0, ix, ix,
+            },
+            {
+                ix, ix, 5,
+                3, ix, 5,
+                3, ix, ix,
+            },
+        },  // clang-format on
+        1);
 }