Support arbitrary phase-modulated pulses (#688)

* Remove units from waveform values

* Add Pulse.ArbitraryPhase

* Add tests for parametrized pulses

* Abstract repr support

* rads -> rad

* Addressing review comments

* Fix conversion of phase into detuning

* Calculate phase modulation samples in ChannelSamples

* Support sequence drawing with phase modulation

* Fix typo

Co-authored-by: Antoine Cornillot <61453516+a-corni@users.noreply.github.com>

* Improve docstring format

---------

Co-authored-by: Antoine Cornillot <61453516+a-corni@users.noreply.github.com>

pulser-core/pulser/json/abstract_repr/deserializer.py

@@ -263,6 +263,19 @@ def _deserialize_operation(seq: Sequence, op: dict, vars: dict) -> None:
                 post_phase_shift=post_phase_shift,
             )
 
+        seq.add(
+            pulse=pulse,
+            channel=op["channel"],
+            protocol=op["protocol"],
+        )
+    elif op["op"] == "pulse_arbitrary_phase":
+        pulse = Pulse.ArbitraryPhase(
+            amplitude=_deserialize_waveform(op["amplitude"], vars),
+            phase=_deserialize_waveform(op["phase"], vars),
+            post_phase_shift=_deserialize_parameter(
+                op["post_phase_shift"], vars
+            ),
+        )
         seq.add(
             pulse=pulse,
             channel=op["channel"],

pulser-core/pulser/json/abstract_repr/schemas/sequence-schema.json

@@ -690,6 +690,49 @@
       ],
       "type": "object"
     },
+    "OpPulseArbitraryPhase": {
+      "additionalProperties": false,
+      "properties": {
+        "amplitude": {
+          "$ref": "#/definitions/Waveform",
+          "description": "Pulse amplitude waveform (in rad/µs)"
+        },
+        "channel": {
+          "$ref": "#/definitions/ChannelName",
+          "description": "Device channel to use for this pulse."
+        },
+        "op": {
+          "const": "pulse_arbitrary_phase",
+          "type": "string"
+        },
+        "phase": {
+          "$ref": "#/definitions/Waveform",
+          "description": "The pulse phase waveform (in radians)"
+        },
+        "post_phase_shift": {
+          "$ref": "#/definitions/ParametrizedNum",
+          "description": "A phase shift (in radians) immediately after the end of the pulse"
+        },
+        "protocol": {
+          "description": "A parametrized pulse constructed with an arbitrary phase waveform.",
+          "enum": [
+            "min-delay",
+            "no-delay",
+            "wait-for-all"
+          ],
+          "type": "string"
+        }
+      },
+      "required": [
+        "op",
+        "protocol",
+        "channel",
+        "amplitude",
+        "phase",
+        "post_phase_shift"
+      ],
+      "type": "object"
+    },
     "OpTarget": {
       "additionalProperties": false,
       "description": "Adds a waveform to the pulse.",
@@ -735,6 +778,9 @@
         {
           "$ref": "#/definitions/OpPulse"
         },
+        {
+          "$ref": "#/definitions/OpPulseArbitraryPhase"
+        },
         {
           "$ref": "#/definitions/OpPhaseShift"
         },
@@ -1138,4 +1184,4 @@
       "type": "object"
     }
   }
-}
+}

pulser-core/pulser/json/abstract_repr/serializer.py

@@ -296,7 +296,10 @@ def remove_kwarg_if_default(
                 "channel": data["channel"],
                 "protocol": data["protocol"],
             }
-            op_dict.update(data["pulse"]._to_abstract_repr())
+            pulse_abstract_repr = data["pulse"]._to_abstract_repr()
+            if "detuning" not in pulse_abstract_repr:
+                op_dict["op"] = "pulse_arbitrary_phase"
+            op_dict.update(pulse_abstract_repr)
             operations.append(op_dict)
         elif "phase_shift" in call.name:
             targets = call.args[1:]

pulser-core/pulser/json/abstract_repr/signatures.py

@@ -81,6 +81,9 @@ def all_pos_args(self) -> tuple[str, ...]:
     "Pulse": PulserSignature(
         pos=("amplitude", "detuning", "phase"), keyword=("post_phase_shift",)
     ),
