Coverage for test_layerregrid.py: 100%

94 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 

8@pytest.fixture(scope="module") 

9def test_da(): 

10 nlayer = 2 

11 nrow = 2 

12 ncol = 3 

13 shape = (nlayer, nrow, ncol) 

14 dims = ("layer", "y", "x") 

15 coords = {"layer": [1, 2], "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]} 

16 data = np.arange(np.product(shape), dtype=np.float64).reshape(shape) 

17 source = xr.DataArray(data, coords, dims) 

18 return source 

19 

20 

21def test_layerregridder__mean_1(test_da): 

22 layerregridder = imod.prepare.LayerRegridder(method="mean") 

23 source = test_da 

24 top_src = xr.full_like(source, 0.0) 

25 top_src.data[1, :, :] = -2.0 

26 bottom_src = xr.full_like(source, -2.0) 

27 bottom_src.data[1, :, :] = -4.0 

28 

29 top_dst = xr.full_like(source, 0.0) 

30 top_dst.data[1, :, :] = -3.0 

31 bottom_dst = xr.full_like(source, -3.0) 

32 bottom_dst.data[1, :, :] = -4.0 

33 

34 actual = layerregridder.regrid(source, top_src, bottom_src, top_dst, bottom_dst) 

35 coords = {"layer": [1, 2], "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]} 

36 dims = ("layer", "y", "x") 

37 expected = xr.DataArray( 

38 [[[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]], [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]], 

39 coords, 

40 dims, 

41 ) 

42 assert actual.identical(expected) 

43 

44 

45def test_layerregridder__mean_2(test_da): 

46 layerregridder = imod.prepare.LayerRegridder(method="mean") 

47 source = test_da 

48 top_src = xr.full_like(source, 0.0) 

49 top_src.data[1, :, :] = -2.0 

50 bottom_src = xr.full_like(source, -2.0) 

51 bottom_src.data[1, :, :] = -4.0 

52 

53 coords = {"layer": [1, 2, 3], "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]} 

54 dims = ("layer", "y", "x") 

55 top_dst = xr.DataArray( 

56 [ 

57 [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], 

58 [[-1.0, -1.0, -1.0], [-1.0, -1.0, -1.0]], 

59 [[-3, -3, -3], [-3, -3, -3]], 

60 ], 

61 coords, 

62 dims, 

63 ) 

64 bottom_dst = xr.full_like(top_dst, -1.0) 

65 bottom_dst.data[1, :, :] = -3.0 

66 bottom_dst.data[2, :, :] = -4.0 

67 

68 actual = layerregridder.regrid(source, top_src, bottom_src, top_dst, bottom_dst) 

69 coords = {"layer": [1, 2, 3], "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]} 

70 dims = ("layer", "y", "x") 

71 expected = xr.DataArray( 

72 [ 

73 [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], 

74 [[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]], 

75 [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], 

76 ], 

77 coords, 

78 dims, 

79 ) 

80 assert actual.identical(expected) 

81 

82 

83def test_layerregridder_dst_larger_src(test_da): 

84 layerregridder = imod.prepare.LayerRegridder(method="mean") 

85 source = test_da 

86 top_src = xr.full_like(source, 0.0) 

87 top_src.data[1, :, :] = -2.0 

88 bottom_src = xr.full_like(source, -2.0) 

89 bottom_src.data[1, :, :] = -4.0 

90 

91 coords = {"layer": [1, 2, 3], "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]} 

92 dims = ("layer", "y", "x") 

93 top_dst = xr.DataArray( 

94 [ 

95 [[0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], 

96 [[-1.0, -1.0, -1.0], [-1.0, -1.0, -1.0]], 

97 [[-3.0, -4.0, -3.0], [-3.0, -4.0, -3.0]], 

98 ], 

99 coords, 

100 dims, 

101 ) 

102 bottom_dst = xr.DataArray( 

103 [ 

104 [[-1.0, -1.0, -1.0], [-1.0, -1.0, -1.0]], 

105 [[-3.0, -4.0, -3.0], [-3.0, -4.0, -3.0]], 

106 [[-5.0, -5.0, -3], [-5.0, -5.0, -3]], 

107 ], 

108 coords, 

109 dims, 

110 ) 

111 

112 actual = layerregridder.regrid(source, top_src, bottom_src, top_dst, bottom_dst) 

113 coords = {"layer": [1, 2, 3], "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]} 

114 dims = ("layer", "y", "x") 

115 expected = xr.DataArray( 

116 [ 

117 [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], 

118 [[3.0, 5.0, 5.0], [6.0, 8.0, 8.0]], 

119 [[6.0, np.nan, np.nan], [9.0, np.nan, np.nan]], 

120 ], 

121 coords, 

122 dims, 

123 ) 

124 assert actual.identical(expected) 

125 

126 

127def test_layerregridder__mode(test_da): 

128 layerregridder = imod.prepare.LayerRegridder(method="mode") 

129 source = test_da 

130 top_src = xr.full_like(source, 0.0) 

131 top_src.data[1, :, :] = -2.0 

132 bottom_src = xr.full_like(source, -2.0) 

133 bottom_src.data[1, :, :] = -4.0 

134 

135 top_dst = xr.full_like(source, 0.0) 

136 top_dst.data[1, :, :] = -3.0 

137 bottom_dst = xr.full_like(source, -3.0) 

138 bottom_dst.data[1, :, :] = -4.0 

139 

140 actual = layerregridder.regrid(source, top_src, bottom_src, top_dst, bottom_dst) 

141 coords = {"layer": [1, 2], "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]} 

142 dims = ("layer", "y", "x") 

143 expected = xr.DataArray( 

144 [[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]], 

145 coords, 

146 dims, 

147 ) 

148 assert actual.identical(expected) 

149 

150 

151def test_layerregridder_topbot_nan(): 

152 layerregridder = imod.prepare.LayerRegridder(method="mean") 

153 coords = {"x": [0.5], "y": [0.5], "layer": [1, 2, 3, 4]} 

154 dims = ("layer", "y", "x") 

155 shape = (4, 1, 1) 

156 source = xr.DataArray(np.ones(shape), coords, dims) 

157 top_src = xr.DataArray( 

158 np.array([np.nan, np.nan, 2.0, 1.0]).reshape(shape), coords, dims 

159 ) 

160 bottom_src = xr.DataArray( 

161 np.array([np.nan, np.nan, 1.0, 0.0]).reshape(shape), coords, dims 

162 ) 

163 top_dst = xr.DataArray( 

164 np.array([4.0, 3.0, 2.0, np.nan]).reshape(shape), coords, dims 

165 ) 

166 bottom_dst = xr.DataArray( 

167 np.array([np.nan, 2.0, 1.0, np.nan]).reshape(shape), coords, dims 

168 ) 

169 actual = layerregridder.regrid(source, top_src, bottom_src, top_dst, bottom_dst) 

170 

171 expected = xr.DataArray( 

172 np.array([np.nan, np.nan, 1.0, np.nan]).reshape(shape), coords, dims 

173 ) 

174 assert actual.identical(expected)