Coverage for C:\src\imod-python\imod\flow\wel.py: 76%
68 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
1import numpy as np
3import imod
4from imod.flow.pkgbase import BoundaryCondition
7class Well(BoundaryCondition):
8 """
9 The Well package is used to simulate a specified flux to individual cells
10 and specified in units of length3/time.
12 Parameters
13 ----------
14 id_name: str or list of str
15 name of the well(s).
16 x: float or list of floats
17 x coordinate of the well(s).
18 y: float or list of floats
19 y coordinate of the well(s).
20 rate: float or list of floats.
21 pumping rate in the well(s).
22 layer: "None" or int, optional
23 layer from which the pumping takes place.
24 time: "None" or listlike of np.datetime64, datetime.datetime, pd.Timestamp,
25 cftime.datetime
26 time during which the pumping takes place. Only need to specify if
27 model is transient.
28 """
30 _pkg_id = "wel"
31 _variable_order = ["rate"]
33 def __init__(
34 self,
35 id_name,
36 x,
37 y,
38 rate,
39 layer,
40 time=None,
41 ):
42 super().__init__()
43 variables = {
44 "id_name": id_name,
45 "x": x,
46 "y": y,
47 "rate": rate,
48 "layer": layer,
49 "time": time,
50 }
51 variables = {k: np.atleast_1d(v) for k, v in variables.items() if v is not None}
52 length = max(map(len, variables.values()))
53 index = np.arange(1, length + 1)
54 self.dataset["index"] = index
56 for k, v in variables.items():
57 if v.size == index.size:
58 self.dataset[k] = ("index", v)
59 elif v.size == 1:
60 self.dataset[k] = ("index", np.full(length, v))
61 else:
62 raise ValueError(f"Length of {k} does not match other arguments")
64 def _compose_values_layer(self, varname, directory, nlayer, time=None):
65 values = {}
66 d = {"directory": directory, "name": directory.stem, "extension": ".ipf"}
68 if time is None:
69 if "layer" in self.dataset:
70 for layer in np.unique(self.dataset["layer"]):
71 layer = int(layer)
72 d["layer"] = layer
73 values[layer] = self._compose_path(d)
74 else:
75 for layer in range(1, nlayer + 1): # 1-based indexing
76 values[layer] = self._compose_path(d)
78 else:
79 d["time"] = time
80 if "layer" in self.dataset:
81 # Since the well data is in long table format, it's the only
82 # input that has to be inspected.
83 select = np.argwhere((self["time"] == time).values)
84 for layer in np.unique(self["layer"].values[select]):
85 d["layer"] = layer
86 values[layer] = self._compose_path(d)
87 else:
88 for layer in range(1, nlayer + 1): # 1-based indexing
89 values[layer] = self._compose_path(d)
91 return values
93 def _is_periodic(self):
94 # Periodic stresses are defined for all variables
95 return "stress_periodic" in self.dataset.attrs
97 def _get_runfile_times(self, _, globaltimes, ds_times=None):
98 if ds_times is None:
99 ds_times = np.unique(self["time"].values)
101 da = self.dataset
103 runfile_times, starts = super()._get_runfile_times(
104 da, globaltimes, ds_times=ds_times
105 )
107 return runfile_times, starts
109 def _save_layers(self, df, directory, time=None):
110 d = {"directory": directory, "name": directory.stem, "extension": ".ipf"}
111 d["directory"].mkdir(exist_ok=True, parents=True)
113 if time is not None:
114 d["time"] = time
116 if "layer" in df:
117 for layer, layerdf in df.groupby("layer"):
118 d["layer"] = layer
119 # Ensure right order
120 outdf = layerdf[["x", "y", "rate", "id_name"]]
121 path = self._compose_path(d)
122 imod.ipf.write(path, outdf)
123 else:
124 outdf = df[["x", "y", "rate", "id_name"]]
125 path = self._compose_path(d)
126 imod.ipf.write(path, outdf)
128 def save(self, directory):
129 ds = self.dataset
131 if "time" in ds:
132 for time, timeda in ds.groupby("time"):
133 timedf = timeda.to_dataframe()
134 self._save_layers(timedf, directory, time=time)
135 else:
136 self._save_layers(ds.to_dataframe(), directory)