Coverage for test_interpolate.py: 100%

1 import numpy as np

2 import pytest

3 import xarray as xr

4

5 import imod

6

7

8 def 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)

20

21

22 def 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)

33

34

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

47

48

49@pytest.mark.parametrize("chunksize", [1, 2])

50 def 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)

64

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)

69

70

71@pytest.mark.parametrize("chunksize", [1, 2])

72 def 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)

86

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)

91

92

93@pytest.mark.parametrize("chunksize", [1, 2])

94 def 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)

107

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)

112

113

114@pytest.mark.parametrize("chunksize", [1, 2])

115 def 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)

128

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)

133

134

135@pytest.mark.parametrize("chunksize", [1, 2])

136 def 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)

154

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)

159

160

161@pytest.mark.parametrize("chunksize", [1, 2])

162 def 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)

180

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)

185

186

187# TODO: add tests for checking nodata behaviour in 2D and 3D

188# TODO: maybe do parametrize stuff versus scipy linear interp?