QOL improvements for DisjointLists

dwave/optimization/_model.pyx

@@ -1144,10 +1144,9 @@ cdef class Symbol:
 
             >>> from dwave.optimization.model import Model
             >>> model = Model()
-            >>> lsymbol, lsymbol_lists = model.disjoint_lists(
-            ...     primary_set_size=5,
-            ...     num_disjoint_lists=2)
-            >>> lsymbol_lists[0].equals(next(lsymbol.iter_successors()))
+            >>> x = model.binary()
+            >>> y = x + 5
+            >>> y.equals(next(x.iter_successors()))
             True
         """
         cdef vector[cppNode.SuccessorView].const_iterator it = self.node_ptr.successors().begin()
@@ -1233,17 +1232,17 @@ cdef class Symbol:
 
             >>> from dwave.optimization import Model
             >>> model = Model()
-            >>> lsymbol, lsymbol_lists = model.disjoint_lists(primary_set_size=5, num_disjoint_lists=2)
+            >>> lsymbol = model.disjoint_lists_symbol(primary_set_size=5, num_disjoint_lists=2)
             >>> with model.lock():
             ...     model.states.resize(2)
             ...     lsymbol.set_state(0, [[0, 4], [1, 2, 3]])
             ...     lsymbol.set_state(1, [[3, 4], [0, 1, 2]])
-            ...     print(f"state 0: {lsymbol_lists[0].state(0)} and {lsymbol_lists[1].state(0)}")
-            ...     print(f"state 1: {lsymbol_lists[0].state(1)} and {lsymbol_lists[1].state(1)}")
+            ...     print(f"state 0: {lsymbol[0].state(0)} and {lsymbol[1].state(0)}")
+            ...     print(f"state 1: {lsymbol[0].state(1)} and {lsymbol[1].state(1)}")
             ...     lsymbol.reset_state(0)
             ...     print("After reset:")
-            ...     print(f"state 0: {lsymbol_lists[0].state(0)} and {lsymbol_lists[1].state(0)}")
-            ...     print(f"state 1: {lsymbol_lists[0].state(1)} and {lsymbol_lists[1].state(1)}")
+            ...     print(f"state 0: {lsymbol[0].state(0)} and {lsymbol[1].state(0)}")
+            ...     print(f"state 1: {lsymbol[0].state(1)} and {lsymbol[1].state(1)}")
             state 0: [0. 4.] and [1. 2. 3.]
             state 1: [3. 4.] and [0. 1. 2.]
             After reset:

dwave/optimization/generators.py

@@ -267,7 +267,7 @@ def capacitated_vehicle_routing(demand: numpy.typing.ArrayLike,
         A model encoding the CVRP problem.
 
     Notes:
-        The model uses a :class:`~dwave.optimization.model.Model.disjoint_lists`
+        The model uses a :class:`~dwave.optimization.symbols.DisjointLists`
         class as the decision variable being optimized, with permutations of its
         sublist representing various itineraries for each vehicle.
     """
@@ -381,7 +381,7 @@ class as the decision variable being optimized, with permutations of its
     capacity = model.constant(vehicle_capacity)
 
     # Add the decision variable
