Recreate cosine window PR (#4704)

### Before submitting

Please complete the following checklist when submitting a PR:

- [ ] All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to
the
      test directory!

- [ ] All new functions and code must be clearly commented and
documented.
If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [ ] Ensure that the test suite passes, by running `make test`.

- [ ] Add a new entry to the `doc/releases/changelog-dev.md` file,
summarizing the
      change, and including a link back to the PR.

- [ ] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.


------------------------------------------------------------------------------------------------------------

**Context:** This PR is a recreation of the work "Cosine Window
StatePrep"

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

---------

Co-authored-by: Guillermo Alonso-Linaje <[email protected]>
Co-authored-by: Jay Soni <[email protected]>
Co-authored-by: soranjh <[email protected]>

doc/introduction/templates.rst

@@ -114,6 +114,10 @@ state preparation is typically used as the first operation.
     :description: :doc:`ArbitraryStatePreparation <../code/api/pennylane.ArbitraryStatePreparation>`
     :figure: _static/templates/subroutines/arbitrarystateprep.png
 
+.. gallery-item::
+    :description: :doc:`CosineWindow <../code/api/pennylane.CosineWindow>`
+    :figure: _static/templates/state_preparations/thumbnail_cosine_window.png
+
 .. raw:: html
 
         <div style='clear:both'></div>

doc/releases/changelog-dev.md

@@ -148,6 +148,54 @@
   [(#4620)](https://github.com/PennyLaneAI/pennylane/pull/4620)
   [(#4632)](https://github.com/PennyLaneAI/pennylane/pull/4632)
 
+* The `CosineWindow` template has been added to prepare an initial state based on a cosine wave function.
+  [(#4683)](https://github.com/PennyLaneAI/pennylane/pull/4683)
+
+  ```python
+  import pennylane as qml
+  import matplotlib.pyplot as plt
+
+  dev = qml.device('default.qubit', wires=4)
+
+  @qml.qnode(dev)
+  def example_circuit():
+        qml.CosineWindow(wires=range(4))
+        return qml.state()
+  output = example_circuit()
+
+  # Graph showing state amplitudes
+  plt.bar(range(len(output)), output)
+  plt.show()
+  ```
+
+  We can show how this operator is built:
+
+  ```python
+  import pennylane as qml
+
+  dev = qml.device("default.qubit", wires=5)
+
+  op = qml.CosineWindow(wires=range(5))
+
+  @qml.qnode(dev)
+  def circuit():
+      op.decomposition()
+      return qml.state()
+
+  print(qml.draw(circuit)())
+
+  ```
+
+  ```pycon
+
+  0: ──────────────╭QFT†──Rϕ(1.57)─┤  State
+  1: ──────────────├QFT†──Rϕ(0.79)─┤  State
+  2: ──────────────├QFT†──Rϕ(0.39)─┤  State
+  3: ──────────────├QFT†──Rϕ(0.20)─┤  State
+  4: ──H──RZ(3.14)─╰QFT†──Rϕ(0.10)─┤  State
+
+  ```  
+
 <h3>Improvements 🛠</h3>
 
 * `pennylane.devices.preprocess` now offers the transforms `decompose`, `validate_observables`, `validate_measurements`,
@@ -496,6 +544,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Guillermo Alonso,
 Utkarsh Azad,
 Jack Brown,
 Stepan Fomichev,

pennylane/templates/state_preparations/__init__.py

@@ -20,3 +20,4 @@
 from .basis import BasisStatePreparation
 from .arbitrary_state_preparation import ArbitraryStatePreparation
 from .basis_qutrit import QutritBasisStatePreparation
+from .cosine_window import CosineWindow

pennylane/templates/state_preparations/cosine_window.py

@@ -0,0 +1,129 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+r"""
+Contains the CosineWindow template.
+"""
+import numpy as np
+import pennylane as qml
+from pennylane.operation import StatePrepBase
+from pennylane import math
+from pennylane.wires import Wires, WireError
+
+
+class CosineWindow(StatePrepBase):
+    r"""CosineWindow(wires)
+    Prepares an initial state with a cosine wave function.
+
+    The wave function is defined below where :math:`m` is the number of wires.
+
+    .. math::
+
+        |\psi\rangle = \sqrt{2^{1-m}} \sum_{k=0}^{2^m-1} \cos(\frac{\pi k}{2^m} - \frac{\pi}{2}) |k\rangle,
+
+    .. figure:: ../../_static/templates/state_preparations/cosine_window.png
+        :align: center
+        :width: 65%
+        :target: javascript:void(0);
+
+    .. note::
+
+        The wave function is shifted by :math:`\frac{\pi}{2}` units so that the window is centered.
+
+    For more details see `Phys. Rev. D 106 (2022) <https://journals.aps.org/prd/abstract/10.1103/PhysRevD.106.034503>`_.
+
+    .. seealso:: :class:`~.QuantumPhaseEstimation` and :class:`~.QFT`.
+
+    Args:
+        wires (Sequence[int] or int): the wire(s) the operation acts on
+
+    **Example**
+
+    >>> dev = qml.device('default.qubit', wires=2)
+    >>> @qml.qnode(dev)
+    ... def example_circuit():
+    ...     qml.CosineWindow(wires=range(2))
+    ...     return qml.probs()
+    >>> print(example_circuit())
+    [1.87469973e-33 2.50000000e-01 5.00000000e-01 2.50000000e-01]
+    """
+
+    @staticmethod
+    def compute_decomposition(wires):  # pylint: disable=arguments-differ,unused-argument
+        r"""Representation of the operator as a product of other operators (static method).
+        It is efficiently decomposed from one QFT over all qubits and one-qubit rotation gates.
+
+        Args:
+            wires (Iterable, Wires): the wire(s) the operation acts on
+
+        Returns:
+            list[Operator]: decomposition into lower level operations
+        """
+
+        decomp_ops = []
+
+        decomp_ops.append(qml.Hadamard(wires=wires[-1]))
+        decomp_ops.append(qml.RZ(np.pi, wires=wires[-1]))
+        decomp_ops.append(qml.adjoint(qml.QFT)(wires=wires))
+
+        for ind, wire in enumerate(wires):
+            decomp_ops.append(qml.PhaseShift(np.pi * 2 ** (-ind - 1), wires=wire))
+
+        return decomp_ops
+
+    def label(self, decimals=None, base_label=None, cache=None):
+        return "CosineWindow"
+
+    def state_vector(self, wire_order=None):  # pylint: disable=arguments-differ,unused-argument
+        r"""Calculation of the state vector generated by the cosine window.
+
+        Args:
+            wire_order (Iterable, Wires): Custom order of wires for the returned state vector.
+
+        Raises:
+            WireError: Custom wire_order must contain all wires.
+
+        Returns:
+            TensorLike[complex]: output state
+        """
+
+        num_op_wires = len(self.wires)
+        op_vector_shape = (2,) * num_op_wires
+        vector = np.array(
+            [
+                np.sqrt(2 ** (1 - num_op_wires))
+                * np.cos(-np.pi / 2 + np.pi * x / 2**num_op_wires)
+                for x in range(2**num_op_wires)
+            ]
+        )
+        op_vector = math.reshape(vector, op_vector_shape)
+
+        if wire_order is None or Wires(wire_order) == self.wires:
+            return op_vector
+
+        wire_order = Wires(wire_order)
+        if not wire_order.contains_wires(self.wires):
+            raise WireError(f"Custom wire_order must contain all {self.name} wires")
+
+        indices = tuple([Ellipsis] + [slice(None)] * num_op_wires)
+
+        ket_shape = [2] * num_op_wires
+        ket = np.zeros(ket_shape, dtype=np.complex128)
+        ket[indices] = op_vector
+
+        if self.wires != wire_order[:num_op_wires]:
+            current_order = self.wires + list(Wires.unique_wires([wire_order, self.wires]))
+            desired_order = [current_order.index(w) for w in wire_order]
+            ket = ket.transpose(desired_order)
+
+        return math.convert_like(ket, op_vector)

tests/templates/test_state_preparations/test_cosine_window.py

@@ -0,0 +1,119 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the CosineWindow template.
+"""
+# pylint: disable=too-few-public-methods
+import pytest
+import numpy as np
+import pennylane as qml
+from pennylane.wires import WireError
+
+
+class TestDecomposition:
+    """Tests that the template defines the correct decomposition."""
+
+    def test_correct_gates_single_wire(self):
+        """Test that the correct gates are applied."""
+
+        op = qml.CosineWindow(wires=[0])
+        queue = op.expand().operations
+
+        assert queue[0].name == "Hadamard"
+        assert queue[1].name == "RZ"
+        assert queue[2].name == "Adjoint(QFT)"
+        assert queue[3].name == "PhaseShift"
+
+        assert np.isclose(queue[3].data[0], np.pi / 2)
+
+    def test_correct_gates_many_wires(self):
+        """Test that the correct gates are applied on two wires."""
+
+        op = qml.CosineWindow(wires=[0, 1, 2, 3, 4])
+        queue = op.decomposition()
+
+        assert queue[0].name == "Hadamard"
+        assert queue[1].name == "RZ"
+        assert queue[2].name == "Adjoint(QFT)"
+
+        for ind, q in enumerate(queue[3:]):
+            assert q.name == "PhaseShift"
+            assert np.isclose(q.data[0], np.pi / 2 ** (ind + 1))
+
+    def test_custom_wire_labels(self):
+        """Test that template can deal with non-numeric, nonconsecutive wire labels."""
+
+        dev = qml.device("default.qubit", wires=3)
+        dev2 = qml.device("default.qubit", wires=["z", "a", "k"])
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.CosineWindow(wires=range(3))
+            return qml.expval(qml.Identity(0)), qml.state()
+
+        @qml.qnode(dev2)
+        def circuit2():
+            qml.CosineWindow(wires=["z", "a", "k"])
+            return qml.expval(qml.Identity("z")), qml.state()
+
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
+
+        assert np.allclose(res1, res2)
+        assert np.allclose(state1, state2)
+
+
+class TestRepresentation:
+    """Test id and label."""
+
+    def test_id(self):
+        """Tests that the id attribute can be set."""
+        wires = [0, 1, 2]
+        template = qml.CosineWindow(wires=wires, id="a")
+        assert template.id == "a"
+        assert template.wires == qml.wires.Wires(wires)
+
+    def test_label(self):
+        """Test label method returns CosineWindow"""
+        op = qml.CosineWindow(wires=[0, 1])
+        assert op.label() == "CosineWindow"
+
+
+class TestStateVector:
+    """Test the state_vector() method of various CosineWindow operations."""
+
+    def test_CosineWindow_state_vector(self):
+        """Tests that the state vector is correct for a single wire."""
+        op = qml.CosineWindow(wires=[0])
+        res = op.state_vector()
+        expected = np.array([0.0, 1.0])
+        assert np.allclose(res, expected)
+
+        op = qml.CosineWindow(wires=[0, 1])
+        res = np.reshape(op.state_vector() ** 2, (-1,))
+        expected = np.array([0.0, 0.25, 0.5, 0.25])
+        assert np.allclose(res, expected)
+
+    def test_CosineWindow_state_vector_bad_wire_order(self):
+        """Tests that the provided wire_order must contain the wires in the operation."""
+        qsv_op = qml.CosineWindow(wires=[0, 1])
+        with pytest.raises(WireError, match="wire_order must contain all CosineWindow wires"):
+            qsv_op.state_vector(wire_order=[1, 2])
+
+    def test_CosineWindow_state_vector_wire_order(self):
+        """Tests that the state vector works with a different order of wires."""
+        op = qml.CosineWindow(wires=[0, 1])
+        res = np.reshape(op.state_vector(wire_order=[1, 0]) ** 2, (-1,))
+        expected = np.array([0.0, 0.5, 0.25, 0.25])
+        assert np.allclose(res, expected)