Add mask for entangling gates

main.py

@@ -4,7 +4,7 @@
 import pennylane as qml
 from pennylane import numpy as np
 
-from maskit.masks import MaskedCircuit, PerturbationAxis, PerturbationMode
+from maskit.masks import Mask, MaskedCircuit, PerturbationAxis, PerturbationMode
 from maskit.iris import load_iris
 from maskit.utils import cross_entropy, check_params
 from maskit.circuits import variational_circuit, iris_circuit
@@ -61,6 +61,7 @@ def init_parameters(
         layers=layers,
         wires=wires,
         default_value=default_value,
+        entangling_mask=Mask(shape=(layers, wires - 1)),
         dynamic_parameters=dynamic_parameters,
     )
     mc.layer_mask[current_layers:] = True

maskit/circuits.py

@@ -24,8 +24,18 @@ def basic_variational_circuit(params, rotations, masked_circuit: MaskedCircuit):
             rotation(full_parameters[layer][wire], wires=wire)
 
         for wire in range(0, wires - 1, 2):
+            if (
+                masked_circuit.entangling_mask is not None
+                and masked_circuit.entangling_mask[layer, wire]
+            ):
+                continue
             qml.CZ(wires=[wire, wire + 1])
         for wire in range(1, wires - 1, 2):
+            if (
+                masked_circuit.entangling_mask is not None
+                and masked_circuit.entangling_mask[layer, wire]
+            ):
+                continue
             qml.CZ(wires=[wire, wire + 1])
 
 

maskit/masks.py

@@ -13,6 +13,8 @@ class PerturbationAxis(Enum):
     LAYERS = 1
     #: Perturbation affects random locations in parameter mask
     RANDOM = 2
+    #: Perturbation affects entangling gates
+    ENTANGLING = 3
 
 
 class PerturbationMode(Enum):
@@ -199,12 +201,15 @@ class MaskedCircuit(object):
     :param parameter_mask: Initialization values of paramater mask, defaults to `None`
     :param layer_mask: Initialization values of layer mask, defaults to `None`
     :param wire_mask: Initialization values of wire mask, defaults to `None`
+    :param entangling_mask: The mask to apply for entangling gates within the circuit,
+        defaults to None
     """
 
     __slots__ = (
         "_layer_mask",
         "_wire_mask",
         "_parameter_mask",
+        "_entangling_mask",
         "parameters",
         "default_value",
         "_dynamic_parameters",
@@ -220,6 +225,7 @@ def __init__(
         parameter_mask: Optional[np.ndarray] = None,
         layer_mask: Optional[np.ndarray] = None,
         wire_mask: Optional[np.ndarray] = None,
+        entangling_mask: Optional[Mask] = None,
     ):
         assert (
             layers == parameters.shape[0]
@@ -234,6 +240,9 @@ def __init__(
         self._layer_mask = Mask(shape=(layers,), parent=self, mask=layer_mask)
         self._wire_mask = Mask(shape=(wires,), parent=self, mask=wire_mask)
         self.default_value = default_value
+        if entangling_mask is not None:
+            assert layers == entangling_mask.shape[0]
+        self._entangling_mask = entangling_mask
         self._dynamic_parameters = dynamic_parameters
 
     @property
@@ -266,7 +275,10 @@ def mask(self) -> np.ndarray:
         return mask
 
     def active(self) -> int:
-        """Number of active gates in the circuit."""
+        """
+        Number of active gates in the circuit based on layer, wire, and parameter mask.
+        Entangling gates are not included.
+        """
         mask = self.mask
         return mask.size - np.sum(mask)
 
@@ -285,6 +297,11 @@ def parameter_mask(self):
         """Returns the encapsulated parameter mask."""
         return self._parameter_mask
 
+    @property
+    def entangling_mask(self):
+        """Returns the encapsulated mask of entangling gates."""
+        return self._entangling_mask
+
     def perturb(
         self,
         axis: PerturbationAxis = PerturbationAxis.RANDOM,
@@ -315,6 +332,9 @@ def perturb(
             self._wire_mask.perturb(amount=amount, mode=mode)
         elif axis == PerturbationAxis.RANDOM:  # Axis is on parameters
             self._parameter_mask.perturb(amount=amount, mode=mode)
+        elif axis == PerturbationAxis.ENTANGLING:
+            if self._entangling_mask:
+                self._entangling_mask.perturb(amount=amount, mode=mode)
         else:
             raise NotImplementedError(f"The perturbation {axis} is not supported")
 
@@ -325,6 +345,9 @@ def shrink(self, axis: PerturbationAxis = PerturbationAxis.LAYERS, amount: int =
             self._wire_mask.shrink(amount)
         elif axis == PerturbationAxis.RANDOM:
             self._parameter_mask.shrink(amount)
+        elif axis == PerturbationAxis.ENTANGLING:
+            if self._entangling_mask:
+                self._entangling_mask.shrink(amount)
         else:
             raise NotImplementedError(f"The perturbation {axis} is not supported")
 
@@ -333,6 +356,8 @@ def clear(self):
         self._layer_mask.clear()
         self._wire_mask.clear()
         self._parameter_mask.clear()
+        if self.entangling_mask is not None:
+            self._entangling_mask.clear()
 
     def apply_mask(self, values: np.ndarray):
         """
