Repository: dwavesystems/dimod
Pull request #1382 (2024-08-02 17:00:21)

Add labels keyword argument to CQM.iter_constraint_data()
See https://github.com/dwavesystems/dimod/issues/1381. This addresses one of the options laid out.


Patch:
diff --git a/dimod/constrained/constrained.py b/dimod/constrained/constrained.py
index 3704a8a2..04f8a102 100644
--- a/dimod/constrained/constrained.py
+++ b/dimod/constrained/constrained.py
@@ -1144,7 +1144,11 @@ class ConstrainedQuadraticModel(cyConstrainedQuadraticModel):
 
         return cqm
 
-    def iter_constraint_data(self, sample_like: SamplesLike) -> Iterator[ConstraintData]:
+    def iter_constraint_data(self,
+                             sample_like: SamplesLike,
+                             *,
+                             labels: typing.Optional[typing.Iterable[typing.Hashable]] = None,
+                             ) -> Iterator[ConstraintData]:
         """Yield information about the constraints for the given sample.
 
         Note that this method iterates over constraints in the same order as
@@ -1153,6 +1157,7 @@ class ConstrainedQuadraticModel(cyConstrainedQuadraticModel):
         Args:
             sample_like: A sample. `sample-like` is an extension of
                 NumPy's array_like structure. See :func:`.as_samples`.
+            labels: A subset of the constraint labels over which to iterate.
 
         Yields:
             A :class:`collections.namedtuple` with the following fields.
@@ -1185,14 +1190,23 @@ class ConstrainedQuadraticModel(cyConstrainedQuadraticModel):
         """
         # We go ahead and coerce to Variables for performance, since .energies() prefers
         # that format
-        sample, labels = as_samples(sample_like, labels_type=Variables)
+        sample, variable_labels = as_samples(sample_like, labels_type=Variables)
 
         if sample.shape[0] != 1:
             raise ValueError("sample_like should be a single sample, "
                              f"received {sample.shape[0]} samples")
 
-        for label, constraint in self.constraints.items():
-            lhs = constraint.lhs.energy((sample, labels))
+        # by default iterate over all constraints in the model
+        if labels is None:
+            labels = self.constraint_labels
+
+        for label in labels:
+            try:
+                constraint = self.constraints[label]
+            except KeyError as err:
+                raise ValueError(f"unknown constraint label: {label!r}") from err
+
+            lhs = constraint.lhs.energy((sample, variable_labels))
             rhs = constraint.rhs
             sense = constraint.sense
 
