Coverage for C:\src\imod-python\imod\flow\wel.py: 76%

1 import numpy as np

2

3 import imod

4 from imod.flow.pkgbase import BoundaryCondition

5

6

7 class 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.

11

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 """

29

30 _pkg_id = "wel"

31 _variable_order = ["rate"]

32

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

55

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")

63

64 def _compose_values_layer(self, varname, directory, nlayer, time=None):

65 values = {}

66 d = {"directory": directory, "name": directory.stem, "extension": ".ipf"}

67

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)

77

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)

90

91 return values

92

93 def _is_periodic(self):

94 # Periodic stresses are defined for all variables

95 return "stress_periodic" in self.dataset.attrs

96

97 def _get_runfile_times(self, _, globaltimes, ds_times=None):

98 if ds_times is None:

99 ds_times = np.unique(self["time"].values)

100

101 da = self.dataset

102

103 runfile_times, starts = super()._get_runfile_times(

104 da, globaltimes, ds_times=ds_times

105 )

106

107 return runfile_times, starts

108

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)

112

113 if time is not None:

114 d["time"] = time

115

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)

127

128 def save(self, directory):

129 ds = self.dataset

130

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)