Coverage for test_layerregrid.py: 100%
94 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
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
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
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
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)
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
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
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)
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
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 )
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)
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
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
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)
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)
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)