Coverage for C:\src\imod-python\imod\wq\rch.py: 90%
77 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
1import abc
3import jinja2
4import numpy as np
6from imod.wq.pkgbase import BoundaryCondition
9class Recharge(BoundaryCondition, abc.ABC):
10 _pkg_id = "rch"
12 _mapping = (("rech", "rate"),)
14 _keywords = {"save_budget": {False: 0, True: 1}}
16 _template = jinja2.Template(
17 "[rch]\n"
18 " nrchop = {{recharge_option}}\n"
19 " irchcb = {{save_budget}}\n"
20 " {%- for name, dictname in mapping -%}"
21 " {%- for time, timedict in dicts[dictname].items() -%}"
22 " {%- for layer, value in timedict.items() %}\n"
23 " {{name}}_p{{time}} = {{value}}\n"
24 " {%- endfor -%}\n"
25 " {%- endfor -%}"
26 " {%- endfor -%}"
27 )
29 def _render(self, directory, globaltimes, nlayer, *args, **kwargs):
30 d = {
31 "mapping": self._mapping,
32 "save_budget": self.dataset["save_budget"].values,
33 "recharge_option": self._option,
34 }
35 self._replace_keyword(d, "save_budget")
37 dicts = {}
38 for _, name in self._mapping:
39 dicts[name] = self._compose_values_timelayer(
40 name, globaltimes, directory, nlayer=nlayer
41 )
42 d["dicts"] = dicts
44 return self._template.render(d)
46 def _pkgcheck(self, ibound=None):
47 pass
49 def repeat_stress(self, rate=None, concentration=None, use_cftime=False):
50 varnames = ["rate", "concentration"]
51 values = [rate, concentration]
52 for varname, value in zip(varnames, values):
53 self._repeat_stress(varname, value, use_cftime)
56class RechargeTopLayer(Recharge):
57 """
58 The Recharge package is used to simulate a specified flux distributed over
59 the top of the model and specified in units of length/time (usually m/d). Within MODFLOW,
60 these rates are multiplied by the horizontal area of the cells to which they
61 are applied to calculate the volumetric flux rates. In this class the
62 Recharge gets applied to the top grid layer (NRCHOP=1).
64 Parameters
65 ----------
66 rate: float or xr.DataArray of floats
67 is the amount of recharge.
68 concentration: float or xr.DataArray of floats
69 is the concentration of the recharge
70 save_budget: bool, optional
71 flag indicating if the budget needs to be saved.
72 Default is False.
73 """
75 _option = 1
77 def __init__(self, rate, concentration, save_budget=False):
78 super().__init__()
79 self["rate"] = rate
80 self["concentration"] = concentration
81 self["save_budget"] = save_budget
83 def _set_ssm_layers(self, ibound):
84 self._ssm_layers = np.array([1])
87class RechargeLayers(Recharge):
88 """
89 The Recharge package is used to simulate a specified flux distributed over
90 the top of the model and specified in units of length/time (usually m/d). Within MODFLOW,
91 these rates are multiplied by the horizontal area of the cells to which they
92 are applied to calculate the volumetric flux rates. In this class the
93 Recharge gets applied to a specific, specified, layer (NRCHOP=2).
95 Parameters
96 ----------
97 rate: float or xr.DataArray of floats
98 is the amount of recharge.
99 recharge_layer: int or xr.DataArray of integers
100 layer number variable that defines the layer in each vertical column
101 where recharge is applied
102 concentration: float or xr.DataArray of floats
103 is the concentration of the recharge
104 save_budget: bool, optional
105 flag indicating if the budget needs to be saved.
106 Default is False.
107 """
109 _option = 2
111 _mapping = (("rech", "rate"), ("irch", "recharge_layer"))
113 def __init__(self, rate, recharge_layer, concentration, save_budget=False):
114 super().__init__()
115 self["rate"] = rate
116 self["recharge_layer"] = recharge_layer
117 self["concentration"] = concentration
118 self["save_budget"] = save_budget
120 def _set_ssm_layers(self, ibound):
121 unique_layers = np.unique(self.dataset["recharge_layer"].values)
122 unique_layers = unique_layers[~np.isnan(unique_layers)]
123 self._ssm_layers = unique_layers.astype(np.int32)
126class RechargeHighestActive(Recharge):
127 """
128 The Recharge package is used to simulate a specified flux distributed over
129 the top of the model and specified in units of length/time (usually m/d). Within MODFLOW,
130 these rates are multiplied by the horizontal area of the cells to which they
131 are applied to calculate the volumetric flux rates. In this class the
132 Recharge gets applied to the highest active cell in each vertical column
133 (NRCHOP=3).
135 Parameters
136 ----------
137 rate: float or xr.DataArray of floats
138 is the amount of recharge.
139 concentration: float or xr.DataArray of floats
140 is the concentration of the recharge
141 save_budget: bool, optional
142 flag indicating if the budget needs to be saved.
143 Default is False.
144 """
146 _option = 3
148 def __init__(self, rate, concentration, save_budget=False):
149 super().__init__()
151 rate_scalar = np.ndim(rate) == 0
152 conc_scalar = np.ndim(concentration) == 0
153 if rate_scalar and (not conc_scalar):
154 raise ValueError("Rate cannot be scalar if concentration is non-scalar.")
156 self["rate"] = rate
157 self["concentration"] = concentration
158 self["save_budget"] = save_budget
160 def _set_ssm_layers(self, ibound):
161 rate = self.dataset["rate"]
162 rate_idf = ("x" in rate.dims) and ("y" in rate.dims)
163 conc_scalar = np.ndim(self.dataset["concentration"]) == 0
164 if rate_idf and conc_scalar:
165 rch_active = (rate != 0.0) & rate.notnull()
166 if "time" in rch_active.dims:
167 rch_active = rch_active.any("time")
168 rch_active = rch_active & (ibound > 0)
169 else:
170 rch_active = ibound > 0
172 top_layer = ibound["layer"].where(rch_active).min("layer")
173 top_layer = top_layer.where((ibound > 0).any("layer"))
174 unique_layers = np.unique(top_layer.values)
175 unique_layers = unique_layers[~np.isnan(unique_layers)]
176 self._ssm_layers = unique_layers.astype(np.int32)
178 def repeat_stress(self, rate=None, concentration=None, use_cftime=False):
179 varnames = ["rate", "concentration"]
180 values = [rate, concentration]
181 for varname, value in zip(varnames, values):
182 self._repeat_stress(varname, value, use_cftime)