Coverage for tests/stars.py: 95%

166 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-12-08 10:31 +0900

1import pytest 

2import airball 

3import airball.units as u 

4import numpy as np 

5 

6################################################ 

7################################################ 

8########## INITIALIZATION TESTS ############## 

9################################################ 

10################################################ 

11 

12def test_initialize_Stars_with_lists(): 

13 stars = airball.Stars(m=[1,2,3], b=[4,5,6], v=[7,8,9]) 

14 assert np.all(stars.m == np.array([1,2,3]) * u.solMass) 

15 assert np.all(stars.b == np.array([4,5,6]) * u.au) 

16 assert np.all(stars.v == np.array([7,8,9]) * u.km/u.s) 

17 

18def test_initialize_Stars_with_ndarrays(): 

19 stars = airball.Stars(m=np.array([1,2,3]), b=np.array([4,5,6]), v=np.array([7,8,9])) 

20 assert np.all(stars.m == np.array([1,2,3]) * u.solMass) 

21 assert np.all(stars.b == np.array([4,5,6]) * u.au) 

22 assert np.all(stars.v == np.array([7,8,9]) * u.km/u.s) 

23 

24def test_initialize_Stars_with_lists_and_units(): 

25 stars = airball.Stars(m=[1,2,3]*u.solMass, b=[4,5,6]*u.au, v=[7,8,9]*u.km/u.s) 

26 assert np.all(stars.m == np.array([1,2,3]) * u.solMass) 

27 assert np.all(stars.b == np.array([4,5,6]) * u.au) 

28 assert np.all(stars.v == np.array([7,8,9]) * u.km/u.s) 

29 

30 stars = airball.Stars(m=[1*u.solMass, 2000*u.jupiterMass, 1e6*u.earthMass], b=[4*u.au, 5*u.au, 6*u.au], v=[7*u.km/u.s, 8*u.km/u.s, 9*u.km/u.s]) 

31 assert np.all(stars.m == np.array([1, 1.90918846793865, 3.003489348850793]) * u.solMass) 

32 assert np.all(stars.b == np.array([4,5,6]) * u.au) 

33 assert np.all(stars.v == np.array([7,8,9]) * u.km/u.s) 

34 

35def test_initialize_Stars_with_ndarrays_and_units(): 

36 stars = airball.Stars(m=np.array([1,2,3])*u.solMass, b=np.array([4,5,6])*u.au, v=np.array([7,8,9])*u.km/u.s) 

37 assert np.all(stars.m == np.array([1,2,3]) * u.solMass) 

38 assert np.all(stars.b == np.array([4,5,6]) * u.au) 

39 assert np.all(stars.v == np.array([7,8,9]) * u.km/u.s) 

40 

41def test_initialize_Stars_with_lists_and_shapes(): 

42 stars = airball.Stars(m=[[1,2,3],[1,2,3]], b=[[4,5,6], [4,5,6]], v=[[7,8,9], [7,8,9]]) 

43 assert np.all(stars.m == np.array([[1,2,3],[1,2,3]]) * u.solMass) 

44 assert np.all(stars.b == np.array([[4,5,6], [4,5,6]]) * u.au) 

45 assert np.all(stars.v == np.array([[7,8,9], [7,8,9]]) * u.km/u.s) 

46 

47def test_initialize_Stars_with_lists_and_odd_shapes(): 

48 with pytest.raises(airball.stars.ListLengthException): 

49 airball.Stars(m=[[1,2,3],[1,2,3]], b=[[4,5], [4,6]], v=[7,8,9]) 

50 

51def test_initialize_Stars_with_ndarray_shapes(): 

52 stars = airball.Stars(m=np.array([[1,2,3],[1,2,3]])*u.solMass, b=np.array([[4,5,6], [4,5,6]])*u.au, v=np.array([[7,8,9], [7,8,9]])*u.km/u.s) 

53 assert np.all(stars.m == np.array([[1,2,3],[1,2,3]]) * u.solMass) 

54 assert np.all(stars.b == np.array([[4,5,6], [4,5,6]]) * u.au) 

55 assert np.all(stars.v == np.array([[7,8,9], [7,8,9]]) * u.km/u.s) 

56 

57def test_initialize_Stars_with_ndarray_odd_shapes(): 

58 with pytest.raises(airball.stars.ListLengthException): 

59 airball.Stars(m=np.array([[1,2,3],[1,2,3]])*u.solMass, b=np.array([[4,5], [4,6]])*u.au, v=np.array([7,8,9])*u.km/u.s) 

60 with pytest.raises(airball.stars.ListLengthException): 

61 airball.Stars(m=np.array([[1,2,3],[1,2,3],[1,2,3]])*u.solMass, b=np.array([[4,5], [4,6]])*u.au, v=np.array([7,8])*u.km/u.s) 

