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

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, Literal, TypeVar, 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 DataArray, Dataset, DatasetError, Variable 

13from .data_group import DataGroup 

14from .unary import to_unit 

15 

16_T = TypeVar('_T', Variable, DataGroup[object]) 

17_VarDa = TypeVar('_VarDa', Variable, DataArray) 

18_DsDg = TypeVar('_DsDg', Dataset, DataGroup[object]) 

19_E1 = TypeVar('_E1') 

20_E2 = TypeVar('_E2') 

21 

22 

23def islinspace(x: _T, dim: str | None = None) -> _T: 

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

25 

26 Parameters 

27 ---------- 

28 x: 

29 Variable to check. 

30 dim: 

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

32 

33 Returns 

34 ------- 

35 : 

36 Variable of value True if the variable contains regularly 

37 spaced values, variable of value False otherwise. 

38 """ 

39 if dim is None: 

40 return _call_cpp_func(_cpp.islinspace, x) # type: ignore[return-value] 

41 else: 

42 return _call_cpp_func(_cpp.islinspace, x, dim) # type: ignore[return-value] 

43 

44 

45def issorted( 

46 x: _T, dim: str, order: Literal['ascending', 'descending'] = 'ascending' 

47) -> _T: 

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

49 

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

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

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

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

54 

55 Parameters 

56 ---------- 

57 x: 

58 Variable to check. 

59 dim: 

60 Dimension along which order is checked. 

61 order: 

62 Sorting order. 

63 

64 Returns 

65 ------- 

66 : 

67 Variable containing one less dim than the original 

68 variable with the corresponding boolean value for whether 

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

70 

71 See Also 

72 -------- 

73 scipp.allsorted 

74 """ 

75 return _call_cpp_func(_cpp.issorted, x, dim, order) # type: ignore[return-value] 

76 

77 

78@overload 

79def allsorted( 

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

81) -> bool: ... 

82 

83 

84@overload 

85def allsorted( 

86 x: DataGroup[object], 

87 dim: str, 

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

89) -> DataGroup[object]: ... 

90 

91 

92def allsorted( 

93 x: Variable | DataGroup[object], 

94 dim: str, 

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

96) -> bool | DataGroup[object]: 

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

98 

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

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

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

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

103 

104 Parameters 

105 ---------- 

106 x: 

107 Variable to check. 

108 dim: 

109 Dimension along which order is checked. 

110 order: 

111 Sorting order. 

112 

113 Returns 

114 ------- 

115 : 

116 True if the variable values are monotonously ascending or 

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

118 

119 See Also 

120 -------- 

121 scipp.issorted 

122 """ 

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

124 

125 

126def sort( 

127 x: VariableLikeType, 

128 key: str | Variable, 

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

130) -> VariableLikeType: 

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

132 

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

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

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

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

137 

138 Parameters 

139 ---------- 

140 x: scipp.typing.VariableLike 

141 Data to be sorted. 

142 key: 

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

144 order: 

145 Sorting order. 

146 

147 Returns 

148 ------- 

149 : scipp.typing.VariableLike 

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

151 

152 Raises 

153 ------ 

154 scipp.DimensionError 

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

156 """ 

157 return _call_cpp_func(_cpp.sort, x, key, order) # type: ignore[return-value] 

158 

159 

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

161 """Return the input without variances. 

162 

163 Parameters 

164 ---------- 

165 x: scipp.typing.VariableLike 

166 Input data. 

167 

168 Returns 

169 ------- 

170 : scipp.typing.VariableLike 

171 The same as the input but without variances. 

172 

173 See Also 

174 -------- 

175 scipp.variances, scipp.stddevs 

176 """ 

177 return _call_cpp_func(_cpp.values, x) # type: ignore[return-value] 

178 

179 

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

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

182 

183 Parameters 

184 ---------- 

185 x: scipp.typing.VariableLike 

186 Input data with variances. 

187 

188 Returns 

189 ------- 

190 : scipp.typing.VariableLike 

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

192 and without variances itself. 

193 

194 See Also 

195 -------- 

196 scipp.values, scipp.stddevs 

197 """ 

198 return _call_cpp_func(_cpp.variances, x) # type: ignore[return-value] 

199 

200 

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

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

203 

204 Parameters 

205 ---------- 

206 x: scipp.typing.VariableLike 

207 Input data with variances. 

208 

209 Returns 

210 ------- 

211 : scipp.typing.VariableLike 

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

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

214 

215 See Also 

216 -------- 

217 scipp.values, scipp.variances 

218 """ 

219 return _call_cpp_func(_cpp.stddevs, x) # type: ignore[return-value] 

220 

221 

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

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

224 

225 Parameters 

226 ---------- 

227 condition: 

228 Variable with dtype=bool. 

229 x: 

230 Variable with values from which to choose. 

231 y: 

