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

75 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 11:25 +0200

1from typing import List, Optional, Tuple 

2 

3import numpy as np 

4import pandas as pd 

5 

6from imod.logging import init_log_decorator 

7from imod.mf6.interfaces.imaskingsettings import IMaskingSettings 

8from imod.mf6.interfaces.iregridpackage import IRegridPackage 

9from imod.mf6.package import Package 

10from imod.mf6.utilities.regrid import RegridderType 

11from imod.mf6.validation import DisBottomSchema 

12from imod.mf6.write_context import WriteContext 

13from imod.schemata import ( 

14 AllValueSchema, 

15 AnyValueSchema, 

16 DimsSchema, 

17 DTypeSchema, 

18 IdentityNoDataSchema, 

19 IndexesSchema, 

20) 

21 

22 

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