Repository: pasqal-io/Pulser

Issue:
Special operators for waveforms
I think it would be a nice improvement to have some special operators for waveforms. I came up with these ones:
 - Negation: Doing `-wf` would return a Waveform with all it's samples' sign flipped
 - Multiplication: Doing `wf * 2` would multiply all it's values by 2. It doesn't make sense to multiply two waveforms.
 
 Not so sure these make sense:
 - Addition: Doing `wf + 2` would add 2 to all values. `wf1 + wf2` would concatenate two waveforms.
 - Absolute value
 - Power
 - Modulo

Feel free to discuss the pertinence of each or add more suggestions. 

Comments:

Patch:
diff --git a/pulser/waveforms.py b/pulser/waveforms.py
index f8f6c08..9a126c1 100644
--- a/pulser/waveforms.py
+++ b/pulser/waveforms.py
@@ -80,6 +80,19 @@ class Waveform(ABC):
     def __repr__(self):
         pass
 
+    @abstractmethod
+    def __mul__(self, other):
+        pass
+
+    def __neg__(self):
+        return self.__mul__(-1)
+
+    def __truediv__(self, other):
+        if other == 0:
+            raise ZeroDivisionError("Can't divide a waveform by zero.")
+        else:
+            return self.__mul__(1/other)
+
     def __eq__(self, other):
         if not isinstance(other, Waveform):
             return False
@@ -188,6 +201,9 @@ class CompositeWaveform(Waveform):
     def __repr__(self):
         return f'CompositeWaveform({self.duration} ns, {self._waveforms!r})'
 
+    def __mul__(self, other):
+        return CompositeWaveform(*(wf * other for wf in self._waveforms))
+
 
 class CustomWaveform(Waveform):
     """A custom waveform.
@@ -200,7 +216,7 @@ class CustomWaveform(Waveform):
 
     def __init__(self, samples):
         """Initializes a custom waveform."""
-        samples_arr = np.array(samples)
+        samples_arr = np.array(samples, dtype=float)
         self._samples = samples_arr
         with warnings.catch_warnings():
             warnings.filterwarnings("error")
@@ -231,6 +247,9 @@ class CustomWaveform(Waveform):
     def __repr__(self):
         return f'CustomWaveform({self.duration} ns, {self.samples!r})'
 
+    def __mul__(self, other):
+        return CustomWaveform(self._samples * float(other))
+
 
 class ConstantWaveform(Waveform):
     """A waveform of constant value.
@@ -276,6 +295,9 @@ class ConstantWaveform(Waveform):
         return (f"ConstantWaveform({self._duration} ns, "
                 + f"{self._value:.3g} rad/µs)")
 
+    def __mul__(self, other):
+        return ConstantWaveform(self._duration, self._value * float(other))
+
 
 class RampWaveform(Waveform):
     """A linear ramp waveform.
@@ -328,6 +350,10 @@ class RampWaveform(Waveform):
         return (f"RampWaveform({self._duration} ns, " +
                 f"{self._start:.3g}->{self._stop:.3g} rad/µs)")
 
+    def __mul__(self, other):
+        k = float(other)
+        return RampWaveform(self._duration, self._start * k, self._stop * k)
+
 
 class BlackmanWaveform(Waveform):
     """A Blackman window of a specified duration and area.
@@ -402,3 +428,6 @@ class BlackmanWaveform(Waveform):
 
     def __repr__(self):
         return f"BlackmanWaveform({self._duration} ns, Area: {self._area:.3g})"
+
+    def __mul__(self, other):
+        return BlackmanWaveform(self._duration, self._area * float(other))


Test patch:
diff --git a/pulser/tests/test_waveforms.py b/pulser/tests/test_waveforms.py
index 94eab68..51e6997 100644
--- a/pulser/tests/test_waveforms.py
+++ b/pulser/tests/test_waveforms.py
@@ -122,9 +122,10 @@ def test_composite():
 
 
 def test_custom():
-    wf = CustomWaveform(np.arange(16))
+    data = np.arange(16, dtype=float)
+    wf = CustomWaveform(data)
     assert wf.__str__() == 'Custom'
-    assert wf.__repr__() == f'CustomWaveform(16 ns, {np.arange(16)!r})'
+    assert wf.__repr__() == f'CustomWaveform(16 ns, {data!r})'
 
 
 def test_ramp():
@@ -149,3 +150,13 @@ def test_blackman():
     wf = BlackmanWaveform.from_max_val(-10, -np.pi)
     assert np.isclose(wf.integral, -np.pi)
     assert np.min(wf.samples) > -10
+
+
+def test_ops():
+    assert -constant == ConstantWaveform(100, 3)
+    assert ramp * 2 == RampWaveform(2e3, 10, 38)
+    assert --custom == custom
+    assert blackman / 2 == BlackmanWaveform(40, np.pi / 2)
+    assert composite * 1 == composite
+    with pytest.raises(ZeroDivisionError):
+        constant / 0