Coverage for test_benchmark.py: 100%

108 statements  

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

1# keep track of performance using pytest-benchmark 

2# https://pytest-benchmark.readthedocs.io/en/stable/ 

3 

4import flopy 

5import numpy as np 

6import pytest 

7import xarray as xr 

8 

9import imod 

10 

11 

12def setup_mf6_basic_simulation_flopy(): 

13 """ 

14 Defines a basic FloPy Modflow 6 model. 

15 

16 The code is based on this example notebook from the FloPy documentation: 

17 https://github.com/modflowpy/flopy/blob/develop/examples/Notebooks/flopy3_mf6_A_simple-model.ipynb 

18 

19 The only real change is to increase the number of layers from 10 to 100 to get 1 million cells. 

20 """ 

21 # For this example, we will set up a model workspace. 

22 # Model input files and output files will reside here. 

23 

24 name = "mf6lake" 

25 h1 = 100.0 

26 h2 = 90.0 

27 Nlay = 10 

28 N = 101 

29 L = 400.0 

30 H = 50.0 

31 k = 1.0 

32 

33 # Create the Flopy simulation object 

34 sim = flopy.mf6.MFSimulation(sim_name=name, exe_name="mf6", version="mf6") 

35 

36 # Create the Flopy temporal discretization object 

37 flopy.mf6.modflow.mftdis.ModflowTdis( 

38 sim, pname="tdis", time_units="DAYS", nper=1, perioddata=[(1.0, 1, 1.0)] 

39 ) 

40 

41 # Create the Flopy groundwater flow (gwf) model object 

42 model_nam_file = f"{name}.nam" 

43 gwf = flopy.mf6.ModflowGwf(sim, modelname=name, model_nam_file=model_nam_file) 

44 

45 # Create the Flopy iterative model solver (ims) Package object 

46 flopy.mf6.modflow.mfims.ModflowIms(sim, pname="ims", complexity="SIMPLE") 

47 

48 # Create the discretization package 

49 bot = np.linspace(-H / Nlay, -H, Nlay) 

50 delrow = delcol = L / (N - 1) 

51 flopy.mf6.modflow.mfgwfdis.ModflowGwfdis( 

52 gwf, 

53 pname="dis", 

54 nlay=Nlay, 

55 nrow=N, 

56 ncol=N, 

57 delr=delrow, 

58 delc=delcol, 

59 top=0.0, 

60 botm=bot, 

61 ) 

62 

63 # Create the initial conditions package 

64 start = h1 * np.ones((Nlay, N, N)) 

65 flopy.mf6.modflow.mfgwfic.ModflowGwfic(gwf, pname="ic", strt=start) 

66 

67 # Create the node property flow package 

68 flopy.mf6.modflow.mfgwfnpf.ModflowGwfnpf( 

69 gwf, pname="npf", icelltype=1, k=k, save_flows=True 

70 ) 

71 

72 # Create the constant head package. 

73 chd_rec = [] 

74 chd_rec.append(((0, int(N / 4), int(N / 4)), h2)) 

75 for layer in range(0, Nlay): 

76 for row_col in range(0, N): 

77 chd_rec.append(((layer, row_col, 0), h1)) 

78 chd_rec.append(((layer, row_col, N - 1), h1)) 

79 if row_col != 0 and row_col != N - 1: 

80 chd_rec.append(((layer, 0, row_col), h1)) 

81 chd_rec.append(((layer, N - 1, row_col), h1)) 

82 flopy.mf6.modflow.mfgwfchd.ModflowGwfchd( 

83 gwf, 

84 pname="chd", 

85 maxbound=len(chd_rec), 

86 stress_period_data=chd_rec, 

87 save_flows=True, 

88 ) 

89 

90 # Create the output control package 

91 headfile = f"{name}.hds" 

92 head_filerecord = [headfile] 

93 budgetfile = f"{name}.cbc" 

94 budget_filerecord = [budgetfile] 

95 saverecord = [("HEAD", "ALL"), ("BUDGET", "ALL")] 

96 printrecord = [("HEAD", "LAST")] 

97 flopy.mf6.modflow.mfgwfoc.ModflowGwfoc( 

98 gwf, 

99 pname="oc", 

100 saverecord=saverecord, 

101 head_filerecord=head_filerecord, 

102 budget_filerecord=budget_filerecord, 

103 printrecord=printrecord, 

104 ) 

105 return sim 

106 

107 

108def setup_mf6_basic_simulation_imod(): 

109 """ 

110 Defines a basic Modflow 6 model, to benchmark against the FloPy version above. 

111 """ 

