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

feihoo87 / waveforms / 16717725805

04 Aug 2025 08:13AM UTC coverage: 52.583% (-0.3%) from 52.861%
16717725805

push

github

feihoo87
Update version number to 2.2.1 to reflect recent changes and improvements in the waveform module.

1 of 1 new or added line in 1 file covered. (100.0%)

362 existing lines in 2 files now uncovered.

1323 of 2516 relevant lines covered (52.58%)

6.31 hits per line

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

45.53
/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
×
UNCOV
25
    x = Fraction(spec * num).limit_denominator(1000000000)
×
UNCOV
26
    if x.denominator <= 24:
×
UNCOV
27
        return True, x, -1
×
28
    return False, x, 0
×
29

30

31
def _spec_num_latex(num):
12✔
UNCOV
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:
×
UNCOV
42
                spec_latex = f"\\frac{{{1}}}{{{spec_latex}}}"
×
43
            if x.denominator == 1:
×
UNCOV
44
                if x.numerator == 1:
×
45
                    return f"{spec_latex}"
×
46
                else:
UNCOV
47
                    return f"{x.numerator:g}{spec_latex}"
×
48
            else:
49
                if x.numerator < 0:
×
UNCOV
50
                    return f"-\\frac{{{-x.numerator}}}{{{x.denominator}}}{spec_latex}"
×
51
                else:
UNCOV
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:
×
UNCOV
69
        a, n = s.split("e")
×
UNCOV
70
        n = float(n)
×
UNCOV
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 = ""
×
UNCOV
81
        elif shift[0] != '-':
×
UNCOV
82
            shift = "+" + shift
×
UNCOV
83
        return r"\mathrm{Func}" + f"{funID}(t{shift}, ...)"
×
UNCOV
84
    return _baseFunc_latex[funID](shift, *args)
×
85

86

87
def _wav_latex(wav):
12✔
88

UNCOV
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)
×
UNCOV
103
        for fun, n in zip(*mul):
×
104
            fun_expr = _fun_latex(fun)
×
105
            if n != 1:
×
UNCOV
106
                mul_expr.append(fun_expr + "^{" + f"{n}" + "}")
×
107
            else:
108
                mul_expr.append(fun_expr)
×
109
        sum_expr.append(''.join(mul_expr))
×
110

UNCOV
111
    ret = sum_expr[0]
×
112
    for expr in sum_expr[1:]:
×
113
        if expr[0] == '-':
×
UNCOV
114
            ret += expr
×
115
        else:
UNCOV
116
            ret += "+" + expr
×
UNCOV
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:
×
UNCOV
139
                if i == 0:
×
UNCOV
140
                    return -inf
×
UNCOV
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:
×
UNCOV
149
                if i == 0:
×
UNCOV
150
                    return inf
×
UNCOV
151
                return bounds[N - i - 1]
×
152
        return -inf
×
153

154
    @property
12✔
155
    def begin(self):
12✔
UNCOV
156
        if self.start is None:
×
UNCOV
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✔
UNCOV
163
        if self.stop is None:
×
UNCOV
164
            return self._end(self.bounds, self.seq)
×
165
        else:
UNCOV
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✔
UNCOV
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✔
UNCOV
191
                    sos = np.array(sos)
×
192
                elif not sos.flags.writeable:
12✔
UNCOV
193
                    sos = sos.copy()
×
194
                if initial:
12✔
UNCOV
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:
UNCOV
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:
×
UNCOV
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
×
UNCOV
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)
×
UNCOV
223
                size = round((stop - start) * sample_rate)
×
224
            else:
225
                stop = start + length
×
226
                size = chunk_size
×
UNCOV
227
            x = np.linspace(start, stop, size, endpoint=False)
×
228

UNCOV
229
            if filters is None:
×
UNCOV
230
                if out is not None:
×
UNCOV
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)
×
UNCOV
245
                if initial:
×
246
                    sig += initial
×
247
                if out is not None:
×
UNCOV
248
                    out[start_n:start_n + size] = sig
×
UNCOV
249
                yield cast(np.ndarray, sig)
×
250

UNCOV
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✔
UNCOV
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✔
UNCOV
279
            except:
×
UNCOV
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✔
UNCOV
385
                seq.pop()
×
UNCOV
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✔
UNCOV
392
        seq = []
×
UNCOV
393
        for expr in self.seq:
×
UNCOV
394
            seq.append(filter(expr, low, high, eps))
×
UNCOV
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✔
UNCOV
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)
×
UNCOV
419
        w = self.marker + other.marker
×
420

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

UNCOV
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)
×
UNCOV
435
        w = self.marker + other.marker
×
436

UNCOV
437
        def _and(a, b):
×
438
            if a != _zero and b != _zero:
×
UNCOV
439
                return _one
×
440
            else:
UNCOV
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

UNCOV
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:
×
UNCOV
469
                in_wave = False
×
470
                b = b - edge
×
471
                if b > bounds[-1]:
×
472
                    bounds.append(b)
×
UNCOV
473
                    seq.append(_zero)
×
474
                else:
UNCOV
475
                    bounds.pop()
×
UNCOV
476
                    bounds.append(b)
×
UNCOV
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:
UNCOV
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✔
UNCOV
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✔
UNCOV
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
UNCOV
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✔
UNCOV
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✔
UNCOV
545
            elif not accumulate:
×
546
                out *= 0
×
547
            self._fill_parts(parts, out)
12✔
548
        else:
UNCOV
549
            if out is None:
