Coverage for test_interpolate.py: 100%

146 statements  

« 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 

4 

5import imod 

6 

7 

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) 

20 

21 

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) 

33 

34 

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) 

47 

48 

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) 

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]) 

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) 

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]) 

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) 

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]) 

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) 

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]) 

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) 

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]) 

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) 

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?