Coverage for install/scipp/core/operations.py: 64%

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

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

3# @author Matthew Andrew 

4from __future__ import annotations 

5 

6from typing import Any, Dict, Literal, Optional, Union, overload 

7 

8from .._scipp import core as _cpp 

9from ..typing import ScippIndex, VariableLikeType 

10from ._cpp_wrapper_util import call_func as _call_cpp_func 

11from .comparison import identical 

12from .cpp_classes import Dataset, DatasetError, Variable 

13from .data_group import DataGroup 

14from .unary import to_unit 

15 

16 

17def islinspace(x: Variable, dim: Optional[str] = None) -> Variable: 

18 """Check if the values of a variable are evenly spaced. 

19 

20 Parameters 

21 ---------- 

22 x: 

23 Variable to check. 

24 dim: 

25 Optional variable for the dim to check from the Variable. 

26 

27 Returns 

28 ------- 

29 : 

30 Variable of value True if the variable contains regularly 

31 spaced values, variable of value False otherwise. 

32 """ 

33 if dim is None: 

34 return _call_cpp_func(_cpp.islinspace, x) 

35 else: 

36 return _call_cpp_func(_cpp.islinspace, x, dim) 

37 

38 

39def issorted( 

40 x: Variable, dim: str, order: Literal['ascending', 'descending'] = 'ascending' 

41) -> Variable: 

42 """Check if the values of a variable are sorted. 

43 

44 - If ``order`` is 'ascending', 

45 check if values are non-decreasing along ``dim``. 

46 - If ``order`` is 'descending', 

47 check if values are non-increasing along ``dim``. 

48 

49 Parameters 

50 ---------- 

51 x: 

52 Variable to check. 

53 dim: 

54 Dimension along which order is checked. 

55 order: 

56 Sorting order. 

57 

58 Returns 

59 ------- 

60 : 

61 Variable containing one less dim than the original 

62 variable with the corresponding boolean value for whether 

63 it was sorted along the given dim for the other dimensions. 

64 

65 See Also 

66 -------- 

67 scipp.allsorted 

68 """ 

69 return _call_cpp_func(_cpp.issorted, x, dim, order) 

70 

71 

72def allsorted( 

73 x: Variable, dim: str, order: Literal['ascending', 'descending'] = 'ascending' 

74) -> bool: 

75 """Check if all values of a variable are sorted. 

76 

77 - If ``order`` is 'ascending', 

78 check if values are non-decreasing along ``dim``. 

79 - If ``order`` is 'descending', 

80 check if values are non-increasing along ``dim``. 

81 

82 Parameters 

83 ---------- 

84 x: 

85 Variable to check. 

86 dim: 

87 Dimension along which order is checked. 

88 order: 

89 Sorting order. 

90 

91 Returns 

92 ------- 

93 : 

94 True if the variable values are monotonously ascending or 

95 descending (depending on the requested order), False otherwise. 

96 

97 See Also 

98 -------- 

99 scipp.issorted 

100 """ 

101 return _call_cpp_func(_cpp.allsorted, x, dim, order) 

102 

103 

104def sort( 

105 x: VariableLikeType, 

106 key: Union[str, Variable], 

107 order: Literal['ascending', 'descending'] = 'ascending', 

108) -> VariableLikeType: 

109 """Sort variable along a dimension by a sort key or dimension label. 

110 

111 - If ``order`` is 'ascending', 

112 sort such that values are non-decreasing according to ``key``. 

113 - If ``order`` is 'descending', 

114 sort such that values are non-increasing according to ``key``. 

115 

116 Parameters 

117 ---------- 

118 x: scipp.typing.VariableLike 

119 Data to be sorted. 

120 key: 

121 Either a 1D variable sort key or a dimension label. 

122 order: 

123 Sorting order. 

124 

125 Returns 

126 ------- 

127 : scipp.typing.VariableLike 

128 The sorted equivalent of the input with the same type. 

129 

130 Raises 

131 ------ 

132 scipp.DimensionError 

133 If the key is a Variable that does not have exactly 1 dimension. 

134 """ 

