Coverage for fixtures\msw_model_fixture.py: 10%
121 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
1import numpy as np
2import pandas as pd
3import pytest
4import xarray as xr
5from numpy import nan
7from imod import mf6, msw
10def make_coupled_mf6_model():
11 x = [1.0, 2.0, 3.0]
12 y = [3.0, 2.0, 1.0]
13 dz = np.array([1, 10, 100])
14 layer = [1, 2, 3]
15 dx = 1.0
16 dy = -1.0
18 nlay = len(layer)
19 nrow = len(y)
20 ncol = len(x)
22 freq = "D"
23 times = pd.date_range(start="1/1/1971", end="1/3/1971", freq=freq)
25 like = xr.DataArray(
26 data=np.ones((nlay, nrow, ncol)),
27 dims=("layer", "y", "x"),
28 coords={"layer": layer, "y": y, "x": x, "dx": dx, "dy": dy},
29 )
30 idomain = like.astype(np.int32)
32 top = 0.0
33 bottom = top - xr.DataArray(
34 np.cumsum(layer * dz), coords={"layer": layer}, dims="layer"
35 )
37 head = xr.full_like(like, np.nan)
38 head[..., 0] = -1.0
39 head[..., -1] = -1.0
41 head = head.expand_dims(time=times)
43 gwf_model = mf6.GroundwaterFlowModel()
44 gwf_model["dis"] = mf6.StructuredDiscretization(
45 idomain=idomain, top=top, bottom=bottom
46 )
47 gwf_model["chd"] = mf6.ConstantHead(
48 head, print_input=True, print_flows=True, save_flows=True
49 )
51 icelltype = xr.full_like(bottom, 0, dtype=int)
52 k = xr.DataArray(np.ones((nlay)), {"layer": layer}, ("layer",))
53 k33 = xr.DataArray(np.ones((nlay)), {"layer": layer}, ("layer",))
54 gwf_model["npf"] = mf6.NodePropertyFlow(
55 icelltype=icelltype,
56 k=k,
57 k33=k33,
58 variable_vertical_conductance=True,
59 dewatered=False,
60 perched=False,
61 save_flows=True,
62 )
64 gwf_model["ic"] = mf6.InitialConditions(start=0.5)
65 gwf_model["sto"] = mf6.SpecificStorage(1e-3, 0.1, True, 0)
67 recharge = xr.zeros_like(idomain.sel(layer=1), dtype=float)
68 recharge[:, 0] = np.nan
69 recharge[:, -1] = np.nan
71 gwf_model["rch_msw"] = mf6.Recharge(recharge)
73 gwf_model["oc"] = mf6.OutputControl(save_head="last", save_budget="last")
75 # Create wells
76 wel_layer = 3
78 ix = np.tile(np.arange(ncol) + 1, nrow)
79 iy = np.repeat(np.arange(nrow) + 1, ncol)
80 rate = np.zeros(ix.shape)
81 layer = np.full_like(ix, wel_layer)
83 gwf_model["wells_msw"] = mf6.WellDisStructured(
84 layer=layer, row=iy, column=ix, rate=rate
85 )
87 # Attach it to a simulation
88 simulation = mf6.Modflow6Simulation("test")
89 simulation["GWF_1"] = gwf_model
90 # Define solver settings
91 simulation["solver"] = mf6.Solution(
92 modelnames=["GWF_1"],
93 print_option="summary",
94 csv_output=False,
95 no_ptc=True,
96 outer_dvclose=1.0e-4,
97 outer_maximum=500,
98 under_relaxation=None,
99 inner_dvclose=1.0e-4,
100 inner_rclose=0.001,
101 inner_maximum=100,
102 linear_acceleration="cg",
103 scaling_method=None,
104 reordering_method=None,
105 relaxation_factor=0.97,
106 )
107 # Collect time discretization
108 simulation.create_time_discretization(additional_times=times)
110 return simulation
113def make_msw_model():
114 unsaturated_database = "./unsat_database"
116 x = [1.0, 2.0, 3.0]
117 y = [3.0, 2.0, 1.0]
118 subunit = [0, 1]
119 dx = 1.0
120 dy = -1.0
122 nrow = len(y)
123 ncol = len(x)
125 freq = "D"
126 times = pd.date_range(start="1/1/1971", end="1/3/1971", freq=freq)
128 # fmt: off
129 area = xr.DataArray(
130 np.array(
131 [
132 [[0.5, 0.5, 0.5],
133 [nan, nan, nan],
134 [1.0, 1.0, 1.0]],
136 [[0.5, 0.5, 0.5],
137 [1.0, 1.0, 1.0],
138 [nan, nan, nan]],
139 ]
140 ),
141 dims=("subunit", "y", "x"),
142 coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
143 )
145 active = xr.DataArray(
146 np.array(
147 [[False, True, False],
148 [False, True, False],
149 [False, True, False]]),
150 dims=("y", "x"),
151 coords={"y": y, "x": x}
152 )
153 # fmt: on
154 msw_grid = xr.ones_like(active, dtype=float)
156 precipitation = msw_grid.expand_dims(time=times[:-1])
157 evapotranspiration = msw_grid.expand_dims(time=times[:-1]) * 1.2
159 # %% Vegetation
160 day_of_year = np.arange(1, 367)
161 vegetation_index = np.arange(1, 4)
163 coords = {"day_of_year": day_of_year, "vegetation_index": vegetation_index}
165 soil_cover = xr.DataArray(
166 data=np.zeros(day_of_year.shape + vegetation_index.shape),
167 coords=coords,
168 dims=("day_of_year", "vegetation_index"),
169 )
170 soil_cover[132:254, :] = 1.0
171 leaf_area_index = soil_cover * 3
173 vegetation_factor = xr.zeros_like(soil_cover)
174 vegetation_factor[132:142, :] = 0.7
175 vegetation_factor[142:152, :] = 0.9
176 vegetation_factor[152:162, :] = 1.0
177 vegetation_factor[162:192, :] = 1.2
178 vegetation_factor[192:244, :] = 1.1
179 vegetation_factor[244:254, :] = 0.7
181 # %% Landuse
182 landuse_index = np.arange(1, 4)
183 names = ["grassland", "maize", "potatoes"]
185 coords = {"landuse_index": landuse_index}
187 landuse_names = xr.DataArray(data=names, coords=coords, dims=("landuse_index",))
188 vegetation_index_da = xr.DataArray(
189 data=vegetation_index, coords=coords, dims=("landuse_index",)
190 )
191 lu = xr.ones_like(vegetation_index_da, dtype=float)
193 # %% Well
195 wel_layer = 3
197 ix = np.tile(np.arange(ncol) + 1, nrow)
198 iy = np.repeat(np.arange(nrow) + 1, ncol)
199 rate = np.zeros(ix.shape)
200 layer = np.full_like(ix, wel_layer)
202 well = mf6.WellDisStructured(layer=layer, row=iy, column=ix, rate=rate)
204 # %% Modflow 6
205 dz = np.array([1, 10, 100])
206 layer = [1, 2, 3]
208 top = 0.0
209 bottom = top - xr.DataArray(
210 np.cumsum(layer * dz), coords={"layer": layer}, dims="layer"
211 )
213 idomain = xr.full_like(msw_grid, 1, dtype=int).expand_dims(layer=layer)
214 dis = mf6.StructuredDiscretization(top=top, bottom=bottom, idomain=idomain)
216 # %% Initiate model
217 msw_model = msw.MetaSwapModel(unsaturated_database=unsaturated_database)
218 msw_model["grid"] = msw.GridData(
219 area,
220 xr.full_like(area, 1, dtype=int),
221 xr.full_like(area, 1.0, dtype=float),
222 xr.full_like(active, 1.0, dtype=float),
223 xr.full_like(active, 1, dtype=int),
224 active,
225 )
227 msw_model["ic"] = msw.InitialConditionsRootzonePressureHead(initial_pF=2.2)
229 # %% Meteo
230 msw_model["meteo_grid"] = msw.MeteoGrid(precipitation, evapotranspiration)
231 msw_model["mapping_prec"] = msw.PrecipitationMapping(precipitation)
232 msw_model["mapping_evt"] = msw.EvapotranspirationMapping(precipitation * 1.5)
234 # %% Sprinkling
235 msw_model["sprinkling"] = msw.Sprinkling(
236 max_abstraction_groundwater=xr.full_like(msw_grid, 100.0),
237 max_abstraction_surfacewater=xr.full_like(msw_grid, 100.0),
238 well=well,
239 )
241 # %% Ponding
242 msw_model["ponding"] = msw.Ponding(
243 ponding_depth=xr.full_like(area, 0.0),
244 runon_resistance=xr.full_like(area, 1.0),
245 runoff_resistance=xr.full_like(area, 1.0),
246 )
248 # %% Scaling Factors
249 msw_model["scaling"] = msw.ScalingFactors(
250 scale_soil_moisture=xr.full_like(area, 1.0),
251 scale_hydraulic_conductivity=xr.full_like(area, 1.0),
252 scale_pressure_head=xr.full_like(area, 1.0),
253 depth_perched_water_table=xr.full_like(msw_grid, 1.0),
254 )
256 # %% Infiltration
257 msw_model["infiltration"] = msw.Infiltration(
258 infiltration_capacity=xr.full_like(area, 1.0),
259 downward_resistance=xr.full_like(msw_grid, -9999.0),
260 upward_resistance=xr.full_like(msw_grid, -9999.0),
261 bottom_resistance=xr.full_like(msw_grid, -9999.0),
262 extra_storage_coefficient=xr.full_like(msw_grid, 0.1),
263 )
265 # %% Vegetation
266 msw_model["crop_factors"] = msw.AnnualCropFactors(
267 soil_cover=soil_cover,
268 leaf_area_index=leaf_area_index,
269 interception_capacity=xr.zeros_like(soil_cover),
270 vegetation_factor=vegetation_factor,
271 interception_factor=xr.ones_like(soil_cover),
272 bare_soil_factor=xr.ones_like(soil_cover),
273 ponding_factor=xr.ones_like(soil_cover),
274 )
276 # %% Landuse options
277 msw_model["landuse_options"] = msw.LanduseOptions(
278 landuse_name=landuse_names,
279 vegetation_index=vegetation_index_da,
280 jarvis_o2_stress=xr.ones_like(lu),
281 jarvis_drought_stress=xr.ones_like(lu),
282 feddes_p1=xr.full_like(lu, 99.0),
283 feddes_p2=xr.full_like(lu, 99.0),
284 feddes_p3h=lu * [-2.0, -4.0, -3.0],
285 feddes_p3l=lu * [-8.0, -5.0, -5.0],
286 feddes_p4=lu * [-80.0, -100.0, -100.0],
287 feddes_t3h=xr.full_like(lu, 5.0),
288 feddes_t3l=xr.full_like(lu, 1.0),
289 threshold_sprinkling=lu * [-8.0, -5.0, -5.0],
290 fraction_evaporated_sprinkling=xr.full_like(lu, 0.05),
291 gift=xr.full_like(lu, 20.0),
292 gift_duration=xr.full_like(lu, 0.25),
293 rotational_period=lu * [10, 7, 7],
294 start_sprinkling_season=lu * [120, 180, 150],
295 end_sprinkling_season=lu * [230, 230, 240],
296 interception_option=xr.ones_like(lu, dtype=int),
297 interception_capacity_per_LAI=xr.zeros_like(lu),
298 interception_intercept=xr.ones_like(lu),
299 )
301 # %% Metaswap Mappings
302 msw_model["mod2svat"] = msw.CouplerMapping(modflow_dis=dis, well=well)
304 # %% Output Control
305 msw_model["oc_idf"] = msw.IdfMapping(area, -9999.0)
306 msw_model["oc_var"] = msw.VariableOutputControl()
307 msw_model["oc_time"] = msw.TimeOutputControl(time=times)
309 return msw_model
312@pytest.fixture(scope="function")
313def coupled_mf6_model():
314 return make_coupled_mf6_model()
317@pytest.fixture(scope="function")
318def msw_model():
319 return make_msw_model()