Add Qudit phase gate, Z_d

unitary/alpha/qudit_gates.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
+
+from typing import List, Dict, Optional, Tuple
+
 import numpy as np
 import cirq
 
@@ -55,6 +59,68 @@ def _circuit_diagram_info_(self, args):
         return f"X({self.source_state}_{self.destination_state})"
 
 
+class QuditRzGate(cirq.EigenGate):
+    """Phase shifts a single state basis of the qudit.
+
+    A generalization of the phase shift gate to qudits.
+    https://en.wikipedia.org/wiki/Quantum_logic_gate#Phase_shift_gates
+
+    Implements Z_d as defined in eqn (5) of https://arxiv.org/abs/2008.00959
+    with the addition of a state parameter for convenience.
+    For a qudit of dimensionality d, shifts the phase of |phased_state> by radians.
+
+    Args:
+        dimension: Dimension of the qudits: for instance, a dimension of 3
+          would be a qutrit.
+        radians: The phase shift applied to the |phased_state>, measured in
+          radians.
+        phased_state: Optional index of the state to be phase shifted. Defaults
+          to phase shifting the state |dimension-1>.
+    """
+
+    _cached_eigencomponents: Dict[int, List[Tuple[float, np.ndarray]]] = {}
+
+    def __init__(
+        self, dimension: int, radians: float = np.pi, phased_state: Optional[int] = None
+    ):
+        super().__init__(exponent=radians / np.pi, global_shift=0)
+        self.dimension = dimension
+        if phased_state is not None:
+            if phased_state >= dimension or phased_state < 0:
+                raise ValueError(
+                    f"state {phased_state} is not valid for a qudit of"
+                    f" dimension {dimension}."
+                )
+            self.phased_state = phased_state
+        else:
+            self.phased_state = self.dimension - 1
+
+    def _qid_shape_(self):
+        return (self.dimension,)
+
+    def _eigen_components(self) -> List[Tuple[float, np.ndarray]]:
+        eigen_key = (self.dimension, self.phased_state)
+        if eigen_key not in QuditRzGate._cached_eigencomponents:
+            components = []
+            for i in range(self.dimension):
+                half_turns = 0
+                m = np.zeros((self.dimension, self.dimension))
+                m[i][i] = 1
+                if i == self.phased_state:
+                    half_turns = 1
+                components.append((half_turns, m))
+            QuditRzGate._cached_eigencomponents[eigen_key] = components
+        return QuditRzGate._cached_eigencomponents[eigen_key]
+
+    def _circuit_diagram_info_(self, args):
+        return cirq.CircuitDiagramInfo(
+            wire_symbols=("Z_d"), exponent=self._format_exponent_as_angle(args)
+        )
+
+    def _with_exponent(self, exponent: float) -> "QuditRzGate":
+        return QuditRzGate(rads=exponent * np.pi)
+
+
 class QuditPlusGate(cirq.Gate):
     """Cycles all the states by `addend` using a permutation gate.
     This gate adds a number to each state. For instance,`QuditPlusGate(dimension=3, addend=1)`

unitary/alpha/qudit_gates_test.py

@@ -252,6 +252,88 @@ def test_iswap(q0: int, q1: int):
     assert np.all(results.measurements["m1"] == q0)
 
 
+@pytest.mark.parametrize("dimension, phase_rads", [(2, np.pi), (3, 1), (4, np.pi * 2)])
+def test_rz_unitary(dimension: float, phase_rads: float):
+    rz = qudit_gates.QuditRzGate(dimension=dimension, radians=phase_rads)
+    expected_unitary = np.identity(n=dimension, dtype=np.complex64)
+
+    # 1j = e ^ ( j * ( pi / 2 )), so we multiply phase_rads by 2 / pi.
+    expected_unitary[dimension - 1][dimension - 1] = 1j ** (phase_rads * 2 / np.pi)
+
+    assert np.isclose(phase_rads / np.pi, rz._exponent)
+    rz_unitary = cirq.unitary(rz)
+    assert np.allclose(cirq.unitary(rz), expected_unitary)
+    assert np.allclose(np.eye(len(rz_unitary)), rz_unitary.dot(rz_unitary.T.conj()))
+
+
+@pytest.mark.parametrize(
+    "phase_1, phase_2, addend, expected_state",
+    [
+        (0, 0, 1, 2),
+        (np.pi * 2 / 3, np.pi * 4 / 3, 0, 2),
+        (np.pi * 4 / 3, np.pi * 2 / 3, 0, 1),
+    ],
+)
+def test_X_HZH_qudit_identity(
+    phase_1: float, phase_2: float, addend: int, expected_state: int
+):
+    # For d=3, there are three identities: one for each swap.
+    # HH is equivalent to swapping |1> with |2>
+    # Applying a 1/3 turn to |1> and a 2/3 turn to |2> results in swapping
+    # |0> and |2>
+    # Applying a 2/3 turn to |1> and a 1/3 turn to |2> results in swapping
+    # |0> and |1>
+    qutrit = cirq.NamedQid("q0", dimension=3)
+    c = cirq.Circuit()
+    c.append(qudit_gates.QuditPlusGate(3, addend=addend)(qutrit))
+    c.append(qudit_gates.QuditHadamardGate(dimension=3)(qutrit))
+    c.append(
+        qudit_gates.QuditRzGate(dimension=3, radians=phase_1, phased_state=1)(qutrit)
+    )
+    c.append(
+        qudit_gates.QuditRzGate(dimension=3, radians=phase_2, phased_state=2)(qutrit)
+    )
+    c.append(qudit_gates.QuditHadamardGate(dimension=3)(qutrit))
+    c.append(cirq.measure(qutrit, key="m"))
+    sim = cirq.Simulator()
+    results = sim.run(c, repetitions=1000)
+    assert np.all(results.measurements["m"] == expected_state)
+
+
+@pytest.mark.parametrize(
+    "phase_1, phase_2, addend, expected_state",
+    [
+        (0, 0, 1, 2),
+        (np.pi * 2 / 3, np.pi * 4 / 3, 0, 2),
+        (np.pi * 4 / 3, np.pi * 2 / 3, 0, 1),
+    ],
+)
+def test_X_HZH_qudit_identity(
+    phase_1: float, phase_2: float, addend: int, expected_state: int
+):
+    # For d=3, there are three identities: one for each swap.
+    # HH is equivalent to swapping |1> with |2>
+    # Applying a 1/3 turn to |1> and a 2/3 turn to |2> results in swapping
+    # |0> and |2>
+    # Applying a 2/3 turn to |1> and a 1/3 turn to |2> results in swapping
+    # |0> and |1>
+    qutrit = cirq.NamedQid("q0", dimension=3)
+    c = cirq.Circuit()
+    c.append(qudit_gates.QuditPlusGate(3, addend=addend)(qutrit))
+    c.append(qudit_gates.QuditHadamardGate(dimension=3)(qutrit))
+    c.append(
+        qudit_gates.QuditRzGate(dimension=3, radians=phase_1, phased_state=1)(qutrit)
+    )
+    c.append(
+        qudit_gates.QuditRzGate(dimension=3, radians=phase_2, phased_state=2)(qutrit)
+    )
+    c.append(qudit_gates.QuditHadamardGate(dimension=3)(qutrit))
+    c.append(cirq.measure(qutrit, key="m"))
+    sim = cirq.Simulator()
+    results = sim.run(c, repetitions=1000)
+    assert np.all(results.measurements["m"] == expected_state)
+
+
 @pytest.mark.parametrize(
     "q0, q1", [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]
 )