Add local RZ gate for pulser backend (#114)

qadence/backends/pulser/pulses.py

@@ -19,7 +19,7 @@
     Interaction,
     WaitBlock,
 )
-from qadence.operations import RX, RY, AnalogEntanglement, OpName
+from qadence.operations import RX, RY, RZ, AnalogEntanglement, OpName
 from qadence.parameters import evaluate
 
 from .channels import GLOBAL_CHANNEL, LOCAL_CHANNEL
@@ -32,6 +32,7 @@
     OpName.ZERO,
     OpName.RX,
     OpName.RY,
+    OpName.RZ,
     OpName.ANALOGENTANG,
     OpName.ANALOGRX,
     OpName.ANALOGRY,
@@ -64,8 +65,9 @@ def add_pulses(
     local_channel = sequence.device.channels["rydberg_local"]
     global_channel = sequence.device.channels["rydberg_global"]
 
-    rx = partial(digital_rot_pulse, channel=local_channel, phase=0, config=config)
-    ry = partial(digital_rot_pulse, channel=local_channel, phase=np.pi / 2, config=config)
+    rx = partial(digital_xy_rot_pulse, channel=local_channel, phase=0, config=config)
+    ry = partial(digital_xy_rot_pulse, channel=local_channel, phase=np.pi / 2, config=config)
+    rz = partial(digital_z_rot_pulse, channel=local_channel, phase=np.pi / 2, config=config)
 
     # TODO: lets move those to `@singledipatch`ed functions
     if isinstance(block, WaitBlock):
@@ -115,10 +117,15 @@ def add_pulses(
             entangle_pulse(t, global_channel, config), GLOBAL_CHANNEL, protocol="wait-for-all"
         )
 
-    elif isinstance(block, (RX, RY)):
+    elif isinstance(block, (RX, RY, RZ)):
         (uuid, p) = block.parameters.uuid_param("parameter")
         angle = evaluate(p) if p.is_number else sequence.declare_variable(uuid)
-        pulse = rx(angle) if isinstance(block, RX) else ry(angle)
+        if isinstance(block, RX):
+            pulse = rx(angle)
+        elif isinstance(block, RY):
+            pulse = ry(angle)
+        elif isinstance(block, RZ):
+            pulse = rz(angle)
         sequence.target(qubit_support, LOCAL_CHANNEL)
         sequence.add(pulse, LOCAL_CHANNEL, protocol="wait-for-all")
 
@@ -197,7 +204,7 @@ def entangle_pulse(
     return Pulse(amplitude=amplitude, detuning=detuning, phase=np.pi / 2)
 
 
-def digital_rot_pulse(
+def digital_xy_rot_pulse(
     angle: TVar | float, phase: float, channel: Channel, config: Configuration | None = None
 ) -> Pulse:
     if config is None:
@@ -214,3 +221,33 @@ def digital_rot_pulse(
     )
 
     return Pulse.ConstantDetuning(amplitude=amplitude_wf, detuning=0, phase=phase)
+
+
+def digital_z_rot_pulse(
+    angle: TVar | float, phase: float, channel: Channel, config: Configuration | None = None
+) -> Pulse:
+    if config is None:
+        max_det = channel.max_abs_detuning
+    else:
+        max_det = config.detuning if config.detuning is not None else channel.max_abs_detuning
+
+    # get pulse duration in ns
+    duration = 1000 * abs(angle) / max_det
+
+    # create amplitude waveform
+    amp_wf = SquareWaveform.from_duration(
+        duration=duration,  # type: ignore
+        max_amp=0.0,  # type: ignore[arg-type]
+        duration_steps=channel.clock_period,  # type: ignore[attr-defined]
+        min_duration=channel.min_duration,
+    )
+
+    # create detuning waveform
+    det_wf = SquareWaveform.from_duration(
+        duration=duration,  # type: ignore
+        max_amp=max_det,  # type: ignore[arg-type]
+        duration_steps=channel.clock_period,  # type: ignore[attr-defined]
+        min_duration=channel.min_duration,
+    )
+
+    return Pulse(amplitude=amp_wf, detuning=det_wf, phase=abs(phase))

tests/backends/pulser_basic/test_pulser_conversion.py

@@ -12,20 +12,21 @@
 from qadence.backends.pulser.backend import make_sequence
 from qadence.backends.pulser.config import Configuration
 from qadence.backends.pulser.devices import Device, RealisticDevice
-from qadence.backends.pulser.pulses import digital_rot_pulse, entangle_pulse
+from qadence.backends.pulser.pulses import digital_xy_rot_pulse, digital_z_rot_pulse, entangle_pulse
 from qadence.blocks import AbstractBlock
 from qadence.blocks.analog import Interaction
 from qadence.circuit import QuantumCircuit
 from qadence.divergences import js_divergence
-from qadence.operations import RX, RY, entangle
+from qadence.operations import RX, RY, RZ, entangle
 from qadence.register import Register as QadenceRegister
 
 
 @pytest.mark.parametrize(
     "Qadence_op, func",
     [
-        (RX(0, 1.5), lambda ch: digital_rot_pulse(1.5, 0, ch)),
-        (RY(1, 1.5), lambda ch: digital_rot_pulse(1.5, np.pi / 2, ch)),
+        (RX(0, 1.5), lambda ch: digital_xy_rot_pulse(1.5, 0, ch)),
+        (RY(1, 1.5), lambda ch: digital_xy_rot_pulse(1.5, np.pi / 2, ch)),
+        (RZ(0, 1.5), lambda ch: digital_z_rot_pulse(1.5, np.pi / 2, ch)),
     ],
 )
 def test_single_qubit_block_conversion(Qadence_op: AbstractBlock, func: Callable) -> None:

tests/backends/test_pulser_pyq_compat.py

@@ -12,7 +12,7 @@
 from qadence.constructors import ising_hamiltonian, total_magnetization
 from qadence.divergences import js_divergence
 from qadence.models import QuantumModel
-from qadence.operations import CNOT, RX, RY, AnalogRX, AnalogRY, H, X, Z, entangle
+from qadence.operations import CNOT, RX, RY, RZ, AnalogRX, AnalogRY, H, X, Z, entangle
 from qadence.parameters import FeatureParameter
 from qadence.types import DiffMode
 
@@ -52,12 +52,14 @@ def test_compatibility_pyqtorch_pulser_entanglement(
 def test_compatibility_pyqtorch_pulser_digital_rot(obs: AbstractBlock) -> None:
     phi = FeatureParameter("phi")
     psi = FeatureParameter("psi")
+    chi = FeatureParameter("chi")
 
     n_qubits = 2
 
     block = chain(
         kron(RX(0, phi), RX(1, phi)),
         kron(RY(0, psi), RY(1, psi)),
+        kron(RZ(0, chi), RZ(1, chi)),
     )
     pyqtorch_circuit = QuantumCircuit(n_qubits, block)
 
@@ -67,7 +69,7 @@ def test_compatibility_pyqtorch_pulser_digital_rot(obs: AbstractBlock) -> None:
     model_pyqtorch = QuantumModel(
         pyqtorch_circuit, backend=BackendName.PYQTORCH, diff_mode=DiffMode.AD, observable=obs
     )
-    conf = {"spacing": LARGE_SPACING, "amplitude_local": 2 * np.pi}
+    conf = {"spacing": LARGE_SPACING, "amplitude_local": 2 * np.pi, "detuning": 2 * np.pi}
     model_pulser = QuantumModel(
         pulser_circuit,
         backend=BackendName.PULSER,
@@ -80,6 +82,7 @@ def test_compatibility_pyqtorch_pulser_digital_rot(obs: AbstractBlock) -> None:
     values = {
         "phi": torch.rand(batch_size),
         "psi": torch.rand(batch_size),
+        "chi": torch.rand(batch_size),
     }
 
     pyqtorch_expval = model_pyqtorch.expectation(values=values)