waveforms 2.1.1__tar.gz → 2.2.1__tar.gz

{waveforms-2.1.1 → waveforms-2.2.1}/PKG-INFO

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

{waveforms-2.1.1 → waveforms-2.2.1}/waveforms/__init__.py

@@ -4,7 +4,8 @@ from .multy_drag import drag_sin, drag_sinx
 from .version import __version__
 from .waveform import (D, Waveform, WaveVStack, chirp, const, cos, cosh,
                        coshPulse, cosPulse, cut, drag, exp, function, gaussian,
-                       general_cosine, hanning, interp, mixing, one, poly,
-                       registerBaseFunc, registerDerivative, samplingPoints,
-                       sign, sin, sinc, sinh, square, step, t, zero)
+                       general_cosine, hanning, interp, mixing, mollifier, one,
+                       poly, registerBaseFunc, registerDerivative,
+                       samplingPoints, sign, sin, sinc, sinh, square, step, t,
+                       zero)
 from .waveform_parser import wave_eval

{waveforms-2.1.1 → waveforms-2.2.1}/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 = ...
@@ -28,7 +29,7 @@ def is_const(x: tuple[tuple, tuple]) -> bool:
     pass
 
 
-def basic_wave(Type, *args, shift=0) -> tuple[tuple, tuple]:
+def basic_wave(Type: int, *args, shift: float = 0) -> tuple[tuple, tuple]:
     pass
 
 
@@ -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[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
 
 
@@ -103,6 +107,8 @@ HYPERBOLICCHIRP: int = ...
 COSH: int = ...
 SINH: int = ...
 DRAG: int = ...
+MOLLIFIER: int = ...
+D_GAUSSIAN: int = ...
 
 
 def simplify(expr: tuple[tuple, tuple], eps: float) -> tuple[tuple, tuple]:

{waveforms-2.1.1 → waveforms-2.2.1}/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)
 
@@ -351,6 +356,19 @@ def _drag(t: np.ndarray, t0: float, freq: float, width: float, delta: float,
     return Omega_x * np.cos(wt) + Omega_y * np.sin(wt)
 
 
+def _mollifier(t: np.ndarray, r: float, d: int):
+    x = t / r
+    if d == 0:
+        return np.exp(1 / (np.abs(x)**2 - 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
+
+
 LINEAR = registerBaseFunc(_LINEAR)
 GAUSSIAN = registerBaseFunc(_GAUSSIAN)
 ERF = registerBaseFunc(_ERF)
@@ -364,6 +382,8 @@ HYPERBOLICCHIRP = registerBaseFunc(_HYPERBOLICCHIRP)
 COSH = registerBaseFunc(_COSH)
 SINH = registerBaseFunc(_SINH)
 DRAG = registerBaseFunc(_drag)
+MOLLIFIER = registerBaseFunc(_mollifier)
+D_GAUSSIAN = registerBaseFunc(_D_GAUSSIAN)
 
 
 def _d_LINEAR(shift, *args):
@@ -433,6 +453,14 @@ def _d_HYPERBOLICCHIRP(shift, f0, k, phi0):
                                          shift)), (-1, 1)), ), (2 * pi * f0, ))
 
 
+def _d_MOLLIFIER(shift, r, d):
+    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
 registerDerivative(LINEAR, _d_LINEAR)
 registerDerivative(GAUSSIAN, _d_GAUSSIAN)
@@ -446,6 +474,8 @@ registerDerivative(SINH, _d_SINH)
 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.1.1 → waveforms-2.2.1}/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.1.1 → waveforms-2.2.1}/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.1.1 → waveforms-2.2.1}/waveforms/version.py

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

{waveforms-2.1.1 → waveforms-2.2.1}/waveforms/waveform.py

@@ -1,16 +1,21 @@
+from __future__ import annotations
+
 from fractions import Fraction
+from typing import Generator, Iterable, cast
 
 import numpy as np
 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,
-                        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):
@@ -124,13 +129,13 @@ class Waveform:
         self.start = None
         self.stop = None
         self.sample_rate = None
-        self.filters = None
+        self.filters: tuple[np.ndarray, float] | None = None
         self.label = None
 
     @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]
@@ -140,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]
@@ -160,12 +165,14 @@ class Waveform:
         else:
             return min(self.stop, self._end(self.bounds, self.seq))
 
-    def sample(self,
-               sample_rate=None,
-               out=None,
-               chunk_size=None,
-               function_lib=None,
-               filters=None):
+    def sample(
+        self,
+        sample_rate=None,
+        out: np.ndarray | None = None,
+        chunk_size=None,
+        function_lib=None,
+        filters: tuple[np.ndarray, float] | None = None
+    ) -> np.ndarray | Iterable[np.ndarray]:
         if sample_rate is None:
             sample_rate = self.sample_rate
         if self.start is None or self.stop is None or sample_rate is None:
@@ -176,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):
@@ -184,27 +192,34 @@ class Waveform:
                 elif not sos.flags.writeable:
                     sos = sos.copy()
                 if initial:
