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
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1from typing import Optional, Tuple
3import numpy as np
4import pandas as pd
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)
22class VerticesDiscretization(Package, IRegridPackage):
23 """
24 Discretization by Vertices (DISV).
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 """
37 _pkg_id = "disv"
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 }
66 _grid_data = {"top": np.float64, "bottom": np.float64, "idomain": np.int32}
67 _keyword_map = {"bottom": "botm"}
68 _template = Package._initialize_template(_pkg_id)
70 _regrid_method = {
71 "top": (RegridderType.OVERLAP, "mean"),
72 "bottom": (RegridderType.OVERLAP, "mean"),
73 "idomain": (RegridderType.OVERLAP, "mode"),
74 }
76 _skip_mask_arrays = ["bottom"]
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)
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
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)
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
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
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")
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")
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
153 def _validate(self, schemata, **kwargs):
154 # Insert additional kwargs
155 kwargs["bottom"] = self["bottom"]
156 errors = super()._validate(schemata, **kwargs)
158 return errors
160 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]:
161 return self._regrid_method