@@ -1219,6 +1233,7 @@ class ConstrainedQuadraticModel(cyConstrainedQuadraticModel):
     def iter_violations(self, sample_like: SamplesLike, *,
                         skip_satisfied: bool = False,
                         clip: bool = False,
+                        labels: typing.Optional[typing.Iterable[typing.Hashable]] = None,
                         ) -> Iterator[Tuple[Hashable, Bias]]:
         """Yield violations for all constraints.
 
@@ -1227,6 +1242,7 @@ class ConstrainedQuadraticModel(cyConstrainedQuadraticModel):
                 NumPy's array_like structure. See :func:`.as_samples`..
             skip_satisfied: If True, does not yield constraints that are satisfied.
             clip: If True, negative violations are rounded up to 0.
+            labels: A subset of the constraint labels over which to iterate.
 
         Yields:
             A 2-tuple containing the constraint label and the amount of the
@@ -1273,14 +1289,14 @@ class ConstrainedQuadraticModel(cyConstrainedQuadraticModel):
         if skip_satisfied:
             # clip doesn't matter in this case
             # todo: feasibility tolerance?
-            for datum in self.iter_constraint_data(sample_like):
+            for datum in self.iter_constraint_data(sample_like, labels=labels):
                 if datum.violation > 0:
                     yield datum.label, datum.violation
         elif clip:
-            for datum in self.iter_constraint_data(sample_like):
+            for datum in self.iter_constraint_data(sample_like, labels=labels):
                 yield datum.label, max(datum.violation, 0.0)
         else:
-            for datum in self.iter_constraint_data(sample_like):
+            for datum in self.iter_constraint_data(sample_like, labels=labels):
                 yield datum.label, datum.violation
 
     def is_almost_equal(self, other: 'ConstrainedQuadraticModel',
diff --git a/releasenotes/notes/feature-CQM.iter_constraint_data-labels-9ffdc5a2db6f09c4.yaml b/releasenotes/notes/feature-CQM.iter_constraint_data-labels-9ffdc5a2db6f09c4.yaml
new file mode 100644
index 00000000..4cf02cac
--- /dev/null
+++ b/releasenotes/notes/feature-CQM.iter_constraint_data-labels-9ffdc5a2db6f09c4.yaml
@@ -0,0 +1,8 @@
+---
+features:
+  - |
+    Add ``labels`` keyword argument to ``ConstrainedQuadraticModel.iter_constraint_data()``
+    and ``ConstrainedQuadraticModel.iter_violations()`` methods.
+    The ``labels`` keyword argument allows a user to generate the data/violations for
+    a subset of the constraints.
+    See `#1381 <https://github.com/dwavesystems/dimod/issues/1381>`_.


Test patch:
diff --git a/tests/test_constrained.py b/tests/test_constrained.py
index e0e2eec5..63ecd031 100644
--- a/tests/test_constrained.py
+++ b/tests/test_constrained.py
@@ -676,6 +676,37 @@ class TestIterViolations(unittest.TestCase):
         sample = {'j': -1, 'i': 1004, 'k': 1000}
         self.assertEqual(cqm.violations(sample, clip=False), {label_ge: -3.0, label_le: -1.0})
 
+    def test_iter_violations(self):
+        cqm = CQM()
+        x, y, z = dimod.Binaries(['x', 'y', 'z'])
+
+        c0 = cqm.add_constraint(x + y <= 1, label="c0")
+        c1 = cqm.add_constraint(x - y >= 1, label="c1")
+        c2 = cqm.add_constraint(z >= 0, label="c2")
+
+        with self.subTest("no other kwargs"):
+            data = list(cqm.iter_violations({"x": 0, "y": 1, "z": 0}, labels=[c0, c1]))
+
+            self.assertEqual(len(data), 2)
+            self.assertEqual(data[0][0], c0)  # ordered by labels
+            self.assertEqual(data[1][0], c1)
+
+        with self.subTest("with clip kwarg"):
+            data = list(cqm.iter_violations({"x": 0, "y": 1, "z": 0}, labels=[c0, c1], clip=True))
+
+            self.assertEqual(len(data), 2)
+            self.assertEqual(data[0][0], c0)  # ordered by labels
+            self.assertEqual(data[1][0], c1)
+
+        with self.subTest("with skip_satisfied kwarg"):
+            data = list(cqm.iter_violations({"x": 0, "y": 1, "z": 0}, labels=[c0, c1], skip_satisfied=True))
+
+            self.assertEqual(len(data), 1)
+            self.assertEqual(data[0][0], c1)
+
+        with self.assertRaisesRegex(ValueError, "unknown constraint label: 'unknown label'"):
+            list(cqm.iter_violations({"x": 0, "y": 1, "z": 0}, labels=[c0, "unknown label"]))
+
 
 class TestIsAlmostEqual(unittest.TestCase):
     def test_simple(self):
@@ -1790,6 +1821,23 @@ class TestIterConstraintData(unittest.TestCase):
         self.assertEqual(list(cqm.constraints.keys()),
             [datum.label for datum in cqm.iter_constraint_data({'x': 1, 'y': 0, 'z': 0})])
 
+    def test_labels(self):
+        cqm = CQM()
+        x, y, z = dimod.Binaries(['x', 'y', 'z'])
+
+        c0 = cqm.add_constraint(x + y <= 1, label="c0")
+        c1 = cqm.add_constraint(x - y >= 1, label="c1")
+        c2 = cqm.add_constraint(z >= .5, label="c2")
+
+        data = list(cqm.iter_constraint_data({"x": 0, "y": 1, "z": 0}, labels=[c0, c1]))
+
+        self.assertEqual(len(data), 2)
+        self.assertEqual(data[0].label, c0)  # ordered by labels
+        self.assertEqual(data[1].label, c1)
+
+        with self.assertRaisesRegex(ValueError, "unknown constraint label: 'unknown label'"):
+            list(cqm.iter_constraint_data({"x": 0, "y": 1, "z": 0}, labels=[c0, "unknown label"]))
+
 
 class TestStr(unittest.TestCase):
     def test_one(self):