Coverage for C:\src\imod-python\imod\mf6\ssm.py: 85%
60 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +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": list(
73 zip(
74 self["package_names"].values,
75 self["concentration_boundary_type"].values,
76 self["auxiliary_variable_name"].values,
77 )
78 ),
79 }
80 for var in ("print_flows", "save_flows"):
81 value = self[var].values[()]
82 if self._valid(value):
83 d[var] = value
84 return self._template.render(d)
86 @staticmethod
87 def from_flow_model(
88 model: GroundwaterFlowModel,
89 species: str,
90 print_flows: bool = False,
91 save_flows: bool = False,
92 validate: bool = True,
93 is_split: bool = False,
94 ):
95 """
96 Derive a Source and Sink Mixing package from a Groundwater Flow model's
97 boundary conditions (e.g. GeneralHeadBoundary). Boundary condition
98 packages which have the ``concentration`` variable set are included.
100 Parameters
101 ----------
102 model: GroundwaterFlowModel
103 Groundwater flow model from which sources & sinks have to be
104 inferred.
105 species: str
106 Name of species to create a transport model for. This name will be
107 looked for in the ``species`` dimensions of the ``concentration``
108 argument.
109 print_flows: ({True, False}, optional)
110 Indicates that the list of general head boundary flow rates will be
111 printed to the listing file for every stress period time step in which
112 "BUDGET PRINT" is specified in Output Control. If there is no Output
113 Control option and PRINT FLOWS is specified, then flow rates are printed
114 for the last time step of each stress period.
115 Default is False.
116 save_flows: ({True, False}, optional)
117 Indicates that general head boundary flow terms will be written to the
118 file specified with "BUDGET FILEOUT" in Output Control.
119 Default is False.
120 validate: ({True, False}, optional)
121 Flag to indicate whether the package should be validated upon
122 initialization. This raises a ValidationError if package input is
123 provided in the wrong manner. Defaults to True.
124 is_split: ({True, False}, optional)
125 Flag to indicate if the simulation has been split into partitions.
126 """
127 if not isinstance(model, GroundwaterFlowModel):
128 raise TypeError(
129 "model must be a GroundwaterFlowModel, received instead: "
130 f"{type(model).__name__}"
131 )
132 names = []
133 boundary_types = []
134 aux_var_names = []
135 for name, package in model.items():
136 if isinstance(package, BoundaryCondition):
137 ds = package.dataset
138 # The package should contain a concentration variable, with a
139 # species coordinate.
140 if "concentration" not in ds.data_vars:
141 raise ValueError(f"concentration not present in package {name}")
142 if "species" not in ds["concentration"].coords:
143 raise ValueError(
144 f"No species coordinate for concentration in package {name}"
145 )
147 # While somewhat far-fetched, it is possible for different
148 # species to have different mixing behavior.
149 type_da = ds["concentration_boundary_type"]
150 if "species" in type_da.dims:
151 type_da = type_da.sel(species=species)
153 names.append(name)
154 boundary_types.append(type_da.values[()])
155 aux_var_names.append(species)
157 if len(names) == 0:
158 msg = "flow model does not contain boundary conditions"
159 if is_split:
160 logger.info(f"{msg}, returning None instead of SourceSinkMixing")
161 return
162 else:
163 raise ValueError(msg)
165 return SourceSinkMixing(
166 names,
167 boundary_types,
168 aux_var_names,
169 print_flows=print_flows,
170 save_flows=save_flows,
171 validate=validate,
172 )
174 def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]:
175 return self._regrid_method