Coverage for fixtures\msw_model_fixture.py: 63%

1 import numpy as np

2 import pandas as pd

3 import pytest

4 import xarray as xr

5 from numpy import nan

6

7 from imod import mf6, msw

8

9

10 def 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 outer_dvclose=1.0e-4,

95 outer_maximum=500,

96 under_relaxation=None,

97 inner_dvclose=1.0e-4,

98 inner_rclose=0.001,

99 inner_maximum=100,

100 linear_acceleration="cg",

101 scaling_method=None,

102 reordering_method=None,

103 relaxation_factor=0.97,

104 )

105 # Collect time discretization

106 simulation.create_time_discretization(additional_times=times)

107

108 return simulation

109

110

111 def make_msw_model():

112 unsaturated_database = "./unsat_database"

113

114 x = [1.0, 2.0, 3.0]

115 y = [3.0, 2.0, 1.0]

116 subunit = [0, 1]

117 dx = 1.0

118 dy = -1.0

119

120 nrow = len(y)

121 ncol = len(x)

122

123 freq = "D"

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

125

126 # fmt: off

127 area = xr.DataArray(

128 np.array(

129 [

130 [[0.5, 0.5, 0.5],

131 [nan, nan, nan],

132 [1.0, 1.0, 1.0]],

133

134 [[0.5, 0.5, 0.5],

135 [1.0, 1.0, 1.0],

136 [nan, nan, nan]],

137 ]

138 ),

139 dims=("subunit", "y", "x"),

140 coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}

141 )

142

143 active = xr.DataArray(

144 np.array(

145 [[False, True, False],

146 [False, True, False],

147 [False, True, False]]),

148 dims=("y", "x"),

149 coords={"y": y, "x": x}

150 )

151 # fmt: on

152 msw_grid = xr.ones_like(active, dtype=float)

153

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

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

156

157 # %% Vegetation

158 day_of_year = np.arange(1, 367)

159 vegetation_index = np.arange(1, 4)

160

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

162

163 soil_cover = xr.DataArray(

164 data=np.zeros(day_of_year.shape + vegetation_index.shape),

165 coords=coords,

166 dims=("day_of_year", "vegetation_index"),

167 )

168 soil_cover[132:254, :] = 1.0

169 leaf_area_index = soil_cover * 3

170

171 vegetation_factor = xr.zeros_like(soil_cover)

172 vegetation_factor[132:142, :] = 0.7

173 vegetation_factor[142:152, :] = 0.9

174 vegetation_factor[152:162, :] = 1.0

175 vegetation_factor[162:192, :] = 1.2

176 vegetation_factor[192:244, :] = 1.1

177 vegetation_factor[244:254, :] = 0.7

178

179 # %% Landuse

180 landuse_index = np.arange(1, 4)

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

182

183 coords = {"landuse_index": landuse_index}

184

185 landuse_names = xr.DataArray(data=names, coords=coords, dims=("landuse_index",))

186 vegetation_index_da = xr.DataArray(

187 data=vegetation_index, coords=coords, dims=("landuse_index",)

188 )

189 lu = xr.ones_like(vegetation_index_da, dtype=float)

190

191 # %% Well

192

193 wel_layer = 3

194

195 ix = np.tile(np.arange(ncol) + 1, nrow)

196 iy = np.repeat(np.arange(nrow) + 1, ncol)

197 rate = np.zeros(ix.shape)

198 layer = np.full_like(ix, wel_layer)

199

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

201

202 # %% Modflow 6

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

204 layer = [1, 2, 3]

205

206 top = 0.0

207 bottom = top - xr.DataArray(

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

209 )

210

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

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

213

214 # %% Initiate model

215 msw_model = msw.MetaSwapModel(unsaturated_database=unsaturated_database)

216 msw_model["grid"] = msw.GridData(

217 area,

218 xr.full_like(area, 1, dtype=int),

219 xr.full_like(area, 1.0, dtype=float),

220 xr.full_like(active, 1.0, dtype=float),

221 xr.full_like(active, 1, dtype=int),

222 active,

223 )

224

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

226

227 # %% Meteo

