Coverage for test_common.py: 100%

241 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 11:41 +0200

1import warnings 

2 

3import numba 

4import numpy as np 

5import xarray as xr 

6 

7import imod 

8 

9 

10def test_starts(): 

11 @numba.njit 

12 def get_starts(src_x, dst_x): 

13 result = [] 

14 for i, j in imod.prepare.common._starts(src_x, dst_x): 

15 result.append((i, j)) 

16 return result 

17 

18 # Complete range 

19 src_x = np.arange(0.0, 11.0, 1.0) 

20 dst_x = np.arange(0.0, 11.0, 2.5) 

21 # Returns tuples with (src_ind, dst_ind) 

22 # List comprehension gives PicklingError 

23 result = get_starts(src_x, dst_x) 

24 assert result == [(0, 0), (1, 2), (2, 5), (3, 7)] 

25 

26 # Partial dst 

27 dst_x = np.arange(5.0, 11.0, 2.5) 

28 result = get_starts(src_x, dst_x) 

29 assert result == [(0, 5), (1, 7)] 

30 

31 # Partial src 

32 src_x = np.arange(5.0, 11.0, 1.0) 

33 dst_x = np.arange(0.0, 11.0, 2.5) 

34 result = get_starts(src_x, dst_x) 

35 assert result == [(0, 0), (1, 0), (2, 0), (3, 2)] 

36 

37 # Irregular grid 

38 src_x = np.array([0.0, 2.5, 7.5, 10.0]) 

39 dst_x = np.array([0.0, 5.0, 10.0]) 

40 result = get_starts(src_x, dst_x) 

41 assert result == [(0, 0), (1, 1)] 

42 

43 # Negative coords 

44 src_x = np.arange(-20.0, -9.0, 1.0) 

45 dst_x = np.arange(-20.0, -9.0, 2.5) 

46 result = get_starts(src_x, dst_x) 

47 assert result == [(0, 0), (1, 2), (2, 5), (3, 7)] 

48 

49 # Mixed coords 

50 src_x = np.arange(-5.0, 6.0, 1.0) 

51 dst_x = np.arange(-5.0, 6.0, 2.5) 

52 result = get_starts(src_x, dst_x) 

53 assert result == [(0, 0), (1, 2), (2, 5), (3, 7)] 

54 

55 

56def test_weights(): 

57 src_x = np.arange(0.0, 11.0, 1.0) 

58 dst_x = np.arange(0.0, 11.0, 2.5) 

59 max_len, (dst_inds, src_inds, weights) = imod.prepare.common._weights_1d( 

60 src_x, dst_x, False 

61 ) 

62 assert max_len == 3 

63 assert np.allclose(dst_inds, np.array([0, 1, 2, 3])) 

64 assert np.allclose(src_inds, np.array([[0, 1, 2], [2, 3, 4], [5, 6, 7], [7, 8, 9]])) 

65 assert np.allclose( 

66 weights, 

67 np.array([[1.0, 1.0, 0.5], [0.5, 1.0, 1.0], [1.0, 1.0, 0.5], [0.5, 1.0, 1.0]]), 

68 ) 

69 

70 # Irregular grid 

71 src_x = np.array([0.0, 2.5, 7.5, 10.0]) 

72 dst_x = np.array([0.0, 5.0, 10.0]) 

73 max_len, (dst_inds, src_inds, weights) = imod.prepare.common._weights_1d( 

74 src_x, dst_x, False 

75 ) 

76 assert max_len == 2 

77 assert np.allclose(dst_inds, np.array([0, 1])) 

78 assert np.allclose(src_inds, np.array([[0, 1], [1, 2]])) 

79 assert np.allclose(weights, np.array([[2.5, 2.5], [2.5, 2.5]])) 

80 

81 # Mixed coords 

82 src_x = np.arange(-5.0, 6.0, 1.0) 

83 dst_x = np.arange(-5.0, 6.0, 2.5) 

84 max_len, (dst_inds, src_inds, weights) = imod.prepare.common._weights_1d( 

85 src_x, dst_x, False 

86 ) 

87 assert max_len == 3 

88 assert np.allclose(dst_inds, np.array([0, 1, 2, 3])) 

89 assert np.allclose(src_inds, np.array([[0, 1, 2], [2, 3, 4], [5, 6, 7], [7, 8, 9]])) 

