Coverage for C:\src\imod-python\imod\mf6\dsp.py: 100%

27 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1from typing import Optional, Tuple 

2 

3import numpy as np 

4 

5from imod.logging import init_log_decorator 

6from imod.mf6.interfaces.iregridpackage import IRegridPackage 

7from imod.mf6.package import Package 

8from imod.mf6.utilities.regrid import RegridderType 

9from imod.mf6.validation import PKG_DIMS_SCHEMA 

10from imod.schemata import ( 

11 CompatibleSettingsSchema, 

12 DimsSchema, 

13 DTypeSchema, 

14 IdentityNoDataSchema, 

15 IndexesSchema, 

16) 

17 

18 

19class Dispersion(Package, IRegridPackage): 

20 """ 

21 Molecular Diffusion and Dispersion. 

22 

23 Parameters 

24 ---------- 

25 diffusion_coefficient: xr.DataArray 

26 effective molecular diffusion coefficient. (DIFFC) 

27 longitudinal_horizontal: xr.DataArray 

28 longitudinal dispersivity in horizontal direction. If flow is strictly 

29 horizontal, then this is the longitudinal dispersivity that will be 

30 used. If flow is not strictly horizontal or strictly vertical, then the 

31 longitudinal dispersivity is a function of both ALH and ALV. If 

32 mechanical dispersion is represented (by specifying any dispersivity 

33 values) then this array is required. (ALH) 

34 transverse_horizontal1: xr.DataArray 

35 transverse dispersivity in horizontal direction. This is the transverse 

36 dispersivity value for the second ellipsoid axis. If flow is strictly 

37 horizontal and directed in the x direction (along a row for a regular 

38 grid), then this value controls spreading in the y direction. 

39 If mechanical dispersion is represented (by specifying any dispersivity 

40 values) then this array is required. (ATH1) 

41 longitudinal_vertical: xr.DataArray, optional 

42 longitudinal dispersivity in vertical direction. If flow is strictly 

43 vertical, then this is the longitudinal dispsersivity value that will be 

44 used. If flow is not strictly horizontal or strictly vertical, then the 

45 longitudinal dispersivity is a function of both ALH and ALV. If this 

46 value is not specified and mechanical dispersion is represented, then 

47 this array is set equal to ALH. (ALV) 

48 transverse_horizontal2: xr.DataArray, optional 

49 transverse dispersivity in horizontal direction. This is the transverse 

50 dispersivity value for the third ellipsoid axis. If flow is strictly 

51 horizontal and directed in the x direction (along a row for a regular 

52 grid), then this value controls spreading in the z direction. If this 

53 value is not specified and mechanical dispersion is represented, then 

54 this array is set equal to ATH1. (ATH2) 

55 tranverse_vertical: xr.DataArray, optional 

56 transverse dispersivity when flow is in vertical direction. If flow is 

57 strictly vertical and directed in the z direction, then this value 

58 controls spreading in the x and y directions. If this value is not 

59 specified and mechanical dispersion is represented, then this array is 

60 set equal to ATH2. (ATV) 

61 xt3d_off: bool, optional 

62 deactivate the xt3d method and use the faster and less accurate 

63 approximation. (XT3D_OFF) 

64 xt3d_rhs: bool, optional 

65 add xt3d terms to right-hand side, when possible. This option uses less 

66 memory, but may require more iterations. (XT3D_RHS) 

67 validate: {True, False} 

68 Flag to indicate whether the package should be validated upon 

69 initialization. This raises a ValidationError if package input is 

70 provided in the wrong manner. Defaults to True. 

71 """ 

72 

73 _pkg_id = "dsp" 

74 _template = Package._initialize_template(_pkg_id) 

75 _grid_data = { 

76 "diffusion_coefficient": np.float64, 

77 "longitudinal_horizontal": np.float64, 

78 "transversal_horizontal1": np.float64, 

79 "longitudinal_vertical": np.float64, 

80 "transversal_horizontal2": np.float64, 

81 "transversal_vertical": np.float64, 

82 } 