+    "Pulse.ArbitraryPhase": PulserSignature(
+        pos=("amplitude", "phase"), keyword=("post_phase_shift",)
+    ),
     # Special case operators
     "truediv": PulserSignature(
         pos=("lhs", "rhs"), extra=dict(expression="div")

pulser-core/pulser/parametrized/paramobj.py

@@ -252,7 +252,11 @@ def _to_abstract_repr(self) -> dict[str, Any]:
                 cls_name = self.args[0].__name__
                 name = f"{cls_name}.{op_name}"
                 signature = SIGNATURES[
-                    "Pulse" if cls_name == "Pulse" else name
+                    (
+                        "Pulse"
+                        if cls_name == "Pulse" and op_name != "ArbitraryPhase"
+                        else name
+                    )
                 ]
                 # No existing classmethod has *args in its signature
                 assert (

pulser-core/pulser/pulse.py

@@ -28,7 +28,12 @@
 from pulser.json.utils import obj_to_dict
 from pulser.parametrized import Parametrized, ParamObj
 from pulser.parametrized.decorators import parametrize
-from pulser.waveforms import ConstantWaveform, Waveform
+from pulser.waveforms import (
+    ConstantWaveform,
+    CustomWaveform,
+    RampWaveform,
+    Waveform,
+)
 
 if TYPE_CHECKING:
     from pulser.channels.base_channel import Channel
@@ -50,18 +55,20 @@ class Pulse:
     If either quantity is constant throughout the entire pulse, use the
     ``ConstantDetuning``, ``ConstantAmplitude`` or ``ConstantPulse`` class
     method to create it.
+    If defining the pulse's phase modulation is preferred over its frequency
+    modulation, use ``Pulse.ArbitraryPhase()``.
 
     Note:
         We define the ``amplitude`` of a pulse to be its Rabi frequency,
         :math:`\Omega`, in rad/µs. Equivalently, the ``detuning`` is
         :math:`\delta`, also in rad/µs.
 
     Args:
-        amplitude: The pulse amplitude waveform.
-        detuning: The pulse detuning waveform.
+        amplitude: The pulse amplitude waveform (in rad/µs).
+        detuning: The pulse detuning waveform (in rad/µs).
         phase: The pulse phase (in radians).
         post_phase_shift: Optionally lets you add a phase
-            shift(in rads) immediately after the end of the pulse. This allows
+            shift(in rad) immediately after the end of the pulse. This allows
             for enconding of arbitrary single-qubit gates into a single pulse
             (see ``Sequence.phase_shift()`` for more information).
     """
@@ -127,11 +134,11 @@ def ConstantDetuning(
         """Creates a Pulse with an amplitude waveform and a constant detuning.
 
         Args:
-            amplitude: The pulse amplitude waveform.
+            amplitude: The pulse amplitude waveform (in rad/µs).
             detuning: The detuning value (in rad/µs).
             phase: The pulse phase (in radians).
             post_phase_shift: Optionally lets you add a
-                phase shift (in rads) immediately after the end of the pulse.
+                phase shift (in rad) immediately after the end of the pulse.
         """
         detuning_wf = ConstantWaveform(
             cast(Waveform, amplitude).duration, detuning
@@ -151,10 +158,10 @@ def ConstantAmplitude(
 
         Args:
             amplitude: The pulse amplitude value (in rad/µs).
-            detuning: The pulse detuning waveform.
+            detuning: The pulse detuning waveform (in rad/µs).
             phase: The pulse phase (in radians).
             post_phase_shift: Optionally lets you add a
-                phase shift (in rads) immediately after the end of the pulse.
+                phase shift (in rad) immediately after the end of the pulse.
         """
         amplitude_wf = ConstantWaveform(
             cast(Waveform, detuning).duration, amplitude
@@ -178,12 +185,68 @@ def ConstantPulse(
             detuning: The detuning value (in rad/µs).
             phase: The pulse phase (in radians).
             post_phase_shift: Optionally lets you add a
-                phase shift (in rads) immediately after the end of the pulse.
+                phase shift (in rad) immediately after the end of the pulse.
         """
         amplitude_wf = ConstantWaveform(duration, amplitude)
         detuning_wf = ConstantWaveform(duration, detuning)
         return cls(amplitude_wf, detuning_wf, phase, post_phase_shift)
 
+    @classmethod
+    @parametrize
+    def ArbitraryPhase(
+        cls,
+        amplitude: Waveform | Parametrized,
+        phase: Waveform | Parametrized,
+        post_phase_shift: float | Parametrized = 0.0,
+    ) -> Pulse:
+        r"""Pulse with an arbitrary phase waveform.
+
+        Args:
+            amplitude: The amplitude waveform (in rad/µs).
+            phase: The phase waveform (in rad).
+            post_phase_shift: Optionally lets you add a
+                phase shift (in rad) immediately after the end of the pulse.
+
+        Note:
+            Due to how the Hamiltonian is defined in Pulser, the phase and
+            detuning are related by
+
+            .. math:: \phi(t) = \phi_c - \sum_{k=0}^{t} \delta(k)
+
+            where :math:`\phi_c` is the pulse's constant phase offset.
+            From a given phase waveform, we extract the phase offset and
+            detuning waveform that respect this formula for every sample of
+            :math:`\phi(t)` and use these quantities to define the Pulse.
+
+        Warning:
+            Except when the phase waveform is a ``ConstantWaveform`` or a
+            ``RampWaveform``, the extracted detuning waveform will be a
+            ``CustomWaveform``. This makes the Pulse uncapable of automatically
+            extending its duration to fit a channel's clock period.
+
+        Returns:
+            A regular Pulse, with the phase waveform translated into a
+            detuning waveform and a constant phase offset.
+        """
+        if not isinstance(phase, Waveform):
+            raise TypeError(
+                f"'phase' must be a waveform, not of type {type(phase)}."
+            )
+        detuning: Waveform
+        if isinstance(phase, ConstantWaveform):
+            detuning = ConstantWaveform(phase.duration, 0.0)
+        elif isinstance(phase, RampWaveform):
+            detuning = ConstantWaveform(phase.duration, -phase.slope * 1e3)
+        else:
+            detuning_samples = -np.diff(phase.samples) * 1e3  # rad/ns->rad/µs
+            # Use the same value in the first two detuning samples
+            detuning = CustomWaveform(
+                np.pad(detuning_samples, (1, 0), mode="edge")
+            )
+        # Adjust phase_c to incorporate the first detuning sample
+        phase_c = phase.first_value + detuning.first_value * 1e-3
+        return cls(amplitude, detuning, phase_c, post_phase_shift)
+
     def draw(self) -> None:
         """Draws the pulse's amplitude and frequency waveforms."""
         fig, ax1 = plt.subplots()
@@ -254,15 +317,17 @@ def _to_abstract_repr(self) -> dict[str, Any]:
 
     def __str__(self) -> str:
         return (
-            f"Pulse(Amp={self.amplitude!s}, Detuning={self.detuning!s}, "
+            f"Pulse(Amp={self.amplitude!s} rad/µs, "
+            f"Detuning={self.detuning!s} rad/µs, "
             f"Phase={self.phase:.3g})"
         )
 
     def __repr__(self) -> str:
         return (
-            f"Pulse(amp={self.amplitude!r}, detuning={self.detuning!r}, "
-            + f"phase={self.phase:.3g}, "
-            + f"post_phase_shift={self.post_phase_shift:.3g})"
+            f"Pulse(amp={self.amplitude!r} rad/µs, "
+            f"detuning={self.detuning!r} rad/µs, "
+            f"phase={self.phase:.3g}, "
+            f"post_phase_shift={self.post_phase_shift:.3g})"
         )
 
     def __eq__(self, other: Any) -> bool:

pulser-core/pulser/sampler/samples.py

@@ -98,9 +98,15 @@ class ChannelSamples:
     eom_start_buffers: list[tuple[int, int]] = field(default_factory=list)
     eom_end_buffers: list[tuple[int, int]] = field(default_factory=list)
     target_time_slots: list[_TimeSlot] = field(default_factory=list)
+    _centered_phase: np.ndarray | None = None
 
     def __post_init__(self) -> None:
-        assert len(self.amp) == len(self.det) == len(self.phase)
+        assert (
+            len(self.amp)
+            == len(self.det)
+            == len(self.phase)
+            == len(self.centered_phase)
+        )
         self.duration = len(self.amp)
 
         for t in self.slots:
@@ -117,6 +123,28 @@ def initial_targets(self) -> set[QubitId]:
             else set()
         )
 
+    @property
+    def centered_phase(self) -> np.ndarray:
+        """The phase samples centered in ]-π, π]."""
+        if self._centered_phase is not None:
+            return self._centered_phase
+        phase_ = self.phase.copy() % (2 * np.pi)
+        phase_[phase_ > np.pi] -= 2 * np.pi
+        return phase_
+
+    @property
+    def phase_modulation(self) -> np.ndarray:
+        r"""The phase modulation samples (in rad).
+
+        Constructed by combining the integral of the detuning samples with the
+        phase offset samples according to
+
+        .. math:: \phi(t) = \phi_c(t) - \sum_{k=0}^{t} \delta(k)
+        """
+        return cast(
+            np.ndarray, self.centered_phase - np.cumsum(self.det * 1e-3)
+        )
+
     def extend_duration(self, new_duration: int) -> ChannelSamples:
         """Extends the duration of the samples.
 
@@ -151,7 +179,21 @@ def extend_duration(self, new_duration: int) -> ChannelSamples:
             (0, extension),
             mode="edge" if self.phase.size > 0 else "constant",
         )
-        return replace(self, amp=new_amp, det=new_detuning, phase=new_phase)
+        _new_centered_phase = None
+        if self._centered_phase is not None:
+            _new_centered_phase = np.pad(
+                self._centered_phase,
+                (0, extension),
+                mode="edge" if self._centered_phase.size > 0 else "constant",
+            )
+
+        return replace(
+            self,
+            amp=new_amp,
+            det=new_detuning,
+            phase=new_phase,
+            _centered_phase=_new_centered_phase,
+        )
 
     def is_empty(self) -> bool:
         """Whether the channel is effectively empty.
@@ -372,6 +414,9 @@ def masked(
             new_samples["det"] = channel_obj.modulate(self.det, keep_ends=True)
 
         new_samples["phase"] = channel_obj.modulate(self.phase, keep_ends=True)
+        new_samples["_centered_phase"] = channel_obj.modulate(
+            self.centered_phase, keep_ends=True
+        )
         for key in new_samples:
             new_samples[key] = new_samples[key][slice(0, max_duration)]
         return replace(self, **new_samples)