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

1 from typing import List, Optional, Tuple

2

3 import numpy as np

4 import pandas as pd

5

6 from imod.logging import init_log_decorator

7 from imod.mf6.interfaces.imaskingsettings import IMaskingSettings

8 from imod.mf6.interfaces.iregridpackage import IRegridPackage

9 from imod.mf6.package import Package

10 from imod.mf6.utilities.regrid import RegridderType

11 from imod.mf6.validation import DisBottomSchema

12 from imod.mf6.write_context import WriteContext

13 from imod.schemata import (

14 AllValueSchema,

15 AnyValueSchema,

16 DimsSchema,

17 DTypeSchema,

18 IdentityNoDataSchema,

19 IndexesSchema,

20)

21

22

23 class VerticesDiscretization(Package, IRegridPackage, IMaskingSettings):

24 """

25 Discretization by Vertices (DISV).

26

27 Parameters

28 ----------

29 top: array of floats (xu.UgridDataArray)

30 bottom: array of floats (xu.UgridDataArray)

31 idomain: array of integers (xu.UgridDataArray)

32 validate: {True, False}

33 Flag to indicate whether the package should be validated upon

34 initialization. This raises a ValidationError if package input is

35 provided in the wrong manner. Defaults to True.

36 """

37

38 _pkg_id = "disv"

39

40 _init_schemata = {

41 "top": [

42 DTypeSchema(np.floating),

43 DimsSchema("{face_dim}") | DimsSchema(),

44 IndexesSchema(),

45 ],

46 "bottom": [

47 DTypeSchema(np.floating),

48 DimsSchema("layer", "{face_dim}") | DimsSchema("layer"),

49 IndexesSchema(),

50 ],

51 "idomain": [

52 DTypeSchema(np.integer),

53 DimsSchema("layer", "{face_dim}"),

54 IndexesSchema(),

55 ],

56 }

57 _write_schemata = {

58 "idomain": (AnyValueSchema(">", 0),),

59 "top": (

60 AllValueSchema(">", "bottom", ignore=("idomain", "==", -1)),

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

62 # No need to check coords: dataset ensures they align with idomain.

63 ),

64 "bottom": (DisBottomSchema(other="idomain", is_other_notnull=(">", 0)),),

65 }

66

67 _grid_data = {"top": np.float64, "bottom": np.float64, "idomain": np.int32}

68 _keyword_map = {"bottom": "botm"}

69 _template = Package._initialize_template(_pkg_id)

70

71 _regrid_method = {

72 "top": (RegridderType.OVERLAP, "mean"),

73 "bottom": (RegridderType.OVERLAP, "mean"),

74 "idomain": (RegridderType.OVERLAP, "mode"),

75 }

76

77 @property

78 def skip_variables(self) -> List[str]:

79 return ["bottom"]

80

81 @init_log_decorator()

82 def __init__(self, top, bottom, idomain, validate: bool = True):

83 dict_dataset = {

84 "idomain": idomain,

85 "top": top,

86 "bottom": bottom,

87 }

88 super().__init__(dict_dataset)

89 self._validate_init_schemata(validate)

90

91 def render(self, directory, pkgname, binary):

92 disdirectory = directory / pkgname

93 d = {}

94 grid = self.dataset.ugrid.grid

95 d["xorigin"] = grid.node_x.min()

96 d["yorigin"] = grid.node_y.min()

97 d["nlay"] = self.dataset["idomain"].coords["layer"].size

98 facedim = grid.face_dimension

99 d["ncpl"] = self.dataset["idomain"].coords[facedim].size

100 d["nvert"] = grid.node_x.size

101

102 _, d["top"] = self._compose_values(

103 self.dataset["top"], disdirectory, "top", binary=binary

104 )

105 d["botm_layered"], d["botm"] = self._compose_values(

106 self["bottom"], disdirectory, "botm", binary=binary

107 )

108 d["idomain_layered"], d["idomain"] = self._compose_values(

109 self["idomain"], disdirectory, "idomain", binary=binary

110 )

111 return self._template.render(d)

112

113 def _verts_dataframe(self) -> pd.DataFrame:

114 grid = self.dataset.ugrid.grid

115 df = pd.DataFrame(grid.node_coordinates)

116 df.index += 1

117 return df

118

119 def _cell2d_dataframe(self) -> pd.DataFrame:

120 grid = self.dataset.ugrid.grid

121 df = pd.DataFrame(grid.face_coordinates)

122 df.index += 1

123 # modflow requires clockwise; ugrid requires ccw

124 face_nodes = grid.face_node_connectivity[:, ::-1]

125 df[2] = (face_nodes != grid.fill_value).sum(axis=1)

126 for i, column in enumerate(face_nodes.T):

127 # Use extension array to write empty values

128 # Should be more efficient than mixed column?

129 df[3 + i] = pd.arrays.IntegerArray(

130 values=column + 1,

131 mask=(column == grid.fill_value),

132 )

133 return df

134

135 def write_blockfile(self, pkgname, globaltimes, write_context: WriteContext):

136 dir_for_render = write_context.get_formatted_write_directory()

137 content = self.render(dir_for_render, pkgname, write_context.use_binary)

138 filename = write_context.write_directory / f"{pkgname}.{self._pkg_id}"

139 with open(filename, "w") as f:

140 f.write(content)

141 f.write("\n\n")

142

143 f.write("begin vertices\n")

144 self._verts_dataframe().to_csv(

145 f, header=False, sep=" ", lineterminator="\n"

146 )

147 f.write("end vertices\n\n")

148

149 f.write("begin cell2d\n")

150 self._cell2d_dataframe().to_csv(

151 f, header=False, sep=" ", lineterminator="\n"

152 )

153 f.write("end cell2d\n")

154 return

155

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

157 # Insert additional kwargs

158 kwargs["bottom"] = self["bottom"]

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

160

161 return errors

162

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

164 return self._regrid_method