@@ -372,6 +397,10 @@ def copy(self: Self) -> Self:
         clone._wire_mask = self._wire_mask.copy(clone)
         clone.parameters = self.parameters.copy()
         clone.default_value = self.default_value
+        if self._entangling_mask is not None:
+            clone._entangling_mask = self._entangling_mask.copy(clone)
+        else:
+            clone._entangling_mask = None
         clone._dynamic_parameters = self._dynamic_parameters
         return clone
 
@@ -480,8 +509,15 @@ def __init__(
         layers: int,
         wires: int,
         default_value: Optional[float] = None,
+        entangling_mask: Optional[Mask] = None,
     ):
-        super().__init__(parameters, layers, wires, default_value=default_value)
+        super().__init__(
+            parameters,
+            layers,
+            wires,
+            default_value=default_value,
+            entangling_mask=entangling_mask,
+        )
         self._parameter_freeze_mask = Mask(shape=parameters.shape)
         self._layer_freeze_mask = Mask(shape=(layers,))
         self._wire_freeze_mask = Mask(shape=(wires,))

tests/test_masks.py

@@ -1,4 +1,5 @@
-from tests.utils import cost, create_freezable_circuit, variational_circuit, device
+from maskit.circuits import variational_circuit as original_variational_circuit
+from tests.utils import cost, create_freezable_circuit, device, variational_circuit
 import random
 import pytest
 
@@ -47,6 +48,13 @@ def test_init(self):
                 layers=size,
                 wires=size + 1,
             )
+        with pytest.raises(AssertionError):
+            MaskedCircuit(
+                parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
+                layers=size,
+                wires=size,
+                entangling_mask=Mask(shape=(size + 1, size)),
+            )
         mc = MaskedCircuit(
             parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
             layers=size,
@@ -56,40 +64,64 @@ def test_init(self):
         assert pnp.array_equal(mc.wire_mask, pnp.ones((size,), dtype=bool))
 
     def test_wrong_mode(self):
-        mp = self._create_circuit(3)
+        mp = self._create_circuit_with_entangling_gates(3)
         with pytest.raises(AssertionError):
             mp.perturb(axis=PerturbationAxis.LAYERS, mode=10)
 
     def test_wrong_axis(self):
-        mp = self._create_circuit(3)
+        mp = self._create_circuit_with_entangling_gates(3)
         with pytest.raises(NotImplementedError):
             mp.perturb(axis=10)
 
     @pytest.mark.parametrize("axis", list(PerturbationAxis))
     def test_clear(self, axis):
         size = 3
-        mp = self._create_circuit(size)
+        mp = self._create_circuit_with_entangling_gates(size)
         mp.perturb(axis=axis, amount=size)
         assert (
-            pnp.sum(mp.layer_mask) + pnp.sum(mp.wire_mask) + pnp.sum(mp.parameter_mask)
+            pnp.sum(mp.layer_mask)
+            + pnp.sum(mp.wire_mask)
+            + pnp.sum(mp.parameter_mask)
+            + pnp.sum(mp.entangling_mask)
             == size
         )
         mp.clear()
+        assert (
+            pnp.sum(mp.layer_mask)
+            + pnp.sum(mp.wire_mask)
+            + pnp.sum(mp.parameter_mask)
+            + pnp.sum(mp.entangling_mask)
+            == 0
+        )
+
+    def test_perturb_entangling(self):
+        size = 3
+        mp = self._create_circuit(size)
+        mp.perturb(axis=PerturbationAxis.ENTANGLING, amount=1)
+        # ensure nothing has happened as entangling mask is None
+        assert mp.entangling_mask is None
         assert (
             pnp.sum(mp.layer_mask) + pnp.sum(mp.wire_mask) + pnp.sum(mp.parameter_mask)
             == 0
         )
 
+        mp = self._create_circuit_with_entangling_gates(size)
+        mp.perturb(axis=PerturbationAxis.ENTANGLING, amount=1)
+        assert pnp.sum(mp.entangling_mask) == 1
+
     @pytest.mark.parametrize("axis", list(PerturbationAxis))
     def test_zero_amount(self, axis):
-        mp = self._create_circuit(3)
+        mp = self._create_circuit_with_entangling_gates(3)
         pre_sum = (
-            pnp.sum(mp.wire_mask) + pnp.sum(mp.layer_mask) + pnp.sum(mp.parameter_mask)
+            pnp.sum(mp.wire_mask)
+            + pnp.sum(mp.layer_mask)
+            + pnp.sum(mp.parameter_mask)
+            + pnp.sum(mp.entangling_mask)
         )
         mp.perturb(axis=axis, amount=0)
         assert pre_sum == pnp.sum(mp.wire_mask) + pnp.sum(mp.layer_mask) + pnp.sum(
             mp.parameter_mask
-        )
+        ) + pnp.sum(mp.entangling_mask)
 
     def test_apply_mask(self):
         size = 3
