waveforms 2.2.0__tar.gz → 2.2.2__tar.gz

{waveforms-2.2.0 → waveforms-2.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waveforms
-Version: 2.2.0
+Version: 2.2.2
 Summary: Edit waveforms used in experiment
 Author-email: feihoo87 <feihoo87@gmail.com>
 Maintainer-email: feihoo87 <feihoo87@gmail.com>

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms/_waveform.pyi

@@ -2,6 +2,7 @@ from typing import Callable
 
 import numpy as np
 from numpy import e, inf, pi
+from numpy.typing import NDArray
 
 NDIGITS: int = ...
 __TypeIndex: int = ...
@@ -48,12 +49,15 @@ def pow(x: tuple[tuple, tuple], n: int) -> tuple[tuple, tuple]:
     pass
 
 
-def calc_parts(bounds: tuple,
-               seq: tuple,
-               x: np.ndarray,
-               function_lib: dict,
-               min=-inf,
-               max=inf) -> tuple[list[tuple[int, int, np.ndarray]], type]:
+def calc_parts(
+    bounds: tuple,
+    seq: tuple,
+    x: np.ndarray,
+    function_lib: dict,
+    min=-inf,
+    max=inf
+) -> tuple[list[tuple[int, int, NDArray[np.float64 | np.complex128] | int
+                      | float | complex]], type]:
     pass
 
 
@@ -104,6 +108,7 @@ COSH: int = ...
 SINH: int = ...
 DRAG: int = ...
 MOLLIFIER: int = ...
+D_GAUSSIAN: int = ...
 
 
 def simplify(expr: tuple[tuple, tuple], eps: float) -> tuple[tuple, tuple]:

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms/_waveform.pyx

@@ -295,6 +295,11 @@ def _GAUSSIAN(t, std_sq2):
     return np.exp(-(t / std_sq2)**2)
 
 
+def _D_GAUSSIAN(t, std_sq2, n):
+    return (-1)**n / std_sq2**n * special.hermite(n)(
+        t / std_sq2) * np.exp(-(t / std_sq2)**2)
+
+
 def _ERF(t, std_sq2):
     return special.erf(t / std_sq2)
 
@@ -353,14 +358,17 @@ def _drag(t: np.ndarray, t0: float, freq: float, width: float, delta: float,
 
 def _mollifier(t: np.ndarray, r: float, d: int):
     x = t / r
+    xx_1 = np.abs(x)**2 - 1
     if d == 0:
-        return np.exp(1 / (np.abs(x)**2 - 1) + 1)
+        return np.where(xx_1 == 0, 0, np.exp(1 / xx_1 + 1))
     else:
         p = np.poly1d([-2, 0])
         for n in range(1, d):
             p = np.poly1d([1, 0, -2, 0, 1]) * p.deriv() + np.poly1d(
                 [-4 * n, 0, 4 * n - 2, 0]) * p
-        return np.exp(1 / (np.abs(x)**2 - 1) + 1) * p(x) / (1 - x**2)**(2 * d) / r**d
+        return np.where(xx_1 == 0, 0,
+                        np.exp(1 / xx_1 + 1) /
+                        (-xx_1)**(2 * d)) * p(x) / r**d
 
 
 LINEAR = registerBaseFunc(_LINEAR)
@@ -377,6 +385,7 @@ COSH = registerBaseFunc(_COSH)
 SINH = registerBaseFunc(_SINH)
 DRAG = registerBaseFunc(_drag)
 MOLLIFIER = registerBaseFunc(_mollifier)
+D_GAUSSIAN = registerBaseFunc(_D_GAUSSIAN)
 
 
 def _d_LINEAR(shift, *args):
@@ -447,7 +456,11 @@ def _d_HYPERBOLICCHIRP(shift, f0, k, phi0):
 
 
 def _d_MOLLIFIER(shift, r, d):
-    return (((((MOLLIFIER, r, d+1, shift), ), (1, )), ), (1, ))
+    return (((((MOLLIFIER, r, d + 1, shift), ), (1, )), ), (1, ))
+
+
+def _d_D_GAUSSIAN(shift, std_sq2, n):
+    return (((((D_GAUSSIAN, std_sq2, n + 1, shift), ), (1, )), ), (1, ))
 
 
 # register derivative
@@ -464,6 +477,7 @@ registerDerivative(LINEARCHIRP, _d_LINEARCHIRP)
 registerDerivative(EXPONENTIALCHIRP, _d_EXPONENTIALCHIRP)
 registerDerivative(HYPERBOLICCHIRP, _d_HYPERBOLICCHIRP)
 registerDerivative(MOLLIFIER, _d_MOLLIFIER)
+registerDerivative(D_GAUSSIAN, _d_D_GAUSSIAN)
 
 
 def _cos_power_n(x, n):

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms/distortion.py

@@ -1,8 +1,9 @@
 import warnings
-from itertools import repeat, zip_longest
-from typing import Sequence
+from itertools import zip_longest
+from typing import Sequence, cast
 
 import numpy as np
+from numpy.typing import NDArray
 from scipy.fftpack import fft, fftfreq, ifft, ifftshift
 from scipy.optimize import curve_fit
 from scipy.signal import fftconvolve, lfilter, lfiltic, tf2zpk, zpk2sos, zpk2tf
@@ -55,7 +56,7 @@ def zDistortKernel(dt: float, params: Sequence[tuple]) -> np.ndarray:
     for tau, A in params:
         H += (1j * A * omega * tau) / (1j * omega * tau + 1)
 
-    ker = ifftshift(ifft(1 / H)).real
+    ker = cast(NDArray[np.complex128], ifftshift(ifft(1 / H))).real
     return ker
 
 
@@ -98,11 +99,14 @@ def exp_decay_filter_old(amp, tau, sample_rate):
     return b, a
 
 
-def exp_decay_filter(amp: float | Sequence[float],
-                     tau: float | Sequence[float],
-                     sample_rate: float,
-                     inv: bool = False,
-                     output='ba') -> tuple[np.ndarray, np.ndarray]:
+def exp_decay_filter(
+    amp: float | Sequence[float],
+    tau: float | Sequence[float],
+    sample_rate: float,
+    inv: bool = False,
+    output='ba'
+) -> NDArray[np.float64] | tuple[NDArray[np.float64], NDArray[
+        np.float64]] | tuple[NDArray[np.float64], NDArray[np.float64], float]:
     """
     exp decay filter
 
@@ -137,7 +141,9 @@ def exp_decay_filter(amp: float | Sequence[float],
 
     if isinstance(amp, (int, float, complex)):
         amp = [amp]
-        tau = [tau]
+        tau = [cast(float, tau)]
+    amp = cast(Sequence[float], amp)
+    tau = cast(Sequence[float], tau)
     numerator, denominator = np.poly1d([0.0]), np.poly1d([1.0])
     for i, (A, t) in enumerate(zip(amp, tau)):
         denominator = denominator * np.poly1d([1, -1 / t])
@@ -148,18 +154,22 @@ def exp_decay_filter(amp: float | Sequence[float],
         numerator = numerator + n
     numerator = numerator + denominator
 
-    z = np.exp(-numerator.roots / sample_rate)
-    p = np.exp(-denominator.roots / sample_rate)
+    z = cast(NDArray[np.float64], np.exp(-numerator.roots / sample_rate))
+    p = cast(NDArray[np.float64], np.exp(-denominator.roots / sample_rate))
     if inv:
         z, p = p, z
-    k = numerator(0) / denominator(0) * np.prod(1 - p) / np.prod(1 - z)
+    k = cast(float,
+             numerator(0) / denominator(0) * np.prod(1 - p) / np.prod(1 - z))
 
     if output == 'sos':
-        return zpk2sos(z, p, k)
+        return cast(NDArray[np.float64], zpk2sos(z, p, k))
     elif output == 'ba':
-        return zpk2tf(z, p, k)
+        return cast(tuple[NDArray[np.float64], NDArray[np.float64]],
+                    zpk2tf(z, p, k))
     elif output == 'zpk':
         return z, p, k
+    else:
+        raise ValueError(f"Invalid output type: {output}")
 
 
 def reflection_filter(f, A, tau):
@@ -251,7 +261,8 @@ def stable_filter(exp_decay_filters: list, sample_rate: float):
     """
     filters = []
     for amp, tau in exp_decay_filters:
-        a, b = exp_decay_filter(amp, tau, sample_rate)
+        a, b = cast(tuple[NDArray[np.float64], NDArray[np.float64]],
+                    exp_decay_filter(amp, tau, sample_rate))
         filters.append((b, a))
 
     b, a = combine_filters(filters)
@@ -262,14 +273,15 @@ def stable_filter(exp_decay_filters: list, sample_rate: float):
         return False
 
 
-def predistort(sig: np.ndarray,
-               filters: list = None,
-               ker: np.ndarray = None,
-               initial: float = 0.0,
-               initial_x: np.ndarray | None = None,
-               initial_y: np.ndarray | None = None,
-               zi: np.ndarray | None = None,
-               return_zf: bool = False) -> np.ndarray:
+def predistort(
+        sig: np.ndarray,
+        filters: list | None = None,
+        ker: np.ndarray | None = None,
+        initial: float = 0.0,
+        initial_x: np.ndarray | None = None,
+        initial_y: np.ndarray | None = None,
+        zi: np.ndarray | None = None,
+        return_zf: bool = False) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
     if filters is not None:
         b, a = combine_filters(filters)
         z, p, k = tf2zpk(b, a)
@@ -315,7 +327,8 @@ def predistort(sig: np.ndarray,
 def distort(points, params, sample_rate, initial=0.0):
     filters = []
     for amp, tau in np.asarray(params).reshape(-1, 2):
-        b, a = exp_decay_filter(amp, abs(tau), sample_rate)
+        b, a = cast(tuple[NDArray[np.float64], NDArray[np.float64]],
+                    exp_decay_filter(amp, abs(tau), sample_rate))
         filters.append((b, a))
     return predistort(points, filters, initial=initial)
 
@@ -342,6 +355,7 @@ def phase_curve(t, params, df_dphi, pulse_width, start, wav, sample_rate):
 
 if __name__ == '__main__':
     import matplotlib.pyplot as plt
+
     from waveforms import square
 
     data = np.load('Z_distortion.npz')
@@ -354,7 +368,7 @@ if __name__ == '__main__':
     wav = 0.1 * (square(2e-6) << 1e-6)
 
     def f(t, *params):
-        return phase_curve(t, params, df_dphi, 10e-9, 25e-9)
+        return phase_curve(t, params, df_dphi, 10e-9, 25e-9, wav, sample_rate)
 
     params = [-0.03, 0.1e-6, 0.02, 0.3e-6]
     popt, pcov = curve_fit(f, x, y, p0=params)

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms/multy_drag.py

@@ -230,19 +230,3 @@ def drag_sinx(freq,
                     bounds=(round(t0,
                                   NDIGITS), round(t0 + width + plateau,
                                                   NDIGITS), +inf))
-
-
-def _mollifier(x, width):
-    return np.exp(1 / ((x / width)**2 - 1))
-
-
-def _format_MOLLIFIER(shift, *args):
-    width, *_ = args
-    return f'\\exp\\frac{{1}}{{\\left(\\frac{{t-{shift:.{NDIGITS}f}}}{{{width:.{NDIGITS}f}}}\\right)^2-1}}'
-
-
-MOLLIFIER = registerBaseFunc(_mollifier)
-registerDerivative(
-    MOLLIFIER,
-    lambda x, width: 2 * x / width**2 * _mollifier(x, width),
-)

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms/version.py

@@ -1,2 +1,2 @@
 """Define version number here and read it from setup.py automatically"""
-__version__ = "2.2.0"
+__version__ = "2.2.2"

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms/waveform.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from fractions import Fraction
 from typing import Generator, Iterable, cast
 
@@ -6,12 +8,14 @@ from numpy import e, inf, pi
 from numpy.typing import NDArray
 from scipy.signal import sosfilt
 
-from ._waveform import (
-    _D, COS, COSH, DRAG, ERF, EXP, EXPONENTIALCHIRP, GAUSSIAN, HYPERBOLICCHIRP,
-    INTERP, LINEAR, LINEARCHIRP, MOLLIFIER, NDIGITS, SINC, SINH, _baseFunc,
-    _baseFunc_latex, _const, _half, _one, _zero, add, basic_wave, calc_parts,
-    filter, is_const, merge_waveform, mul, pow, registerBaseFunc,
-    registerBaseFuncLatex, registerDerivative, shift, simplify, wave_sum)
+from ._waveform import (_D, COS, COSH, D_GAUSSIAN, DRAG, ERF, EXP,
+                        EXPONENTIALCHIRP, GAUSSIAN, HYPERBOLICCHIRP, INTERP,
+                        LINEAR, LINEARCHIRP, MOLLIFIER, NDIGITS, SINC, SINH,
+                        _baseFunc, _baseFunc_latex, _const, _half, _one, _zero,
+                        add, basic_wave, calc_parts, filter, is_const,
+                        merge_waveform, mul, pow, registerBaseFunc,
+                        registerBaseFuncLatex, registerDerivative, shift,
+                        simplify, wave_sum)
 
 
 def _test_spec_num(num, spec):
@@ -131,7 +135,7 @@ class Waveform:
     @staticmethod
     def _begin(bounds, seq):
         for i, s in enumerate(seq):
-            if s is not _zero:
+            if s != _zero:
                 if i == 0:
                     return -inf
                 return bounds[i - 1]
@@ -141,7 +145,7 @@ class Waveform:
     def _end(bounds, seq):
         N = len(bounds)
         for i, s in enumerate(seq[::-1]):
-            if s is not _zero:
+            if s != _zero:
                 if i == 0:
                     return inf
                 return bounds[N - i - 1]
@@ -179,7 +183,8 @@ class Waveform:
             filters = self.filters
         if chunk_size is None:
             x = np.arange(self.start, self.stop, 1 / sample_rate)
-            sig = self.__call__(x, out=out, function_lib=function_lib)
+            sig = cast(np.ndarray,
+                       self.__call__(x, out=out, function_lib=function_lib))
             if filters is not None:
                 sos, initial = filters
                 if not isinstance(sos, np.ndarray):
@@ -393,22 +398,22 @@ class Waveform:
         return Waveform(*merge_waveform(self.bounds, self.seq, other.bounds,
                                         other.seq, oper))
 
-    def __pow__(self, n):
+    def __pow__(self, n) -> Waveform:
         return Waveform(self.bounds, tuple(pow(w, n) for w in self.seq))
 
-    def __add__(self, other):
+    def __add__(self, other) -> Waveform:
         if isinstance(other, Waveform):
             return self._comb(other, add)
         else:
             return self + const(other)
 
-    def __radd__(self, v):
+    def __radd__(self, v) -> Waveform:
         return const(v) + self
 
-    def __ior__(self, other):
+    def __ior__(self, other) -> Waveform:
         return self | other
 
-    def __or__(self, other):
+    def __or__(self, other) -> Waveform:
         if isinstance(other, (int, float, complex)):
             other = const(other)
         w = self.marker + other.marker
@@ -421,10 +426,10 @@ class Waveform:
 
         return self._comb(other, _or)
 
-    def __iand__(self, other):
+    def __iand__(self, other) -> Waveform:
         return self & other
 
-    def __and__(self, other):
+    def __and__(self, other) -> Waveform:
         if isinstance(other, (int, float, complex)):
             other = const(other)
         w = self.marker + other.marker
@@ -443,7 +448,7 @@ class Waveform:
         return Waveform(w.bounds,
                         tuple(_zero if s == _zero else _one for s in w.seq))
 
-    def mask(self, edge=0):
+    def mask(self, edge: float = 0) -> Waveform:
         w = self.marker
         in_wave = w.seq[0] == _zero
         bounds = []
@@ -471,31 +476,31 @@ class Waveform:
                     bounds.append(b)
         return Waveform(tuple(bounds), tuple(seq))
 
-    def __mul__(self, other):
+    def __mul__(self, other) -> Waveform:
         if isinstance(other, Waveform):
             return self._comb(other, mul)
         else:
             return self * const(other)
 
-    def __rmul__(self, v):
+    def __rmul__(self, v) -> Waveform:
         return const(v) * self
 
-    def __truediv__(self, other):
+    def __truediv__(self, other) -> Waveform:
         if isinstance(other, Waveform):
             raise TypeError('division by waveform')
         else:
             return self * const(1 / other)
 
-    def __neg__(self):
+    def __neg__(self) -> Waveform:
         return -1 * self
 
-    def __sub__(self, other):
+    def __sub__(self, other) -> Waveform:
         return self + (-other)
 
-    def __rsub__(self, v):
+    def __rsub__(self, v) -> Waveform:
         return v + (-self)
 
-    def __rshift__(self, time):
+    def __rshift__(self, time) -> Waveform:
         return Waveform(
             tuple(round(bound + time, NDIGITS) for bound in self.bounds),
             tuple(shift(expr, time) for expr in self.seq))
@@ -520,11 +525,12 @@ class Waveform:
         self,
         x,
         frag=False,
-        out: np.ndarray | None = None,
+        out: np.ndarray | list | None = None,
         accumulate=False,
         function_lib=None
-    ) -> NDArray[np.float64] | list[tuple[int, int,
-                                          NDArray[np.float64]]] | np.float64:
+    ) -> NDArray[np.float64 | np.complex128] | list[
+            tuple[int, int, NDArray[np.float64 | np.complex128]] | int
+            | float | complex] | np.float64:
         if function_lib is None:
             function_lib = _baseFunc
         if isinstance(x, (int, float, complex)):
@@ -541,8 +547,9 @@ class Waveform:
             self._fill_parts(parts, out)
         else:
             if out is None:
-                return parts
+                return cast(list, parts)
             else:
+                out = cast(list, out)
                 if not accumulate:
                     out.clear()
                     out.extend(parts)
@@ -666,7 +673,8 @@ class WaveVStack(Waveform):
 
     def __call__(self, x, frag=False, out=None, function_lib=None):
         assert frag is False, 'WaveVStack does not support frag mode'
-        out = np.full_like(x, self.offset, dtype=complex)
+        out = np.full_like(x, self.offset, dtype=np.complex128)
+        out = cast(NDArray[np.complex128], out)
         if self.shift != 0:
             x = x - self.shift
         if function_lib is None:
@@ -731,6 +739,8 @@ class WaveVStack(Waveform):
         wav.start = self.start
         wav.stop = self.stop
         wav.sample_rate = self.sample_rate
+        wav.filters = self.filters
+        wav.label = self.label
         return wav
 
     @staticmethod
@@ -749,9 +759,11 @@ class WaveVStack(Waveform):
         ret.stop = self.stop
         ret.shift = self.shift + time
         ret.offset = self.offset
+        ret.filters = self.filters
+        ret.label = self.label
         return ret
 
-    def __add__(self, other):
+    def __add__(self, other) -> WaveVStack:
         ret = WaveVStack()
         ret.wlist.extend(self.wlist)
         if isinstance(other, WaveVStack):
@@ -767,28 +779,34 @@ class WaveVStack(Waveform):
         else:
             # ret.wlist.append(((+inf, ), (_const(1.0 * other), )))
             ret.offset += other
+        ret.filters = self.filters
+        ret.label = self.label
         return ret
 
-    def __radd__(self, v):
+    def __radd__(self, v) -> WaveVStack:
         return self + v
 
-    def __mul__(self, other):
+    def __mul__(self, other) -> WaveVStack:
         if isinstance(other, Waveform):
             other = other.simplify() << self.shift
             ret = WaveVStack([Waveform(*w) * other for w in self.wlist])
             if self.offset != 0:
                 w = other * self.offset
                 ret.wlist.append((w.bounds, w.seq))
+            ret.filters = self.filters
+            ret.label = self.label
             return ret
         else:
             ret = WaveVStack([Waveform(*w) * other for w in self.wlist])
             ret.offset = self.offset * other
+            ret.filters = self.filters
+            ret.label = self.label
             return ret
 
-    def __rmul__(self, v):
+    def __rmul__(self, v) -> WaveVStack:
         return self * v
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         if self.wlist:
             return False
         else:
@@ -999,6 +1017,23 @@ def _format_MOLLIFIER(shift, *args):
         return f"\\mathrm{{Mollifier}}^{{({d})}}\\left(t{shift_str}, r={r}\\right)"
 
 
+def _format_D_GAUSSIAN(shift, *args):
+    sigma = _num_latex(args[0] / np.sqrt(2))
+    d = args[1]
+    shift_str = _num_latex(-shift)
+    if shift_str == '0':
+        shift_str = ''
+    elif shift_str[0] != '-':
+        shift_str = '+' + shift_str
+
+    if d == 0:
+        return f"\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
+    elif d == 1:
+        return f"\\frac{{\\mathrm{{d}}}}{{\\mathrm{{d}}t}}\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
+    else:
+        return f"\\frac{{\\mathrm{{d}}^{{{d}}}}}{{\\mathrm{{d}}t^{{{d}}}}}\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
+
+
 registerBaseFuncLatex(LINEAR, _format_LINEAR)
 registerBaseFuncLatex(GAUSSIAN, _format_GAUSSIAN)
 registerBaseFuncLatex(ERF, _format_ERF)
@@ -1009,6 +1044,7 @@ registerBaseFuncLatex(COSH, _format_COSH)
 registerBaseFuncLatex(SINH, _format_SINH)
 registerBaseFuncLatex(DRAG, _format_DRAG)
 registerBaseFuncLatex(MOLLIFIER, _format_MOLLIFIER)
+registerBaseFuncLatex(D_GAUSSIAN, _format_D_GAUSSIAN)
 
 
 def D(wav: Waveform, d: int = 1) -> Waveform:
@@ -1066,7 +1102,7 @@ def step(edge, type='erf'):
                         seq=(_zero, rise, _one))
 
 
