• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

feihoo87 / waveforms / 17799669430

17 Sep 2025 01:42PM UTC coverage: 52.692%. First build
17799669430

push

github

feihoo87
Update version number to 2.2.2 and propagate filters and labels in WaveVStack class methods for improved functionality.

9 of 11 new or added lines in 2 files covered. (81.82%)

1331 of 2526 relevant lines covered (52.69%)

6.32 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

45.88
/waveforms/waveform.py
1
from __future__ import annotations
12✔
2

3
from fractions import Fraction
12✔
4
from typing import Generator, Iterable, cast
12✔
5

6
import numpy as np
12✔
7
from numpy import e, inf, pi
12✔
8
from numpy.typing import NDArray
12✔
9
from scipy.signal import sosfilt
12✔
10

11
from ._waveform import (_D, COS, COSH, D_GAUSSIAN, DRAG, ERF, EXP,
12✔
12
                        EXPONENTIALCHIRP, GAUSSIAN, HYPERBOLICCHIRP, INTERP,
13
                        LINEAR, LINEARCHIRP, MOLLIFIER, NDIGITS, SINC, SINH,
14
                        _baseFunc, _baseFunc_latex, _const, _half, _one, _zero,
15
                        add, basic_wave, calc_parts, filter, is_const,
16
                        merge_waveform, mul, pow, registerBaseFunc,
17
                        registerBaseFuncLatex, registerDerivative, shift,
18
                        simplify, wave_sum)
19

20

21
def _test_spec_num(num, spec):
12✔
22
    x = Fraction(num / spec).limit_denominator(1000000000)
×
23
    if x.denominator <= 24:
×
24
        return True, x, 1
×
25
    x = Fraction(spec * num).limit_denominator(1000000000)
×
26
    if x.denominator <= 24:
×
27
        return True, x, -1
×
28
    return False, x, 0
×
29

30

31
def _spec_num_latex(num):
12✔
32
    for spec, spec_latex in [(1, ''), (np.sqrt(2), '\\sqrt{2}'),
×
33
                             (np.sqrt(3), '\\sqrt{3}'),
34
                             (np.sqrt(5), '\\sqrt{5}'),
35
                             (np.log(2), '\\log{2}'), (np.log(3), '\\log{3}'),
36
                             (np.log(5), '\\log{5}'), (np.e, 'e'),
37
                             (np.pi, '\\pi'), (np.pi**2, '\\pi^2'),
38
                             (np.sqrt(np.pi), '\\sqrt{\\pi}')]:
39
        flag, x, sign = _test_spec_num(num, spec)
×
40
        if flag:
×
41
            if sign < 0:
×
42
                spec_latex = f"\\frac{{{1}}}{{{spec_latex}}}"
×
43
            if x.denominator == 1:
×
44
                if x.numerator == 1:
×
45
                    return f"{spec_latex}"
×
46
                else:
47
                    return f"{x.numerator:g}{spec_latex}"
×
48
            else:
49
                if x.numerator < 0:
×
50
                    return f"-\\frac{{{-x.numerator}}}{{{x.denominator}}}{spec_latex}"
×
51
                else:
52
                    return f"\\frac{{{x.numerator}}}{{{x.denominator}}}{spec_latex}"
×
53
    return f"{num:g}"
×
54

55

56
def _num_latex(num):
12✔
57
    if num == -np.inf:
×
58
        return r"-\infty"
×
59
    elif num == np.inf:
×
60
        return r"\infty"
×
61
    if num.imag > 0:
×
62
        return f"\\left({_num_latex(num.real)}+{_num_latex(num.imag)}j\\right)"
×
63
    elif num.imag < 0:
×
64
        return f"\\left({_num_latex(num.real)}-{_num_latex(-num.imag)}j\\right)"
×
65
    s = _spec_num_latex(num.real)
×
66
    if s == '' and round(num.real) == 1:
×
67
        return '1'
×
68
    if "e" in s:
×
69
        a, n = s.split("e")
×
70
        n = float(n)
×
71
        s = f"{a} \\times 10^{{{n:g}}}"
×
72
    return s
×
73

74

75
def _fun_latex(fun):
12✔
76
    funID, *args, shift = fun
×
77
    if _baseFunc_latex[funID] is None:
×
78
        shift = _num_latex(shift)
×
79
        if shift == "0":
×
80
            shift = ""
×
81
        elif shift[0] != '-':
×
82
            shift = "+" + shift
×
83
        return r"\mathrm{Func}" + f"{funID}(t{shift}, ...)"
×
84
    return _baseFunc_latex[funID](shift, *args)
×
85

86

87
def _wav_latex(wav):
12✔
88

89
    if wav == _zero:
×
90
        return "0"
×
91
    elif is_const(wav):
×
92
        return f"{wav[1][0]}"
×
93

94
    sum_expr = []
×
95
    for mul, amp in zip(*wav):
×
96
        if mul == ((), ()):
×
97
            sum_expr.append(_num_latex(amp))
×
98
            continue
×
99
        mul_expr = []
×
100
        amp = _num_latex(amp)
×
101
        if amp != "1":
×
102
            mul_expr.append(amp)
×
103
        for fun, n in zip(*mul):
×
104
            fun_expr = _fun_latex(fun)
×
105
            if n != 1:
×
106
                mul_expr.append(fun_expr + "^{" + f"{n}" + "}")
×
107
            else:
108
                mul_expr.append(fun_expr)
×
109
        sum_expr.append(''.join(mul_expr))
×
110

111
    ret = sum_expr[0]
×
112
    for expr in sum_expr[1:]:
×
113
        if expr[0] == '-':
×
114
            ret += expr
×
115
        else:
116
            ret += "+" + expr
×
117
    return ret
×
118

119

120
class Waveform:
12✔
121
    __slots__ = ('bounds', 'seq', 'max', 'min', 'start', 'stop', 'sample_rate',
12✔
122
                 'filters', 'label')
123

124
    def __init__(self, bounds=(+inf, ), seq=(_zero, ), min=-inf, max=inf):
12✔
125
        self.bounds = bounds
12✔
126
        self.seq = seq
12✔
127
        self.max = max
12✔
128
        self.min = min
12✔
129
        self.start = None
12✔
130
        self.stop = None
12✔
131
        self.sample_rate = None
12✔
132
        self.filters: tuple[np.ndarray, float] | None = None
12✔
133
        self.label = None
12✔
134

135
    @staticmethod
12✔
136
    def _begin(bounds, seq):
12✔
137
        for i, s in enumerate(seq):
×
138
            if s != _zero:
×
139
                if i == 0:
×
140
                    return -inf
×
141
                return bounds[i - 1]
×
142
        return inf
×
143

144
    @staticmethod
12✔
145
    def _end(bounds, seq):
12✔
146
        N = len(bounds)
×
147
        for i, s in enumerate(seq[::-1]):
×
148
            if s != _zero:
×
149
                if i == 0:
×
150
                    return inf
×
151
                return bounds[N - i - 1]
×
152
        return -inf
×
153

154
    @property
12✔
155
    def begin(self):
12✔
156
        if self.start is None:
×
157
            return self._begin(self.bounds, self.seq)
×
158
        else:
159
            return max(self.start, self._begin(self.bounds, self.seq))
×
160

161
    @property
12✔
162
    def end(self):
12✔
163
        if self.stop is None:
×
164
            return self._end(self.bounds, self.seq)
×
165
        else:
166
            return min(self.stop, self._end(self.bounds, self.seq))
×
167

168
    def sample(
12✔
169
        self,
170
        sample_rate=None,
171
        out: np.ndarray | None = None,
172
        chunk_size=None,
173
        function_lib=None,
174
        filters: tuple[np.ndarray, float] | None = None
175
    ) -> np.ndarray | Iterable[np.ndarray]:
176
        if sample_rate is None:
12✔
177
            sample_rate = self.sample_rate
12✔
178
        if self.start is None or self.stop is None or sample_rate is None:
12✔
179
            raise ValueError(
×
180
                f'Waveform is not initialized. {self.start=}, {self.stop=}, {sample_rate=}'
181
            )
182
        if filters is None:
12✔
183
            filters = self.filters
12✔
184
        if chunk_size is None:
12✔
185
            x = np.arange(self.start, self.stop, 1 / sample_rate)
12✔
186
            sig = cast(np.ndarray,
12✔
187
                       self.__call__(x, out=out, function_lib=function_lib))