@@ -122,6 +154,15 @@ def test_copy(self):
         assert pnp.sum(new_mp.wire_mask) == 0
         assert pnp.sum(new_mp.layer_mask) == pnp.sum(mp.layer_mask)
         assert pnp.sum(new_mp.parameter_mask) == 0
+        assert new_mp.entangling_mask is None
+
+        # also test copying of existing entanglement mask
+        mp = self._create_circuit_with_entangling_gates(size)
+        assert mp.entangling_mask is not None
+        new_mp = mp.copy()
+        mp.entangling_mask[0, 0] = True
+        assert pnp.sum(mp.entangling_mask) == 1
+        assert pnp.sum(new_mp.entangling_mask) == 0
 
     def test_parameters(self):
         size = 3
@@ -155,6 +196,19 @@ def test_shrink_parameter(self):
         mp.shrink(amount=1, axis=PerturbationAxis.RANDOM)
         assert pnp.sum(mp.mask) == mp.mask.size - 1
 
+    def test_shrink_entangling(self):
+        size = 3
+        mp = self._create_circuit_with_entangling_gates(size)
+        mp.entangling_mask[:] = True
+        mp.shrink(amount=1, axis=PerturbationAxis.ENTANGLING)
+        assert pnp.sum(mp.entangling_mask) == mp.entangling_mask.size - 1
+
+        # also test in case no mask is set
+        mp = self._create_circuit(size)
+        mp.shrink(amount=1, axis=PerturbationAxis.ENTANGLING)
+        assert mp.entangling_mask is None
+        assert pnp.sum(mp.mask) == 0  # also ensure that nothing else was shrunk
+
     def test_shrink_wrong_axis(self):
         mp = self._create_circuit(3)
         with pytest.raises(NotImplementedError):
@@ -259,10 +313,34 @@ def test_dynamic_parameters(self):
         for i in range(size):
             assert circuit.parameters[0, i] != 1
 
+    def test_entangling_mask_application(self):
+        size = 4
+        mp = self._create_circuit_with_entangling_gates(size)
+        rotations = [pnp.random.choice([0, 1, 2]) for _ in range(size * size)]
+        circuit = qml.QNode(original_variational_circuit, device(mp.wire_mask.size))
+
+        circuit(mp.differentiable_parameters, rotations, mp)
+        assert circuit.specs["gate_types"]["CZ"] == 12
+
+        mp.perturb(
+            axis=PerturbationAxis.ENTANGLING, mode=PerturbationMode.ADD, amount=6
+        )
+        circuit(mp.differentiable_parameters, rotations, mp)
+        assert circuit.specs["gate_types"]["CZ"] == 6
+
     def _create_circuit(self, size):
         parameters = pnp.random.uniform(low=-pnp.pi, high=pnp.pi, size=(size, size))
         return MaskedCircuit(parameters=parameters, layers=size, wires=size)
 
+    def _create_circuit_with_entangling_gates(self, size):
+        parameters = pnp.random.uniform(low=-pnp.pi, high=pnp.pi, size=(size, size))
+        return MaskedCircuit(
+            parameters=parameters,
+            layers=size,
+            wires=size,
+            entangling_mask=Mask(shape=(size, size - 1)),
+        )
+
 
 class TestFreezableMaskedCircuit:
     def test_init(self):