Coverage for fixtures\mf6_small_models_fixture.py: 100%

94 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-13 11:15 +0200

1from copy import deepcopy 

2from typing import Callable, Union 

3 

4import numpy as np 

5import pytest 

6import xarray as xr 

7import xugrid as xu 

8 

9import imod 

10 

11 

12def grid_data_structured( 

13 dtype: type, value: Union[int, float], cellsize: float 

14) -> xr.DataArray: 

15 """ 

16 This function creates a dataarray with scalar values for a grid of configurable cell size. 

17 """ 

18 horizontal_range = 10 

19 y = np.arange(horizontal_range, -cellsize, -cellsize) 

20 x = np.arange(0, horizontal_range + cellsize, cellsize) 

21 

22 nlayer = 3 

23 

24 shape = nlayer, len(x), len(y) 

25 dims = ("layer", "y", "x") 

26 layer = np.arange(1, nlayer + 1) 

27 

28 coords = {"layer": layer, "y": y, "x": x, "dx": cellsize, "dy": cellsize} 

29 

30 structured_grid_data = xr.DataArray( 

31 np.ones(shape, dtype=dtype) * value, coords=coords, dims=dims 

32 ) 

33 

34 return structured_grid_data 

35 

36 

37def grid_data_structured_layered( 

38 dtype: type, value: Union[int, float], cellsize: float 

39) -> xr.DataArray: 

40 """ 

41 This function creates a dataarray with scalar values for a grid of configurable cell size. The values are 

42 multiplied with the layer index. 

43 """ 

44 unstructured_grid_data = grid_data_structured(dtype, value, cellsize) 

45 nlayer = unstructured_grid_data.coords["layer"].max().values[()] 

46 for ilayer in range(1, nlayer + 1): 

47 layer_value = ilayer * value 

48 unstructured_grid_data.loc[{"layer": ilayer}] = layer_value 

49 return unstructured_grid_data 

50 

51 

52def grid_data_unstructured( 

53 dtype: type, value: Union[int, float], cellsize: float 

54) -> xu.UgridDataArray: 

55 """ 

56 This function creates a dataarray with scalar values for a grid of configurable cell size. 

57 First a regular grid is constructed and then this is converted to an ugrid dataarray. 

58 """ 

59 return xu.UgridDataArray.from_structured( 

60 grid_data_structured(dtype, value, cellsize) 

61 ) 

62 

63 

64def grid_data_unstructured_layered( 

65 dtype: type, value: Union[int, float], cellsize: float 

66) -> xu.UgridDataArray: 

67 """ 

68 This function creates a dataarray with scalar values for a grid of configurable cell size. The values are 

69 multiplied with the layer index. First a regular grid is constructed and then this is converted to an ugrid dataarray. 

70 """ 

71 return xu.UgridDataArray.from_structured( 

72 grid_data_structured_layered(dtype, value, cellsize) 

73 ) 

74 

75 

76def _make_model( 

77 grid_data_function: ( 

78 Callable[[type, Union[int, float], float], xr.DataArray] 

79 | Callable[[type, Union[int, float], float], xu.UgridDataArray] 

80 ), 

81 cellsize: float, 

82) -> imod.mf6.GroundwaterFlowModel: 

83 gwf_model = imod.mf6.GroundwaterFlowModel() 

84 

85 grid_data_function(np.float64, 1, cellsize) 

86 

87 # Create a constant head field. Some cells should not be constant head in order to have a meaningfull simulation 

88 constant_head = grid_data_function(np.float64, 1, cellsize) 

89 gwf_model["chd"] = imod.mf6.ConstantHead( 

90 constant_head, print_input=True, print_flows=True, save_flows=True 

91 ) 

92 gwf_model["ic"] = imod.mf6.InitialConditions(start=0.0) 

93 

94 icelltype = grid_data_function(np.int32, 1, cellsize) 

95 k = grid_data_function(np.float64, 1.23, cellsize) 

96 k33 = k.copy() 

97 

98 gwf_model["npf"] = imod.mf6.NodePropertyFlow( 

99 icelltype=icelltype, 

100 k=k, 

101 k33=k33, 

102 save_flows=True, 

103 ) 

104 gwf_model["sto"] = imod.mf6.SpecificStorage( 

105 specific_storage=grid_data_function(np.float64, 0.002, cellsize), 

106 specific_yield=0.15, 

107 transient=False, 

108 convertible=0, 

109 ) 

110 gwf_model["oc"] = imod.mf6.OutputControl(save_head="all", save_budget="all") 

111 rch_rate_all = grid_data_function(np.float64, 0.002, cellsize) 

112 rch_rate = rch_rate_all.sel(layer=[1]) 

113 

114 gwf_model["rch"] = imod.mf6.Recharge(rch_rate) 

