Enhancement/Sum constraints on NumberNode

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

@@ -28,6 +28,55 @@ namespace dwave::optimization {
 /// A contiguous block of numbers.
 class NumberNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
  public:
+    /// Stateless sum constraint information.
+    ///
+    /// A sum constraint constrains the sum of values within slices of the array.
+    /// The slices are defined along `axis` when `axis` has a value. If
+    /// `axis == std::nullopt`, the constraint is applied to the entire array,
+    /// which is treated as a flat array with a single slice.
+    ///
+    /// Constraints may be defined either:
+    /// - for ALL slices (the `operators` and `bounds` vectors have length 1), or
+    /// - PER slice (their lengths equal the number of slices along `axis`).
+    ///
+    /// Each slice sum is constrained by an `Operator` and a corresponding `bound`.
+    struct SumConstraint {
+     public:
+        /// Allowable operators.
+        enum class Operator { Equal, LessEqual, GreaterEqual };
+
+        /// To reduce the # of `IntegerNode` and `BinaryNode` constructors, we
+        /// allow only one constructor.
+        SumConstraint(std::optional<ssize_t> axis, std::vector<Operator> operators,
+                      std::vector<double> bounds);
+
+        /// Return the axis along which slices are defined.
+        /// If `std::nullopt`, the sum constraint applies to the entire array.
+        std::optional<ssize_t> axis() const { return axis_; };
+
+        /// Obtain the bound associated with a given slice.
+        double get_bound(const ssize_t slice) const;
+
+        /// Obtain the operator associated with a given slice.
+        Operator get_operator(const ssize_t slice) const;
+
+        /// The number of bounds.
+        ssize_t num_bounds() const { return bounds_.size(); };
+
+        /// The number of operators.
+        ssize_t num_operators() const { return operators_.size(); };
+
+     private:
+        /// Axis along which slices are defined (`std::nullopt` = whole array).
+        std::optional<ssize_t> axis_ = std::nullopt;
+        /// Operator for ALL slices (vector has length one) or operators PER
+        /// slice (length of vector is equal to the number of slices).
+        std::vector<Operator> operators_;
+        /// Bound for ALL slices (vector has length one) or bounds PER slice
+        /// (length of vector is equal to the number of slices).
+        std::vector<double> bounds_;
+    };
+
     NumberNode() = delete;
 
     // Overloads needed by the Array ABC **************************************
@@ -68,6 +117,11 @@ class NumberNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
     // Initialize the state of the node randomly
     template <std::uniform_random_bit_generator Generator>
     void initialize_state(State& state, Generator& rng) const {
+        // Currently do not support random node initialization with sum constraints.
+        if (sum_constraints_.size() > 0) {
+            throw std::invalid_argument("Cannot randomly initialize_state with sum constraints.");
+        }
+
         std::vector<double> values;
         const ssize_t size = this->size();
         values.reserve(size);
@@ -86,6 +140,9 @@ class NumberNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
         return initialize_state(state, std::move(values));
     }
 
+    /// @copydoc Node::propagate()
+    void propagate(State& state) const override;
+
     // NumberNode methods *****************************************************
 
     // In the given state, swap the value of index i with the value of index j.
@@ -106,21 +163,40 @@ class NumberNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
     // in a given index.
     void clip_and_set_value(State& state, ssize_t index, double value) const;
 
+    /// Return the stateless sum constraints.
+    const std::vector<SumConstraint>& sum_constraints() const;
+
+    /// If the node is subject to sum constraints, we track the state
+    /// dependent sum of the values within each slice per constraint. The
+    /// returned vector is indexed in the same ordering as the constraints
+    /// given by `sum_constraints()`.
+    const std::vector<std::vector<double>>& sum_constraints_lhs(const State& state) const;
+
+    std::vector<std::vector<double>>& sum_constraints_lhs(State& state) const;
+
  protected:
     explicit NumberNode(std::span<const ssize_t> shape, std::vector<double> lower_bound,
-                        std::vector<double> upper_bound);
+                        std::vector<double> upper_bound,
+                        std::vector<SumConstraint> sum_constraints = {});
 
     // Return truth statement: 'value is valid in a given index'.
     virtual bool is_valid(ssize_t index, double value) const = 0;
 
     // Default value in a given index.
     virtual double default_value(ssize_t index) const = 0;
 
+    /// Statelss global minimum and maximum of the values stored in NumberNode.
     double min_;
     double max_;
 
+    /// Stateless index-wise upper and lower bounds.
     std::vector<double> lower_bounds_;
     std::vector<double> upper_bounds_;
+
+    /// Stateless sum constraints.
+    std::vector<SumConstraint> sum_constraints_;
+    /// Indicator variable that all sum constraint operators are "==".
+    bool sum_constraints_all_equals_;
 };
 
 /// A contiguous block of integer numbers.
