Enhancement/axis wise bounds NumberNode

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

@@ -28,6 +28,35 @@ 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 {
+        /// 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);
+
+        /// 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;
+
+        /// 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;
+    };
+
     NumberNode() = delete;
 
     // Overloads needed by the Array ABC **************************************
@@ -68,6 +97,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);
@@ -106,21 +140,39 @@ 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 hyperslice
+    /// 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_;
 };
 
 /// A contiguous block of integer numbers.
@@ -134,33 +186,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 +254,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,31 @@ 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_opertors, 
+                      vector[double] axis_bounds)
+            Py_ssize_t axis
+            vector[Operator] operators
+            vector[double] bounds
+
         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

dwave/optimization/model.py

@@ -165,7 +165,9 @@ def objective(self, value: ArraySymbol):
 
     def binary(self, shape: None | _ShapeLike = None,
                lower_bound: None | np.typing.ArrayLike = None,
-               upper_bound: None | np.typing.ArrayLike = None) -> BinaryVariable:
+               upper_bound: None | np.typing.ArrayLike = None,
+               subject_to: None | list[tuple(int, str | list[str], float |
+                                       list[float])] = None) -> BinaryVariable:
         r"""Create a binary symbol as a decision variable.
 
         Args:
@@ -178,6 +180,17 @@ def binary(self, shape: None | _ShapeLike = None,
                 scalar (one bound for all variables) or an array (one bound for
                 each variable). Non-boolean values are rounded down to the domain
                 [0,1]. If None, the default value of 1 is used.
+            subject_to (optional): Axis-wise bounds applied to the symbol. Must be an
+                array of tuples where each tuple has the form: (axis, operators, bounds)
+                - axis (int): The axis along which the bounds are applied.
+                - operators (str | array[str]): The operator(s) ("<=", "==", or ">="). 
+                A single operator applies to all slices along the axis; an
+                array specifies one operator per slice.
+                - bounds (float | array[float]): The bound value(s). A single value
+                applies to all slices; an array specifies one bound per slice.
+                If provided, the sum of values within each slice along the
+                specified axis must satisfy the corresponding operator–bound
+                pair. Note: At most one axis-wise bound may be provided.
 
         Returns:
             A binary symbol.
@@ -215,15 +228,33 @@ def binary(self, shape: None | _ShapeLike = None,
             >>> np.all([1, 0] == b.upper_bound())
             True
 
+            This example adds a :math:`(2x3)`-sized binary symbol with
+            index-wise lower bounds and an axis-wise bound along axis 1. Let
+            x_i (int i : 0 <= i <= 2) denote the sum of the values within
+            slice i along axis 1. For each state defined for this symbol:
+            (x_0 <= 0), (x_1 == 2), and (x_2 >= 1).
+
+            >>> from dwave.optimization.model import Model
+            >>> import numpy as np
+            >>> model = Model()
+            >>> b = model.binary([2, 3], lower_bound=[[0, 1, 1], [0, 1, 0]],
+            ... subject_to=[(1, ["<=", "==", ">="], [0, 2, 1])])
+            >>> np.all(n.axis_wise_bounds() == [(1, ["<=", "==", ">="], [0, 2, 1])])
+            True
+
         See Also:
             :class:`~dwave.optimization.symbols.numbers.BinaryVariable`: equivalent symbol.
 
         .. versionchanged:: 0.6.7
-            Beginning in version 0.6.7, user-defined bounds and index-wise
-            bounds are supported.
+            Beginning in version 0.6.7, user-defined index-wise bounds are
+            supported.
+
+        .. versionchanged:: 0.6.12
+            Beginning in version 0.6.12, user-defined axis-wise bounds are
+            supported.
         """
         from dwave.optimization.symbols import BinaryVariable  # avoid circular import
-        return BinaryVariable(self, shape, lower_bound, upper_bound)
+        return BinaryVariable(self, shape, lower_bound, upper_bound, subject_to)
 
     def constant(self, array_like: numpy.typing.ArrayLike) -> Constant:
         r"""Create a constant symbol.
@@ -478,7 +509,8 @@ def integer(
             shape: None | _ShapeLike = None,
             lower_bound: None | numpy.typing.ArrayLike = None,
             upper_bound: None | numpy.typing.ArrayLike = None,
-            ) -> IntegerVariable:
+               subject_to: None | list[tuple(int, str | list[str], float |
+                                       list[float])] = None) -> IntegerVariable:
         r"""Create an integer symbol as a decision variable.
 
         Args:
@@ -491,6 +523,17 @@ def integer(
                 scalar (one bound for all variables) or an array (one bound for
                 each variable). Non-integer values are down up. If None, the
                 default value is used.
+            subject_to (optional): Axis-wise bounds applied to the symbol. Must be an
+                array of tuples where each tuple has the form: (axis, operators, bounds)
+                - axis (int): The axis along which the bounds are applied.
+                - operators (str | array[str]): The operator(s) ("<=", "==", or ">="). 
+                A single operator applies to all slice along the axis; an array
+                specifies one operator per slice.
+                - bounds (float | array[float]): The bound value(s). A single value
+                applies to all slices; an array specifies one bound per slice.
+                If provided, the sum of values within each slice along the
+                specified axis must satisfy the corresponding operator–bound
+                pair. Note: At most one axis-wise bound may be provided.
 
         Returns:
             An integer symbol.
@@ -529,15 +572,33 @@ def integer(
             >>> np.all([1, 2] == i.upper_bound())
             True
 
+            This example adds a :math:`(2x3)`-sized integer symbol with general
+            lower and upper bounds and an axis-wise bound along axis 1. Let x_i
+            (int i : 0 <= i <= 2) denote the sum of the values within
+            slice i along axis 1. For each state defined for this symbol:
+            (x_0 <= 2), (x_1 <= 4), and (x_2 <= 5).
+
+            >>> from dwave.optimization.model import Model
+            >>> import numpy as np
+            >>> model = Model()
+            >>> i = model.integer([2, 3], lower_bound=1, upper_bound=3,
+            ... subject_to=[(1, "<=", [2, 4, 5])])
+            >>> np.all(i.axis_wise_bounds() == [(1, ["<="], [2, 4, 5])])
+            True
+
         See Also:
             :class:`~dwave.optimization.symbols.numbers.IntegerVariable`: equivalent symbol.
 
         .. versionchanged:: 0.6.7
             Beginning in version 0.6.7, user-defined index-wise bounds are
             supported.
+
+        .. versionchanged:: 0.6.12
+            Beginning in version 0.6.12, user-defined axis-wise bounds are
+            supported.
         """
         from dwave.optimization.symbols import IntegerVariable  # avoid circular import
-        return IntegerVariable(self, shape, lower_bound, upper_bound)
+        return IntegerVariable(self, shape, lower_bound, upper_bound, subject_to)
 
     def list(self,
             n: int,