×
550
                return cast(list, parts)
×
551
            else:
UNCOV
552
                out = cast(list, out)
×
UNCOV
553
                if not accumulate:
×
554
                    out.clear()
×
UNCOV
555
                    out.extend(parts)
×
556
                else:
UNCOV
557
                    self._merge_parts(parts, out)
×
558
        return out
12✔
559

560
    def __hash__(self):
12✔
UNCOV
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✔
UNCOV
577
        parts = []
×
578
        start = -np.inf
×
579
        for end, wav in zip(self.bounds, self.seq):
×
580
            e_str = _wav_latex(wav)
×
UNCOV
581
            start_str = _num_latex(start)
×
582
            end_str = _num_latex(end)
×
UNCOV
583
            parts.append(e_str + r",~~&t\in" + f"({start_str},{end_str}" +
×
584
                         (']' if end < np.inf else ')'))
UNCOV
585
            start = end
×
586
        if len(parts) == 1:
×
UNCOV
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
            ])
UNCOV
593
        return "$$\n{}\n$$".format(expr)
×
594

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

598
        CHUNK = 1024
×
599
        RATE = 48000
×
600

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

604
        p = pyaudio.PyAudio()
×
UNCOV
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:
×
UNCOV
615
                        dynamic_volume = 1.0 / lim
×
616
                        amp = 2**15 * 0.99 * volume * dynamic_volume
×
UNCOV
617
                    data = (amp * data).astype(np.int16)
×
UNCOV
618
                    stream.write(bytes(data.data))
×
619
            finally:
620
                stream.stop_stream()
×
UNCOV
621
                stream.close()
×
622
        finally:
623
            p.terminate()
×
624

625
    def play(self, time_unit=1, volume=1.0):
12✔
UNCOV
626
        import multiprocessing as mp
×
UNCOV
627
        p = mp.Process(target=self._play,
×
628
                       args=(time_unit, volume),
629
                       daemon=True)
UNCOV
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✔
UNCOV
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✔
UNCOV
669
        if self.stop is None:
×
UNCOV
670
            return self.__end()
×
671
        else:
UNCOV
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:
UNCOV
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✔
UNCOV
735
            wav += self.offset
×
736
        if self.shift != 0:
12✔
UNCOV
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
        return wav
12✔
743

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

752
    def __rshift__(self, time):
12✔
753
        ret = WaveVStack()
12✔
754
        ret.wlist = self.wlist
12✔
755
        ret.sample_rate = self.sample_rate
12✔
756
        ret.start = self.start
12✔
757
        ret.stop = self.stop
12✔
758
        ret.shift = self.shift + time
12✔
759
        ret.offset = self.offset
12✔
760
        return ret
12✔
761

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

780
    def __radd__(self, v) -> WaveVStack:
12✔
781
        return self + v
×
782

783
    def __mul__(self, other) -> WaveVStack:
12✔
784
        if isinstance(other, Waveform):
12✔
785
            other = other.simplify() << self.shift
12✔
786
            ret = WaveVStack([Waveform(*w) * other for w in self.wlist])
12✔
787
            if self.offset != 0:
12✔
UNCOV
788
                w = other * self.offset
×
789
                ret.wlist.append((w.bounds, w.seq))
×
790
            return ret
12✔
791
        else:
792
            ret = WaveVStack([Waveform(*w) * other for w in self.wlist])
×
793
            ret.offset = self.offset * other
×
UNCOV
794
            return ret
×
795

796
    def __rmul__(self, v) -> WaveVStack:
12✔
UNCOV
797
        return self * v
×
798

799
    def __eq__(self, other) -> bool:
12✔
UNCOV
800
        if self.wlist:
×
801
            return False
×
802
        else:
803
            return zero() == other
×
804

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

808
    def __getstate__(self) -> tuple:
12✔
UNCOV
809
        function_lib = self.function_lib
×
UNCOV
810
        if function_lib:
×
UNCOV
811
            try:
×
UNCOV
812
                import dill
×
813
                function_lib = dill.dumps(function_lib)
×
UNCOV
814
            except:
×
815
                function_lib = None
×
816
        return (self.wlist, self.start, self.stop, self.sample_rate,
×
817
                self.offset, self.shift, self.filters, self.label,
818
                function_lib)
819

820
    def __setstate__(self, state: tuple) -> None:
12✔
821
        (self.wlist, self.start, self.stop, self.sample_rate, self.offset,
×
822
         self.shift, self.filters, self.label, function_lib) = state
UNCOV
823
        if function_lib:
×
UNCOV
824
            try:
×
825
                import dill
×
UNCOV
826
                function_lib = dill.loads(function_lib)
×
827
            except:
×
UNCOV
828
                function_lib = None
×
829
        self.function_lib = function_lib
×
830

831

832
def play(data, rate=48000):
12✔
833
    import io
×
834

UNCOV
835
    import pyaudio
×
836

837
    CHUNK = 1024
×
838

UNCOV
839
    max_amp = np.max(np.abs(data))
×
840

841
    if max_amp > 1:
×
UNCOV
842
        data /= max_amp
×
843

UNCOV
844
    data = np.array(2**15 * 0.999 * data, dtype=np.int16)
×
845
    buff = io.BytesIO(data.data)
×
846
    p = pyaudio.PyAudio()
×
847

848
    try:
×
849
        stream = p.open(format=pyaudio.paInt16,
×
850
                        channels=1,
851
                        rate=rate,
852
                        output=True)
