Add Size symbol

dwave/optimization/include/dwave-optimization/nodes/indexing.hpp

@@ -285,4 +285,33 @@ class ReshapeNode : public ArrayOutputMixin<ArrayNode> {
     const Array* array_ptr_;
 };
 
+class SizeNode : public ScalarOutputMixin<ArrayNode> {
+ public:
+    explicit SizeNode(ArrayNode* node_ptr);
+
+    double const* buff(const State& state) const override;
+
+    void commit(State& state) const override;
+
+    std::span<const Update> diff(const State&) const override;
+
+    void initialize_state(State& state) const override;
+
+    // SizeNode's value is always a non-negative integer.
+    bool integral() const override { return true; }
+
+    double max() const override;
+
+    double min() const override;
+
+    void propagate(State& state) const override;
+
+    void revert(State& state) const override;
+
+ private:
+    // we could dynamically cast each time, but it's easier to just keep separate
+    // pointer to the "array" part of the predecessor
+    const Array* array_ptr_;
+};
+
 }  // namespace dwave::optimization

dwave/optimization/libcpp/nodes.pxd

@@ -79,6 +79,9 @@ cdef extern from "dwave-optimization/nodes/indexing.hpp" namespace "dwave::optim
     cdef cppclass ReshapeNode(ArrayNode):
         pass
 
+    cdef cppclass SizeNode(ArrayNode):
+        pass
+
 
 cdef extern from "dwave-optimization/nodes/mathematical.hpp" namespace "dwave::optimization" nogil:
     cdef cppclass AbsoluteNode(ArrayNode):

dwave/optimization/model.pyi

@@ -174,7 +174,7 @@ class ArraySymbol(Symbol):
     def prod(self) -> Prod: ...
     def reshape(self, shape: _ShapeLike) -> Reshape: ...
     def shape(self) -> typing.Tuple[int, ...]: ...
-    def size(self) -> int: ...
+    def size(self) -> typing.Union[int, Size]: ...
     def state(self, index: int = 0, *, copy: bool = True) -> numpy.ndarray: ...
     def state_size(self) -> int: ...
     def strides(self) -> typing.Tuple[int, ...]: ...

dwave/optimization/model.pyx

@@ -1877,7 +1877,8 @@ cdef class ArraySymbol(Symbol):
     def size(self):
         r"""Return the number of elements in the symbol.
 
-        ``-1`` indicates a variable number of elements.
+        If the symbol has a fixed size, returns that size as an integer.
+        Otherwise, returns a :class:`~dwave.optimization.symbols.Size` symbol.
 
         Examples:
             This example checks the size of a :math:`2 \times 3`
@@ -1888,7 +1889,12 @@ cdef class ArraySymbol(Symbol):
             >>> x = model.binary((2, 3))
             >>> x.size()
             6
+
         """
+        if self.array_ptr.dynamic():
+            from dwave.optimization.symbols import Size
+            return Size(self)
+
         return self.array_ptr.size()
 
     def state(self, Py_ssize_t index = 0, *, bool copy = True):

dwave/optimization/src/nodes/indexing.cpp

@@ -1733,4 +1733,65 @@ std::span<const Update> ReshapeNode::diff(const State& state) const {
 
 void ReshapeNode::revert(State& state) const {}  // stateless node
 
+// SizeNode *******************************************************************
+
+struct SizeNodeData : NodeStateData {
+    explicit SizeNodeData(std::integral auto value) : values(0, value, value) {}
+
+    double const* buff() const { return &values.value; }
+    void commit() { values.old = values.value; }
+    std::span<const Update> diff() const {
+        return std::span<const Update>(&values, values.old != values.value);
+    }
+    void revert() { values.value = values.old; }
+    void update(std::integral auto value) { values.value = value; }
+
+    Update values;
+};
+
+SizeNode::SizeNode(ArrayNode* node_ptr) : array_ptr_(node_ptr) { this->add_predecessor(node_ptr); }
+
+double const* SizeNode::buff(const State& state) const {
+    return data_ptr<SizeNodeData>(state)->buff();
+}
+
+void SizeNode::commit(State& state) const { return data_ptr<SizeNodeData>(state)->commit(); }
+
+std::span<const Update> SizeNode::diff(const State& state) const {
+    return data_ptr<SizeNodeData>(state)->diff();
+}
+
+void SizeNode::initialize_state(State& state) const {
+    int index = this->topological_index();
+    assert(index >= 0 && "must be topologically sorted");
+    assert(static_cast<int>(state.size()) > index && "unexpected state length");
+    assert(state[index] == nullptr && "already initialized state");
+
+    state[index] = std::make_unique<SizeNodeData>(array_ptr_->size(state));
+}
+
+double SizeNode::max() const {
+    // exactly the size of fixed-length predecessors
+    if (!array_ptr_->dynamic()) return array_ptr_->size();
+
+    // Ask the predecessor for its size, though in some cases it doesn't know
+    // so fall back on max of ssize_t
+    return array_ptr_->sizeinfo().max.value_or(std::numeric_limits<ssize_t>::max());
+}
+
+double SizeNode::min() const {
+    // exactly the size of fixed-length predecessors
+    if (!array_ptr_->dynamic()) return array_ptr_->size();
+
+    // Ask the predecessor for its size, though in some cases it doesn't know
+    // so fall back on 0
+    return array_ptr_->sizeinfo().min.value_or(0);
+}
+
+void SizeNode::propagate(State& state) const {
+    return data_ptr<SizeNodeData>(state)->update(array_ptr_->size(state));
+}
+
+void SizeNode::revert(State& state) const { return data_ptr<SizeNodeData>(state)->revert(); }
+
 }  // namespace dwave::optimization