135 return _call_cpp_func(_cpp.sort, x, key, order) 

136 

137 

138def values(x: VariableLikeType) -> VariableLikeType: 

139 """Return the input without variances. 

140 

141 Parameters 

142 ---------- 

143 x: scipp.typing.VariableLike 

144 Input data. 

145 

146 Returns 

147 ------- 

148 : scipp.typing.VariableLike 

149 The same as the input but without variances. 

150 

151 See Also 

152 -------- 

153 scipp.variances, scipp.stddevs 

154 """ 

155 return _call_cpp_func(_cpp.values, x) 

156 

157 

158def variances(x: VariableLikeType) -> VariableLikeType: 

159 """Return the input's variances as values. 

160 

161 Parameters 

162 ---------- 

163 x: scipp.typing.VariableLike 

164 Input data with variances. 

165 

166 Returns 

167 ------- 

168 : scipp.typing.VariableLike 

169 The same as the input but with values set to the input's variances 

170 and without variances itself. 

171 

172 See Also 

173 -------- 

174 scipp.values, scipp.stddevs 

175 """ 

176 return _call_cpp_func(_cpp.variances, x) 

177 

178 

179def stddevs(x: VariableLikeType) -> VariableLikeType: 

180 """Return the input's standard deviations as values. 

181 

182 Parameters 

183 ---------- 

184 x: scipp.typing.VariableLike 

185 Input data with variances. 

186 

187 Returns 

188 ------- 

189 : scipp.typing.VariableLike 

190 The same as the input but with values set to standard deviations computed 

191 from the input's variances and without variances itself. 

192 

193 See Also 

194 -------- 

195 scipp.values, scipp.variances 

196 """ 

197 return _call_cpp_func(_cpp.stddevs, x) 

198 

199 

200def where(condition: Variable, x: Variable, y: Variable) -> Variable: 

201 """Return elements chosen from x or y depending on condition. 

202 

203 Parameters 

204 ---------- 

205 condition: 

206 Variable with dtype=bool. 

207 x: 

208 Variable with values from which to choose. 

209 y: 

210 Variable with values from which to choose. 

211 

212 Returns 

213 ------- 

214 : 

215 Variable with elements from x where condition is True 

216 and elements from y elsewhere. 

217 """ 

218 return _call_cpp_func(_cpp.where, condition, x, y) 

219 

220 

221def to( 

222 var: VariableLikeType, 

223 *, 

224 unit: Optional[Union[_cpp.Unit, str]] = None, 

225 dtype: Optional[Any] = None, 

226 copy: bool = True, 

227) -> VariableLikeType: 

228 """Converts a Variable or DataArray to a different dtype and/or a different unit. 

229 

230 If the dtype and unit are both unchanged and ``copy`` is `False`, 

231 the object is returned without making a deep copy. 

232 

233 This method will choose whether to do the dtype or units translation first, by 

234 using the following rules in order: 

235 

236 - If either the input or output dtype is float64, the unit translation will be done 

237 on the float64 type 

238 - If either the input or output dtype is float32, the unit translation will be done 

239 on the float32 type 

240 - If both the input and output dtypes are integer types, the unit translation will 

241 be done on the larger type 

242 - In other cases, the dtype is converted first and then the unit translation is done 

243 

244 Parameters 

245 ---------- 

246 unit: 

247 Target unit. If ``None``, the unit is unchanged. 

248 dtype: 

249 Target dtype. If ``None``, the dtype is unchanged. 

250 copy: 

251 If ``False``, return the input object if possible. 

252 If ``True``, the function always returns a new object. 

253 

254 Returns 

255 ------- 

256 : Same as input 

257 New object with specified dtype and unit. 

258 

259 Raises 

260 ------ 

261 scipp.DTypeError 

262 If the input cannot be converted to the given dtype. 

263 scipp.UnitError 

264 If the input cannot be converted to the given unit. 

265 

266 See Also 

267 -------- 

268 scipp.to_unit, scipp.DataArray.astype, scipp.Variable.astype 

269 """ 

