Coverage for C:\src\imod-python\imod\msw\meteo_grid.py: 29%

68 statements  

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

1import csv 

2from pathlib import Path 

3from typing import Optional, Union 

4 

5import numpy as np 

6import pandas as pd 

7import xarray as xr 

8 

9import imod 

10from imod.msw.pkgbase import MetaSwapPackage 

11from imod.msw.timeutil import to_metaswap_timeformat 

12 

13 

14class MeteoGrid(MetaSwapPackage): 

15 """ 

16 This contains the meteorological grid data. Grids are written to ESRI ASCII 

17 files. The meteorological data requires a time coordinate. Next to a 

18 MeteoGrid instance, instances of PrecipitationMapping and 

19 EvapotranspirationMapping are required as well to specify meteorological 

20 information to MetaSWAP. 

21 

22 This class is responsible for `mete_grid.inp`. 

23 

24 Parameters 

25 ---------- 

26 precipitation: array of floats (xr.DataArray) 

27 Contains the precipitation grids in mm/d. A time coordinate is required. 

28 evapotranspiration: array of floats (xr.DataArray) 

29 Contains the evapotranspiration grids in mm/d. A time coordinate is 

30 required. 

31 """ 

32 

33 _file_name = "mete_grid.inp" 

34 _meteo_dirname = "meteo_grids" 

35 

36 def __init__(self, precipitation: xr.DataArray, evapotranspiration: xr.DataArray): 

37 super().__init__() 

38 

39 self.dataset["precipitation"] = precipitation 

40 self.dataset["evapotranspiration"] = evapotranspiration 

41 

42 self._pkgcheck() 

43 

44 def write_free_format_file(self, path: Union[str, Path], dataframe: pd.DataFrame): 

45 """ 

46 Write free format file. The mete_grid.inp file is free format. 

47 """ 

48 

49 columns = list(self.dataset.data_vars) 

50 

51 dataframe.loc[:, columns] = '"' + dataframe[columns] + '"' 

52 # Add required columns, which we will not use. 

53 # These are only used when WOFOST is used 

54 # TODO: Add support for temperature to allow WOFOST support 

55 wofost_columns = [ 

56 "minimum_day_temperature", 

57 "maximum_day_temperature", 

58 "mean_temperature", 

59 ] 

60 dataframe.loc[:, wofost_columns] = '"NoValue"' 

61 

62 self.check_string_lengths(dataframe) 

63 

64 dataframe.to_csv( 

65 path, header=False, quoting=csv.QUOTE_NONE, float_format="%.4f", index=False 

66 ) 

67 

68 def _compose_filename( 

69 self, d: dict, directory: Path, pattern: Optional[str] = None 

70 ): 

71 """ 

72 Construct a filename, following the iMOD conventions. 

73 

74 

75 Parameters 

76 ---------- 

77 d : dict 

78 dict of parts (time, layer) for filename. 

79 pattern : string or re.pattern 

80 Format to create pattern for. 

81 

82 Returns 

83 ------- 

84 str 

85 Absolute path. 

86 

87 """ 

88 return str(directory / imod.util.path.compose(d, pattern)) 

89 

90 def _is_grid(self, varname: str): 

91 coords = self.dataset[varname].coords 

92 

93 if "y" not in coords and "x" not in coords: 

94 return False 

95 else: 

96 return True 

97 

98 def _compose_dataframe(self, times: np.array): 

99 dataframe = pd.DataFrame(index=times) 

100 

101 year, time_since_start_year = to_metaswap_timeformat(times) 

102 

103 dataframe["time_since_start_year"] = time_since_start_year 

104 dataframe["year"] = year 

105 

106 # Data dir is always relative to model dir, so don't use model directory 

107 # here 

108 data_dir = Path(".") / self._meteo_dirname 

109 

110 for varname in self.dataset.data_vars: 

111 # If grid, we have to add the filename of the .asc to be written 

112 if self._is_grid(varname): 

113 dataframe[varname] = [ 

114 self._compose_filename( 

115 dict(time=time, name=varname, extension=".asc"), 

116 directory=data_dir, 

117 ) 

118 for time in times 

119 ] 

120 else: 

121 dataframe[varname] = self.dataset[varname].values.astype(str) 

122 

123 return dataframe 

124 

125 def check_string_lengths(self, dataframe: pd.DataFrame): 

126 """ 

127 Check if strings lengths do not exceed 256 characters. 

128 With absolute paths this might be an issue. 

129 """ 

130 

131 # Because two quote marks are added later. 

132 character_limit = 254 

133 

134 columns = list(self.dataset.data_vars) 

135 

136 str_too_long = [ 

137 np.any(dataframe[varname].str.len() > character_limit) 

138 for varname in columns 

139 ] 

140 

141 if any(str_too_long): 

142 indexes_true = np.where(str_too_long)[0] 

143 too_long_columns = list(np.array(columns)[indexes_true]) 

144 raise ValueError( 

145 f"Encountered strings longer than 256 characters in columns: {too_long_columns}" 

146 ) 

147 

148 def write(self, directory: Union[str, Path], *args): 

149 """ 

150 Write mete_grid.inp and accompanying ASCII grid files. 

151 

152 Parameters 

153 ---------- 

154 directory: str or Path 

155 directory to write file in. 

156 """ 

157 

158 directory = Path(directory) 

159 

160 times = self.dataset["time"].values 

161 

162 dataframe = self._compose_dataframe(times) 

163 self.write_free_format_file(directory / self._file_name, dataframe) 

164 

165 # Write grid data to ESRI ASCII files 

166 for varname in self.dataset.data_vars: 

167 if self._is_grid(varname): 

168 path = (directory / self._meteo_dirname / varname).with_suffix(".asc") 

169 imod.rasterio.save(path, self.dataset[varname], nodata=-9999.0) 

170 

171 def _pkgcheck(self): 

172 for varname in self.dataset.data_vars: 

173 coords = self.dataset[varname].coords 

174 if "time" not in coords: 

175 raise ValueError(f"No 'time' coordinate included in {varname}") 

176 

177 allowed_dims = ["time", "y", "x"] 

178 

179 excess_dims = set(self.dataset[varname].dims) - set(allowed_dims) 

180 if len(excess_dims) > 0: 

181 raise ValueError( 

182 f"Received excess dims {excess_dims} in {self.__class__} for " 

183 f"{varname}, please provide data with {allowed_dims}" 

184 )