diff --git a/src/pymodaq_data/data.py b/src/pymodaq_data/data.py index 6e0f492f..37536c2c 100644 --- a/src/pymodaq_data/data.py +++ b/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 diff --git a/src/pymodaq_data/numpy_func.py b/src/pymodaq_data/numpy_func.py index 1d260eee..71cf87cf 100644 --- a/src/pymodaq_data/numpy_func.py +++ b/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): diff --git a/src/pymodaq_data/plotting/plotter/plotter.py b/src/pymodaq_data/plotting/plotter/plotter.py index d64b8032..220c3b53 100644 --- a/src/pymodaq_data/plotting/plotter/plotter.py +++ b/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())) diff --git a/tests/data_test.py b/tests/data_test.py index a2d1bff2..b287084e 100644 --- a/tests/data_test.py +++ b/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))