brimfile.data

  1import numpy as np
  2import asyncio
  3
  4import warnings
  5from typing import Any
  6from numpy.typing import NDArray
  7
  8from .file_abstraction import FileAbstraction, sync, _async_getitem, _gather_sync
  9from .utils import concatenate_paths, list_objects_matching_pattern_async, get_object_name, set_object_name
 10from .utils import np_array_to_smallest_int_type, _determine_chunk_size
 11
 12from .metadata import Metadata
 13from .metadata.types import MetadataItem
 14
 15from numbers import Number
 16
 17from . import units
 18from .analysis_results import AnalysisResults
 19from .calibration import Calibration
 20from .constants import brim_obj_names
 21
 22__docformat__ = "google"
 23
 24
 25class Data:
 26    """
 27    Represents a data group within the brim file.
 28    """
 29    # make AnalysisResults available as an attribute of Data
 30    AnalysisResults = AnalysisResults
 31
 32    def __init__(self, file: FileAbstraction, path: str, *, 
 33                 newly_created: bool = False, _initialize: bool = True):
 34        """
 35        Initialize the Data object. This constructor should not be called directly.
 36
 37        Args:
 38            file (File): The parent File object.
 39            path (str): The path to the data group within the file.
 40            newly_created (bool): Whether this data group is being created as new.
 41                            If True, the constructor will not attempt to load spatial mapping.
 42            _initialize (bool): FOR INTERNAL USE ONLY. Whether to automatically initialize the current data group. 
 43                Set to False if you want to initialize them manually later using the _init_async() method. Default is True.
 44        """
 45        self._file = file
 46        self._path = path
 47        
 48
 49        if _initialize:
 50            sync(self._init_async(newly_created=newly_created))        
 51    
 52    async def _init_async(self, newly_created: bool = False) -> None:
 53        """
 54        See __init__() for the description of the arguments.
 55        """
 56        self._group = await self._file.open_group(self._path)
 57
 58        self._sparse = await self._load_sparse_flag_async()
 59        # the _spatial_map is None for non sparse data but the _spatial_map_px_size should always be valid
 60        self._spatial_map, self._spatial_map_px_size = await self._load_spatial_mapping_async() if not newly_created else (None, None)
 61
 62    def get_name(self):
 63        """
 64        Returns the name of the data group.
 65        """
 66        return sync(get_object_name(self._file, self._path))
 67    
 68    def get_index(self):
 69        """
 70        Returns the index of the data group.
 71        """
 72        return int(self._path.split('/')[-1].split('_')[-1])
 73
 74    async def _load_sparse_flag_async(self) -> bool:
 75        """
 76        Load the 'Sparse' flag for the data group.
 77
 78        Returns:
 79            bool: The value of the 'Sparse' flag, or False if the attribute is not found or invalid.
 80        """
 81        try:
 82            sparse = await self._file.get_attr(self._group, 'Sparse')
 83            if isinstance(sparse, bool):
 84                return sparse
 85            else:
 86                warnings.warn(
 87                    f"Invalid value for 'Sparse' attribute in {self._path}. Expected a boolean, got {type(sparse)}. Defaulting to False.")
 88                return False
 89        except Exception:
 90            # if the attribute is not found, return the default value False
 91            return False
 92
 93    async def _load_spatial_mapping_async(self, load_in_memory: bool=True) -> tuple:
 94        """
 95        Load a spatial mapping in the same format as 'Cartesian visualisation',
 96        irrespectively on whether 'Spatial_map' is defined instead.
 97        -1 is used for "empty" pixels in the image
 98        Args:
 99            load_in_memory (bool): Specify whether the map should be forced to load in memory or just opened as a dataset.
100        Returns:
101            The spatial map and the corresponding pixel size as a tuple of 3 Metadata.Item, both in the order z, y, x.
102            If the spatial mapping is not defined in the file, returns None for the spatial map.
103            The pixel size is read from the data group for non-sparse data.
104        """
105        cv = None
106        px_size = 3*(Metadata.Item(value=1, units=None),)
107
108        cv_path = concatenate_paths(
109            self._path, brim_obj_names.data.cartesian_visualisation)
110        sm_path = concatenate_paths(
111            self._path, brim_obj_names.data.spatial_map)
112        
113        if await self._file.object_exists(cv_path):
114            cv = await self._file.open_dataset(cv_path)
115
116            #read the pixel size from the 'Cartesian visualisation' dataset
117            px_size_val = None
118            px_size_units = None
119            try:
120                px_size_val = await self._file.get_attr(cv, 'element_size')
121                if px_size_val is None or len(px_size_val) != 3:
122                    raise ValueError(
123                        "The 'element_size' attribute of 'Cartesian_visualisation' must be a tuple of 3 elements")
124            except Exception:
125                px_size_val = 3*(1,)
126                warnings.warn(
127                    "No pixel size defined for Cartesian visualisation")            
128            px_size_units = await units.of_attribute(
129                    self._file, cv, 'element_size')
130            px_size = ()
131            for i in range(3):
132                # if px_size_val[i] is not a number, set it to 1 and px_size_units to None
133                if isinstance(px_size_val[i], Number):
134                    px_size += (Metadata.Item(px_size_val[i], px_size_units), )
135                else:
136                    px_size += (Metadata.Item(1, None), )
137                    
138
139            if load_in_memory:
140                cv = await cv.to_np_array()  # load the spatial map in memory as a numpy array
141                cv = np_array_to_smallest_int_type(cv)
142
143        elif await self._file.object_exists(sm_path):
144            async def load_spatial_map_from_file():
145                async def load_coordinate_from_sm(coord: str):
146                    res = np.empty(0)  # empty array
147                    try:
148                        res = await self._file.open_dataset(
149                            concatenate_paths(sm_path, coord))
150                        res = await res.to_np_array()
151                        res = np.squeeze(res)  # remove single-dimensional entries
152                    except Exception as e:
153                        # if the coordinate does not exist, return an empty array
154                        pass
155                    if len(res.shape) > 1:
156                        raise ValueError(
157                            f"The 'Spatial_map/{coord}' dataset is not a 1D array as expected")
158                    return res
159
160                def check_coord_array(arr, size):
161                    if arr.size == 0:
162                        return np.zeros(size)
163                    elif arr.size != size:
164                        raise ValueError(
165                            "The 'Spatial_map' dataset is invalid")
166                    return arr
167
168                x, y, z = await asyncio.gather(
169                    load_coordinate_from_sm('x'),
170                    load_coordinate_from_sm('y'),
171                    load_coordinate_from_sm('z')
172                    )
173                size = max([x.size, y.size, z.size])
174                if size == 0:
175                    raise ValueError("The 'Spatial_map' dataset is empty")
176                x = check_coord_array(x, size)
177                y = check_coord_array(y, size)
178                z = check_coord_array(z, size)
179                return x, y, z
180
181            def calculate_step(x):
182                n = len(np.unique(x))
183                if n == 1:
184                    d = None
185                else:
186                    d = (np.max(x)-np.min(x))/(n-1)
187                return n, d
188
189            x, y, z = await load_spatial_map_from_file()
190
191            # TODO extend the reconstruction to non-cartesian cases
192
193            nX, dX = calculate_step(x)
194            nY, dY = calculate_step(y)
195            nZ, dZ = calculate_step(z)
196
197            indices = np_array_to_smallest_int_type(np.lexsort((x, y, z)))
198            cv = np.reshape(indices, (nZ, nY, nX))
199
200            px_size_units = await units.of_object(self._file, sm_path)
201            px_size = ()
202            for i in range(3):
203                px_sz = (dZ, dY, dX)[i]
204                px_unit = px_size_units
205                if px_sz is None:
206                    px_sz = 1
207                    px_unit = None
208                px_size += (Metadata.Item(px_sz, px_unit),)
209        elif not self._sparse:
210            try:
211                px_sz = await self._file.get_attr(self._group, 'element_size')
212                if len(px_sz) != 3:
213                    raise ValueError(
214                        "The 'element_size' attribute must be a tuple of 3 elements")
215                px_unit = None
216                try:
217                    px_unit = await units.of_attribute(self._file, self._group, 'element_size')
218                except Exception:
219                    warnings.warn("Pixel size unit is not provided for non-sparse data.")
220                px_size = tuple(Metadata.Item(el, px_unit) for el in px_sz)
221            except Exception:
222                warnings.warn("Pixel size is not provided for non-sparse data.")
223
224        return cv, px_size
225
226    def get_PSD(self) -> tuple:
227        """
228        LOW LEVEL FUNCTION
229
230        Retrieve the Power Spectral Density (PSD) and frequency from the current data group.
231        Note: this function exposes the internals of the brim file and thus the interface might change in future versions.
232        Use only if more specialized functions are not working for your application!
233        Returns:
234            tuple: (PSD, frequency, PSD_units, frequency_units)
235                - PSD: A 2D (or more) numpy array containing all the spectra (see [specs](https://github.com/brillouin-imaging/Brillouin-standard-file/blob/main/docs/brim_file_specs.md) for more details).
236                - frequency: A numpy array representing the frequency data (see [specs](https://github.com/brillouin-imaging/Brillouin-standard-file/blob/main/docs/brim_file_specs.md) for more details).
237                - PSD_units: The units of the PSD.
238                - frequency_units: The units of the frequency.
239        """
240        warnings.warn(
241            "Data.get_PSD is deprecated and will be removed in a future release. "
242            "Use Data.get_PSD_as_spatial_map instead.",
243            DeprecationWarning,
244            stacklevel=2,
245        )
246        PSD, frequency = _gather_sync(
247            self._file.open_dataset(concatenate_paths(
248                self._path, brim_obj_names.data.PSD)),
249            self._file.open_dataset(concatenate_paths(
250                self._path, brim_obj_names.data.frequency))
251        )
252        # retrieve the units of the PSD and frequency
253        PSD_units, frequency_units = _gather_sync(
254            units.of_object(self._file, PSD),
255            units.of_object(self._file, frequency)
256        )
257
258        return PSD, frequency, PSD_units, frequency_units
259    
260    def get_PSD_as_spatial_map(self, *, broadcast_frequency: bool = True) -> tuple:
261        """
262        Retrieve the Power Spectral Density (PSD) as a spatial map and the frequency from the current data group.
263        Arguments:
264            broadcast_frequency (bool): Whether to broadcast the frequency array to match the shape of the PSD if they have different shapes. 
265                This is useful when the frequency is the same for all spectra and thus stored as a 1D array, while the PSD has a spatial dimension. 
266                If False, the function will return a 1D array for the frequency, if the frequency is the same for all spectra.
267        Returns:
268            tuple: (PSD, frequency, PSD_units, frequency_units)
269                - PSD: A 4D (or more) numpy array containing all the spectra. Dimensions are z, y, x, [parameters], spectrum.
270                - frequency: A numpy array representing the frequency data, which has the same shape as PSD or a 1D array (see `broadcast_frequency`).
271                - PSD_units: The units of the PSD.
272                - frequency_units: The units of the frequency.
273        """
274        PSD, frequency = _gather_sync(
275            self._file.open_dataset(concatenate_paths(
276                self._path, brim_obj_names.data.PSD)),        
277            self._file.open_dataset(concatenate_paths(
278                self._path, brim_obj_names.data.frequency))
279            )        
280        # retrieve the units of the PSD and frequency
281        PSD_units, frequency_units = _gather_sync(
282            units.of_object(self._file, PSD),
283            units.of_object(self._file, frequency)
284        )
285
286        # ensure PSD and frequency are numpy arrays
287        PSD = np.array(PSD)  
288        frequency = np.array(frequency)  # ensure it's a numpy array
289        
290        # if the frequency is not the same for all spectra, broadcast it to match the shape of PSD
291        # if it is the same for all spectra, broadcast_frequency determines whether to return it as a 1D array or broadcast it to match the shape of PSD
292        if frequency.ndim > 1 or (broadcast_frequency and frequency.shape != PSD.shape):
293            frequency = np.broadcast_to(frequency, PSD.shape)
294        
295        if self._sparse:
296            if self._spatial_map is None:
297                raise ValueError("The data is defined as sparse, but no spatial mapping is provided.")
298            sm = np.array(self._spatial_map)
299            # reshape the PSD and frequency to have the spatial dimensions first      
300            PSD = PSD[sm, ...]
301            # reshape the frequency only if it is not the same for all spectra
302            if frequency.ndim > 1:
303                frequency = frequency[sm, ...]
304
305        return PSD, frequency, PSD_units, frequency_units
306
307    def _get_spectrum(self, index: int | tuple[int, int, int]) -> tuple:
308        """
309        Synchronous wrapper for `_get_spectrum_async` (see doc for `brimfile.data.Data._get_spectrum_async`)
310        """
311        return sync(self._get_spectrum_async(index))
312    async def _get_spectrum_async(self, index: int | tuple[int, int, int]) -> tuple:
313        """
314        Retrieve a spectrum from the data group by its index or coordinates.
315
316        Args:
317            index (int | tuple[int, int, int]): The index (for sparse data) or z, y, x coordinates (for non-sparse data) of the spectrum to retrieve.
318
319        Returns:
320            tuple: (PSD, frequency, PSD_units, frequency_units) for the specified index. 
321                    PSD can be 1D or more (if there are additional parameters);
322                    frequency has the same size as PSD
323        Raises:
324            IndexError: If the index is out of range for the PSD dataset.
325        """
326        if self._sparse and not isinstance(index, int):
327            raise ValueError("For sparse data, index must be an integer.")
328        elif not self._sparse and not (isinstance(index, tuple) and len(index) == 3):
329            raise ValueError("For non-sparse data, index must be a tuple of (z, y, x) coordinates.")
330            
331        # index = -1 corresponds to no spectrum
332        if self._sparse and index < 0:
333            return None, None, None, None
334        elif not self._sparse and any(i < 0 for i in index):
335            return None, None, None, None
336        PSD, frequency = await asyncio.gather(
337            self._file.open_dataset(concatenate_paths(
338                self._path, brim_obj_names.data.PSD)),                       
339            self._file.open_dataset(concatenate_paths(
340                self._path, brim_obj_names.data.frequency))
341            )
342        if self._sparse and index >= PSD.shape[0]:
343            raise IndexError(
344                f"index {index} out of range for PSD with shape {PSD.shape}")
345        elif not self._sparse and any(i >= PSD.shape[j] for j, i in enumerate(index)):
346            raise IndexError(
347                f"index {index} out of range for PSD with shape {PSD.shape}")
348        # retrieve the units of the PSD and frequency
349        PSD_units, frequency_units = await asyncio.gather(
350            units.of_object(self._file, PSD),
351            units.of_object(self._file, frequency)
352        )
353        # add ellipsis to the index to select the spectrum and the corresponding frequency
354        if self._sparse:
355            index = (index, ...)
356        else:
357            index = index + (..., )
358        # map index to the frequency array, considering the broadcasting rules
359        index_frequency = index
360        if frequency.ndim < PSD.ndim:
361            if self._sparse:
362                # given the definition of the brim file format,
363                # if the frequency has less dimensions that PSD,
364                # it can only be because it is the same for all the spatial position (first dimension)
365                index_frequency = (..., )
366            else:
367                unassigned_indices = PSD.ndim - frequency.ndim
368                if unassigned_indices == 3:
369                    # if the frequency has no spatial dimension, it is the same for all the spatial positions
370                    index_frequency = (..., )
371                else:
372                    # if the frequency has some spatial dimensions but not all, we need to add the corresponding indices to the index of the frequency
373                    index_frequency = index[-unassigned_indices:] + (..., )
374        #get the spectrum and the corresponding frequency at the specified index
375        PSD, frequency = await asyncio.gather(
376            _async_getitem(PSD, index),
377            _async_getitem(frequency, index_frequency)
378        )
379        #broadcast the frequency to match the shape of PSD if needed
380        if frequency.ndim < PSD.ndim:
381            frequency = np.broadcast_to(frequency, PSD.shape)
382        return PSD, frequency, PSD_units, frequency_units
383
384    def get_spectrum_in_image(self, coor: tuple) -> tuple:
385        """
386        Retrieve a spectrum from the data group using spatial coordinates.
387
388        Args:
389            coor (tuple): A tuple containing the z, y, x coordinates of the spectrum to retrieve.
390
391        Returns:
392            tuple: A tuple containing the PSD, frequency, PSD_units, frequency_units for the specified coordinates. See `Data._get_spectrum_async` for details.
393        """
394        if len(coor) != 3:
395            raise ValueError("coor must contain 3 values for z, y, x")
396
397        if self._sparse:
398            index = int(self._spatial_map[coor])
399            return self._get_spectrum(index)
400        else:
401            return self._get_spectrum(coor)
402          
403    async def get_spectrum_and_all_quantities_in_image_async(self, ar: 'Data.AnalysisResults', coor: tuple, index_peak: int = 0) -> tuple[tuple, dict]:
404        """
405        Retrieve the spectrum and all available quantities from the analysis results at a specific spatial coordinate.
406
407        Args:
408            ar (Data.AnalysisResults): The analysis results object to retrieve quantities from.
409            coor (tuple): A tuple containing the z, y, x coordinates in the image.
410            index_peak (int, optional): The index of the peak to retrieve (for multi-peak fits). Defaults to 0.
411
412        Returns:
413            tuple: A tuple containing:
414                - spectrum (tuple): (PSD, frequency, PSD_units, frequency_units) at the specified coordinate
415                - quantities (dict): Dictionary of Metadata.Item in the form result[quantity.name][peak.name]
416        """
417        if len(coor) != 3:
418            raise ValueError("coor must contain 3 values for z, y, x")
419        index = coor
420        if self._sparse:
421            index = int(self._spatial_map[coor])
422        spectrum, quantities = await asyncio.gather(
423            self._get_spectrum_async(index),
424            ar._get_all_quantities_at_index(index, index_peak)
425        )
426        return spectrum, quantities
427    def get_spectrum_and_all_quantities_in_image(self, ar: 'Data.AnalysisResults', coor: tuple, index_peak: int = 0) -> tuple[tuple, dict]:
428        """
429        Synchronous wrapper for `get_spectrum_and_all_quantities_in_image_async` (see doc for `brimfile.data.Data.get_spectrum_and_all_quantities_in_image_async`)
430        """
431        return sync(self.get_spectrum_and_all_quantities_in_image_async(ar, coor, index_peak))
432
433    def get_metadata(self):
434        """
435        Returns the metadata associated with the current Data group
436        Note that this contains both the general metadata stored in the file (which might be redifined by the specific data group)
437        and the ones specific for this data group
438        """
439        return Metadata(self._file, self._path)
440
441    def get_num_parameters(self) -> tuple:
442        """
443        Retrieves the number of parameters
444
445        Returns:
446            tuple: The shape of the parameters if they exist, otherwise an empty tuple.
447        """
448        pars, _ = self.get_parameters()
449        return pars.shape if pars is not None else ()
450
451    def get_parameters(self) -> list:
452        """
453        Retrieves the parameters  and their associated names.
454
455        If PSD.ndims > 2, the parameters are stored in a separate dataset.
456
457        Returns:
458            list: A tuple containing the parameters and their names if there are any, otherwise None.
459        """
460        pars_full_path = concatenate_paths(
461            self._path, brim_obj_names.data.parameters)
462        if sync(self._file.object_exists(pars_full_path)):
463            pars = sync(self._file.open_dataset(pars_full_path))
464            pars_names = sync(self._file.get_attr(pars, 'Name'))
465            return (pars, pars_names)
466        return (None, None)
467
468    def create_calibration_group(self, *, index: NDArray[np.integer] | None = None, calibration_data: list[dict[str, Any]] | None = None,
469                                 timestamp: list[NDArray[Any]] | None = None, same_as: int | None = None, attributes: dict[str, MetadataItem] = None,
470                                 compression: FileAbstraction.Compression = FileAbstraction.Compression()) -> Calibration:
471        """
472        Create a new calibration group in the current data group.
473        For more details on the expected format of the calibration data, see https://github.com/brillouin-imaging/Brillouin-standard-file/blob/main/docs/brim_file_specs.md.
474
475        Parameters:
476            index (np.array | None, optional): Index array for the calibration spectra. For sparse data,
477                this must be 1D; for non-sparse data, this must be 3D.  
478                It can be omitted if each element in `calibration_data` contains only one spectrum.
479            calibration_data (list[dict[str, Any]] | None, optional): Calibration entries to store.
480                Each dictionary must contain `spectra` and `shift` keys, and may provide `shift_units`.                
481            timestamp (list[NDArray[Any]] | None, optional): Timestamp arrays corresponding to each calibration
482                entry. If provided, its length must match `calibration_data`. Defaults to None.
483            same_as (int | None, optional): If provided, links this calibration group to an existing
484                calibration via the `Same_as` attribute. When set, the other data arguments are ignored.
485                Defaults to None.
486            attributes (dict[str, MetadataItem], optional): Additional attributes to attach to the calibration group.
487                Can be one of ('Datetime', 'Description', 'Temperature', 'FSR') with the relative units (when relevant).
488            compression (FileAbstraction.Compression, optional): Compression settings used for created
489                datasets. Defaults to FileAbstraction.Compression().
490
491        Returns:
492            Calibration: The newly created calibration group.
493
494        Raises:
495            ValueError: If the provided calibration data, index, or timestamp is invalid or inconsistent.
496        """
497        calibration_path = concatenate_paths(self._path, brim_obj_names.data.calibration)
498        calibration_group = sync(self._file.create_group(calibration_path))
499
500        # if same_as is provided, create the 'Same_as' attribute to link the calibration group to an existing one
501        if same_as is not None:
502            sync(self._file.create_attr(calibration_group, 'Same_as', same_as))
503        else: # if same_as is provided, the other parameters are ignored
504            # check that calibration_data is provided and valid
505            if calibration_data is None:
506                raise ValueError("'calibration_data' is required when 'same_as' is not provided")
507            if not isinstance(calibration_data, (list, tuple)):
508                calibration_data = [calibration_data,]
509            # check that index is valid if provided
510            if index is not None:
511                # TODO: check of the shape of 'index' is compatible with PSD
512                if self._sparse and index.ndim != 1:
513                    raise ValueError("'index' must be a 1D array for sparse data")
514                if not self._sparse and index.ndim != 3:
515                    raise ValueError("'index' must be a 3D array for non-sparse data")
516            # check that timestamp is valid if provided
517            if timestamp is not None: 
518                if not isinstance(timestamp, (list, tuple)):
519                    timestamp = [timestamp,]
520                if len(timestamp) != len(calibration_data):
521                    raise ValueError("If 'timestamp' is provided, it must have the same length as 'calibration_data'")            
522
523            for m, calib in enumerate(calibration_data):
524                # check that each element in calibration_data is a dictionary containing 'spectra' and 'shift' keys
525                if not isinstance(calib, dict):
526                    raise ValueError("Each element in 'calibration_data' must be a dictionary")
527                if 'spectra' not in calib.keys() or 'shift' not in calib.keys():
528                    raise ValueError("Each calibration data dictionary must contain 'spectra' and 'shift' keys")
529                # retrieve the spectra, shift and shift_units from the calibration data and check that they are valid
530                cal_spectra = np.array(calib['spectra'])
531                if cal_spectra.ndim != 2:
532                    raise ValueError("'spectra' in calibration data must be a 2D array. If only one spectrum is provided, set the first dimension to 1.")
533                cal_shift = calib['shift']
534                cal_shift_units = calib.get('shift_units', None)
535                if cal_shift_units is None:
536                    cal_shift_units = 'GHz'
537                    warnings.warn("No units provided for 'shift' in calibration data, defaulting to GHz")
538                # check that index is compatible with the shape of the spectra
539                if index is None and cal_spectra.shape[0] != 1:
540                    raise ValueError("If 'index' is not provided, each element in 'calibration_data' must contain only one spectrum (i.e. have shape (1, n))")
541                if index is not None and np.max(index) >= cal_spectra.shape[0]:
542                    raise ValueError("If 'index' is provided, its maximum value must be less than the number of spectra in each calibration data element")
543                # add the m arrays together with their attributes to the file
544                spectra_dataset = sync(self._file.create_dataset(calibration_group, f'{m}', cal_spectra, chunk_size=_determine_chunk_size(cal_spectra), compression=compression))
545                sync(self._file.create_attr(spectra_dataset, 'Shift', cal_shift))
546                units.add_to_attribute(self._file, spectra_dataset, 'Shift', cal_shift_units)
547                if timestamp is not None:
548                    timestamp_array = np.array(timestamp[m])
549                    if timestamp_array.ndim != 1 or timestamp_array.shape[0] != cal_spectra.shape[0]:
550                        raise ValueError("Each element in 'timestamp' must be a 1D array with the same length as the number of spectra in each calibration data element")
551                    sync(self._file.create_dataset(calibration_group, f'Timestamp_{m}', timestamp_array, compression=compression))
552            # add the index array to the file
553            if index is not None:
554                sync(self._file.create_dataset(calibration_group, 'Index', index, compression=compression))
555        
556        from .calibration import _STANDARD_ATTRIBUTES
557        # add any additional attributes to the calibration group, checking that they do not overwrite the standard
558        if attributes is not None:
559            for key, value in attributes.items():
560                if key not in _STANDARD_ATTRIBUTES:
561                    warnings.warn(f"Attribute '{key}' is not a standard attribute for calibration groups.\
562                                   Standard attributes are: {', '.join(_STANDARD_ATTRIBUTES)}. \
563                                   Make sure this is intentional!")
564                if not isinstance(value, MetadataItem):
565                    value = MetadataItem(value)
566                sync(self._file.create_attr(calibration_group, key, value.value))
567                if value.units is not None:
568                    units.add_to_attribute(self._file, calibration_group, key, value.units)
569
570        return Calibration(self._file, calibration_path, data_group=self)
571    
572    def get_calibration(self) -> Calibration:
573        """
574        Synchronous wrapper for `get_calibration_async` (see doc for `brimfile.data.Data.get_calibration_async`)
575        """
576        return sync(self.get_calibration_async())
577
578    async def get_calibration_async(self) -> Calibration:
579        """
580        Retrieve the calibration group associated with the current data group.
581
582        Returns:
583            Calibration: The calibration group associated with the current data group.
584
585        Raises:
586            ValueError: If no calibration group is found in the current data group or the referenced calibration group does not exist.
587        """
588        calibration_path = concatenate_paths(self._path, brim_obj_names.data.calibration)
589        if not await self._file.object_exists(calibration_path):
590            raise ValueError(f"No calibration group found in {self._path}")
591        same_as = None
592        try:
593            same_as = await self._file.get_attr(calibration_path, 'Same_as')
594        except Exception:
595            pass #  same_as attribute is optional, if it does not exist we just ignore it
596        # if the 'Same_as' attribute exists, find the calibration group with the corresponding index
597        if same_as is not None:
598            try:
599                d_m = await Data.from_existing_async(self._file, same_as)
600                return await d_m.get_calibration_async()
601            except IndexError:
602                raise ValueError(f"Calibration group in {self._path} references non-existing calibration index {same_as} in the file")
603        cal_group = Calibration(self._file, calibration_path, data_group=self, _initialize=False)
604        await cal_group._init_async()
605        return cal_group
606
607    def create_analysis_results_group(self, data_AntiStokes, data_Stokes=None, *,
608                                          index: int = None, name: str = None, fit_model: 'Data.AnalysisResults.FitModel' = None) -> AnalysisResults:
609        """
610        Adds a new AnalysisResults entry to the current data group.
611        Parameters:
612            data_AntiStokes (dict or list[dict]): see documentation for `brimfile.analysis_results.AnalysisResults.add_data`
613            data_Stokes (dict or list[dict]): same as data_AntiStokes for the Stokes peaks.
614            index (int, optional): The index for the new data entry. If None, the next available index is used. Defaults to None.
615            name (str, optional): The name for the new Analysis group. Defaults to None.
616            fit_model (Data.AnalysisResults.FitModel, optional): The fit model used for the analysis. Defaults to None (no attribute is set).
617        Returns:
618            AnalysisResults: The newly created AnalysisResults object.
619        Raises:
620            IndexError: If the specified index already exists in the dataset.
621            ValueError: If any of the data provided is not valid or consistent
622        """
623        if index is not None:
624            try:
625                self.get_analysis_results(index)
626            except IndexError:
627                pass
628            else:
629                # If the group already exists, raise an error
630                raise IndexError(
631                    f"Analysis {index} already exists in {self._path}")
632        else:
633            ar_groups = self.list_AnalysisResults()
634            indices = [ar['index'] for ar in ar_groups]
635            indices.sort()
636            index = indices[-1] + 1 if indices else 0  # Next available index
637
638        ar = Data.AnalysisResults._create_new(self, index=index, sparse=self._sparse)
639        if name is not None:
640            set_object_name(self._file, ar._path, name)
641        ar.add_data(data_AntiStokes, data_Stokes, fit_model=fit_model)
642
643        return ar
644
645    def list_AnalysisResults(self, retrieve_custom_name=False) -> list:
646        """
647        List all AnalysisResults groups in the current data group. The list is ordered by index.
648
649        Returns:
650            list: A list of dictionaries, each containing:
651                - 'name' (str): The name of the AnalysisResults group.
652                - 'index' (int): The index extracted from the group name.
653                - 'custom_name' (str, optional): if retrieve_custom_name==True, it contains the name of the AnalysisResults group as returned from utils.get_object_name.
654        """
655
656        analysis_results_groups = []
657
658        matched_objs = sync(list_objects_matching_pattern_async(
659            self._file, self._group, brim_obj_names.data.analysis_results + r"_(\d+)$"))
660        async def _make_dict_item(matched_obj, retrieve_custom_name):
661            name = matched_obj[0]
662            index = int(matched_obj[1])
663            curr_obj_dict = {'name': name, 'index': index}
664            if retrieve_custom_name:
665                ar_path = concatenate_paths(self._path, name)
666                custom_name = await get_object_name(self._file, ar_path)
667                curr_obj_dict['custom_name'] = custom_name
668            return curr_obj_dict
669        coros = [_make_dict_item(matched_obj, retrieve_custom_name) for matched_obj in matched_objs]
670        dicts = _gather_sync(*coros)
671        for dict_item in dicts:
672            analysis_results_groups.append(dict_item)
673        # Sort the data groups by index
674        analysis_results_groups.sort(key=lambda x: x['index'])
675
676        return analysis_results_groups
677
678    def get_analysis_results(self, index: int = 0) -> AnalysisResults:
679        """
680        Returns the AnalysisResults at the specified index
681
682        Args:
683            index (int)                
684
685        Raises:
686            IndexError: If there is no analysis with the corresponding index
687        """
688        name = None
689        ls = self.list_AnalysisResults()
690        for el in ls:
691            if el['index'] == index:
692                name = el['name']
693                break
694        if name is None:
695            raise IndexError(f"Analysis {index} not found")
696        path = concatenate_paths(self._path, name)
697        return Data.AnalysisResults(self._file, path, data_group_path=self._path,
698                                    spatial_map=self._spatial_map, spatial_map_px_size=self._spatial_map_px_size, sparse=self._sparse)
699
700    def _add_data(self, PSD: np.ndarray, frequency: np.ndarray, *, scanning: dict = None, freq_units='GHz',
701                  timestamp: np.ndarray = None, compression: FileAbstraction.Compression = FileAbstraction.Compression()):
702        """
703        Add data to the current data group.
704
705        This method adds the provided PSD, frequency, and scanning data to the HDF5 group 
706        associated with this `Data` object. It validates the inputs to ensure they meet 
707        the required specifications before adding them.
708
709        Args:
710            PSD (np.ndarray): A 2D numpy array representing the Power Spectral Density (PSD) data. The last dimension contains the spectra.
711            frequency (np.ndarray): A 1D or 2D numpy array representing the frequency data. 
712                It must be broadcastable to the shape of the PSD array.
713            scanning (dict, optional): A dictionary containing scanning-related data. 
714                Required for sparse data (sparse=True), optional for non-sparse data.
715                For sparse data, must include at least one of 'Spatial_map' or 'Cartesian_visualisation'.
716                It may include the following keys:
717                - 'Spatial_map' (optional): A dictionary containing coordinate arrays:
718                    - 'x', 'y', 'z' (optional): 1D numpy arrays of same length with coordinate values
719                    - 'units' (optional): string with the unit (e.g., 'um')
720                - 'Cartesian_visualisation' (optional): A 3D numpy array (z, y, x) with integer values 
721                   mapping spatial positions to spectra indices. Values must be -1 (invalid/empty pixel) 
722                   or between 0 and PSD.shape[0]-1.
723                - 'Cartesian_visualisation_pixel' (recommended with Cartesian_visualisation): 
724                   Tuple/list of 3 float values (z, y, x) representing pixel size. Unused dimensions can be None.
725                - 'Cartesian_visualisation_pixel_unit' (optional): String for pixel size unit (default: 'um').
726            timestamp (np.ndarray, optional): Timestamps in milliseconds for each spectrum.
727                Must be a 1D array with length equal to PSD.shape[0].
728
729
730        Raises:
731            ValueError: If any of the data provided is not valid or consistent
732        """
733
734        # Check if frequency is broadcastable to PSD
735        try:
736            np.broadcast_shapes(tuple(frequency.shape), tuple(PSD.shape))
737        except ValueError as e:
738            raise ValueError(f"frequency (shape: {frequency.shape}) is not broadcastable to PSD (shape: {PSD.shape}): {e}")
739
740        # Check if at least one of 'Spatial_map' or 'Cartesian_visualisation' is present in the scanning dictionary
741        # This is required for sparse data to establish the spatial mapping
742        has_spatial_mapping = False
743        if scanning is not None:
744            if 'Spatial_map' in scanning:
745                sm = scanning['Spatial_map']
746                size = 0
747
748                def check_coor(coor: str):
749                    if coor in sm:
750                        sm[coor] = np.array(sm[coor])
751                        size1 = sm[coor].size
752                        if size1 != size and size != 0:
753                            raise ValueError(
754                                f"'{coor}' in 'Spatial_map' is invalid!")
755                        return size1
756                    return size
757                size = check_coor('x')
758                size = check_coor('y')
759                size = check_coor('z')
760                if size == 0:
761                    raise ValueError(
762                        "'Spatial_map' should contain at least one x, y or z")
763                has_spatial_mapping = True
764            if 'Cartesian_visualisation' in scanning:
765                cv = scanning['Cartesian_visualisation']
766                if not isinstance(cv, np.ndarray) or cv.ndim != 3:
767                    raise ValueError(
768                        "Cartesian_visualisation must be a 3D numpy array")
769                if not np.issubdtype(cv.dtype, np.integer) or np.min(cv) < -1 or np.max(cv) >= PSD.shape[0]:
770                    raise ValueError(
771                        "Cartesian_visualisation values must be integers between -1 and PSD.shape[0]-1")
772                if 'Cartesian_visualisation_pixel' in scanning:
773                    if len(scanning['Cartesian_visualisation_pixel']) != 3:
774                        raise ValueError(
775                            "Cartesian_visualisation_pixel must always contain 3 values for z, y, x (set to None if not used)")
776                else:
777                    warnings.warn(
778                        "It is recommended to include 'Cartesian_visualisation_pixel' in the scanning dictionary to define pixel size for proper spatial calibration")
779                has_spatial_mapping = True
780        if not has_spatial_mapping and self._sparse:
781            raise ValueError("For sparse data, 'scanning' must be provided and must contain at least one of 'Spatial_map' or 'Cartesian_visualisation'")
782
783        if timestamp is not None:
784            if not isinstance(timestamp, np.ndarray) or timestamp.ndim != 1 or len(timestamp) != PSD.shape[0]:
785                raise ValueError("timestamp is not compatible with PSD")
786
787        # TODO: add and validate additional datasets (i.e. 'Parameters', 'Calibration_index', etc.)
788
789        # Add datasets to the group
790        sync(self._file.create_dataset(
791            self._group, brim_obj_names.data.PSD, data=PSD,
792            chunk_size=_determine_chunk_size(PSD), compression=compression))
793        freq_ds = sync(self._file.create_dataset(
794            self._group,  brim_obj_names.data.frequency, data=frequency,
795            chunk_size=_determine_chunk_size(frequency), compression=compression))
796        units.add_to_object(self._file, freq_ds, freq_units)
797
798        if scanning is not None:
799            if 'Spatial_map' in scanning:
800                sm = scanning['Spatial_map']
801                sm_group = sync(self._file.create_group(concatenate_paths(
802                    self._path, brim_obj_names.data.spatial_map)))
803                if 'units' in sm:
804                    units.add_to_object(self._file, sm_group, sm['units'])
805
806                def add_sm_dataset(coord: str):
807                    if coord in sm:
808                        sync(self._file.create_dataset(
809                            sm_group, coord, data=sm[coord], compression=compression))
810
811                add_sm_dataset('x')
812                add_sm_dataset('y')
813                add_sm_dataset('z')
814            if 'Cartesian_visualisation' in scanning:
815                # convert the Cartesian_visualisation to the smallest integer type
816                cv_arr = np_array_to_smallest_int_type(scanning['Cartesian_visualisation'])
817                cv = sync(self._file.create_dataset(self._group, brim_obj_names.data.cartesian_visualisation,
818                                            data=cv_arr, compression=compression))
819                if 'Cartesian_visualisation_pixel' in scanning:
820                    sync(self._file.create_attr(
821                        cv, 'element_size', scanning['Cartesian_visualisation_pixel']))
822                    if 'Cartesian_visualisation_pixel_unit' in scanning:
823                        px_unit = scanning['Cartesian_visualisation_pixel_unit']
824                    else:
825                        warnings.warn(
826                            "No unit provided for Cartesian_visualisation_pixel, defaulting to 'um'")
827                        px_unit = 'um'
828                    units.add_to_attribute(self._file, cv, 'element_size', px_unit)
829
830        self._spatial_map, self._spatial_map_px_size = sync(self._load_spatial_mapping_async())
831
832        if timestamp is not None:
833            sync(self._file.create_dataset(
834                self._group, 'Timestamp', data=timestamp, compression=compression))
835
836    @staticmethod
837    def list_data_groups(file: FileAbstraction, retrieve_custom_name=False) -> list:
838        """
839        Synchronous wrapper for `list_data_groups_async` (see doc for `brimfile.data.Data.list_data_groups_async`)
840        """
841        return sync(Data.list_data_groups_async(file, retrieve_custom_name))
842
843    @staticmethod
844    async def list_data_groups_async(file: FileAbstraction, retrieve_custom_name=False) -> list:
845        """
846        List all data groups in the brim file. The list is ordered by index.
847
848        Returns:
849            list: A list of dictionaries, each containing:
850                - 'name' (str): The name of the data group in the file.
851                - 'index' (int): The index extracted from the group name.
852                - 'custom_name' (str, optional): if retrieve_custom_name==True, it contains the name of the data group as returned from utils.get_object_name.
853        """
854
855        data_groups = []
856
857        matched_objs = await list_objects_matching_pattern_async(
858            file, brim_obj_names.Brillouin_base_path, brim_obj_names.data.base_group + r"_(\d+)$")
859        
860        async def _make_dict_item(matched_obj, retrieve_custom_name):
861            name = matched_obj[0]
862            index = int(matched_obj[1])
863            curr_obj_dict = {'name': name, 'index': index}
864            if retrieve_custom_name:
865                path = concatenate_paths(
866                    brim_obj_names.Brillouin_base_path, name)
867                custom_name = await get_object_name(file, path)
868                curr_obj_dict['custom_name'] = custom_name
869            return curr_obj_dict
870        
871        coros = [_make_dict_item(matched_obj, retrieve_custom_name) for matched_obj in matched_objs]
872        dicts = await asyncio.gather(*coros)
873        for dict_item in dicts:
874            data_groups.append(dict_item)        
875        # Sort the data groups by index
876        data_groups.sort(key=lambda x: x['index'])
877
878        return data_groups
879
880    @staticmethod
881    async def _get_existing_group_name_async(file: FileAbstraction, index: int) -> str:
882        """
883        Get the name of an existing data group by index.
884
885        Args:
886            file (File): The parent File object.
887            index (int): The index of the data group.
888
889        Returns:
890            str: The name of the data group, or None if not found.
891        """
892        group_name: str = None
893        data_groups = await Data.list_data_groups_async(file)
894        for dg in data_groups:
895            if dg['index'] == index:
896                group_name = dg['name']
897                break
898        return group_name
899    
900    @classmethod
901    async def from_existing_async(cls, file: FileAbstraction, index: int) -> 'Data':
902        """ 
903        Create a Data object from an existing data group in the file.
904        Args:
905            file (File): The parent File object.
906            index (int): The index of the existing data group.      
907        Returns:
908            Data: A Data object corresponding to the existing data group.   
909        Raises:
910            IndexError: If no data group with the specified index is found in the file.
911        """
912        group_name: str = await cls._get_existing_group_name_async(file, index)
913        if group_name is None:
914            raise IndexError(f"No data group with index {index} found in the file")
915        dg = cls(file, concatenate_paths(brim_obj_names.Brillouin_base_path, group_name), _initialize=False) 
916        await dg._init_async()
917        return dg
918    
919    @classmethod
920    def _create_new(cls, file: FileAbstraction, index: int, sparse: bool = False, name: str = None) -> 'Data':
921        """
922        Create a new data group with the specified index.
923
924        Args:
925            file (File): The parent File object.
926            index (int): The index for the new data group.
927            sparse (bool): Whether the data is sparse. See https://github.com/brillouin-imaging/Brillouin-standard-file/blob/main/docs/brim_file_specs.md for details. Defaults to False.
928            name (str, optional): The name for the new data group. Defaults to None.
929
930        Returns:
931            Data: The newly created Data object.
932        """
933        group_name = Data._generate_group_name(index)
934        group = sync(file.create_group(concatenate_paths(
935            brim_obj_names.Brillouin_base_path, group_name)))
936        sync(file.create_attr(group, 'Sparse', sparse))
937        if name is not None:
938            set_object_name(file, group, name)
939        return cls(file, concatenate_paths(brim_obj_names.Brillouin_base_path, group_name), newly_created=True)
940
941    @staticmethod
942    def _generate_group_name(index: int, n_digits: int = None) -> str:
943        """
944        Generate a name for a data group based on the index.
945
946        Args:
947            index (int): The index for the data group.
948            n_digits (int, optional): The number of digits to pad the index with. If None no padding is applied. Defaults to None.
949
950        Returns:
951            str: The generated group name.
952
953        Raises:
954            ValueError: If the index is negative.
955        """
956        if index < 0:
957            raise ValueError("index must be positive")
958        num = str(index)
959        if n_digits is not None:
960            num = num.zfill(n_digits)
961        return f"{brim_obj_names.data.base_group}_{num}"