83 _keyword_map = { 

84 "diffusion_coefficient": "diffc", 

85 "longitudinal_horizontal": "alh", 

86 "transversal_horizontal1": "ath1", 

87 "longitudinal_vertical": "alv", 

88 "transversal_horizontal2": "ath2", 

89 "transversal_vertical": "atv", 

90 } 

91 _init_schemata = { 

92 "diffusion_coefficient": [ 

93 DTypeSchema(np.floating), 

94 IndexesSchema(), 

95 PKG_DIMS_SCHEMA, 

96 ], 

97 "longitudinal_horizontal": [ 

98 DTypeSchema(np.floating), 

99 IndexesSchema(), 

100 PKG_DIMS_SCHEMA, 

101 ], 

102 "transversal_horizontal1": [ 

103 DTypeSchema(np.floating), 

104 IndexesSchema(), 

105 PKG_DIMS_SCHEMA, 

106 ], 

107 "longitudinal_vertical": [ 

108 DTypeSchema(np.floating), 

109 IndexesSchema(), 

110 PKG_DIMS_SCHEMA, 

111 ], 

112 "transversal_horizontal2": [ 

113 DTypeSchema(np.floating), 

114 IndexesSchema(), 

115 PKG_DIMS_SCHEMA, 

116 ], 

117 "transversal_vertical": [ 

118 DTypeSchema(np.floating), 

119 IndexesSchema(), 

120 PKG_DIMS_SCHEMA, 

121 ], 

122 "xt3d_off": [DTypeSchema(np.bool_), DimsSchema()], 

123 "xt3d_rhs": [ 

124 DTypeSchema(np.bool_), 

125 DimsSchema(), 

126 CompatibleSettingsSchema(other="xt3d_off", other_value=False), 

127 ], 

128 } 

129 

130 _write_schemata = { 

131 "diffusion_coefficient": ( 

132 IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)), 

133 ), 

134 "longitudinal_horizontal": ( 

135 IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)), 

136 ), 

137 "transversal_horizontal1": ( 

138 IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)), 

139 ), 

140 "longitudinal_vertical": ( 

141 IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)), 

142 ), 

143 "transversal_horizontal2": ( 

144 IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)), 

145 ), 

146 "transversal_vertical": ( 

147 IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)), 

148 ), 

149 } 

150 

151 _regrid_method = { 

152 "diffusion_coefficient": (RegridderType.OVERLAP, "mean"), 

153 "longitudinal_horizontal": (RegridderType.OVERLAP, "mean"), 

154 "transversal_horizontal1": ( 

155 RegridderType.OVERLAP, 

156 "mean", 

157 ), 

158 "longitudinal_vertical": ( 

159 RegridderType.OVERLAP, 

160 "mean", 

161 ), 

162 "transversal_horizontal2": (RegridderType.OVERLAP, "mean"), 

163 "transversal_vertical": (RegridderType.OVERLAP, "mean"), 

164 } 

165 

166 @init_log_decorator() 

167 def __init__( 

168 self, 

169 diffusion_coefficient, 

170 longitudinal_horizontal, 

171 transversal_horizontal1, 

172 longitudinal_vertical=None, 

173 transversal_horizontal2=None, 

174 transversal_vertical=None, 

175 xt3d_off=False, 

176 xt3d_rhs=False, 

177 validate: bool = True, 

178 ): 

179 dict_dataset = { 

180 "xt3d_off": xt3d_off, 

181 "xt3d_rhs": xt3d_rhs, 

182 "diffusion_coefficient": diffusion_coefficient, 

183 "longitudinal_horizontal": longitudinal_horizontal, 

184 "transversal_horizontal1": transversal_horizontal1, 

185 "longitudinal_vertical": longitudinal_vertical, 

186 "transversal_horizontal2": transversal_horizontal2, 

187 "transversal_vertical": transversal_vertical, 

188 } 

189 super().__init__(dict_dataset) 

190 self._validate_init_schemata(validate) 

191 

192 def _validate(self, schemata, **kwargs): 

193 # Insert additional kwargs 

194 kwargs["xt3d_off"] = self["xt3d_off"] 

195 errors = super()._validate(schemata, **kwargs) 

196 

197 return errors 

198 

199 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: 

200 return self._regrid_method