Coverage for test_interpolate.py: 100%
146 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1import numpy as np
2import pytest
3import xarray as xr
5import imod
8def test_interpolate_1d():
9 data = np.arange(5.0)
10 x = np.arange(5.0)
11 dst_x = np.array([2.5, 3.5])
12 dims = ("x",)
13 source = xr.DataArray(data, {"x": x}, dims)
14 like = xr.DataArray([np.nan, np.nan], {"x": dst_x}, dims)
15 interpolator_1d = imod.prepare.Regridder(method="multilinear")
16 actual = interpolator_1d.regrid(source, like)
17 expected = like.copy(data=[2.5, 3.5])
18 # Note: xr.identical() fails, as the regridder as a dx coord for cellsizes.
19 assert np.allclose(actual.values, expected.values, equal_nan=True)
22def test_interpolate_1d__reversed():
23 data = np.arange(5.0)[::-1]
24 x = np.arange(5.0)[::-1]
25 dst_x = np.array([2.5, 3.5])[::-1]
26 dims = ("x",)
27 source = xr.DataArray(data, {"x": x}, dims)
28 like = xr.DataArray([np.nan, np.nan], {"x": dst_x}, dims)
29 interpolator_1d = imod.prepare.Regridder(method="multilinear")
30 actual = interpolator_1d.regrid(source, like)
31 expected = like.copy(data=[3.5, 2.5])
32 assert np.allclose(actual.values, expected.values, equal_nan=True)
35def test_interpolate_1d__beyond_egdes():
36 data = [0.0, 1.0]
37 x = [0.5, 1.5]
38 dst_x = [-1.25, -0.75, -0.25, 0.25, 0.75, 1.25, 1.75, 2.25]
39 dims = ("x",)
40 source = xr.DataArray(data, {"x": x}, dims)
41 like_data = np.full(len(dst_x), np.nan)
42 like = xr.DataArray(like_data, {"x": dst_x}, dims)
43 interpolator_1d = imod.prepare.Regridder(method="multilinear")
44 actual = interpolator_1d.regrid(source, like)
45 expected = like.copy(data=[np.nan] * 3 + [0.0, 0.25, 0.75, 1.0, np.nan])
46 assert np.allclose(actual.values, expected.values, equal_nan=True)
49@pytest.mark.parametrize("chunksize", [1, 2])
50def test_interpolate_2d(chunksize):
51 data = np.array([[0.0, 0.0], [1.0, 1.0]])
52 x = [0.5, 1.5]
53 dst_x = [0.75, 1.25]
54 y = [0.0, 1.0]
55 dst_y = [0.25, 0.75]
56 dims = ("y", "x")
57 source = xr.DataArray(data, {"x": x, "y": y}, dims)
58 dst_data = np.full_like(data, np.nan)
59 like = xr.DataArray(dst_data, {"y": dst_y, "x": dst_x}, dims)
60 interpolator_2d = imod.prepare.Regridder(method="multilinear")
61 actual = interpolator_2d.regrid(source, like)
62 expected = like.copy(data=[[0.25, 0.25], [0.75, 0.75]])
63 assert np.allclose(actual.values, expected.values, equal_nan=True)
65 # Now with chunks
66 source = source.chunk({"x": chunksize, "y": chunksize})
67 actual = interpolator_2d.regrid(source, like)
68 assert np.allclose(actual.values, expected.values, equal_nan=True)
71@pytest.mark.parametrize("chunksize", [1, 2])
72def test_interpolate_2d__reversed_y(chunksize):
73 data = np.array([[0.0, 0.0], [1.0, 1.0]])
74 x = [0.5, 1.5]
75 dst_x = [0.75, 1.25]
76 y = list(reversed([0.0, 1.0]))
77 dst_y = list(reversed([0.25, 0.75]))
78 dims = ("y", "x")
79 source = xr.DataArray(data, {"x": x, "y": y}, dims)
80 dst_data = np.full_like(data, np.nan)
81 like = xr.DataArray(dst_data, {"y": dst_y, "x": dst_x}, dims)
82 interpolator_2d = imod.prepare.Regridder(method="multilinear")
83 actual = interpolator_2d.regrid(source, like)
84 expected = like.copy(data=[[0.25, 0.25], [0.75, 0.75]])
85 assert np.allclose(actual.values, expected.values, equal_nan=True)
87 # Now with chunks
88 source = source.chunk({"x": chunksize})
89 actual = interpolator_2d.regrid(source, like)
90 assert np.allclose(actual.values, expected.values, equal_nan=True)
93@pytest.mark.parametrize("chunksize", [1, 2])
94def test_interpolate_1d__nan_withstartingedge(chunksize):
95 data = [np.nan, 1.0, 1.0]
96 x = [0.5, 1.5, 2.5]
97 # 1.0 is on the starting edge, and should get a value
98 dst_x = [0.25, 0.50, 0.75, 1.0]
99 dims = ("x",)
100 source = xr.DataArray(data, {"x": x}, dims)
101 like_data = np.full(len(dst_x), np.nan)
102 like = xr.DataArray(like_data, {"x": dst_x}, dims)
103 interpolator_1d = imod.prepare.Regridder(method="multilinear")
104 actual = interpolator_1d.regrid(source, like)
105 expected = like.copy(data=[np.nan] * 3 + [1.0])
106 assert np.allclose(actual.values, expected.values, equal_nan=True)
108 # Now with chunks
109 source = source.chunk({"x": chunksize})
110 actual = interpolator_1d.regrid(source, like)
111 assert np.allclose(actual.values, expected.values, equal_nan=True)
114@pytest.mark.parametrize("chunksize", [1, 2])
115def test_interpolate_1d__nan_withendingedge(chunksize):
116 data = [1.0, 1.0, np.nan]
117 x = [0.5, 1.5, 2.5]
118 # 3.0 is on the starting edge, and should get a value
119 dst_x = [1.0, 1.5, 2.0, 2.5]
120 dims = ("x",)
121 source = xr.DataArray(data, {"x": x}, dims)
122 like_data = np.full(len(dst_x), np.nan)
123 like = xr.DataArray(like_data, {"x": dst_x}, dims)
124 interpolator_1d = imod.prepare.Regridder(method="multilinear")
125 actual = interpolator_1d.regrid(source, like)
126 expected = like.copy(data=[1.0] * 2 + [np.nan] * 2)
127 assert np.allclose(actual.values, expected.values, equal_nan=True)
129 # Now with chunks
130 source = source.chunk({"x": chunksize})
131 actual = interpolator_1d.regrid(source, like)
132 assert np.allclose(actual.values, expected.values, equal_nan=True)
135@pytest.mark.parametrize("chunksize", [1, 2])
136def test_interpolate_2d__over_z(chunksize):
137 data = np.array([[[0.0, 0.0], [1.0, 1.0]], [[1.0, 1.0], [2.0, 2.0]]])
138 x = [0.5, 1.5]
139 dst_x = [0.75, 1.25]
140 y = [0.0, 1.0]
141 dst_y = [0.25, 0.75]
142 z = [0.0, 1.0]
143 dst_z = [0.0, 1.0]
144 dims = ("z", "y", "x")
145 source = xr.DataArray(data, {"x": x, "y": y, "z": z}, dims)
146 dst_data = np.full_like(data, np.nan)
147 like = xr.DataArray(dst_data, {"y": dst_y, "x": dst_x, "z": dst_z}, dims)
148 interpolator_2d = imod.prepare.Regridder(method="multilinear")
149 actual = interpolator_2d.regrid(source, like)
150 expected = like.copy(
151 data=[[[0.25, 0.25], [0.75, 0.75]], [[1.25, 1.25], [1.75, 1.75]]]
152 )
153 assert np.allclose(actual.values, expected.values, equal_nan=True)
155 # Now with chunks
156 source = source.chunk({"x": chunksize, "y": chunksize})
157 actual = interpolator_2d.regrid(source, like)
158 assert np.allclose(actual.values, expected.values, equal_nan=True)
161@pytest.mark.parametrize("chunksize", [1, 2])
162def test_interpolate_3d__over_xyz(chunksize):
163 data = np.array([[[0.0, 0.0], [1.0, 1.0]], [[1.0, 1.0], [2.0, 2.0]]])
164 x = [0.5, 1.5]
165 dst_x = [0.75, 1.25]
166 y = [0.0, 1.0]
167 dst_y = [0.25, 0.75]
168 z = [0.0, 1.0]
169 dst_z = [0.05, 0.95]
170 dims = ("z", "y", "x")
171 source = xr.DataArray(data, {"x": x, "y": y, "z": z}, dims)
172 dst_data = np.full_like(data, np.nan)
173 like = xr.DataArray(dst_data, {"y": dst_y, "x": dst_x, "z": dst_z}, dims)
174 interpolator_3d = imod.prepare.Regridder(method="multilinear")
175 actual = interpolator_3d.regrid(source, like)
176 expected = like.copy(data=[[[0.3, 0.3], [0.8, 0.8]], [[1.2, 1.2], [1.7, 1.7]]])
177 # This fails for some reason, different ordering of coords (not dims!):
178 # assert actual.identical(expected)
179 assert np.allclose(actual.values, expected.values, equal_nan=True)
181 # Now with chunks
182 source = source.chunk({"x": chunksize, "y": chunksize, "z": chunksize})
183 actual = interpolator_3d.regrid(source, like)
184 assert np.allclose(actual.values, expected.values, equal_nan=True)
187# TODO: add tests for checking nodata behaviour in 2D and 3D
188# TODO: maybe do parametrize stuff versus scipy linear interp?