Coverage for C:\src\imod-python\imod\mf6\ssm.py: 85%
60 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
5from imod.logging import init_log_decorator, logger
6from imod.mf6 import GroundwaterFlowModel
7from imod.mf6.boundary_condition import BoundaryCondition
8from imod.mf6.interfaces.iregridpackage import IRegridPackage
9from imod.mf6.utilities.regrid import RegridderType
10from imod.schemata import DTypeSchema
13def with_index_dim(array_like):
14 # At least1d will also extract the values if array_like is a DataArray.
15 arr1d = np.atleast_1d(array_like)
16 if arr1d.ndim > 1:
17 raise ValueError("array must be 1d")
18 return ("index", arr1d)
21class SourceSinkMixing(BoundaryCondition, IRegridPackage):
22 """
23 Parameters
24 ----------
25 package_names: array_like of str
26 concentration_boundary_type: array_like of str
27 auxiliary_variable_name: array_like of str
28 print_flows: bool
29 save_flows: bool
30 validate: bool
31 """
33 _pkg_id = "ssm"
34 _template = BoundaryCondition._initialize_template(_pkg_id)
36 _init_schemata = {
37 "package_names": [DTypeSchema(np.str_)],
38 "concentration_boundary_type": [DTypeSchema(np.str_)],
39 "auxiliary_variable_name": [DTypeSchema(np.str_)],
40 "print_flows": [DTypeSchema(np.bool_)],
41 "save_flows": [DTypeSchema(np.bool_)],
42 }
44 _write_schemata = {}
46 _regrid_method: dict[str, Tuple[RegridderType, str]] = {}
48 @init_log_decorator()
49 def __init__(
50 self,
51 package_names,
52 concentration_boundary_type,
53 auxiliary_variable_name,
54 print_flows: bool = False,
55 save_flows: bool = False,
56 validate: bool = True,
57 ):
58 dict_dataset = {
59 # By sharing the index, this will raise an error if lengths do not
60 # match.
61 "package_names": with_index_dim(package_names),
62 "concentration_boundary_type": with_index_dim(concentration_boundary_type),
63 "auxiliary_variable_name": with_index_dim(auxiliary_variable_name),
64 "print_flows": print_flows,
65 "save_flows": save_flows,
66 }
67 super().__init__(dict_dataset)
68 self._validate_init_schemata(validate)
70 def render(self, directory, pkgname, globaltimes, binary):
71 d = {
72 "sources": [
73 (a, b, c)
74 for a, b, c in zip(
75 self["package_names"].values,
76 self["concentration_boundary_type"].values,
77 self["auxiliary_variable_name"].values,
78 )
79 ],
80 }
81 for var in ("print_flows", "save_flows"):
82 value = self[var].values[()]
83 if self._valid(value):
84 d[var] = value
85 return self._template.render(d)
87 @staticmethod
88 def from_flow_model(
89 model: GroundwaterFlowModel,
90 species: str,
91 print_flows: bool = False,
92 save_flows: bool = False,
93 validate: bool = True,
94 is_split: bool = False,
95 ):
96 """
97 Derive a Source and Sink Mixing package from a Groundwater Flow model's
98 boundary conditions (e.g. GeneralHeadBoundary). Boundary condition
99 packages which have the ``concentration`` variable set are included.
101 Parameters
102 ----------
103 model: GroundwaterFlowModel
104 Groundwater flow model from which sources & sinks have to be
105 inferred.
106 species: str
107 Name of species to create a transport model for. This name will be
108 looked for in the ``species`` dimensions of the ``concentration``
109 argument.
110 print_flows: ({True, False}, optional)
111 Indicates that the list of general head boundary flow rates will be
112 printed to the listing file for every stress period time step in which
113 "BUDGET PRINT" is specified in Output Control. If there is no Output
114 Control option and PRINT FLOWS is specified, then flow rates are printed
115 for the last time step of each stress period.
116 Default is False.
117 save_flows: ({True, False}, optional)
118 Indicates that general head boundary flow terms will be written to the
119 file specified with "BUDGET FILEOUT" in Output Control.
120 Default is False.
121 validate: ({True, False}, optional)
122 Flag to indicate whether the package should be validated upon
123 initialization. This raises a ValidationError if package input is
124 provided in the wrong manner. Defaults to True.
125 is_split: ({True, False}, optional)
126 Flag to indicate if the simulation has been split into partitions.
127 """
128 if not isinstance(model, GroundwaterFlowModel):
129 raise TypeError(
130 "model must be a GroundwaterFlowModel, received instead: "
131 f"{type(model).__name__}"
132 )
133 names = []
134 boundary_types = []
135 aux_var_names = []
136 for name, package in model.items():
137 if isinstance(package, BoundaryCondition):
138 ds = package.dataset
139 # The package should contain a concentration variable, with a
140 # species coordinate.
141 if "concentration" not in ds.data_vars:
142 raise ValueError(f"concentration not present in package {name}")
143 if "species" not in ds["concentration"].coords:
144 raise ValueError(
145 f"No species coordinate for concentration in package {name}"
146 )
148 # While somewhat far-fetched, it is possible for different
149 # species to have different mixing behavior.
150 type_da = ds["concentration_boundary_type"]
151 if "species" in type_da.dims:
152 type_da = type_da.sel(species=species)
154 names.append(name)
155 boundary_types.append(type_da.values[()])
156 aux_var_names.append(species)
158 if len(names) == 0:
159 msg = "flow model does not contain boundary conditions"
160 if is_split:
161 logger.info(f"{msg}, returning None instead of SourceSinkMixing")
162 return
163 else:
164 raise ValueError(msg)
166 return SourceSinkMixing(
167 names,
168 boundary_types,
169 aux_var_names,
170 print_flows=print_flows,
171 save_flows=save_flows,
172 validate=validate,
173 )
175 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]:
176 return self._regrid_method