@@ -134,33 +210,45 @@ class IntegerNode : public NumberNode {
     // Default to a single scalar integer with default bounds
     IntegerNode() : IntegerNode({}) {}
 
-    // Create an integer array with the user-defined bounds.
-    // Defaulting to the specified default bounds.
+    // Create an integer array with the user-defined index-wise bounds and sum
+    // constraints. Index-wise bounds default to the specified default bounds.
+    // By default, there are no sum constraints.
     IntegerNode(std::span<const ssize_t> shape,
                 std::optional<std::vector<double>> lower_bound = std::nullopt,
-                std::optional<std::vector<double>> upper_bound = std::nullopt);
+                std::optional<std::vector<double>> upper_bound = std::nullopt,
+                std::vector<SumConstraint> sum_constraints = {});
     IntegerNode(std::initializer_list<ssize_t> shape,
                 std::optional<std::vector<double>> lower_bound = std::nullopt,
-                std::optional<std::vector<double>> upper_bound = std::nullopt);
+                std::optional<std::vector<double>> upper_bound = std::nullopt,
+                std::vector<SumConstraint> sum_constraints = {});
     IntegerNode(ssize_t size, std::optional<std::vector<double>> lower_bound = std::nullopt,
-                std::optional<std::vector<double>> upper_bound = std::nullopt);
+                std::optional<std::vector<double>> upper_bound = std::nullopt,
+                std::vector<SumConstraint> sum_constraints = {});
 
     IntegerNode(std::span<const ssize_t> shape, double lower_bound,
-                std::optional<std::vector<double>> upper_bound = std::nullopt);
+                std::optional<std::vector<double>> upper_bound = std::nullopt,
+                std::vector<SumConstraint> sum_constraints = {});
     IntegerNode(std::initializer_list<ssize_t> shape, double lower_bound,
-                std::optional<std::vector<double>> upper_bound = std::nullopt);
+                std::optional<std::vector<double>> upper_bound = std::nullopt,
+                std::vector<SumConstraint> sum_constraints = {});
     IntegerNode(ssize_t size, double lower_bound,
-                std::optional<std::vector<double>> upper_bound = std::nullopt);
+                std::optional<std::vector<double>> upper_bound = std::nullopt,
+                std::vector<SumConstraint> sum_constraints = {});
 
     IntegerNode(std::span<const ssize_t> shape, std::optional<std::vector<double>> lower_bound,
-                double upper_bound);
+                double upper_bound, std::vector<SumConstraint> sum_constraints = {});
     IntegerNode(std::initializer_list<ssize_t> shape,
-                std::optional<std::vector<double>> lower_bound, double upper_bound);
-    IntegerNode(ssize_t size, std::optional<std::vector<double>> lower_bound, double upper_bound);
-
-    IntegerNode(std::span<const ssize_t> shape, double lower_bound, double upper_bound);
-    IntegerNode(std::initializer_list<ssize_t> shape, double lower_bound, double upper_bound);
-    IntegerNode(ssize_t size, double lower_bound, double upper_bound);
+                std::optional<std::vector<double>> lower_bound, double upper_bound,
+                std::vector<SumConstraint> sum_constraints = {});
+    IntegerNode(ssize_t size, std::optional<std::vector<double>> lower_bound, double upper_bound,
+                std::vector<SumConstraint> sum_constraints = {});
+
+    IntegerNode(std::span<const ssize_t> shape, double lower_bound, double upper_bound,
+                std::vector<SumConstraint> sum_constraints = {});
+    IntegerNode(std::initializer_list<ssize_t> shape, double lower_bound, double upper_bound,
+                std::vector<SumConstraint> sum_constraints = {});
+    IntegerNode(ssize_t size, double lower_bound, double upper_bound,
+                std::vector<SumConstraint> sum_constraints = {});
 
     // Overloads needed by the Node ABC ***************************************
 