188
            if filters is not None:
12✔
189
                sos, initial = filters
12✔
190
                if not isinstance(sos, np.ndarray):
12✔
191
                    sos = np.array(sos)
×
192
                elif not sos.flags.writeable:
12✔
193
                    sos = sos.copy()
×
194
                if initial:
12✔
195
                    sig = cast(np.ndarray, sosfilt(sos,
×
196
                                                   sig - initial)) + initial
197
                else:
198
                    sig = cast(np.ndarray, sosfilt(sos, sig))
12✔
199
            return cast(np.ndarray, sig)
12✔
200
        else:
201
            return self._sample_iter(sample_rate, chunk_size, out,
×
202
                                     function_lib, filters)
203

204
    def _sample_iter(
12✔
205
        self, sample_rate, chunk_size, out: np.ndarray | None, function_lib,
206
        filters: tuple[np.ndarray, float] | None
207
    ) -> Generator[np.ndarray, None, None]:
208
        start = cast(float, self.start)
×
209
        start_n = 0
×
210
        if filters is not None:
×
211
            sos, initial = filters
×
212
            if not isinstance(sos, np.ndarray):
×
213
                sos = np.array(sos)
×
214
            elif not sos.flags.writeable:
×
215
                sos = sos.copy()
×
216
            # zi = sosfilt_zi(sos)
217
            zi = np.zeros((sos.shape[0], 2))
×
218
        length = chunk_size / sample_rate
×
219
        while start < cast(float, self.stop):
×
220
            if start + length > cast(float, self.stop):
×
221
                length = cast(float, self.stop) - start
×
222
                stop = cast(float, self.stop)
×
223
                size = round((stop - start) * sample_rate)
×
224
            else:
225
                stop = start + length
×
226
                size = chunk_size
×
227
            x = np.linspace(start, stop, size, endpoint=False)
×
228

229
            if filters is None:
×
230
                if out is not None:
×
231
                    yield cast(
×
232
                        np.ndarray,
233
                        self.__call__(x,
234
                                      out=out[start_n:],
235
                                      function_lib=function_lib))
236
                else:
237
                    yield cast(np.ndarray,
×
238
                               self.__call__(x, function_lib=function_lib))
239
            else:
240
                sig = cast(np.ndarray,
×
241
                           self.__call__(x, function_lib=function_lib))
242
                if initial:
×
243
                    sig -= initial
×
244
                sig, zi = sosfilt(sos, sig, zi=zi)
×
245
                if initial:
×
246
                    sig += initial
×
247
                if out is not None:
×
248
                    out[start_n:start_n + size] = sig
×
249
                yield cast(np.ndarray, sig)
×
250

251
            start = stop
×
252
            start_n += chunk_size
×
253

254
    @staticmethod
12✔
255
    def _tolist(bounds, seq, ret=None):
12✔
256
        if ret is None:
12✔
257
            ret = []
×
258
        ret.append(len(bounds))
12✔
259
        for seq, b in zip(seq, bounds):
12✔
260
            ret.append(b)
12✔
261
            tlist, amplist = seq
12✔
262
            ret.append(len(amplist))
12✔
263
            for t, amp in zip(tlist, amplist):
12✔
264
                ret.append(amp)
12✔
265
                mtlist, nlist = t
12✔
266
                ret.append(len(nlist))
12✔
267
                for fun, n in zip(mtlist, nlist):
12✔
268
                    ret.append(n)
12✔
269
                    ret.append(len(fun))
12✔
270
                    ret.extend(fun)
12✔
271
        return ret
12✔
272

273
    @staticmethod
12✔
274
    def _fromlist(l, pos=0):
12✔
275

276
        def _read(l, pos, size):
12✔
277
            try:
12✔
278
                return tuple(l[pos:pos + size]), pos + size
12✔
279
            except:
×
280
                raise ValueError('Invalid waveform format')
×
281

282
        (nseg, ), pos = _read(l, pos, 1)
12✔
283
        bounds = []
12✔
284
        seq = []
12✔
285
        for _ in range(nseg):
12✔
286
            (b, nsum), pos = _read(l, pos, 2)
12✔
287
            bounds.append(b)
12✔
288
            amp = []
12✔
289
            t = []
12✔
290
            for _ in range(nsum):
12✔
291
                (a, nmul), pos = _read(l, pos, 2)
12✔
292
                amp.append(a)
12✔
293
                nlst = []
12✔
294
                mt = []
12✔
295
                for _ in range(nmul):
12✔
296
                    (n, nfun), pos = _read(l, pos, 2)
12✔
297
                    nlst.append(n)
12✔
298
                    fun, pos = _read(l, pos, nfun)
12✔
299
                    mt.append(fun)
12✔
300
                t.append((tuple(mt), tuple(nlst)))
12✔
301
            seq.append((tuple(t), tuple(amp)))
12✔
302

303
        return tuple(bounds), tuple(seq), pos
12✔
304

305
    def tolist(self):
12✔
306
        l = [self.max, self.min, self.start, self.stop, self.sample_rate]
12✔
307
        if self.filters is None:
12✔
308
            l.append(None)
12✔
309
        else:
310
            sos, initial = self.filters
12✔
311
            sos = list(sos.reshape(-1))
12✔
312
            l.append(len(sos))
12✔
313
            l.extend(sos)
12✔
314
            l.append(initial)
12✔
315

316
        return self._tolist(self.bounds, self.seq, l)
12✔
317

318
    @classmethod
12✔
319
    def fromlist(cls, l):
12✔
320
        w = cls()
12✔
321
        pos = 6
12✔
322
        (w.max, w.min, w.start, w.stop, w.sample_rate, sos_size) = l[:pos]
12✔
323
        if sos_size is not None:
12✔
324
            sos = np.array(l[pos:pos + sos_size]).reshape(-1, 6)
12✔
325
            pos += sos_size
12✔
326
            initial = l[pos]
12✔
327
            pos += 1
12✔
328
            w.filters = sos, initial
12✔
329

330
        w.bounds, w.seq, pos = cls._fromlist(l, pos)
12✔
331
        return w
12✔
332

333
    def totree(self):
12✔
334
        if self.filters is None:
12✔
335
            header = (self.max, self.min, self.start, self.stop,
12✔
336
                      self.sample_rate, None)
337
        else:
338
            header = (self.max, self.min, self.start, self.stop,
12✔
339
                      self.sample_rate, self.filters)
340
        body = []
12✔
341

342
        for seq, b in zip(self.seq, self.bounds):
12✔
343
            tlist, amplist = seq
12✔
344
            new_seq = []
12✔
345
            for t, amp in zip(tlist, amplist):
12✔
346
                mtlist, nlist = t
12✔
347
                new_t = []
12✔
348
                for fun, n in zip(mtlist, nlist):
12✔
349
                    new_t.append((n, fun))
12✔
350
                new_seq.append((amp, tuple(new_t)))
12✔
351
            body.append((b, tuple(new_seq)))
12✔
352
        return header, tuple(body)
12✔
353

354
    @staticmethod
12✔
355
    def fromtree(tree):
12✔
356
        w = Waveform()
12✔
357
        header, body = tree
12✔
358

359
        (w.max, w.min, w.start, w.stop, w.sample_rate, w.filters) = header
12✔
360
        bounds = []
12✔
361
        seqs = []
12✔
362
        for b, seq in body:
12✔
363
            bounds.append(b)
12✔
364
            amp_list = []
12✔
365
            t_list = []
12✔
366
            for amp, t in seq:
12✔
367
                amp_list.append(amp)
12✔
368
                n_list = []
12✔
369
                mt_list = []
12✔
370
                for n, mt in t:
12✔
371
                    n_list.append(n)
12✔
372
                    mt_list.append(mt)
12✔
373
                t_list.append((tuple(mt_list), tuple(n_list)))
12✔
374
            seqs.append((tuple(t_list), tuple(amp_list)))
12✔
375
        w.bounds = tuple(bounds)
12✔
376
        w.seq = tuple(seqs)
12✔
377
        return w
12✔
378

379
    def simplify(self, eps=1e-15):
12✔
380
        seq = [simplify(self.seq[0], eps)]
12✔
381
        bounds = [self.bounds[0]]
12✔
382
        for expr, b in zip(self.seq[1:], self.bounds[1:]):
12✔
383
            expr = simplify(expr, eps)
