Coverage for fixtures\mf6_lake_package_fixture.py: 100%

79 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 

4import xugrid as xu 

5 

6from imod.mf6.lak import Lake, LakeData, OutletManning 

7from imod.mf6.write_context import WriteContext 

8 

9 

10def create_lake_table( 

11 number_rows, starting_stage, starting_sarea, starting_volume, starting_barea=None 

12): 

13 """ 

14 creates a lake table. The mandatory columns are stage, area and volume. Optionally a barea column 

15 is present. The number of rows of the table is an input argument. The first row of the table will 

16 have the starting values specified in the input arguments. Each next row will have the values of the 

17 previous row, plus one. 

18 """ 

19 rownumbers = np.arange(number_rows) 

20 

21 lake_table = xr.Dataset(coords={"row": rownumbers}) 

22 stage_data = (rownumbers + starting_stage).astype(np.float64) 

23 volume_data = (rownumbers + starting_volume).astype(np.float64) 

24 sarea_data = (rownumbers + starting_sarea).astype(np.float64) 

25 

26 lake_table["stage"] = xr.DataArray( 

27 coords={"row": rownumbers}, 

28 data=stage_data, 

29 ) 

30 lake_table["volume"] = xr.DataArray(coords={"row": rownumbers}, data=volume_data) 

31 

32 lake_table["sarea"] = xr.DataArray(coords={"row": rownumbers}, data=sarea_data) 

33 

34 if starting_barea is not None: 

35 barea_data = (rownumbers + starting_barea).astype(np.float64) 

36 lake_table["barea"] = xr.DataArray(coords={"row": rownumbers}, data=barea_data) 

37 

38 lake_table_array = lake_table.to_array() 

39 lake_table_array = lake_table_array.rename({"variable": "column"}) 

40 return lake_table_array 

41 

42 

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

44def naardermeer(basic_dis): 

45 def _naardermeer(has_lake_table=False): 

46 idomain, _, _ = basic_dis 

47 is_lake = xr.full_like(idomain, False, dtype=bool) 

48 is_lake[0, 1, 1] = True 

49 is_lake[0, 1, 2] = True 

50 is_lake[0, 2, 2] = True 

51 lake_table = None 

52 if has_lake_table: 

53 lake_table = create_lake_table(3, 10, 20, 30) 

54 return create_lake_data_structured( 

55 is_lake, starting_stage=11.0, name="Naardermeer", lake_table=lake_table 

56 ) 

57 

58 return _naardermeer 

59 

60 

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

62def ijsselmeer(basic_dis): 

63 def _ijsselmeer(has_lake_table=False): 

64 idomain, _, _ = basic_dis 

65 is_lake = xr.full_like(idomain, False, dtype=bool) 

66 is_lake[0, 4, 4] = True 

67 is_lake[0, 4, 5] = True 

68 is_lake[0, 5, 5] = True 

69 lake_table = None 

70 if has_lake_table: 

71 lake_table = create_lake_table(6, 8, 9, 10, 11) 

72 return create_lake_data_structured( 

73 is_lake, starting_stage=15.0, name="IJsselmeer", lake_table=lake_table 

74 ) 

75 

76 return _ijsselmeer 

77 

78 

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

80def lake_package(naardermeer, ijsselmeer): 

81 outlet1 = OutletManning("Naardermeer", "IJsselmeer", 23.0, 24.0, 25.0, 26.0) 

82 outlet2 = OutletManning("IJsselmeer", "Naardermeer", 27.0, 28.0, 29.0, 30.0) 

83 return Lake.from_lakes_and_outlets( 

84 [naardermeer(), ijsselmeer()], [outlet1, outlet2] 

85 ) 

86 

87 

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

89def lake_table(): 

90 return create_lake_table(5, 2.1, 3.1415, 100) 

91 

92 

93def create_lake_data_structured( 

94 is_lake, 

95 starting_stage, 

96 name, 

97 status=None, 

98 stage=None, 

99 rainfall=None, 

100 evaporation=None, 

101 runoff=None, 

102 inflow=None, 

103 withdrawal=None, 

104 auxiliary=None, 

105 lake_table=None, 

106): 

107 HORIZONTAL = 0 

108 connection_type = xr.full_like(is_lake, HORIZONTAL, dtype=np.float64).where(is_lake) 

109 bed_leak = xr.full_like(is_lake, 0.2, dtype=np.float64).where(is_lake) 

110 top_elevation = xr.full_like(is_lake, 0.3, dtype=np.float64).where(is_lake) 

111 bot_elevation = xr.full_like(is_lake, 0.4, dtype=np.float64).where(is_lake) 

112 connection_length = xr.full_like(is_lake, 0.5, dtype=np.float64).where(is_lake) 

113 connection_width = xr.full_like(is_lake, 0.6, dtype=np.float64).where(is_lake) 

114 return LakeData( 

115 starting_stage=starting_stage, 

116 boundname=name, 

117 connection_type=connection_type, 

118 bed_leak=bed_leak, 

119 top_elevation=top_elevation, 

120 bot_elevation=bot_elevation, 

121 connection_length=connection_length, 

122 connection_width=connection_width, 

123 status=status, 

124 stage=stage, 

125 rainfall=rainfall, 

126 evaporation=evaporation, 

127 runoff=runoff, 

128 inflow=inflow, 

129 withdrawal=withdrawal, 

130 auxiliary=auxiliary, 

131 lake_table=lake_table, 

132 ) 

133 

134 

135def write_and_read(package, path, filename, globaltimes=None) -> str: 

136 write_context = WriteContext(write_directory=path) 

137 package.write(filename, globaltimes, write_context) 

138 with open(path / f"{filename}.lak") as f: 

139 actual = f.read() 

140 return actual 

141 

142 

143def create_lake_data_unstructured( 

144 is_lake, 

145 starting_stage, 

146 name, 

147 status=None, 

148 stage=None, 

149 rainfall=None, 

150 evaporation=None, 

151 runoff=None, 

152 inflow=None, 

153 withdrawal=None, 

154 auxiliary=None, 

155 lake_table=None, 

156): 

157 HORIZONTAL = 0 

158 connection_type = xu.full_like(is_lake, HORIZONTAL, dtype=np.float64).where(is_lake) 

159 bed_leak = xu.full_like(is_lake, 0.2, dtype=np.float64).where(is_lake) 

160 top_elevation = xu.full_like(is_lake, 0.3, dtype=np.float64).where(is_lake) 

161 bot_elevation = xu.full_like(is_lake, 0.4, dtype=np.float64).where(is_lake) 

162 connection_length = xu.full_like(is_lake, 0.5, dtype=np.float64).where(is_lake) 

163 connection_width = xu.full_like(is_lake, 0.6, dtype=np.float64).where(is_lake) 

164 return LakeData( 

165 starting_stage=starting_stage, 

166 boundname=name, 

167 connection_type=connection_type, 

168 bed_leak=bed_leak, 

169 top_elevation=top_elevation, 

170 bot_elevation=bot_elevation, 

171 connection_length=connection_length, 

172 connection_width=connection_width, 

173 status=status, 

174 stage=stage, 

175 rainfall=rainfall, 

176 evaporation=evaporation, 

177 runoff=runoff, 

178 inflow=inflow, 

179 withdrawal=withdrawal, 

180 auxiliary=auxiliary, 

181 lake_table=lake_table, 

182 )