diff --git a/dwave/optimization/include/dwave-optimization/nodes/flow.hpp b/dwave/optimization/include/dwave-optimization/nodes/flow.hpp index a8e01cc9..9aab61c1 100644 --- a/dwave/optimization/include/dwave-optimization/nodes/flow.hpp +++ b/dwave/optimization/include/dwave-optimization/nodes/flow.hpp @@ -85,6 +85,74 @@ class ExtractNode : public ArrayOutputMixin> { const SizeInfo sizeinfo_; }; +/// Return the indices of the non-zero elements of an array. +/// +/// Equivalent to ``np.transpose(np.nonzero(arr))`` (i.e. ``np.argwhere(arr)``). +/// Whereas ``np.nonzero()`` returns a tuple of ``arr.ndim`` 1d index arrays, this +/// node always outputs a single ``(num_nonzero, arr.ndim)`` array. The rows are +/// ordered by the C-order (row-major) flattened index of the non-zero elements. +/// +/// The predecessor array must have ``ndim >= 1``; scalar (0d) inputs are not +/// supported. +class ArgWhereNode : public ArrayOutputMixin> { + public: + explicit ArgWhereNode(ArrayNode* array_ptr); + + /// @copydoc Array::buff() + double const* buff(const State& state) const override; + + /// @copydoc Node::commit() + void commit(State& state) const override; + + /// @copydoc Array::diff() + std::span diff(const State& state) const override; + + /// @copydoc Node::initialize_state() + void initialize_state(State& state) const override; + + /// @copydoc Array::integral() + bool integral() const override; + + /// @copydoc Array::max() + double max() const override; + + /// @copydoc Array::min() + double min() const override; + + /// @copydoc Node::propagate() + void propagate(State& state) const override; + + /// @copydoc Node::revert() + void revert(State& state) const override; + + using Array::shape; + + /// @copydoc Array::shape() + std::span shape(const State& state) const override; + + using Array::size; + + /// @copydoc Array::size() + ssize_t size(const State& state) const override; + + /// @copydoc Array::size_diff() + ssize_t size_diff(const State& state) const override; + + /// @copydoc Array::sizeinfo() + SizeInfo sizeinfo() const override; + + protected: + void replace_predecessor_(ssize_t index, Node* node_ptr) override; + + private: + const Array* array_ptr_; + + const SizeInfo sizeinfo_; + + // The minimum and maximum index value that can appear in the output. + const std::pair minmax_; +}; + /// Choose elements from x or y depending on condition. /// /// `condition` must be either a scalar array or the same shape as `x` and `y`. diff --git a/dwave/optimization/src/nodes/flow.cpp b/dwave/optimization/src/nodes/flow.cpp index 19f3e71c..fa2db188 100644 --- a/dwave/optimization/src/nodes/flow.cpp +++ b/dwave/optimization/src/nodes/flow.cpp @@ -15,9 +15,13 @@ #include "dwave-optimization/nodes/flow.hpp" #include +#include +#include +#include #include #include #include +#include #include #include "_state.hpp" @@ -172,6 +176,193 @@ ssize_t ExtractNode::size_diff(const State& state) const { SizeInfo ExtractNode::sizeinfo() const { return this->sizeinfo_; } +/// ArgWhereNode + +// Validate the predecessor and return the (static) output shape (num_nonzero, ndim). +std::vector argwhere_shape(const ArrayNode* array_ptr) { + if (!array_ptr) throw std::invalid_argument("node pointer cannot be nullptr"); + if (array_ptr->ndim() < 1) { + throw std::invalid_argument( + "cannot take the non-zero indices of a scalar (0d) array" + ); + } + // The number of non-zero elements is state-dependent, so axis 0 is dynamic. + return {Array::DYNAMIC_SIZE, array_ptr->ndim()}; +} + +// The output holds one row of `ndim` indices per non-zero element, so its size +// is `ndim` times the number of non-zero elements, which ranges from 0 to the +// size of the predecessor. +SizeInfo argwhere_sizeinfo(const Array* self, const Array* array_ptr) { + const std::optional arr_max = array_ptr->sizeinfo().max; + std::optional max = std::nullopt; + if (arr_max.has_value()) max = arr_max.value() * array_ptr->ndim(); + return SizeInfo(self, 0, max); +} + +// The output values are indices into the predecessor. The smallest is 0. The +// largest is one less than the size of the predecessor's largest dimension. +std::pair argwhere_minmax(const Array* array_ptr) { + const std::span shape = array_ptr->shape(); + + // the product of the fixed (non-axis-0) dimensions + const ssize_t rest = std::reduce(shape.begin() + 1, shape.end(), 1, std::multiplies()); + + // the maximum possible length of the (possibly dynamic) axis 0 + ssize_t axis0; + if (rest == 0) { + axis0 = 0; // the array is empty, so there are never any indices + } else if (const std::optional arr_max = array_ptr->sizeinfo().max; + arr_max.has_value()) { + axis0 = arr_max.value() / rest; + } else { + axis0 = std::numeric_limits::max(); + } + + double max = static_cast(axis0) - 1; + for (const ssize_t& dim : shape | std::views::drop(1)) { + max = std::max(max, static_cast(dim) - 1); + } + return {0.0, std::max(0.0, max)}; +} + +// Append the multi-index of a single element - the C-order flat `index` +// unravelled according to `shape` - to the flattened output buffer. +void argwhere_emplace_multi_index(std::vector& out, ssize_t index, + std::span shape) { + const ssize_t ndim = shape.size(); + const ssize_t offset = out.size(); + out.resize(offset + ndim); + for (ssize_t axis = ndim - 1; axis >= 0; --axis) { + // shape[axis] > 0 here: if any dimension were 0 the array would be + // empty and this function would not be called. + out[offset + axis] = index % shape[axis]; + index /= shape[axis]; + } +} + +// The state holds the flattened (num_nonzero, ndim) index buffer as well as the +// state-dependent shape. +struct ArgWhereNodeData : ArrayNodeStateData { + ArgWhereNodeData(std::vector&& values, ssize_t ndim) noexcept : + ArrayNodeStateData(std::move(values)), shape_{0, ndim} { + update_shape(); + } + + std::unique_ptr copy() const override { + return std::make_unique(*this); + } + + // Recompute the (dynamic) number of rows from the current buffer size. + void update_shape() { shape_[0] = this->size() / shape_[1]; } + + std::span shape() const { + return std::span(shape_.data(), shape_.size()); + } + + // shape_[0] is the number of non-zero indices (dynamic), shape_[1] is the + // (fixed) number of dimensions of the predecessor. + std::array shape_; +}; + +ArgWhereNode::ArgWhereNode(ArrayNode* array_ptr) : + ArrayOutputMixin(argwhere_shape(array_ptr)), + array_ptr_(array_ptr), + sizeinfo_(argwhere_sizeinfo(this, array_ptr)), + minmax_(argwhere_minmax(array_ptr)) { + add_predecessor_(array_ptr); +} + +double const* ArgWhereNode::buff(const State& state) const { + return data_ptr_(state)->buff(); +} + +void ArgWhereNode::commit(State& state) const { data_ptr_(state)->commit(); } + +std::span ArgWhereNode::diff(const State& state) const { + return data_ptr_(state)->diff(); +} + +void ArgWhereNode::initialize_state(State& state) const { + const std::span shape = array_ptr_->shape(state); + const std::ranges::view auto arr = array_ptr_->view(state); + + std::vector values; + ssize_t index = 0; + for (auto it = arr.begin(); it != std::default_sentinel; ++it, ++index) { + if (static_cast(*it)) argwhere_emplace_multi_index(values, index, shape); + } + + emplace_data_ptr_(state, std::move(values), array_ptr_->ndim()); +} + +bool ArgWhereNode::integral() const { return true; } + +double ArgWhereNode::max() const { return minmax_.second; } + +double ArgWhereNode::min() const { return minmax_.first; } + +void ArgWhereNode::propagate(State& state) const { + const std::span arr_diff = array_ptr_->diff(state); + if (arr_diff.empty()) return; + + auto node_data = data_ptr_(state); + + // Find the smallest flat index in the predecessor that changed. Every + // element before it keeps both its truthiness and its (fixed) multi-index, + // so the corresponding output rows are unchanged. + const ssize_t min_changed = std::ranges::min( + arr_diff | std::views::transform([](const Update& update) { return update.index; }) + ); + + const std::span shape = array_ptr_->shape(state); + const std::ranges::view auto arr = array_ptr_->view(state); + const ssize_t ndim = array_ptr_->ndim(); + + // Count the non-zero elements strictly before min_changed. Each contributes + // one already-correct row of `ndim` values to the front of the buffer. + auto is_nonzero = [](double value) { return static_cast(value); }; + const ssize_t count = std::count_if(arr.begin(), arr.begin() + min_changed, is_nonzero); + + // Recompute the rows for every element from min_changed onwards. + std::vector values; + ssize_t index = min_changed; + for (auto it = arr.begin() + min_changed; it != std::default_sentinel; ++it, ++index) { + if (is_nonzero(*it)) argwhere_emplace_multi_index(values, index, shape); + } + + node_data->assign(std::move(values), count * ndim); + node_data->update_shape(); +} + +void ArgWhereNode::replace_predecessor_(ssize_t index, Node* node_ptr) { + Node::replace_predecessor_(index, node_ptr); + + assert(index == 0); + array_ptr_ = dynamic_cast(node_ptr); + assert(array_ptr_ != nullptr); +} + +void ArgWhereNode::revert(State& state) const { + auto node_data = data_ptr_(state); + node_data->revert(); + node_data->update_shape(); +} + +std::span ArgWhereNode::shape(const State& state) const { + return data_ptr_(state)->shape(); +} + +ssize_t ArgWhereNode::size(const State& state) const { + return data_ptr_(state)->size(); +} + +ssize_t ArgWhereNode::size_diff(const State& state) const { + return data_ptr_(state)->size_diff(); +} + +SizeInfo ArgWhereNode::sizeinfo() const { return sizeinfo_; } + /// WhereNode struct WhereNodeData : ArrayNodeStateData { diff --git a/tests/cpp/nodes/test_flow.cpp b/tests/cpp/nodes/test_flow.cpp index bebe3d75..77977239 100644 --- a/tests/cpp/nodes/test_flow.cpp +++ b/tests/cpp/nodes/test_flow.cpp @@ -246,6 +246,200 @@ TEST_CASE("ExtractNode") { } } +TEST_CASE("ArgWhereNode") { + auto graph = Graph(); + + GIVEN("A scalar integer") { + auto scalar_ptr = graph.emplace_node(std::vector{}, -5, 5); + + THEN("We cannot construct a ArgWhereNode from it") { + CHECK_THROWS_AS(graph.emplace_node(scalar_ptr), std::invalid_argument); + } + } + + GIVEN("A 1d integer array and its ArgWhereNode") { + auto arr_ptr = graph.emplace_node(std::vector{5}, -5, 5); + auto nz_ptr = graph.emplace_node(arr_ptr); + + THEN("The ArgWhereNode has the shape/properties we expect") { + CHECK(nz_ptr->ndim() == 2); + CHECK(nz_ptr->dynamic()); + CHECK(std::ranges::equal(nz_ptr->shape(), std::vector{-1, 1})); + CHECK(nz_ptr->integral()); + CHECK(nz_ptr->min() == 0); + CHECK(nz_ptr->max() == 4); // largest index into an array of length 5 + } + + THEN("The ArgWhereNode has the sizeinfo we expect") { + auto sizeinfo = nz_ptr->sizeinfo(); + CHECK(sizeinfo.min == 0); + CHECK(sizeinfo.max == 5); // ndim (1) * size (5) + } + + WHEN("We initialize a state with some non-zero values") { + auto state = graph.empty_state(); + arr_ptr->initialize_state(state, {0, 3, 0, -2, 1}); + graph.initialize_state(state); + + THEN("The output is the transpose of the non-zero indices") { + CHECK(std::ranges::equal(nz_ptr->shape(state), std::vector{3, 1})); + CHECK(nz_ptr->size(state) == 3); + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{1, 3, 4})); + } + + AND_WHEN("We change some values and propagate") { + arr_ptr->set_value(state, 0, 4); // becomes non-zero + arr_ptr->set_value(state, 3, 0); // becomes zero + graph.propagate(state, graph.descendants(state, {arr_ptr})); + + THEN("The output is updated") { + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{0, 1, 4})); + } + + AND_WHEN("We commit") { + graph.commit(state, graph.descendants(state, {arr_ptr})); + + THEN("The output is retained") { + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{0, 1, 4})); + } + } + + AND_WHEN("We revert") { + graph.revert(state, graph.descendants(state, {arr_ptr})); + + THEN("The output returns to its original value") { + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{1, 3, 4})); + CHECK(std::ranges::equal(nz_ptr->shape(state), std::vector{3, 1})); + } + } + } + } + } + + GIVEN("A 2d integer array and its ArgWhereNode") { + auto arr_ptr = graph.emplace_node(std::vector{2, 3}, -5, 5); + auto nz_ptr = graph.emplace_node(arr_ptr); + + THEN("The ArgWhereNode has the shape/properties we expect") { + CHECK(std::ranges::equal(nz_ptr->shape(), std::vector{-1, 2})); + CHECK(nz_ptr->min() == 0); + CHECK(nz_ptr->max() == 2); // largest index is into the length-3 axis + } + + WHEN("We initialize a state") { + auto state = graph.empty_state(); + // [[0, 5, 0], + // [0, 0, 2]] + arr_ptr->initialize_state(state, {0, 5, 0, 0, 0, 2}); + graph.initialize_state(state); + + THEN("The output holds the (row, col) index of each non-zero") { + CHECK(std::ranges::equal(nz_ptr->shape(state), std::vector{2, 2})); + // (0, 1) and (1, 2) + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{0, 1, 1, 2})); + } + } + } + + GIVEN("A 1d dynamic array and its ArgWhereNode, tracked for consistency") { + auto dyn_ptr = graph.emplace_node( + std::initializer_list{-1}, -5, 5, true + ); + auto nz_ptr = graph.emplace_node(dyn_ptr); + graph.emplace_node(nz_ptr); + + auto state = graph.empty_state(); + dyn_ptr->initialize_state(state, {0, 3, 0, 2}); + graph.initialize_state(state); + + THEN("The output starts correct") { + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{1, 3})); + } + + WHEN("We flip a value, grow, and shrink") { + dyn_ptr->set(state, 0, 5); // index 0 becomes non-zero + dyn_ptr->grow(state, {0, 7}); // append indices 4 (zero) and 5 (non-zero) + dyn_ptr->shrink(state); // drop index 5 + graph.propose(state, {dyn_ptr}); + + THEN("The output is correct") { + // array is now [5, 3, 0, 2, 0], non-zero at 0, 1, 3 + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{0, 1, 3})); + } + } + + WHEN("We make changes but reject them") { + dyn_ptr->set(state, 1, 0); + dyn_ptr->grow(state, {9}); + graph.propose(state, {dyn_ptr}, [](const Graph&, State&) { return false; }); + + THEN("The output is unchanged") { + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{1, 3})); + } + } + } + + GIVEN("A 2d dynamic array and its ArgWhereNode, tracked for consistency") { + auto dyn_ptr = graph.emplace_node( + std::initializer_list{-1, 2}, -5, 5, true + ); + auto nz_ptr = graph.emplace_node(dyn_ptr); + graph.emplace_node(nz_ptr); + + auto state = graph.empty_state(); + // [[0, 0], [1, 0]] + dyn_ptr->initialize_state(state, {0, 0, 1, 0}); + graph.initialize_state(state); + + THEN("The output starts correct") { + CHECK(std::ranges::equal(nz_ptr->shape(state), std::vector{1, 2})); + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{1, 0})); + } + + WHEN("We grow a row with a non-zero and propagate") { + dyn_ptr->grow(state, {0, 3}); // append row [0, 3] + graph.propose(state, {dyn_ptr}); + + THEN("The new (row, col) index appears") { + // non-zero at (1, 0) and (2, 1) + CHECK(std::ranges::equal(nz_ptr->view(state), std::vector{1, 0, 2, 1})); + } + } + } + + SECTION("equality") { + auto* a0_ptr = graph.emplace_node(std::vector{3}, -5, 5); + auto* a1_ptr = graph.emplace_node(std::vector{3}, -5, 5); + + Node* a_ptr = graph.emplace_node(a0_ptr); + Node* b_ptr = graph.emplace_node(a0_ptr); + Node* c_ptr = graph.emplace_node(a1_ptr); + + CHECK(a_ptr->equal_to(*a_ptr)); + CHECK(a_ptr->equal_to(*b_ptr)); + CHECK(not a_ptr->equal_to(*c_ptr)); + CHECK(not a_ptr->equal_to(*a0_ptr)); + } + + SECTION("predecessor replacement") { + auto* a0_ptr = graph.emplace_node(std::vector{3}, -5, 5); + auto* a1_ptr = graph.emplace_node(std::vector{3}, -5, 5); + + auto* nz_ptr = graph.emplace_node(a0_ptr); + + a1_ptr->take_successors(*a0_ptr); + + CHECK_THAT(nz_ptr->predecessors(), RangeEquals({a1_ptr})); + + auto state = graph.empty_state(); + a0_ptr->initialize_state(state, {0, 0, 0}); + a1_ptr->initialize_state(state, {1, 0, 2}); + graph.initialize_state(state); + + CHECK_THAT(nz_ptr->view(state), RangeEquals({0, 2})); + } +} + TEST_CASE("WhereNode") { auto graph = Graph();