62 with pytest.raises(airball.stars.ListLengthException): 

63 airball.Stars(m=np.array([[1,2,3],[1,2,3]])*u.solMass, b=np.array([[4,5], [4,6]])*u.au, v=np.array([7,8])*u.km/u.s) 

64 

65def test_initialize_Stars_with_mixed_lists(): 

66 stars = airball.Stars(m=[1,2,3], b=np.array([4,5,6]), v=[7,8,9]*u.km/u.s) 

67 assert np.all(stars.m == np.array([1,2,3]) * u.solMass) 

68 assert np.all(stars.b == np.array([4,5,6]) * u.au) 

69 assert np.all(stars.v == np.array([7,8,9]) * u.km/u.s) 

70 

71def test_initialize_Stars_with_mixed_lists_and_floats(): 

72 stars = airball.Stars(m=np.array([4,5,6]), b=1, v=7*u.km/u.s) 

73 assert np.all(stars.m == np.array([4,5,6]) * u.solMass) 

74 assert np.all(stars.b == np.array([1,1,1]) * u.au) 

75 assert np.all(stars.v == np.array([7,7,7]) * u.km/u.s) 

76 

77 stars = airball.Stars(m=1*u.solMass, b=np.array([4,5,6]), v=7) 

78 assert np.all(stars.m == np.array([1,1,1]) * u.solMass) 

79 assert np.all(stars.b == np.array([4,5,6]) * u.au) 

80 assert np.all(stars.v == np.array([7,7,7]) * u.km/u.s) 

81 

82def test_initialize_Stars_with_floats_and_size(): 

83 stars = airball.Stars(m=1, b=3.0, v=7*u.km/u.s, size=3) 

84 assert np.all(stars.m == np.array([1,1,1]) * u.solMass) 

85 assert np.all(stars.b == np.array([3,3,3]) * u.au) 

86 assert np.all(stars.v == np.array([7,7,7]) * u.km/u.s) 

87 

88def test_initialize_Stars_with_uneven_lists(): 

89 with pytest.raises(airball.stars.ListLengthException): 

90 airball.Stars(m=[1], b=[4,5,6], v=[7,8,9]) 

91 with pytest.raises(airball.stars.ListLengthException): 

92 airball.Stars(m=[1,2,3], b=[4], v=[7,8,9]) 

93 with pytest.raises(airball.stars.ListLengthException): 

94 airball.Stars(m=[1,2,3], b=[4,5,6], v=[7]) 

95 

96def test_initialize_Stars_with_unspecified_parameters(): 

97 with pytest.raises(airball.stars.UnspecifiedParameterException): 

98 airball.Stars(b=[4,5,6], v=[7,8,9]) 

99 with pytest.raises(airball.stars.UnspecifiedParameterException): 

100 airball.Stars(m=[1,2,3], v=[7,8,9]) 

101 with pytest.raises(airball.stars.UnspecifiedParameterException): 

102 airball.Stars(m=[1,2,3], b=[4,5,6]) 

103 

104def test_initialize_Stars_with_floats_and_unspecified_size(): 

105 with pytest.raises(airball.stars.UnspecifiedParameterException): 

106 airball.Stars(m=1, b=3, v=7) 

107 

108def test_initialize_Stars_with_lists_and_specified_size(): 

109 with pytest.raises(airball.stars.OverspecifiedParametersException): 

110 airball.Stars(m=[1,2], v=[7,8,9], size=4) 

111 

112def test_initialize_Stars_with_custom_object(): 

113 # Test using a custom object that defines __len__. 

114 class CustomObject: 

115 def __init__(self, data): self.data = data 

116 def __len__(self): return len(self.data) 

117 @property 

118 def shape(self): return (len(self.data),) 

119 

120 # Create an instance of the custom object 

121 my_object = CustomObject([1, 2, 3]) 

122 with pytest.raises(airball.stars.IncompatibleListException): 

123 airball.Stars(m=my_object, b=[4,5,6], v=[7,8,9]) 

124 airball.Stars(m=[1,2,3], b=my_object, v=[7,8,9]) 

125 airball.Stars(m=[1,2,3], b=[4,5,6], v=my_object) 

126 

127def test_initialize_Stars_with_custom_object_no_shape(): 

128 # Test using a custom object that defines __len__. 

129 class CustomObject: 

130 def __init__(self, data): self.data = data 

131 def __len__(self): return len(self.data) 

132 

133 # Create an instance of the custom object 

134 my_object = CustomObject([1, 2, 3]) 

135 with pytest.raises(airball.stars.IncompatibleListException): 

136 airball.Stars(m=my_object, b=[4,5,6], v=[7,8,9]) 

137 airball.Stars(m=[1,2,3], b=my_object, v=[7,8,9]) 

138 airball.Stars(m=[1,2,3], b=[4,5,6], v=my_object) 

139 