853
        try:
×
854
            while True:
×
UNCOV
855
                data = buff.read(CHUNK)
×
856
                if data:
×
UNCOV
857
                    stream.write(data)
×
858
                else:
UNCOV
859
                    break
×
860
        finally:
UNCOV
861
            stream.stop_stream()
×
UNCOV
862
            stream.close()
×
863
    finally:
UNCOV
864
        p.terminate()
×
865

866

867
_zero_waveform = Waveform()
12✔
868
_one_waveform = Waveform(seq=(_one, ))
12✔
869

870

871
def zero():
12✔
872
    return _zero_waveform
12✔
873

874

875
def one():
12✔
876
    return _one_waveform
12✔
877

878

879
def const(c):
12✔
880
    return Waveform(seq=(_const(1.0 * c), ))
12✔
881

882

883
# register base function
884
def _format_LINEAR(shift, *args):
12✔
UNCOV
885
    if shift != 0:
×
UNCOV
886
        shift = _num_latex(-shift)
×
UNCOV
887
        if shift[0] == '-':
×
888
            return f"(t{shift})"
×
889
        else:
890
            return f"(t+{shift})"
×
891
    else:
892
        return 't'
×
893

894

895
def _format_GAUSSIAN(shift, *args):
12✔
UNCOV
896
    sigma = _num_latex(args[0] / np.sqrt(2))
×
897
    shift = _num_latex(-shift)
×
UNCOV
898
    if shift != '0':
×
UNCOV
899
        if shift[0] != '-':
×
900
            shift = '+' + shift
×
901
        if sigma == '1':
×
UNCOV
902
            return ('\\exp\\left[-\\frac{\\left(t' + shift +
×
903
                    '\\right)^2}{2}\\right]')
904
        else:
UNCOV
905
            return ('\\exp\\left[-\\frac{1}{2}\\left(\\frac{t' + shift + '}{' +
×
906
                    sigma + '}\\right)^2\\right]')
907
    else:
908
        if sigma == '1':
×
909
            return ('\\exp\\left(-\\frac{t^2}{2}\\right)')
×
910
        else:
911
            return ('\\exp\\left[-\\frac{1}{2}\\left(\\frac{t}{' + sigma +
×
912
                    '}\\right)^2\\right]')
913

914

915
def _format_SINC(shift, *args):
12✔
916
    shift = _num_latex(-shift)
×
UNCOV
917
    bw = _num_latex(args[0])
×
918
    if shift != '0':
×
919
        if shift[0] != '-':
×
UNCOV
920
            shift = '+' + shift
×
921
        if bw == '1':
×
UNCOV
922
            return '\\mathrm{sinc}(t' + shift + ')'
×
923
        else:
UNCOV
924
            return '\\mathrm{sinc}[' + bw + '(t' + shift + ')]'
×
925
    else:
926
        if bw == '1':
×
927
            return '\\mathrm{sinc}(t)'
×
928
        else:
929
            return '\\mathrm{sinc}(' + bw + 't)'
×
930

931

932
def _format_COSINE(shift, *args):
12✔
933
    freq = args[0] / 2 / np.pi
×
934
    phase = -shift * freq
×
935
    freq = _num_latex(freq)
×
936
    if freq == '1':
×
937
        freq = ''
×
938
    phase = _num_latex(phase)
×
UNCOV
939
    if phase == '0':
×
940
        phase = ''
×
UNCOV
941
    elif phase[0] != '-':
×
UNCOV
942
        phase = '+' + phase
×
UNCOV
943
    if phase != '':
×
944
        return f'\\cos\\left[2\\pi\\left({freq}t{phase}\\right)\\right]'
×
945
    elif freq != '':
×
946
        return f'\\cos\\left(2\\pi\\times {freq}t\\right)'
×
947
    else:
UNCOV
948
        return '\\cos\\left(2\\pi t\\right)'
×
949

950

951
def _format_ERF(shift, *args):
12✔
UNCOV
952
    if shift > 0:
×
953
        return '\\mathrm{erf}(\\frac{t-' + f"{_num_latex(shift)}" + '}{' + f'{args[0]:g}' + '})'
×
954
    elif shift < 0:
×
955
        return '\\mathrm{erf}(\\frac{t+' + f"{_num_latex(-shift)}" + '}{' + f'{args[0]:g}' + '})'
×
956
    else:
UNCOV
957
        return '\\mathrm{erf}(\\frac{t}{' + f'{args[0]:g}' + '})'
×
958

959

960
def _format_COSH(shift, *args):
12✔
UNCOV
961
    if shift > 0:
×
962
        return '\\cosh(\\frac{t-' + f"{_num_latex(shift)}" + '}{' + f'{1/args[0]:g}' + '})'
×
963
    elif shift < 0:
×
964
        return '\\cosh(\\frac{t+' + f"{_num_latex(-shift)}" + '}{' + f'{1/args[0]:g}' + '})'
×
965
    else:
UNCOV
966
        return '\\cosh(\\frac{t}{' + f'{1/args[0]:g}' + '})'
×
967

968

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

977

978
def _format_EXP(shift, *args):
12✔
UNCOV
979
    if _num_latex(shift) and shift > 0:
×
980
        return '\\exp\\left(-' + f'{args[0]:g}' + '\\left(t-' + f"{_num_latex(shift)}" + '\\right)\\right)'
×
UNCOV
981
    elif _num_latex(-shift) and shift < 0:
