Set operations and tests added

doc/releases/changelog-dev.md

@@ -32,6 +32,9 @@
 * `QuantumScript.hash` is now cached, leading to performance improvements.
   [(#5919)](https://github.com/PennyLaneAI/pennylane/pull/5919)
 
+* Set operations are now supported by Wires.
+  [(#5983)](https://github.com/PennyLaneAI/pennylane/pull/5983)
+
 * The representation for `Wires` has now changed to be more copy-paste friendly.
   [(#5958)](https://github.com/PennyLaneAI/pennylane/pull/5958)
 

pennylane/wires.py

@@ -493,6 +493,186 @@ def unique_wires(list_of_wires):
 
         return Wires(tuple(unique), _override=True)
 
+    def union(self, other):
+        """Return the union of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the union with.
+
+        Returns:
+            Wires: A new Wires object representing the union of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([3, 4, 5])
+        >>> wires1.union(wires2)
+        Wires([1, 2, 3, 4, 5])
+        """
+        return Wires(set(self.labels) | set(other.labels))
+
+    def __or__(self, other):
+        """Return the union of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the union with.
+
+        Returns:
+            Wires: A new Wires object representing the union of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([3, 4, 5])
+        >>> wires1 | wires2
+        Wires([1, 2, 3, 4, 5])
+        """
+        if not isinstance(other, Wires):
+            raise TypeError(f"Can only do the union of Wires with Wires. Got {type(other)}.")
+        return self.union(other)
+
+    def intersection(self, other):
+        """Return the intersection of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the intersection with.
+
+        Returns:
+            Wires: A new Wires object representing the intersection of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([2, 3, 4])
+        >>> wires1.intersection(wires2)
+        Wires([2, 3])
+        """
+        return Wires(set(self.labels) & set(other.labels))
+
+    def __and__(self, other):
+        """Return the intersection of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the intersection with.
+
+        Returns:
+            Wires: A new Wires object representing the intersection of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([2, 3, 4])
+        >>> wires1 & wires2
+        Wires([2, 3])
+        """
+        if not isinstance(other, Wires):
+            raise TypeError(f"Can only do the intersection of Wires with Wires. Got {type(other)}.")
+        return self.intersection(other)
+
+    def difference(self, other):
+        """Return the difference of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the difference with.
+
+        Returns:
+            Wires: A new Wires object representing the difference of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([2, 3, 4])
+        >>> wires1.difference(wires2)
+        Wires([1])
+        """
+        return Wires(set(self.labels) - set(other.labels))
+
+    def __sub__(self, other):
+        """Return the difference of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the difference with.
+
+        Returns:
+            Wires: A new Wires object representing the difference of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([2, 3, 4])
+        >>> wires1 - wires2
+        Wires([1])
+        """
+        if not isinstance(other, Wires):
+            raise TypeError(
+                f"Can only do the difference of Wires with other Wires. Got {type(other)}."
+            )
+        return self.difference(other)
+
+    def symmetric_difference(self, other):
+        """Return the symmetric difference of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the symmetric difference with.
+
+        Returns:
+            Wires: A new Wires object representing the symmetric difference of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([3, 4, 5])
+        >>> wires1.symmetric_difference(wires2)
+        Wires([1, 2, 4, 5])
+        """
+        return Wires(set(self.labels) ^ set(other.labels))
+
+    def __xor__(self, other):
+        """Return the symmetric difference of the current Wires object and another Wires object.
+
+        Args:
+            other (Wires): Another Wires object to perform the symmetric difference with.
+
+        Returns:
+            Wires: A new Wires object representing the symmetric difference of the two Wires objects.
+
+        Raises:
+            TypeError: If `other` is not an instance of Wires.
+
+        **Example**
+
+        >>> wires1 = Wires([1, 2, 3])
+        >>> wires2 = Wires([3, 4, 5])
+        >>> wires1 ^ wires2
+        Wires([1, 2, 4, 5])
+        """
+        if not isinstance(other, Wires):
+            raise TypeError(
+                f"Can only the symmetric difference of Wires with other Wires. Got {type(other)}."
+            )
+        return self.symmetric_difference(other)
+
 
 # Register Wires as a PyTree-serializable class
 register_pytree(Wires, Wires._flatten, Wires._unflatten)  # pylint: disable=protected-access

tests/test_wires.py

@@ -377,3 +377,138 @@ def test_wires_pytree(self, source):
         wires2 = tree_unflatten(tree, wires_flat)
         assert isinstance(wires2, Wires), f"{wires2} is not Wires"
         assert wires == wires2, f"{wires} != {wires2}"
+
+    @pytest.mark.parametrize(
+        "wire_a, wire_b, expected",
+        [
+            (Wires([0, 1]), Wires([2, 3]), Wires([0, 1, 2, 3])),
+            (Wires([]), Wires([1, 2, 3]), Wires([1, 2, 3])),
+            (Wires([1, 2]), Wires([1, 2, 3]), Wires([1, 2, 3])),
+        ],
+    )
+    def test_union(self, wire_a, wire_b, expected):
+        """
+        Test the union operation (|) between two Wires objects.
+        """
+        assert wire_a | wire_b == expected
+        assert wire_a.union(wire_b) == expected
+
+    @pytest.mark.parametrize(
+        "wire_a, wire_b, expected",
+        [
+            (Wires([0, 1, 2]), Wires([2, 3, 4]), Wires([2])),
+            (Wires([1, 2, 3]), Wires([]), Wires([])),
+            (Wires([1, 2, 3]), Wires([1, 2, 3, 4]), Wires([1, 2, 3])),
+        ],
+    )
+    def test_intersection(self, wire_a, wire_b, expected):
+        """
+        Test the intersection operation (&) between two Wires objects.
+        """
+        assert wire_a & wire_b == expected
+        assert wire_a.intersection(wire_b) == expected
+
+    @pytest.mark.parametrize(
+        "wire_a, wire_b, expected",
+        [
+            (Wires([0, 1, 2, 3]), Wires([2, 3, 4]), Wires([0, 1])),
+            (Wires([1, 2, 3]), Wires([]), Wires([1, 2, 3])),
+            (Wires([1, 2, 3]), Wires([1, 2, 3, 4]), Wires([])),
+        ],
+    )
+    def test_difference(self, wire_a, wire_b, expected):
+        """
+        Test the difference operation (-) between two Wires objects.
+        """
+        assert wire_a - wire_b == expected
+        assert wire_a.difference(wire_b) == expected
+
+    @pytest.mark.parametrize(
+        "wire_a, wire_b, expected",
+        [
+            (Wires([0, 1, 2]), Wires([2, 3, 4]), Wires([0, 1, 3, 4])),
+            (Wires([]), Wires([1, 2, 3]), Wires([1, 2, 3])),
+            (Wires([1, 2, 3]), Wires([1, 2, 3]), Wires([])),
+        ],
+    )
+    def test_symmetric_difference(self, wire_a, wire_b, expected):
+        """
+        Test the symmetric difference operation (^) between two Wires objects.
+        """
+        assert wire_a ^ wire_b == expected
+        assert wire_a.symmetric_difference(wire_b) == expected
+
+    @pytest.mark.parametrize(
+        "wire_a, wire_b, wire_c, wire_d, expected",
+        [
+            (
+                Wires([0, 1]),
+                Wires([2, 3]),
+                Wires([4, 5]),
+                Wires([6, 7]),
+                Wires([0, 1, 2, 3, 4, 5, 6, 7]),
+            ),
+            (Wires([0, 1]), Wires([1, 2]), Wires([2, 3]), Wires([3, 4]), Wires([0, 1, 2, 3, 4])),
+            (Wires([]), Wires([1, 2]), Wires([2, 3]), Wires([3, 4, 5]), Wires([1, 2, 3, 4, 5])),
+        ],
+    )
+    # pylint: disable=too-many-arguments
+    def test_multiple_union(self, wire_a, wire_b, wire_c, wire_d, expected):
+        """
+        Test the union operation (|) with multiple Wires objects.
+        """
+        result = wire_a | wire_b | wire_c | wire_d
+        assert result == expected
+        assert wire_a.union(wire_b.union(wire_c.union(wire_d))) == expected
+
+    def test_complex_operation(self):
+        """
+        Test a complex operation involving multiple set operations.
+        This test combines union, intersection, difference, and symmetric difference operations.
+        """
+        wire_a = Wires([0, 1, 2, 3])
+        wire_b = Wires([2, 3, 4, 5])
+        wire_c = Wires([4, 5, 6, 7])
+        wire_d = Wires([6, 7, 8, 9])
+
+        # ((A ∪ B) ∩ (C ∪ D)) ^ ((A - D) ∪ (C - B))
+        result = ((wire_a | wire_b) & (wire_c | wire_d)) ^ ((wire_a - wire_d) | (wire_c - wire_b))
+        assert (wire_a | wire_b) & (wire_c | wire_d) == Wires([4, 5])
+        assert (wire_a - wire_d) | (wire_c - wire_b) == Wires([0, 1, 2, 3, 6, 7])
+
+        expected = Wires([0, 1, 2, 3, 4, 5, 6, 7])
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "operation, error_message",
+        [
+            (lambda a, b: a | b, "Can only do the union of Wires with Wires"),
+            (lambda a, b: a & b, "Can only do the intersection of Wires with Wires"),
+            (lambda a, b: a - b, "Can only do the difference of Wires with other Wires"),
+            (lambda a, b: a ^ b, "Can only the symmetric difference of Wires with other Wires"),
+        ],
+    )
+    def test_type_error_raised(self, operation, error_message):
+        """
+        Test that TypeError is raised with the correct error message when performing
+        set operations between a Wires object and an incompatible type.
+
+        This test covers union (|), intersection (&), difference (-),
+        and symmetric difference (^) operations.
+        """
+        wire = Wires([1, 2, 3])
+        invalid_operands = [
+            42,
+            "string",
+            [1, 2, 3],
+            {1, 2, 3},
+            (1, 2, 3),
+        ]
+
+        for invalid_operand in invalid_operands:
+            with pytest.raises(TypeError, match=error_message):
+                operation(wire, invalid_operand)
+
+            # Ensure the reverse also has a TypeError
+            with pytest.raises(TypeError, match="unsupported operand"):
+                operation(invalid_operand, wire)