waveforms 2.1.0__tar.gz → 2.2.0__tar.gz

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

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

{waveforms-2.1.0 → waveforms-2.2.0}/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.0 → waveforms-2.2.0}/waveforms/_waveform.pyi

@@ -28,7 +28,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
 
 
@@ -53,7 +53,7 @@ def calc_parts(bounds: tuple,
                x: np.ndarray,
                function_lib: dict,
                min=-inf,
-               max=inf) -> tuple[list[np.ndarray], type]:
+               max=inf) -> tuple[list[tuple[int, int, np.ndarray]], type]:
     pass
 
 
@@ -103,6 +103,7 @@ HYPERBOLICCHIRP: int = ...
 COSH: int = ...
 SINH: int = ...
 DRAG: int = ...
+MOLLIFIER: int = ...
 
 
 def simplify(expr: tuple[tuple, tuple], eps: float) -> tuple[tuple, tuple]:

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

@@ -351,6 +351,18 @@ 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 +376,7 @@ HYPERBOLICCHIRP = registerBaseFunc(_HYPERBOLICCHIRP)
 COSH = registerBaseFunc(_COSH)
 SINH = registerBaseFunc(_SINH)
 DRAG = registerBaseFunc(_drag)
+MOLLIFIER = registerBaseFunc(_mollifier)
 
 
 def _d_LINEAR(shift, *args):
@@ -433,6 +446,10 @@ 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, ))
+
+
 # register derivative
 registerDerivative(LINEAR, _d_LINEAR)
 registerDerivative(GAUSSIAN, _d_GAUSSIAN)
@@ -446,6 +463,7 @@ registerDerivative(SINH, _d_SINH)
 registerDerivative(LINEARCHIRP, _d_LINEARCHIRP)
 registerDerivative(EXPONENTIALCHIRP, _d_EXPONENTIALCHIRP)
 registerDerivative(HYPERBOLICCHIRP, _d_HYPERBOLICCHIRP)
+registerDerivative(MOLLIFIER, _d_MOLLIFIER)
 
 
 def _cos_power_n(x, n):

{waveforms-2.1.0 → waveforms-2.2.0}/waveforms/utils.py

@@ -1,6 +1,6 @@
 from itertools import repeat
 from types import MappingProxyType
-from typing import Optional, Sequence
+from typing import Optional, Sequence, cast
 
 import numpy as np
 import scipy.sparse as sp
@@ -20,13 +20,13 @@ def freeze(x):
     elif isinstance(x, (np.ndarray, np.matrix)):
         x.flags.writeable = False
     elif isinstance(x, sp.spmatrix):
-        x.data.flags.writeable = False
-        if x.format in {'csr', 'csc', 'bsr'}:
-            x.indices.flags.writeable = False
-            x.indptr.flags.writeable = False
-        elif x.format == 'coo':
-            x.row.flags.writeable = False
-            x.col.flags.writeable = False
+        cast(np.ndarray, getattr(x, 'data')).flags.writeable = False
+        if getattr(x, 'format') in {'csr', 'csc', 'bsr'}:
+            cast(np.ndarray, getattr(x, 'indices')).flags.writeable = False
+            cast(np.ndarray, getattr(x, 'indptr')).flags.writeable = False
+        elif getattr(x, 'format') == 'coo':
+            cast(np.ndarray, getattr(x, 'row')).flags.writeable = False
+            cast(np.ndarray, getattr(x, 'col')).flags.writeable = False
     elif isinstance(x, bytearray):
         x = bytes(x)
     return x
@@ -72,12 +72,14 @@ def getFTMatrix(fList: Sequence[float],
     if weight is None or len(weight) == 0:
         weight = np.full(numOfPoints, 2 / numOfPoints)
     if phaseList is None or len(phaseList) == 0:
-        phaseList = np.zeros_like(fList)
+        phase_list = np.zeros_like(fList)
+    else:
+        phase_list = phaseList
     if weight.ndim == 1:
-        weightList = repeat(weight)
+        weight_list = repeat(weight)
     else:
-        weightList = weight
-    for f, phase, weight in zip(fList, phaseList, weightList):
+        weight_list = weight
+    for f, phase, weight in zip(fList, phase_list, weight_list):
         e.append(weight * np.exp(-1j * (2 * np.pi * f * t + phase)))
     return np.asarray(e).T
 

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

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

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

@@ -1,16 +1,17 @@
 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, 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,7 +125,7 @@ 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
@@ -160,12 +161,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:
@@ -179,28 +182,39 @@ class Waveform:
             sig = self.__call__(x, out=out, function_lib=function_lib)
             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()
                 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
@@ -209,13 +223,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)
@@ -223,7 +241,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
@@ -498,16 +516,21 @@ 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 | None = None,
+        accumulate=False,
+        function_lib=None
+    ) -> NDArray[np.float64] | list[tuple[int, int,
+                                          NDArray[np.float64]]] | 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:
@@ -957,6 +980,25 @@ 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)"
+
+
 registerBaseFuncLatex(LINEAR, _format_LINEAR)
 registerBaseFuncLatex(GAUSSIAN, _format_GAUSSIAN)
 registerBaseFuncLatex(ERF, _format_ERF)
@@ -966,12 +1008,26 @@ registerBaseFuncLatex(EXP, _format_EXP)
 registerBaseFuncLatex(COSH, _format_COSH)
 registerBaseFuncLatex(SINH, _format_SINH)
 registerBaseFuncLatex(DRAG, _format_DRAG)
+registerBaseFuncLatex(MOLLIFIER, _format_MOLLIFIER)
 
 
-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):
@@ -1181,6 +1237,40 @@ def slepian(duration, *arg):
     return wav * square(duration)
 
 
+def mollifier(width, plateau: float = 0.0, d: int = 0):
+    """
+    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 + ...
@@ -1376,7 +1466,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.0 → waveforms-2.2.0}/waveforms/waveform_parser.py

@@ -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.0 → waveforms-2.2.0}/waveforms.egg-info/PKG-INFO

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