-                    sig = sosfilt(sos, sig - initial) + initial
+                    sig = cast(np.ndarray, sosfilt(sos,
+                                                   sig - initial)) + initial
                 else:
-                    sig = sosfilt(sos, sig)
-            return sig
+                    sig = cast(np.ndarray, sosfilt(sos, sig))
+            return cast(np.ndarray, sig)
         else:
             return self._sample_iter(sample_rate, chunk_size, out,
                                      function_lib, filters)
 
-    def _sample_iter(self, sample_rate, chunk_size, out, function_lib,
-                     filters):
-        start = self.start
+    def _sample_iter(
+        self, sample_rate, chunk_size, out: np.ndarray | None, function_lib,
+        filters: tuple[np.ndarray, float] | None
+    ) -> Generator[np.ndarray, None, None]:
+        start = cast(float, self.start)
         start_n = 0
         if filters is not None:
             sos, initial = filters
+            if not isinstance(sos, np.ndarray):
+                sos = np.array(sos)
+            elif not sos.flags.writeable:
+                sos = sos.copy()
             # zi = sosfilt_zi(sos)
             zi = np.zeros((sos.shape[0], 2))
         length = chunk_size / sample_rate
-        while start < self.stop:
-            if start + length > self.stop:
-                length = self.stop - start
-                stop = self.stop
+        while start < cast(float, self.stop):
+            if start + length > cast(float, self.stop):
+                length = cast(float, self.stop) - start
+                stop = cast(float, self.stop)
                 size = round((stop - start) * sample_rate)
             else:
                 stop = start + length
@@ -213,13 +228,17 @@ class Waveform:
 
             if filters is None:
                 if out is not None:
-                    yield self.__call__(x,
-                                        out=out[start_n:],
-                                        function_lib=function_lib)
+                    yield cast(
+                        np.ndarray,
+                        self.__call__(x,
+                                      out=out[start_n:],
+                                      function_lib=function_lib))
                 else:
-                    yield self.__call__(x, function_lib=function_lib)
+                    yield cast(np.ndarray,
+                               self.__call__(x, function_lib=function_lib))
             else:
-                sig = self.__call__(x, function_lib=function_lib)
+                sig = cast(np.ndarray,
+                           self.__call__(x, function_lib=function_lib))
                 if initial:
                     sig -= initial
                 sig, zi = sosfilt(sos, sig, zi=zi)
@@ -227,7 +246,7 @@ class Waveform:
                     sig += initial
                 if out is not None:
                     out[start_n:start_n + size] = sig
-                yield sig
+                yield cast(np.ndarray, sig)
 
             start = stop
             start_n += chunk_size
@@ -379,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
@@ -407,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
@@ -429,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 = []
@@ -457,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))
@@ -502,16 +521,22 @@ class Waveform:
         for start, stop, part in parts:
             out[start:stop] += part
 
-    def __call__(self,
-                 x,
-                 frag=False,
-                 out=None,
-                 accumulate=False,
-                 function_lib=None):
+    def __call__(
+        self,
+        x,
+        frag=False,
+        out: np.ndarray | list | None = None,
+        accumulate=False,
+        function_lib=None
+    ) -> 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)):
-            return self.__call__(np.array([x]), function_lib=function_lib)[0]
+            return cast(
+                NDArray[np.float64],
+                self.__call__(np.array([x]), function_lib=function_lib))[0]
         parts, dtype = calc_parts(self.bounds, self.seq, x, function_lib,
                                   self.min, self.max)
         if not frag:
@@ -522,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)
@@ -647,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:
@@ -732,7 +759,7 @@ class WaveVStack(Waveform):
         ret.offset = self.offset
         return ret
 
-    def __add__(self, other):
+    def __add__(self, other) -> WaveVStack:
         ret = WaveVStack()
         ret.wlist.extend(self.wlist)
         if isinstance(other, WaveVStack):
@@ -750,10 +777,10 @@ class WaveVStack(Waveform):
             ret.offset += other
         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])
@@ -766,10 +793,10 @@ class WaveVStack(Waveform):
             ret.offset = self.offset * other
             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:
@@ -961,6 +988,42 @@ def _format_DRAG(shift, *args):
     return f"DRAG(...)"
 
 
+def _format_MOLLIFIER(shift, *args):
+    r = _num_latex(args[0])
+    d = _num_latex(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{{Mollifier}}\\left(t{shift_str}, r={r}\\right)"
+    elif d == '1':
+        return f"\\mathrm{{Mollifier}}'\\left(t{shift_str}, r={r}\\right)"
+    elif d == '2':
+        return f"\\mathrm{{Mollifier}}''\\left(t{shift_str}, r={r}\\right)"
+    else:
+        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)
@@ -970,12 +1033,27 @@ registerBaseFuncLatex(EXP, _format_EXP)
 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):
+def D(wav: Waveform, d: int = 1) -> Waveform:
     """derivative
+
+    Parameters
+    ----------
+    wav : Waveform
+        The waveform to take the derivative of.
+    d : int, optional
+        The order of the derivative, by default 1.
     """
