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
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1"""
2Assign wells to layers.
3"""
5from typing import Optional, Union
7import numpy as np
8import pandas as pd
9import xarray as xr
10import xugrid as xu
12import imod
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 )
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
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 )
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()
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")
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]
76 return id_in_bounds, xy_top, xy_bottom, xy_k
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.
93 Wells located outside of the grid are removed.
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
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 """
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}")
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 )
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)
134 if k is None:
135 k = 1.0
136 else:
137 k = xy_k.values.ravel()
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()
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())]
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))
173 indexes = ["id"]
174 for dim in ["species", "time"]:
175 if dim in wells_in_bounds:
176 indexes.append(dim)
177 columns.remove(dim)
179 df[columns] = 1 # N.B. integer!
181 assigned = (
182 wells_in_bounds.set_index(indexes) * df.set_index(["id", "layer"])
183 ).reset_index()
184 return assigned