-    routes_decision, routes = model.disjoint_lists(
+    routes = model.disjoint_lists_symbol(
         primary_set_size=num_customers,
         num_disjoint_lists=number_of_vehicles)
 
@@ -455,7 +455,7 @@ def capacitated_vehicle_routing_with_time_windows(demand: numpy.typing.ArrayLike
         A model encoding the CVRPTW problem.
 
     Notes:
-        The model uses a :class:`~dwave.optimization.model.Model.disjoint_lists`
+        The model uses a :class:`~dwave.optimization.symbols.DisjointLists`
         class as the decision variable being optimized, with permutations of its
         sublist representing various itineraries for each vehicle.
     """
@@ -561,7 +561,7 @@ class as the decision variable being optimized, with permutations of its
     one = model.constant(1)
 
     # Add the decision variable
-    routes_decision, routes = model.disjoint_lists(
+    routes = model.disjoint_lists_symbol(
         primary_set_size=num_customers,
         num_disjoint_lists=number_of_vehicles)
 

dwave/optimization/model.py

@@ -32,6 +32,7 @@
 import numpy as np
 import tempfile
 import typing
+import warnings
 
 from dwave.optimization._model import ArraySymbol, _Graph, Symbol
 from dwave.optimization.states import States
@@ -331,11 +332,65 @@ def disjoint_lists(
             >>> from dwave.optimization.model import Model
             >>> model = Model()
             >>> destinations, routes = model.disjoint_lists(10, 4)
+
+        .. deprecated:: 0.6.7
+
+            The return behavior of this method will be changed in
+            dwave.optimization 0.8.0. Use :meth:`.disjoint_lists_symbol`.
+        """
+
+        warnings.warn(
+            "The return behavior of Model.disjoint_lists() is deprecated "
+            "since dwave.optimization 0.6.7 and will be changed to the "
+            "behavior of Model.disjoint_lists_symbol() in 0.8.0. Use "
+            "Model.disjoint_lists_symbol().",
+            DeprecationWarning,
+        )
+
+        disjoint_lists = self.disjoint_lists_symbol(
+            primary_set_size, num_disjoint_lists
+        )
+        return disjoint_lists, list(disjoint_lists)
+
+    def disjoint_lists_symbol(
+            self,
+            primary_set_size: int,
+            num_disjoint_lists: int,
+            ) -> DisjointLists:
+        """Create a disjoint-lists symbol as a decision variable.
+
+        Divides a set of the elements of ``range(primary_set_size)`` into
+        ``num_disjoint_lists`` ordered partitions.
+
+        Args:
+            primary_set_size: Number of elements in the primary set to
+                be partitioned into disjoint lists.
+            num_disjoint_lists: Number of disjoint lists.
+
+        Returns:
+            A disjoint-lists symbol.
+
+        Examples:
+            This example creates a symbol of 10 elements that is divided
+            into 4 lists.
+
+            >>> from dwave.optimization.model import Model
+            >>> model = Model()
+            >>> disjoint_lists = model.disjoint_lists_symbol(10, 4)
+            >>> disjoint_lists.primary_set_size()
+            10
+            >>> disjoint_lists.num_disjoint_lists()
+            4
         """
         from dwave.optimization.symbols import DisjointLists, DisjointList  # avoid circular import
-        main = DisjointLists(self, primary_set_size, num_disjoint_lists)
-        lists = [DisjointList(main, i) for i in range(num_disjoint_lists)]
-        return main, lists
+        disjoint_lists = DisjointLists(self, primary_set_size, num_disjoint_lists)
+
+        # create the DisjointList symbols, which will create the successor nodes, even
+        # though we won't use them directly here
+        for i in range(num_disjoint_lists):
+            DisjointList(disjoint_lists, i)
+
+        return disjoint_lists
 
     def feasible(self, index: int = 0) -> bool:
         """Check the feasibility of the state at the input index.

dwave/optimization/symbols/collections.pyx

@@ -305,7 +305,7 @@ cdef class DisjointList(ArraySymbol):
     """
     def __init__(self, DisjointLists parent, Py_ssize_t list_index):
         if list_index < 0 or list_index >= parent.num_disjoint_lists():
-            raise ValueError(
+            raise IndexError(
                 "`list_index` must be less than the number of disjoint sets of the parent"
             )
 
@@ -411,6 +411,9 @@ cdef class DisjointLists(Symbol):
 
         self.initialize_node(model, self.ptr)
 
+    def __getitem__(self, index: int):
+        return DisjointList(self, index)
+
     @classmethod
     def _from_symbol(cls, Symbol symbol):
         cdef DisjointListsNode* ptr = dynamic_cast_ptr[DisjointListsNode](symbol.node_ptr)
@@ -491,21 +494,21 @@ cdef class DisjointLists(Symbol):
                 Index of the state to set
             state:
                 Assignment of values for the state.
-        
+
         Examples:
             This example sets the state of a disjoint-lists symbol. You can
             inspect the state of each list individually.
-            
+
             >>> from dwave.optimization.model import Model
             >>> model = Model()
-            >>> lists_symbol, lists_array = model.disjoint_lists(
+            >>> lists_symbol = model.disjoint_lists_symbol(
             ...     primary_set_size=5,
             ...     num_disjoint_lists=3
             ... )
             >>> with model.lock():
             ...     model.states.resize(1)
             ...     lists_symbol.set_state(0, [[0, 1, 2, 3], [4], []])
-            ...     for index, disjoint_list in enumerate(lists_array):
+            ...     for index, disjoint_list in enumerate(lists_symbol):
             ...         print(f"DisjointList {index}:")
             ...         print(disjoint_list.state(0))
             DisjointList 0:
@@ -569,6 +572,10 @@ cdef class DisjointLists(Symbol):
         """Return the number of disjoint lists in the symbol."""
         return self.ptr.num_disjoint_lists()
 
+    def primary_set_size(self):
+        """Return the size of primary set of elements that the lists contain."""
+        return self.ptr.primary_set_size()
+
     # An observing pointer to the C++ DisjointListsNode
     cdef DisjointListsNode* ptr
 

releasenotes/notes/improve-model-disjoint-lists-interface-462997c98ae71684.yaml

@@ -0,0 +1,14 @@
+---
+features:
+  - |
+    ``DisjointLists`` symbols are now indexable, returning the corresponding
+    ``DisjointList`` (singular) symbol for the given index. This allows the
+    method of creating disjoint lists on the model to be simplified, returning
+    only a ``DisjointLists`` symbol instead of a tuple of the symbol and a
+    list of ``DisjointList`` symbols. A new method
+    ``Model.disjoint_lists_symbol()`` has been added to the Model class which
+    implements this.
+deprecations:
+  - |
+    The ``Model.disjoint_lists()`` method has been deprecated. Use
+    ``Model.disjoint_lists_symbol()`` instead.

tests/test_model.py

@@ -433,19 +433,18 @@ def test_remove_unused_symbols(self):
         with self.subTest("disjoint lists"):
             model = Model()
 
-            base, lists = model.disjoint_lists(10, 4)
+            disjoint_lists = model.disjoint_lists_symbol(10, 4)
 
             # only use some of the lists
-            model.minimize(lists[0].sum())
-            model.add_constraint(lists[1].sum() <= model.constant(3))
+            model.minimize(disjoint_lists[0].sum())
+            model.add_constraint(disjoint_lists[1].sum() <= model.constant(3))
 
-            lists[2].prod()  # this one will hopefully be removed
+            disjoint_lists[2].prod()  # this one will hopefully be removed
 
             self.assertEqual(model.num_symbols(), 10)
 
             # make sure they aren't being kept alive by other objects
-            del lists
-            del base
+            del disjoint_lists
 
             num_removed = model.remove_unused_symbols()
 

tests/test_symbols.py

@@ -1274,50 +1274,80 @@ def test_inequality(self):
 
     def generate_symbols(self):
         model = Model()
-        d, ds = model.disjoint_lists(10, 4)
+        d = model.disjoint_lists_symbol(10, 4)
         model.lock()
         yield d
-        yield from ds
+        yield from d
 
     def test(self):
         model = Model()
 
-        model.disjoint_lists(10, 4)
+        dls = model.disjoint_lists_symbol(10, 4)
+
+        self.assertEqual(dls.primary_set_size(), 10)
+        self.assertEqual(dls.num_disjoint_lists(), 4)
+
+    def test_deprecated_creation_method(self):
+        model = Model()
+        with self.assertWarnsRegex(
+            DeprecationWarning,
+            r"The return behavior of Model.disjoint_lists\(\) is deprecated"
+        ):
+            d, dls = model.disjoint_lists(10, 4)
+
+        self.assertIsInstance(d, dwave.optimization.symbols.DisjointLists)
+        self.assertEqual(len(dls), 4)
+        self.assertIsInstance(dls[0], dwave.optimization.symbols.DisjointList)
+
+    def test_indexing(self):
+        model = Model()
+
+        dls = model.disjoint_lists_symbol(10, 4)
+
+        self.assertEqual(len(list(dls)), 4)
+        self.assertIsInstance(dls[0], dwave.optimization.symbols.DisjointList)
+        self.assertIsInstance(dls[3], dwave.optimization.symbols.DisjointList)
+
+        with self.assertRaises(IndexError):
+            dls[4]
 
     def test_construction(self):
         model = Model()
 
         with self.assertRaises(ValueError):
-            model.disjoint_lists(-5, 1)
+            model.disjoint_lists_symbol(-5, 1)
         with self.assertRaises(ValueError):
-            model.disjoint_lists(1, -5)
+            model.disjoint_lists_symbol(1, -5)
 
         model.states.resize(1)
 
-        ds, (x,) = model.disjoint_lists(0, 1)
-        self.assertEqual(x.shape(), (-1,))  # todo: handle this special case
+        ds = model.disjoint_lists_symbol(0, 1)
+        self.assertEqual(ds[0].shape(), (-1,))  # todo: handle this special case
 
     def test_num_returned_nodes(self):
         model = Model()
 
-        d, ds = model.disjoint_lists(10, 4)
+        model.disjoint_lists_symbol(10, 4)
+
+        # One DisjointListsNode, and one node for each of the 4 successor lists
+        self.assertEqual(model.num_nodes(), 5)
 
     def test_set_state(self):
         with self.subTest("array-like output lists"):
             model = Model()
             model.states.resize(1)
-            x, ys = model.disjoint_lists(5, 3)
+            x = model.disjoint_lists_symbol(5, 3)
             model.lock()
 
             x.set_state(0, [[0, 1], [2, 3], [4]])
 
-            np.testing.assert_array_equal(ys[0].state(), [0, 1])
-            np.testing.assert_array_equal(ys[1].state(), [2, 3])
-            np.testing.assert_array_equal(ys[2].state(), [4])
+            np.testing.assert_array_equal(x[0].state(), [0, 1])
+            np.testing.assert_array_equal(x[1].state(), [2, 3])
+            np.testing.assert_array_equal(x[2].state(), [4])
 
         with self.subTest("invalid state index"):
             model = Model()
-            x, _ = model.disjoint_lists(5, 3)
+            x = model.disjoint_lists_symbol(5, 3)
 
             state = [[0, 1, 2, 3, 4], [], []]
 
@@ -1338,16 +1368,16 @@ def test_set_state(self):
             # gets translated into integer according to NumPy rules
             model = Model()
             model.states.resize(1)
-            x, ys = model.disjoint_lists(5, 3)
+            x = model.disjoint_lists_symbol(5, 3)
             model.lock()
 
             x.set_state(0, [[4.5, 3, 2, 1, 0], [], []])
-            np.testing.assert_array_equal(ys[0].state(), [4, 3, 2, 1, 0])
+            np.testing.assert_array_equal(x[0].state(), [4, 3, 2, 1, 0])
 
         with self.subTest("invalid"):
             model = Model()
             model.states.resize(1)
-            x, ys = model.disjoint_lists(5, 3)
+            x = model.disjoint_lists_symbol(5, 3)
             model.lock()
 
             with self.assertRaisesRegex(
@@ -1376,10 +1406,10 @@ def test_set_state(self):
     def test_state_size(self):
         model = Model()
 
-        d, ds = model.disjoint_lists(10, 4)
+        d = model.disjoint_lists_symbol(10, 4)
 
         self.assertEqual(d.state_size(), 0)
-        for s in ds:
+        for s in d:
             self.assertEqual(s.state_size(), 10 * 8)