Coverage for tests/test_spectrum.py: 95%

1 import itertools

2

3 from hypothesis import given, strategies as st

4 import hypothesis.extra.numpy as npst

5 import primefac

6 import pytest

7

8 from flexfrac1d.model.model import DiscreteSpectrum

9

10 float_kw = {"allow_nan": False, "allow_infinity": False}

11 number = st.one_of(

12 st.floats(**float_kw),

13 st.integers(),

14)

15

16 non_neg_float_kw = float_kw | {"min_value": 0}

17 non_negative_number = st.one_of(

18 st.floats(**non_neg_float_kw),

19 st.integers(min_value=0),

20)

21

22 pos_float_kw = non_neg_float_kw | {"allow_subnormal": False, "exclude_min": True}

23 positive_number = st.one_of(

24 st.floats(**pos_float_kw),

25 st.integers(min_value=1),

26)

27

28

29 def get_optional_kwargs(*args):

30 combinations = [

31 _c for n in range(len(args) + 1) for _c in itertools.combinations(args, n)

32 ]

33 return st.sampled_from(combinations)

34

35

36@st.composite

37 def comp_shapes(draw, size, max_dims=None):

38 if size == 0:

39 return tuple()

40 if size in (1, 2, 3, 5, 7, 11, 13, 17, 19, 23):

41 return (size,)

42

43 factors = draw(st.permutations(list(primefac.primefac(size))))

44 if max_dims is None: 44 ↛ 47line 44 didn't jump to line 47, because the condition on line 44 was never false

45 max_dims = len(factors)

46

47 ndim = draw(st.integers(min_value=1, max_value=max_dims))

48 while len(factors) > ndim:

49 factors[0] = factors[0] * factors[-1]

50 factors = factors[:-1]

51

52 return tuple(factors)

53

54

55@st.composite

56 def broadcastable(draw):

57 size = draw(st.integers(min_value=0, max_value=255))

58 shape_st = comp_shapes(size)

59 ds_kw = dict()

60

61 amplitudes = draw(get_number_or_array(shape_st, False, "pos"))

62 frequencies = draw(get_number_or_array(shape_st, False, "pos"))

63 opt = draw(get_optional_kwargs("phases", "shapes"))

64 if "phases" in opt:

65 ds_kw["phases"] = draw(get_number_or_array(shape_st, False))

66

67 return amplitudes, frequencies, ds_kw

68

69

70 def get_number_or_array(shape, strict=False, constraints=None):

71 if constraints == "non_neg": 71 ↛ 72line 71 didn't jump to line 72, because the condition on line 71 was never true

72 strategy = npst.arrays(

73 npst.floating_dtypes(),

74 shape=shape,

75 elements=non_neg_float_kw,

76 )

77 if not strict:

78 return non_negative_number | strategy

79 elif constraints == "pos":

80 strategy = npst.arrays(

81 npst.floating_dtypes(),

82 shape=shape,

83 elements=pos_float_kw,

84 )

85 if not strict:

86 return positive_number | strategy

87 else:

88 strategy = npst.arrays(

89 npst.floating_dtypes(),

90 shape=shape,

91 elements=float_kw,

92 )

93 if not strict:

94 return number | strategy

95 return strategy

96

97

98@st.composite

99 def not_broadcastable(draw):

100 # 1 is excluded as it will be compatible with anything, and we want to include 0

101 sizes = draw(

102 st.lists(

103 st.integers(min_value=0, max_value=255).filter(lambda n: n != 1),

104 min_size=2,

105 max_size=3,

106 unique=True,

107 )

108 )

109 if len(sizes) == 2:

110 with_phase = draw(st.booleans())

111 if with_phase:

112 idx_scalar = draw(st.integers(min_value=0, max_value=2))

113 if idx_scalar == 0:

114 amplitudes = draw(get_number_or_array(1, constraints="pos"))

115 frequencies = draw(

116 get_number_or_array(sizes[0], True, constraints="pos")

117 )

118 phases = draw(get_number_or_array(sizes[1], True))

119 elif idx_scalar == 1:

120 amplitudes = draw(

121 get_number_or_array(sizes[0], True, constraints="pos")

122 )

123 frequencies = draw(get_number_or_array(1, constraints="pos"))

124 phases = draw(get_number_or_array(sizes[1], True))

125 else:

126 amplitudes = draw(

127 get_number_or_array(sizes[0], True, constraints="pos")

128 )

129 frequencies = draw(

130 get_number_or_array(sizes[1], True, constraints="pos")

131 )

132 phases = draw(get_number_or_array(1, False))

133 else:

134 amplitudes = draw(get_number_or_array(sizes[0], True, constraints="pos"))

135 frequencies = draw(get_number_or_array(sizes[1], True, constraints="pos"))

136 phases = None

137 else:

138 amplitudes = draw(get_number_or_array(sizes[0], True, constraints="pos"))

139 frequencies = draw(get_number_or_array(sizes[1], True, constraints="pos"))

140 phases = draw(get_number_or_array(sizes[2], True))

141

142 return amplitudes, frequencies, phases

143

144

145@given(args=broadcastable())

146 def test_sanitised(args):

147 amplitudes, frequencies, kwargs = args

148 DiscreteSpectrum(amplitudes, frequencies, **kwargs)

149

150

151@given(args=not_broadcastable())

152 def test_unsanitised(args):

153 amplitudes, frequencies, phases = args

154 with pytest.raises(ValueError):

155 if phases is None:

156 DiscreteSpectrum(amplitudes, frequencies)

157 else:

158 DiscreteSpectrum(amplitudes, frequencies, phases)