Coverage for C:\src\imod-python\imod\prepare\wells.py: 100%

67 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1""" 

2Assign wells to layers. 

3""" 

4 

5from typing import Optional, Union 

6 

7import numpy as np 

8import pandas as pd 

9import xarray as xr 

10import xugrid as xu 

11 

12import imod 

13 

14 

15def vectorized_overlap(bounds_a, bounds_b): 

16 """ 

17 Vectorized overlap computation. 

18 Compare with: 

19 overlap = max(0, min(a[1], b[1]) - max(a[0], b[0])) 

20 """ 

21 return np.maximum( 

22 0.0, 

23 np.minimum(bounds_a[:, 1], bounds_b[:, 1]) 

24 - np.maximum(bounds_a[:, 0], bounds_b[:, 0]), 

25 ) 

26 

27 

28def compute_overlap(wells, top, bottom): 

29 # layer bounds shape of (n_well, n_layer, 2) 

30 layer_bounds = np.stack((bottom, top), axis=-1) 

31 well_bounds = np.broadcast_to( 

32 np.stack( 

33 (wells["bottom"].to_numpy(), wells["top"].to_numpy()), 

34 axis=-1, 

35 )[np.newaxis, :, :], 

36 layer_bounds.shape, 

37 ) 

38 overlap = vectorized_overlap( 

39 well_bounds.reshape((-1, 2)), 

40 layer_bounds.reshape((-1, 2)), 

41 ) 

42 return overlap 

43 

44 

45def locate_wells( 

46 wells: pd.DataFrame, 

47 top: Union[xr.DataArray, xu.UgridDataArray], 

48 bottom: Union[xr.DataArray, xu.UgridDataArray], 

49 k: Union[xr.DataArray, xu.UgridDataArray, None], 

50): 

51 if not isinstance(top, (xu.UgridDataArray, xr.DataArray)): 

52 raise TypeError( 

53 "top and bottom should be DataArray or UgridDataArray, received: " 

54 f"{type(top).__name__}" 

55 ) 

56 

57 # Default to a xy_k value of 1.0: weigh every layer equally. 

58 xy_k = 1.0 

59 first = wells.groupby("id").first() 

60 x = first["x"].to_numpy() 

61 y = first["y"].to_numpy() 

62 

63 xy_top = imod.select.points_values(top, x=x, y=y, out_of_bounds="ignore") 

64 xy_bottom = imod.select.points_values(bottom, x=x, y=y, out_of_bounds="ignore") 

65 if k is not None: 

66 xy_k = imod.select.points_values(k, x=x, y=y, out_of_bounds="ignore") 

67 

68 # Discard out-of-bounds wells. 

69 index = xy_top["index"] 

70 if not np.array_equal(xy_bottom["index"], index): 

71 raise ValueError("bottom grid does not match top grid") 

72 if k is not None and not np.array_equal(xy_k["index"], index): 

73 raise ValueError("k grid does not match top grid") 

74 id_in_bounds = first.index[index] 

75 

76 return id_in_bounds, xy_top, xy_bottom, xy_k 

77 

78 

79def assign_wells( 

80 wells: pd.DataFrame, 

81 top: Union[xr.DataArray, xu.UgridDataArray], 

82 bottom: Union[xr.DataArray, xu.UgridDataArray], 

83 k: Optional[Union[xr.DataArray, xu.UgridDataArray]] = None, 

84 minimum_thickness: Optional[float] = 0.05, 

85 minimum_k: Optional[float] = 1.0, 

86) -> pd.DataFrame: 

87 """ 

88 Distribute well pumping rate according to filter length when ``k=None``, or 

89 to transmissivity of the sediments surrounding the filter. Minimum 

90 thickness and minimum k should be set to avoid placing wells in clay 

91 layers. 

92 

93 Wells located outside of the grid are removed. 

94 

95 Parameters 

96 ---------- 

97 wells: pd.DataFrame 

98 Should contain columns x, y, id, top, bottom, rate. 

99 top: xr.DataArray or xu.UgridDataArray 

100 Top of the model layers. 

101 bottom: xr.DataArray or xu.UgridDataArray 

102 Bottom of the model layers. 

103 k: xr.DataArray or xu.UgridDataArray, optional 

104 Horizontal conductivity of the model layers. 

105 minimum_thickness: float, optional, default: 0.01 

106 minimum_k: float, optional, default: 1.0 

107 Minimum conductivity 

108 

109 Returns 

110 ------- 

111 placed_wells: pd.DataFrame 

112 Wells with rate subdivided per layer. Contains the original columns of 

113 ``wells``, as well as layer, overlap, transmissivity. 

114 """ 

115 

116 names = {"x", "y", "id", "top", "bottom", "rate"} 

117 missing = names.difference(wells.columns) 

118 if missing: 

119 raise ValueError(f"Columns are missing in wells dataframe: {missing}") 

120 

121 types = [type(arg) for arg in (top, bottom, k) if arg is not None] 

122 if len(set(types)) != 1: 

123 members = ",".join([t.__name__ for t in types]) 

124 raise TypeError( 

125 "top, bottom, and optionally k should be of the same type, " 

126 f"received: {members}" 

127 ) 

128 

129 id_in_bounds, xy_top, xy_bottom, xy_k = locate_wells(wells, top, bottom, k) 

130 wells_in_bounds = wells.set_index("id").loc[id_in_bounds].reset_index() 

131 first = wells_in_bounds.groupby("id").first() 

132 overlap = compute_overlap(first, xy_top, xy_bottom) 

133 

134 if k is None: 

135 k = 1.0 

136 else: 

137 k = xy_k.values.ravel() 

138 

139 # Distribute rate according to transmissivity. 

140 n_layer, n_well = xy_top.shape 

141 df = pd.DataFrame( 

142 index=pd.Index(np.tile(first.index, n_layer), name="id"), 

143 data={ 

144 "layer": np.repeat(top["layer"], n_well), 

145 "overlap": overlap, 

146 "k": k, 

147 "transmissivity": overlap * k, 

148 }, 

149 ) 

150 # remove entries 

151 # -in very thin layers or when the wellbore penetrates the layer very little 

152 # -in low conductivity layers 

153 df = df.loc[(df["overlap"] >= minimum_thickness) & (df["k"] >= minimum_k)] 

154 df["rate"] = df["transmissivity"] / df.groupby("id")["transmissivity"].transform( 

155 "sum" 

156 ) 

157 # Create a unique index for every id-layer combination. 

158 df["index"] = np.arange(len(df)) 

159 df = df.reset_index() 

160 

161 # Get rid of those that are removed because of minimum thickness or 

162 # transmissivity. 

163 wells_in_bounds = wells_in_bounds.loc[wells_in_bounds["id"].isin(df["id"].unique())] 

164 

165 # Use pandas multi-index broadcasting. 

166 # Maintain all other columns as-is. 

167 wells_in_bounds["index"] = 1 # N.B. integer! 

168 wells_in_bounds["overlap"] = 1.0 

169 wells_in_bounds["k"] = 1.0 

170 wells_in_bounds["transmissivity"] = 1.0 

171 columns = list(set(wells_in_bounds.columns).difference(df.columns)) 

172 

173 indexes = ["id"] 

174 for dim in ["species", "time"]: 

175 if dim in wells_in_bounds: 

176 indexes.append(dim) 

177 columns.remove(dim) 

178 

179 df[columns] = 1 # N.B. integer! 

180 

181 assigned = ( 

182 wells_in_bounds.set_index(indexes) * df.set_index(["id", "layer"]) 

183 ).reset_index() 

184 return assigned