# -*- coding: utf-8 -*- """ Abstract :class:`~sidpy.io.dataset.Dataset` base-class Created on Tue Nov 3 15:07:16 2015 @author: Gerd Duscher Modified by Mani Valleti. Look up dask source code to understand how numerical functions are implemented starting code from: https://scikit-allel.readthedocs.io/en/v0.21.1/_modules/allel/model/dask.html """ from __future__ import division, print_function, absolute_import, unicode_literals from hashlib import new from functools import wraps from re import A import sys from collections.abc import Iterable, Iterator, Mapping import warnings import ase import dask.array.core import numpy as np import matplotlib.pylab as plt import string import dask.array as da import h5py from enum import Enum from numbers import Number from .dimension import Dimension, DimensionType from ..base.num_utils import get_slope from ..base.dict_utils import print_nested_dict from ..viz.dataset_viz import CurveVisualizer, ImageVisualizer, ImageStackVisualizer from ..viz.dataset_viz import SpectralImageVisualizer, FourDimImageVisualizer, ComplexSpectralImageVisualizer from ..viz.dataset_viz import PointCloudVisualizer # from ..hdf.hdf_utils import is_editable_h5 from .dimension import DimensionType from copy import deepcopy, copy from sidpy.base.string_utils import validate_single_string_arg import logging def is_simple_list(lst): if isinstance(lst, list): return any(hasattr(item, '__getitem__') for item in lst) return False class DataType(Enum): UNKNOWN = -1 SPECTRUM = 1 LINE_PLOT = 2 LINE_PLOT_FAMILY = 3 IMAGE = 4 IMAGE_MAP = 5 IMAGE_STACK = 6 # 3d SPECTRAL_IMAGE = 7 IMAGE_4D = 8 POINT_CLOUD = 9 def view_subclass(dask_array, cls): """ View a dask Array as an instance of a dask Array sub-class. Parameters ---------- dask_array cls Returns ------- cls: sidpy.Dataset """ return cls(dask_array.dask, name=dask_array.name, chunks=dask_array.chunks, dtype=dask_array.dtype, shape=dask_array.shape) class Dataset(da.Array): """ ..autoclass::Dataset To instantiate from an existing array-like object, use :func:`Dataset.from_array` - requires numpy array, list or tuple This dask array is extended to have the following attributes: -data_type: DataTypes ('image', 'image_stack', spectral_image', ... -units: str -quantity: str what kind of data ('intensity', 'height', ..) -title: title of the data set -modality: character of data such as 'STM, 'AFM', 'TEM', 'SEM', 'DFT', 'simulation', ..) -source: origin of data such as acquisition instrument ('Nion US100', 'VASP', ..) -_axes: dictionary of Dimensions one for each data dimension (the axes are dimension datasets with name, label, units, and 'dimension_type' attributes). -metadata: dictionary of additional metadata -original_metadata: dictionary of original metadata of file, -labels: returns labels of all dimensions. -data_descriptor: returns a label for the colorbar in matplotlib and such functions: -from_array(data, title): constructs the dataset form an array like object (numpy array, dask array, ...) -like_data(data,title): constructs the dataset form an array like object and copies attributes and metadata from parent dataset -copy() -plot(): plots dataset dependent on data_type and dimension_types. -get_extent(): extent to be used with imshow function of matplotlib -set_dimension(axis, dimensions): set a Dimension to a specific axis -rename_dimension(dimension, name): renames attribute of dimension -view_metadata: pretty plot of metadata dictionary -view_original_metadata: pretty plot of original_metadata dictionary """ def __init__(self, *args, **kwargs): """ Initializes Dataset object which is essentially a Dask array underneath Attributes ---------- self.quantity : str Physical quantity. E.g. - current self.units : str Physical units. E.g. - amperes self.data_type : enum Type of data such as Image, Spectrum, Spectral Image etc. self.title : str Title for Dataset self._structures : dict dictionary of ase.Atoms objects to represent structures, can be given a name self.view : Visualizer Instance of class appropriate for visualizing this object self.data_descriptor : str Description of this dataset self.modality : str character of data such as 'STM', 'TEM', 'DFT' self.source : str Source of this dataset. Such as instrument, analysis, etc.? self.h5_dataset : h5py.Dataset Reference to HDF5 Dataset object from which this Dataset was created self._axes : dict Dictionary of Dimension objects per dimension of the Dataset self.meta_data : dict Metadata to store relevant additional information for the dataset. self.original_metadata : dict Metadata from the original source of the dataset. This dictionary often contains the vendor-specific metadata or internal attributes of the analysis algorithm """ # TODO: Consider using python package - pint for quantities super().__init__() self._units = '' self._quantity = '' self._title = '' self._data_type = DataType.UNKNOWN self._modality = '' self._source = '' self._structures = {} self._h5_dataset = None self._metadata = {} self._original_metadata = {} self._axes = {} self.view = None # this will hold the figure and axis reference for a plot self.__protected = set() # a set to keep track of protected attributes self.point_cloud = None # attribute to store coordinates and base_image for point_cloud datatype self._variance = None # to save variance dask.array def __repr__(self): rep = 'sidpy.Dataset of type {} with:\n '.format(self.data_type.name) rep = rep + super(Dataset, self).__repr__() rep = rep + '\n data contains: {} ({})'.format(self.quantity, self.units) rep = rep + '\n and Dimensions: ' for key in self._axes: rep = rep + '\n' + self._axes[key].__repr__() if hasattr(self, 'metadata'): if len(self.metadata) > 0: rep = rep + '\n with metadata: {}'.format(list(self.metadata.keys())) return rep def hdf_close(self): if self.h5_dataset is not None: self.h5_dataset.file.close() print(self.h5_dataset) def __setattr__(self, key, value): if not hasattr(self, '_Dataset__protected'): super().__setattr__(key, value) else: # if key is in __protected, only Dimension and numpy.ndarray instances are allowed to be set if key != 'none' and key in self._Dataset__protected: if not isinstance(value, Dimension): raise AttributeError('The attribute "{}" is reserved to represent a dimension'.format(key)) else: if getattr(self, key).name == value.name and len(getattr(self, key)) == len(value): cur_ind = [i for i in self._axes if self._axes[i].name == key][0] self.del_dimension(cur_ind) self._axes[cur_ind] = value self.__dict__[key] = value self.__dict__['dim_{}'.format(cur_ind)] = value self.__protected.add(key) self.__protected.add('dim_{}'.format(cur_ind)) else: raise NotImplementedError("The new dimension's name or length does not " "match with the existing dimension.") else: super().__setattr__(key, value) @classmethod def from_array(cls, x, title='generic', chunks='auto', lock=False, datatype='UNKNOWN', units='generic', quantity='generic', modality='generic', source='generic', coordinates=None, variance=None, **kwargs): """ Initializes a sidpy dataset from an array-like object (i.e. numpy array) All meta-data will be set to be generically. Parameters ---------- x: array-like object the values which will populate this dataset chunks: optional integer or list of integers the shape of the chunks to be loaded title: optional string the title of this dataset lock: boolean datatype: str or sidpy.DataType data type of set: i.e.: 'image', spectrum', .. units: str units of dataset i.e. counts, A quantity: str quantity of dataset like intensity modality: str modality of dataset like source: str source of dataset like what kind of microscope or function coordinates: numpy array, optional coordinates for point cloud point_cloud: dict or None dict with coordinates and base_image for point_cloud data_type variance: array-like object the variance values of the x array Returns ------- sidpy dataset """ # create vanilla dask array if isinstance(x, da.Array) and not np.any(np.isnan(x.shape)): dask_array = x else: dask_array = da.from_array(np.array(x), chunks=chunks, lock=lock) # view as subclass sid_dataset = view_subclass(dask_array, cls) sid_dataset.data_type = datatype sid_dataset.units = units sid_dataset.title = title sid_dataset.quantity = quantity sid_dataset.modality = modality sid_dataset.source = source sid_dataset._axes = {} for dim in range(sid_dataset.ndim): # TODO: add parent to dimension to set attribute if name changes sid_dataset.set_dimension(dim, Dimension(np.arange(sid_dataset.shape[dim]), string.ascii_lowercase[dim])) sid_dataset.metadata = {} sid_dataset.original_metadata = {} sid_dataset.variance = variance # add coordinates for point_cloud datatype if coordinates is not None: sid_dataset.point_cloud = {'coordinates': coordinates} else: sid_dataset.point_cloud = None return sid_dataset def like_data(self, data, title=None, chunks='auto', lock=False, coordinates=None, variance=None, **kwargs): """ Returns sidpy.Dataset of new values but with metadata of this dataset - if dimension of new dataset is different from this dataset and the scale is linear, then this scale will be applied to the new dataset (naming and units will stay the same), otherwise the dimension will be generic. -Additional functionality to override numeric functions Parameters ---------- data: array like values of new sidpy dataset title: optional string title of new sidpy dataset chunks: optional list of integers size of chunks for dask array lock: optional boolean for dask array coordinates: array like coordinates for point cloud variance: numpy array, optional variance of dataset Returns ------- sidpy dataset """ title_suffix = kwargs.get('title_suffix', '') title_prefix = kwargs.get('title_prefix', '') reset_quantity = kwargs.get('reset_quantity', False) reset_units = kwargs.get('reset_units', False) checkdims = kwargs.get('checkdims', True) # if coordinates is None: # coordinates = self.point_cloud['coordinates'] new_data = self.from_array(data, chunks=chunks, lock=lock, variance =variance) new_data.data_type = self.data_type # if variance is None: # if new_data.shape == self.shape: # new_data.variance = self.variance # units if reset_units: new_data.units = 'generic' else: new_data.units = self.units if title is not None: new_data.title = title else: if title_prefix or title_suffix: new_data.title = self.title else: new_data.title = self.title + '_new' new_data.title = title_prefix + new_data.title + title_suffix # quantity if reset_quantity: new_data.quantity = 'generic' else: new_data.quantity = self.quantity new_data.modality = self.modality new_data.source = self.source if checkdims: for dim in range(new_data.ndim): # TODO: add parent to dimension to set attribute if name changes if len(self._axes[dim].values) == new_data.shape[dim]: new_data.set_dimension(dim, self._axes[dim]) else: # assuming the axis scale is equidistant try: scale = get_slope(self._axes[dim]) # axis = self._axes[dim].copy() axis = Dimension(np.arange(new_data.shape[dim]) * scale, self._axes[dim].name) axis.quantity = self._axes[dim].quantity axis.units = self._axes[dim].units axis.dimension_type = self._axes[dim].dimension_type new_data.set_dimension(dim, axis) except ValueError: print('using generic parameters for dimension ', dim) new_data.metadata = dict(self.metadata).copy() new_data.original_metadata = {} return new_data def __reduce_dimensions(self, new_dataset, axes, keepdims=False): new_dataset.del_dimension() if not keepdims: i = 0 for key, dim in self._axes.items(): new_dim = dim.copy() if key not in axes: new_dataset.set_dimension(i, new_dim) i += 1 if keepdims: for key, dim in self._axes.items(): new_dim = dim.copy() if key in axes: new_dim = Dimension(np.arange(1), name=new_dim.name, quantity=new_dim.quantity, units=new_dim.units, dimension_type=new_dim.dimension_type) new_dataset.set_dimension(key, new_dim) return new_dataset def __rearrange_axes(self, new_dataset, new_order=None): """Rearranges the dimension order of the current instance Parameters: new_order: list or tuple of integers All the dimensions that are not in the new_order are deleted """ new_dataset.del_dimension() for i, dim in enumerate(new_order): new_dataset.set_dimension(i, self._axes[dim]) return new_dataset def copy(self): """ Returns a deep copy of this dataset. Returns ------- sidpy dataset """ dataset_copy = Dataset.from_array(self, self.title, self.chunks) dataset_copy.title = self.title dataset_copy.units = self.units dataset_copy.quantity = self.quantity dataset_copy.data_type = self.data_type dataset_copy.modality = self.modality dataset_copy.source = self.source dataset_copy.point_cloud = self.point_cloud dataset_copy.variance = self.variance dataset_copy.del_dimension() for dim in self._axes: dataset_copy.set_dimension(dim, self._axes[dim]) dataset_copy.metadata = dict(self.metadata).copy() return dataset_copy def __validate_dim(self, ind, name): """ Validates the provided index for a Dimension object Parameters ---------- ind : int Index of the dimension Raises ------- TypeError : if ind is not an integer IndexError : if ind is less than 0 or greater than maximum allowed index for Dimension ValueError: if name is not 'none' and is already used. """ if not isinstance(ind, int): raise TypeError('Dimension must be an integer') if (0 > ind) or (ind >= self.ndim): raise IndexError('Dimension must be an integer between 0 and {}' ''.format(self.ndim - 1)) for key, dim in self._axes.items(): if key != ind: if name != 'none' and name == dim.name: raise ValueError('name: {} already used, but must be unique'.format(name)) def rename_dimension(self, ind, name): """ Renames Dimension at the specified index Parameters ---------- ind : int Index of the dimension name : str New name for Dimension """ self.__validate_dim(ind, name) if not isinstance(name, str): raise TypeError('New Dimension name must be a string') if hasattr(self, self._axes[ind].name): delattr(self, self._axes[ind].name) if self._axes[ind].name in self.__protected: self.__protected.remove(self._axes[ind].name) if hasattr(self, 'dim_{}'.format(ind)): delattr(self, 'dim_{}'.format(ind)) self.__protected.remove('dim_{}'.format(ind)) self._axes[ind]._name = validate_single_string_arg(name, 'name') # protected attribute name setattr(self, name, self._axes[ind]) self.__protected.add(name) setattr(self, 'dim_{}'.format(ind), self._axes[ind]) self.__protected.add('dim_{}'.format(ind)) def set_dimension(self, ind, dimension): """ sets the dimension for the dataset including new name and updating the axes dictionary Parameters ---------- ind: int Index of dimension dimension: sidpy.Dimension Dimension object describing this dimension of the Dataset Returns ------- """ if not isinstance(dimension, Dimension): raise TypeError('dimension needs to be a sidpy.Dimension object') self.__validate_dim(ind, dimension.name) if len(dimension.values) != self.shape[ind]: raise ValueError('The length of the dimension array does not match the shape of the ' 'dataset at {}th dimension. {} != {}'.format(ind, len(dimension.values), self.shape[ind]) ) dim = dimension.copy() try: if hasattr(self, self._axes[ind].name): delattr(self, self._axes[ind].name) if self._axes[ind].name in self.__protected: self.__protected.remove(self._axes[ind].name) except KeyError: pass setattr(self, dimension.name, dim) self.__protected.add(dimension.name) if hasattr(self, 'dim_{}'.format(ind)): delattr(self, 'dim_{}'.format(ind)) if 'dim_{}'.format(ind) in self.__protected: self.__protected.remove('dim_{}'.format(ind)) # we don't need this. But I am trying to be consistent setattr(self, 'dim_{}'.format(ind), dim) self._axes[ind] = dim self.__protected.add('dim_{}'.format(ind)) def del_dimension(self, ind=None): """ Deletes the dimension attached to axis 'ind'. """ if isinstance(ind, int): ind = [ind] elif ind is None: ind = list(np.arange(self.ndim)) else: ind = list(ind) for i in ind: # Delete the attribute with the format dim_0 if hasattr(self, 'dim_{}'.format(i)): delattr(self, 'dim_{}'.format(i)) if 'dim_{}'.format(i) in self.__protected: self.__protected.remove('dim_{}'.format(i)) if i in self._axes.keys(): # Deleting the dataset attribute that has the dimension's name if hasattr(self, self._axes[i].name): delattr(self, self._axes[i].name) if self._axes[i].name in self.__protected: self.__protected.remove(self._axes[i].name) # Deleting the key-value pair from the _axes dictionary del self._axes[i] def view_metadata(self): """ Prints the metadata to stdout Returns ------- None """ if isinstance(self.metadata, dict): print_nested_dict(self.metadata) def view_original_metadata(self): """ Prints the original_metadata dictionary to stdout Returns ------- None """ if isinstance(self.original_metadata, dict): print_nested_dict(self.original_metadata) def plot(self, verbose=False, figure=None, **kwargs): """ Plots the dataset according to the - shape of the sidpy Dataset, - data_type of the sidpy Dataset and - dimension_type of dimensions of sidpy Dataset the dimension_type 'spatial' or 'spectral' determines how a dataset is plotted. Recognized data_types are: 1D: any keyword, but 'spectrum' or 'line_plot' are encouraged 2D: 'image' or one of ['spectrum_family', 'line_family', 'line_plot_family', 'spectra'] 3D: 'image', 'image_map', 'image_stack', 'spectrum_image' 4D: not implemented yet, but will be similar to spectrum_image. Parameters ---------- verbose: boolean kwargs: dictionary for additional plotting parameters additional keywords (besides the matplotlib ones) for plotting are: - scale_bar: for images to replace axis with a scale bar inside the image figure: matplotlib figure object define figure to which this datset will be plotted Returns ------- self.view.fig: matplotlib figure reference """ if verbose: print('Shape of dataset is: ', self.shape) if self.data_type.value < 0: raise NameError('Datasets with UNKNOWN data_types cannot be plotted') if len(self.shape) == 1: if verbose: print('1D dataset') self.view = CurveVisualizer(self, figure=figure, **kwargs) # plt.show() elif len(self.shape) == 2: # this can be an image or a set of line_plots if verbose: print('2D dataset') if self.data_type == DataType.IMAGE: self.view = ImageVisualizer(self, figure=figure, **kwargs) elif self.data_type.value <= DataType['LINE_PLOT'].value: # self.data_type in ['spectrum_family', 'line_family', 'line_plot_family', 'spectra']: self.view = CurveVisualizer(self, figure=figure, **kwargs) elif self.data_type == DataType.POINT_CLOUD: self.view = PointCloudVisualizer(self, figure=figure, **kwargs) else: raise NotImplementedError('Datasets with data_type {} cannot be plotted, yet.'.format(self.data_type)) elif len(self.shape) == 3: if verbose: print('3D dataset:', self.data_type) if self.data_type == DataType.IMAGE: self.view = ImageVisualizer(self, figure=figure, **kwargs) elif self.data_type == DataType.IMAGE_MAP: pass elif self.data_type == DataType.IMAGE_STACK: self.view = ImageStackVisualizer(self, figure=figure, **kwargs) elif self.data_type == DataType.SPECTRAL_IMAGE: if 'complex' in self.dtype.name: self.view = ComplexSpectralImageVisualizer(self, figure=figure, **kwargs) else: self.view = SpectralImageVisualizer(self, figure=figure, **kwargs) elif self.data_type.name == 'SPECTRAL_IMAGE': print('spec3') if 'complex' in self.dtype.name: self.view = ComplexSpectralImageVisualizer(self, figure=figure, **kwargs) else: self.view = SpectralImageVisualizer(self, figure=figure, **kwargs) elif self.data_type == DataType.POINT_CLOUD: self.view = PointCloudVisualizer(self, figure=figure, **kwargs) else: raise NotImplementedError('Datasets with data_type {} cannot be plotted, yet.'.format(self.data_type)) elif len(self.shape) == 4: if verbose: print('4D dataset') if self.data_type == DataType.IMAGE: self.view = ImageVisualizer(self, **kwargs) plt.show() elif self.data_type == DataType.IMAGE_MAP: pass elif self.data_type == DataType.IMAGE_STACK: self.view = ImageStackVisualizer(self, figure=figure, **kwargs) plt.show() elif self.data_type == DataType.SPECTRAL_IMAGE: if 'complex' in self.dtype.name: self.view = ComplexSpectralImageVisualizer(self, figure=figure, **kwargs) else: self.view = SpectralImageVisualizer(self, figure=figure, **kwargs) plt.show() elif self.data_type == DataType.IMAGE_4D: self.view = FourDimImageVisualizer(self, figure=figure, **kwargs) plt.show() if verbose: print('4D dataset') else: raise NotImplementedError('Datasets with data_type {} cannot be plotted, yet.'.format(self.data_type)) else: raise NotImplementedError('Datasets with data_type {} cannot be plotted, yet.'.format(self.data_type)) return self.view.fig def set_thumbnail(self, figure=None, thumbnail_size=128): """ Creates a thumbnail which is stored in thumbnail attribute of sidpy Dataset Thumbnail data is saved to Thumbnail group of associated h5_file if it exists Parameters ---------- thumbnail_size: int size of icon in pixels (length of square) Returns ------- thumbnail: numpy.ndarray """ import imageio # Thumbnail configurations for matplotlib kwargs = {'figsize': (1, 1), 'colorbar': False, 'set_title': False} view = self.plot(figure=figure, **kwargs) for axis in view.axes: axis.set_axis_off() # Creating Thumbnail as png image view.savefig('thumb.png', dpi=thumbnail_size) self.thumbnail = imageio.imread('thumb.png') # Writing thumbnail to h5_file if it exists if self.h5_dataset is not None: if 'Thumbnail' not in self.h5_dataset.file: thumb_group = self.h5_dataset.file.create_group("Thumbnail") else: thumb_group = self.h5_dataset.file["Thumbnail"] if "Thumbnail" in thumb_group: del thumb_group["Thumbnail"] thumb_dset = thumb_group.create_dataset("Thumbnail", data=self.thumbnail) return self.thumbnail def get_extent(self, dimensions): """ get image extents as needed i.e. in matplotlib's imshow function. This function works for equi- or non-equi spaced axes and is suitable for subpixel accuracy of positions Parameters ---------- dimensions: list of dimensions Returns ------- list of floats """ extent = [] for ind, dim in enumerate(dimensions): temp = self._axes[dim].values start = temp[0] - (temp[1] - temp[0]) / 2 end = temp[-1] + (temp[-1] - temp[-2]) / 2 if ind == 1: extent.append(end) # y-axis starts on top extent.append(start) else: extent.append(start) extent.append(end) return extent def get_dimension_slope(self, dim): axis = None if isinstance(dim, int): axis = self._axes[dim] elif isinstance(dim, Dimension): axis = dim return get_slope(axis) def get_dimension_by_number(self, dims_in): if isinstance(dims_in, int): dims_in = [dims_in] for i in range(len(dims_in)): if not isinstance(dims_in[i], int): raise ValueError('Input dimensions must be integers') out_dim = [] for dim in dims_in: out_dim.append(self._axes[dim]) return out_dim def get_dimensions_types(self): out_types = [] for dim, axis in self._axes.items(): out_types.append(axis.dimension_type) return out_types def get_dimensions_by_type(self, dims_in, return_axis=False): """ get dimension by dimension_type name Parameter --------- dims_in: dimension_type/str or list of dimension_types/string Returns ------- dims_out: list of [index] the kind of dimensions specified in input in numerical order of the dataset, not the input! """ if isinstance(dims_in, (str, DimensionType)): dims_in = [dims_in] for i in range(len(dims_in)): if isinstance(dims_in[i], str): dims_in[i] = DimensionType[dims_in[i].upper()] dims_out = [] for dim, axis in self._axes.items(): if axis.dimension_type in dims_in: if return_axis: dims_out.append(axis) else: dims_out.append(dim) # , self._axes[dim]]) return dims_out def get_image_dims(self, return_axis=False): """Get all spatial dimensions""" return self.get_dimensions_by_type(DimensionType.SPATIAL, return_axis=return_axis) def get_spectral_dims(self, return_axis=False): """Get all spectral dimensions""" return self.get_dimensions_by_type(DimensionType.SPECTRAL, return_axis=return_axis) def _griddata_transform(self, **kwargs): """ Interpolate unstructured point cloud for the visualization to 3D/4D sidpy.Dataset Parameters ---------- kwards: parameters to reduce dataset dimentions to 2D (number of point, spectral data) Returns ------- sidpy.Dataset with data_type = SPECTRAL_IMAGE """ from scipy.interpolate import griddata if 'coordinates' in self.metadata.keys(): coord = self.metadata['coordinates'] else: raise NotImplementedError('Datasets with data_type POINT_CLOUD must contain coordinates in metadata.') if 'spacial_units' in self.metadata.keys(): sp_units = self.metadata['spacial_units'] else: sp_units = 'a.u.' im_size = max(50, coord.shape[0]) _x0, _x1 = np.min(coord, axis=0)[0], np.max(coord, axis=0)[0] _y0, _y1 = np.min(coord, axis=0)[1], np.max(coord, axis=0)[1] _delta_x = _x1 - _x0 _delta_y = _y1 - _y0 # to extend filed of view _x0, _x1 = _x0 - 0.05*_delta_x, _x1 + 0.05*_delta_x _y0, _y1 = _y0 - 0.05*_delta_y, _y1 + 0.05 * _delta_y _px_x = np.array((coord[:, 0] - _x0) * im_size/(_x1 - _x0)).astype(int) _px_y = np.array((coord[:, 1] - _y0) * im_size/(_y1 - _y0)).astype(int) grid_x, grid_y = np.mgrid[_x0: _x1: (_x1 - _x0)/im_size, _y0: _y1: (_y1 - _y0)/im_size] grid_z = griddata(coord, self, (grid_x, grid_y), method='nearest') # transpform to 3D _dset = Dataset.from_array(grid_z) _dset.data_type = 'point_cloud' _dset.units = self.units _dset.quantity = self.quantity _dset.title = self.title _dset.set_dimension(0, Dimension(grid_x[:, 0], 'x')) _dset.x.dimension_type = 'spatial' _dset.x.units = sp_units _dset.x.quantity = 'distance' _dset.set_dimension(1, Dimension(grid_y[0], 'y')) _dset.y.dimension_type = 'spatial' _dset.y.units = sp_units _dset.y.quantity = 'distance' _dset.set_dimension(2, self.get_dimension_by_number(1)[0]) if len(self.shape) == 3: _dset.set_dimension(3, self.get_dimension_by_number(2)[0]) _dset.metadata = {'coord': np.array([_px_x, _px_y]).T} if 'variance' in self.metadata.keys(): grid_z_var = griddata(coord, self.metadata['variance'], (grid_x, grid_y), method='nearest') _dset.metadata['variance'] = grid_z_var return _dset @staticmethod def _min_dist(array): _sort_ar = np.sort(array) return np.min(_sort_ar[1:] - _sort_ar[:-1]) @staticmethod def _closest_point(array_coord, point): diff = array_coord - point return np.argmin(diff[:, 0]**2 + diff[:, 1]**2) @property def labels(self): labels = [] for key, dim in self._axes.items(): labels.append('{} ({})'.format(dim.quantity, dim.units)) return labels @property def title(self): return self._title @title.setter def title(self, value): if isinstance(value, str): self._title = value else: raise ValueError('title needs to be a string') @property def structures(self): return self._structures def add_structure(self, atoms, title=None): if isinstance(atoms, ase.Atoms): if title is None: title = atoms.get_chemical_formula() self._structures.update({title: atoms}) else: raise ValueError('structure not an ase.Atoms object') @property def units(self): return self._units @units.setter def units(self, value): if isinstance(value, str): self._units = value else: raise ValueError('units needs to be a string') @property def quantity(self): return self._quantity @quantity.setter def quantity(self, value): if isinstance(value, str): self._quantity = value else: raise ValueError('quantity needs to be a string') @property def data_type(self): return self._data_type @data_type.setter def data_type(self, value): if isinstance(value, str): if value.upper() in DataType._member_names_: self._data_type = DataType[value.upper()] else: self._data_type = DataType.UNKNOWN raise Warning('Supported data_types for plotting are only: ', DataType._member_names_) elif isinstance(value, DataType): self._data_type = value else: raise ValueError('data_type needs to be a string') @property def modality(self): return self._modality @modality.setter def modality(self, value): if isinstance(value, str): self._modality = value else: raise ValueError('modality needs to be a string') @property def source(self): return self._source @source.setter def source(self, value): if isinstance(value, str): self._source = value else: raise ValueError('source needs to be a string') @property def h5_dataset(self): return self._h5_dataset @h5_dataset.setter def h5_dataset(self, value): if isinstance(value, h5py.Dataset): self._h5_dataset = value elif value is None: self.hdf_close() else: raise TypeError('h5_dataset needs to be a hdf5 Dataset') @property def metadata(self): return self._metadata @metadata.setter def metadata(self, value): if isinstance(value, dict): if sys.getsizeof(value) < 64000: self._metadata = value else: raise ValueError('metadata dictionary too large, please use attributes for ' 'large additional data sets') else: raise ValueError('metadata needs to be a python dictionary') @property def original_metadata(self): return self._original_metadata @original_metadata.setter def original_metadata(self, value): if isinstance(value, dict): if sys.getsizeof(value) < 64000: self._original_metadata = value else: raise ValueError('original_metadata dictionary too large, please use attributes for ' 'large additional data sets') else: raise ValueError('original_metadata needs to be a python dictionary') @property def data_descriptor(self): return '{} ({})'.format(self.quantity, self.units) @property def variance(self): return self._variance @variance.setter def variance(self, value): if value is None: self._variance = None else: if np.array(value).shape != np.array(self).shape: raise ValueError('Variance array must have the same dimensionality as the dataset') if isinstance(value, da.Array) and not np.any(np.isnan(value.shape)): self._variance = value else: self._variance = da.from_array(np.array(value)) def fft(self, dimension_type=None): """ Gets the FFT of a sidpy.Dataset of any size The data_type of the sidpy.Dataset determines the dimension_type over which the fourier transform is performed over, if the dimension_type is not set explicitly. The fourier transformed dataset is automatically shifted to center of dataset. Parameters ---------- dimension_type: None, str, or sidpy.DimensionType - optional dimension_type over which fourier transform is performed, if None an educated guess will determine that from dimensions of sidpy.Dataset Returns ------- fft_dset: 2D or 3D complex sidpy.Dataset (not tested for higher dimensions) 2 or 3 dimensional matrix arranged in the same way as input Example ------- >> fft_dataset = sidpy_dataset.fft() >> fft_dataset.plot() """ if dimension_type is None: # test for data_type of sidpy.Dataset if self.data_type.name in ['IMAGE_MAP', 'IMAGE_STACK', 'SPECTRAL_IMAGE', 'IMAGE_4D']: dimension_type = self.dim_2.dimension_type else: dimension_type = self.dim_0.dimension_type if isinstance(dimension_type, str): dimension_type = DimensionType[dimension_type.upper()] if not isinstance(dimension_type, DimensionType): raise TypeError('Could not identify a dimension_type to perform Fourier transform on') axes = self.get_dimensions_by_type(dimension_type) if dimension_type.name in ['SPATIAL', 'RECIPROCAL']: if len(axes) != 2: raise TypeError('sidpy dataset of type', self.data_type, ' has no obvious dimension over which to perform fourier transform, please specify') if dimension_type.name == 'SPATIAL': new_dimension_type = DimensionType.RECIPROCAL else: new_dimension_type = DimensionType.SPATIAL elif dimension_type.name == 'SPECTRAL': if len(axes) != 1: raise TypeError('sidpy dataset of type', self.data_type, ' has no obvious dimension over which to perform fourier transform, please specify') new_dimension_type = DimensionType.SPECTRAL else: raise NotImplementedError('fourier transform not implemented for dimension_type ', dimension_type.name) fft_transform = np.fft.fftshift(da.fft.fftn(self, axes=axes)) fft_dset = self.like_data(fft_transform) fft_dset.units = 'a.u.' fft_dset.modality = 'fft' units_x = '1/' + self._axes[axes[0]].units fft_dset.set_dimension(axes[0], Dimension(np.fft.fftshift(np.fft.fftfreq(self.shape[axes[0]], d=get_slope(self._axes[axes[0]].values))), name='u', units=units_x, dimension_type=new_dimension_type, quantity='reciprocal')) if len(axes) > 1: units_y = '1/' + self._axes[axes[1]].units fft_dset.set_dimension(axes[1], Dimension(np.fft.fftshift(np.fft.fftfreq(self.shape[axes[1]], d=get_slope(self._axes[axes[1]].values))), name='v', units=units_y, dimension_type=new_dimension_type, quantity='reciprocal_length')) return fft_dset def flatten_complex(self): """ This function returns a dataset with real and imaginary components that have been flattened This is necessary for scenarios such as fitting of complex functions Must be a 2D or 1D dataset to begin with Output: - ouput_arr: sidpy.Dataset object """ assert self.ndim < 3, "flatten_complex() only works on 1D or 2D datasets, current dataset has {}".format( self.ndim) # Only the second dimension needs to be changed # Because we are stacking real and imaginary, this means we just tile the existing axis values if len(self._axes) == 1: index_ax = 0 elif len(self._axes) == 2: index_ax = 1 new_ax_values = np.tile(self._axes[index_ax].values, 2) output_arr = self.like_data(dask.array.hstack([self.real, self.imag])) output_arr.set_dimension(index_ax, Dimension(new_ax_values, name=output_arr._axes[index_ax].name, units=output_arr._axes[index_ax].units, dimension_type=output_arr._axes[index_ax].dimension_type, quantity=output_arr._axes[index_ax].quantity)) return output_arr # ##################################################### # Original dask.array functions replaced # ################################################## def __eq__(self, other): # TODO: Test __eq__ if not isinstance(other, Dataset): return False # if (self.__array__() == other.__array__()).all(): if (self.__array__().__eq__(other.__array__())).all(): if self._units != other._units: return False if self._quantity != other._quantity: return False if self._source != other._source: return False if self._data_type != other._data_type: return False if self._modality != other._modality: return False if self._axes != other._axes: return False if (self._variance is not None) and (other._variance is not None): if not (self._variance.__eq__(other._variance)).all(): return False elif (self._variance is not None) or (other._variance is not None): return False return True return False @property def T(self): return self.transpose() def abs(self): return self.like_data(super().__abs__(), title_suffix='_absolute_value') ###################################################### # Original dask.array functions handed through ################################################## @property def real(self): result = self.like_data(super().real) if self._variance is not None: result._variance = self._variance.real return result @property def imag(self): result = self.like_data(super().imag) if self._variance is not None: result._variance = self._variance.image return result # This is wrapper method for the methods that reduce dimensions def reduce_dims(original_method): @wraps(original_method) def wrapper_method(self, *args, **kwargs): result, arguments = original_method(self, *args, **kwargs) axis, keepdims = arguments.get('axis'), arguments.get('keepdims', False) if axis is None and not keepdims: return result.compute() if axis is None: axes = list(np.arange(self.ndim)) elif isinstance(axis, int): axes = [axis] else: axes = list(axis) return self.__reduce_dimensions(result, axes, keepdims) return wrapper_method @reduce_dims def all(self, axis=None, keepdims=False, split_every=None, out=None): result = self.like_data(super().all(axis=axis, keepdims=keepdims, split_every=split_every, out=out), title_prefix='all_aggregate_', checkdims=False) if self._variance is not None: result._variance = self._variance.all(axis=axis, keepdims=keepdims, split_every=split_every, out=out) return result, locals().copy() @reduce_dims def any(self, axis=None, keepdims=False, split_every=None, out=None): result = self.like_data(super().any(axis=axis, keepdims=keepdims, split_every=split_every, out=out), title_prefix='any_aggregate_', checkdims=False) if self._variance is not None: result._variance = self._variance.any(axis=axis, keepdims=keepdims, split_every=split_every, out=out) return result, locals().copy() @reduce_dims def min(self, axis=None, keepdims=False, split_every=None, out=None): result = self.like_data(super().min(axis=axis, keepdims=keepdims, split_every=split_every, out=out), title_prefix='min_aggregate_', checkdims=False) if self._variance is not None: if axis is not None: _min_ind_axis = super().argmin(axis=axis, split_every=split_every, out=out) _coords = np.array(list(np.ndindex(_min_ind_axis.shape))) #list? _inds = np.insert(_coords, axis, np.array(_min_ind_axis).flatten(), axis=1) _extracted_points = da.take(self._variance.flatten(), np.ravel_multi_index(_inds.T, (self._variance.shape))) result._variance = _extracted_points.reshape(result.shape).rechunk(result.chunksize) else: _ind = np.unravel_index(super().min(), self._variance.shape) result._variance = self._variance[_ind] return result, locals().copy() @reduce_dims def max(self, axis=None, keepdims=False, split_every=None, out=None): result = self.like_data(super().max(axis=axis, keepdims=keepdims, split_every=split_every, out=out), title_prefix='max_aggregate_', checkdims=False) if self._variance is not None: if axis is not None: _max_ind_axis = super().argmax(axis=axis, split_every=split_every, out=out) _coords = np.array(list(np.ndindex(_max_ind_axis.shape))) #list? _inds = np.insert(_coords, axis, np.array(_max_ind_axis).flatten(), axis=1) _extracted_points = da.take(self._variance.flatten(), np.ravel_multi_index(_inds.T, (self._variance.shape))) result._variance = _extracted_points.reshape(result.shape).rechunk(result.chunksize) return result, locals().copy() @reduce_dims def sum(self, axis=None, dtype=None, keepdims=False, split_every=None, out=None): result = self.like_data(super().sum(axis=axis, dtype=dtype, keepdims=keepdims, split_every=split_every, out=out), title_prefix='sum_aggregate_', checkdims=False) if self._variance is not None: result._variance = abs(self._variance).sum(axis=axis, dtype=dtype, keepdims=keepdims, split_every=split_every, out=out) #TODO imaginary return result, locals().copy() @reduce_dims def mean(self, axis=None, dtype=None, keepdims=False, split_every=None, out=None): result = self.like_data(super().mean(axis=axis, dtype=dtype, keepdims=keepdims, split_every=split_every, out=out), title_prefix='mean_aggregate_', checkdims=False) if (self._variance is not None) and (axis is not None): if type(axis) is tuple: sh = np.prod(np.array(self._variance.shape, dtype=int)[list(axis)]) else: sh = axis result._variance = self._variance.sum(axis=axis, dtype=dtype, keepdims=keepdims, split_every=split_every, out=out)/sh**2 return result, locals().copy() @reduce_dims def std(self, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None): result = self.like_data(super().std(axis=axis, dtype=dtype, keepdims=keepdims, ddof=0, split_every=split_every, out=out), title_prefix='std_aggregate_', checkdims=False) return result, locals().copy() @reduce_dims def var(self, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None): result = self.like_data(super().var(axis=axis, dtype=dtype, keepdims=keepdims, ddof=ddof, split_every=split_every, out=out), title_prefix='var_aggregate_', checkdims=False) return result, locals().copy() @reduce_dims def argmin(self, axis=None, split_every=None, out=None): result = self.like_data(super().argmin(axis=axis, split_every=split_every, out=out), title_prefix='argmin_aggregate_', reset_units=True, reset_quantity=True, check_dims=False) return result, locals().copy() @reduce_dims def argmax(self, axis=None, split_every=None, out=None): result = self.like_data(super().argmax(axis=axis, split_every=split_every, out=out), title_prefix='argmax_aggregate_', reset_units=True, reset_quantity=True, check_dims=False) return result, locals().copy() def angle(self, deg=False): result = self.like_data(da.angle(self, deg=deg), reset_units=True, reset_quantity=True, title_prefix='angle_', checkdims=True) if deg: result.units = 'degrees' else: result.units = 'radians' return result def conj(self): return self.like_data(super().conj(), reset_units=True, reset_quantity=True, title_prefix='conj_', checkdims=True) def astype(self, dtype, **kwargs): return self.like_data(super().astype(dtype=dtype, **kwargs), variance=self._variance) def flatten(self): result = self.like_data(super().flatten(), title_prefix='flattened_', check_dims=False) if self._variance is not None: result.variance = self._variance.flatten() return result def ravel(self): return self.flatten() def clip(self, min=None, max=None): return self.like_data(super().clip(min=min, max=max), reset_quantity=True, title_prefix='clipped_') def compute_chunk_sizes(self): return self.like_data(super().compute_chunk_sizes()) def cumprod(self, axis, dtype=None, out=None, method='sequential'): if axis is None: self = self.flatten() axis = 0 return self.like_data(super().cumprod(axis=axis, dtype=dtype, out=out, method=method), title_prefix='cumprod_', reset_quantity=True) def cumsum(self, axis, dtype=None, out=None, method='sequential'): if axis is None: self = self.flatten() axis = 0 return self.like_data(super().cumsum(axis=axis, dtype=dtype, out=out, method=method), title_prefix='cumsum_', reset_quantity=True) # What happens to the dimensions?? def dot(self, other): return self.from_array(super().dot(other)) def squeeze(self, axis=None): result = self.like_data(super().squeeze(axis=axis), title_prefix='Squeezed_', checkdims=False) if self._variance is not None: result._variance = self._variance.squeeze(axis=axis) if axis is None: shape_list = list(self.shape) axes = [i for i in range(self.ndim) if shape_list[i] == 1] elif isinstance(axis, int): axes = [axis] else: axes = list(axis) return self.__reduce_dimensions(result, axes, keepdims=False) def swapaxes(self, axis1, axis2): result = self.like_data(super().swapaxes(axis1, axis2), title_prefix='Swapped_axes_', checkdims=False) if self._variance is not None: result._variance = self._variance.swapaxes(axis1, axis2) new_order = np.arange(self.ndim) new_order[axis1] = axis2 new_order[axis2] = axis1 return self.__rearrange_axes(result, new_order) def transpose(self, *axes): result = self.like_data(super().transpose(*axes), title_prefix='Transposed_', checkdims=False) if self._variance is not None: result._variance = self._variance.transpose(*axes) if not axes: new_axes_order = range(self.ndim)[::-1] elif len(axes) == 1 and isinstance(axes[0], Iterable): new_axes_order = axes[0] else: new_axes_order = axes return self.__rearrange_axes(result, new_axes_order) def round(self, decimals=0): return self.like_data(super().round(decimals=decimals), title_prefix='Rounded_') def reshape(self, shape, merge_chunks=True, limit=None): # This somehow adds an extra dimension at the end # Will come back to this warnings.warn('Dimensional information will be lost.\ Please use fold, unfold to combine dimensions') if len(shape) == 1 and isinstance(shape[0], Iterable): new_shape = shape[0] else: new_shape = shape return super().reshape(*new_shape, merge_chunks) @reduce_dims def prod(self, axis=None, dtype=None, keepdims=False, split_every=None, out=None): result = self.like_data(super().prod(axis=axis, dtype=dtype, keepdims=keepdims, split_every=split_every, out=out), title_prefix='prod_aggregate', reset_units=True, reset_quantity=True, checkdims=False) return result, locals().copy() @reduce_dims def trace(self, offset=0, axis1=0, axis2=1, dtype=None): if self.ndim == 2: axes = None result = (super().trace(offset=offset)) else: axes = [axis1, axis2] result = self.like_data(super().trace(offset=offset, axis1=axis1, axis2=axis2, dtype=None), title_prefix='Trace_', checkdims=False) local_args = locals().copy() local_args['axis'] = axes return result, local_args def repeat(self, repeats, axis=None): result = self.like_data(super().repeat(repeats=repeats, axis=axis), title_prefix='Repeated_', checkdims=False) # result._axes = {} for i, dim in self._axes.items(): if axis != i: new_dim = dim.copy() else: new_dim = Dimension(np.repeat(dim.values, repeats=repeats), name=dim.name, quantity=dim.quantity, units=dim.units, dimension_type=dim.dimension_type) result.set_dimension(i, new_dim) return result @reduce_dims def moment(self, order, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None): result = self.like_data(super().moment(order=order, axis=axis, dtype=dtype, keepdims=keepdims, ddof=0, split_every=split_every, out=out), title_prefix='moment_aggregate_', checkdims=False) return result, locals().copy() def persist(self, **kwargs): return self.like_data(super().persist(**kwargs), title_prefix='persisted_') def rechunk(self, chunks='auto', threshold=None, block_size_limit=None, balance=False): return self.like_data(super().rechunk(chunks=chunks, threshold=threshold, block_size_limit=block_size_limit, balance=balance), title_prefix='Rechunked_') def fold(self, dim_order=None, method=None): """ This method collapses the dimensions of the sidpy dataset """ """ Parameters ---------- dim_order: List of lists or tuple of tuples -Each element corresponds to the order of axes in the corresponding new axis after the collapse -Default: None method: str -'spaspec': collapses the original dataset to a 2D dataset, where spatail dimensions form the zeroth axis and spectral dimensions form the first axis -'spa': combines all the spatial dimensions into a single dimension and the combined dimension will be first -'spec': combines all the spectral dimensions into a single dimension and the combined dimension will be last -Uses the user defined dim_order when set to None -Default: None Returns ------- Collapsed sidpy.Dataset object whose number of dimensions equals two if method=='spaspec' or len(dim_order) """ if method is None: if dim_order is None: raise NotImplementedError("Specify the dim_order or set the\ method to 'spaspec'") if not (isinstance(dim_order, list) or isinstance(dim_order, tuple)): raise NotImplementedError("dim_order should be a List or a Tuple") dim_order_list = [list(x) for x in dim_order] # Book-keeping for unfolding fold_attr = {'_axes': self._axes.copy()} if method == 'spaspec': dim_order_list = [[], []] for dim, axis in self._axes.items(): if axis.dimension_type == DimensionType.SPATIAL: dim_order_list[0].extend([dim]) elif axis.dimension_type == DimensionType.SPECTRAL: dim_order_list[1].extend([dim]) else: warnings.warn('One of the dimensions is neither Spatial\ nor Spectral Type and is considered to be a \ part of the last collapsed dimension') dim_order_list[1].extend([dim]) if method == 'spa': dim_order_list = [[]] for dim, axis in self._axes.items(): if axis.dimension_type == DimensionType.SPATIAL: dim_order_list[0].extend([dim]) else: dim_order_list.append([dim]) if len(dim_order_list[0]) == 0: raise NotImplementedError("No spatial dimensions found and the method is set to 'spa' ") if len(dim_order_list[0]) == 1: warnings.warn('Only one spatial dimension found\ Folding returns the original dataset') if method == 'spec': dim_order_list = [[]] for dim, axis in self._axes.items(): if axis.dimension_type == DimensionType.SPECTRAL: dim_order_list[-1].extend([dim]) else: dim_order_list.insert(-1, [dim]) if len(dim_order_list[-1]) == 0: raise NotImplementedError("No spectral dimensions found and the method is set to 'spec'") if len(dim_order_list[-1]) == 1: warnings.warn('Only one spatial dimension found\ Folding returns the original dataset') # We need the flattened list to transpose the original array dim_order_flattened = [item for sublist in dim_order_list for item in sublist] # Check if all the dimensions are accounted for, if len(dim_order_flattened) != len(self.shape): warnings.warn('All the dimensions that are not present in the dim_order \ are considered to be a part of last collapsed dimension') left_dims = set(np.arange(0, self.ndim)) - set(dim_order_flattened) dim_order_list[-1].extend(list(left_dims)) dim_order_flattened.extend(list(left_dims)) fold_attr['dim_order_flattened'] = dim_order_flattened fold_attr['dim_order'] = dim_order_list # Get the shape of the collapsed array new_shape = np.ones(len(dim_order_list)).astype(int) for i, dim in enumerate(dim_order_list): for d in dim: new_shape[i] *= self.shape[d] # Collapsed dask array transposed_dset = self.transpose(dim_order_flattened) folded_dset = self.like_data(da.reshape(transposed_dset, tuple(new_shape), merge_chunks=True), title_prefix='folded_', checkdims=False) fold_attr['shape_transposed'] = [self.shape[i] for i in dim_order_flattened] # Setting the dimensions for spaspec method if method == 'spaspec': folded_dset._axes[0].dimension_type = DimensionType.SPATIAL folded_dset._axes[1].dimension_type = DimensionType.SPECTRAL folded_dset.metadata['fold_attr'] = fold_attr # Setting the dimensions for a general case for i, dim in enumerate(dim_order_list): dim_types = [self._axes[d].dimension_type for d in dim] if dim_types.count(dim_types[0]) == len(dim_types): folded_dset._axes[i].dimension_type = dim_types[0] return folded_dset def unfold(self): try: shape_transposed = self.metadata['fold_attr']['shape_transposed'] dim_order_flattened = self.metadata['fold_attr']['dim_order_flattened'] old_axes = self.metadata['fold_attr']['_axes'] except: raise NotImplementedError('unfold only works on the dataset that was collapsed/folded by' ' the fold method') reshaped_dset = da.reshape(self, shape_transposed, merge_chunks=True) old_order = [dim_order_flattened.index(d) for d in range(len(dim_order_flattened))] unfolded_dset = self.like_data(da.transpose(reshaped_dset, old_order), title=self.title.replace('folded_', ''), checkdims=False) unfolded_dset._axes = {} for i, dim in old_axes.items(): unfolded_dset.set_dimension(i, dim.copy()) del unfolded_dset.metadata['fold_attr'] return unfolded_dset # Following methods are to be edited def adjust_axis(self, result, axis, title='', keepdims=False): if not keepdims: dim = 0 dataset = self.from_array(result) if isinstance(axis, int): axis = [axis] # for ax, dimension in self._axes.items(): # if int(ax) not in axis: # delattr(self, dimension.name) # delattr(self, f'dim_{ax}') # del self._axes[ax] for ax, dimension in self._axes.items(): if int(ax) not in axis: dataset.set_dimension(dim, dimension) dim += 1 else: dataset = self.like_data(result) dataset.title = title + self.title dataset.modality = f'sum axis {axis}' dataset.quantity = self.quantity dataset.source = self.source dataset.units = self.units return dataset def choose(self, choices): return self.like_data(super().choose(choices)) def __abs__(self): return self.like_data(super().__abs__(), title_suffix='_absolute_value') def __add__(self, other): return self.like_data(super().__add__(other)) def __radd__(self, other): return self.like_data(super().__radd__(other)) def __and__(self, other): return self.like_data(super().__and__(other)) def __rand__(self, other): return self.like_data(super().__rand__(other)) def __div__(self, other): return self.like_data(super().__div__(other)) def __rdiv__(self, other): return self.like_data(super().__rdiv__(other)) def __gt__(self, other): return self.like_data(super().__gt__(other)) def __ge__(self, other): return self.like_data(super().__ge__(other)) def __getitem__(self, idx): # Here we need to modify the dimensions of the sliced dataset using the argument index if not isinstance(idx, tuple): # This comes into play when slicing is done using 'None' or just integers. # For example: dset[4] or dset[None] idx = tuple([idx]) # The following line creates a new sidpy dataset with generic dimensions and .. # all the other attributes copied from 'self' aka parent dataset. sliced = self.like_data(super().__getitem__(idx), checkdims=False) # Delete the dimensions created by like_data sliced.del_dimension() old_dims = copy(self._axes) j, k = 0, 0 # j is for self (old_dims) and k is for the sliced dataset (new dimensions) for ind in idx: if ind is None: # Add a new dimension sliced.set_dimension(k, Dimension(1)) k += 1 elif isinstance(ind, (int, np.integer)): j += 1 elif isinstance(ind, (slice, list)): old_dim = old_dims[j] sliced.set_dimension(k, Dimension(old_dim[ind].values, name=old_dim.name, quantity=old_dim.quantity, units=old_dim.units, dimension_type=old_dim.dimension_type)) j += 1 k += 1 elif isinstance(ind, (np.ndarray, da.Array)): if not ind.ndim == 1: raise NotImplementedError('Multi Dimensional Slicing of sidpy Dataset' 'is not available at this moment, please' 'raise an issue on out GitHub page') old_dim = old_dims[j] sliced.set_dimension(k, Dimension(old_dim[np.array(ind)].values, name=old_dim.name, quantity=old_dim.quantity, units=old_dim.units, dimension_type=old_dim.dimension_type)) j += 1 k += 1 elif ind is Ellipsis: start_dim = idx.index(Ellipsis) ellipsis_dims = sliced.ndim - (len(idx) - 1) stop_dim = start_dim + ellipsis_dims for l in range(start_dim, stop_dim): old_dim = old_dims[j] sliced.set_dimension(k, old_dim) j += 1 k += 1 # Adding the rest of the dimensions for k in range(k, sliced.ndim): old_dim = old_dims[j] sliced.set_dimension(k, Dimension(old_dim.values, name=old_dim.name, quantity=old_dim.quantity, units=old_dim.units, dimension_type=old_dim.dimension_type)) j += 1 k += 1 return sliced def __invert__(self): return self.like_data(super().__invert__()) def __lshift__(self, other): return self.like_data(super().__lshift__(other)) def __rlshift__(self, other): return self.like_data(super().__rlshift__(other)) def __lt__(self, other): return self.like_data(super().__lt__(other)) def __le__(self, other): return self.like_data(super().__lt__(other)) def __mod__(self, other): return self.like_data(super().__lshift__(other)) def __rmod__(self, other): return self.like_data(super().__rmod__(other)) def __mul__(self, other): return self.like_data(super().__mul__(other)) def __rmul__(self, other): return self.like_data(super().__rmul__(other)) def __ne__(self, other): return self.like_data(super().__ne__(other)) def __neg__(self): return self.like_data(super().__neg__()) def __or__(self, other): return self.like_data(super().__or__(other)) def __ror__(self, other): return self.like_data(super().__ror__(other)) def __pos__(self): return self.like_data(super().__pos__()) def __pow__(self, other): return self.like_data(super().__pow__(other)) def __rpow__(self, other): return self.like_data(super().__rpow__(other)) def __rshift__(self, other): return self.like_data(super().__rshift__(other)) def __rrshift__(self, other): return self.like_data(super().__rrshift__(other)) def __sub__(self, other): return self.like_data(super().__sub__(other)) def __rsub__(self, other): return self.like_data(super().__rsub__(other)) def __truediv__(self, other): return self.like_data(super().__truediv__(other)) def __rtruediv__(self, other): return self.like_data(super().__rtruediv__(other)) def __floordiv__(self, other): return self.like_data(super().__floordiv__(other)) def __rfloordiv__(self, other): return self.like_data(super().__rfloordiv__(other)) def __xor__(self, other): return self.like_data(super().__xor__(other)) def __rxor__(self, other): return self.like_data(super().__rxor__(other)) def __matmul__(self, other): return self.like_data(super().__matmul__(other)) def __rmatmul__(self, other): return self.like_data(super().__rmatmul__(other)) def __array_ufunc__(self, numpy_ufunc, method, *inputs, **kwargs): out = kwargs.get("out", ()) if method == "__call__": # if numpy_ufunc is np.matmul: # from dask.array.routines import matmul # # # special case until apply_gufunc handles optional dimensions # return self.like_data(matmul(*inputs, **kwargs)) if numpy_ufunc.signature is not None: from dask.array.gufunc import apply_gufunc return self.like_data(apply_gufunc( numpy_ufunc, numpy_ufunc.signature, *inputs, **kwargs)) if numpy_ufunc.nout > 1: from dask.array import ufunc try: da_ufunc = getattr(ufunc, numpy_ufunc.__name__) except AttributeError: return NotImplemented return self.like_data(da_ufunc(*inputs, **kwargs)) else: return self.like_data(dask.array.core.elemwise(numpy_ufunc, *inputs, **kwargs)) elif method == "outer": from dask.array import ufunc try: da_ufunc = getattr(ufunc, numpy_ufunc.__name__) except AttributeError: return NotImplemented return self.like_data(da_ufunc.outer(*inputs, **kwargs)) else: return NotImplemented def convert_hyperspy(s): """ imports a hyperspy signal object into sidpy.Dataset Parameters ---------- s: hyperspy dataset Return ------ dataset: sidpy.Dataset """ try: import hyperspy.api as hs except ModuleNotFoundError: raise ModuleNotFoundError("Hyperspy is not installed") if not isinstance(s, (hs.signals.Signal1D, hs.signals.Signal2D)): raise TypeError('This is not a hyperspy signal object') dataset = Dataset.from_array(s, name=s.metadata.General.title) # Add dimension info axes = s.axes_manager.as_dictionary() if isinstance(s, hs.signals.Signal1D): if s.data.ndim < 2: dataset.data_type = 'spectrum' elif s.data.ndim > 1: if s.data.ndim == 2: dataset = Dataset.from_array(np.expand_dims(s, 2), title=s.metadata.General.title) dataset.set_dimension(2, Dimension([0], name='y', units='pixel', quantity='distance', dimension_type='spatial')) dataset.data_type = DataType.SPECTRAL_IMAGE for key, axis in axes.items(): if axis['navigate']: dimension_type = 'spatial' else: dimension_type = 'spectral' dim_array = np.arange(axis['size']) * axis['scale'] + axis['offset'] if axis['units'] == '': axis['units'] = 'frame' dataset.set_dimension(int(key[-1]), Dimension(dim_array, name=axis['name'], units=axis['units'], quantity=axis['name'], dimension_type=dimension_type)) elif isinstance(s, hs.signals.Signal2D): if s.data.ndim < 4: if s.data.ndim == 2: dataset.data_type = 'image' elif s.data.ndim == 3: dataset.data_type = 'image_stack' for key, axis in axes.items(): if axis['navigate']: dimension_type = 'temporal' else: dimension_type = 'spatial' dim_array = np.arange(axis['size']) * axis['scale'] + axis['offset'] if axis['units'] == '': axis['units'] = 'pixel' dataset.set_dimension(int(key[-1]), Dimension(dim_array, name=axis['name'], units=axis['units'], quantity=axis['name'], dimension_type=dimension_type)) elif s.data.ndim == 4: dataset.data_type = 'IMAGE_4D' for key, axis in axes.items(): if axis['navigate']: dimension_type = 'spatial' else: dimension_type = 'reciprocal' dim_array = np.arange(axis['size']) * axis['scale'] + axis['offset'] dataset.set_dimension(int(key[-1]), Dimension(dim_array, name=axis['name'], units=axis['units'], quantity=axis['name'], dimension_type=dimension_type)) dataset.metadata = dict(s.metadata) dataset.original_metadata = dict(s.original_metadata) dataset.title = dataset.metadata['General']['title'] dataset.units = dataset.metadata['Signal']['quantity '].split('(')[-1][:-1] dataset.quantity = dataset.metadata['Signal']['quantity '].split('(')[0] dataset.source = 'hyperspy' return dataset