-def square(width, edge=0, type='erf'):
+def square(width: float, edge: float = 0, type: str = 'erf') -> Waveform:
     if width <= 0:
         return zero()
     if edge == 0:
@@ -1079,7 +1115,9 @@ def square(width, edge=0, type='erf'):
                 (step(edge, type=type) >> width / 2))
 
 
-def gaussian(width, plateau=0.0):
+def gaussian(width: float,
+             plateau: float = 0.0,
+             d: int | None = None) -> Waveform:
     if width <= 0 and plateau <= 0.0:
         return zero()
     # width is two times FWHM
@@ -1087,26 +1125,27 @@ def gaussian(width, plateau=0.0):
     std_sq2 = width / 3.3302184446307908
     # std is set to give total pulse area same as a square
     # std_sq2 = width/np.sqrt(np.pi)
+    if d is None:
+        base = lambda shift: basic_wave(GAUSSIAN, std_sq2, shift=shift)
+    else:
+        base = lambda shift: basic_wave(D_GAUSSIAN, std_sq2, d, shift=shift)
+
     if round(0.5 * plateau, NDIGITS) <= 0.0:
         return Waveform(bounds=(round(-0.75 * width,
                                       NDIGITS), round(0.75 * width,
                                                       NDIGITS), +inf),
-                        seq=(_zero, basic_wave(GAUSSIAN, std_sq2), _zero))
+                        seq=(_zero, base(0), _zero))
     else:
         return Waveform(bounds=(round(-0.75 * width - 0.5 * plateau,
                                       NDIGITS), round(-0.5 * plateau, NDIGITS),
                                 round(0.5 * plateau, NDIGITS),
                                 round(0.75 * width + 0.5 * plateau,
                                       NDIGITS), +inf),
-                        seq=(_zero,
-                             basic_wave(GAUSSIAN,
-                                        std_sq2,
-                                        shift=-0.5 * plateau), _one,
-                             basic_wave(GAUSSIAN, std_sq2,
-                                        shift=0.5 * plateau), _zero))
+                        seq=(_zero, base(-0.5 * plateau), _one,
+                             base(0.5 * plateau), _zero))
 
 
-def cos(w, phi=0):
+def cos(w: float, phi: float = 0) -> Waveform:
     if w == 0:
         return const(np.cos(phi))
     if w < 0:
