Coverage for python/lsst/meas/photoz/base/utils.py: 85%

1# This file is part of meas_photoz_base. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = [ 

25 "chunk_iterator", 

26] 

27 

28from collections.abc import Generator 

29 

30import numpy as np 

31 

32 

33def chunk_iterator( 

34 chunk_size: int, 

35 data: dict[str, np.ndarray], 

36 n_values: int, 

37) -> Generator[tuple[dict, int]]: 

38 """Yield chunks of a dict of equal-length arrays in a fixed chunk size. 

39 

40 Parameters 

41 ---------- 

42 chunk_size : int 

43 The size of each chunk. 

44 data : dict[str, np.ndarray] 

45 The input data. 

46 n_values : int 

47 Length of each array in data. 

48 

49 Yields 

50 ------ 

51 a_chunk : dict[str, np.array] 

52 A chunk of data. 

53 

54 chunk_size: int 

55 The number of objects in the chunk. 

56 """ 

57 if not chunk_size > 0: 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true

58 raise ValueError(f"{chunk_size=} must be >0") 

59 for key, values in data.items(): 

60 if (n_key := len(values)) != n_values: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true

61 raise ValueError(f"{key=} has len={n_key} != {n_values=}") 

62 

63 chunk_start = 0 

64 chunk_stop = 0 

65 

66 while chunk_stop < n_values: 

67 chunk_stop = min(chunk_start + chunk_size, n_values) 

68 this_chunk_size = chunk_stop - chunk_start 

69 this_chunk = {key: val[chunk_start:chunk_stop] for key, val in data.items()} 

70 chunk_start += this_chunk_size 

71 yield this_chunk, this_chunk_size