| Philipp Le | f9a597a | 2022-05-25 00:59:19 +0200 | [diff] [blame] | 1 | # SPDX-License-Identifier: MPL-2.0 |
| 2 | # Copyright (c) 2022 Philipp Le <philipp@philipple.de>. |
| 3 | # This Source Code Form is subject to the terms of the Mozilla Public |
| 4 | # License, v. 2.0. If a copy of the MPL was not distributed with this |
| 5 | # file, You can obtain one at https://mozilla.org/MPL/2.0/. |
| 6 | |
| 7 | from __future__ import annotations |
| 8 | |
| 9 | from typing import List, Tuple, Callable |
| 10 | import numpy as np |
| 11 | import scipy.signal |
| 12 | import copy |
| 13 | from dataclasses import dataclass |
| 14 | |
| 15 | from .fft import abs_log_fft, swap_freq |
| 16 | |
| 17 | |
| 18 | @dataclass(init=False) |
| 19 | class Signal: |
| 20 | t: np.ndarray |
| 21 | signal: np.ndarray |
| 22 | sample_rate: float |
| 23 | |
| 24 | def __init__(self, t: np.ndarray, signal: np.ndarray, sample_rate: float): |
| 25 | assert len(t) == len(signal), 't must have the same length like signal' |
| 26 | assert np.var(np.max(t[1:] - t[:len(t)-1])) < 1e-9, 't must be equidistant' |
| 27 | assert (((np.max(t) - np.min(t)) / (len(t) - 1)) - sample_rate) < 1e-9, 't must be spaced with sample period' |
| 28 | self.t = t |
| 29 | self.signal = signal |
| 30 | self.sample_rate = sample_rate |
| 31 | |
| 32 | def fft(self, window_fn: Callable[[int], np.ndarray] = scipy.signal.windows.boxcar) -> Tuple[np.ndarray, np.ndarray]: |
| 33 | f = swap_freq(np.fft.fftfreq(self.t.shape[-1], 1.0/self.sample_rate)) |
| 34 | x = abs_log_fft(self.signal * window_fn(len(self.signal))) |
| 35 | return f, x |
| 36 | |
| 37 | def split(self, equiv_length: int) -> List[Signal]: |
| 38 | parts = [] |
| 39 | t = copy.deepcopy(self.t) |
| 40 | signal = copy.deepcopy(self.signal) |
| 41 | while len(t) > equiv_length: |
| 42 | parts.append( |
| 43 | Signal( |
| 44 | t=t[:equiv_length], |
| 45 | signal=signal[:equiv_length], |
| 46 | sample_rate=self.sample_rate |
| 47 | ) |
| 48 | ) |
| 49 | t = t[equiv_length:] |
| 50 | signal = signal[equiv_length:] |
| 51 | if len(t) > 0: |
| 52 | parts.append( |
| 53 | Signal( |
| 54 | t=t[:equiv_length], |
| 55 | signal=signal[:equiv_length], |
| 56 | sample_rate=self.sample_rate |
| 57 | ) |
| 58 | ) |
| 59 | return parts |
| 60 | |
| 61 | def __add__(self, other: Signal) -> Signal: |
| 62 | assert self.t == other.t, 'Time vectors must be the same' |
| 63 | assert self.sample_rate == other.sample_rate, 'Sampling rate must be equal' |
| 64 | return Signal( |
| 65 | t=copy.deepcopy(self.t), |
| 66 | signal=self.signal + other.signal, |
| 67 | sample_rate=self.sample_rate |
| 68 | ) |