12✔
384
            if expr == seq[-1]:
12✔
385
                seq.pop()
×
386
                bounds.pop()
×
387
            seq.append(expr)
12✔
388
            bounds.append(b)
12✔
389
        return Waveform(tuple(bounds), tuple(seq))
12✔
390

391
    def filter(self, low=0, high=inf, eps=1e-15):
12✔
392
        seq = []
×
393
        for expr in self.seq:
×
394
            seq.append(filter(expr, low, high, eps))
×
395
        return Waveform(self.bounds, tuple(seq))
×
396

397
    def _comb(self, other, oper):
12✔
398
        return Waveform(*merge_waveform(self.bounds, self.seq, other.bounds,
12✔
399
                                        other.seq, oper))
400

401
    def __pow__(self, n) -> Waveform:
12✔
402
        return Waveform(self.bounds, tuple(pow(w, n) for w in self.seq))
12✔
403

404
    def __add__(self, other) -> Waveform:
12✔
405
        if isinstance(other, Waveform):
12✔
406
            return self._comb(other, add)
12✔
407
        else:
408
            return self + const(other)
12✔
409

410
    def __radd__(self, v) -> Waveform:
12✔
411
        return const(v) + self
×
412

413
    def __ior__(self, other) -> Waveform:
12✔
414
        return self | other
×
415

416
    def __or__(self, other) -> Waveform:
12✔
417
        if isinstance(other, (int, float, complex)):
×
418
            other = const(other)
×
419
        w = self.marker + other.marker
×
420

421
        def _or(a, b):
×
422
            if a != _zero or b != _zero:
×
423
                return _one
×
424
            else:
425
                return _zero
×
426

427
        return self._comb(other, _or)
×
428

429
    def __iand__(self, other) -> Waveform:
12✔
430
        return self & other
×
431

432
    def __and__(self, other) -> Waveform:
12✔
433
        if isinstance(other, (int, float, complex)):
×
434
            other = const(other)
×
435
        w = self.marker + other.marker
×
436

437
        def _and(a, b):
×
438
            if a != _zero and b != _zero:
×
439
                return _one
×
440
            else:
441
                return _zero
×
442

443
        return self._comb(other, _and)
×
444

445
    @property
12✔
446
    def marker(self):
12✔
447
        w = self.simplify()
×
448
        return Waveform(w.bounds,
×
449
                        tuple(_zero if s == _zero else _one for s in w.seq))
450

451
    def mask(self, edge: float = 0) -> Waveform:
12✔
452
        w = self.marker
×
453
        in_wave = w.seq[0] == _zero
×
454
        bounds = []
×
455
        seq = []
×
456

457
        if w.seq[0] == _zero:
×
458
            in_wave = False
×
459
            b = w.bounds[0] - edge
×
460
            bounds.append(b)
×
461
            seq.append(_zero)
×
462

463
        for b, s in zip(w.bounds[1:], w.seq[1:]):
×
464
            if not in_wave and s != _zero:
×
465
                in_wave = True
×
466
                bounds.append(b + edge)
×
467
                seq.append(_one)
×
468
            elif in_wave and s == _zero:
×
469
                in_wave = False
×
470
                b = b - edge
×
471
                if b > bounds[-1]:
×
472
                    bounds.append(b)
×
473
                    seq.append(_zero)
×
474
                else:
475
                    bounds.pop()
×
476
                    bounds.append(b)
×
477
        return Waveform(tuple(bounds), tuple(seq))
×
478

479
    def __mul__(self, other) -> Waveform:
12✔
480
        if isinstance(other, Waveform):
12✔
481
            return self._comb(other, mul)
12✔
482
        else:
483
            return self * const(other)
×
484

485
    def __rmul__(self, v) -> Waveform:
12✔
486
        return const(v) * self
12✔
487

488
    def __truediv__(self, other) -> Waveform:
12✔
489
        if isinstance(other, Waveform):
12✔
490
            raise TypeError('division by waveform')
×
491
        else:
492
            return self * const(1 / other)
12✔
493

494
    def __neg__(self) -> Waveform:
12✔
495
        return -1 * self
12✔
496

497
    def __sub__(self, other) -> Waveform:
12✔
498
        return self + (-other)
12✔
499

500
    def __rsub__(self, v) -> Waveform:
12✔
501
        return v + (-self)
×
502

503
    def __rshift__(self, time) -> Waveform:
12✔
504
        return Waveform(
12✔
505
            tuple(round(bound + time, NDIGITS) for bound in self.bounds),
506
            tuple(shift(expr, time) for expr in self.seq))
507

508
    def __lshift__(self, time):
12✔
509
        return self >> (-time)
12✔
510

511
    @staticmethod
12✔
512
    def _merge_parts(
12✔
513
        parts: list[tuple[int, int, np.ndarray | int | float | complex]],
514
        out: list[tuple[int, int, np.ndarray | int | float | complex]]
515
    ) -> list[tuple[int, int, np.ndarray | int | float | complex]]:
516
        # TODO: merge parts
517
        raise NotImplementedError
×
518

519
    @staticmethod
12✔
520
    def _fill_parts(parts, out):
12✔
521
        for start, stop, part in parts:
12✔
522
            out[start:stop] += part
12✔
523

524
    def __call__(
12✔
525
        self,
526
        x,
527
        frag=False,
528
        out: np.ndarray | list | None = None,
529
        accumulate=False,
530
        function_lib=None
531
    ) -> NDArray[np.float64 | np.complex128] | list[
532
            tuple[int, int, NDArray[np.float64 | np.complex128]] | int
533
            | float | complex] | np.float64:
534
        if function_lib is None:
12✔
535
            function_lib = _baseFunc
12✔
536
        if isinstance(x, (int, float, complex)):
12✔
537
            return cast(
×
538
                NDArray[np.float64],
539
                self.__call__(np.array([x]), function_lib=function_lib))[0]
540
        parts, dtype = calc_parts(self.bounds, self.seq, x, function_lib,
12✔
541
                                  self.min, self.max)
542
        if not frag:
12✔
543
            if out is None:
12✔
544
                out = np.zeros_like(x, dtype=dtype)
12✔
545
            elif not accumulate:
×
546
                out *= 0
×
547
            self._fill_parts(parts, out)
12✔
548
        else:
549
            if out is None:
×
550
                return cast(list, parts)
×
551
            else:
552
                out = cast(list, out)
×
553
                if not accumulate:
×
554
                    out.clear()
×
555
                    out.extend(parts)
×
556
                else:
557
                    self._merge_parts(parts, out)
×
558
        return out
12✔
559

560
    def __hash__(self):
12✔
561
        return hash((self.max, self.min, self.start, self.stop,
×
562
                     self.sample_rate, self.bounds, self.seq))
563

564
    def __eq__(self, o: object) -> bool:
12✔
565
        if isinstance(o, (int, float, complex)):
12✔
566
            return self == const(o)
12✔
567
        elif isinstance(o, Waveform):
12✔
568
            a = self.simplify()
12✔
569
            b = o.simplify()
12✔
570
            return a.seq == b.seq and a.bounds == b.bounds and (
12✔
571
                a.max, a.min, a.start, a.stop) == (b.max, b.min, b.start,
572
                                                   b.stop)
573
        else:
574
            return False
×
575

576
    def _repr_latex_(self):
12✔
577
        parts = []
×
578
        start = -np.inf
×
579
        for end, wav in zip(self.bounds, self.seq):
×
580
            e_str = _wav_latex(wav)
×
581
            start_str = _num_latex(start)
×
582
            end_str = _num_latex(end)
×
583
            parts.append(e_str + r",~~&t\in" + f"({start_str},{end_str}" +
×
584
                         (']' if end < np.inf else ')'))
585
            start = end
×
586
        if len(parts) == 1:
×
587
            expr = ''.join(['f(t)=', *parts[0].split('&')])
×
588
        else:
589
            expr = '\n'.join([
×
590
                r"f(t)=\begin{cases}", (r"\\" + '\n').join(parts),
591
                r"\end{cases}"
592
            ])
593
        return "$$\n{}\n$$".format(expr)
×
594

595
    def _play(self, time_unit, volume=1.0):
12✔
596
        import pyaudio
×
597

598
        CHUNK = 1024
×
599
        RATE = 48000
×
600

601
        dynamic_volume = 1.0
×
602
        amp = 2**15 * 0.999 * volume * dynamic_volume
