Coverage for install/scipp/utils/collapse_and_slices.py: 89%

1# SPDX-License-Identifier: BSD-3-Clause 

2# Copyright (c) 2023 Scipp contributors (https://github.com/scipp) 

3# @author Neil Vaytet 

4 

5# Other imports 

6import numpy as np 

7 

8 

9def _to_slices(scipp_obj, slice_dims, slice_shape, volume): 

10 # Create container to collect all 1D slices as 1D variables 

11 all_slices = {} 

12 

13 # Go through the dims that need to be collapsed, and create an array that 

14 # holds the range of indices for each dimension 

15 # Say we have [Y, 5], and [Z, 3], then dim_list will contain 

16 # [[0, 1, 2, 3, 4], [0, 1, 2]] 

17 dim_list = [np.arange(slice_shape[dim], dtype=np.int32) for dim in slice_dims] 

18 # Next create a grid of indices 

19 # grid will contain 

20 # [ [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], 

21 # [[0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [2, 2, 2, 2, 2]] ] 

22 grid = np.meshgrid(*list(dim_list)) 

23 # Reshape the grid to have a 2D array of length volume, i.e. 

24 # [ [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4], 

25 # [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2] ] 

26 res = np.reshape(grid, (len(slice_dims), volume)) 

27 # Now make a master array which also includes the dimension labels, i.e. 

28 # [ [Y, Y, Y, Y, Y, Y, Y, Y, Y, Y, Y, Y, Y, Y, Y], 

29 # [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4], 

30 # [Z, Z, Z, Z, Z, Z, Z, Z, Z, Z, Z, Z, Z, Z, Z], 

31 # [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2] ] 

32 slice_list = [] 

33 for i, dim in enumerate(slice_dims): 

34 slice_list.append([dim] * volume) 

35 slice_list.append(res[i]) 

36 # Finally reshape the master array to look like 

37 # [ [[Y, 0], [Z, 0]], [[Y, 1], [Z, 0]], 

38 # [[Y, 2], [Z, 0]], [[Y, 3], [Z, 0]], 

39 # [[Y, 4], [Z, 0]], [[Y, 0], [Z, 1]], 

40 # [[Y, 1], [Z, 1]], [[Y, 2], [Z, 1]], 

41 # [[Y, 3], [Z, 1]], 

42 # ... 

43 # ] 

44 slice_list = np.reshape( 

45 np.transpose(np.array(slice_list, dtype=np.dtype('O'))), 

46 (volume, len(slice_dims), 2), 

47 ) 

48 

49 # Extract each entry from the slice_list 

50 for line in slice_list: 

51 vslice = scipp_obj 

52 key = "" 

53 for s in line: 

54 vslice = vslice[s[0], s[1]] 

55 key += f"{s[0]}:{s[1]}-" 

56 all_slices[key[:-1]] = vslice 

57 

58 return all_slices 

59 

60 

61def collapse(scipp_obj, keep): 

62 """ 

63 Slice down the input object until only the supplied `keep` dimension is 

64 left (effectively 'collapsing' all but one dimension), and return a 

65 `dict` of 1D slices. A common use for this is plotting spectra from 

66 detectors where most pixels contain noise, but one specific channel 

67 contains a strong signal. The `plot` function accepts a `dict` of data 

68 arrays. 

69 

70 :param [scipp_obj]: Dataset or DataArray to be split into slices. 

71 :type [scipp_obj]: Dataset or DataArray 

72 :param [keep]: Dimension to be preserved. 

73 :type [dim]: str 

74 :return: A dictionary holding 1D slices of the input object. 

75 :rtype: dict 

76 """ 

77 

78 dims = scipp_obj.dims 

79 shape = scipp_obj.shape 

80 

81 # Gather list of dimensions that are to be collapsed 

82 slice_dims = [] 

83 volume = 1 

84 slice_shape = {} 

85 for d, size in zip(dims, shape, strict=True): 

86 if d != keep: 

87 slice_dims.append(d) 

88 slice_shape[d] = size 

89 volume *= size 

90 

91 return _to_slices(scipp_obj, slice_dims, slice_shape, volume) 

92 

93 

94def slices(scipp_obj, dim): 

95 """ 

96 Slice input along given dim, and return all the slices in a `dict`. 

97 

98 :param [scipp_obj]: Dataset or DataArray to be split into slices. 

99 :type [scipp_obj]: Dataset or DataArray 

100 :param [dim]: Dimension along which to slice. 

101 :type [dim]: str 

102 :return: A dictionary holding slices of the input object. 

103 :rtype: dict 

104 """ 

105 

106 slice_dims = [dim] 

107 volume = scipp_obj.shape[scipp_obj.dims.index(dim)] 

108 slice_shape = {dim: volume} 

109 

110 return _to_slices(scipp_obj, slice_dims, slice_shape, volume)