Coverage for C:\src\imod-python\imod\mf6\buy.py: 91%
46 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, Sequence
3import numpy as np
4import xarray as xr
6from imod.logging import init_log_decorator
7from imod.mf6.package import Package
8from imod.schemata import DTypeSchema
11def assign_index(arg):
12 if isinstance(arg, xr.DataArray):
13 arg = arg.values
14 elif not isinstance(arg, (np.ndarray, list, tuple)):
15 raise TypeError("should be a tuple, list, or numpy array")
17 arr = np.array(arg)
18 if arr.ndim != 1:
19 raise ValueError("should be 1D")
21 return xr.DataArray(arr, dims=("index",))
24class Buoyancy(Package):
25 """
26 Buoyancy package. This package must be included when performing variable
27 density simulation.
29 Note that ``reference_density`` is a single value, but
30 ``density_concentration_slope``, ``reference_concentration`` and
31 ``modelname`` require an entry for every active species. Please refer to the
32 examples.
34 Parameters
35 ----------
36 reference_density: float,
37 fluid reference density used in the equation of state.
38 density_concentration_slope: sequence of floats
39 Slope of the (linear) density concentration line used in the density
40 equation of state.
41 reference_concentration: sequence of floats
42 Reference concentration used in the density equation of state.
43 modelname: sequence of strings,
44 Name of the GroundwaterTransport (GWT) model used for the
45 concentrations.
46 species: sequence of str,
47 Name of the species used to calculate a density value.
48 hhformulation_rhs: bool, optional.
49 use the variable-density hydraulic head formulation and add off-diagonal
50 terms to the right-hand. This option will prevent the BUY Package from
51 adding asymmetric terms to the flow matrix. Default value is ``False``.
52 densityfile:
53 name of the binary output file to write density information. The density
54 file has the same format as the head file. Density values will be
55 written to the density file whenever heads are written to the binary
56 head file. The settings for controlling head output are contained in the
57 Output Control option.
58 validate: {True, False}
59 Flag to indicate whether the package should be validated upon
60 initialization. This raises a ValidationError if package input is
61 provided in the wrong manner. Defaults to True.
63 Examples
64 --------
66 The buoyancy input for a single species called "salinity", which is
67 simulated by a GWT model called "gwt-1" are specified as follows:
69 >>> buy = imod.mf6.Buoyance(
70 ... reference_density=1000.0,
71 ... density_concentration_slope=[0.025],
72 ... reference_concentration=[0.0],
73 ... modelname=["gwt-1"],
74 ... species=["salinity"],
75 ... )
77 Multiple species can be specified by presenting multiple values with an
78 associated species coordinate. Two species, called "c1" and "c2", simulated
79 by the GWT models "gwt-1" and "gwt-2" are specified as:
81 >>> coords = {"species": ["c1", "c2"]}
82 >>> buy = imod.mf6.Buoyance(
83 ... reference_density=1000.0,
84 ... density_concentration_slope=[0.025, 0.01],
85 ... reference_concentration=[0.0, 0.0],
86 ... modelname=["gwt-1", "gwt-2"],
87 ... species=["c1", "c2"],
88 ... )
89 """
91 _pkg_id = "buy"
92 _template = Package._initialize_template(_pkg_id)
94 _init_schemata = {
95 "reference_density": [DTypeSchema(np.floating)],
96 "density_concentration_slope": [DTypeSchema(np.floating)],
97 "reference_concentration": [DTypeSchema(np.floating)],
98 }
100 _write_schemata = {}
102 @init_log_decorator()
103 def __init__(
104 self,
105 reference_density: float,
106 density_concentration_slope: Sequence[float],
107 reference_concentration: Sequence[float],
108 modelname: Sequence[str],
109 species: Sequence[str],
110 hhformulation_rhs: bool = False,
111 densityfile: Optional[str] = None,
112 validate: bool = True,
113 ):
114 dict_dataset = {
115 "reference_density": reference_density,
116 # Assign a shared index: this also forces equal lenghts
117 "density_concentration_slope": assign_index(density_concentration_slope),
118 "reference_concentration": assign_index(reference_concentration),
119 "modelname": assign_index(modelname),
120 "species": assign_index(species),
121 "hhformulation_rhs": hhformulation_rhs,
122 "densityfile": densityfile,
123 }
124 super().__init__(dict_dataset)
125 self._validate_init_schemata(validate)
127 def render(self, directory, pkgname, globaltimes, binary):
128 ds = self.dataset
129 packagedata = []
131 for i, (a, b, c, d) in enumerate(
132 zip(
133 ds["density_concentration_slope"].values,
134 ds["reference_concentration"].values,
135 ds["modelname"].values,
136 ds["species"].values,
137 )
138 ):
139 packagedata.append((i + 1, a, b, c, d))
141 d = {
142 "nrhospecies": self.dataset["species"].size,
143 "packagedata": packagedata,
144 }
146 for varname in ["hhformulation_rhs", "reference_density", "densityfile"]:
147 value = self.dataset[varname].values[()]
148 if self._valid(value):
149 d[varname] = value
151 return self._template.render(d)
153 def update_transport_models(self, new_modelnames: Sequence[str]):
154 """
155 The names of the transport models can change in some cases, for example
156 when partitioning. Use this function to update the names of the
157 transport models.
158 """
159 transport_model_names = self.get_transport_model_names()
160 if not len(transport_model_names) == len(new_modelnames):
161 raise ValueError("the number of transport models cannot be changed.")
162 for modelname, new_modelname in zip(transport_model_names, new_modelnames):
163 if not modelname in new_modelname:
164 raise ValueError(
165 "new transport model names do not match the old ones. The new names should be equal to the old ones, with a suffix."
166 )
167 self.dataset["modelname"] = assign_index(new_modelnames)
169 def get_transport_model_names(self) -> list[str]:
170 """
171 Returns the names of the transport models used by this buoyancy package.
172 """
173 return list(self.dataset["modelname"].values)