×
UNCOV
982
        return '\\exp\\left(-' + f'{args[0]:g}' + '\\left(t+' + f"{_num_latex(-shift)}" + '\\right)\\right)'
×
983
    else:
984
        return '\\exp\\left(-' + f'{args[0]:g}' + 't\\right)'
×
985

986

987
def _format_DRAG(shift, *args):
12✔
988
    return f"DRAG(...)"
×
989

990

991
def _format_MOLLIFIER(shift, *args):
12✔
992
    r = _num_latex(args[0])
×
993
    d = _num_latex(args[1])
×
994
    shift_str = _num_latex(-shift)
×
995
    if shift_str == '0':
×
996
        shift_str = ''
×
997
    elif shift_str[0] != '-':
×
UNCOV
998
        shift_str = '+' + shift_str
×
999

UNCOV
1000
    if d == '0':
×
UNCOV
1001
        return f"\\mathrm{{Mollifier}}\\left(t{shift_str}, r={r}\\right)"
×
UNCOV
1002
    elif d == '1':
×
UNCOV
1003
        return f"\\mathrm{{Mollifier}}'\\left(t{shift_str}, r={r}\\right)"
×
UNCOV
1004
    elif d == '2':
×
UNCOV
1005
        return f"\\mathrm{{Mollifier}}''\\left(t{shift_str}, r={r}\\right)"
×
1006
    else:
UNCOV
1007
        return f"\\mathrm{{Mollifier}}^{{({d})}}\\left(t{shift_str}, r={r}\\right)"
×
1008

1009

1010
def _format_D_GAUSSIAN(shift, *args):
12✔
UNCOV
1011
    sigma = _num_latex(args[0] / np.sqrt(2))
×
UNCOV
1012
    d = args[1]
×
UNCOV
1013
    shift_str = _num_latex(-shift)
×
UNCOV
1014
    if shift_str == '0':
×
UNCOV
1015
        shift_str = ''
×
UNCOV
1016
    elif shift_str[0] != '-':
×
UNCOV
1017
        shift_str = '+' + shift_str
×
1018

UNCOV
1019
    if d == 0:
×
UNCOV
1020
        return f"\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
×
UNCOV
1021
    elif d == 1:
×
UNCOV
1022
        return f"\\frac{{\\mathrm{{d}}}}{{\\mathrm{{d}}t}}\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
×
1023
    else:
1024
        return f"\\frac{{\\mathrm{{d}}^{{{d}}}}}{{\\mathrm{{d}}t^{{{d}}}}}\\mathrm{{Gaussian}}\\left(t{shift_str}, \\sigma={sigma}\\right)"
×
1025

1026

1027
registerBaseFuncLatex(LINEAR, _format_LINEAR)
12✔
1028
registerBaseFuncLatex(GAUSSIAN, _format_GAUSSIAN)
12✔
1029
registerBaseFuncLatex(ERF, _format_ERF)
12✔
1030
registerBaseFuncLatex(COS, _format_COSINE)
12✔
1031
registerBaseFuncLatex(SINC, _format_SINC)
12✔
1032
registerBaseFuncLatex(EXP, _format_EXP)
12✔
1033
registerBaseFuncLatex(COSH, _format_COSH)
12✔
1034
registerBaseFuncLatex(SINH, _format_SINH)
12✔
1035
registerBaseFuncLatex(DRAG, _format_DRAG)
12✔
1036
registerBaseFuncLatex(MOLLIFIER, _format_MOLLIFIER)
12✔
1037
registerBaseFuncLatex(D_GAUSSIAN, _format_D_GAUSSIAN)
12✔
1038

1039

1040
def D(wav: Waveform, d: int = 1) -> Waveform:
12✔
1041
    """derivative
1042

1043
    Parameters
1044
    ----------
1045
    wav : Waveform
1046
        The waveform to take the derivative of.
1047
    d : int, optional
1048
        The order of the derivative, by default 1.
1049
    """
1050
    assert d >= 0 and isinstance(d, int), "d must be a non-negative integer"
×
UNCOV
1051
    if d == 0:
×
UNCOV
1052
        return wav
×
UNCOV
1053
    elif d == 1:
×
UNCOV
1054
        return Waveform(bounds=wav.bounds, seq=tuple(_D(x) for x in wav.seq))
×
1055
    else:
UNCOV
1056
        return D(D(wav, d - 1), 1)
×
1057

1058

1059
def convolve(a, b):
12✔
UNCOV
1060
    pass
×
1061

1062

1063
def sign():
12✔
1064
    return Waveform(bounds=(0, +inf), seq=(_const(-1), _one))
×
1065

1066

1067
def step(edge, type='erf'):
12✔
1068
    """
1069
    type: "erf", "cos", "linear"
1070
    """
1071
    if edge == 0:
12✔
1072
        return Waveform(bounds=(0, +inf), seq=(_zero, _one))
12✔
1073
    if type == 'cos':
12✔
UNCOV
1074
        rise = add(_half,
×
1075
                   mul(_half, basic_wave(COS, pi / edge, shift=0.5 * edge)))
UNCOV
1076
        return Waveform(bounds=(round(-edge / 2,
×
1077
                                      NDIGITS), round(edge / 2,
1078
                                                      NDIGITS), +inf),
1079
                        seq=(_zero, rise, _one))
1080
    elif type == 'linear':
12✔
1081
        rise = add(_half, mul(_const(1 / edge), basic_wave(LINEAR)))
