diff --git a/src/science_signal/signal.py b/src/science_signal/signal.py index c1bbc50..878ae05 100644 --- a/src/science_signal/signal.py +++ b/src/science_signal/signal.py @@ -1,7 +1,19 @@ from scipy.interpolate import CubicSpline # pyright: ignore[reportMissingTypeStubs] from numpy.typing import ArrayLike, NDArray -from numpy import float64, linspace, array - +from numpy import ( + append, + cos, + float64, + linspace, + array, + logical_and, + sin, + where, + arange, + zeros, +) +from scipy.signal import detrend, welch +from scipy.fft import rfftfreq, rfft, irfft from science_signal import interpolate @@ -45,6 +57,100 @@ class Signal: """ return CubicSpline(self.x, self.y) + def cos(self) -> "Signal": + """ + Return the cosinus of this signal + """ + return Signal( + self.x, + cos(self.y), + ) + + def cut(self, start: float, end: float) -> "Signal": + """ + Cut signal from a start x to an end x + """ + indexes = where(logical_and(self.x > start, self.x < end)) + return Signal( + self.x[indexes], + self.y[indexes], + ) + + def detrend(self, type: str = "linear") -> "Signal": + """ + Short alias for scipy.detrend with default value + """ + return Signal( + self.x, + detrend(self.y, type=type), + ) + + def extend( + self, start: None | float = None, end: None | float = None + ) -> "Signal": + """ + Extend a signal to a start x or an end x + """ + x = self.x + y = self.y + if start is not None: + x = append(arange(start, x[0], 1 / self.rate), x) + y = append(zeros(len(x) - len(y)), y) + if end is not None: + x = append(x, arange(x[-1], end, 1 / self.rate)) + y = append(y, zeros(len(x) - len(y))) + return Signal( + x, + y, + ) + + def filter( + self, start: None | float, end: None | float + ) -> "Signal": + freq_x = rfftfreq(len(self), self.sampling) + freq_y = rfft(self.y) + + index_to_remove = where( + logical_and(abs(freq_x) < start, abs(freq_x) > end) + ) + freq_y[index_to_remove] = 0 + y = irfft(freq_y) + return Signal( + self.x[: len(y)], + y[: len(self.y)], + ) + + def psd(self, fft_length: int = 10) -> "Signal": + """ + Compute psd of a given signal + """ + psd, freq = welch( + self.y, + self.rate, + nperseg=fft_length * self.rate, + detrend="linear", + ) + + return Signal(freq, psd) + + def sin(self) -> "Signal": + """ + Return the sinus of this signal + """ + return Signal( + self.x, + sin(self.y), + ) + + def sqrt(self) -> "Signal": + """ + Return the root square of this signal + """ + return Signal( + self.x, + self.y ** (1 / 2), + ) + def operator_signal( self, other: "Signal", operator: str ) -> "Signal": @@ -180,6 +286,12 @@ class Signal: return other.operator_signal(self, "/") return self.roperator_arraylike(other, "/") + def __pow__(self, other: float | int) -> "Signal": + return Signal( + self.x, + self.y**other, + ) + def __array__(self) -> NDArray[float64]: return array([self.x, self.y])