Enhancement/axis wise bounds NumberNode

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

@@ -28,6 +28,43 @@ namespace dwave::optimization {
 /// A contiguous block of numbers.
 class NumberNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
  public:
+    /// Struct for stateless axis-wise bound information. Given an `axis`,
+    /// define constraints on the sum of the values in each slice along `axis`.
+    /// Constraints can be defined for ALL slices along `axis` or PER slice
+    /// along `axis`. Allowable operators are defined by `Operator`.
+    struct AxisBound {
+     public:
+        /// Allowable axis-wise bound operators.
+        enum class Operator { Equal, LessEqual, GreaterEqual };
+
+        /// To reduce the # of `IntegerNode` and `BinaryNode` constructors, we
+        /// allow only one constructor.
+        AxisBound(ssize_t axis, std::vector<Operator> axis_operators,
+                  std::vector<double> axis_bounds);
+
+        ssize_t axis() const { return axis_; };
+
+        /// Obtain the bound associated with a given slice along `axis`.
+        double get_bound(const ssize_t slice) const;
+
+        /// Obtain the operator associated with a given slice along `axis`.
+        Operator get_operator(const ssize_t slice) const;
+
+        ssize_t num_bounds() const { return bounds_.size(); };
+
+        ssize_t num_operators() const { return operators_.size(); };
+
+     private:
+        /// The bound axis
+        ssize_t axis_;
+        /// Operator for ALL axis 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 axis 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 +105,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 bound axes.
+        if (bound_axes_info_.size() > 0) {
+            throw std::invalid_argument("Cannot randomly initialize_state with bound axes.");
+        }
+
         std::vector<double> values;
         const ssize_t size = this->size();
         values.reserve(size);
@@ -86,6 +128,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 +151,41 @@ 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 axis-wise bound information i.e. bound_axes_info_.
+    const std::vector<AxisBound>& axis_wise_bounds() const;
+
+    /// Return the state-dependent sum of the values within each slice
+    /// along each bound axis. The returned vector is indexed by the
+    /// bound axes in the same ordering that `axis_wise_bounds()` returns.
+    const std::vector<std::vector<double>>& bound_axis_sums(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<AxisBound> bound_axes = {});
 
     // 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;
 
+    /// Update the running bound axis sums where the value stored at `index` is
+    /// changed by `value_change` in a given state.
+    void update_bound_axis_slice_sums(State& state, const ssize_t index,
+                                      const double value_change) const;
+
+    /// 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 information on each bound axis.
+    std::vector<AxisBound> bound_axes_info_;
+    /// Indicator variable that all axis-wise bound operators are "==".
+    bool bound_axis_ops_all_equals_;
 };
 
 /// A contiguous block of integer numbers.
@@ -134,33 +199,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- and axis-wise bounds.
+    // Index-wise bounds default to the specified default bounds. By default,
+    // there are no axis-wise bounds.
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
 
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
 
     IntegerNode(std::span<const ssize_t> shape, std::optional<std::vector<double>> lower_bound,
-                double upper_bound);
+                double upper_bound, std::vector<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
+    IntegerNode(ssize_t size, std::optional<std::vector<double>> lower_bound, double upper_bound,
+                std::vector<AxisBound> bound_axes = {});
+
+    IntegerNode(std::span<const ssize_t> shape, double lower_bound, double upper_bound,
+                std::vector<AxisBound> bound_axes = {});
+    IntegerNode(std::initializer_list<ssize_t> shape, double lower_bound, double upper_bound,
+                std::vector<AxisBound> bound_axes = {});
+    IntegerNode(ssize_t size, double lower_bound, double upper_bound,
+                std::vector<AxisBound> bound_axes = {});
 
     // Overloads needed by the Node ABC ***************************************
 
@@ -190,33 +267,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- and axis-wise bounds.
+    // Index-wise bounds default to lower_bound = 0.0 and upper_bound = 1.0. By
+    // default, there are no axis-wise bounds.
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
 
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
 
     BinaryNode(std::span<const ssize_t> shape, std::optional<std::vector<double>> lower_bound,
-               double upper_bound);
+               double upper_bound, std::vector<AxisBound> bound_axes = {});
     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<AxisBound> bound_axes = {});
+    BinaryNode(ssize_t size, std::optional<std::vector<double>> lower_bound, double upper_bound,
+               std::vector<AxisBound> bound_axes = {});
+
+    BinaryNode(std::span<const ssize_t> shape, double lower_bound, double upper_bound,
+               std::vector<AxisBound> bound_axes = {});
+    BinaryNode(std::initializer_list<ssize_t> shape, double lower_bound, double upper_bound,
+               std::vector<AxisBound> bound_axes = {});
+    BinaryNode(ssize_t size, double lower_bound, double upper_bound,
+               std::vector<AxisBound> bound_axes = {});
 
     // 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

@@ -19,16 +19,34 @@ 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 AxisBound:
+            # 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::AxisBound::Operator":
+                Equal
+                LessEqual
+                GreaterEqual
+
+            AxisBound(Py_ssize_t axis, vector[Operator] axis_operators,
+                      vector[double] axis_bounds)
+
+            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[AxisBound] axis_wise_bounds()
+
+    cdef cppclass IntegerNode(NumberNode):
+        pass
+
+    cdef cppclass BinaryNode(IntegerNode):
+        pass