×
603

604
        p = pyaudio.PyAudio()
×
605
        try:
×
606
            stream = p.open(format=pyaudio.paInt16,
×
607
                            channels=1,
608
                            rate=RATE,
609
                            output=True)
610
            try:
×
611
                for data in self.sample(sample_rate=RATE / time_unit,
×
612
                                        chunk_size=CHUNK):
613
                    lim = np.abs(data).max()
×
614
                    if lim > 0 and dynamic_volume > 1.0 / lim:
×
615
                        dynamic_volume = 1.0 / lim
×
616
                        amp = 2**15 * 0.99 * volume * dynamic_volume
×
617
                    data = (amp * data).astype(np.int16)
×
618
                    stream.write(bytes(data.data))
×
619
            finally:
620
                stream.stop_stream()
×
621
                stream.close()
×
622
        finally:
623
            p.terminate()
×
624

625
    def play(self, time_unit=1, volume=1.0):
12✔
626
        import multiprocessing as mp
×
627
        p = mp.Process(target=self._play,
×
628
                       args=(time_unit, volume),
629
                       daemon=True)
630
        p.start()
×
631

632

633
class WaveVStack(Waveform):
12✔
634

635
    def __init__(self, wlist: list[Waveform] = []):
12✔
636
        self.wlist = [(w.bounds, w.seq) for w in wlist]
12✔
637
        self.start = None
12✔
638
        self.stop = None
12✔
639
        self.sample_rate = None
12✔
640
        self.offset = 0
12✔
641
        self.shift = 0
12✔
642
        self.filters = None
12✔
643
        self.label = None
12✔
644
        self.function_lib = None
12✔
645

646
    def __begin(self):
12✔
647
        if self.wlist:
×
648
            v = [self._begin(bounds, seq) for bounds, seq in self.wlist]
×
649
            return min(v)
×
650
        else:
651
            return -inf
×
652

653
    def __end(self):
12✔
654
        if self.wlist:
×
655
            v = [self._end(bounds, seq) for bounds, seq in self.wlist]
×
656
            return max(v)
×
657
        else:
658
            return inf
×
659

660
    @property
12✔
661
    def begin(self):
12✔
662
        if self.start is None:
×
663
            return self.__begin()
×
664
        else:
665
            return max(self.start, self.__begin())
×
666

667
    @property
12✔
668
    def end(self):
12✔
669
        if self.stop is None:
×
670
            return self.__end()
×
671
        else:
672
            return min(self.stop, self.__end())
×
673

674
    def __call__(self, x, frag=False, out=None, function_lib=None):
12✔
675
        assert frag is False, 'WaveVStack does not support frag mode'
12✔
676
        out = np.full_like(x, self.offset, dtype=np.complex128)
12✔
677
        out = cast(NDArray[np.complex128], out)
12✔
678
        if self.shift != 0:
12✔
679
            x = x - self.shift
12✔
680
        if function_lib is None:
12✔
681
            if self.function_lib is None:
12✔
682
                function_lib = _baseFunc
12✔
683
            else:
684
                function_lib = self.function_lib
×
685
        for bounds, seq in self.wlist:
12✔
686
            parts, dtype = calc_parts(bounds, seq, x, function_lib)
12✔
687
            self._fill_parts(parts, out)
12✔
688
        return out.real
12✔
689

690
    def tolist(self):
12✔
691
        l = [
12✔
692
            self.start,
693
            self.stop,
694
            self.offset,
695
            self.shift,
696
            self.sample_rate,
697
        ]
698
        if self.filters is None:
12✔
699
            l.append(None)
12✔
700
        else:
701
            sos, initial = self.filters
12✔
702
            sos = list(sos.reshape(-1))
12✔
703
            l.append(len(sos))
12✔
704
            l.extend(sos)
12✔
705
            l.append(initial)
12✔
706
        l.append(len(self.wlist))
12✔
707
        for bounds, seq in self.wlist:
12✔
708
            self._tolist(bounds, seq, l)
12✔
709
        return l
12✔
710

711
    @classmethod
12✔
712
    def fromlist(cls, l):
12✔
713
        w = cls()
12✔
714
        pos = 6
12✔
715
        w.start, w.stop, w.offset, w.shift, w.sample_rate, sos_size = l[:pos]
12✔
716
        if sos_size is not None:
12✔
717
            sos = np.array(l[pos:pos + sos_size]).reshape(-1, 6)
12✔
718
            pos += sos_size
12✔
719
            initial = l[pos]
12✔
720
            pos += 1
12✔
721
            w.filters = sos, initial
12✔
722
        n = l[pos]
12✔
723
        pos += 1
12✔
724
        for _ in range(n):
12✔
725
            bounds, seq, pos = cls._fromlist(l, pos)
12✔
726
            w.wlist.append((bounds, seq))
12✔
727
        return w
12✔
728

729
    def simplify(self, eps=1e-15):
12✔
730
        if not self.wlist:
12✔
731
            return zero()
12✔
732
        bounds, seq = wave_sum(self.wlist)
12✔
733
        wav = Waveform(bounds=bounds, seq=seq)
12✔
734
        if self.offset != 0:
12✔
735
            wav += self.offset
×
736
        if self.shift != 0:
12✔
737
            wav >>= self.shift
×
738
        wav = wav.simplify(eps)
12✔
739
        wav.start = self.start
12✔
740
        wav.stop = self.stop
12✔
741
        wav.sample_rate = self.sample_rate
12✔
742
        wav.filters = self.filters
12✔
743
        wav.label = self.label
12✔
744
        return wav
12✔
745

746
    @staticmethod
12✔
747
    def _rshift(wlist, time):
12✔
748
        if time == 0:
×
749
            return wlist
×
750
        return [(tuple(round(bound + time, NDIGITS) for bound in bounds),
×
751
                 tuple(shift(expr, time) for expr in seq))
752
                for bounds, seq in wlist]
753

754
    def __rshift__(self, time):
12✔
755
        ret = WaveVStack()
12✔
756
        ret.wlist = self.wlist
12✔
757
        ret.sample_rate = self.sample_rate
12✔
758
        ret.start = self.start
12✔
759
        ret.stop = self.stop
12✔
760
        ret.shift = self.shift + time
12✔
761
        ret.offset = self.offset
12✔
762
        ret.filters = self.filters
12✔
763
        ret.label = self.label
12✔
764
        return ret
12✔
765

766
    def __add__(self, other) -> WaveVStack:
12✔
767
        ret = WaveVStack()
12✔
768
        ret.wlist.extend(self.wlist)
12✔
769
        if isinstance(other, WaveVStack):
12✔
770
            if other.shift != self.shift:
×
771
                ret.wlist = self._rshift(ret.wlist, self.shift)
×
772
                ret.wlist.extend(self._rshift(other.wlist, other.shift))
×
773
            else:
774
                ret.wlist.extend(other.wlist)
×
775
            ret.offset = self.offset + other.offset
×
776
        elif isinstance(other, Waveform):
12✔
777
            other <<= self.shift
12✔
778
            ret.wlist.append((other.bounds, other.seq))
12✔
779
        else:
780
            # ret.wlist.append(((+inf, ), (_const(1.0 * other), )))
781
            ret.offset += other
12✔
782
        ret.filters = self.filters
12✔
783
        ret.label = self.label
12✔
784
        return ret
12✔
785

786
    def __radd__(self, v) -> WaveVStack:
12✔
787
        return self + v
×
788

789
    def __mul__(self, other) -> WaveVStack:
12✔
790
        if isinstance(other, Waveform):
12✔
791
            other = other.simplify() << self.shift
12✔
792
            ret = WaveVStack([Waveform(*w) * other for w in self.wlist])
12✔
793
            if self.offset != 0:
12✔
794
                w = other * self.offset
×
795
                ret.wlist.append((w.bounds, w.seq))
×
796
            ret.filters = self.filters
12✔
797
            ret.label = self.label
12✔
798
            return ret
12✔
799
        else:
800
            ret = WaveVStack([Waveform(*w) * other for w in self.wlist])
×
801
            ret.offset = self.offset * other
×
NEW
802
            ret.filters = self.filters
×
NEW
803
            ret.label = self.label
×
804
            return ret
×
805

806
    def __rmul__(self, v) -> WaveVStack:
12✔
807
        return self * v
×
808

809
    def __eq__(self, other) -> bool:
12✔
810
        if self.wlist:
