waveforms 3.0.0__tar.gz → 3.1.0__tar.gz

{waveforms-3.0.0/waveforms.egg-info → waveforms-3.1.0}/PKG-INFO

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

{waveforms-3.0.0 → waveforms-3.1.0}/tests/test_waveform.py

@@ -13,7 +13,8 @@ import waveforms as wf
 
 PUBLIC_NAMES = {
     "D", "Waveform", "WaveVStack", "chirp", "const", "cos", "cosh",
-    "coshPulse", "cosPulse", "cut", "drag", "exp", "function",
+    "coshPulse", "cosPulse", "cut", "drag", "drag_sin", "drag_sinx",
+    "exp", "function",
     "gaussian", "general_cosine", "get_time_resolution", "hanning", "interp", "mixing",
     "mollifier", "one", "poly", "registerBaseFunc", "registerDerivative",
     "samplingPoints", "set_time_resolution", "sign", "sin", "sinc", "sinh", "square", "step",
@@ -134,6 +135,96 @@ def test_sampling_points_chirps_and_mixing():
     assert np.all(np.isfinite(q(x)))
 
 
+@pytest.mark.parametrize("constructor,extra", [
+    (wf.drag_sin, {}),
+    (wf.drag_sinx, {"tab": 0.47}),
+])
+def test_multi_frequency_drag_roundtrip_shift_and_parser(constructor, extra):
+    parameters = dict(
+        freq=5e9,
+        width=22.22e-9,
+        plateau=3e-9,
+        delta=-13.7e6,
+        block_freq=(-91e6, 37e6, 124e6),
+        phase=0.31,
+        t0=7e-9,
+        **extra,
+    )
+    wav = constructor(**parameters)
+    x = np.linspace(-2e-9, 40e-9, 8193)
+    values = wav(x)
+    assert np.all(np.isfinite(values))
+    assert np.all(values[(x < parameters["t0"])
+                         | (x >= parameters["t0"] + parameters["width"]
+                            + parameters["plateau"])] == 0)
+
+    restored = wf.Waveform.from_bytes(wav.to_bytes())
+    assert restored == wav
+    assert np.array_equal(restored(x), values)
+    delay = 11e-9
+    assert np.allclose((wav >> delay)(x + delay), values)
+
+    expression = (
+        f"{constructor.__name__}(5e9, 22.22e-9, plateau=3e-9, "
+        "delta=-13.7e6, block_freq=(-91e6, 37e6, 124e6), "
+        f"phase=0.31, t0=7e-9{', tab=0.47' if extra else ''})"
+    )
+    assert wf.wave_eval(expression) == wav
+
+
+def test_drag_sin_scalar_block_frequency_and_argument_validation():
+    scalar = wf.drag_sin(5e9, 20e-9, block_freq=80e6)
+    sequence = wf.drag_sin(5e9, 20e-9, block_freq=(80e6,))
+    assert scalar.to_bytes() == sequence.to_bytes()
+    with pytest.raises(ValueError):
+        wf.drag_sin(5e9, 0)
+    with pytest.raises(ValueError):
+        wf.drag_sinx(5e9, 20e-9, tab=0)
+
+
+def test_multi_frequency_drag_matches_v2_numerical_behavior():
+    parameters = dict(
+        freq=5e9,
+        width=22.22e-9,
+        plateau=3e-9,
+        delta=-13.7e6,
+        block_freq=(-91e6, 37e6, 124e6),
+        phase=0.31,
+        t0=7e-9,
+    )
+    fractions = np.array([0, 0.07, 0.19, 0.37, 0.51,
+                          0.68, 0.83, 0.96, 0.999])
+    x = parameters["t0"] + fractions * (
+        parameters["width"] + parameters["plateau"]
+    )
+    expected_sin = np.array([
+        0,
+        0.20532545807257174,
+        0.6208384978963238,
+        -0.4333357461117501,
+        -0.08840393960186489,
+        -0.8620379323126761,
+        -0.5748645543753418,
+        -0.07489572316435908,
+        -0.00309835789254227,
+    ])
+    expected_sinx = np.array([
+        0,
+        0.11971095177791738,
+        0.36196762048509773,
+        -0.03689065391238458,
+        -0.05154217041569324,
+        -0.19970252200405683,
+        -0.3351634210081106,
+        -0.04366647169944161,
+        -0.00180643635595232,
+    ])
+    assert np.allclose(wf.drag_sin(**parameters)(x), expected_sin,
+                       rtol=2e-13, atol=2e-13)
+    assert np.allclose(wf.drag_sinx(**parameters, tab=0.47)(x), expected_sinx,
+                       rtol=2e-13, atol=2e-13)
+
+
 def test_filters_and_chunked_sampling():
     sample_rate = 1000
     b, a = butter(3, 4.0, "lowpass", fs=sample_rate)

{waveforms-3.0.0 → waveforms-3.1.0}/waveforms/__init__.py

@@ -13,6 +13,8 @@ from .waveform import (
     cosPulse,
     cut,
     drag,
+    drag_sin,
+    drag_sinx,
     exp,
     function,
     gaussian,
@@ -41,7 +43,8 @@ from .waveform import (
 
 __all__ = [
     "D", "Waveform", "WaveVStack", "chirp", "const", "cos", "cosh",
-    "coshPulse", "cosPulse", "cut", "drag", "e", "exp", "function",
+    "coshPulse", "cosPulse", "cut", "drag", "drag_sin", "drag_sinx",
+    "e", "exp", "function",
     "gaussian", "general_cosine", "get_time_resolution", "hanning",
     "interp", "mixing", "mollifier", "one", "pi", "poly",
     "registerBaseFunc", "registerDerivative", "samplingPoints",

{waveforms-3.0.0 → waveforms-3.1.0}/waveforms/_waveform.pyi

@@ -18,6 +18,8 @@ SINH: int
 DRAG: int
 MOLLIFIER: int
 D_GAUSSIAN: int
+DRAG_SIN: int
+DRAG_SINX: int
 
 
 class PackedWaveform:

{waveforms-3.0.0 → waveforms-3.1.0}/waveforms/_waveform.pyx

@@ -12,7 +12,7 @@ import struct
 from bisect import bisect_left
 from functools import lru_cache
 from itertools import chain, product
-from math import comb
+from math import comb, factorial
 
 import numpy as np
 import scipy.special as special
@@ -36,6 +36,8 @@ SINH = 12
 DRAG = 13
 MOLLIFIER = 14
 D_GAUSSIAN = 15
+DRAG_SIN = 16
+DRAG_SINX = 17
 
 _zero = ((), ())
 _one = ((((), ()),), (1.0,))
@@ -169,6 +171,179 @@ def _drag(t, t0, freq, width, delta, block_freq, phase):
     return omega_x * np.cos(wt) + omega_y * np.sin(wt)
 
 
+def _b_series_matrix(block_coefficients):
+    matrix = np.zeros((len(block_coefficients) + 1, 2, 2))
+    matrix[0] = np.identity(2)
+    for coefficient in block_coefficients:
+        rotation = np.array([[0, coefficient], [-coefficient, 0]])
+        matrix[1:] += matrix[:-1] @ rotation
+    return matrix
+
+
+def _sin_power_derivatives(int power, int order, double angular_frequency):
+    matrix = np.zeros((order + 1, power + 1))
+    matrix[0, power] = 1
+    for derivative in range(1, order + 1):
+        if derivative % 2:
+            matrix[derivative, :-1] = (
+                matrix[derivative - 1, 1:] * np.arange(1, power + 1)
+                * angular_frequency
+            )
+        else:
+            matrix[derivative, :-2] = (
+                matrix[derivative - 2, 2:] * np.arange(1, power)
+                * np.arange(2, power + 1)
+            )
+            matrix[derivative] -= (
+                matrix[derivative - 2] * np.arange(power + 1) ** 2
+            )
+            matrix[derivative] *= angular_frequency ** 2
+    return matrix
+
+
+def _edge_polynomial(derivatives, double x):
+    values = np.asarray(derivatives, dtype=float).copy()
+    values[0] -= 1
+    size = values.shape[0]
+    matrix = np.zeros((size, size))
+    for derivative in range(size):
+        for coefficient in range(size):
+            matrix[derivative, coefficient] = (
+                x ** (size + coefficient - derivative)
+                * factorial(size + coefficient)
+                / factorial(size + coefficient - derivative)
+            )
+    coefficients = np.linalg.solve(matrix, values)
+    return np.poly1d((*np.flip(coefficients), *np.zeros_like(values[:-1]), 1))
+
+
+@lru_cache(maxsize=128)
+def _drag_sin_plan(double width, double delta, tuple block_freq):
+    frequencies = np.asarray(block_freq, dtype=float)
+    block_coefficients = (
+        np.empty(0) if frequencies.size == 0
+        else 1 / (2 * pi * (frequencies - delta))
+    )
+    power = max(((len(block_coefficients) + 2) >> 1) << 1, 2)
+    transform = _b_series_matrix(block_coefficients)
+    derivatives = _sin_power_derivatives(
+        power, len(block_coefficients), pi / width
+    )
+    peak_basis = np.ones(power + 1)
+    peak_basis[1::2] = 0
+    peak_derivatives = derivatives @ peak_basis
+    transformed_peak = np.einsum(
+        "ijk,ki->j",
+        transform,
+        np.array([peak_derivatives, np.zeros_like(peak_derivatives)]),
+    )
+    normalization = np.sqrt(np.sum(np.abs(transformed_peak) ** 2))
+    return transform, derivatives, np.arange(power + 1), normalization
+
+
+def _drag_sin_derivatives(t, double t0, double width, double plateau,
+                          derivatives, powers):
+    midpoint = t0 + width / 2
+    plateau_stop = midpoint + plateau
+    plateau_mask = (t > midpoint) & (t < plateau_stop)
+    adjusted = np.where(t >= plateau_stop, t - plateau, t)
+    angle = pi / width * (adjusted - t0)
+    sine = np.sin(angle)
+    cosine = np.cos(angle)
+    if np.any(plateau_mask):
+        sine = sine.copy()
+        cosine = cosine.copy()
+        sine[plateau_mask] = 0
+        cosine[plateau_mask] = 0
+    basis = sine ** powers.reshape((-1, 1))
+    basis[1::2] *= cosine
+    values = derivatives @ basis
+    if np.any(plateau_mask):
+        values[0, plateau_mask] = 1
+    return values
+
+
+def _drag_omega_sin(t, double t0, double width, double delta,
+                    tuple block_freq, double plateau):
+    transform, derivatives, powers, normalization = _drag_sin_plan(
+        width, delta, block_freq
+    )
+    values = _drag_sin_derivatives(
+        t, t0, width, plateau, derivatives, powers
+    )
+    components = np.array([values, np.zeros_like(values)])
+    return np.einsum("ijk,kim->jm", transform, components) / normalization
+
+
+@lru_cache(maxsize=128)
+def _drag_sinx_plan(double width, double delta, tuple block_freq, double tab):
+    transform, derivatives, powers, _ = _drag_sin_plan(
+        width, delta, block_freq
+    )
+    angular_frequency = pi / width
+
+    left_basis = np.sin(angular_frequency * (1 - tab) * width / 2) ** powers
+    left_basis[1::2] *= np.cos(angular_frequency * (1 - tab) * width / 2)
+    left_polynomial = _edge_polynomial(
+        derivatives @ left_basis, -tab * width / 2
+    )
+
+    right_basis = np.sin(angular_frequency * (1 + tab) * width / 2) ** powers
+    right_basis[1::2] *= np.cos(angular_frequency * (1 + tab) * width / 2)
+    right_polynomial = _edge_polynomial(
+        derivatives @ right_basis, tab * width / 2
+    )
+
+    order = len(block_freq) + 1
+    left_derivatives = tuple(
+        np.polyder(left_polynomial, derivative) for derivative in range(order)
+    )
+    right_derivatives = tuple(
+        np.polyder(right_polynomial, derivative) for derivative in range(order)
+    )
+    return (transform, derivatives, powers, left_derivatives,
+            right_derivatives)
+
+
+def _drag_omega_sinx(t, double t0, double width, double delta,
+                     tuple block_freq, double plateau, double tab):
+    (transform, derivatives, powers, left_derivatives,
+     right_derivatives) = _drag_sinx_plan(width, delta, block_freq, tab)
+    values = _drag_sin_derivatives(
+        t, t0, width, plateau, derivatives, powers
+    )
+    midpoint = t0 + width / 2
+    plateau_stop = midpoint + plateau
+    left_mask = (t >= midpoint - tab * width / 2) & (t <= midpoint)
+    right_mask = ((t >= plateau_stop)
+                  & (t <= plateau_stop + tab * width / 2))
+    for derivative in range(len(block_freq) + 1):
+        values[derivative, left_mask] = left_derivatives[derivative](
+            t[left_mask] - midpoint
+        )
+        values[derivative, right_mask] = right_derivatives[derivative](
+            t[right_mask] - plateau_stop
+        )
+    components = np.array([values, np.zeros_like(values)])
+    return np.einsum("ijk,kim->jm", transform, components)
+
+
+def _drag_sin(t, t0, freq, width, delta, block_freq, phase, plateau):
+    omega_x, omega_y = _drag_omega_sin(
+        t, t0, width, delta, block_freq, plateau
+    )
+    wt = 2 * pi * (freq + delta) * t - (2 * pi * delta * t0 + phase)
+    return omega_x * np.cos(wt) + omega_y * np.sin(wt)
+
+
+def _drag_sinx(t, t0, freq, width, delta, block_freq, phase, plateau, tab):
+    omega_x, omega_y = _drag_omega_sinx(
+        t, t0, width, delta, block_freq, plateau, tab
+    )
+    wt = 2 * pi * (freq + delta) * t - (2 * pi * delta * t0 + phase)
+    return omega_x * np.cos(wt) + omega_y * np.sin(wt)
+
+
 @lru_cache(maxsize=64)
 def _mollifier_poly(d):
     polynomial = np.poly1d([-2, 0])
@@ -205,6 +380,8 @@ _base_functions = {
     DRAG: _drag,
     MOLLIFIER: _MOLLIFIER,
     D_GAUSSIAN: _D_GAUSSIAN,
+    DRAG_SIN: _drag_sin,
+    DRAG_SINX: _drag_sinx,
 }
 
 
@@ -280,6 +457,18 @@ def _encode_builtin_args(bytearray out, int opcode, args):
                                float(values[1]), _time_to_tick(values[2]),
                                float(values[3]), float(values[4]),
                                float(values[5]), _time_to_tick(values[6])))
+    elif opcode in (DRAG_SIN, DRAG_SINX):
+        t0, freq, width, delta, block_freq, phase, plateau = args[:7]
+        block_freq = tuple(float(value) for value in block_freq)
+        out.extend(struct.pack("<qdqddq", _time_to_tick(t0), float(freq),
+                               _time_to_tick(width), float(delta),
+                               float(phase), _time_to_tick(plateau)))
+        if opcode == DRAG_SINX:
+            out.extend(struct.pack("<d", float(args[7])))
+        out.extend(struct.pack("<I", len(block_freq)))
+        if block_freq:
+            out.extend(struct.pack(f"<{len(block_freq)}d", *block_freq))
+        out.extend(struct.pack("<q", _time_to_tick(args[-1])))
     elif opcode in (MOLLIFIER, D_GAUSSIAN):
         out.extend(struct.pack("<qiq", _time_to_tick(args[0]), int(args[1]),
                                _time_to_tick(args[2])))
@@ -318,6 +507,26 @@ def _decode_builtin_args(bytes data, Py_ssize_t pos, int opcode):
             values[4] = None
         values[-1] = _tick_to_time(values[-1])
         return tuple(values), pos + 56
+    if opcode in (DRAG_SIN, DRAG_SINX):
+        t0, freq, width, delta, phase, plateau = struct.unpack_from(
+            "<qdqddq", data, pos
+        )
+        pos += 48
+        tab = None
+        if opcode == DRAG_SINX:
+            tab = struct.unpack_from("<d", data, pos)[0]
+            pos += 8
+        count = struct.unpack_from("<I", data, pos)[0]
+        pos += 4
+        block_freq = struct.unpack_from(f"<{count}d", data, pos)
+        pos += 8 * count
+        shift = struct.unpack_from("<q", data, pos)[0]
+        pos += 8
+        values = (_tick_to_time(t0), freq, _tick_to_time(width), delta,
+                  block_freq, phase, _tick_to_time(plateau))
+        if opcode == DRAG_SINX:
+            values += (tab,)
+        return values + (_tick_to_time(shift),), pos
     if opcode in (MOLLIFIER, D_GAUSSIAN):
         first, order, shift = struct.unpack_from("<qiq", data, pos)
         return (_tick_to_time(first), order, _tick_to_time(shift)), pos + 20

{waveforms-3.0.0 → waveforms-3.1.0}/waveforms/version.py

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

{waveforms-3.0.0 → waveforms-3.1.0}/waveforms/waveform.py

@@ -16,7 +16,8 @@ from numpy import e, inf, pi
 from scipy.signal import sosfilt
 
 from ._waveform import (
-    COS, COSH, D_GAUSSIAN, DRAG, ERF, EXP, EXPONENTIALCHIRP, GAUSSIAN,
+    COS, COSH, D_GAUSSIAN, DRAG, DRAG_SIN, DRAG_SINX, ERF, EXP,
+    EXPONENTIALCHIRP, GAUSSIAN,
     HYPERBOLICCHIRP, INTERP, LINEAR, LINEARCHIRP, MOLLIFIER, SINC, SINH,
     PackedStack, PackedWaveform, basic, constant, get_time_resolution,
     piecewise, quantize_time, registerBaseFunc, registerDerivative,
@@ -813,6 +814,42 @@ def drag(freq, width, plateau=0, delta=0, block_freq=None, phase=0, t0=0):
     )
 
 
+def _block_frequencies(block_freq):
+    if block_freq is None:
+        return ()
+    if np.isscalar(block_freq):
+        return (float(block_freq),)
+    return tuple(float(value) for value in block_freq)
+
+
+def drag_sin(freq, width, plateau=0, delta=0, block_freq=None, phase=0,
+             t0=0):
+    """Return a sine-power DRAG pulse with spectral blocking constraints."""
+    if width <= 0:
+        raise ValueError("width must be positive")
+    phase += pi * delta * (width + plateau)
+    core = _scalar(
+        DRAG_SIN, t0, freq, width, delta, _block_frequencies(block_freq),
+        phase, plateau,
+    )
+    return _piecewise((t0, t0 + width + plateau, inf), 0, core, 0)
+
+
+def drag_sinx(freq, width, plateau=0, delta=0, block_freq=None, phase=0,
+              t0=0, tab=0.618):
+    """Return a flat-top sine-power DRAG pulse with polynomial joins."""
+    if width <= 0:
+        raise ValueError("width must be positive")
+    if not 0 < tab <= 1:
+        raise ValueError("tab must be in the interval (0, 1]")
+    phase += pi * delta * (width + plateau)
+    core = _scalar(
+        DRAG_SINX, t0, freq, width, delta, _block_frequencies(block_freq),
+        phase, plateau, tab,
+    )
+    return _piecewise((t0, t0 + width + plateau, inf), 0, core, 0)
+
+
 def chirp(f0, f1, T, phi0=0, type="linear"):
     if f0 == f1:
         return sin(f0, phi0)
@@ -941,7 +978,8 @@ def wave_eval(expr: str) -> Waveform:
 
 __all__ = [
     "D", "Waveform", "WaveVStack", "chirp", "const", "cos", "cosh",
-    "coshPulse", "cosPulse", "cut", "drag", "exp", "function",
+    "coshPulse", "cosPulse", "cut", "drag", "drag_sin", "drag_sinx",
+    "exp", "function",
     "gaussian", "general_cosine", "get_time_resolution", "hanning",
     "interp", "mixing", "mollifier", "one", "play", "poly",
     "registerBaseFunc", "registerDerivative", "samplingPoints",

{waveforms-3.0.0 → waveforms-3.1.0}/waveforms/waveform_parser.py

@@ -32,7 +32,7 @@ class WaveformVisitor:
         self.extra_modules = tuple(extra_modules)
         self.functions = [
             'D', 'chirp', 'const', 'cos', 'cosh', 'coshPulse', 'cosPulse',
-            'cut', 'drag', 'exp', 'gaussian',
+            'cut', 'drag', 'drag_sin', 'drag_sinx', 'exp', 'gaussian',
             'general_cosine', 'hanning', 'interp', 'mixing', 'mollifier',
             'one', 'poly', 'samplingPoints', 'sign', 'sin', 'sinc', 'sinh',
             'square', 'step', 't', 'zero'

{waveforms-3.0.0 → waveforms-3.1.0/waveforms.egg-info}/PKG-INFO

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