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

1 from typing import 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.iregridpackage import IRegridPackage

8 from imod.mf6.package import Package

9 from imod.mf6.utilities.regrid import RegridderType

10 from imod.mf6.validation import DisBottomSchema

11 from imod.mf6.write_context import WriteContext

12 from imod.schemata import (

13 AllValueSchema,

14 AnyValueSchema,

15 DimsSchema,

16 DTypeSchema,

17 IdentityNoDataSchema,

18 IndexesSchema,

19)

20

21

22 class VerticesDiscretization(Package, IRegridPackage):

23 """

24 Discretization by Vertices (DISV).

25

26 Parameters

27 ----------

28 top: array of floats (xu.UgridDataArray)

29 bottom: array of floats (xu.UgridDataArray)

30 idomain: array of integers (xu.UgridDataArray)

31 validate: {True, False}

32 Flag to indicate whether the package should be validated upon

33 initialization. This raises a ValidationError if package input is

34 provided in the wrong manner. Defaults to True.

35 """

36

37 _pkg_id = "disv"

38

39 _init_schemata = {

40 "top": [

41 DTypeSchema(np.floating),

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

43 IndexesSchema(),

44 ],

45 "bottom": [

46 DTypeSchema(np.floating),

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

48 IndexesSchema(),

49 ],

50 "idomain": [

51 DTypeSchema(np.integer),

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

53 IndexesSchema(),

54 ],

55 }

56 _write_schemata = {

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

58 "top": (

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

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

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

62 ),

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

64 }

65

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

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

68 _template = Package._initialize_template(_pkg_id)

69

70 _regrid_method = {

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

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

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

74 }

75

76 _skip_mask_arrays = ["bottom"]

77

78 @init_log_decorator()

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

80 dict_dataset = {

81 "idomain": idomain,

82 "top": top,

83 "bottom": bottom,

84 }

85 super().__init__(dict_dataset)

86 self._validate_init_schemata(validate)

87

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

89 disdirectory = directory / pkgname

90 d = {}

91 grid = self.dataset.ugrid.grid

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

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

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

95 facedim = grid.face_dimension

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

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

98

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

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

101 )

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

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

104 )

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

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

107 )

108 return self._template.render(d)

109

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

111 grid = self.dataset.ugrid.grid

112 df = pd.DataFrame(grid.node_coordinates)

113 df.index += 1

114 return df

115

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

117 grid = self.dataset.ugrid.grid

118 df = pd.DataFrame(grid.face_coordinates)

119 df.index += 1

120 # modflow requires clockwise; ugrid requires ccw

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

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

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

124 # Use extension array to write empty values

125 # Should be more efficient than mixed column?

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

127 values=column + 1,

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

129 )

130 return df

131

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

133 dir_for_render = write_context.get_formatted_write_directory()

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

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

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

137 f.write(content)

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

139

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

141 self._verts_dataframe().to_csv(

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

143 )

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

145

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

147 self._cell2d_dataframe().to_csv(

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

149 )

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

151 return

152

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

154 # Insert additional kwargs

155 kwargs["bottom"] = self["bottom"]

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

157

158 return errors

159

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

161 return self._regrid_method