270 if unit is None and dtype is None: 

271 raise ValueError("Must provide dtype or unit or both") 

272 

273 if dtype is None: 

274 return to_unit(var, unit, copy=copy) 

275 

276 if unit is None: 

277 return var.astype(dtype, copy=copy) 

278 

279 if dtype == _cpp.DType.float64: 

280 convert_dtype_first = True 

281 elif var.dtype == _cpp.DType.float64: 

282 convert_dtype_first = False 

283 elif dtype == _cpp.DType.float32: 

284 convert_dtype_first = True 

285 elif var.dtype == _cpp.DType.float32: 

286 convert_dtype_first = False 

287 elif var.dtype == _cpp.DType.int64 and dtype == _cpp.DType.int32: 

288 convert_dtype_first = False 

289 elif var.dtype == _cpp.DType.int32 and dtype == _cpp.DType.int64: 

290 convert_dtype_first = True 

291 else: 

292 convert_dtype_first = True 

293 

294 if convert_dtype_first: 

295 return to_unit(var.astype(dtype, copy=copy), unit=unit, copy=False) 

296 else: 

297 return to_unit(var, unit=unit, copy=copy).astype(dtype, copy=False) 

298 

299 

300@overload 

301def merge(lhs: Dataset, rhs: Dataset) -> Dataset: ... 

302 

303 

304@overload 

305def merge(lhs: DataGroup, rhs: DataGroup) -> DataGroup: ... 

306 

307 

308def merge(lhs, rhs): 

309 """Merge two datasets or data groups into one. 

310 

311 If an item appears in both inputs, it must have an identical value in both. 

312 

313 Parameters 

314 ---------- 

315 lhs: 

316 First dataset or data group. 

317 rhs: 

318 Second dataset or data group. 

319 

320 Returns 

321 ------- 

322 : 

323 A new object that contains the union of all data items, 

324 coords, masks and attributes. 

325 

326 Raises 

327 ------ 

328 scipp.DatasetError 

329 If there are conflicting items with different content. 

330 """ 

331 if isinstance(lhs, Dataset) or isinstance(rhs, Dataset): 

332 return _call_cpp_func(_cpp.merge, lhs, rhs) 

333 return _merge_data_group(lhs, rhs) 

334 

335 

336def _generic_identical(a: Any, b: Any) -> bool: 

337 try: 

338 return identical(a, b) 

339 except TypeError: 

340 from numpy import array_equal 

341 

342 try: 

343 return array_equal(a, b) 

344 except TypeError: 

345 return a == b 

346 

347 

348def _merge_data_group(lhs: DataGroup, rhs: DataGroup) -> DataGroup: 

349 res = DataGroup(dict(lhs)) 

350 for k, v in rhs.items(): 

351 if k in res and not _generic_identical(res[k], v): 

352 raise DatasetError(f"Cannot merge data groups. Mismatch in item {k}") 

353 res[k] = v 

354 return res 

355 

356 

357def label_based_index_to_positional_index( 

358 sizes: Dict[str, int], 

359 coord: Variable, 

360 index: Union[slice[Optional[Variable]], Variable], 

361) -> ScippIndex: 

362 """Returns the positional index equivalent to label based indexing 

363 the coord with values.""" 

364 dim, *inds = _call_cpp_func( 

365 _cpp.label_based_index_to_positional_index, 

366 list(sizes.keys()), 

367 list(sizes.values()), 

368 coord, 

369 index, 

370 ) 

371 return (dim, inds[0] if len(inds) == 1 else slice(*inds))