V4 - fix: Attempt to reconstruct TemplateWaveforms from quil_rs.WaveformInvocations

pyquil/quilatom.py

@@ -15,6 +15,7 @@
 ##############################################################################
 
 from fractions import Fraction
+import inspect
 from numbers import Number
 from typing import (
     Any,
@@ -1041,6 +1042,8 @@ def fdel(self: "TemplateWaveform") -> None:
     reason="The TemplateWaveform class will be removed, consider using WaveformInvocation instead.",
 )
 class TemplateWaveform(quil_rs.WaveformInvocation, QuilAtom):
+    NAME: ClassVar[str]
+
     def __new__(
         cls,
         name: str,
@@ -1102,6 +1105,42 @@ def samples(self, rate: float) -> np.ndarray:
 
     @classmethod
     def _from_rs_waveform_invocation(cls, waveform: quil_rs.WaveformInvocation) -> "TemplateWaveform":
+        """
+        The ``quil`` package has no equivalent to ``TemplateWaveform``s, this function checks the name and
+        properties of a ``quil`` ``WaveformInvocation`` to see if they potentially match a subclass of
+        ``TemplateWaveform``. If a match is found and construction succeeds, then that type is returned.
+        Otherwise, a generic ``WaveformInvocation`` is returned.
+        """
+        from pyquil.quiltwaveforms import (
+            FlatWaveform,
+            GaussianWaveform,
+            DragGaussianWaveform,
+            HrmGaussianWaveform,
+            ErfSquareWaveform,
+            BoxcarAveragerKernel,
+        )
+
+        template: Type["TemplateWaveform"]  # mypy needs a type annotation here to understand this.
+        for template in [
+            FlatWaveform,
+            GaussianWaveform,
+            DragGaussianWaveform,
+            HrmGaussianWaveform,
+            ErfSquareWaveform,
+            BoxcarAveragerKernel,
+        ]:
+            if template.NAME != waveform.name:
+                continue
+            parameter_names = [
+                parameter[0] for parameter in inspect.getmembers(template, lambda a: isinstance(a, property))
+            ]
+            if set(waveform.parameters.keys()).issubset(parameter_names):
+                try:
+                    parameters = {key: value.inner() for key, value in waveform.parameters.items()}
+                    return template(**parameters)  # type: ignore[arg-type]
+                except TypeError:
+                    break
+
         return super().__new__(cls, waveform.name, waveform.parameters)
 
     def __str__(self) -> str:

pyquil/quiltwaveforms.py

@@ -28,6 +28,8 @@ class FlatWaveform(TemplateWaveform):
     A flat (constant) waveform.
     """
 
+    NAME = "flat"
+
     def __new__(
         cls,
         duration: float,
@@ -36,7 +38,7 @@ def __new__(
         phase: Optional[float] = None,
         detuning: Optional[float] = None,
     ) -> Self:
-        return super().__new__(cls, "flat", duration=duration, iq=iq, scale=scale, phase=phase, detuning=detuning)
+        return super().__new__(cls, cls.NAME, duration=duration, iq=iq, scale=scale, phase=phase, detuning=detuning)
 
     iq = _template_waveform_property("iq", doc="A raw IQ value.")
     scale = _template_waveform_property("scale", doc="An optional global scaling factor.", dtype=float)
@@ -55,6 +57,8 @@ def samples(self, rate: float) -> np.ndarray:
 class GaussianWaveform(TemplateWaveform):
     """A Gaussian pulse."""
 
+    NAME = "gaussian"
+
     def __new__(
         cls,
         duration: float,
@@ -65,7 +69,7 @@ def __new__(
         detuning: Optional[float] = None,
     ) -> Self:
         return super().__new__(
-            cls, "gaussian", duration=duration, fwhm=fwhm, t0=t0, scale=scale, phase=phase, detuning=detuning
+            cls, cls.NAME, duration=duration, fwhm=fwhm, t0=t0, scale=scale, phase=phase, detuning=detuning
         )
 
     fwhm = _template_waveform_property("fwhm", doc="The Full-Width-Half-Max of the Gaussian (seconds).", dtype=float)
@@ -92,6 +96,8 @@ def samples(self, rate: float) -> np.ndarray:
 class DragGaussianWaveform(TemplateWaveform):
     """A DRAG Gaussian pulse."""
 
+    NAME = "drag_gaussian"
+
     def __new__(
         cls,
         duration: float,
@@ -105,7 +111,7 @@ def __new__(
     ) -> Self:
         return super().__new__(
             cls,
-            "drag_gaussian",
+            cls.NAME,
             duration=duration,
             fwhm=fwhm,
             t0=t0,
@@ -151,6 +157,8 @@ class HrmGaussianWaveform(TemplateWaveform):
         10.1063/1.447644
     """
 
+    NAME = "hrm_gaussian"
+
     def __new__(
         cls,
         duration: float,
@@ -165,7 +173,7 @@ def __new__(
     ) -> Self:
         return super().__new__(
             cls,
-            "hrm_gaussian",
+            cls.NAME,
             duration=duration,
             fwhm=fwhm,
             t0=t0,
@@ -220,6 +228,8 @@ def samples(self, rate: float) -> np.ndarray:
 class ErfSquareWaveform(TemplateWaveform):
     """A pulse with a flat top and edges that are error functions (erf)."""
 
+    NAME = "erf_square"
+
     def __new__(
         cls,
         duration: float,
@@ -232,7 +242,7 @@ def __new__(
     ) -> Self:
         return super().__new__(
             cls,
-            "erf_square",
+            cls.NAME,
             duration=duration,
             risetime=risetime,
             pad_left=pad_left,
@@ -278,14 +288,17 @@ def samples(self, rate: float) -> np.ndarray:
     reason="The TemplateWaveform class and its subclasses will be removed, consider using WaveformInvocation instead.",
 )
 class BoxcarAveragerKernel(TemplateWaveform):
+
+    NAME = "boxcar_kernel"
+
     def __new__(
         cls,
         duration: float,
         scale: Optional[float] = None,
         phase: Optional[float] = None,
         detuning: Optional[float] = None,
     ) -> Self:
-        return super().__new__(cls, "boxcar_kernel", duration=duration, scale=scale, phase=phase, detuning=detuning)
+        return super().__new__(cls, cls.NAME, duration=duration, scale=scale, phase=phase, detuning=detuning)
 
     scale = _template_waveform_property("scale", doc="An optional global scaling factor.", dtype=float)
 

test/unit/__snapshots__/test_quilbase.ambr

@@ -536,9 +536,24 @@
 # name: TestPulse.test_out[Blocking]
   'PULSE 123 q "FRAMEX" WAVEFORMY'
 # ---
+# name: TestPulse.test_out[BoxcarAveragerKernel]
+  'NONBLOCKING PULSE 123 q "FRAMEX" boxcar_kernel(duration: 2.5, scale: 1)'
+# ---
+# name: TestPulse.test_out[DragGaussianWaveform]
+  'NONBLOCKING PULSE 123 q "FRAMEX" drag_gaussian(alpha: 1, anh: 0.1, duration: 2.5, fwhm: 1, t0: 1)'
+# ---
+# name: TestPulse.test_out[ErfSquareWaveform]
+  'NONBLOCKING PULSE 123 q "FRAMEX" erf_square(duration: 2.5, pad_left: 1, pad_right: 0.1, risetime: 1, scale: 1)'
+# ---
 # name: TestPulse.test_out[FlatWaveform]
   'NONBLOCKING PULSE 123 q "FRAMEX" flat(duration: 2.5, iq: 1+2.0i)'
 # ---
+# name: TestPulse.test_out[GaussianWaveform]
+  'NONBLOCKING PULSE 123 q "FRAMEX" gaussian(duration: 2.5, fwhm: 1, phase: 0.1, t0: 1)'
+# ---
+# name: TestPulse.test_out[HrmGaussianWaveform]
+  'NONBLOCKING PULSE 123 q "FRAMEX" hrm_gaussian(alpha: 1, anh: 0.1, duration: 2.5, fwhm: 1, second_order_hrm_coeff: 0.5, t0: 1)'
+# ---
 # name: TestPulse.test_out[NonBlocking]
   'NONBLOCKING PULSE 123 q "FRAMEX" WAVEFORMY'
 # ---

test/unit/test_quilbase.py

@@ -71,9 +71,16 @@
     Nop,
 )
 from pyquil.paulis import PauliSum, PauliTerm
-from pyquil.quilatom import BinaryExp, Mul, Frame, Qubit, Expression, Waveform, WaveformReference
+from pyquil.quilatom import BinaryExp, Mul, Frame, Qubit, Expression, Waveform, WaveformReference, TemplateWaveform
 from pyquil.api._compiler import QPUCompiler
-from pyquil.quiltwaveforms import FlatWaveform
+from pyquil.quiltwaveforms import (
+    FlatWaveform,
+    GaussianWaveform,
+    DragGaussianWaveform,
+    HrmGaussianWaveform,
+    ErfSquareWaveform,
+    BoxcarAveragerKernel,
+)
 
 
 @pytest.mark.parametrize(
@@ -875,8 +882,42 @@ def test_nonblocking(self, capture: Capture, nonblocking: bool):
             FlatWaveform(duration=2.5, iq=complex(1.0, 2.0)),
             True,
         ),
+        (
+            Frame([Qubit(123), FormalArgument("q")], "FRAMEX"),
+            GaussianWaveform(duration=2.5, fwhm=1.0, t0=1.0, phase=0.1),
+            True,
+        ),
+        (
+            Frame([Qubit(123), FormalArgument("q")], "FRAMEX"),
+            DragGaussianWaveform(duration=2.5, fwhm=1.0, t0=1.0, anh=0.1, alpha=1.0),
+            True,
+        ),
+        (
+            Frame([Qubit(123), FormalArgument("q")], "FRAMEX"),
+            HrmGaussianWaveform(duration=2.5, fwhm=1.0, t0=1.0, anh=0.1, alpha=1.0, second_order_hrm_coeff=0.5),
+            True,
+        ),
+        (
+            Frame([Qubit(123), FormalArgument("q")], "FRAMEX"),
+            ErfSquareWaveform(duration=2.5, risetime=1.0, pad_left=1.0, pad_right=0.1, scale=1.0),
+            True,
+        ),
+        (
+            Frame([Qubit(123), FormalArgument("q")], "FRAMEX"),
+            BoxcarAveragerKernel(duration=2.5, scale=1.0),
+            True,
+        ),
     ],
-    ids=("Blocking", "NonBlocking", "FlatWaveform"),
+    ids=(
+        "Blocking",
+        "NonBlocking",
+        "FlatWaveform",
+        "GaussianWaveform",
+        "DragGaussianWaveform",
+        "HrmGaussianWaveform",
+        "ErfSquareWaveform",
+        "BoxcarAveragerKernel",
+    ),
 )
 class TestPulse:
     @pytest.fixture
@@ -893,6 +934,9 @@ def test_frame(self, pulse: Pulse, frame: Frame):
 
     def test_waveform(self, pulse: Pulse, waveform: Waveform):
         assert pulse.waveform == waveform
+        if isinstance(waveform, TemplateWaveform):
+            assert isinstance(pulse.waveform, type(waveform))
+            pulse.waveform.samples(0.5)
         pulse.waveform = WaveformReference("new-waveform")
         assert pulse.waveform == WaveformReference("new-waveform")