12✔
1082
        return Waveform(bounds=(round(-edge / 2,
12✔
1083
                                      NDIGITS), round(edge / 2,
1084
                                                      NDIGITS), +inf),
1085
                        seq=(_zero, rise, _one))
1086
    else:
UNCOV
1087
        std_sq2 = edge / 5
×
1088
        # rise = add(_half, mul(_half, basic_wave(ERF, std_sq2)))
UNCOV
1089
        rise = ((((), ()), (((ERF, std_sq2, 0), ), (1, ))), (0.5, 0.5))
×
UNCOV
1090
        return Waveform(bounds=(-round(edge, NDIGITS), round(edge,
×
1091
                                                             NDIGITS), +inf),
1092
                        seq=(_zero, rise, _one))
1093

1094

1095
def square(width: float, edge: float = 0, type: str = 'erf') -> Waveform:
12✔
1096
    if width <= 0:
12✔
UNCOV
1097
        return zero()
×
1098
    if edge == 0:
12✔
1099
        return Waveform(bounds=(round(-0.5 * width,
12✔
1100
                                      NDIGITS), round(0.5 * width,
1101
                                                      NDIGITS), +inf),
1102
                        seq=(_zero, _one, _zero))
1103
    else:
1104
        return ((step(edge, type=type) << width / 2) -
12✔
1105
                (step(edge, type=type) >> width / 2))
1106

1107

1108
def gaussian(width: float,
12✔
1109
             plateau: float = 0.0,
1110
             d: int | None = None) -> Waveform:
1111
    if width <= 0 and plateau <= 0.0:
12✔
UNCOV
1112
        return zero()
×
1113
    # width is two times FWHM
1114
    # std_sq2 = width / (4 * np.sqrt(np.log(2)))
1115
    std_sq2 = width / 3.3302184446307908
12✔
1116
    # std is set to give total pulse area same as a square
1117
    # std_sq2 = width/np.sqrt(np.pi)
1118
    if d is None:
12✔
1119
        base = lambda shift: basic_wave(GAUSSIAN, std_sq2, shift=shift)
12✔
1120
    else:
UNCOV
1121
        base = lambda shift: basic_wave(D_GAUSSIAN, std_sq2, d, shift=shift)
×
1122

1123
    if round(0.5 * plateau, NDIGITS) <= 0.0:
12✔
1124
        return Waveform(bounds=(round(-0.75 * width,
12✔
1125
                                      NDIGITS), round(0.75 * width,
1126
                                                      NDIGITS), +inf),
1127
                        seq=(_zero, base(0), _zero))
1128
    else:
UNCOV
1129
        return Waveform(bounds=(round(-0.75 * width - 0.5 * plateau,
×
1130
                                      NDIGITS), round(-0.5 * plateau, NDIGITS),
1131
                                round(0.5 * plateau, NDIGITS),
1132
                                round(0.75 * width + 0.5 * plateau,
1133
                                      NDIGITS), +inf),
1134
                        seq=(_zero, base(-0.5 * plateau), _one,
1135
                             base(0.5 * plateau), _zero))
1136

1137

1138
def cos(w: float, phi: float = 0) -> Waveform:
12✔
1139
    if w == 0:
12✔
1140
        return const(np.cos(phi))
×
1141
    if w < 0:
12✔
UNCOV
1142
        phi = -phi
×
UNCOV
1143
        w = -w
×
1144
    return Waveform(seq=(basic_wave(COS, w, shift=-phi / w), ))
12✔
1145

1146

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

1155

1156
def exp(alpha: float | complex) -> Waveform:
12✔
1157
    if isinstance(alpha, complex):
12✔
1158
        if alpha.real == 0:
12✔
UNCOV
1159
            return cos(alpha.imag) + 1j * sin(alpha.imag)
×
1160
        else:
1161
            return exp(alpha.real) * (cos(alpha.imag) + 1j * sin(alpha.imag))
12✔
1162
    else:
1163
        return Waveform(seq=(basic_wave(EXP, alpha), ))
12✔
1164

1165

1166
def sinc(bw: float) -> Waveform:
12✔
UNCOV
1167
    if bw <= 0:
×
UNCOV
1168
        return zero()
×
1169
    width = 100 / bw
×
UNCOV
1170
    return Waveform(bounds=(round(-0.5 * width,
×
1171
                                  NDIGITS), round(0.5 * width, NDIGITS), +inf),
1172
                    seq=(_zero, basic_wave(SINC, bw), _zero))
1173

1174

1175
def cosPulse(width: float, plateau: float = 0.0) -> Waveform:
12✔
1176
    # cos = basic_wave(COS, 2*np.pi/width)
1177
    # pulse = mul(add(cos, _one), _half)
UNCOV
1178
    if round(0.5 * plateau, NDIGITS) > 0:
×
UNCOV
1179
        return square(plateau + 0.5 * width, edge=0.5 * width, type='cos')
×
UNCOV
1180
    if width <= 0:
×
UNCOV
1181
        return zero()
×
UNCOV
1182
    pulse = ((((), ()), (((COS, 6.283185307179586 / width, 0), ), (1, ))),
×
1183
             (0.5, 0.5))
UNCOV
1184
    return Waveform(bounds=(round(-0.5 * width,
×
1185
                                  NDIGITS), round(0.5 * width, NDIGITS), +inf),
1186
                    seq=(_zero, pulse, _zero))
1187

1188

1189
def hanning(width: float, plateau: float = 0.0) -> Waveform:
12✔
UNCOV
1190
    return cosPulse(width, plateau=plateau)
