Feature/pad arrays

src/pymodaq_data/data.py

@@ -8,6 +8,8 @@
 
 from abc import ABCMeta, abstractmethod, abstractproperty
 import numbers
+from copy import deepcopy
+
 import numpy as np
 from numpy.lib.mixins import NDArrayOperatorsMixin
 from typing import List, Tuple, Union, Any, Callable
@@ -1081,6 +1083,7 @@ def fliplr(self):
         new_data.data = [np.fliplr(dat) for dat in new_data]
         return new_data
 
+
     def append(self, data: DataWithAxes):
         """Append data content if the underlying arrays have the same shape and compatible units"""
         for dat in data:
@@ -2272,6 +2275,16 @@ def interp(self,  new_axis_data: Union[Axis, np.ndarray], **kwargs) -> DataWithA
                                   labels=self.labels)
         return new_data
 
+    def pad(self, pad_width: Union[int, Tuple[int, int], Iterable[Tuple[int, int]]], **kwargs):
+        """ Pad the data arrays using the numpy pad function
+
+        The accepted pad_witdh type is the same than the numpy pad function
+
+
+        see numpy.pad method for the signature and possible named arguments
+        """
+        return np.pad(self, pad_width, **kwargs)
+
     def ft(self, axis: int = 0, axis_label: str = None,
            axis_units: str = None, labels: List[str] = None) -> DataWithAxes:
         """Process the Fourier Transform of the data on the specified axis and returns the new data
@@ -2557,6 +2570,25 @@ def create_missing_axes(self):
                         axes.append(ax)
         self.axes = axes
 
+    def rot90(self, k=1, axes=(0, 1)):
+        """ Rotate an array by 90 degrees in the plane specified by axes.
+
+        Valid only for 2D data
+        """
+        if self.dim == DataDim.Data2D:
+            new_data: DataWithAxes = copy.copy(self)
+            new_data.data = [np.rot90(dat, k, axes) for dat in new_data]
+            new_axes = []
+            axis = new_data.get_axis_from_index(0)[0]
+            axis.index = 1
+            new_axes.append(axis)
+            axis = new_data.get_axis_from_index(1)[0]
+            axis.index = 0
+            new_data.axes = new_axes
+            return new_data
+        else:
+            return self
+
     def _compute_slices(self, slices, is_navigation=True, is_index=True):
         """Compute the total slice to apply to the data
 

src/pymodaq_data/numpy_func.py

@@ -1,16 +1,21 @@
 from typing import Union, List, TYPE_CHECKING, Iterable, Optional, Callable
 import numbers
+from copy import deepcopy
 
 import numpy as np
 from pint.facets.numpy.numpy_func import HANDLED_UFUNCS  # imported by the data module
 from pymodaq_data import Q_
 from pymodaq_data import data as data_mod
 
+from pymodaq_utils.logger import set_logger, get_module_name
+
 if TYPE_CHECKING:
     from pymodaq_data.data import DataBase, DataWithAxes
 
 HANDLED_FUNCTIONS = {}
 
+logger = set_logger(get_module_name(__file__))
+
 
 def process_arguments_for_ufuncs(input: 'DataBase',
                       inputs: List[Union[numbers.Number, Q_, np.ndarray, 'DataBase']]):
@@ -161,6 +166,38 @@ def _roll(dwa: 'DataWithAxes', *args, **kwargs):
     dwa_func.name += f"_{'roll'}"
     return dwa_func
 
+
+@implements('pad')
+def _pad(dwa: 'DataWithAxes', pad_width, mode = 'constant', **kwargs):
+    dwa.create_missing_axes()
+    for axis in dwa.axes:
+        if not axis.is_axis_linear():
+            raise TypeError('Could not pad data with non linear axes')
+    if isinstance(pad_width, int):
+            pad_width = [(pad_width, pad_width) for _ in range(len(dwa.shape))]
+    elif len(pad_width) == 1:
+        if hasattr(pad_width[0], '__len__') and len(pad_width[0]) == 2:
+            pad_width = [pad_width[0] for _ in range(len(dwa.shape))]
+        else:
+            pad_width = pad_width[0]
+            pad_width = [(pad_width, pad_width) for _ in range(len(dwa.shape))]
+    elif len(pad_width) == 2 and not hasattr(pad_width[0], '__len__'):
+        pad_width = [pad_width for _ in range(len(dwa.shape))]
+    elif len(pad_width) == len(dwa.shape):
+        if not hasattr(pad_width[0], '__len__'):
+            raise TypeError('Could not pad data with the given argument')
+    else:
+        raise TypeError(f'Could not pad data with the given arguments: {pad_width}')
+    dwa_func = dwa.deepcopy_with_new_data(data=[np.pad(array, pad_width, mode, **kwargs) for array in dwa])
+    dwa_func.axes = []
+    for axis in deepcopy(dwa.axes):
+        axis.offset -= pad_width[axis.index][0] * axis.scaling
+        axis.size += pad_width[axis.index][0] + pad_width[axis.index][1]
+        dwa_func.axes.append(axis)
+    dwa_func.name += f"_{'pad'}"
+    return dwa_func
+
+
 # ******** functions that return booleans ***********
 @implements('all')
 def _all(dwa: 'DataWithAxes', *args, axis: Optional[Union[int, Iterable[int]]] = None, **kwargs):

src/pymodaq_data/plotting/plotter/plotter.py

