Coverage for test_common.py: 100%
241 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-16 11:41 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-16 11:41 +0200
1import warnings
3import numba
4import numpy as np
5import xarray as xr
7import imod
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
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)]
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)]
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)]
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)]
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)]
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)]
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 )
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]]))
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 )
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 )
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]]))
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]]))
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
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 )
140 values[1] = 3.0
141 assert np.allclose(imod.prepare.common.mode(values, weights), 3.0)
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))
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)
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))
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)
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
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)
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)
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)
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)
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)
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
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
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
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
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
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",))
272 actual = imod.prepare.common._slice_src(src, like, 0)
273 expected = src.isel(x=slice(0, 2))
274 assert actual.equals(expected)
276 actual = imod.prepare.common._slice_src(src, like, 1)
277 expected = src.isel(x=slice(0, 3))
278 assert actual.equals(expected)
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)
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)]
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)]
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)]
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)]
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)]
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 ]
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
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
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)
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)