Coverage for test_voxelize.py: 100%
82 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_voxelize__mean_1(test_da):
22 voxelizer = imod.prepare.Voxelizer(method="mean")
23 source = test_da
24 z = [-0.5, -1.5, -2.5, -3.5]
25 like = xr.DataArray(np.arange(4), {"z": z}, ["z"])
27 top = xr.full_like(source, 0.0)
28 top.data[1, :, :] = -2.0
29 bottom = xr.full_like(source, -2.0)
30 bottom.data[1, :, :] = -4.0
32 actual = voxelizer.voxelize(source, top, bottom, like)
33 coords = {"z": z, "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]}
34 coords["z"] = z
35 dims = ("z", "y", "x")
36 expected = xr.DataArray(
37 [
38 [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]],
39 [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]],
40 [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]],
41 [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]],
42 ],
43 coords,
44 dims,
45 )
46 assert actual.identical(expected)
49def test_voxelize__mean_2(test_da):
50 voxelizer = imod.prepare.Voxelizer(method="mean")
51 source = test_da
52 z = [-0.5, -2.0, -3.5]
53 dz = [-1.0, -2.0, -1.0]
54 like = xr.DataArray(np.arange(3), {"z": z, "dz": ("z", dz)}, ["z"])
56 top = xr.full_like(source, 0.0)
57 top.data[1, :, :] = -2.0
58 bottom = xr.full_like(source, -2.0)
59 bottom.data[1, :, :] = -4.0
61 actual = voxelizer.voxelize(source, top, bottom, like)
62 coords = {"z": z, "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]}
63 coords["z"] = z
64 dims = ("z", "y", "x")
65 expected = xr.DataArray(
66 [
67 [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]],
68 [[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],
69 [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]],
70 ],
71 coords,
72 dims,
73 )
74 assert actual.identical(expected)
77def test_voxelize__max_overlap_1(test_da):
78 voxelizer = imod.prepare.Voxelizer(method="max_overlap")
79 source = test_da
80 z = [-0.5, -2.5]
81 dz = [-1.0, -3.0]
82 like = xr.DataArray(np.arange(2), {"z": z, "dz": ("z", dz)}, ["z"])
84 top = xr.full_like(source, 0.0)
85 top.data[1, :, :] = -2.0
86 bottom = xr.full_like(source, -2.0)
87 bottom.data[1, :, :] = -4.0
89 actual = voxelizer.voxelize(source, top, bottom, like)
90 coords = {"z": z, "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]}
91 coords["z"] = z
92 dims = ("z", "y", "x")
93 expected = xr.DataArray(
94 [[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]],
95 coords,
96 dims,
97 )
98 assert actual.identical(expected)
101def test_voxelize__max_overlap_2(test_da):
102 voxelizer = imod.prepare.Voxelizer(method="max_overlap")
103 source = test_da
104 z = [-0.5, -2.5]
105 dz = [-1.0, -3.0]
106 like = xr.DataArray(np.arange(2), {"z": z, "dz": ("z", dz)}, ["z"])
108 top = xr.full_like(source, -2.0)
109 top.data[0, :, :] = 0.0
110 top.data[0, 0, 0] = np.nan
111 top.data[0, 1, 1] = np.nan
112 bottom = xr.full_like(source, -4.0)
113 bottom.data[0, :, :] = -2.0
114 bottom.data[0, 0, 0] = np.nan
115 bottom.data[0, 1, 1] = np.nan
117 actual = voxelizer.voxelize(source, top, bottom, like)
118 coords = {"z": z, "y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]}
119 coords["z"] = z
120 dims = ("z", "y", "x")
121 expected = xr.DataArray(
122 [
123 [[np.nan, 1.0, 2.0], [3.0, np.nan, 5.0]],
124 [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]],
125 ],
126 coords,
127 dims,
128 )
129 assert actual.identical(expected)