90 assert np.allclose( 

91 weights, 

92 np.array([[1.0, 1.0, 0.5], [0.5, 1.0, 1.0], [1.0, 1.0, 0.5], [0.5, 1.0, 1.0]]), 

93 ) 

94 

95 

96def test_relative_weights(): 

97 # In the test above, the absolute weights are the same as the relative weights 

98 # To have a test case, we simply multiply coordinates by two, while the 

99 # relative weights should remain the same. 

100 src_x = np.arange(0.0, 11.0, 1.0) * 2.0 

101 dst_x = np.arange(0.0, 11.0, 2.5) * 2.0 

102 max_len, (dst_inds, src_inds, weights) = imod.prepare.common._weights_1d( 

103 src_x, dst_x, True 

104 ) 

105 assert max_len == 3 

106 assert np.allclose(dst_inds, np.array([0, 1, 2, 3])) 

107 assert np.allclose(src_inds, np.array([[0, 1, 2], [2, 3, 4], [5, 6, 7], [7, 8, 9]])) 

108 assert np.allclose( 

109 weights, 

110 np.array([[1.0, 1.0, 0.5], [0.5, 1.0, 1.0], [1.0, 1.0, 0.5], [0.5, 1.0, 1.0]]), 

111 ) 

112 

113 # Something non-equidistant 

114 src_x = np.array([0.0, 1.5]) 

115 dst_x = np.array([0.0, 3.0]) 

116 max_len, (dst_inds, src_inds, weights) = imod.prepare.common._weights_1d( 

117 src_x, dst_x, True 

118 ) 

119 assert np.allclose(weights, np.array([[1.0]])) 

120 

121 src_x = np.array([0.0, 3.0]) 

122 dst_x = np.array([0.0, 1.5]) 

123 max_len, (dst_inds, src_inds, weights) = imod.prepare.common._weights_1d( 

124 src_x, dst_x, True 

125 ) 

126 assert np.allclose(weights, np.array([[0.5]])) 

127 

128 

129def test_area_weighted_methods(): 

130 values = np.arange(5.0) 

131 weights = np.arange(0.0, 50.0, 10.0) 

132 values[0] = np.nan 

133 

134 assert np.allclose(imod.prepare.common.mean(values, weights), 3.0) 

135 assert np.allclose(imod.prepare.common.harmonic_mean(values, weights), 2.5) 

136 assert np.allclose( 

137 imod.prepare.common.geometric_mean(values, weights), 2.780778340631819 

138 ) 

139 

140 values[1] = 3.0 

141 assert np.allclose(imod.prepare.common.mode(values, weights), 3.0) 

142 

143 # Check if no issues arise with all nan 

144 values[:] = np.nan 

145 assert np.isnan(imod.prepare.common.mean(values, weights)) 

146 assert np.isnan(imod.prepare.common.harmonic_mean(values, weights)) 

147 assert np.isnan(imod.prepare.common.geometric_mean(values, weights)) 

148 assert np.isnan(imod.prepare.common.mode(values, weights)) 

149 

150 

151def test_methods(): 

152 values = np.arange(5.0) 

153 weights = np.arange(0.0, 50.0, 10.0) 

154 values[0] = np.nan 

155 assert np.allclose(imod.prepare.common.sum(values, weights), 10.0) 

156 assert np.allclose(imod.prepare.common.minimum(values, weights), 1.0) 

157 assert np.allclose(imod.prepare.common.maximum(values, weights), 4.0) 

158 assert np.allclose(imod.prepare.common.median(values, weights), 2.5) 

159 assert np.allclose(imod.prepare.common.conductance(values, weights), 300.0) 

160 assert np.allclose(imod.prepare.common.max_overlap(values, weights), 4.0) 

161 

162 # Check if no issues arise with all nan 

163 values[:] = np.nan 

164 with warnings.catch_warnings(): 

165 warnings.filterwarnings("ignore", "All-NaN slice encountered") 

166 assert np.isnan(imod.prepare.common.sum(values, weights)) 

167 assert np.isnan(imod.prepare.common.minimum(values, weights)) 

168 assert np.isnan(imod.prepare.common.maximum(values, weights)) 

169 assert np.isnan(imod.prepare.common.median(values, weights)) 

170 assert np.isnan(imod.prepare.common.conductance(values, weights)) 

171 assert np.isnan(imod.prepare.common.max_overlap(values, weights)) 

172 

173 

174def test_methods_zeros(): 

175 values = np.zeros(5) 

176 weights = np.arange(0.0, 50.0, 10.0) 

