Coverage for test_voxelize.py: 100%

82 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_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"]) 

26 

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 

31 

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) 

47 

48 

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

55 

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 

60 

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) 

75 

76 

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

83 

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 

88 

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) 

99 

100 

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

107 

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 

116 

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)