×
811
            return False
×
812
        else:
813
            return zero() == other
×
814

815
    def _repr_latex_(self):
12✔
816
        return r"\sum_{i=1}^{" + f"{len(self.wlist)}" + r"}" + r"f_i(t)"
×
817

818
    def __getstate__(self) -> tuple:
12✔
819
        function_lib = self.function_lib
×
820
        if function_lib:
×
821
            try:
×
822
                import dill
×
823
                function_lib = dill.dumps(function_lib)
×
824
            except:
×
825
                function_lib = None
×
826
        return (self.wlist, self.start, self.stop, self.sample_rate,
×
827
                self.offset, self.shift, self.filters, self.label,
828
                function_lib)
829

830
    def __setstate__(self, state: tuple) -> None:
12✔
831
        (self.wlist, self.start, self.stop, self.sample_rate, self.offset,
×
832
         self.shift, self.filters, self.label, function_lib) = state
833
        if function_lib:
×
834
            try:
×
835
                import dill
×
836
                function_lib = dill.loads(function_lib)
×
837
            except:
×
838
                function_lib = None
×
839
        self.function_lib = function_lib
×
840

841

842
def play(data, rate=48000):
12✔
843
    import io
×
844

845
    import pyaudio
×
846

847
    CHUNK = 1024
×
848

849
    max_amp = np.max(np.abs(data))
×
850

851
    if max_amp > 1:
×
852
        data /= max_amp
×
853

854
    data = np.array(2**15 * 0.999 * data, dtype=np.int16)
×
855
    buff = io.BytesIO(data.data)
×
856
    p = pyaudio.PyAudio()
×
857

858
    try:
×
859
        stream = p.open(format=pyaudio.paInt16,
×
860
                        channels=1,
861
                        rate=rate,
862
                        output=True)
863
        try:
×
864
            while True:
×
865
                data = buff.read(CHUNK)
×
866
                if data:
×
867
                    stream.write(data)
×
868
                else:
869
                    break
×
870
        finally:
871
            stream.stop_stream()
×
872
            stream.close()
×
873
    finally:
874
        p.terminate()
×
875

876

877
_zero_waveform = Waveform()
12✔
878
_one_waveform = Waveform(seq=(_one, ))
12✔
879

880

881
def zero():
12✔
882
    return _zero_waveform
12✔
883

884

885
def one():
12✔
886
    return _one_waveform
12✔
887

888

889
def const(c):
12✔
890
    return Waveform(seq=(_const(1.0 * c), ))
12✔
891

892

893
# register base function
894
def _format_LINEAR(shift, *args):
12✔
895
    if shift != 0:
×
896
        shift = _num_latex(-shift)
×
897
        if shift[0] == '-':
×
898
            return f"(t{shift})"
×
899
        else:
900
            return f"(t+{shift})"
×
901
    else:
902
        return 't'
×
903

904

905
def _format_GAUSSIAN(shift, *args):
12✔
906
    sigma = _num_latex(args[0] / np.sqrt(2))
×
907
    shift = _num_latex(-shift)
×
908
    if shift != '0':
×
909
        if shift[0] != '-':
×
910
            shift = '+' + shift
×
911
        if sigma == '1':
×
912
            return ('\\exp\\left[-\\frac{\\left(t' + shift +
×
913
                    '\\right)^2}{2}\\right]')
914
        else:
915
            return ('\\exp\\left[-\\frac{1}{2}\\left(\\frac{t' + shift + '}{' +
×
916
                    sigma + '}\\right)^2\\right]')
917
    else:
918
        if sigma == '1':
×
919
            return ('\\exp\\left(-\\frac{t^2}{2}\\right)')
×
920
        else:
921
            return ('\\exp\\left[-\\frac{1}{2}\\left(\\frac{t}{' + sigma +
×
922
                    '}\\right)^2\\right]')
923

924

925
def _format_SINC(shift, *args):
12✔
926
    shift = _num_latex(-shift)
×
927
    bw = _num_latex(args[0])
×
928
    if shift != '0':
×
929
        if shift[0] != '-':
×
930
            shift = '+' + shift
×
931
        if bw == '1':
×
932
            return '\\mathrm{sinc}(t' + shift + ')'
×
933
        else:
934
            return '\\mathrm{sinc}[' + bw + '(t' + shift + ')]'
×
935
    else:
936
        if bw == '1':
×
937
            return '\\mathrm{sinc}(t)'
×
938
        else:
939
            return '\\mathrm{sinc}(' + bw + 't)'
×
940

941

942
def _format_COSINE(shift, *args):
12✔
943
    freq = args[0] / 2 / np.pi
×
944
    phase = -shift * freq
×
945
    freq = _num_latex(freq)
×
946
    if freq == '1':
×
947
        freq = ''
×
948
    phase = _num_latex(phase)
×
949
    if phase == '0':
×
950
        phase = ''
×
951
    elif phase[0] != '-':
×
952
        phase = '+' + phase
×
953
    if phase != '':
×
954
        return f'\\cos\\left[2\\pi\\left({freq}t{phase}\\right)\\right]'
×
955
    elif freq != '':
×
956
        return f'\\cos\\left(2\\pi\\times {freq}t\\right)'
×
957
    else:
958
        return '\\cos\\left(2\\pi t\\right)'
×
959

960

961
def _format_ERF(shift, *args):
12✔
962
    if shift > 0:
×
963
        return '\\mathrm{erf}(\\frac{t-' + f"{_num_latex(shift)}" + '}{' + f'{args[0]:g}' + '})'
×
964
    elif shift < 0:
×
965
        return '\\mathrm{erf}(\\frac{t+' + f"{_num_latex(-shift)}" + '}{' + f'{args[0]:g}' + '})'
×
966
    else:
967
        return '\\mathrm{erf}(\\frac{t}{' + f'{args[0]:g}' + '})'
×
968

969

970
def _format_COSH(shift, *args):
12✔
971
    if shift > 0:
×
972
        return '\\cosh(\\frac{t-' + f"{_num_latex(shift)}" + '}{' + f'{1/args[0]:g}' + '})'
×
973
    elif shift < 0:
×
974
        return '\\cosh(\\frac{t+' + f"{_num_latex(-shift)}" + '}{' + f'{1/args[0]:g}' + '})'
×
975
    else:
976
        return '\\cosh(\\frac{t}{' + f'{1/args[0]:g}' + '})'
×
977

978

979
def _format_SINH(shift, *args):
12✔
980
    if shift > 0:
×
981
        return '\\sinh(\\frac{t-' + f"{_num_latex(shift)}" + '}{' + f'{args[0]:g}' + '})'
×
982
    elif shift < 0:
×
983
        return '\\sinh(\\frac{t+' + f"{_num_latex(-shift)}" + '}{' + f'{args[0]:g}' + '})'
×
984
    else:
985
        return '\\sinh(\\frac{t}{' + f'{args[0]:g}' + '})'
×
986

987

988
def _format_EXP(shift, *args):
12✔
989
    if _num_latex(shift) and shift > 0:
×
990
        return '\\exp\\left(-' + f'{args[0]:g}' + '\\left(t-' + f"{_num_latex(shift)}" + '\\right)\\right)'
×
991
    elif _num_latex(-shift) and shift < 0:
×
992
        return '\\exp\\left(-' + f'{args[0]:g}' + '\\left(t+' + f"{_num_latex(-shift)}" + '\\right)\\right)'
×
993
    else:
994
        return '\\exp\\left(-' + f'{args[0]:g}' + 't\\right)'
×
995

996

997
def _format_DRAG(shift, *args):
12✔
998
    return f"DRAG(...)"
×
999

1000

1001
def _format_MOLLIFIER(shift, *args):
12✔
1002
    r = _num_latex(args[0])
×
1003
    d = _num_latex(args[1])
×
1004
    shift_str = _num_latex(-shift)
×
1005
    if shift_str == '0':
×
1006
        shift_str = ''
×
1007
    elif shift_str[0] != '-':
×
1008
        shift_str = '+' + shift_str
×
1009

1010
    if d == '0':
×
1011
        return f"\\mathrm{{Mollifier}}\\left(t{shift_str}, r={r}\\right)"
×
1012
    elif d == '1':
×
1013
        return f"\\mathrm{{Mollifier}}'\\left(t{shift_str}, r={r}\\right)"
×
1014
    elif d == '2':
