Coverage for tests/unit/math/test_bins.py: 100%

1from __future__ import annotations 

2 

3import math 

4from typing import Final 

5 

6import numpy as np 

7import pytest 

8from numpy.testing import assert_allclose 

9 

10from muutils.math.bins import Bins 

11 

12 

13_LOG_MIN: Final[float] = 1e-3 

14 

15 

16def _expected_log_edges( 

17 start: float, 

18 stop: float, 

19 n_bins: int, 

20 *, 

21 log_min: float = _LOG_MIN, 

22) -> np.ndarray: 

23 """Mimic the precise branch logic of `Bins.edges` for log scale.""" 

24 if math.isclose(start, 0.0): 

25 return np.concatenate( 

26 [ 

27 np.array([0.0]), 

28 np.logspace(math.log10(log_min), math.log10(stop), n_bins), 

29 ] 

30 ) 

31 if start < log_min: 

32 return np.concatenate( 

33 [np.array([0.0]), np.logspace(math.log10(start), math.log10(stop), n_bins)] 

34 ) 

35 return np.logspace(math.log10(start), math.log10(stop), n_bins + 1) 

36 

37 

38@pytest.mark.parametrize( 

39 "n_bins,start,stop", 

40 [(1, 0.0, 1.0), (4, -3.0, 5.0), (10, 2.5, 7.5)], 

41) 

42def test_edges_linear(n_bins: int, start: float, stop: float) -> None: 

43 bins = Bins(n_bins=n_bins, start=start, stop=stop, scale="lin") 

44 expected = np.linspace(start, stop, n_bins + 1) 

45 assert_allclose(bins.edges, expected) 

46 assert bins.edges.shape == (n_bins + 1,) 

47 

48 

49@pytest.mark.parametrize("n_bins", [1, 4, 16]) 

50def test_centers_linear(n_bins: int) -> None: 

51 bins = Bins(n_bins=n_bins, start=0.0, stop=1.0, scale="lin") 

52 expected_centers = (bins.edges[:-1] + bins.edges[1:]) / 2 

53 assert_allclose(bins.centers, expected_centers) 

54 assert bins.centers.shape == (n_bins,) 

55 

56 

57@pytest.mark.parametrize("start", [1e-2, 0.1, 1.0]) 

58def test_edges_log_standard(start: float) -> None: 

59 n_bins, stop = 8, 10.0 

60 bins = Bins(n_bins=n_bins, start=start, stop=stop, scale="log") 

61 expected = _expected_log_edges(start, stop, n_bins) 

62 assert_allclose(bins.edges, expected) 

63 assert bins.edges.shape == (n_bins + 1,) 

64 

65 

66def test_edges_log_start_zero() -> None: 

67 n_bins, stop = 6, 1.0 

68 bins = Bins(n_bins=n_bins, start=0.0, stop=stop, scale="log") 

69 expected = _expected_log_edges(0.0, stop, n_bins) 

70 assert_allclose(bins.edges, expected) 

71 assert bins.edges[0] == 0.0 

72 

73 

74def test_edges_log_small_start_include_zero() -> None: 

75 n_bins, start, stop = 5, 1e-4, 1.0 

76 bins = Bins(n_bins=n_bins, start=start, stop=stop, scale="log") 

77 expected = _expected_log_edges(start, stop, n_bins) 

78 assert_allclose(bins.edges, expected) 

79 assert bins.edges[0] == 0.0 

80 

81 

82def test_log_negative_start_raises() -> None: 

83 with pytest.raises(ValueError): 

84 _ = Bins(n_bins=3, start=-0.1, stop=1.0, scale="log").edges 

85 

86 

87def test_invalid_scale_raises() -> None: 

88 with pytest.raises(ValueError): 

89 _ = Bins(n_bins=3, start=0.0, stop=1.0, scale="strange").edges # type: ignore[arg-type] 

90 

91 

92def test_changed_n_bins_copy() -> None: 

93 original = Bins(n_bins=4, start=0.0, stop=1.0, scale="lin") 

94 new_bins = original.changed_n_bins_copy(10) 

95 

96 assert new_bins is not original 

97 assert new_bins.n_bins == 10 

98 for attr in ("start", "stop", "scale", "_log_min", "_zero_in_small_start_log"): 

99 assert getattr(new_bins, attr) == getattr(original, attr) 

100 

101 assert new_bins.edges.shape == (11,) 

102 

103 

104@pytest.mark.parametrize("n_bins,scale", [(1, "lin"), (3, "lin"), (7, "log")]) 

105def test_edges_shape(n_bins: int, scale: str) -> None: 

106 bins = Bins( 

107 n_bins=n_bins, 

108 start=0.1 if scale == "log" else 0.0, 

109 stop=2.0, 

110 scale=scale, # type: ignore 

111 ) 

112 assert bins.edges.shape == (n_bins + 1,)