Some documentation rework and resampling notebook example with CM

Signed-off-by: Alexis Jeandet <alexis.jeandet@member.fsf.org>

speasy/products/variable.py

@@ -217,6 +217,11 @@ def __mul__(self, other):
     def __rmul__(self, other):
         return self.__mul__(other)
 
+    def __pow__(self, power, modulo=None):
+        res = self.copy()
+        np.power(self.__values_container.values, power, out=res.__values_container.values)
+        return res
+
     def __add__(self, other):
         if type(other) is SpeasyVariable:
             if self.__values_container.shape != other.__values_container.shape:
@@ -263,9 +268,24 @@ def __truediv__(self, other):
             res = self.copy()
             np.divide(self.__values_container.values, float(other), out=res.__values_container.values)
             return res
+        if type(other) is SpeasyVariable:
+            return np.divide(self, other)
         raise TypeError(
             f"Can't divide SpeasyVariable by {type(other)}")
 
+    def __rtruediv__(self, other):
+        if type(other) in (int, float):
+            res = self.copy()
+            np.divide(float(other), self.__values_container.values, out=res.__values_container.values)
+            return res
+        if type(other) is SpeasyVariable:
+            res = self.copy()
+            np.divide(other.__values_container.values, self.__values_container.values,
+                      out=res.__values_container.values)
+            return np.divide(other, self)
+        raise TypeError(
+            f"Can't divide {type(other)} by SpeasyVariable")
+
     def __array_function__(self, func, types, args, kwargs):
         if func.__name__ in SpeasyVariable.__LIKE_NP_FUNCTIONS__:
             return SpeasyVariable.__dict__[func.__name__](self)

speasy/signal/filtering/__init__.py

@@ -1,28 +1,65 @@
 from scipy import signal
-from typing import Callable
+from typing import Callable, Union, Collection
 from speasy.products import SpeasyVariable
 import numpy as np
 
 
-def apply_sos_filter(sos: np.ndarray, filter_function: Callable, var: SpeasyVariable) -> SpeasyVariable:
-    res = np.empty_like(var)
+def _apply_filter(filter_function: Callable, sos: np.ndarray, var: SpeasyVariable) -> SpeasyVariable:
+    res = SpeasyVariable.reserve_like(var)
     res.values[:] = filter_function(sos, var.values, axis=0)
     return res
 
 
-def sosfiltfilt(sos: np.ndarray, var: SpeasyVariable) -> SpeasyVariable:
-    """Apply an IIR filter to the data using :func:`scipy.signal.sosfiltfilt`.
+def apply_sos_filter(sos: np.ndarray, filter_function: Callable,
+                     var: Union[SpeasyVariable, Collection[SpeasyVariable]]) -> Union[
+    SpeasyVariable, Collection[SpeasyVariable]]:
+    """Apply an IIR filter to the variable(s) using the given filter function. This function just applies the filter to the
+    values of the variable without any resampling, it assumes that the variable has a regular time axis.
+
+    Parameters
+    ----------
+    sos: np.ndarray
+        Second-order sections representation of the filter, as returned by :func:`scipy.signal.iirfilter` with `output='sos'` for example.
+    filter_function: Callable
+        The filter function to use (e.g. :func:`scipy.signal.sosfiltfilt`)
+    var: SpeasyVariable or Collection[SpeasyVariable]
+        The variable(s) to filter
+
+    Returns
+    -------
+    SpeasyVariable or Collection[SpeasyVariable]
+        The filtered variable(s)
+
+    Notes
+    -----
+    It only supports 1D variables.
+    """
+
+    if isinstance(var, SpeasyVariable):
+        return _apply_filter(filter_function, sos, var)
+    else:
+        return [_apply_filter(filter_function, sos, v) for v in var]
+
+
+def sosfiltfilt(sos: np.ndarray, var: Union[SpeasyVariable, Collection[SpeasyVariable]]) -> Union[
+    SpeasyVariable, Collection[SpeasyVariable]]:
+    """Apply an IIR filter to the variable(s) using :func:`scipy.signal.sosfiltfilt`. This function just applies the filter to
+    the values of the variable without any resampling, it assumes that the variable has a regular time axis.
 
     Parameters
     ----------
     sos: np.ndarray
         Second-order sections representation of the filter