@@ -1115,7 +1154,7 @@ def cos(w, phi=0):
     return Waveform(seq=(basic_wave(COS, w, shift=-phi / w), ))
 
 
-def sin(w, phi=0):
+def sin(w: float, phi: float = 0) -> Waveform:
     if w == 0:
         return const(np.sin(phi))
     if w < 0:
@@ -1124,7 +1163,7 @@ def sin(w, phi=0):
     return Waveform(seq=(basic_wave(COS, w, shift=(pi / 2 - phi) / w), ))
 
 
-def exp(alpha):
+def exp(alpha: float | complex) -> Waveform:
     if isinstance(alpha, complex):
         if alpha.real == 0:
             return cos(alpha.imag) + 1j * sin(alpha.imag)
@@ -1134,7 +1173,7 @@ def exp(alpha):
         return Waveform(seq=(basic_wave(EXP, alpha), ))
 
 
-def sinc(bw):
+def sinc(bw: float) -> Waveform:
     if bw <= 0:
         return zero()
     width = 100 / bw
@@ -1143,7 +1182,7 @@ def sinc(bw):
                     seq=(_zero, basic_wave(SINC, bw), _zero))
 
 
-def cosPulse(width, plateau=0.0):
+def cosPulse(width: float, plateau: float = 0.0) -> Waveform:
     # cos = basic_wave(COS, 2*np.pi/width)
     # pulse = mul(add(cos, _one), _half)
     if round(0.5 * plateau, NDIGITS) > 0:
@@ -1157,19 +1196,21 @@ def cosPulse(width, plateau=0.0):
                     seq=(_zero, pulse, _zero))
 
 
-def hanning(width, plateau=0.0):
+def hanning(width: float, plateau: float = 0.0) -> Waveform:
     return cosPulse(width, plateau=plateau)
 
 
-def cosh(w):
+def cosh(w: float) -> Waveform:
     return Waveform(seq=(basic_wave(COSH, w), ))
 
 
-def sinh(w):
+def sinh(w: float) -> Waveform:
     return Waveform(seq=(basic_wave(SINH, w), ))
 
 
-def coshPulse(width, eps=1.0, plateau=0.0):
+def coshPulse(width: float,
+              eps: float = 1.0,
+              plateau: float = 0.0) -> Waveform:
     """Cosine hyperbolic pulse with the following im
 
     pulse edge shape:
@@ -1219,25 +1260,25 @@ def coshPulse(width, eps=1.0, plateau=0.0):
                         seq=(_zero, raising, _one, falling, _zero))
 
 
-def general_cosine(duration, *arg):
+def general_cosine(duration: float, *arg: float) -> Waveform:
     wav = zero()
-    arg = np.asarray(arg)
-    arg /= arg[::2].sum()
-    for i, a in enumerate(arg, start=1):
+    arg_ = np.asarray(arg)
+    arg_ /= arg_[::2].sum()
+    for i, a in enumerate(arg_, start=1):
         wav += a / 2 * (1 - (-1)**i * cos(i * 2 * pi / duration))
     return wav * square(duration)
 
 
-def slepian(duration, *arg):
+def slepian(duration: float, *arg: float) -> Waveform:
     wav = zero()
-    arg = np.asarray(arg)
-    arg /= arg[::2].sum()
-    for i, a in enumerate(arg, start=1):
+    arg_ = np.asarray(arg)
+    arg_ /= arg_[::2].sum()
+    for i, a in enumerate(arg_, start=1):
         wav += a / 2 * (1 - (-1)**i * cos(i * 2 * pi / duration))
     return wav * square(duration)
 
 
-def mollifier(width, plateau: float = 0.0, d: int = 0):
+def mollifier(width: float, plateau: float = 0.0, d: int = 0) -> Waveform:
     """
     Mollifier function is a smooth function that is 1 at the origin and 0 outside a certain radius.
     It is defined as:
