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

1import numpy as np 

2import pandas as pd 

3import pytest 

4import xarray as xr 

5from numpy import nan 

6 

7from imod import mf6, msw 

8 

9 

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 

17 

18 nlay = len(layer) 

19 nrow = len(y) 

20 ncol = len(x) 

21 

22 freq = "D" 

23 times = pd.date_range(start="1/1/1971", end="1/3/1971", freq=freq) 

24 

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) 

31 

32 top = 0.0 

33 bottom = top - xr.DataArray( 

34 np.cumsum(layer * dz), coords={"layer": layer}, dims="layer" 

35 ) 

36 

37 head = xr.full_like(like, np.nan) 

38 head[..., 0] = -1.0 

39 head[..., -1] = -1.0 

40 

41 head = head.expand_dims(time=times) 

42 

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 ) 

50 

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 ) 

63 

64 gwf_model["ic"] = mf6.InitialConditions(start=0.5) 

65 gwf_model["sto"] = mf6.SpecificStorage(1e-3, 0.1, True, 0) 

66 

67 recharge = xr.zeros_like(idomain.sel(layer=1), dtype=float) 

68 recharge[:, 0] = np.nan 

69 recharge[:, -1] = np.nan 

70 

71 gwf_model["rch_msw"] = mf6.Recharge(recharge) 

72 

73 gwf_model["oc"] = mf6.OutputControl(save_head="last", save_budget="last") 

74 

75 # Create wells 

76 wel_layer = 3 

77 

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) 

82 

83 gwf_model["wells_msw"] = mf6.WellDisStructured( 

84 layer=layer, row=iy, column=ix, rate=rate 

85 ) 

86 

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) 

109 

110 return simulation 

111 

112 

113def make_msw_model(): 

114 unsaturated_database = "./unsat_database" 

115 

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 

121 

122 nrow = len(y) 

123 ncol = len(x) 

124 

125 freq = "D" 

126 times = pd.date_range(start="1/1/1971", end="1/3/1971", freq=freq) 

127 

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]], 

135 

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 ) 

144 

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) 

155 

156 precipitation = msw_grid.expand_dims(time=times[:-1]) 

157 evapotranspiration = msw_grid.expand_dims(time=times[:-1]) * 1.2 

158 

159 # %% Vegetation 

160 day_of_year = np.arange(1, 367) 

161 vegetation_index = np.arange(1, 4) 

162 

163 coords = {"day_of_year": day_of_year, "vegetation_index": vegetation_index} 

164 

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 

172 

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 

180 

181 # %% Landuse 

182 landuse_index = np.arange(1, 4) 

183 names = ["grassland", "maize", "potatoes"] 

184 

185 coords = {"landuse_index": landuse_index} 

186 

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) 

192 

193 # %% Well 

194 

195 wel_layer = 3 

196 

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) 

201 

202 well = mf6.WellDisStructured(layer=layer, row=iy, column=ix, rate=rate) 

203 

204 # %% Modflow 6 

205 dz = np.array([1, 10, 100]) 

206 layer = [1, 2, 3] 

207 

208 top = 0.0 

209 bottom = top - xr.DataArray( 

210 np.cumsum(layer * dz), coords={"layer": layer}, dims="layer" 

211 ) 

212 

213 idomain = xr.full_like(msw_grid, 1, dtype=int).expand_dims(layer=layer) 

214 dis = mf6.StructuredDiscretization(top=top, bottom=bottom, idomain=idomain) 

215 

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 ) 

226 

227 msw_model["ic"] = msw.InitialConditionsRootzonePressureHead(initial_pF=2.2) 

228 

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) 

233 

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 ) 

240 

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 ) 

247 

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 ) 

255 

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 ) 

264 

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 ) 

275 

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 ) 

300 

301 # %% Metaswap Mappings 

302 msw_model["mod2svat"] = msw.CouplerMapping(modflow_dis=dis, well=well) 

303 

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) 

308 

309 return msw_model 

310 

311 

312@pytest.fixture(scope="function") 

313def coupled_mf6_model(): 

314 return make_coupled_mf6_model() 

315 

316 

317@pytest.fixture(scope="function") 

318def msw_model(): 

319 return make_msw_model()