from typing import Optional, Sequence
import numpy as np
import xarray as xr
from imod.logging import init_log_decorator
from imod.mf6.package import Package
from imod.schemata import DTypeSchema
def assign_index(arg):
if isinstance(arg, xr.DataArray):
arg = arg.values
elif not isinstance(arg, (np.ndarray, list, tuple)):
raise TypeError("should be a tuple, list, or numpy array")
arr = np.array(arg)
if arr.ndim != 1:
raise ValueError("should be 1D")
return xr.DataArray(arr, dims=("index",))
[docs]
class Buoyancy(Package):
"""
Buoyancy package. This package must be included when performing variable
density simulation.
Note that ``reference_density`` is a single value, but
``density_concentration_slope``, ``reference_concentration`` and
``modelname`` require an entry for every active species. Please refer to the
examples.
Parameters
----------
reference_density: float,
fluid reference density used in the equation of state.
density_concentration_slope: sequence of floats
Slope of the (linear) density concentration line used in the density
equation of state.
reference_concentration: sequence of floats
Reference concentration used in the density equation of state.
modelname: sequence of strings,
Name of the GroundwaterTransport (GWT) model used for the
concentrations.
species: sequence of str,
Name of the species used to calculate a density value.
hhformulation_rhs: bool, optional.
use the variable-density hydraulic head formulation and add off-diagonal
terms to the right-hand. This option will prevent the BUY Package from
adding asymmetric terms to the flow matrix. Default value is ``False``.
densityfile:
name of the binary output file to write density information. The density
file has the same format as the head file. Density values will be
written to the density file whenever heads are written to the binary
head file. The settings for controlling head output are contained in the
Output Control option.
validate: {True, False}
Flag to indicate whether the package should be validated upon
initialization. This raises a ValidationError if package input is
provided in the wrong manner. Defaults to True.
Examples
--------
The buoyancy input for a single species called "salinity", which is
simulated by a GWT model called "gwt-1" are specified as follows:
>>> buy = imod.mf6.Buoyance(
... reference_density=1000.0,
... density_concentration_slope=[0.025],
... reference_concentration=[0.0],
... modelname=["gwt-1"],
... species=["salinity"],
... )
Multiple species can be specified by presenting multiple values with an
associated species coordinate. Two species, called "c1" and "c2", simulated
by the GWT models "gwt-1" and "gwt-2" are specified as:
>>> coords = {"species": ["c1", "c2"]}
>>> buy = imod.mf6.Buoyance(
... reference_density=1000.0,
... density_concentration_slope=[0.025, 0.01],
... reference_concentration=[0.0, 0.0],
... modelname=["gwt-1", "gwt-2"],
... species=["c1", "c2"],
... )
"""
_pkg_id = "buy"
_template = Package._initialize_template(_pkg_id)
_init_schemata = {
"reference_density": [DTypeSchema(np.floating)],
"density_concentration_slope": [DTypeSchema(np.floating)],
"reference_concentration": [DTypeSchema(np.floating)],
}
_write_schemata = {}
[docs]
@init_log_decorator()
def __init__(
self,
reference_density: float,
density_concentration_slope: Sequence[float],
reference_concentration: Sequence[float],
modelname: Sequence[str],
species: Sequence[str],
hhformulation_rhs: bool = False,
densityfile: Optional[str] = None,
validate: bool = True,
):
dict_dataset = {
"reference_density": reference_density,
# Assign a shared index: this also forces equal lenghts
"density_concentration_slope": assign_index(density_concentration_slope),
"reference_concentration": assign_index(reference_concentration),
"modelname": assign_index(modelname),
"species": assign_index(species),
"hhformulation_rhs": hhformulation_rhs,
"densityfile": densityfile,
}
super().__init__(dict_dataset)
self._validate_init_schemata(validate)
def render(self, directory, pkgname, globaltimes, binary):
ds = self.dataset
packagedata = []
for i, (a, b, c, d) in enumerate(
zip(
ds["density_concentration_slope"].values,
ds["reference_concentration"].values,
ds["modelname"].values,
ds["species"].values,
)
):
packagedata.append((i + 1, a, b, c, d))
d = {
"nrhospecies": self.dataset["species"].size,
"packagedata": packagedata,
}
for varname in ["hhformulation_rhs", "reference_density", "densityfile"]:
value = self.dataset[varname].values[()]
if self._valid(value):
d[varname] = value
return self._template.render(d)
def update_transport_models(self, new_modelnames: Sequence[str]):
"""
The names of the transport models can change in some cases, for example
when partitioning. Use this function to update the names of the
transport models.
"""
transport_model_names = self.get_transport_model_names()
if not len(transport_model_names) == len(new_modelnames):
raise ValueError("the number of transport models cannot be changed.")
for modelname, new_modelname in zip(transport_model_names, new_modelnames):
if modelname not in new_modelname:
raise ValueError(
"new transport model names do not match the old ones. The new names should be equal to the old ones, with a suffix."
)
self.dataset["modelname"] = assign_index(new_modelnames)
def get_transport_model_names(self) -> list[str]:
"""
Returns the names of the transport models used by this buoyancy package.
"""
return list(self.dataset["modelname"].values)