×
1191

1192

1193
def cosh(w: float) -> Waveform:
12✔
UNCOV
1194
    return Waveform(seq=(basic_wave(COSH, w), ))
×
1195

1196

1197
def sinh(w: float) -> Waveform:
12✔
1198
    return Waveform(seq=(basic_wave(SINH, w), ))
×
1199

1200

1201
def coshPulse(width: float,
12✔
1202
              eps: float = 1.0,
1203
              plateau: float = 0.0) -> Waveform:
1204
    """Cosine hyperbolic pulse with the following im
1205

1206
    pulse edge shape:
1207
            cosh(eps / 2) - cosh(eps * t / T)
1208
    f(t) = -----------------------------------
1209
                  cosh(eps / 2) - 1
1210
    where T is the pulse width and eps is the pulse edge steepness.
1211
    The pulse is defined for t in [-T/2, T/2].
1212

1213
    In case of plateau > 0, the pulse is defined as:
1214
           | f(t + plateau/2)   if t in [-T/2 - plateau/2, - plateau/2]
1215
    g(t) = | 1                  if t in [-plateau/2, plateau/2]
1216
           | f(t - plateau/2)   if t in [plateau/2, T/2 + plateau/2]
1217

1218
    Parameters
1219
    ----------
1220
    width : float
1221
        Pulse width.
1222
    eps : float
1223
        Pulse edge steepness.
1224
    plateau : float
1225
        Pulse plateau.
1226
    """
1227
    if width <= 0 and plateau <= 0:
×
1228
        return zero()
×
UNCOV
1229
    w = eps / width
×
UNCOV
1230
    A = np.cosh(eps / 2)
×
1231

1232
    if plateau == 0.0 or round(-0.5 * plateau, NDIGITS) == round(
×
1233
            0.5 * plateau, NDIGITS):
1234
        pulse = ((((), ()), (((COSH, w, 0), ), (1, ))), (A / (A - 1),
×
1235
                                                         -1 / (A - 1)))
1236
        return Waveform(bounds=(round(-0.5 * width,
×
1237
                                      NDIGITS), round(0.5 * width,
1238
                                                      NDIGITS), +inf),
1239
                        seq=(_zero, pulse, _zero))
1240
    else:
UNCOV
1241
        raising = ((((), ()), (((COSH, w, -0.5 * plateau), ), (1, ))),
×
1242
                   (A / (A - 1), -1 / (A - 1)))
UNCOV
1243
        falling = ((((), ()), (((COSH, w, 0.5 * plateau), ), (1, ))),
×
1244
                   (A / (A - 1), -1 / (A - 1)))
UNCOV
1245
        return Waveform(bounds=(round(-0.5 * width - 0.5 * plateau,
×
1246
                                      NDIGITS), round(-0.5 * plateau, NDIGITS),
1247
                                round(0.5 * plateau, NDIGITS),
1248
                                round(0.5 * width + 0.5 * plateau,
1249
                                      NDIGITS), +inf),
1250
                        seq=(_zero, raising, _one, falling, _zero))
1251

1252

1253
def general_cosine(duration: float, *arg: float) -> Waveform:
12✔
UNCOV
1254
    wav = zero()
×
1255
    arg_ = np.asarray(arg)
×
1256
    arg_ /= arg_[::2].sum()
×
UNCOV
1257
    for i, a in enumerate(arg_, start=1):
×
UNCOV
1258
        wav += a / 2 * (1 - (-1)**i * cos(i * 2 * pi / duration))
×
UNCOV
1259
    return wav * square(duration)
×
1260

1261

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

1270

1271
def mollifier(width: float, plateau: float = 0.0, d: int = 0) -> Waveform:
12✔
1272
    """
1273
    Mollifier function is a smooth function that is 1 at the origin and 0 outside a certain radius.
1274
    It is defined as:
1275

1276
    f(x) = exp(1 / ((x / r) ^ 2 - 1) + 1)  in case |x| < r
1277
         = 0                           in case |x| >= r
1278
    where r = width / 2 is the radius of the mollifier.
1279

1280
    The parameter plateau is the width of the plateau.
1281
    The parameter d is the order of the derivative.
1282
    """
UNCOV
1283
    assert d >= 0 and isinstance(d, int), "d must be a non-negative integer"
×
UNCOV
1284
    assert width > 0, "width must be positive"
×
1285

UNCOV
1286
    if plateau <= 0:
×
UNCOV
1287
        return Waveform(bounds=(-0.5 * width, 0.5 * width, inf),
×
1288
                        seq=(_zero, basic_wave(MOLLIFIER, width / 2,
1289
                                               d), _zero))
1290
    else:
UNCOV
1291
        return Waveform(bounds=(-0.5 * width - 0.5 * plateau, -0.5 * plateau,
×
1292
                                0.5 * plateau, 0.5 * width + 0.5 * plateau,
1293
                                inf),
1294
                        seq=(_zero,
1295
                             basic_wave(MOLLIFIER,
1296
                                        width / 2,
1297
                                        d,
1298
                                        shift=-0.5 * plateau), _one,
1299
                             basic_wave(MOLLIFIER,
1300
                                        width / 2,
1301
                                        d,
1302
                                        shift=0.5 * plateau), _zero))
1303

1304

1305
def _poly(*a):
12✔
1306
    """
1307
    a[0] + a[1] * t + a[2] * t**2 + ...
1308
    """
1309
    t = []