@@ -190,33 +278,44 @@ class BinaryNode : public IntegerNode {
     /// A binary scalar variable with lower_bound = 0.0 and upper_bound = 1.0
     BinaryNode() : BinaryNode({}) {}
 
-    // Create a binary array with the user-defined bounds.
-    // Defaulting to lower_bound = 0.0 and upper_bound = 1.0
+    // Create a binary array with the user-defined index-wise bounds and sum
+    // constraints. Index-wise bounds default to lower_bound = 0.0 and
+    // upper_bound = 1.0. By default, there are no sum constraints.
     BinaryNode(std::span<const ssize_t> shape,
                std::optional<std::vector<double>> lower_bound = std::nullopt,
-               std::optional<std::vector<double>> upper_bound = std::nullopt);
+               std::optional<std::vector<double>> upper_bound = std::nullopt,
+               std::vector<SumConstraint> sum_constraints = {});
     BinaryNode(std::initializer_list<ssize_t> shape,
                std::optional<std::vector<double>> lower_bound = std::nullopt,
-               std::optional<std::vector<double>> upper_bound = std::nullopt);
+               std::optional<std::vector<double>> upper_bound = std::nullopt,
+               std::vector<SumConstraint> sum_constraints = {});
     BinaryNode(ssize_t size, std::optional<std::vector<double>> lower_bound = std::nullopt,
-               std::optional<std::vector<double>> upper_bound = std::nullopt);
+               std::optional<std::vector<double>> upper_bound = std::nullopt,
+               std::vector<SumConstraint> sum_constraints = {});
 
     BinaryNode(std::span<const ssize_t> shape, double lower_bound,
-               std::optional<std::vector<double>> upper_bound = std::nullopt);
+               std::optional<std::vector<double>> upper_bound = std::nullopt,
+               std::vector<SumConstraint> sum_constraints = {});
     BinaryNode(std::initializer_list<ssize_t> shape, double lower_bound,
-               std::optional<std::vector<double>> upper_bound = std::nullopt);
+               std::optional<std::vector<double>> upper_bound = std::nullopt,
+               std::vector<SumConstraint> sum_constraints = {});
     BinaryNode(ssize_t size, double lower_bound,
-               std::optional<std::vector<double>> upper_bound = std::nullopt);
+               std::optional<std::vector<double>> upper_bound = std::nullopt,
+               std::vector<SumConstraint> sum_constraints = {});
 
     BinaryNode(std::span<const ssize_t> shape, std::optional<std::vector<double>> lower_bound,
-               double upper_bound);
+               double upper_bound, std::vector<SumConstraint> sum_constraints = {});
     BinaryNode(std::initializer_list<ssize_t> shape, std::optional<std::vector<double>> lower_bound,
-               double upper_bound);
-    BinaryNode(ssize_t size, std::optional<std::vector<double>> lower_bound, double upper_bound);
-
-    BinaryNode(std::span<const ssize_t> shape, double lower_bound, double upper_bound);
-    BinaryNode(std::initializer_list<ssize_t> shape, double lower_bound, double upper_bound);
-    BinaryNode(ssize_t size, double lower_bound, double upper_bound);
+               double upper_bound, std::vector<SumConstraint> sum_constraints = {});
+    BinaryNode(ssize_t size, std::optional<std::vector<double>> lower_bound, double upper_bound,
+               std::vector<SumConstraint> sum_constraints = {});
+
+    BinaryNode(std::span<const ssize_t> shape, double lower_bound, double upper_bound,
+               std::vector<SumConstraint> sum_constraints = {});
+    BinaryNode(std::initializer_list<ssize_t> shape, double lower_bound, double upper_bound,
+               std::vector<SumConstraint> sum_constraints = {});
+    BinaryNode(ssize_t size, double lower_bound, double upper_bound,
+               std::vector<SumConstraint> sum_constraints = {});
 
     // Flip the value (0 -> 1 or 1 -> 0) at index i in the given state.
     void flip(State& state, ssize_t i) const;

dwave/optimization/libcpp/nodes/numbers.pxd

@@ -12,23 +12,42 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 
+from libcpp.optional cimport optional
 from libcpp.vector cimport vector
 
 from dwave.optimization.libcpp.graph cimport ArrayNode
 from dwave.optimization.libcpp.state cimport State
 
 
 cdef extern from "dwave-optimization/nodes/numbers.hpp" namespace "dwave::optimization" nogil:
-    cdef cppclass IntegerNode(ArrayNode):
-        void initialize_state(State&, vector[double]) except+
-        double lower_bound(Py_ssize_t index)
-        double upper_bound(Py_ssize_t index)
-        double lower_bound() except+
-        double upper_bound() except+
 
-    cdef cppclass BinaryNode(ArrayNode):
+    cdef cppclass NumberNode(ArrayNode):
+        struct SumConstraint:
+            # It appears Cython automatically assumes all (standard) enums are "public".
+            # Because of this, we use this very explict override.
+            enum class Operator "dwave::optimization::NumberNode::SumConstraint::Operator":
+                Equal
+                LessEqual
+                GreaterEqual
+
+            SumConstraint(optional[Py_ssize_t] axis, vector[Operator] operators,
+                          vector[double] bounds)
+
+            optional[Py_ssize_t] axis()
+            double get_bound(Py_ssize_t slice)
+            Operator get_operator(Py_ssize_t slice)
+            Py_ssize_t num_bounds()
+            Py_ssize_t num_operators()
+
         void initialize_state(State&, vector[double]) except+
         double lower_bound(Py_ssize_t index)
         double upper_bound(Py_ssize_t index)
         double lower_bound() except+
         double upper_bound() except+
+        const vector[SumConstraint] sum_constraints()
+
+    cdef cppclass IntegerNode(NumberNode):
+        pass
+
+    cdef cppclass BinaryNode(IntegerNode):
+        pass