Coverage for tests/stars.py: 95%

1 import pytest

2 import airball

3 import airball.units as u

4 import numpy as np

5

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

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

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

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

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

11

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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