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
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-16 11:25 +0200
1from typing import List, Optional, Tuple
3import numpy as np
4import pandas as pd
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)
23class VerticesDiscretization(Package, IRegridPackage, IMaskingSettings):
24 """
25 Discretization by Vertices (DISV).
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 """
38 _pkg_id = "disv"
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 }
67 _grid_data = {"top": np.float64, "bottom": np.float64, "idomain": np.int32}
68 _keyword_map = {"bottom": "botm"}
69 _template = Package._initialize_template(_pkg_id)
71 _regrid_method = {
72 "top": (RegridderType.OVERLAP, "mean"),
73 "bottom": (RegridderType.OVERLAP, "mean"),
74 "idomain": (RegridderType.OVERLAP, "mode"),
75 }
77 @property
78 def skip_variables(self) -> List[str]:
79 return ["bottom"]
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)
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
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)
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
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
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")
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")
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
156 def _validate(self, schemata, **kwargs):
157 # Insert additional kwargs
158 kwargs["bottom"] = self["bottom"]
159 errors = super()._validate(schemata, **kwargs)
161 return errors
163 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]:
164 return self._regrid_method