Coverage for tests/test_spectrum.py: 95%

1import itertools 

2 

3from hypothesis import given, strategies as st 

4import hypothesis.extra.numpy as npst 

5import primefac 

6import pytest 

7 

8from flexfrac1d.model.model import DiscreteSpectrum 

9 

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

11number = st.one_of( 

12 st.floats(**float_kw), 

13 st.integers(), 

14) 

15 

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

17non_negative_number = st.one_of( 

18 st.floats(**non_neg_float_kw), 

19 st.integers(min_value=0), 

20) 

21 

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

23positive_number = st.one_of( 

24 st.floats(**pos_float_kw), 

25 st.integers(min_value=1), 

26) 

27 

28 

29def 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 

37def 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 

56def 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 

70def 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 

99def 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()) 

146def test_sanitised(args): 

147 amplitudes, frequencies, kwargs = args 

148 DiscreteSpectrum(amplitudes, frequencies, **kwargs) 

149 

150 

151@given(args=not_broadcastable()) 

152def 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)