177 assert np.allclose(imod.prepare.common.mean(values, weights), 0.0) 

178 

179 

180def test_overlap(): 

181 assert imod.prepare.common._overlap((0.0, 1.0), (0.0, 2.0)) == 1.0 

182 assert imod.prepare.common._overlap((-1.0, 1.0), (0.0, 2.0)) == 1.0 

183 assert imod.prepare.common._overlap((-1.0, 3.0), (0.0, 2.0)) == 2.0 

184 assert imod.prepare.common._overlap((-1.0, 3.0), (-2.0, 2.0)) == 3.0 

185 

186 

187def test_reshape(): 

188 src = np.zeros((3, 5)) 

189 dst = np.zeros((3, 2)) 

190 iter_src, iter_dst = imod.prepare.common._reshape(src, dst, ndim_regrid=1) 

191 assert iter_src.shape == (3, 5) 

192 assert iter_dst.shape == (3, 2) 

193 

194 src = np.zeros((2, 4, 3, 5)) 

195 dst = np.zeros((2, 4, 3, 2)) 

196 iter_src, iter_dst = imod.prepare.common._reshape(src, dst, ndim_regrid=1) 

197 assert iter_src.shape == (24, 5) 

198 assert iter_dst.shape == (24, 2) 

199 

200 src = np.zeros((3, 5)) 

201 dst = np.zeros((3, 2)) 

202 iter_src, iter_dst = imod.prepare.common._reshape(src, dst, ndim_regrid=2) 

203 assert iter_src.shape == (1, 3, 5) 

204 assert iter_dst.shape == (1, 3, 2) 

205 

206 src = np.zeros((2, 4, 3, 5)) 

207 dst = np.zeros((2, 4, 3, 2)) 

208 iter_src, iter_dst = imod.prepare.common._reshape(src, dst, ndim_regrid=3) 

209 assert iter_src.shape == (2, 4, 3, 5) 

210 assert iter_dst.shape == (2, 4, 3, 2) 

211 

212 

213def test_is_subset(): 

214 # increasing 

215 a1 = np.array([0.0, 1.0, 2.0, 3.0]) 

216 a2 = np.array([0.0, 1.0]) 

217 assert imod.prepare.common._is_subset(a1, a2) 

218 a2 = np.array([0.0, 1.0, 3.0]) 

219 assert not imod.prepare.common._is_subset(a1, a2) 

220 # decreasing 

221 a1 = np.array([0.0, 1.0, 2.0, 3.0])[::-1] 

222 a2 = np.array([0.0, 1.0])[::-1] 

223 assert imod.prepare.common._is_subset(a1, a2) 

224 a2 = np.array([0.0, 1.0, 3.0])[::-1] 

225 assert not imod.prepare.common._is_subset(a1, a2) 

226 # Not a contiguous subset 

227 a1 = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) 

228 a2 = np.array([0.0, 2.0, 4.0]) 

229 assert not imod.prepare.common._is_subset(a1, a2) 

230 assert not imod.prepare.common._is_subset(a2, a1) 

231 

232 

233def test_selection_indices(): 

234 # Left-inclusive 

235 # Vertices 

236 src_x = np.array([0.0, 20.0, 40.0, 60.0, 80.0, 100.0]) 

237 xmin, xmax = 0.0, 20.0 

238 i0, i1 = imod.prepare.common._selection_indices(src_x, xmin, xmax, 0) 

239 assert i0 == 0 

240 assert i1 == 1 

241 

242 xmin, xmax = 0.0, 21.0 

243 i0, i1 = imod.prepare.common._selection_indices(src_x, xmin, xmax, 0) 

244 assert i0 == 0 

245 assert i1 == 2 

246 

247 xmin, xmax = 0.0, 40.0 

248 i0, i1 = imod.prepare.common._selection_indices(src_x, xmin, xmax, 0) 

249 assert i0 == 0 

250 assert i1 == 2 

251 

252 xmin, xmax = 0.0, 40.0 

253 i0, i1 = imod.prepare.common._selection_indices(src_x, xmin, xmax, 1) 

254 assert i0 == 0 

255 assert i1 == 3 

256 

257 xmin, xmax = 20.0, 40.0 

258 i0, i1 = imod.prepare.common._selection_indices(src_x, xmin, xmax, 1) 

259 assert i0 == 0 

260 assert i1 == 3 

261 

262 

263def test_slice_src(): 

264 # dx of 100.0 

