Repository: PennyLaneAI/pennylane
Pull request #1144 (2021-03-17 18:19:40)

Mottonen state prep edit
**Context:**
The ``MottonenStatePreparation`` template uses a cascade of RY rotations and RZ rotations to bring a qubit register to a desired state. When the desired state vector only has non-negative, real coefficients, the RZ rotations are not necessary. In this case, applying the cascade of RZ rotations adds unnecessary CNOT gates at the end of the circuit.

**Description of the Change:**
We make the application of the RZ cascade conditional. If the planned RZ rotations are all near 0, we do not apply the RZ cascade.

**Benefits:**
Reduced circuit depth for state vectors that only have non-negative, real elements (e.g. stochastic probability distributions)

**Possible Drawbacks:**
Must test for incorrect phases
Must test for differentiability

**Related GitHub Issues:**


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144?src=pr&el=h1) Report
> Merging [#1144](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144?src=pr&el=desc) (68a76ca) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/34954b560a95125c0852886346ca8c7915142e11?el=desc) (34954b5) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1144      +/-   ##
==========================================
+ Coverage   97.97%   97.99%   +0.01%     
==========================================
  Files         142      142              
  Lines       10289    10377      +88     
==========================================
+ Hits        10081    10169      +88     
  Misses        208      208              
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pennylane/templates/state\_preparations/mottonen.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdGF0ZV9wcmVwYXJhdGlvbnMvbW90dG9uZW4ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/devices/tf\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvdGZfb3BzLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/devices/autograd\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvYXV0b2dyYWRfb3BzLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/devices/default\_mixed.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9taXhlZC5weQ==) | `100.00% <0.00%> (ø)` | |
| [pennylane/devices/default\_qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `100.00% <0.00%> (ø)` | |
| [pennylane/devices/default\_qubit\_tf.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF90Zi5weQ==) | `90.66% <0.00%> (ø)` | |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `94.02% <0.00%> (ø)` | |
| [pennylane/devices/default\_qubit\_autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9hdXRvZ3JhZC5weQ==) | `100.00% <0.00%> (ø)` | |
| [pennylane/ops/qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL29wcy9xdWJpdC5weQ==) | `98.02% <0.00%> (+0.11%)` | :arrow_up: |
| [pennylane/devices/jax\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvamF4X29wcy5weQ==) | `96.36% <0.00%> (+1.24%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144?src=pr&el=footer). Last update [34954b5...68a76ca](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1144?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

trbromley: Minor comment - but don't forget to update the PR when it's ready for review:
- If draft PR, can mark as "ready for review"
- Remove any WIP tags
- Remove WIP from title
- Remove the "Before submitting..." in the PR main comment

Not a worry for this PR as it's already ready to go :+1: 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index ee66d45f4..a4eb07b99 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -397,6 +397,32 @@
 
 <h3>Improvements</h3>
 
+* Edited the ``MottonenStatePreparation`` template to improve performance on states with only real amplitudes
+  by reducing the number of redundant CNOT gates at the end of a circuit.
+
+  ```python
+  dev = qml.device("default.qubit", wires=2)
+  
+  inputstate = [np.sqrt(0.2), np.sqrt(0.3), np.sqrt(0.4), np.sqrt(0.1)]
+  
+  @qml.qnode(dev)
+  def circuit():
+    mottonen.MottonenStatePreparation(inputstate,wires=[0, 1])
+    return qml.expval(qml.PauliZ(0))
+  ```
+  Previously returned:
+  ```pycon
+  >>> print(qml.draw(circuit)())
+  0: ──RY(1.57)──╭C─────────────╭C──╭C──╭C──┤ ⟨Z⟩ 
+  1: ──RY(1.35)──╰X──RY(0.422)──╰X──╰X──╰X──┤   
+  ```
+  Now returns:
+  ```pycon
+  >>> print(qml.draw(circuit)())
+  0: ──RY(1.57)──╭C─────────────╭C──┤ ⟨Z⟩ 
+  1: ──RY(1.35)──╰X──RY(0.422)──╰X──┤   
+  ```
+
 - The `QAOAEmbedding` and `BasicEntanglerLayers` are now classes inheriting 
   from `Operation`, and define the ansatz in their `expand()` method. This 
   change does not affect the user interface. 
diff --git a/pennylane/templates/state_preparations/mottonen.py b/pennylane/templates/state_preparations/mottonen.py
index d61aee72b..17b43c3a4 100644
--- a/pennylane/templates/state_preparations/mottonen.py
+++ b/pennylane/templates/state_preparations/mottonen.py
@@ -293,10 +293,11 @@ def MottonenStatePreparation(state_vector, wires):
         target = wires_reverse[k - 1]
         _uniform_rotation_dagger(qml.RY, alpha_y_k, control, target)
 
-    # Apply inverse z rotation cascade to prepare correct phases of amplitudes
-    for k in range(len(wires_reverse), 0, -1):
-        alpha_z_k = _get_alpha_z(omega, len(wires_reverse), k)
-        control = wires_reverse[k:]
-        target = wires_reverse[k - 1]
-        if len(alpha_z_k) > 0:
-            _uniform_rotation_dagger(qml.RZ, alpha_z_k, control, target)
+    # If necessary, apply inverse z rotation cascade to prepare correct phases of amplitudes
+    if not qml.math.allclose(omega, 0):
+        for k in range(len(wires_reverse), 0, -1):
+            alpha_z_k = _get_alpha_z(omega, len(wires_reverse), k)
+            control = wires_reverse[k:]
+            target = wires_reverse[k - 1]
+            if len(alpha_z_k) > 0:
+                _uniform_rotation_dagger(qml.RZ, alpha_z_k, control, target)


Test patch:
diff --git a/tests/templates/test_state_preparations.py b/tests/templates/test_state_preparations.py
index e8ef54d96..fe5a8dc54 100644
--- a/tests/templates/test_state_preparations.py
+++ b/tests/templates/test_state_preparations.py
@@ -19,14 +19,18 @@ Integration tests should be placed into ``test_templates.py``.
 
 import math
 from unittest.mock import patch
-import numpy as np
+from pennylane import numpy as np
 import pytest
 import pennylane as qml
-from pennylane.templates.state_preparations import (BasisStatePreparation,
-                                                    MottonenStatePreparation,
-                                                    ArbitraryStatePreparation)
+from pennylane.templates.state_preparations import (
+    BasisStatePreparation,
+    MottonenStatePreparation,
+    ArbitraryStatePreparation,
+)
 from pennylane.templates.state_preparations.mottonen import gray_code
-from pennylane.templates.state_preparations.arbitrary_state_preparation import _state_preparation_pauli_words
+from pennylane.templates.state_preparations.arbitrary_state_preparation import (
+    _state_preparation_pauli_words,
+)
 from pennylane.templates.state_preparations.mottonen import _get_alpha_y
 from pennylane.wires import Wires
 
@@ -47,11 +51,32 @@ class TestHelperFunctions:
 
         assert gray_code(rank) == expected_gray_code
 
-    @pytest.mark.parametrize("num_wires,expected_pauli_words", [
-        (1, ["X", "Y"]),
-        (2, ["XI", "YI", "IX", "IY", "XX", "XY"]),
-        (3, ["XII", "YII", "IXI", "IYI", "IIX", "IIY", "IXX", "IXY", "XXI", "XYI", "XIX", "XIY", "XXX", "XXY"]),
-    ])
+    @pytest.mark.parametrize(
+        "num_wires,expected_pauli_words",
+        [
+            (1, ["X", "Y"]),
+            (2, ["XI", "YI", "IX", "IY", "XX", "XY"]),
+            (
+                3,
+                [
+                    "XII",
+                    "YII",
+                    "IXI",
+                    "IYI",
+                    "IIX",
+                    "IIY",
+                    "IXX",
+                    "IXY",
+                    "XXI",
+                    "XYI",
+                    "XIX",
+                    "XIY",
+                    "XXX",
+                    "XXY",
+                ],
+            ),
+        ],
+    )
     def test_state_preparation_pauli_words(self, num_wires, expected_pauli_words):
         """Test that the correct Pauli words are returned."""
         for idx, pauli_word in enumerate(_state_preparation_pauli_words(num_wires)):
@@ -214,7 +239,9 @@ class TestMottonenStatePreparation:
         ([1/2, 0, 1j/2, 1j/math.sqrt(2)], [0, 1], [1/2, 0, 0, 0, 1j/2, 0, 1j/math.sqrt(2), 0]),
     ])
     # fmt: on
-    def test_state_preparation_fidelity(self, tol, qubit_device_3_wires, state_vector, wires, target_state):
+    def test_state_preparation_fidelity(
+        self, tol, qubit_device_3_wires, state_vector, wires, target_state
+    ):
         """Tests that the template MottonenStatePreparation integrates correctly with PennyLane
         and produces states with correct fidelity."""
 
@@ -227,7 +254,7 @@ class TestMottonenStatePreparation:
         circuit()
 
         state = circuit.device.state.ravel()
-        fidelity = abs(np.vdot(state, target_state))**2
+        fidelity = abs(np.vdot(state, target_state)) ** 2
 
         # We test for fidelity here, because the vector themselves will hardly match
         # due to imperfect state preparation
@@ -278,7 +305,9 @@ class TestMottonenStatePreparation:
         ([1/2, 0, 1j/2, 1j/math.sqrt(2)], [0, 1], [1/2, 0, 0, 0, 1j/2, 0, 1j/math.sqrt(2), 0]),
     ])
     # fmt: on
-    def test_state_preparation_probability_distribution(self, tol, qubit_device_3_wires, state_vector, wires, target_state):
+    def test_state_preparation_probability_distribution(
+        self, tol, qubit_device_3_wires, state_vector, wires, target_state
+    ):
         """Tests that the template MottonenStatePreparation integrates correctly with PennyLane
         and produces states with correct probability distribution."""
 
@@ -292,8 +321,8 @@ class TestMottonenStatePreparation:
 
         state = circuit.device.state.ravel()
 
-        probabilities = np.abs(state)**2
-        target_probabilities = np.abs(target_state)**2
+        probabilities = np.abs(state) ** 2
+        target_probabilities = np.abs(target_state) ** 2
 
         assert np.allclose(probabilities, target_probabilities, atol=tol, rtol=0)
 
@@ -324,11 +353,14 @@ class TestMottonenStatePreparation:
         with pytest.raises(ValueError, match="State vector must be of (length|shape)"):
             MottonenStatePreparation(state_vector, wires)
 
-    @pytest.mark.parametrize("current_qubit, expected", [
-        (1, np.array([0, 0, 0, 1.23095942])),
-        (2, np.array([2.01370737, 3.14159265])),
-        (3, np.array([1.15927948])),
-    ])
+    @pytest.mark.parametrize(
+        "current_qubit, expected",
+        [
+            (1, np.array([0, 0, 0, 1.23095942])),
+            (2, np.array([2.01370737, 3.14159265])),
+            (3, np.array([1.15927948])),
+        ],
+    )
     def test_get_alpha_y(self, current_qubit, expected, tol):
         """Test the _get_alpha_y helper function."""
 
@@ -359,6 +391,52 @@ class TestMottonenStatePreparation:
             state_vector = np.array([0, 2, 0, 0])
             circuit(state_vector)
 
+    # fmt: off
+    @pytest.mark.parametrize("state_vector, n_wires", [
+        ([1/2, 1/2, 1/2, 1/2], 2),
+        ([1, 0, 0, 0], 2),
+        ([0, 1, 0, 0], 2),
+        ([0, 0, 0, 1], 2),
+        ([0, 1, 0, 0, 0, 0, 0, 0], 3),
+        ([0, 0, 0, 0, 1, 0, 0, 0], 3),
+        ([2/3, 0, 0, 0, 1/3, 0, 0, 2/3], 3),
+        ([1/2, 0, 0, 0, 1/2, 1/2, 1/2, 0], 3),
+        ([1/3, 0, 0, 0, 2/3, 2/3, 0, 0], 3),
+        ([2/3, 0, 0, 0, 1/3, 0, 0, 2/3], 3),
+    ])
+    # fmt: on
+    def test_RZ_skipped(self, state_vector, n_wires):
+        """Tests whether the cascade of RZ gates is skipped for real-valued states"""
+
+        n_CNOT = 2 ** n_wires - 2
+
+        dev = qml.device("default.qubit", wires=n_wires)
+
+        @qml.qnode(dev)
+        def circuit(state_vector):
+            MottonenStatePreparation(state_vector, wires=range(n_wires))
+            return qml.expval(qml.PauliX(wires=0))
+
+        # when the RZ cascade is skipped, CNOT gates should only be those required for RY cascade
+        circuit(state_vector)
+
+        assert circuit.qtape.get_resources()["CNOT"] == n_CNOT
+
+    @pytest.mark.parametrize(
+        "state_vector", [np.array([0.70710678, 0.70710678]), np.array([0.70710678, 0.70710678j])]
+    )
+    def test_gradient_evaluated(self, state_vector):
+        """Test that the gradient is successfully calculated for a simple example. This test only
+        checks that the gradient is calculated without an error."""
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(state_vector):
+            MottonenStatePreparation(state_vector, wires=range(1))
+            return qml.expval(qml.PauliZ(0))
+
+        qml.grad(circuit)(state_vector)
+
 
 class TestArbitraryStatePreparation:
     """Test the ArbitraryStatePreparation template."""
@@ -421,7 +499,7 @@ class TestArbitraryStatePreparation:
 
     def test_GHZ_generation(self, qubit_device_3_wires, tol):
         """Test that the template prepares a GHZ state."""
-        GHZ_state = np.array([1/math.sqrt(2), 0, 0, 0, 0, 0, 0, 1/math.sqrt(2)])
+        GHZ_state = np.array([1 / math.sqrt(2), 0, 0, 0, 0, 0, 0, 1 / math.sqrt(2)])
 
         weights = np.zeros(14)
         weights[13] = math.pi / 2
@@ -438,7 +516,7 @@ class TestArbitraryStatePreparation:
 
     def test_even_superposition_generation(self, qubit_device_3_wires, tol):
         """Test that the template prepares a even superposition state."""
-        even_superposition_state = np.ones(8)/math.sqrt(8)
+        even_superposition_state = np.ones(8) / math.sqrt(8)
 
         weights = np.zeros(14)
         weights[1] = math.pi / 2