feat: add phase RX gate (#945)

.coveragerc

@@ -23,9 +23,15 @@ exclude_lines =
     # Have to re-enable the standard pragma
     pragma: no cover
 
+    # Skipping import testing
+    from importlib.metadata import entry_points
+
     # Don't complain if tests don't hit defensive assertion code:
     raise NotImplementedError
 
+    # Avoid situation where system version causes coverage issues
+    if sys.version_info.minor == 9:
+
 [html]
 directory = build/coverage
 

src/braket/circuits/angled_gate.py

@@ -433,26 +433,27 @@ def get_angle(gate: AngledGate, **kwargs: FreeParameterExpression | str) -> Angl
 
 
 def _get_angles(
-    gate: TripleAngledGate, **kwargs: FreeParameterExpression | str
-) -> TripleAngledGate:
+    gate: DoubleAngledGate | TripleAngledGate, **kwargs: FreeParameterExpression | str
+) -> DoubleAngledGate | TripleAngledGate:
     """Gets the angle with all values substituted in that are requested.
 
     Args:
-        gate (TripleAngledGate): The subclass of TripleAngledGate for which the angle is being
-            obtained.
+        gate (DoubleAngledGate | TripleAngledGate): The subclass of multi angle AngledGate for
+            which the angle is being obtained.
         **kwargs (FreeParameterExpression | str): The named parameters that are being filled
             for a particular gate.
 
     Returns:
-        TripleAngledGate: A new gate of the type of the AngledGate originally used with all angles
-        updated.
+        DoubleAngledGate | TripleAngledGate: A new gate of the type of the AngledGate
+        originally used with all angles updated.
     """
