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

1 from typing import Optional, Tuple

2

3 import numpy as np

4

5 from imod.logging import init_log_decorator

6 from imod.mf6.interfaces.iregridpackage import IRegridPackage

7 from imod.mf6.package import Package

8 from imod.mf6.utilities.regrid import RegridderType

9 from imod.mf6.validation import PKG_DIMS_SCHEMA

10 from imod.schemata import (

11 CompatibleSettingsSchema,

12 DimsSchema,

13 DTypeSchema,

14 IdentityNoDataSchema,

15 IndexesSchema,

16)

17

18

19 class 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