×
1015
        return f"\\mathrm{{Mollifier}}''\\left(t{shift_str}, r={r}\\right)"
×
1016
    else:
1017
        return f"\\mathrm{{Mollifier}}^{{({d})}}\\left(t{shift_str}, r={r}\\right)"
×
1018

1019

1020
def _format_D_GAUSSIAN(shift, *args):
12✔
1021
    sigma = _num_latex(args[0] / np.sqrt(2))
×
1022
    d = args[1]
×
1023
    shift_str = _num_latex(-shift)
×
1024
    if shift_str == '0':
×
1025
        shift_str = ''
×
1026
    elif shift_str[0] != '-':
×
1027
        shift_str = '+' + shift_str
×
1028

1029
    if d == 0:
×
1030
        return f"\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
×
1031
    elif d == 1:
×
1032
        return f"\\frac{{\\mathrm{{d}}}}{{\\mathrm{{d}}t}}\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
×
1033
    else:
1034
        return f"\\frac{{\\mathrm{{d}}^{{{d}}}}}{{\\mathrm{{d}}t^{{{d}}}}}\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
×
1035

1036

1037
registerBaseFuncLatex(LINEAR, _format_LINEAR)
12✔
1038
registerBaseFuncLatex(GAUSSIAN, _format_GAUSSIAN)
12✔
1039
registerBaseFuncLatex(ERF, _format_ERF)
12✔
1040
registerBaseFuncLatex(COS, _format_COSINE)
12✔
1041
registerBaseFuncLatex(SINC, _format_SINC)
12✔
1042
registerBaseFuncLatex(EXP, _format_EXP)
12✔
1043
registerBaseFuncLatex(COSH, _format_COSH)
12✔
1044
registerBaseFuncLatex(SINH, _format_SINH)
12✔
1045
registerBaseFuncLatex(DRAG, _format_DRAG)
12✔
1046
registerBaseFuncLatex(MOLLIFIER, _format_MOLLIFIER)
12✔
1047
registerBaseFuncLatex(D_GAUSSIAN, _format_D_GAUSSIAN)
12✔
1048

1049

1050
def D(wav: Waveform, d: int = 1) -> Waveform:
12✔
1051
    """derivative
1052

1053
    Parameters
1054
    ----------
1055
    wav : Waveform
1056
        The waveform to take the derivative of.
1057
    d : int, optional
1058
        The order of the derivative, by default 1.
1059
    """
1060
    assert d >= 0 and isinstance(d, int), "d must be a non-negative integer"
×
1061
    if d == 0:
×
1062
        return wav
×
1063
    elif d == 1:
×
1064
        return Waveform(bounds=wav.bounds, seq=tuple(_D(x) for x in wav.seq))
×
1065
    else:
1066
        return D(D(wav, d - 1), 1)
×
1067

1068

1069
def convolve(a, b):
12✔
1070
    pass
×
1071

1072

1073
def sign():
12✔
1074
    return Waveform(bounds=(0, +inf), seq=(_const(-1), _one))
×
1075

1076

1077
def step(edge, type='erf'):
12✔
1078
    """
1079
    type: "erf", "cos", "linear"
1080
    """
1081
    if edge == 0:
12✔
1082
        return Waveform(bounds=(0, +inf), seq=(_zero, _one))
12✔
1083
    if type == 'cos':
12✔
1084
        rise = add(_half,
×
1085
                   mul(_half, basic_wave(COS, pi / edge, shift=0.5 * edge)))
1086
        return Waveform(bounds=(round(-edge / 2,
×
1087
                                      NDIGITS), round(edge / 2,
1088
                                                      NDIGITS), +inf),
1089
                        seq=(_zero, rise, _one))
1090
    elif type == 'linear':
12✔
1091
        rise = add(_half, mul(_const(1 / edge), basic_wave(LINEAR)))
12✔
1092
        return Waveform(bounds=(round(-edge / 2,
12✔
1093
                                      NDIGITS), round(edge / 2,
1094
                                                      NDIGITS), +inf),
1095
                        seq=(_zero, rise, _one))
1096
    else:
1097
        std_sq2 = edge / 5
×
1098
        # rise = add(_half, mul(_half, basic_wave(ERF, std_sq2)))
1099
        rise = ((((), ()), (((ERF, std_sq2, 0), ), (1, ))), (0.5, 0.5))
×
1100
        return Waveform(bounds=(-round(edge, NDIGITS), round(edge,
×
1101
                                                             NDIGITS), +inf),
1102
                        seq=(_zero, rise, _one))
1103

1104

1105
def square(width: float, edge: float = 0, type: str = 'erf') -> Waveform:
12✔
1106
    if width <= 0:
12✔
1107
        return zero()
×
1108
    if edge == 0:
12✔
1109
        return Waveform(bounds=(round(-0.5 * width,
12✔
1110
                                      NDIGITS), round(0.5 * width,
1111
                                                      NDIGITS), +inf),
1112
                        seq=(_zero, _one, _zero))
1113
    else:
1114
        return ((step(edge, type=type) << width / 2) -
12✔
1115
                (step(edge, type=type) >> width / 2))
1116

1117

1118
def gaussian(width: float,
12✔
1119
             plateau: float = 0.0,
1120
             d: int | None = None) -> Waveform:
1121
    if width <= 0 and plateau <= 0.0:
12✔
1122
        return zero()
×
1123
    # width is two times FWHM
1124
    # std_sq2 = width / (4 * np.sqrt(np.log(2)))
1125
    std_sq2 = width / 3.3302184446307908
12✔
1126
    # std is set to give total pulse area same as a square
1127
    # std_sq2 = width/np.sqrt(np.pi)
1128
    if d is None:
12✔
1129
        base = lambda shift: basic_wave(GAUSSIAN, std_sq2, shift=shift)
12✔
1130
    else:
1131
        base = lambda shift: basic_wave(D_GAUSSIAN, std_sq2, d, shift=shift)
×
1132

1133
    if round(0.5 * plateau, NDIGITS) <= 0.0:
12✔
1134
        return Waveform(bounds=(round(-0.75 * width,
12✔
1135
                                      NDIGITS), round(0.75 * width,
1136
                                                      NDIGITS), +inf),
1137
                        seq=(_zero, base(0), _zero))
1138
    else:
1139
        return Waveform(bounds=(round(-0.75 * width - 0.5 * plateau,
×
1140
                                      NDIGITS), round(-0.5 * plateau, NDIGITS),
1141
                                round(0.5 * plateau, NDIGITS),
1142
                                round(0.75 * width + 0.5 * plateau,
1143
                                      NDIGITS), +inf),
1144
                        seq=(_zero, base(-0.5 * plateau), _one,
1145
                             base(0.5 * plateau), _zero))
1146

1147

1148
def cos(w: float, phi: float = 0) -> Waveform:
12✔
1149
    if w == 0:
12✔
1150
        return const(np.cos(phi))
×
1151
    if w < 0:
12✔
1152
        phi = -phi
×
1153
        w = -w
×
1154
    return Waveform(seq=(basic_wave(COS, w, shift=-phi / w), ))
12✔
1155

1156

1157
def sin(w: float, phi: float = 0) -> Waveform:
12✔
1158
    if w == 0:
12✔
1159
        return const(np.sin(phi))
×
1160
    if w < 0:
12✔
1161
        phi = -phi + pi
×
1162
        w = -w
×
1163
    return Waveform(seq=(basic_wave(COS, w, shift=(pi / 2 - phi) / w), ))
12✔
1164

1165

1166
def exp(alpha: float | complex) -> Waveform:
12✔
1167
    if isinstance(alpha, complex):
12✔
1168
        if alpha.real == 0:
12✔
1169
            return cos(alpha.imag) + 1j * sin(alpha.imag)
×
1170
        else:
1171
            return exp(alpha.real) * (cos(alpha.imag) + 1j * sin(alpha.imag))
12✔
1172
    else:
1173
        return Waveform(seq=(basic_wave(EXP, alpha), ))
12✔
1174

1175

1176
def sinc(bw: float) -> Waveform:
12✔
1177
    if bw <= 0:
×
1178
        return zero()
×
1179
    width = 100 / bw
×
1180
    return Waveform(bounds=(round(-0.5 * width,
×
1181
                                  NDIGITS), round(0.5 * width, NDIGITS), +inf),
1182
                    seq=(_zero, basic_wave(SINC, bw), _zero))