265 # midpoints 

266 src_x = np.array([50.0, 150.0, 250.0, 350.0, 450.0]) 

267 # dx of 50.0 

268 like_x = np.array([75.0, 125.0, 175.0]) 

269 src = xr.DataArray(np.ones(src_x.size), {"x": src_x}, ("x",)) 

270 like = xr.DataArray(np.ones(like_x.size), {"x": like_x}, ("x",)) 

271 

272 actual = imod.prepare.common._slice_src(src, like, 0) 

273 expected = src.isel(x=slice(0, 2)) 

274 assert actual.equals(expected) 

275 

276 actual = imod.prepare.common._slice_src(src, like, 1) 

277 expected = src.isel(x=slice(0, 3)) 

278 assert actual.equals(expected) 

279 

280 like_x = np.array([125.0, 175.0, 225.0]) 

281 like = xr.DataArray(np.ones(like_x.size), {"x": like_x}, ("x",)) 

282 actual = imod.prepare.common._slice_src(src, like, 1) 

283 expected = src.isel(x=slice(0, 4)) 

284 assert actual.equals(expected) 

285 

286 

287def test_define_single_dim_slices(): 

288 # Simplest first, no chunk in dimension 

289 src_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) 

290 dst_x = np.array([0.0, 2.0, 4.0]) 

291 chunksizes = (4,) 

292 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

293 assert dst_slices == [slice(None, None, None)] 

294 

295 # Clean cuts 

296 chunksizes = (2, 2) 

297 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

298 assert dst_slices == [slice(0, 1, None), slice(1, 2, None)] 

299 

300 # Mixed cut 

301 src_x = np.arange(13.0) 

302 dst_x = np.arange(0.0, 15.0, 2.5) 

303 chunksizes = (4, 4, 4) 

304 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

305 assert dst_slices == [slice(0, 2, None), slice(2, 4, None), slice(4, 5)] 

306 

307 src_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) 

308 dst_x = np.array([0.0, 2.5, 5.0]) 

309 chunksizes = (3, 2) 

310 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

311 assert dst_slices == [slice(0, 2, None)] 

312 

313 # dst larger than src 

314 src_x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) 

315 dst_x = np.array([-1.0, 2.5, 6.0]) 

316 chunksizes = (3, 2) 

317 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

318 assert dst_slices == [slice(0, 2, None)] 

319 

320 src_x = np.arange(13.0) 

321 dst_x = np.arange(0.0, 15.0, 2.5) 

322 chunksizes = (3, 3, 3, 3) 

323 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

324 assert dst_slices == [ 

325 slice(0, 2, None), 

326 slice(2, 3, None), 

327 slice(3, 4, None), 

328 slice(4, 5, None), 

329 ] 

330 

331 src_x = np.arange(0.0, 1010.0, 10.0) 

332 dst_x = np.arange(0.0, 1025.0, 25.0) 

333 chunksizes = (10,) * 10 

334 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

335 assert len(dst_slices) == 10 

336 

337 src_x = np.arange(0.0, 400, 25.0) 

338 dst_x = np.arange(0.0, 390.0, 10.0) 

339 chunksizes = (3,) * 5 

340 dst_slices = imod.prepare.common._define_single_dim_slices(src_x, dst_x, chunksizes) 

341 assert len(dst_slices) == 5 

342 

343 

344def test_sel_chunks(): 

345 src_x = np.arange(5.0) + 0.5 

346 dst_x = np.arange(0.0, 6.0, 2.0) + 1.0 

347 src = xr.DataArray(np.ones(5), {"x": src_x}, ("x",)) 

348 src = src.chunk({"x": (2, 2, 1)}) 

349 like = xr.DataArray(np.ones(3), {"x": dst_x}, ("x",)) 

350 dst_slices, chunks_shape = imod.prepare.common._define_slices(src, like) 

351 assert len(dst_slices) == np.product(chunks_shape) 

352 

353 # 2D 

354 src = xr.DataArray(np.ones((5, 5)), {"y": src_x, "x": src_x}, ("y", "x")) 

355 src = src.chunk({"x": (2, 2, 1), "y": (2, 2, 1)}) 

356 like = xr.DataArray(np.ones((3, 3)), {"y": dst_x, "x": dst_x}, ("y", "x")) 

357 dst_slices, chunks_shape = imod.prepare.common._define_slices(src, like) 

358 assert len(dst_slices) == np.product(chunks_shape)