dwave/optimization/symbols.pyi

@@ -172,6 +172,10 @@ class SetVariable(ArraySymbol):
     def set_state(self, index: int, state: numpy.typing.ArrayLike): ...
 
 
+class Size(ArraySymbol):
+    ...
+
+
 class Square(ArraySymbol):
     ...
 

dwave/optimization/symbols.pyx

@@ -80,6 +80,7 @@ from dwave.optimization.libcpp.nodes cimport (
     QuadraticModelNode as cppQuadraticModelNode,
     ReshapeNode as cppReshapeNode,
     SetNode as cppSetNode,
+    SizeNode as cppSizeNode,
     SubtractNode as cppSubtractNode,
     SquareNode as cppSquareNode,
     SumNode as cppSumNode,
@@ -126,6 +127,7 @@ __all__ = [
     "Reshape",
     "Subtract",
     "SetVariable",
+    "Size",
     "Square",
     "Sum",
     "Where",
@@ -2669,6 +2671,30 @@ cdef class SetVariable(ArraySymbol):
 _register(SetVariable, typeid(cppSetNode))
 
 
+cdef class Size(ArraySymbol):
+    def __init__(self, ArraySymbol array):
+        cdef Model model = array.model
+
+        self.ptr = model._graph.emplace_node[cppSizeNode](array.array_ptr)
+        self.initialize_arraynode(array.model, self.ptr)
+
+    @staticmethod
+    def _from_symbol(Symbol symbol):
+        cdef cppSizeNode* ptr = dynamic_cast_ptr[cppSizeNode](symbol.node_ptr)
+        if not ptr:
+            raise TypeError("given symbol cannot be used to construct a Size")
+
+        cdef Size x = Size.__new__(Size)
+        x.ptr = ptr
+        x.initialize_arraynode(symbol.model, ptr)
+        return x
+
+    # An observing pointer to the C++ SizeNode
+    cdef cppSizeNode* ptr
+
+_register(Size, typeid(cppSizeNode))
+
+
 cdef class Square(ArraySymbol):
     """Squares element-wise of a symbol.
 

releasenotes/notes/feature-Len-328f804a0f28ba37.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - Add ``SizeNode`` C++ class. ``SizeNode`` reports the size of its predecessor. See `#48 <https://github.com/dwavesystems/dwave-optimization/issues/48>`_.
+  - Add ``Size`` Python class. ``Size`` reports the size of its predecessor. See `#48 <https://github.com/dwavesystems/dwave-optimization/issues/48>`_.

tests/cpp/tests/test_nodes_indexing.cpp

@@ -2520,4 +2520,132 @@ TEST_CASE("ReshapeNode") {
     }
 }
 
+TEST_CASE("SizeNode") {
+    GIVEN("A 0d node") {
+        auto C = ConstantNode(5);
+
+        WHEN("We create a node accessing the length") {
+            auto len = SizeNode(&C);
+
+            THEN("The SizeNode is a scalar output") {
+                CHECK(len.size() == 1);
+                CHECK(len.ndim() == 0);
+                CHECK(!len.dynamic());
+            }
+
+            THEN("The output is a integer and we already know the min/max") {
+                CHECK(len.integral());
+                CHECK(len.min() == C.size());
+                CHECK(len.max() == C.size());
+            }
+        }
+    }
+
+    GIVEN("A 1D node with a fixed size") {
+        auto C = ConstantNode(std::vector{0, 1, 2});
+
+        WHEN("We create a node accessing the length") {
+            auto len = SizeNode(&C);
+
+            THEN("The SizeNode is a scalar output") {
+                CHECK(len.size() == 1);
+                CHECK(len.ndim() == 0);
+                CHECK(!len.dynamic());
+            }
+
+            THEN("The output is a integer and we already know the min/max") {
+                CHECK(len.integral());
+                CHECK(len.min() == C.size());
+                CHECK(len.max() == C.size());
+            }
+        }
+    }
+
+    GIVEN("A dynamic node") {
+        auto S = SetNode(5, 2, 4);
+
+        WHEN("We create a node accessing the length") {
+            auto len = SizeNode(&S);
+
+            THEN("The SizeNode is a scalar output") {
+                CHECK(len.size() == 1);
+                CHECK(len.ndim() == 0);
+                CHECK(!len.dynamic());
+            }
+
+            THEN("The output is a integer and we already know the min/max") {
+                CHECK(len.integral());
+                CHECK(len.min() == 2);
+                CHECK(len.max() == 4);
+            }
+        }
+
+        AND_WHEN("We create a node indirectly accessing the length") {
+            auto T = BasicIndexingNode(&S, Slice(0, -1));  // also dynamic, but indirect
+            auto len = SizeNode(&T);
+
+            THEN("The SizeNode is a scalar output") {
+                CHECK(len.size() == 1);
+                CHECK(len.ndim() == 0);
+                CHECK(!len.dynamic());
+            }
+
+            THEN("The output is a integer and we already know the min/max") {
+                CHECK(len.integral());
+
+                // these should always be true
+                CHECK(len.min() >= 0);
+                CHECK(len.max() <= std::numeric_limits<ssize_t>::max());
+
+                // these are an implementation detail and could change in the future
+                CHECK(len.min() == 0);
+                CHECK(len.max() == std::numeric_limits<ssize_t>::max());
+            }
+        }
+    }
+
+    GIVEN("A graph with a SetNode and a corresponding SizeNode, and a state") {
+        auto graph = Graph();
+
+        auto s_ptr = graph.emplace_node<SetNode>(5);
+        auto len_ptr = graph.emplace_node<SizeNode>(s_ptr);
+
+        auto state = graph.empty_state();
+        s_ptr->initialize_state(state, {});  // default to empty
+        graph.initialize_state(state);
+
+        THEN("The state of the SizeNode is the same as the size of the SetNode") {
+            CHECK(*(len_ptr->buff(state)) == s_ptr->size(state));
+        }
+
+        WHEN("We update the set and propagate") {
+            s_ptr->grow(state);
+            s_ptr->propagate(state);
+            len_ptr->propagate(state);
+
+            THEN("The state of the SizeNode is the same as the size of the SetNode") {
+                CHECK(*(len_ptr->buff(state)) == s_ptr->size(state));
+            }
+
+            AND_WHEN("We commit") {
+                s_ptr->commit(state);
+                len_ptr->commit(state);
+
+                THEN("The state of the SizeNode is the same as the size of the SetNode") {
+                    CHECK(*(len_ptr->buff(state)) == s_ptr->size(state));
+                }
+            }
+
+            AND_WHEN("We revert") {
+                s_ptr->revert(state);
+                len_ptr->revert(state);
+
+                THEN("The state of the SizeNode is the same as the size of the SetNode") {
+                    CHECK(*(len_ptr->buff(state)) == s_ptr->size(state));
+                }
+            }
+        }
+    }
+}
+
 }  // namespace dwave::optimization

tests/test_symbols.py

@@ -1737,7 +1737,6 @@ def test_shape(self):
 
         s = model.set(10)
         self.assertEqual(s.shape(), (-1,))
-        self.assertEqual(s.size(), -1)
         self.assertEqual(s.strides(), (np.dtype(np.double).itemsize,))
 
         t = model.set(5, 5)  # this is exactly range(5)
@@ -1810,6 +1809,43 @@ def test_state_size(self):
         self.assertEqual(model.set(10, max_size=5).state_size(), 5 * 8)
 
 
+class TestSize(utils.SymbolTests):
+    def generate_symbols(self):
+        model = Model()
+
+        a = dwave.optimization.symbols.Size(model.constant(5))
+        b = dwave.optimization.symbols.Size(model.constant([0, 1, 2]))
+        c = model.set(5).size()
+
+        with model.lock():
+            yield a
+            yield b
+            yield c
+
+    def test_dynamic(self):
+        model = Model()
+        model.states.resize(2)
+
+        set_ = model.set(5)
+        length = set_.size()
+
+        set_.set_state(0, [])
+        set_.set_state(1, [0, 2, 3])
+
+        with model.lock():
+            self.assertEqual(length.state(0), 0)
+            self.assertEqual(length.state(1), 3)
+
+    def test_scalar(self):
+        model = Model()
+        model.states.resize(1)
+
+        length = dwave.optimization.symbols.Size(model.constant(1))
+
+        with model.lock():
+            self.assertEqual(length.state(), 1)
+
+
 class TestSubtract(utils.BinaryOpTests):
     def generate_symbols(self):
         model = Model()