228 msw_model["meteo_grid"] = msw.MeteoGrid(precipitation, evapotranspiration)

229 msw_model["mapping_prec"] = msw.PrecipitationMapping(precipitation)

230 msw_model["mapping_evt"] = msw.EvapotranspirationMapping(precipitation * 1.5)

231

232 # %% Sprinkling

233 msw_model["sprinkling"] = msw.Sprinkling(

234 max_abstraction_groundwater=xr.full_like(msw_grid, 100.0),

235 max_abstraction_surfacewater=xr.full_like(msw_grid, 100.0),

236 well=well,

237 )

238

239 # %% Ponding

240 msw_model["ponding"] = msw.Ponding(

241 ponding_depth=xr.full_like(area, 0.0),

242 runon_resistance=xr.full_like(area, 1.0),

243 runoff_resistance=xr.full_like(area, 1.0),

244 )

245

246 # %% Scaling Factors

247 msw_model["scaling"] = msw.ScalingFactors(

248 scale_soil_moisture=xr.full_like(area, 1.0),

249 scale_hydraulic_conductivity=xr.full_like(area, 1.0),

250 scale_pressure_head=xr.full_like(area, 1.0),

251 depth_perched_water_table=xr.full_like(msw_grid, 1.0),

252 )

253

254 # %% Infiltration

255 msw_model["infiltration"] = msw.Infiltration(

256 infiltration_capacity=xr.full_like(area, 1.0),

257 downward_resistance=xr.full_like(msw_grid, -9999.0),

258 upward_resistance=xr.full_like(msw_grid, -9999.0),

259 bottom_resistance=xr.full_like(msw_grid, -9999.0),

260 extra_storage_coefficient=xr.full_like(msw_grid, 0.1),

261 )

262

263 # %% Vegetation

264 msw_model["crop_factors"] = msw.AnnualCropFactors(

265 soil_cover=soil_cover,

266 leaf_area_index=leaf_area_index,

267 interception_capacity=xr.zeros_like(soil_cover),

268 vegetation_factor=vegetation_factor,

269 interception_factor=xr.ones_like(soil_cover),

270 bare_soil_factor=xr.ones_like(soil_cover),

271 ponding_factor=xr.ones_like(soil_cover),

272 )

273

274 # %% Landuse options

275 msw_model["landuse_options"] = msw.LanduseOptions(

276 landuse_name=landuse_names,

277 vegetation_index=vegetation_index_da,

278 jarvis_o2_stress=xr.ones_like(lu),

279 jarvis_drought_stress=xr.ones_like(lu),

280 feddes_p1=xr.full_like(lu, 99.0),

281 feddes_p2=xr.full_like(lu, 99.0),

282 feddes_p3h=lu * [-2.0, -4.0, -3.0],

283 feddes_p3l=lu * [-8.0, -5.0, -5.0],

284 feddes_p4=lu * [-80.0, -100.0, -100.0],

285 feddes_t3h=xr.full_like(lu, 5.0),

286 feddes_t3l=xr.full_like(lu, 1.0),

287 threshold_sprinkling=lu * [-8.0, -5.0, -5.0],

288 fraction_evaporated_sprinkling=xr.full_like(lu, 0.05),

289 gift=xr.full_like(lu, 20.0),

290 gift_duration=xr.full_like(lu, 0.25),

291 rotational_period=lu * [10, 7, 7],

292 start_sprinkling_season=lu * [120, 180, 150],

293 end_sprinkling_season=lu * [230, 230, 240],

294 interception_option=xr.ones_like(lu, dtype=int),

295 interception_capacity_per_LAI=xr.zeros_like(lu),

296 interception_intercept=xr.ones_like(lu),

297 )

298

299 # %% Metaswap Mappings

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

301

302 # %% Output Control

303 msw_model["oc_idf"] = msw.IdfMapping(area, -9999.0)

304 msw_model["oc_var"] = msw.VariableOutputControl()

305 msw_model["oc_time"] = msw.TimeOutputControl(time=times)

306

307 return msw_model

308

309

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

311 def coupled_mf6_model():

312 return make_coupled_mf6_model()

313

314

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

316 def msw_model():

317 return make_msw_model()