-    new_angles = [
-        (
+    angles = [f"angle_{i + 1}" for i in range(len(gate._parameters))]
+    new_angles_args = {
+        angle: (
             getattr(gate, angle).subs(kwargs)
             if isinstance(getattr(gate, angle), FreeParameterExpression)
             else getattr(gate, angle)
         )
-        for angle in ("angle_1", "angle_2", "angle_3")
-    ]
-    return type(gate)(angle_1=new_angles[0], angle_2=new_angles[1], angle_3=new_angles[2])
+        for angle in angles
+    }
+    return type(gate)(**new_angles_args)

src/braket/circuits/gates.py

@@ -24,6 +24,7 @@
 from braket.circuits import circuit
 from braket.circuits.angled_gate import (
     AngledGate,
+    DoubleAngledGate,
     TripleAngledGate,
     _get_angles,
     _multi_angled_ascii_characters,
@@ -3298,6 +3299,124 @@ def gpi(
 Gate.register_gate(GPi)
 
 
+class PRx(DoubleAngledGate):
+    r"""Phase Rx gate.
+
+    Unitary matrix:
+
+        .. math:: \mathtt{PRx}(\theta,\phi) = \begin{bmatrix}
+                \cos{(\theta / 2)} & -i e^{-i \phi} \sin{(\theta / 2)} \\
+                -i e^{i \phi} \sin{(\theta / 2)} & \cos{(\theta / 2)}
+            \end{bmatrix}.
+
+    Args:
+        angle_1 (Union[FreeParameterExpression, float]): The first angle of the gate in
+            radians or expression representation.
+        angle_2 (Union[FreeParameterExpression, float]): The second angle of the gate in
+            radians or expression representation.
+    """
+
+    def __init__(
+        self,
+        angle_1: Union[FreeParameterExpression, float],
+        angle_2: Union[FreeParameterExpression, float],
+    ):
+        super().__init__(
+            angle_1=angle_1,
+            angle_2=angle_2,
+            qubit_count=None,
+            ascii_symbols=[_multi_angled_ascii_characters("PRx", angle_1, angle_2)],
+        )
+
+    @property
+    def _qasm_name(self) -> str:
+        return "prx"
+
+    def to_matrix(self) -> np.ndarray:
+        """Returns a matrix representation of this gate.
+
+        Returns:
+            np.ndarray: The matrix representation of this gate.
+        """
+        theta = self.angle_1
+        phi = self.angle_2
+        return np.array(
+            [
+                [
+                    np.cos(theta / 2),
+                    -1j * np.exp(-1j * phi) * np.sin(theta / 2),
+                ],
+                [
+                    -1j * np.exp(1j * phi) * np.sin(theta / 2),
+                    np.cos(theta / 2),
+                ],
+            ]
+        )
+
+    def adjoint(self) -> list[Gate]:
+        return [PRx(-self.angle_1, self.angle_2)]
+
+    @staticmethod
+    def fixed_qubit_count() -> int:
+        return 1
+
+    def bind_values(self, **kwargs) -> PRx:
+        return _get_angles(self, **kwargs)
+
+    @staticmethod
+    @circuit.subroutine(register=True)
+    def prx(
+        target: QubitSetInput,
+        angle_1: Union[FreeParameterExpression, float],
+        angle_2: Union[FreeParameterExpression, float],
+        *,
+        control: Optional[QubitSetInput] = None,
+        control_state: Optional[BasisStateInput] = None,
+        power: float = 1,
+    ) -> Iterable[Instruction]:
+        r"""PhaseRx gate.
+
+        .. math:: \mathtt{PRx}(\theta,\phi) = \begin{bmatrix}
+                \cos{(\theta / 2)} & -i e^{-i \phi} \sin{(\theta / 2)} \\
+                -i e^{i \phi} \sin{(\theta / 2)} & \cos{(\theta / 2)}
+            \end{bmatrix}.
+
+        Args:
+            target (QubitSetInput): Target qubit(s).
+            angle_1 (Union[FreeParameterExpression, float]): First angle in radians.
+            angle_2 (Union[FreeParameterExpression, float]): Second angle in radians.
+            control (Optional[QubitSetInput]): Control qubit(s). Default None.
+            control_state (Optional[BasisStateInput]): Quantum state on which to control the
+                operation. Must be a binary sequence of same length as number of qubits in
+                `control`. Will be ignored if `control` is not present. May be represented as a
+                string, list, or int. For example "0101", [0, 1, 0, 1], 5 all represent
+                controlling on qubits 0 and 2 being in the \\|0⟩ state and qubits 1 and 3 being
+                in the \\|1⟩ state. Default "1" * len(control).
+            power (float): Integer or fractional power to raise the gate to. Negative
+                powers will be split into an inverse, accompanied by the positive power.
+                Default 1.
+
+        Returns:
+            Iterable[Instruction]: PhaseRx instruction.
+
+        Examples:
+            >>> circ = Circuit().prx(0, 0.15, 0.25)
+        """
+        return [
+            Instruction(
+                PRx(angle_1, angle_2),
+                target=qubit,
+                control=control,
+                control_state=control_state,
+                power=power,
+            )
+            for qubit in QubitSet(target)
+        ]
+
+
+Gate.register_gate(PRx)
+
+
 class GPi2(AngledGate):
     r"""IonQ GPi2 gate.
 

src/braket/circuits/translations.py

@@ -67,6 +67,7 @@
     "cswap": braket_gates.CSwap,
     "gpi": braket_gates.GPi,
     "gpi2": braket_gates.GPi2,
+    "prx": braket_gates.PRx,
     "ms": braket_gates.MS,
     "unitary": braket_gates.Unitary,
 }

test/unit_tests/braket/circuits/test_circuit.py

@@ -2536,6 +2536,7 @@ def test_to_unitary_with_global_phase():
         (Circuit().cphaseshift00(0, 1, 0.15), gates.CPhaseShift00(0.15).to_matrix()),
         (Circuit().cphaseshift01(0, 1, 0.15), gates.CPhaseShift01(0.15).to_matrix()),
         (Circuit().cphaseshift10(0, 1, 0.15), gates.CPhaseShift10(0.15).to_matrix()),
+        (Circuit().prx(0, 1, 0.15), gates.PRx(1, 0.15).to_matrix()),
         (Circuit().cy(0, 1), gates.CY().to_matrix()),
         (Circuit().cz(0, 1), gates.CZ().to_matrix()),
         (Circuit().xx(0, 1, 0.15), gates.XX(0.15).to_matrix()),

test/unit_tests/braket/circuits/test_gates.py

@@ -39,6 +39,10 @@ class NoTarget:
     pass
 
 
+class DoubleAngle:
+    pass
+
+
 class TripleAngle:
     pass
 
@@ -103,6 +107,7 @@ class SingleNegControlModifier:
     (Gate.ZZ, "zz", ir.ZZ, [DoubleTarget, Angle], {}),
     (Gate.GPi, "gpi", None, [SingleTarget, Angle], {}),
     (Gate.GPi2, "gpi2", None, [SingleTarget, Angle], {}),
+    (Gate.PRx, "prx", None, [SingleTarget, DoubleAngle], {}),
     (Gate.MS, "ms", None, [DoubleTarget, TripleAngle], {}),
     (
         Gate.Unitary,
@@ -145,9 +150,11 @@ class SingleNegControlModifier:
     Gate.CPhaseShift10,
     Gate.GPi,
     Gate.GPi2,
+    Gate.PRx,
     Gate.MS,
 ]
 
+
 invalid_unitary_matrices = [
     (np.array([[1]])),
     (np.array([1])),
@@ -179,6 +186,10 @@ def angle_valid_input(**kwargs):
     return {"angle": 0.123}
 
 
+def double_angle_valid_input(**kwargs):
+    return {"angle_1": 0.123, "angle_2": 3.567}
+
+
 def triple_angle_valid_input(**kwargs):
     return {"angle_1": 0.123, "angle_2": 4.567, "angle_3": 8.910}
 
@@ -217,6 +228,7 @@ def two_dimensional_matrix_valid_input(**kwargs):
     "SingleTarget": single_target_valid_input,
     "DoubleTarget": double_target_valid_ir_input,
     "Angle": angle_valid_input,
+    "DoubleAngle": double_angle_valid_input,
     "TripleAngle": triple_angle_valid_input,
     "SingleControl": single_control_valid_input,
     "SingleNegControlModifier": single_neg_control_valid_input,
@@ -273,7 +285,7 @@ def create_valid_target_input(irsubclasses):
             control_state = list(single_neg_control_valid_input()["control_state"])
         elif subclass == DoubleControl:
             qubit_set = list(double_control_valid_ir_input().values()) + qubit_set
-        elif subclass in (Angle, TwoDimensionalMatrix, TripleAngle):
+        elif subclass in (Angle, TwoDimensionalMatrix, DoubleAngle, TripleAngle):
             pass
         else:
             raise ValueError("Invalid subclass")
@@ -287,6 +299,8 @@ def create_valid_gate_class_input(irsubclasses, **kwargs):
     input = {}
     if Angle in irsubclasses:
         input.update(angle_valid_input())
+    if DoubleAngle in irsubclasses:
+        input.update(double_angle_valid_input())
     if TripleAngle in irsubclasses:
         input.update(triple_angle_valid_input())
     if TwoDimensionalMatrix in irsubclasses:
@@ -313,7 +327,7 @@ def calculate_qubit_count(irsubclasses):
             qubit_count += 2
         elif subclass == MultiTarget:
             qubit_count += 3
-        elif subclass in (NoTarget, Angle, TwoDimensionalMatrix, TripleAngle):
+        elif subclass in (NoTarget, Angle, TwoDimensionalMatrix, DoubleAngle, TripleAngle):
             pass
         else:
             raise ValueError("Invalid subclass")
@@ -847,6 +861,18 @@ def test_ir_gate_level(testclass, subroutine_name, irclass, irsubclasses, kwargs
             OpenQASMSerializationProperties(qubit_reference_type=QubitReferenceType.PHYSICAL),
             "gpi2(0.17) $4;",
         ),
+        (
+            Gate.PRx(angle_1=0.17, angle_2=3.45),
+            [4],
+            OpenQASMSerializationProperties(qubit_reference_type=QubitReferenceType.VIRTUAL),
+            "prx(0.17, 3.45) q[4];",
+        ),
+        (
+            Gate.PRx(angle_1=0.17, angle_2=3.45),
+            [4],
+            OpenQASMSerializationProperties(qubit_reference_type=QubitReferenceType.PHYSICAL),
+            "prx(0.17, 3.45) $4;",
+        ),
         (
             Gate.MS(angle_1=0.17, angle_2=3.45),
             [4, 5],
@@ -902,6 +928,8 @@ def test_gate_subroutine(testclass, subroutine_name, irclass, irsubclasses, kwar
         subroutine_input = {"target": multi_targets}
         if Angle in irsubclasses:
             subroutine_input.update(angle_valid_input())
+        if DoubleAngle in irsubclasses:
+            subroutine_input.update(double_angle_valid_input())
         if TripleAngle in irsubclasses:
             subroutine_input.update(triple_angle_valid_input())
         assert subroutine(**subroutine_input) == Circuit(instruction_list)
@@ -1012,8 +1040,13 @@ def test_large_unitary():
 
 @pytest.mark.parametrize("gate", parameterizable_gates)
 def test_bind_values(gate):
+    double_angled = gate.__name__ in ["PRx"]
     triple_angled = gate.__name__ in ("MS", "U")
-    num_params = 3 if triple_angled else 1
+    num_params = 1
+    if triple_angled:
+        num_params = 3
+    elif double_angled:
+        num_params = 2
     thetas = [FreeParameter(f"theta_{i}") for i in range(num_params)]
     mapping = {f"theta_{i}": i for i in range(num_params)}
     param_gate = gate(*thetas)
@@ -1024,6 +1057,9 @@ def test_bind_values(gate):
     if triple_angled:
         for angle in new_gate.angle_1, new_gate.angle_2, new_gate.angle_3:
             assert isinstance(angle, float)
+    elif double_angled:
+        for angle in new_gate.angle_1, new_gate.angle_2:
+            assert isinstance(angle, float)
     else:
         assert isinstance(new_gate.angle, float)