112 

113 h1 = 100.0 

114 h2 = 90.0 

115 nlay = 10 

116 nrow = 101 

117 ncol = 101 

118 shape = (nlay, nrow, ncol) 

119 dy = -4.0 

120 dx = 4.0 

121 height = 50.0 

122 k = 1.0 

123 

124 xmax = dx * ncol 

125 ymax = abs(dy) * nrow 

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

127 layer = np.arange(1, nlay + 1) 

128 y = np.arange(ymax, 0.0, dy) + 0.5 * dy 

129 x = np.arange(0.0, xmax, dx) + 0.5 * dx 

130 coords = {"layer": layer, "y": y, "x": x} 

131 

132 # Discretization data 

133 like = xr.DataArray(np.ones(shape), coords=coords, dims=dims) 

134 idomain = like.astype(np.int32) 

135 bottom = xr.DataArray( 

136 np.linspace(-height / nlay, -height, nlay), {"layer": layer}, ("layer",) 

137 ) 

138 

139 # Constant head 

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

141 # set all side edges to h1 

142 head[:, [0, -1], :] = h1 

143 head[:, :, [0, -1]] = h1 

144 # set a single cell to h2 

145 head[0, nrow // 4, ncol // 4] = h2 

146 

147 # Create and fill the groundwater model. 

148 gwf_model = imod.mf6.GroundwaterFlowModel() 

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

150 top=0.0, bottom=bottom, idomain=idomain 

151 ) 

152 gwf_model["chd"] = imod.mf6.ConstantHead(head, save_flows=True) 

153 gwf_model["ic"] = imod.mf6.InitialConditions(start=100.0) 

154 gwf_model["npf"] = imod.mf6.NodePropertyFlow(icelltype=1, k=k, save_flows=True) 

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

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

157 specific_storage=1.0e-15, 

158 specific_yield=0.15, 

159 convertible=0, 

160 transient=False, 

161 ) 

162 

163 # Attach it to a simulation 

164 simulation = imod.mf6.Modflow6Simulation("mf6basic") 

165 simulation["gwf"] = gwf_model 

166 # Define solver settings 

167 simulation["solver"] = imod.mf6.SolutionPresetSimple( 

168 modelnames=["gwf"], 

169 print_option="summary", 

170 ) 

171 

172 # Collect time discretization 

173 simulation.create_time_discretization(additional_times=["2000-01-01", "2000-01-02"]) 

174 return simulation 

175 

176 

177@pytest.fixture(scope="module") 

178def mf6_basic_simulation_flopy(): 

179 return setup_mf6_basic_simulation_flopy() 

180 

181 

182@pytest.fixture(scope="module") 

183def mf6_basic_simulation_imod(): 

184 return setup_mf6_basic_simulation_imod() 

185 

186 

187def write_basic_flopy(mf6_basic_simulation_flopy, tmp_path): 

188 simulation = mf6_basic_simulation_flopy 

189 simulation.set_sim_path(str(tmp_path)) 

190 simulation.write_simulation() 

191 

192 

193def write_basic_imod_binary(mf6_basic_simulation_imod, tmp_path): 

194 simulation = mf6_basic_simulation_imod 

195 modeldir = tmp_path / "mf6-basic" 

196 simulation.write(modeldir) 

197 

198 

199def write_basic_imod_text(mf6_basic_simulation_imod, tmp_path): 

200 simulation = mf6_basic_simulation_imod 

201 modeldir = tmp_path / "mf6-basic" 

202 simulation.write(modeldir, binary=False) 

203 

204 

205def test_setup_basic_flopy(benchmark, mf6_basic_simulation_flopy, tmp_path): 

206 benchmark(setup_mf6_basic_simulation_flopy) 

207 

208 

209def test_setup_basic_imod(benchmark, mf6_basic_simulation_imod, tmp_path): 

210 benchmark(setup_mf6_basic_simulation_imod) 

211 

212 

213def test_write_basic_flopy(benchmark, mf6_basic_simulation_flopy, tmp_path): 

214 benchmark(write_basic_flopy, mf6_basic_simulation_flopy, tmp_path) 

215 

216 

217def test_write_basic_imod_binary(benchmark, mf6_basic_simulation_imod, tmp_path): 

218 benchmark(write_basic_imod_binary, mf6_basic_simulation_imod, tmp_path) 

219 

220 

221def test_write_basic_imod_text(benchmark, mf6_basic_simulation_imod, tmp_path): 

222 benchmark(write_basic_imod_text, mf6_basic_simulation_imod, tmp_path)