@@ -1298,7 +1339,13 @@ def t():
     return Waveform(seq=((((LINEAR, 0), ), (1, )), (1, )))
 
 
-def drag(freq, width, plateau=0, delta=0, block_freq=None, phase=0, t0=0):
+def drag(freq: float,
+         width: float,
+         plateau: float = 0,
+         delta: float = 0,
+         block_freq: float | None = None,
+         phase: float = 0,
+         t0: float = 0) -> Waveform:
     phase += pi * delta * (width + plateau)
     if plateau <= 0:
         return Waveform(seq=(_zero,
@@ -1327,7 +1374,11 @@ def drag(freq, width, plateau=0, delta=0, block_freq=None, phase=0, t0=0):
                                           NDIGITS), +inf))
 
 
-def chirp(f0, f1, T, phi0=0, type='linear'):
+def chirp(f0: float,
+          f1: float,
+          T: float,
+          phi0: float = 0,
+          type: str = 'linear') -> Waveform:
     """
     A chirp is a signal in which the frequency increases (up-chirp)
     or decreases (down-chirp) with time. In some sources, the term
@@ -1366,7 +1417,7 @@ def chirp(f0, f1, T, phi0=0, type='linear'):
         raise ValueError(f'unknown type {type}')
 
 
-def interp(x, y):
+def interp(x: NDArray[np.float64], y: NDArray[np.float64]) -> Waveform:
     seq, bounds = [_zero], [x[0]]
     for x1, x2, y1, y2 in zip(x[:-1], x[1:], y[:-1], y[1:]):
         if x2 == x1:
@@ -1384,12 +1435,20 @@ def interp(x, y):
                                  for b in bounds)).simplify()
 
 
-def cut(wav, start=None, stop=None, head=None, tail=None, min=None, max=None):
+def cut(wav: Waveform,
+        start: float | None = None,
+        stop: float | None = None,
+        head: float | None = None,
+        tail: float | None = None,
+        min: float | None = None,
+        max: float | None = None) -> Waveform:
     offset = 0
     if start is not None and head is not None:
-        offset = head - wav(np.array([1.0 * start]))[0]
+        offset = head - cast(NDArray[np.float64], wav(np.array([1.0 * start
+                                                                ])))[0]
     elif stop is not None and tail is not None:
-        offset = tail - wav(np.array([1.0 * stop]))[0]
+        offset = tail - cast(NDArray[np.float64], wav(np.array([1.0 * stop
+                                                                ])))[0]
     wav = wav + offset
 
     if start is not None:
@@ -1420,15 +1479,15 @@ def samplingPoints(start, stop, points):
                                            tuple(points)), _zero))
 
 
-def mixing(I,
-           Q=None,
+def mixing(I: Waveform,
+           Q: Waveform | None = None,
            *,
-           phase=0.0,
-           freq=0.0,
-           ratioIQ=1.0,
-           phaseDiff=0.0,
-           block_freq=None,
-           DRAGScaling=None):
+           phase: float = 0.0,
+           freq: float = 0.0,
+           ratioIQ: float = 1.0,
+           phaseDiff: float = 0.0,
+           block_freq: float | None = None,
+           DRAGScaling: float | None = None) -> tuple[Waveform, Waveform]:
     """SSB or envelope mixing
     """
     if Q is None:
@@ -1442,8 +1501,8 @@ def mixing(I,
         Qout = -I * sin(w, -phase + phaseDiff) + Q * cos(w, -phase + phaseDiff)
     else:
         # envelope mixing
-        Iout = I * np.cos(-phase) + Q * np.sin(-phase)
-        Qout = -I * np.sin(-phase) + Q * np.cos(-phase)
+        Iout = cast(Waveform, I * np.cos(-phase) + Q * np.sin(-phase))
+        Qout = cast(Waveform, -I * np.sin(-phase) + Q * np.cos(-phase))
 
     # apply DRAG
     if block_freq is not None and block_freq != freq:

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms/waveform_parser.py

@@ -4,7 +4,7 @@ from ast import literal_eval
 from functools import lru_cache
 from pathlib import Path
 
-from antlr4 import *
+from antlr4 import CommonTokenStream, InputStream
 from antlr4.error.ErrorListener import ErrorListener
 
 from . import multy_drag, waveform

{waveforms-2.2.0 → waveforms-2.2.2}/waveforms.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waveforms
-Version: 2.2.0
+Version: 2.2.2
 Summary: Edit waveforms used in experiment
 Author-email: feihoo87 <feihoo87@gmail.com>
 Maintainer-email: feihoo87 <feihoo87@gmail.com>