Coverage for install/scipp/core/dimensions.py: 69%

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

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

3# @author Simon Heybrock 

4 

5from typing import TypeVar 

6 

7from .._scipp.core import CoordError, DataArray, Dataset, Variable 

8from .argument_handlers import combine_dict_args 

9from .bins import bins 

10from .variable import scalar 

11 

12_T = TypeVar('_T', Variable, DataArray, Dataset) 

13 

14 

15def _rename_dims( 

16 self: _T, dims_dict: dict[str, str] | None = None, /, **names: str 

17) -> _T: 

18 """Rename dimensions. 

19 

20 The renaming can be defined: 

21 

22 - using a dict mapping the old to new names, e.g. 

23 ``rename_dims({'x': 'a', 'y': 'b'})`` 

24 - using keyword arguments, e.g. ``rename_dims(x='a', y='b')`` 

25 

26 In both cases, x is renamed to a and y to b. 

27 

28 Dimensions not specified in either input are unchanged. 

29 

30 This function only renames dimensions. 

31 See the ``rename`` method to also rename coordinates and attributes. 

32 

33 Parameters 

34 ---------- 

35 dims_dict: 

36 Dictionary mapping old to new names. 

37 names: 

38 Mapping of old to new names as keyword arguments. 

39 

40 Returns 

41 ------- 

42 : 

43 A new object with renamed dimensions. 

44 """ 

45 return self._rename_dims(combine_dict_args(dims_dict, names)) 

46 

47 

48def _rename_variable( 

49 var: Variable, dims_dict: dict[str, str] | None = None, /, **names: str 

50) -> Variable: 

51 """Rename dimension labels. 

52 

53 The renaming can be defined: 

54 

55 - using a dict mapping the old to new names, e.g. ``rename({'x': 'a', 'y': 'b'})`` 

56 - using keyword arguments, e.g. ``rename(x='a', y='b')`` 

57 

58 In both cases, x is renamed to a and y to b. 

59 

60 Dimensions not specified in either input are unchanged. 

61 

62 Parameters 

63 ---------- 

64 dims_dict: 

65 Dictionary mapping old to new names. 

66 names: 

67 Mapping of old to new names as keyword arguments. 

68 

69 Returns 

70 ------- 

71 : 

72 A new variable with renamed dimensions which shares a buffer with the input. 

73 

74 See Also 

75 -------- 

76 scipp.Variable.rename_dims: 

77 Equivalent for ``Variable`` but differs for ``DataArray`` and ``Dataset``. 

78 """ 

79 return var.rename_dims(combine_dict_args(dims_dict, names)) 

80 

81 

82def _rename_data_array( 

83 da: DataArray, dims_dict: dict[str, str] | None = None, /, **names: str 

84) -> DataArray: 

85 """Rename the dimensions, coordinates, and attributes. 

86 

87 The renaming can be defined: 

88 

89 - using a dict mapping the old to new names, e.g. ``rename({'x': 'a', 'y': 'b'})`` 

90 - using keyword arguments, e.g. ``rename(x='a', y='b')`` 

91 

92 In both cases, x is renamed to a and y to b. 

93 

94 Names not specified in either input are unchanged. 

95 

96 Parameters 

97 ---------- 

98 dims_dict: 

99 Dictionary mapping old to new names. 

100 names: 

101 Mapping of old to new names as keyword arguments. 

102 

103 Returns 

104 ------- 

105 : 

106 A new data array with renamed dimensions, coordinates, and attributes. 

107 Buffers are shared with the input. 

108 

109 See Also 

110 -------- 

111 scipp.DataArray.rename_dims: 

112 Only rename dimensions, not coordinates and attributes. 

113 """ 

114 renaming_dict = combine_dict_args(dims_dict, names) 

115 out = da.rename_dims(renaming_dict) 

116 if out.bins is not None: 

117 out.data = bins(**out.bins.constituents) 

118 for old, new in renaming_dict.items(): 

119 if new in out.deprecated_meta: 

120 raise CoordError( 

121 f"Cannot rename '{old}' to '{new}', since a coord or attr named {new} " 

122 "already exists." 

123 ) 

124 for meta in (out.coords, out.deprecated_attrs): 

125 if old in meta: 

126 meta[new] = meta.pop(old) 

127 if out.bins is not None: 

128 for meta in (out.bins.coords, out.bins.deprecated_attrs): 

129 if old in meta: 

130 meta[new] = meta.pop(old) 

131 return out 

132 

133 

134def _rename_dataset( 

135 ds: Dataset, dims_dict: dict[str, str] | None = None, /, **names: str 

136) -> Dataset: 

137 """Rename the dimensions, coordinates and attributes of all the items. 

138 

139 The renaming can be defined: 

140 

141 - using a dict mapping the old to new names, e.g. ``rename({'x': 'a', 'y': 'b'})`` 

142 - using keyword arguments, e.g. ``rename(x='a', y='b')`` 

143 

144 In both cases, x is renamed to a and y to b. 

145 

146 Names not specified in either input are unchanged. 

147 

148 Parameters 

149 ---------- 

150 dims_dict: 

151 Dictionary mapping old to new names. 

152 names: 

153 Mapping of old to new names as keyword arguments. 

154 

155 Returns 

156 ------- 

157 : 

158 A new dataset with renamed dimensions, coordinates, and attributes. 

159 Buffers are shared with the input. 

160 

161 See Also 

162 -------- 

163 scipp.Dataset.rename_dims: 

164 Only rename dimensions, not coordinates and attributes. 

165 """ 

166 dims_dict = combine_dict_args(dims_dict, names) 

167 if len(ds) != 0: 

168 return Dataset( 

169 {key: _rename_data_array(value, dims_dict) for key, value in ds.items()} 

170 ) 

171 # This relies on broadcast and DataArray.__init__ not making copies 

172 # to avoid allocating too much extra memory. 

173 dummy = DataArray( 

174 scalar(0).broadcast(dims=ds.dims, shape=ds.shape), coords=ds.coords 

175 ) 

176 return Dataset(coords=_rename_data_array(dummy, dims_dict).coords)