232 Variable with values from which to choose. 

233 

234 Returns 

235 ------- 

236 : 

237 Variable with elements from x where condition is True 

238 and elements from y elsewhere. 

239 """ 

240 return _call_cpp_func(_cpp.where, condition, x, y) # type: ignore[return-value] 

241 

242 

243def to( 

244 var: Variable, 

245 *, 

246 unit: _cpp.Unit | str | None = None, 

247 dtype: Any | None = None, 

248 copy: bool = True, 

249) -> Variable: 

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

251 

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

253 the object is returned without making a deep copy. 

254 

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

256 using the following rules in order: 

257 

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

259 on the float64 type 

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

261 on the float32 type 

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

263 be done on the larger type 

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

265 

266 Parameters 

267 ---------- 

268 unit: 

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

270 dtype: 

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

272 copy: 

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

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

275 

276 Returns 

277 ------- 

278 : Same as input 

279 New object with specified dtype and unit. 

280 

281 Raises 

282 ------ 

283 scipp.DTypeError 

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

285 scipp.UnitError 

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

287 

288 See Also 

289 -------- 

290 scipp.to_unit, scipp.DataArray.astype, scipp.Variable.astype 

291 """ 

292 if unit is None and dtype is None: 

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

294 

295 if dtype is None: 

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

297 

298 if unit is None: 

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

300 

301 if dtype == _cpp.DType.float64: 

302 convert_dtype_first = True 

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

304 convert_dtype_first = False 

305 elif dtype == _cpp.DType.float32: 

306 convert_dtype_first = True 

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

308 convert_dtype_first = False 

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

310 convert_dtype_first = False 

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

312 convert_dtype_first = True 

313 else: 

314 convert_dtype_first = True 

315 

316 if convert_dtype_first: 

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

318 else: 

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

320 

321 

322def merge(lhs: _DsDg, rhs: _DsDg) -> _DsDg: 

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

324 

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

326 

327 Parameters 

328 ---------- 

329 lhs: Dataset | DataGroup 

330 First dataset or data group. 

331 rhs: Dataset | DataGroup 

332 Second dataset or data group. 

333 

334 Returns 

335 ------- 

336 : Dataset | DataGroup 

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

338 coords, masks and attributes. 

339 

340 Raises 

341 ------ 

342 scipp.DatasetError 

343 If there are conflicting items with different content. 

344 """ 

345 # Check both arguments to make `_cpp.merge` raise TypeError on mismatch. 

346 if isinstance(lhs, Dataset) or isinstance(rhs, Dataset): # type: ignore[redundant-expr] 

347 return _call_cpp_func(_cpp.merge, lhs, rhs) # type: ignore[return-value] 

348 return _merge_data_group(lhs, rhs) 

349 

350 

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

352 try: 

353 return identical(a, b) 

354 except TypeError: 

355 from numpy import array_equal 

356 

357 try: 

358 return array_equal(a, b) 

359 except TypeError: 

360 return a == b # type: ignore[no-any-return] 

361 

362 

363def _merge_data_group(lhs: DataGroup[_E1], rhs: DataGroup[_E2]) -> DataGroup[_E1 | _E2]: 

364 res: DataGroup[_E1 | _E2] = DataGroup(dict(lhs)) 

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

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

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

368 res[k] = v 

369 return res 

370 

371 

372def _raw_positional_index( 

373 sizes: dict[str, int], coord: Variable, index: slice | Variable 

374) -> tuple[str, list[int | Variable]]: 

375 dim, *inds = _call_cpp_func( 

376 _cpp.label_based_index_to_positional_index, 

377 list(sizes.keys()), 

378 list(sizes.values()), 

379 coord, 

380 index, 

381 ) 

382 return dim, inds # type: ignore[return-value] 

383 

384 

385def label_based_index_to_positional_index( 

386 sizes: dict[str, int], 

387 coord: Variable, 

388 index: slice | Variable, 

389) -> ScippIndex: 

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

391 the coord with values.""" 

392 dim, inds = _raw_positional_index(sizes, coord, index) 

393 # Length of inds is 1 if index was a variable 

394 if len(inds) == 1: 

395 if inds[0] < 0 or sizes[dim] <= inds[0]: 

396 # If index is a variable and the coord is a bin-edge coord 

397 # - if the index is less than any edge then inds[0] is -1 

398 # - if the index is greater than any edge then inds[0] is sizes[dim] 

399 raise IndexError( 

400 f"Value {index} is not contained in the bin-edge coord {coord}" 

401 ) 

402 return (dim, inds[0]) 

403 return (dim, slice(*inds)) 

404 

405 

406def as_const(x: _VarDa) -> _VarDa: 

407 """Return a copy with the readonly flag set.""" 

408 return _call_cpp_func(_cpp.as_const, x) # type: ignore[return-value]