@@ -83,13 +83,14 @@ def inner_wrapper(wrapped_class: PlotterBase) -> Callable:
     def create(cls, key, **kwargs) -> PlotterBase:
         builder = cls._builders[cls.__name__].get(key)
         if not builder:
-            raise ValueError(key)
+            raise ValueError(f'{key} is not a valid plotter: {cls.backends()}')
         return builder(**kwargs)
 
     def get(self, backend: str, **kwargs):
         return self.create(backend, **kwargs)
 
-    def backends(self) -> List[str]:
+    @classmethod
+    def backends(cls) -> List[str]:
         """Returns the list of plotter backends, main identifier of a given plotter"""
-        return sorted(list(self.builders[self.__class__.__name__].keys()))
+        return sorted(list(cls._builders[cls.__name__].keys()))
 

tests/data_test.py

@@ -703,6 +703,87 @@ def test_interp(self):
         with pytest.raises(ValueError):
             dwa.interp(new_axis_array)
 
+    def test_pad_1D(self):
+        omega0 = 5
+        Npts_in = 2**10
+        time_axis = data_mod.Axis('time', 's', data=np.linspace(0, 10*2*np.pi, Npts_in))
+        dwa = data_mod.DataRaw('sinus',
+                               data=[
+                                   np.sin(omega0 * time_axis.get_data()),
+                                   np.cos(omega0 * time_axis.get_data()),
+                               ], labels=['sinus', 'cosinus'],
+                               axes=[time_axis])
+        pad_width = 123
+        dwa_padded = dwa.pad(pad_width)
+
+        assert dwa_padded.size == dwa.size + 2 * pad_width
+        assert dwa_padded.axes[0].size == dwa.axes[0].size + 2 * pad_width
+
+        pad_width = (123,)
+        dwa_padded = dwa.pad(pad_width)
+
+        assert dwa_padded.size == dwa.size + 2 * pad_width[0]
+        assert dwa_padded.axes[0].size == dwa.axes[0].size + 2 * pad_width[0]
+
+        pad_width = (123, 256)
+        dwa_padded = dwa.pad(pad_width)
+
+        assert dwa_padded.size == dwa.size + np.sum(pad_width)
+        assert dwa_padded.axes[0].size == dwa.axes[0].size + np.sum(pad_width)
+
+        pad_width = ((123, 256),)
+        dwa_padded = dwa.pad(pad_width)
+
+        assert dwa_padded.size == dwa.size + np.sum(pad_width[0])
+        assert dwa_padded.axes[0].size == dwa.axes[0].size + np.sum(pad_width[0])
+
+        pad_width = ((123, 256), (12, 25),)
+        with pytest.raises(TypeError):
+            dwa_padded = dwa.pad(pad_width)
+
+    def test_pad_2D(self):
+        x_axis = data_mod.Axis('x', 's', data=np.linspace(0, 10*2*np.pi, 2**6), index=1)
+        y_axis = data_mod.Axis('y', 's', data=np.linspace(0, 10*2*np.pi, 2**8), index=0)
+        dwa = data_mod.DataRaw('gauss2D',
+                               data=[
+                                   mutils.gauss2D(x_axis.get_data(), 10, 10, y_axis.get_data(), 20, 6),
+                               ], labels=['gauss2D'],
+                               axes=[x_axis, y_axis])
+        pad_width = 123
+        dwa_padded = dwa.pad(pad_width)
+
+        for ind in range(len(dwa.shape)):
+            assert dwa_padded.shape[ind] == dwa.shape[ind] + 2 * pad_width
+            assert dwa_padded.get_axis_from_index(ind)[0].size == dwa.get_axis_from_index(ind)[0].size + 2 * pad_width
+
+        pad_width = (123,)
+        dwa_padded = dwa.pad(pad_width)
+
+        pad_width = pad_width[0]
+        for ind in range(len(dwa.shape)):
+            assert dwa_padded.shape[ind] == dwa.shape[ind] + 2 * pad_width
+            assert dwa_padded.get_axis_from_index(ind)[0].size == dwa.get_axis_from_index(ind)[0].size + 2 * pad_width
+
+        pad_width = (123, 256)
+        dwa_padded = dwa.pad(pad_width)
+        for ind in range(len(dwa.shape)):
+            assert dwa_padded.shape[ind] == dwa.shape[ind] + np.sum(pad_width)
+            assert dwa_padded.get_axis_from_index(ind)[0].size == dwa.get_axis_from_index(ind)[0].size + np.sum(pad_width)
+
+        pad_width = ((123, 256),)
+        dwa_padded = dwa.pad(pad_width)
+        pad_width = pad_width[0]
+        for ind in range(len(dwa.shape)):
+            assert dwa_padded.shape[ind] == dwa.shape[ind] + np.sum(pad_width)
+            assert dwa_padded.get_axis_from_index(ind)[0].size == dwa.get_axis_from_index(ind)[0].size + np.sum(pad_width)
+
+        pad_width = ((123, 256), (12, 25),)
+        dwa_padded = dwa.pad(pad_width)
+        for ind in range(len(dwa.shape)):
+            assert dwa_padded.shape[ind] == dwa.shape[ind] + np.sum(pad_width[ind])
+            assert dwa_padded.get_axis_from_index(ind)[0].size == dwa.get_axis_from_index(ind)[0].size + np.sum(pad_width[ind])
+
+
     def test_ft_ift(self):
         omega0 = 5
         time_axis = data_mod.Axis('time', 's', data=np.linspace(0, 10*2*np.pi, 2**10))