140def test_initialize_Stars_with_custom_object_different_list_length(): 

141 # Test using a custom object that defines __len__. 

142 class CustomObject: 

143 def __init__(self, data): self.data = data 

144 def __len__(self): return len(self.data) 

145 @property 

146 def shape(self): return (len(self.data),) 

147 

148 # Create an instance of the custom object 

149 my_object = CustomObject([1, 2]) 

150 with pytest.raises(airball.stars.ListLengthException): 

151 airball.Stars(m=my_object, b=[4,5,6], v=[7,8,9]) 

152 airball.Stars(m=[1,2,3], b=my_object, v=[7,8,9]) 

153 airball.Stars(m=[1,2,3], b=[4,5,6], v=my_object) 

154 

155def test_initialize_Stars_with_custom_object_no_shape_different_list_length(): 

156 # Test using a custom object that defines __len__. 

157 class CustomObject: 

158 def __init__(self, data): self.data = data 

159 def __len__(self): return len(self.data) 

160 

161 # Create an instance of the custom object 

162 my_object = CustomObject([1, 2]) 

163 with pytest.raises(airball.stars.ListLengthException): 

164 airball.Stars(m=my_object, b=[4,5,6], v=[7,8,9]) 

165 airball.Stars(m=[1,2,3], b=my_object, v=[7,8,9]) 

166 airball.Stars(m=[1,2,3], b=[4,5,6], v=my_object) 

167 

168def test_initialize_Stars_with_large_ndarrays(): 

169 ones = np.ones(int(1e7)) 

170 stars = airball.Stars(m=ones, b=2*ones, v=3*ones) 

171 assert np.all(stars.m == ones * u.solMass) 

172 assert np.all(stars.b == 2*ones * u.au) 

173 assert np.all(stars.v == 3*ones * u.km/u.s) 

174 

175################################################ 

176################################################ 

177########## INDEXING AND SLICING TESTS ######## 

178################################################ 

179################################################ 

180 

181def test_access_Stars_using_ints(): 

182 stars = airball.Stars(m=[1,2,3], b=[4,5,6], v=[7,8,9], inc=[3,2,1], omega=[6,5,4], Omega=[9,8,7]) 

183 assert stars[0] == airball.Star(m=1, b=4, v=7, inc=3, omega=6, Omega=9) 

184 

185def test_access_multidimensional_Stars_using_ints(): 

186 stars = airball.Stars(m=[[1,2,3],[2,2,3]], b=[[4,5,6],[5,5,6]], v=[[7,8,9],[8,8,9]], inc=[[3,2,1], [4,2,1]], omega=[[6,5,4],[7,5,4]], Omega=[[9,8,7],[10,8,7]]) 

187 assert stars[0] == airball.Stars(m=[1,2,3], b=[4,5,6], v=[7,8,9], inc=[3,2,1], omega=[6,5,4], Omega=[9,8,7]) 

188 

189 

190################################################ 

191################################################ 

192########## EQUALITY TESTS #################### 

193################################################ 

194################################################ 

195 

196def test_Star_equality(): 

197 star = airball.Star(m=1, b=4, v=7, inc=3, omega=6, Omega=9) 

198 assert star == airball.Star(m=1, b=4, v=7, inc=3, omega=6, Omega=9) 

199 assert star != airball.Star(m=1, b=4, v=7, inc=3, omega=6, Omega=0) 

200 assert hash(star) == hash(airball.Star(m=1, b=4, v=7, inc=3, omega=6, Omega=9)) 

201 assert hash(star) != hash(airball.Star(m=1, b=4, v=7, inc=3, omega=6, Omega=0)) 

202 

203def test_Stars_equality(): 

204 stars1 = airball.Stars(m=[1,2,3], b=[4,5,6], v=[7,8,9], inc=[3,2,1], omega=[6,5,4], Omega=[9,8,7]) 

205 stars2 = airball.Stars(m=[1,2,3], b=[4,5,6], v=[7,8,9], inc=[3,2,1], omega=[6,5,4], Omega=[9,8,7]) 

206 assert stars1 == stars2 

207 assert stars1[:2] == stars2[:2] 

208 assert hash(stars1) == hash(stars2) 

209 stars2.sort('m') 

210 assert stars1 == stars2 

211 stars2.sort('inc') 

212 assert stars1 != stars2 

213 assert hash(stars1) != hash(stars2) 

214 

215def test_Stars_from_StellarEnvironment_equality(): 

216 oc = airball.OpenCluster() 

217 stars1 = oc.random_star(size=300) 

218 stars2 = stars1.copy() 

219 assert stars1 == stars2 

220 assert stars1[:2] == stars2[:2] 

221 assert hash(stars1) == hash(stars2) 

222 stars2.sort('m') 

223 assert stars1 != stars2 

224 assert hash(stars1) != hash(stars2) 

225