Repository: cirKITers/masKIT

Issue:
Support dropping of entangling gates
So far we only support dropping parameterised gates. But especially with a focus on lowering expressivity of a circuit it might make sense to also enable dropping entangling gates. This could be realised by adding another mask to our `MaskedCircuit` class.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/32?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) Report
> Merging [#32](https://codecov.io/gh/cirKITers/masKIT/pull/32?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) (dc4fa46) into [main](https://codecov.io/gh/cirKITers/masKIT/commit/ccaff9bd4a5d2733d9f111ac0cdd78a7c9730ae6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) (ccaff9b) will **decrease** coverage by `0.10%`.
> The diff coverage is `72.41%`.

[![Impacted file tree graph](https://codecov.io/gh/cirKITers/masKIT/pull/32/graphs/tree.svg?width=650&height=150&src=pr&token=8Or2uHjmFE&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)](https://codecov.io/gh/cirKITers/masKIT/pull/32?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)

```diff
@@            Coverage Diff             @@
##             main      #32      +/-   ##
==========================================
- Coverage   74.54%   74.44%   -0.11%     
==========================================
  Files          15       15              
  Lines        1096     1119      +23     
  Branches      105      111       +6     
==========================================
+ Hits          817      833      +16     
- Misses        275      280       +5     
- Partials        4        6       +2     
```


| [Impacted Files](https://codecov.io/gh/cirKITers/masKIT/pull/32?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) | Coverage Δ | |
|---|---|---|
| [main.py](https://codecov.io/gh/cirKITers/masKIT/pull/32/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFpbi5weQ==) | `0.00% <0.00%> (ø)` | |
| [maskit/circuits.py](https://codecov.io/gh/cirKITers/masKIT/pull/32/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFza2l0L2NpcmN1aXRzLnB5) | `0.00% <0.00%> (ø)` | |
| [maskit/masks.py](https://codecov.io/gh/cirKITers/masKIT/pull/32/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFza2l0L21hc2tzLnB5) | `98.55% <80.00%> (-1.45%)` | :arrow_down: |
| [tests/test\_masks.py](https://codecov.io/gh/cirKITers/masKIT/pull/32/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-dGVzdHMvdGVzdF9tYXNrcy5weQ==) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/32?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/32?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers). Last update [0313002...dc4fa46](https://codecov.io/gh/cirKITers/masKIT/pull/32?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers).

eileen-kuehn: @maxfischer2781, when reviewing, please have a specific look on how the `entangling_mask` is currently handled. I opted for making it optional what requires to test if it is set in case actions are performed. This will probably fail for the `execute` function if the mask is not properly initialised.
Any ideas on this? Further, I am a bit afraid, that this mask might even have not a so well-defined shape... However, I would like to give it a try already. What do you think?
maxfischer2781: Sorry, I missed your previous comment. TBH I'm not totally sure how exactly *this* change works, so it's hard for me to say how it might affect other things.

Perhaps we should discuss at a literal drawing board. :P
eileen-kuehn: > Do you think you can hit the `continue`s with unittests? I can't really comment on whether they work as expected.

I added another unit test that should exactly test for that :) Now you will have to review again!

Patch:
diff --git a/main.py b/main.py
index 0521fd0..0532b73 100644
--- a/main.py
+++ b/main.py
@@ -4,7 +4,7 @@ import random
 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
diff --git a/maskit/circuits.py b/maskit/circuits.py
index ce72ed5..59d6935 100644
--- a/maskit/circuits.py
+++ b/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])
 
 
diff --git a/maskit/masks.py b/maskit/masks.py
index 45c0fb0..a4ba22c 100644
--- a/maskit/masks.py
+++ b/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 @@ class MaskedCircuit(object):
         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 @@ class MaskedCircuit(object):
         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 @@ class MaskedCircuit(object):
         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 @@ class MaskedCircuit(object):
         """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 @@ class MaskedCircuit(object):
             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 @@ class MaskedCircuit(object):
             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 @@ class MaskedCircuit(object):
         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 @@ class MaskedCircuit(object):
         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 @@ class FreezableMaskedCircuit(MaskedCircuit):
         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,))


Test patch:
diff --git a/tests/test_masks.py b/tests/test_masks.py
index 2f4b8fc..9c752aa 100644
--- a/tests/test_masks.py
+++ b/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 @@ class TestMaskedCircuits:
                 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 @@ class TestMaskedCircuits:
         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 @@ class TestMaskedCircuits:
         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 @@ class TestMaskedCircuits:
         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 @@ class TestMaskedCircuits:
         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):