12✔
1310
    amp = []
12✔
1311
    if a[0] != 0:
12✔
1312
        t.append(((), ()))
12✔
1313
        amp.append(a[0])
12✔
1314
    for n, a_ in enumerate(a[1:], start=1):
12✔
1315
        if a_ != 0:
12✔
1316
            t.append((((LINEAR, 0), ), (n, )))
12✔
1317
            amp.append(a_)
12✔
1318
    return tuple(t), tuple(a)
12✔
1319

1320

1321
def poly(a):
12✔
1322
    """
1323
    a[0] + a[1] * t + a[2] * t**2 + ...
1324
    """
1325
    return Waveform(seq=(_poly(*a), ))
12✔
1326

1327

1328
def t():
12✔
UNCOV
1329
    return Waveform(seq=((((LINEAR, 0), ), (1, )), (1, )))
×
1330

1331

1332
def drag(freq: float,
12✔
1333
         width: float,
1334
         plateau: float = 0,
1335
         delta: float = 0,
1336
         block_freq: float | None = None,
1337
         phase: float = 0,
1338
         t0: float = 0) -> Waveform:
1339
    phase += pi * delta * (width + plateau)
×
UNCOV
1340
    if plateau <= 0:
×
1341
        return Waveform(seq=(_zero,
×
1342
                             basic_wave(DRAG, t0, freq, width, delta,
1343
                                        block_freq, phase), _zero),
1344
                        bounds=(round(t0, NDIGITS), round(t0 + width,
1345
                                                          NDIGITS), +inf))
UNCOV
1346
    elif width <= 0:
×
UNCOV
1347
        w = 2 * pi * (freq + delta)
×
UNCOV
1348
        return Waveform(
×
1349
            seq=(_zero,
1350
                 basic_wave(COS, w,
1351
                            shift=(phase + 2 * pi * delta * t0) / w), _zero),
1352
            bounds=(round(t0, NDIGITS), round(t0 + plateau, NDIGITS), +inf))
1353
    else:
UNCOV
1354
        w = 2 * pi * (freq + delta)
×
UNCOV
1355
        return Waveform(
×
1356
            seq=(_zero,
1357
                 basic_wave(DRAG, t0, freq, width, delta, block_freq, phase),
1358
                 basic_wave(COS, w, shift=(phase + 2 * pi * delta * t0) / w),
1359
                 basic_wave(DRAG, t0 + plateau, freq, width, delta, block_freq,
1360
                            phase - 2 * pi * delta * plateau), _zero),
1361
            bounds=(round(t0, NDIGITS), round(t0 + width / 2, NDIGITS),
1362
                    round(t0 + width / 2 + plateau,
1363
                          NDIGITS), round(t0 + width + plateau,
1364
                                          NDIGITS), +inf))
1365

1366

1367
def chirp(f0: float,
12✔
1368
          f1: float,
1369
          T: float,
1370
          phi0: float = 0,
1371
          type: str = 'linear') -> Waveform:
1372
    """
1373
    A chirp is a signal in which the frequency increases (up-chirp)
1374
    or decreases (down-chirp) with time. In some sources, the term
1375
    chirp is used interchangeably with sweep signal.
1376

1377
    type: "linear", "exponential", "hyperbolic"
1378
    """
1379
    if f0 == f1:
12✔
1380
        return sin(f0, phi0)
×
1381
    if T <= 0:
12✔
1382
        raise ValueError('T must be positive')
×
1383

1384
    if type == 'linear':
12✔
1385
        # f(t) = f1 * (t/T) + f0 * (1 - t/T)
1386
        return Waveform(bounds=(0, round(T, NDIGITS), +inf),
12✔
1387
                        seq=(_zero, basic_wave(LINEARCHIRP, f0, f1, T,
1388
                                               phi0), _zero))
1389
    elif type in ['exp', 'exponential', 'geometric']:
12✔
1390
        # f(t) = f0 * (f1/f0) ** (t/T)
1391
        if f0 == 0:
12✔
1392
            raise ValueError('f0 must be non-zero')
×
1393
        alpha = np.log(f1 / f0) / T
12✔
1394
        return Waveform(bounds=(0, round(T, NDIGITS), +inf),
12✔
1395
                        seq=(_zero,
1396
                             basic_wave(EXPONENTIALCHIRP, f0, alpha,
1397
                                        phi0), _zero))
1398
    elif type in ['hyperbolic', 'hyp']:
12✔
1399
        # f(t) = f0 * f1 / (f0 * (t/T) + f1 * (1-t/T))
1400
        if f0 * f1 == 0:
12✔
1401
            return const(np.sin(phi0))
×
1402
        k = (f0 - f1) / (f1 * T)
12✔
1403
        return Waveform(bounds=(0, round(T, NDIGITS), +inf),
12✔
1404
                        seq=(_zero, basic_wave(HYPERBOLICCHIRP, f0, k,
1405
                                               phi0), _zero))
1406
    else:
1407
        raise ValueError(f'unknown type {type}')
×
1408

1409

1410
def interp(x: NDArray[np.float64], y: NDArray[np.float64]) -> Waveform:
12✔
1411
    seq, bounds = [_zero], [x[0]]
×
1412
    for x1, x2, y1, y2 in zip(x[:-1], x[1:], y[:-1], y[1:]):
×
1413
        if x2 == x1:
×
1414
            continue
