Coverage for tests/test_nb_tools.py: 100%

1from starlord import nb_tools as nt 

2from scipy import interpolate 

3import numpy as np 

4import pytest 

5 

6 

7def test_locate_point_uniform(): 

8 # Test underlying python function for coverage and agreement; valid but may confuse linters 

9 locate_point = nt._locate_point_.py_func # pyright: ignore 

10 process_axis = nt._process_axis_.py_func # pyright: ignore 

11 # Uniform input axis 

12 n = 33 

13 x = np.linspace(5, 15, n) 

14 grid = process_axis(x, n) 

15 assert grid == pytest.approx(nt._process_axis_(x, n)) 

16 for xt in 5. + 10. * np.random.rand(100): 

17 i, weight = nt._locate_point_(xt, grid, n) 

18 assert i == np.where(xt >= x)[0][-1] 

19 assert i < n - 1 

20 assert weight >= 0. 

21 assert weight < 1. or (weight == 1 and i == n - 2) 

22 i2, weight2 = locate_point(xt, grid, n) 

23 assert i2 == i2 

24 assert weight == pytest.approx(weight2) 

25 assert nt._locate_point_(3.2, grid, n)[0] == -1 

26 assert locate_point(3.2, grid, n)[0] == -1 

27 assert nt._locate_point_(16., grid, n)[0] == -1 

28 assert locate_point(16., grid, n)[0] == -1 

29 assert nt._locate_point_(np.nan, grid, n)[0] == -1 

30 assert locate_point(np.nan, grid, n)[0] == -1 

31 assert nt._locate_point_(np.inf, grid, n)[0] == -1 

32 assert locate_point(np.inf, grid, n)[0] == -1 

33 assert nt._locate_point_(15., grid, n) == (n - 2, 1.) 

34 assert locate_point(15., grid, n) == (n - 2, 1.) 

35 

36 

37def test_locate_point_nonuniform(): 

38 # Test underlying python function for coverage and agreement; valid but may confuse linters 

39 locate_point = nt._locate_point_.py_func # pyright: ignore 

40 process_axis = nt._process_axis_.py_func # pyright: ignore 

41 # Non-uniform input axis 

42 n = 57 

43 x = np.logspace(-1, 2, n) 

44 grid = process_axis(x, n) 

45 assert grid == pytest.approx(nt._process_axis_(x, n)) 

46 for xt in .1 + 99.9 * np.random.rand(100): 

47 i, weight = nt._locate_point_(xt, grid, n) 

48 assert i == np.where(xt >= x)[0][-1] 

49 assert i < n - 1 

50 assert weight >= 0. 

51 assert weight < 1. or (weight == 1 and i == n - 2) 

52 i2, weight2 = locate_point(xt, grid, n) 

53 assert i2 == i2 

54 assert weight == pytest.approx(weight2) 

55 assert nt._locate_point_(-2., grid, n)[0] == -1 

56 assert locate_point(-2., grid, n)[0] == -1 

57 assert nt._locate_point_(105., grid, n)[0] == -1 

58 assert locate_point(105., grid, n)[0] == -1 

59 assert nt._locate_point_(np.nan, grid, n)[0] == -1 

60 assert locate_point(np.nan, grid, n)[0] == -1 

61 assert nt._locate_point_(np.inf, grid, n)[0] == -1 

62 assert locate_point(np.inf, grid, n)[0] == -1 

63 assert nt._locate_point_(100., grid, n) == (n - 2, 1.) 

64 assert locate_point(100., grid, n) == (n - 2, 1.) 

65 

66 

67def test_1d_gridding(): 

68 x = np.linspace(-5, 5, 100) 

69 y = np.sin(x) + 3.2 * np.cos(x) + .12*x 

70 grid = nt.pack_interpolator([x], y) 

71 assert grid[0] == nt._magicNumber 

72 assert grid[1] == 1. 

73 ref = interpolate.RegularGridInterpolator([x], y) 

74 for xt in (-5. + 10. * np.random.rand(100)): 

75 assert nt.interp1d(grid, xt) == pytest.approx(ref([xt])[0]) 

76 assert nt.interp1d(grid, -5.) == pytest.approx(ref([-5.])[0]) 

77 assert nt.interp1d(grid, 5.) == pytest.approx(ref([5.])[0]) 

78 for xt in (6. + 1000. * np.random.rand(100)): 

79 assert np.isnan(nt.interp1d(grid, xt)) 

80 for xt in (-6. - 1000. * np.random.rand(100)): 

81 assert np.isnan(nt.interp1d(grid, xt)) 

82 assert np.isnan(nt.interp1d(grid, np.inf)) 

83 assert np.isnan(nt.interp1d(grid, np.nan)) 

84 # Test underlying python function for coverage; valid but may confuse linters 

85 interp1d = nt.interp1d.py_func # pyright: ignore 

86 assert interp1d(grid, 3.5) == pytest.approx(ref([3.5])[0]) 

87 assert np.isnan(interp1d(grid, -15.)) 

88 assert np.isnan(interp1d(grid[::-1], 2.)) 

89 

90 

91def test_2d_gridding(): 

92 x = np.linspace(-5, 5, 73) 

93 y = np.logspace(-1, 1, 51) 

94 z = np.sin(x[:, None] + y[None, :]) + 3.2 * np.cos(y[None, :]) + .12 * x[:, None] 

95 grid = nt.pack_interpolator([x, y], z) 

96 assert grid[0] == nt._magicNumber 

97 assert grid[1] == 2. 

98 ref = interpolate.RegularGridInterpolator([x, y], z) 