1183

1184

1185
def cosPulse(width: float, plateau: float = 0.0) -> Waveform:
12✔
1186
    # cos = basic_wave(COS, 2*np.pi/width)
1187
    # pulse = mul(add(cos, _one), _half)
1188
    if round(0.5 * plateau, NDIGITS) > 0:
×
1189
        return square(plateau + 0.5 * width, edge=0.5 * width, type='cos')
×
1190
    if width <= 0:
×
1191
        return zero()
×
1192
    pulse = ((((), ()), (((COS, 6.283185307179586 / width, 0), ), (1, ))),
×
1193
             (0.5, 0.5))
1194
    return Waveform(bounds=(round(-0.5 * width,
×
1195
                                  NDIGITS), round(0.5 * width, NDIGITS), +inf),
1196
                    seq=(_zero, pulse, _zero))
1197

1198

1199
def hanning(width: float, plateau: float = 0.0) -> Waveform:
12✔
1200
    return cosPulse(width, plateau=plateau)
×
1201

1202

1203
def cosh(w: float) -> Waveform:
12✔
1204
    return Waveform(seq=(basic_wave(COSH, w), ))
×
1205

1206

1207
def sinh(w: float) -> Waveform:
12✔
1208
    return Waveform(seq=(basic_wave(SINH, w), ))
×
1209

1210

1211
def coshPulse(width: float,
12✔
1212
              eps: float = 1.0,
1213
              plateau: float = 0.0) -> Waveform:
1214
    """Cosine hyperbolic pulse with the following im
1215

1216
    pulse edge shape:
1217
            cosh(eps / 2) - cosh(eps * t / T)
1218
    f(t) = -----------------------------------
1219
                  cosh(eps / 2) - 1
1220
    where T is the pulse width and eps is the pulse edge steepness.
1221
    The pulse is defined for t in [-T/2, T/2].
1222

1223
    In case of plateau > 0, the pulse is defined as:
1224
           | f(t + plateau/2)   if t in [-T/2 - plateau/2, - plateau/2]
1225
    g(t) = | 1                  if t in [-plateau/2, plateau/2]
1226
           | f(t - plateau/2)   if t in [plateau/2, T/2 + plateau/2]
1227

1228
    Parameters
1229
    ----------
1230
    width : float
1231
        Pulse width.
1232
    eps : float
1233
        Pulse edge steepness.
1234
    plateau : float
1235
        Pulse plateau.
1236
    """
1237
    if width <= 0 and plateau <= 0:
×
1238
        return zero()
×
1239
    w = eps / width
×
1240
    A = np.cosh(eps / 2)
×
1241

1242
    if plateau == 0.0 or round(-0.5 * plateau, NDIGITS) == round(
×
1243
            0.5 * plateau, NDIGITS):
1244
        pulse = ((((), ()), (((COSH, w, 0), ), (1, ))), (A / (A - 1),
×
1245
                                                         -1 / (A - 1)))
1246
        return Waveform(bounds=(round(-0.5 * width,
×
1247
                                      NDIGITS), round(0.5 * width,
1248
                                                      NDIGITS), +inf),
1249
                        seq=(_zero, pulse, _zero))
1250
    else:
1251
        raising = ((((), ()), (((COSH, w, -0.5 * plateau), ), (1, ))),
×
1252
                   (A / (A - 1), -1 / (A - 1)))
1253
        falling = ((((), ()), (((COSH, w, 0.5 * plateau), ), (1, ))),
×
1254
                   (A / (A - 1), -1 / (A - 1)))
1255
        return Waveform(bounds=(round(-0.5 * width - 0.5 * plateau,
×
1256
                                      NDIGITS), round(-0.5 * plateau, NDIGITS),
1257
                                round(0.5 * plateau, NDIGITS),
1258
                                round(0.5 * width + 0.5 * plateau,
1259
                                      NDIGITS), +inf),
1260
                        seq=(_zero, raising, _one, falling, _zero))
1261

1262

1263
def general_cosine(duration: float, *arg: float) -> Waveform:
12✔
1264
    wav = zero()
×
1265
    arg_ = np.asarray(arg)
×
1266
    arg_ /= arg_[::2].sum()
×
1267
    for i, a in enumerate(arg_, start=1):
×
1268
        wav += a / 2 * (1 - (-1)**i * cos(i * 2 * pi / duration))
×
1269
    return wav * square(duration)
×
1270

1271

1272
def slepian(duration: float, *arg: float) -> Waveform:
12✔
1273
    wav = zero()
×
1274
    arg_ = np.asarray(arg)
×
1275
    arg_ /= arg_[::2].sum()
×
1276
    for i, a in enumerate(arg_, start=1):
×
1277
        wav += a / 2 * (1 - (-1)**i * cos(i * 2 * pi / duration))
×
1278
    return wav * square(duration)
×
1279

1280

1281
def mollifier(width: float, plateau: float = 0.0, d: int = 0) -> Waveform:
12✔
1282
    """
1283
    Mollifier function is a smooth function that is 1 at the origin and 0 outside a certain radius.
1284
    It is defined as:
1285

1286
    f(x) = exp(1 / ((x / r) ^ 2 - 1) + 1)  in case |x| < r
1287
         = 0                           in case |x| >= r
1288
    where r = width / 2 is the radius of the mollifier.
1289

1290
    The parameter plateau is the width of the plateau.
1291
    The parameter d is the order of the derivative.
1292
    """
1293
    assert d >= 0 and isinstance(d, int), "d must be a non-negative integer"
×
1294
    assert width > 0, "width must be positive"
×
1295

1296
    if plateau <= 0:
×
1297
        return Waveform(bounds=(-0.5 * width, 0.5 * width, inf),
×
1298
                        seq=(_zero, basic_wave(MOLLIFIER, width / 2,
1299
                                               d), _zero))
1300
    else:
1301
        return Waveform(bounds=(-0.5 * width - 0.5 * plateau, -0.5 * plateau,
×
1302
                                0.5 * plateau, 0.5 * width + 0.5 * plateau,
1303
                                inf),
1304
                        seq=(_zero,
1305
                             basic_wave(MOLLIFIER,
1306
                                        width / 2,
1307
                                        d,
1308
                                        shift=-0.5 * plateau), _one,
1309
                             basic_wave(MOLLIFIER,
1310
                                        width / 2,
1311
                                        d,
1312
                                        shift=0.5 * plateau), _zero))
1313

1314

1315
def _poly(*a):
12✔
1316
    """
1317
    a[0] + a[1] * t + a[2] * t**2 + ...
1318
    """
1319
    t = []
12✔
1320
    amp = []
12✔
1321
    if a[0] != 0:
12✔
1322
        t.append(((), ()))
12✔
1323
        amp.append(a[0])
12✔
1324
    for n, a_ in enumerate(a[1:], start=1):
12✔
1325
        if a_ != 0:
12✔
1326
            t.append((((LINEAR, 0), ), (n, )))
12✔
1327
            amp.append(a_)
12✔
1328
    return tuple(t), tuple(a)
12✔
1329

1330

1331
def poly(a):
12✔
1332
    """
1333
    a[0] + a[1] * t + a[2] * t**2 + ...
1334
    """
1335
    return Waveform(seq=(_poly(*a), ))
12✔
1336

1337

1338
def t():
12✔
1339
    return Waveform(seq=((((LINEAR, 0), ), (1, )), (1, )))
×
1340

1341

1342
def drag(freq: float,
12✔
1343
         width: float,
1344
         plateau: float = 0,
1345
         delta: float = 0,
1346
         block_freq: float | None = None,
1347
         phase: float = 0,
1348
         t0: float = 0) -> Waveform:
1349
    phase += pi * delta * (width + plateau)
×
1350
    if plateau <= 0:
×
1351
        return Waveform(seq=(_zero,
×
1352
                             basic_wave(DRAG, t0, freq, width, delta,
1353
                                        block_freq, phase), _zero),
1354
                        bounds=(round(t0, NDIGITS), round(t0 + width,
1355
                                                          NDIGITS), +inf))
1356
    elif width <= 0:
×
1357
        w = 2 * pi * (freq + delta)
×
1358
        return Waveform(
×
1359
            seq=(_zero,
1360
                 basic_wave(COS, w,
1361
                            shift=(phase + 2 * pi * delta * t0) / w), _zero),
1362
            bounds=(round(t0, NDIGITS), round(t0 + plateau, NDIGITS), +inf))