-    return Waveform(bounds=wav.bounds, seq=tuple(_D(x) for x in wav.seq))
+    assert d >= 0 and isinstance(d, int), "d must be a non-negative integer"
+    if d == 0:
+        return wav
+    elif d == 1:
+        return Waveform(bounds=wav.bounds, seq=tuple(_D(x) for x in wav.seq))
+    else:
+        return D(D(wav, d - 1), 1)
 
 
 def convolve(a, b):
@@ -1014,7 +1092,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:
@@ -1027,7 +1105,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
@@ -1035,26 +1115,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:
@@ -1063,7 +1144,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:
@@ -1072,7 +1153,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)
@@ -1082,7 +1163,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
@@ -1091,7 +1172,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:
@@ -1105,19 +1186,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:
@@ -1167,24 +1250,58 @@ 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: 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:
+
+    f(x) = exp(1 / ((x / r) ^ 2 - 1) + 1)  in case |x| < r
+         = 0                           in case |x| >= r
+    where r = width / 2 is the radius of the mollifier.
+
+    The parameter plateau is the width of the plateau.
+    The parameter d is the order of the derivative.
+    """
+    assert d >= 0 and isinstance(d, int), "d must be a non-negative integer"
+    assert width > 0, "width must be positive"
+
+    if plateau <= 0:
+        return Waveform(bounds=(-0.5 * width, 0.5 * width, inf),
+                        seq=(_zero, basic_wave(MOLLIFIER, width / 2,
+                                               d), _zero))
+    else:
+        return Waveform(bounds=(-0.5 * width - 0.5 * plateau, -0.5 * plateau,
+                                0.5 * plateau, 0.5 * width + 0.5 * plateau,
+                                inf),
+                        seq=(_zero,
+                             basic_wave(MOLLIFIER,
+                                        width / 2,
+                                        d,
+                                        shift=-0.5 * plateau), _one,
+                             basic_wave(MOLLIFIER,
+                                        width / 2,
+                                        d,
+                                        shift=0.5 * plateau), _zero))
+
+
 def _poly(*a):
     """
     a[0] + a[1] * t + a[2] * t**2 + ...
@@ -1212,7 +1329,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,
@@ -1241,7 +1364,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
@@ -1280,7 +1407,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:
@@ -1298,12 +1425,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:
@@ -1334,15 +1469,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:
@@ -1356,8 +1491,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:
@@ -1380,7 +1515,7 @@ def mixing(I,
 __all__ = [
     'D', 'Waveform', 'chirp', 'const', 'cos', 'cosh', 'coshPulse', 'cosPulse',
     'cut', 'drag', 'exp', 'function', 'gaussian', 'general_cosine', 'hanning',
-    'interp', 'mixing', 'one', 'poly', 'registerBaseFunc',
+    'interp', 'mixing', 'mollifier', 'one', 'poly', 'registerBaseFunc',
     'registerDerivative', 'samplingPoints', 'sign', 'sin', 'sinc', 'sinh',
     'square', 'step', 't', 'zero'
 ]

{waveforms-2.1.1 → waveforms-2.2.1}/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
@@ -30,9 +30,9 @@ class WaveformVisitor:
         self.functions = [
             'D', 'chirp', 'const', 'cos', 'cosh', 'coshPulse', 'cosPulse',
             'cut', 'drag', 'drag_sin', 'drag_sinx', 'exp', 'gaussian',
-            'general_cosine', 'hanning', 'interp', 'mixing', 'one', 'poly',
-            'samplingPoints', 'sign', 'sin', 'sinc', 'sinh', 'square', 'step',
-            't', 'zero'
+            'general_cosine', 'hanning', 'interp', 'mixing', 'mollifier',
+            'one', 'poly', 'samplingPoints', 'sign', 'sin', 'sinc', 'sinh',
+            'square', 'step', 't', 'zero'
         ]
         self.constants = {
             'pi': waveform.pi,
@@ -226,14 +226,13 @@ def _generate_antlr_parser():
 
     # Generate ANTLR files
     try:
-        result = subprocess.run([
-            "antlr4", "-Dlanguage=Python3",
-            str(grammar_file)
-        ],
-                                cwd=str(current_dir),
-                                capture_output=True,
-                                text=True,
-                                check=True)
+        result = subprocess.run(
+            ["antlr4", "-Dlanguage=Python3",
+             str(grammar_file)],
+            cwd=str(current_dir),
+            capture_output=True,
+            text=True,
+            check=True)
     except (subprocess.CalledProcessError, FileNotFoundError) as e:
         # Fall back to java command if antlr4 command is not available
         try:
@@ -258,7 +257,7 @@ def parse_waveform_expression(expr: str) -> waveform.Waveform:
     try:
         # Generate parser files if they don't exist
         # _generate_antlr_parser()
-        
+
         # Import generated ANTLR classes
         from .WaveformLexer import WaveformLexer
         from .WaveformParser import WaveformParser

{waveforms-2.1.1 → waveforms-2.2.1}/waveforms.egg-info/PKG-INFO

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