Repository: cirKITers/masKIT

Issue:
Normalise predictions for cross entropy calculation
As discussed while implementing cross entropy in #38, we should add a normalisation of the predictions that probabilities add up to `1.0`:

> Adding a normalisation should be our best option.

_Originally posted by @cDenius in https://github.com/cirKITers/masKIT/pull/45#r705040666_

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/38?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) Report
> Merging [#38](https://codecov.io/gh/cirKITers/masKIT/pull/38?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) (d6b35e9) 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 **increase** coverage by `1.88%`.
> The diff coverage is `0.00%`.

[![Impacted file tree graph](https://codecov.io/gh/cirKITers/masKIT/pull/38/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/38?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)

```diff
@@            Coverage Diff             @@
##             main      #38      +/-   ##
==========================================
+ Coverage   74.54%   76.43%   +1.88%     
==========================================
  Files          15       15              
  Lines        1096     1379     +283     
  Branches      105      147      +42     
==========================================
+ Hits          817     1054     +237     
- Misses        275      321      +46     
  Partials        4        4              
```


| [Impacted Files](https://codecov.io/gh/cirKITers/masKIT/pull/38?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) | Coverage Δ | |
|---|---|---|
| [maskit/utils.py](https://codecov.io/gh/cirKITers/masKIT/pull/38/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFza2l0L3V0aWxzLnB5) | `0.00% <0.00%> (ø)` | |
| [maskit/masks.py](https://codecov.io/gh/cirKITers/masKIT/pull/38/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFza2l0L21hc2tzLnB5) | `100.00% <0.00%> (ø)` | |
| [tests/test\_masks.py](https://codecov.io/gh/cirKITers/masKIT/pull/38/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% <0.00%> (ø)` | |
| [main.py](https://codecov.io/gh/cirKITers/masKIT/pull/38/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFpbi5weQ==) | `0.69% <0.00%> (+0.69%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/38?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/38?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers). Last update [0313002...d6b35e9](https://codecov.io/gh/cirKITers/masKIT/pull/38?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).


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,))
diff --git a/maskit/utils.py b/maskit/utils.py
index 4097ee8..f5045d7 100644
--- a/maskit/utils.py
+++ b/maskit/utils.py
@@ -8,15 +8,38 @@ def check_params(train_params):
     assert isinstance(train_params["optimizer"], ExtendedOptimizers)
 
 
-def cross_entropy(predictions, targets, epsilon=1e-12):
+def cross_entropy(
+    predictions: np.ndarray, targets: np.ndarray, epsilon: float = 1e-15
+) -> float:
     """
-    Computes cross entropy between targets (encoded as one-hot vectors)
-    and predictions.
-    Input: predictions (N, k) ndarray
-           targets (N, k) ndarray
-    Returns: scalar
+    Cross entropy calculation between :py:attr:`targets` (encoded as one-hot vectors)
+    and :py:attr:`predictions`. Predictions are normalized to sum up to `1.0`.
+
+    .. note::
+
+        The implementation of this function is based on the discussion on
+        `StackOverflow <https://stackoverflow.com/a/47398312/10138546>`_.
+
+        Due to ArrayBoxes that are required for automatic differentiation, we currently
+        use this implementation instead of implementations provided by sklearn for
+        example.
+
+    :param predictions: Predictions in same order as targets. In case predictions for
+        several samples are given, the weighted cross entropy is returned.
+    :param targets: Ground truth labels for supplied samples.
+    :param epsilon: Amount to clip predictions as log is not defined for `0` and `1`.
     """
+    assert (
+        predictions.shape == targets.shape
+    ), f"Shape of predictions {predictions.shape} must match targets {targets.shape}"
+    current_sum = np.sum(predictions, axis=predictions.ndim - 1)
+
+    if predictions.ndim == 1:
+        sample_count = 1
+        predictions = predictions / current_sum
+    else:
+        sample_count = predictions.shape[0]
+        predictions = predictions / current_sum[:, np.newaxis]
+
     predictions = np.clip(predictions, epsilon, 1.0 - epsilon)
-    N = predictions.shape[0]
-    ce = -np.sum(targets * np.log(predictions + 1e-9)) / N
-    return ce
+    return -np.sum(targets * np.log(predictions)) / sample_count
diff --git a/pyproject.toml b/pyproject.toml
index 1b2db15..cea3441 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,24 +1,35 @@
 [build-system]
-requires = ["flit_core >=2,<4"]
+requires = ["flit_core >=3.2,<4"]
 build-backend = "flit_core.buildapi"
 
-[tool.flit.metadata]
-module = "maskit"
-author = "Eileen Kuehn, Christof Wendenius"
-author-email = "eileen.kuehn@kit.edu"
-home-page = "https://github.com/cirKITers/masKIT"
-classifiers = [ "License :: OSI Approved :: MIT License",]
-description-file = "README.md"
-requires = [
+[project]
+name = "maskit"
+authors = [
+    {name = "Eileen Kuehn", email = "eileen.kuehn@kit.edu"},
+    {name = "Christof Wendenius", email = "christof.wendenius@kit.edu"},
+]
+readme = "README.md"
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+]
+requires-python = ">=3.7"
+dynamic = ["version", "description"]
+license = {file = "LICENSE"}
+keywords = ["quantum computing",]
+dependencies = [
     "pennylane ~= 0.16.0",
     "typing_extensions",
 ]
 
-[tool.flit.metadata.requires-extra]
+[project.optional-dependencies]
 test = [
     "pytest >=4.3.0",
     "flake8",
     "flake8-bugbear",
     "black",
+    "scikit-learn",
 ]
 dev = ["pre-commit"]
+
+[project.urls]
+Source = "https://github.com/cirKITers/masKIT"


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):
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..7f97e73
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,30 @@
+import numpy as np
+import pytest
+from sklearn.metrics import log_loss
+
+from maskit.utils import cross_entropy
+
+
+@pytest.mark.parametrize(
+    "predictions,targets",
+    [
+        pytest.param(
+            np.array([[0.25, 0.25, 0.25, 0.25], [0.01, 0.01, 0.01, 0.96]]),
+            np.array([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]]),
+        ),
+        pytest.param(np.array([0.0, 0.0, 1.0, 0.0]), np.array([0.0, 0.0, 1.0, 0.0])),
+        pytest.param(
+            np.array([0.25, 0.25, 0.25, 0.25]), np.array([0.0, 0.0, 1.0, 0.0])
+        ),
+        pytest.param(
+            np.array([0.01, 0.01, 0.01, 0.96]), np.array([0.0, 0.0, 0.0, 1.0])
+        ),
+    ],
+)
+def test_cross_entropy(predictions, targets):
+    if predictions.ndim == 1:
+        expected_value = log_loss([targets], [predictions])
+    else:
+        expected_value = log_loss(targets, predictions)
+    calculated_value = cross_entropy(predictions=predictions, targets=targets)
+    assert np.isclose(expected_value, calculated_value)