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

72 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 

4import pandas as pd 

5 

6from imod.logging import init_log_decorator 

7from imod.mf6.interfaces.iregridpackage import IRegridPackage 

8from imod.mf6.package import Package 

9from imod.mf6.utilities.regrid import RegridderType 

10from imod.mf6.validation import DisBottomSchema 

11from imod.mf6.write_context import WriteContext 

12from imod.schemata import ( 

13 AllValueSchema, 

14 AnyValueSchema, 

15 DimsSchema, 

16 DTypeSchema, 

17 IdentityNoDataSchema, 

18 IndexesSchema, 

19) 

20 

21 

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