99 xtest = np.random.rand(100, 2) * np.array([10., 9.9]) + np.array([-5., .1]) 

100 for xt in xtest: 

101 assert nt.interp2d(grid, xt[0], xt[1]) == pytest.approx(ref(xt)[0]) 

102 assert np.isnan(nt.interp2d(grid, -6, .2)) 

103 assert np.isnan(nt.interp2d(grid, 3., .09999)) 

104 assert np.isfinite(nt.interp2d(grid, -5., .5)) 

105 assert np.isfinite(nt.interp2d(grid, -2., .1)) 

106 assert np.isfinite(nt.interp2d(grid, -5., .1)) 

107 assert np.isfinite(nt.interp2d(grid, 5., 1.)) 

108 assert np.isfinite(nt.interp2d(grid, 3., 1.)) 

109 assert np.isnan(nt.interp2d(grid, np.inf, 2.5)) 

110 assert np.isnan(nt.interp2d(grid, -3.2, np.nan)) 

111 # Test underlying python function for coverage; valid but may confuse linters 

112 interp2d = nt.interp2d.py_func # pyright: ignore 

113 assert interp2d(grid, 3.5, .5) == pytest.approx(ref([3.5, .5])[0]) 

114 assert np.isnan(interp2d(grid, -15, .2)) 

115 assert np.isnan(interp2d(grid[::-1], 2., .5)) 

116 

117 

118def test_3d_gridding(): 

119 x = np.linspace(-5, 5, 73) 

120 y = np.logspace(-1, 1, 51) 

121 z = np.linspace(2, 3.5, 13) 

122 u = np.sin(x[:, None, None] + y[None, :, None]) + 3.2 * np.cos(z[None, None, :]) + .12 * x[:, None, None] 

123 grid = nt.pack_interpolator([x, y, z], u) 

124 assert grid[0] == nt._magicNumber 

125 assert grid[1] == 3. 

126 ref = interpolate.RegularGridInterpolator([x, y, z], u) 

127 xtest = np.random.rand(100, 3) * np.array([10., 9.9, 1.5]) + np.array([-5., .1, 2.]) 

128 for xt in xtest: 

129 assert nt.interp3d(grid, *xt) == pytest.approx(ref(xt)[0]) 

130 assert np.isnan(nt.interp3d(grid, -6, .2, 1.5)) 

131 assert np.isnan(nt.interp3d(grid, 3., .19999, 3.6)) 

132 assert np.isfinite(nt.interp3d(grid, -5., .1, 2.)) 

133 assert np.isfinite(nt.interp3d(grid, -3., .1, 3.5)) 

134 assert np.isfinite(nt.interp3d(grid, -5., 1., 2.)) 

135 assert np.isnan(nt.interp3d(grid, np.inf, 2.5, 3.4)) 

136 assert np.isnan(nt.interp3d(grid, -3.2, 2.5123, np.nan)) 

137 # Test underlying python function for coverage; valid but may confuse linters 

138 interp3d = nt.interp3d.py_func # pyright: ignore 

139 assert interp3d(grid, 3.5, .5, 2.2) == pytest.approx(ref([3.5, .5, 2.2])[0]) 

140 assert np.isnan(interp3d(grid, -4, 15., 3.1)) 

141 assert np.isnan(interp3d(grid, -4, .2, 3.51)) 

142 assert np.isnan(interp3d(grid[::-1], 2., .5, 3.2)) 

143 

144 

145def test_4d_gridding(): 

146 x = np.linspace(-5, 5, 73) 

147 y = np.logspace(-1, 1, 51) 

148 z = np.linspace(2, 3.5, 13) 

149 u = np.array([-5, -3, 0., 1.5, 15.]) 

150 v = np.sin(u[None, None, None, :] + .530 * y[None, :, None, None]) 

151 v = v + 3.2 * np.cos(z[None, None, :, None]) + .12 * x[:, None, None, None] 

152 grid = nt.pack_interpolator([x, y, z, u], v) 

153 assert grid[0] == nt._magicNumber 

154 assert grid[1] == 4. 

155 ref = interpolate.RegularGridInterpolator([x, y, z, u], v) 

156 xtest = np.random.rand(100, 4) * np.array([10., 9.9, 1.5, 20.]) + np.array([-5., .1, 2., -5]) 

157 for xt in xtest: 

158 assert nt.interp4d(grid, *xt) == pytest.approx(ref(xt)[0]) 

159 assert np.isnan(nt.interp4d(grid, -6, .2, 2.5, 16.1)) 

160 assert np.isnan(nt.interp4d(grid, 3., .19999, 3.3, -5.001)) 

161 assert np.isfinite(nt.interp4d(grid, -5., .1, 2., -5)) 

162 assert np.isfinite(nt.interp4d(grid, -3., .15, 2.5, 15.)) 

163 assert np.isfinite(nt.interp4d(grid, 5., 1., 3.5, 15)) 

164 assert np.isnan(nt.interp4d(grid, -4., 2.5, 3.4, np.inf)) 

165 assert np.isnan(nt.interp4d(grid, -3.2, 2.5123, 3.4, np.nan)) 

166 # Test underlying python function for coverage; valid but may confuse linters 

167 interp4d = nt.interp4d.py_func # pyright: ignore 

168 assert interp4d(grid, 3.5, .5, 2.2, 0.1) == pytest.approx(ref([3.5, .5, 2.2, 0.1])[0]) 

169 assert np.isnan(interp4d(grid, -4, 2., 2.9, 15.1)) 

170 assert np.isnan(interp4d(grid, -4, .2, .2, -5.1)) 

171 assert np.isnan(interp4d(grid[::-1], 2., .5, 3.2, 3.5))