Add more operator

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])