Coverage for fixtures\mf6_flow_with_transport_fixture.py: 100%

110 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1import numpy as np 

2import pytest 

3import xarray as xr 

4 

5from imod.mf6 import ( 

6 GroundwaterFlowModel, 

7 InitialConditions, 

8 NodePropertyFlow, 

9 OutputControl, 

10 River, 

11 SpecificStorage, 

12) 

13 

14globaltimes = np.array( 

15 [ 

16 "2000-01-01", 

17 "2000-01-02", 

18 "2000-01-03", 

19 ], 

20 dtype="datetime64[ns]", 

21) 

22 

23 

24class grid_dimensions: 

25 nlay = 3 

26 nrow = 15 

27 ncol = 15 

28 dx = 5000 

29 dy = -5000 

30 xmin = 0 

31 ymin = 0 

32 

33 

34def get_data_array(dimensions, globaltimes): 

35 ntimes = len(globaltimes) 

36 shape = (ntimes, dimensions.nlay, dimensions.nrow, dimensions.ncol) 

37 dims = ("time", "layer", "y", "x") 

38 

39 layer = np.array([1, 2, 3]) 

40 xmax = dimensions.dx * dimensions.ncol 

41 ymax = abs(dimensions.dy) * dimensions.nrow 

42 y = np.arange(ymax, dimensions.ymin, dimensions.dy) + 0.5 * dimensions.dy 

43 x = np.arange(dimensions.xmin, xmax, dimensions.dx) + 0.5 * dimensions.dx 

44 coords = {"time": globaltimes, "layer": layer, "y": y, "x": x} 

45 

46 # Discretization data 

47 return xr.DataArray( 

48 np.ones(shape), 

49 coords=coords, 

50 dims=dims, 

51 ) 

52 

53 

54@pytest.fixture(scope="session") 

55def head_fc(): 

56 idomain = get_data_array(grid_dimensions(), globaltimes) 

57 

58 # Constant head 

59 head = xr.full_like(idomain, np.nan) 

60 return head 

61 

62 

63@pytest.fixture(scope="session") 

64def concentration_fc(): 

65 idomain = get_data_array(grid_dimensions(), globaltimes) 

66 idomain = idomain.expand_dims(species=["salinity", "temperature"]) 

67 

68 concentration = xr.full_like(idomain, np.nan) 

69 return concentration 

70 

71 

72@pytest.fixture(scope="session") 

73def conductance_fc(): 

74 globaltimes = np.array( 

75 [ 

76 "2000-01-01", 

77 "2000-01-02", 

78 "2000-01-03", 

79 ], 

80 dtype="datetime64[ns]", 

81 ) 

82 idomain = get_data_array(grid_dimensions(), globaltimes) 

83 

84 # Constant head 

85 conductance = xr.full_like(idomain, np.nan) 

86 return conductance 

87 

88 

89@pytest.fixture(scope="session") 

90def elevation_fc(): 

91 idomain = get_data_array(grid_dimensions(), globaltimes) 

92 

93 elevation = xr.full_like(idomain, np.nan) 

94 elevation[:, 0, 7, 7:9] = 1.0 

95 

96 return elevation 

97 

98 

99@pytest.fixture(scope="session") 

100def rate_fc(): 

101 idomain = get_data_array(grid_dimensions(), globaltimes) 

102 

103 rate = xr.full_like(idomain, np.nan) 

104 rate[:, 0, 7, 7:9] = 0.001 

105 

106 return rate 

107 

108 

109@pytest.fixture(scope="session") 

110def proportion_rate_fc(): 

111 idomain = get_data_array(grid_dimensions(), globaltimes) 

112 

113 proportion_rate = xr.full_like(idomain, np.nan) 

114 proportion_rate[:, 0, 7, 7:9] = 0.3 

115 return proportion_rate 

116 

117 

118@pytest.fixture(scope="session") 

119def proportion_depth_fc(): 

120 idomain = get_data_array(grid_dimensions(), globaltimes) 

121 

122 proportion_depth = xr.full_like(idomain, np.nan) 

123 proportion_depth[:, 0, 7, 7:9] = 0.4 

124 return proportion_depth 

125 

126 

127@pytest.fixture(scope="session") 

128def porosity_fc(): 

129 idomain = get_data_array(grid_dimensions(), globaltimes) 

130 

131 porosity_fc = xr.full_like(idomain, np.nan).isel(time=0) 

132 return porosity_fc 

133 

134 

135@pytest.fixture(scope="session") 

136def decay_fc(): 

137 idomain = get_data_array(grid_dimensions(), globaltimes) 

138 

139 decay_fc = xr.full_like(idomain, np.nan).isel(time=0) 

140 return decay_fc 

141 

142 

143@pytest.fixture(scope="session") 

144def decay_sorbed_fc(): 

145 idomain = get_data_array(grid_dimensions(), globaltimes) 

146 

147 decay_sorbed_fc = xr.full_like(idomain, np.nan).isel(time=0) 

148 return decay_sorbed_fc 

149 

150 

151@pytest.fixture(scope="session") 

152def bulk_density_fc(): 

153 idomain = get_data_array(grid_dimensions(), globaltimes) 

154 

155 bulk_density_fc = xr.full_like(idomain, np.nan).isel(time=0) 

156 return bulk_density_fc 

157 

158 

159@pytest.fixture(scope="session") 

160def distcoef_fc(): 

161 idomain = get_data_array(grid_dimensions(), globaltimes) 

162 

163 distcoef_fc = xr.full_like(idomain, np.nan).isel(time=0) 

164 return distcoef_fc 

165 

166 

167@pytest.fixture(scope="session") 

168def sp2_fc(): 

169 idomain = get_data_array(grid_dimensions(), globaltimes) 

170 

171 sp2_fc = xr.full_like(idomain, np.nan).isel(time=0) 

172 return sp2_fc 

173 

174 

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

176@pytest.mark.usefixtures("concentration_fc") 

177def flow_model_with_concentration(concentration_fc): 

178 idomain = get_data_array(grid_dimensions(), globaltimes) 

179 cellType = xr.full_like(idomain.isel(time=0), 1, dtype=np.int32) 

180 k = xr.full_like(idomain.isel(time=0), 10.0) 

181 k33 = xr.full_like(idomain.isel(time=0), 10.0) 

182 

183 # River 

184 riv_dict = dict( 

185 stage=idomain.sel(layer=1), 

186 conductance=idomain.sel(layer=1), 

187 bottom_elevation=idomain.sel(layer=1) - 1.0, 

188 concentration=concentration_fc.sel(layer=1), 

189 ) 

190 

191 gwf_model = GroundwaterFlowModel() 

192 

193 gwf_model["npf"] = NodePropertyFlow( 

194 icelltype=cellType, 

195 k=k, 

196 k33=k33, 

197 ) 

198 

199 gwf_model["sto"] = SpecificStorage( 

200 specific_storage=1.0e-5, 

201 specific_yield=0.15, 

202 transient=False, 

203 convertible=0, 

204 ) 

205 gwf_model["ic"] = InitialConditions(start=0.0) 

206 gwf_model["oc"] = OutputControl(save_head="all", save_budget="all") 

207 gwf_model["riv-1"] = River( 

208 concentration_boundary_type="AUX", 

209 **riv_dict, 

210 ) 

211 

212 return gwf_model