×
UNCOV
1415
        seq.append(
×
1416
            add(
1417
                mul(_const((y2 - y1) / (x2 - x1)), basic_wave(LINEAR,
1418
                                                              shift=x1)),
1419
                _const(y1)))
UNCOV
1420
        bounds.append(x2)
×
UNCOV
1421
    bounds.append(inf)
×
UNCOV
1422
    seq.append(_zero)
×
UNCOV
1423
    return Waveform(seq=tuple(seq),
×
1424
                    bounds=tuple(round(b, NDIGITS)
1425
                                 for b in bounds)).simplify()
1426

1427

1428
def cut(wav: Waveform,
12✔
1429
        start: float | None = None,
1430
        stop: float | None = None,
1431
        head: float | None = None,
1432
        tail: float | None = None,
1433
        min: float | None = None,
1434
        max: float | None = None) -> Waveform:
1435
    offset = 0
×
1436
    if start is not None and head is not None:
×
UNCOV
1437
        offset = head - cast(NDArray[np.float64], wav(np.array([1.0 * start
×
1438
                                                                ])))[0]
1439
    elif stop is not None and tail is not None:
×
UNCOV
1440
        offset = tail - cast(NDArray[np.float64], wav(np.array([1.0 * stop
×
1441
                                                                ])))[0]
1442
    wav = wav + offset
×
1443

UNCOV
1444
    if start is not None:
×
1445
        wav = wav * (step(0) >> start)
×
1446
    if stop is not None:
×
UNCOV
1447
        wav = wav * ((1 - step(0)) >> stop)
×
UNCOV
1448
    if min is not None:
×
1449
        wav.min = min
×
1450
    if max is not None:
×
1451
        wav.max = max
×
1452
    return wav
×
1453

1454

1455
def function(fun, *args, start=None, stop=None):
12✔
UNCOV
1456
    TYPEID = registerBaseFunc(fun)
×
1457
    seq = (basic_wave(TYPEID, *args), )
×
1458
    wav = Waveform(seq=seq)
×
1459
    if start is not None:
×
UNCOV
1460
        wav = wav * (step(0) >> start)
×
1461
    if stop is not None:
×
UNCOV
1462
        wav = wav * ((1 - step(0)) >> stop)
×
1463
    return wav
×
1464

1465

1466
def samplingPoints(start, stop, points):
12✔
UNCOV
1467
    return Waveform(bounds=(round(start, NDIGITS), round(stop, NDIGITS), inf),
×
1468
                    seq=(_zero, basic_wave(INTERP, start, stop,
1469
                                           tuple(points)), _zero))
1470

1471

1472
def mixing(I: Waveform,
12✔
1473
           Q: Waveform | None = None,
1474
           *,
1475
           phase: float = 0.0,
1476
           freq: float = 0.0,
1477
           ratioIQ: float = 1.0,
1478
           phaseDiff: float = 0.0,
1479
           block_freq: float | None = None,
1480
           DRAGScaling: float | None = None) -> tuple[Waveform, Waveform]:
1481
    """SSB or envelope mixing
1482
    """
UNCOV
1483
    if Q is None:
×
UNCOV
1484
        I = I
×
UNCOV
1485
        Q = zero()
×
1486

UNCOV
1487
    w = 2 * pi * freq
×
UNCOV
1488
    if freq != 0.0:
×
1489
        # SSB mixing
UNCOV
1490
        Iout = I * cos(w, -phase) + Q * sin(w, -phase)
×
UNCOV
1491
        Qout = -I * sin(w, -phase + phaseDiff) + Q * cos(w, -phase + phaseDiff)
×
1492
    else:
1493
        # envelope mixing
UNCOV
1494
        Iout = cast(Waveform, I * np.cos(-phase) + Q * np.sin(-phase))
×
UNCOV
1495
        Qout = cast(Waveform, -I * np.sin(-phase) + Q * np.cos(-phase))
×
1496

1497
    # apply DRAG
UNCOV
1498
    if block_freq is not None and block_freq != freq:
×
UNCOV
1499
        a = block_freq / (block_freq - freq)
×
UNCOV
1500
        b = 1 / (block_freq - freq)
×
UNCOV
1501
        I = a * Iout + b / (2 * pi) * D(Qout)
×
UNCOV
1502
        Q = a * Qout - b / (2 * pi) * D(Iout)
×
UNCOV
1503
        Iout, Qout = I, Q
×
UNCOV
1504
    elif DRAGScaling is not None and DRAGScaling != 0:
×
1505
        # 2 * pi * scaling * (freq - block_freq) = 1
UNCOV
1506
        I = (1 - w * DRAGScaling) * Iout - DRAGScaling * D(Qout)
×
UNCOV
1507
        Q = (1 - w * DRAGScaling) * Qout + DRAGScaling * D(Iout)
×
UNCOV
1508
        Iout, Qout = I, Q
×
1509

UNCOV
1510
    Qout = ratioIQ * Qout
×
1511

UNCOV
1512
    return Iout, Qout
×
1513

1514

1515
__all__ = [
12✔
1516
    'D', 'Waveform', 'chirp', 'const', 'cos', 'cosh', 'coshPulse', 'cosPulse',
1517
    'cut', 'drag', 'exp', 'function', 'gaussian', 'general_cosine', 'hanning',
1518
    'interp', 'mixing', 'mollifier', 'one', 'poly', 'registerBaseFunc',
1519
    'registerDerivative', 'samplingPoints', 'sign', 'sin', 'sinc', 'sinh',
1520
    'square', 'step', 't', 'zero'
1521
]
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