115 

116 return gwf_model 

117 

118 

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

120def structured_flow_model() -> imod.mf6.GroundwaterFlowModel: 

121 cellsize = 2.0 

122 

123 gwf_model = _make_model(grid_data_structured, cellsize) 

124 

125 bottom = grid_data_structured_layered(np.float64, -1.0, cellsize) 

126 idomain = grid_data_structured(np.int32, 1, cellsize) 

127 

128 # Unassign the constant-head package from some cells to have a more meaningful simulation 

129 constant_head = gwf_model["chd"].dataset["head"] 

130 constant_head.loc[{"x": cellsize, "y": 2 * cellsize, "layer": 1}] = np.nan 

131 constant_head.loc[{"x": 2 * cellsize, "y": cellsize, "layer": 1}] = np.nan 

132 constant_head.loc[{"x": cellsize, "y": 2 * cellsize, "layer": 1}] = np.nan 

133 constant_head.loc[{"x": 2 * cellsize, "y": 2 * cellsize, "layer": 1}] = np.nan 

134 

135 gwf_model["dis"] = imod.mf6.StructuredDiscretization( 

136 top=10.0, bottom=bottom, idomain=idomain 

137 ) 

138 return gwf_model 

139 

140 

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

142def unstructured_flow_model() -> imod.mf6.GroundwaterFlowModel: 

143 cellsize = 2.0 

144 

145 gwf_model = _make_model(grid_data_unstructured, cellsize) 

146 

147 # Unassign the constant-head package from some cells to have a more meaningful simulation 

148 constant_head = gwf_model["chd"].dataset["head"] 

149 constant_head.loc[{"mesh2d_nFaces": 12, "layer": 1}] = np.nan 

150 constant_head.loc[{"mesh2d_nFaces": 13, "layer": 1}] = np.nan 

151 constant_head.loc[{"mesh2d_nFaces": 12, "layer": 2}] = np.nan 

152 constant_head.loc[{"mesh2d_nFaces": 13, "layer": 2}] = np.nan 

153 

154 bottom = grid_data_unstructured_layered(np.float64, -1.0, cellsize) 

155 idomain = grid_data_unstructured(np.int32, 1, cellsize) 

156 gwf_model["disv"] = imod.mf6.VerticesDiscretization( 

157 top=10.0, bottom=bottom, idomain=idomain 

158 ) 

159 return gwf_model 

160 

161 

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

163def solution_settings() -> imod.mf6.Solution: 

164 return imod.mf6.Solution( 

165 modelnames=["flow"], 

166 print_option="summary", 

167 outer_dvclose=1.0e-4, 

168 outer_maximum=500, 

169 under_relaxation=None, 

170 inner_dvclose=1.0e-4, 

171 inner_rclose=0.001, 

172 inner_maximum=100, 

173 linear_acceleration="cg", 

174 scaling_method=None, 

175 reordering_method=None, 

176 relaxation_factor=0.97, 

177 ) 

178 

179 

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

181def structured_flow_simulation( 

182 structured_flow_model: imod.mf6.GroundwaterFlowModel, 

183 solution_settings: imod.mf6.Solution, 

184) -> imod.mf6.Modflow6Simulation: 

185 simulation = imod.mf6.Modflow6Simulation("original_simulation") 

186 simulation["flow"] = structured_flow_model 

187 simulation["solution"] = solution_settings 

188 simulation.create_time_discretization( 

189 additional_times=["2000-01-01T00:00", "2020-01-02T00:00", "2020-01-03T00:00"] 

190 ) 

191 return simulation 

192 

193 

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

195def structured_flow_simulation_2_flow_models( 

196 structured_flow_simulation: imod.mf6.Modflow6Simulation, 

197) -> imod.mf6.Modflow6Simulation: 

198 """Returns transient confined model.""" 

199 other_flow_model = deepcopy(structured_flow_simulation["flow"]) 

200 structured_flow_simulation["flow_copy"] = other_flow_model 

201 structured_flow_simulation["solution"].add_model_to_solution("flow_copy") 

202 

203 return structured_flow_simulation 

204 

205 

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

207def unstructured_flow_simulation( 

208 unstructured_flow_model: imod.mf6.GroundwaterFlowModel, 

209 solution_settings: imod.mf6.Solution, 

210) -> imod.mf6.Modflow6Simulation: 

211 simulation = imod.mf6.Modflow6Simulation("original_simulation") 

212 simulation["flow"] = unstructured_flow_model 

213 simulation["solution"] = solution_settings 

214 simulation.create_time_discretization( 

215 additional_times=["2000-01-01T00:00", "2020-01-02T00:00", "2020-01-03T00:00"] 

216 ) 

217 return simulation