-    var: SpeasyVariable
-        The variable to filter
+    var: SpeasyVariable or Collection[SpeasyVariable]
+        The variable(s) to filter
 
     Returns
     -------
-    SpeasyVariable
-        The filtered variable
+    SpeasyVariable or Collection[SpeasyVariable]
+        The filtered variable(s)
+
+    Notes
+    -----
+    It only supports 1D variables.
     """
     return apply_sos_filter(sos, signal.sosfiltfilt, var)

speasy/signal/resampling/__init__.py

@@ -60,13 +60,13 @@ def _interpolate(ref_time: np.ndarray, var: SpeasyVariable, interpolate_callback
 def resample(var: Union[SpeasyVariable, Collection[SpeasyVariable]], new_dt: Union[float, np.timedelta64],
              interpolate_callback: Optional[Callable] = None,
              *args, **kwargs) -> Union[SpeasyVariable, Collection[SpeasyVariable]]:
-    """Resample a variable to a new time step. The time vector will be generated from the start and stop times of the
+    """Resample a variable(s) to a new time step. The time vector will be generated from the start and stop times of the
     input variable. Uses :func:`numpy.interp` to do the resampling by default.
 
     Parameters
     ----------
     var: SpeasyVariable or Collection[SpeasyVariable]
-        The variable or a collection of variables to resample
+        The variable(s) to resample
     new_dt: float or np.timedelta64
         The new time step in seconds or as a numpy timedelta64
     interpolate_callback: Callable or None
@@ -75,7 +75,11 @@ def resample(var: Union[SpeasyVariable, Collection[SpeasyVariable]], new_dt: Uni
     Returns
     -------
     SpeasyVariable or Collection[SpeasyVariable]
-        The resampled variable or a collection of resampled variables
+        The resampled variable(s) with all metadata preserved except for the new time axis
+
+    Notes
+    -----
+    It only supports 1D variables.
     """
     if type(var) in (list, tuple):
         return [resample(v, new_dt, interpolate_callback, *args, **kwargs) for v in var]
@@ -87,21 +91,26 @@ def resample(var: Union[SpeasyVariable, Collection[SpeasyVariable]], new_dt: Uni
 def interpolate(ref: Union[np.ndarray, SpeasyVariable], var: Union[SpeasyVariable, Collection[SpeasyVariable]],
                 interpolate_callback: Optional[Callable] = None,
                 *args, **kwargs) -> Union[SpeasyVariable, Collection[SpeasyVariable]]:
-    """Interpolate a variable to a new time vector. The time vector will be taken from the reference variable. Uses :func:`numpy.interp` to do the resampling by default.
+    """Interpolate a variable(s) to a new time vector. The time vector will be taken from the reference variable.
+    Uses :func:`numpy.interp` to do the resampling by default.
 
     Parameters
     ----------
     ref: np.ndarray or SpeasyVariable
         The reference time vector
     var: SpeasyVariable or Collection[SpeasyVariable]
-        The variable or a collection of variables to interpolate
+        The variable(s) to interpolate
     interpolate_callback: Callable or None
         The interpolation function to use, defaults to :func:`numpy.interp` (Optional)
 
     Returns
     -------
     SpeasyVariable or Collection[SpeasyVariable]
-        The interpolated variable or a collection of interpolated variables
+        The interpolated variable(s) with all metadata preserved except for the new time axis
+
+    Notes
+    -----
+    It only supports 1D variables.
     """
     if isinstance(ref, SpeasyVariable):
         ref_time = ref.time