1363
    else:
1364
        w = 2 * pi * (freq + delta)
×
1365
        return Waveform(
×
1366
            seq=(_zero,
1367
                 basic_wave(DRAG, t0, freq, width, delta, block_freq, phase),
1368
                 basic_wave(COS, w, shift=(phase + 2 * pi * delta * t0) / w),
1369
                 basic_wave(DRAG, t0 + plateau, freq, width, delta, block_freq,
1370
                            phase - 2 * pi * delta * plateau), _zero),
1371
            bounds=(round(t0, NDIGITS), round(t0 + width / 2, NDIGITS),
1372
                    round(t0 + width / 2 + plateau,
1373
                          NDIGITS), round(t0 + width + plateau,
1374
                                          NDIGITS), +inf))
1375

1376

1377
def chirp(f0: float,
12✔
1378
          f1: float,
1379
          T: float,
1380
          phi0: float = 0,
1381
          type: str = 'linear') -> Waveform:
1382
    """
1383
    A chirp is a signal in which the frequency increases (up-chirp)
1384
    or decreases (down-chirp) with time. In some sources, the term
1385
    chirp is used interchangeably with sweep signal.
1386

1387
    type: "linear", "exponential", "hyperbolic"
1388
    """
1389
    if f0 == f1:
12✔
1390
        return sin(f0, phi0)
×
1391
    if T <= 0:
12✔
1392
        raise ValueError('T must be positive')
×
1393

1394
    if type == 'linear':
12✔
1395
        # f(t) = f1 * (t/T) + f0 * (1 - t/T)
1396
        return Waveform(bounds=(0, round(T, NDIGITS), +inf),
12✔
1397
                        seq=(_zero, basic_wave(LINEARCHIRP, f0, f1, T,
1398
                                               phi0), _zero))
1399
    elif type in ['exp', 'exponential', 'geometric']:
12✔
1400
        # f(t) = f0 * (f1/f0) ** (t/T)
1401
        if f0 == 0:
12✔
1402
            raise ValueError('f0 must be non-zero')
×
1403
        alpha = np.log(f1 / f0) / T
12✔
1404
        return Waveform(bounds=(0, round(T, NDIGITS), +inf),
12✔
1405
                        seq=(_zero,
1406
                             basic_wave(EXPONENTIALCHIRP, f0, alpha,
1407
                                        phi0), _zero))
1408
    elif type in ['hyperbolic', 'hyp']:
12✔
1409
        # f(t) = f0 * f1 / (f0 * (t/T) + f1 * (1-t/T))
1410
        if f0 * f1 == 0:
12✔
1411
            return const(np.sin(phi0))
×
1412
        k = (f0 - f1) / (f1 * T)
12✔
1413
        return Waveform(bounds=(0, round(T, NDIGITS), +inf),
12✔
1414
                        seq=(_zero, basic_wave(HYPERBOLICCHIRP, f0, k,
1415
                                               phi0), _zero))
1416
    else:
1417
        raise ValueError(f'unknown type {type}')
×
1418

1419

1420
def interp(x: NDArray[np.float64], y: NDArray[np.float64]) -> Waveform:
12✔
1421
    seq, bounds = [_zero], [x[0]]
×
1422
    for x1, x2, y1, y2 in zip(x[:-1], x[1:], y[:-1], y[1:]):
×
1423
        if x2 == x1:
×
1424
            continue
×
1425
        seq.append(
×
1426
            add(
1427
                mul(_const((y2 - y1) / (x2 - x1)), basic_wave(LINEAR,
1428
                                                              shift=x1)),
1429
                _const(y1)))
1430
        bounds.append(x2)
×
1431
    bounds.append(inf)
×
1432
    seq.append(_zero)
×
1433
    return Waveform(seq=tuple(seq),
×
1434
                    bounds=tuple(round(b, NDIGITS)
1435
                                 for b in bounds)).simplify()
1436

1437

1438
def cut(wav: Waveform,
12✔
1439
        start: float | None = None,
1440
        stop: float | None = None,
1441
        head: float | None = None,
1442
        tail: float | None = None,
1443
        min: float | None = None,
1444
        max: float | None = None) -> Waveform:
1445
    offset = 0
×
1446
    if start is not None and head is not None:
×
1447
        offset = head - cast(NDArray[np.float64], wav(np.array([1.0 * start
×
1448
                                                                ])))[0]
1449
    elif stop is not None and tail is not None:
×
1450
        offset = tail - cast(NDArray[np.float64], wav(np.array([1.0 * stop
×
1451
                                                                ])))[0]
1452
    wav = wav + offset
×
1453

1454
    if start is not None:
×
1455
        wav = wav * (step(0) >> start)
×
1456
    if stop is not None:
×
1457
        wav = wav * ((1 - step(0)) >> stop)
×
1458
    if min is not None:
×
1459
        wav.min = min
×
1460
    if max is not None:
×
1461
        wav.max = max
×
1462
    return wav
×
1463

1464

1465
def function(fun, *args, start=None, stop=None):
12✔
1466
    TYPEID = registerBaseFunc(fun)
×
1467
    seq = (basic_wave(TYPEID, *args), )
×
1468
    wav = Waveform(seq=seq)
×
1469
    if start is not None:
×
1470
        wav = wav * (step(0) >> start)
×
1471
    if stop is not None:
×
1472
        wav = wav * ((1 - step(0)) >> stop)
×
1473
    return wav
×
1474

1475

1476
def samplingPoints(start, stop, points):
12✔
1477
    return Waveform(bounds=(round(start, NDIGITS), round(stop, NDIGITS), inf),
×
1478
                    seq=(_zero, basic_wave(INTERP, start, stop,
1479
                                           tuple(points)), _zero))
1480

1481

1482
def mixing(I: Waveform,
12✔
1483
           Q: Waveform | None = None,
1484
           *,
1485
           phase: float = 0.0,
1486
           freq: float = 0.0,
1487
           ratioIQ: float = 1.0,
1488
           phaseDiff: float = 0.0,
1489
           block_freq: float | None = None,
1490
           DRAGScaling: float | None = None) -> tuple[Waveform, Waveform]:
1491
    """SSB or envelope mixing
1492
    """
1493
    if Q is None:
×
1494
        I = I
×
1495
        Q = zero()
×
1496

1497
    w = 2 * pi * freq
×
1498
    if freq != 0.0:
×
1499
        # SSB mixing
1500
        Iout = I * cos(w, -phase) + Q * sin(w, -phase)
×
1501
        Qout = -I * sin(w, -phase + phaseDiff) + Q * cos(w, -phase + phaseDiff)
×
1502
    else:
1503
        # envelope mixing
1504
        Iout = cast(Waveform, I * np.cos(-phase) + Q * np.sin(-phase))
×
1505
        Qout = cast(Waveform, -I * np.sin(-phase) + Q * np.cos(-phase))
×
1506

1507
    # apply DRAG
1508
    if block_freq is not None and block_freq != freq:
×
1509
        a = block_freq / (block_freq - freq)
×
1510
        b = 1 / (block_freq - freq)
×
1511
        I = a * Iout + b / (2 * pi) * D(Qout)
×
1512
        Q = a * Qout - b / (2 * pi) * D(Iout)
×
1513
        Iout, Qout = I, Q
×
1514
    elif DRAGScaling is not None and DRAGScaling != 0:
×
1515
        # 2 * pi * scaling * (freq - block_freq) = 1
1516
        I = (1 - w * DRAGScaling) * Iout - DRAGScaling * D(Qout)
×
1517
        Q = (1 - w * DRAGScaling) * Qout + DRAGScaling * D(Iout)
×
1518
        Iout, Qout = I, Q
×
1519

1520
    Qout = ratioIQ * Qout
×
1521

1522
    return Iout, Qout
×
1523

1524

1525
__all__ = [
12✔
1526
    'D', 'Waveform', 'chirp', 'const', 'cos', 'cosh', 'coshPulse', 'cosPulse',
1527
    'cut', 'drag', 'exp', 'function', 'gaussian', 'general_cosine', 'hanning',
1528
    'interp', 'mixing', 'mollifier', 'one', 'poly', 'registerBaseFunc',
1529
    'registerDerivative', 'samplingPoints', 'sign', 'sin', 'sinc', 'sinh',
1530
    'square', 'step', 't', 'zero'
1531
]
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc