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

1import numpy as np 

2 

3import imod 

4from imod.flow.pkgbase import BoundaryCondition 

5 

6 

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. 

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)