Coverage for src/CSET/operators/__init__.py: 100%

92 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-28 17:20 +0000

1# © Crown copyright, Met Office (2022-2026) and CSET contributors. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Subpackage contains all of CSET's operators.""" 

16 

17import inspect 

18import json 

19import logging 

20import os 

21import zipfile 

22from pathlib import Path 

23 

24from iris import FUTURE 

25 

26# Import operators here so they are exported for use by recipes. 

27import CSET.operators 

28from CSET.operators import ( 

29 ageofair, 

30 aggregate, 

31 aviation, 

32 collapse, 

33 constraints, 

34 convection, 

35 ensembles, 

36 feature, 

37 filters, 

38 fluxes, 

39 humidity, 

40 imageprocessing, 

41 mesoscale, 

42 misc, 

43 plot, 

44 power_spectrum, 

45 precipitation, 

46 pressure, 

47 read, 

48 regrid, 

49 scoreswrappers, 

50 temperature, 

51 transect, 

52 wind, 

53 write, 

54) 

55 

56# Exported operators & functions to use elsewhere. 

57__all__ = [ 

58 "ageofair", 

59 "aggregate", 

60 "aviation", 

61 "collapse", 

62 "constraints", 

63 "convection", 

64 "ensembles", 

65 "execute_recipe", 

66 "feature", 

67 "filters", 

68 "fluxes", 

69 "get_operator", 

70 "humidity", 

71 "imageprocessing", 

72 "mesoscale", 

73 "misc", 

74 "plot", 

75 "power_spectrum", 

76 "precipitation", 

77 "pressure", 

78 "read", 

79 "regrid", 

80 "scoreswrappers", 

81 "temperature", 

82 "transect", 

83 "wind", 

84 "write", 

85] 

86 

87logger = logging.getLogger(__name__) 

88 

89# Stop iris giving a warning whenever it loads something. 

90FUTURE.datum_support = True 

91# Stop iris giving a warning whenever it saves something. 

92FUTURE.save_split_attrs = True 

93# Accept microsecond precision in iris times. 

94FUTURE.date_microseconds = True 

95 

96 

97def get_operator(name: str): 

98 """Get an operator by its name. 

99 

100 Parameters 

101 ---------- 

102 name: str 

103 The name of the desired operator. 

104 

105 Returns 

106 ------- 

107 function 

108 The named operator. 

109 

110 Raises 

111 ------ 

112 ValueError 

113 If name is not an operator. 

114 

115 Examples 

116 -------- 

117 >>> CSET.operators.get_operator("read.read_cubes") 

118 <function read_cubes at 0x7fcf9353c8b0> 

119 """ 

120 logger.debug("get_operator(%s)", name) 

121 try: 

122 name_sections = name.split(".") 

123 operator = CSET.operators 

124 for section in name_sections: 

125 operator = getattr(operator, section) 

126 if callable(operator): 

127 return operator 

128 else: 

129 raise TypeError 

130 except (AttributeError, TypeError) as err: 

131 raise ValueError(f"Unknown operator: {name}") from err 

132 

133 

134def _write_metadata(recipe: dict): 

135 """Write a meta.json file in the CWD.""" 

136 metadata = recipe.copy() 

137 # Remove steps, as not needed, and might contain non-serialisable types. 

138 metadata.pop("steps", None) 

139 # To remove long variable names with suffix 

140 if "title" in metadata: 

141 metadata["title"] = metadata["title"].replace("_for_climate_averaging", "") 

142 metadata["title"] = metadata["title"].replace("_radiative_timestep", "") 

143 metadata["title"] = metadata["title"].replace("_maximum_random_overlap", "") 

144 metadata["title"] = metadata["title"].replace( 

145 "_for_surface_roughness_length_for_momentum_in_air", "" 

146 ) 

147 metadata["title"] = metadata["title"].replace( 

148 "sea_surface_wind_wave_mean_period_from_variance_spectral_density_first_frequency_moment", 

149 "wave_mean_period_first_moment", 

150 ) 

151 metadata["title"] = metadata["title"].replace( 

152 "sea_surface_wind_wave_mean_period_from_variance_spectral_density_second_frequency_moment", 

153 "wave_mean_period_second_moment", 

154 ) 

155 with open("meta.json", "wt", encoding="UTF-8") as fp: 

156 json.dump(metadata, fp, indent=2) 

157 

158 

159def _step_parser(step: dict, step_input: any) -> str: 

160 """Execute a recipe step, recursively executing any sub-steps.""" 

161 logger.debug("Executing step: %s", step) 

162 kwargs = {} 

163 for key, value in step.items(): 

164 if key == "operator": 

165 operator = get_operator(value) 

166 logger.info("operator: %s", value) 

167 elif isinstance(value, dict) and "operator" in value: 

168 logger.debug("Recursing into argument: %s", key) 

169 kwargs[key] = _step_parser(value, step_input) 

170 else: 

171 kwargs[key] = value 

172 logger.debug("args: %s", kwargs) 

173 logger.debug("step_input: %s", step_input) 

174 # If first argument of operator is explicitly defined, use that rather 

175 # than step_input. This is known through introspection of the operator. 

176 first_arg = next(iter(inspect.signature(operator).parameters.keys())) 

177 logger.debug("first_arg: %s", first_arg) 

178 if first_arg not in kwargs: 

179 logger.debug("first_arg not in kwargs, using step_input.") 

180 return operator(step_input, **kwargs) 

181 else: 

182 logger.debug("first_arg in kwargs.") 

183 return operator(**kwargs) 

184 

185 

186def create_diagnostic_archive(): 

187 """Create archive for easy download of plots and data.""" 

188 output_directory: Path = Path.cwd() 

189 archive_path = output_directory / "diagnostic.zip" 

190 with zipfile.ZipFile( 

191 archive_path, "w", compression=zipfile.ZIP_DEFLATED 

192 ) as archive: 

193 for file in output_directory.rglob("*"): 

194 # Check the archive doesn't add itself. 

195 if not file.samefile(archive_path): 

196 archive.write(file, arcname=file.relative_to(output_directory)) 

197 

198 

199def execute_recipe( 

200 recipe: dict, 

201 output_directory: Path, 

202 style_file: Path | None = None, 

203 plot_resolution: int | None = None, 

204 skip_write: bool | None = None, 

205) -> None: 

206 """Parse and executes the steps from a recipe file. 

207 

208 Parameters 

209 ---------- 

210 recipe: dict 

211 Parsed recipe. 

212 output_directory: Path 

213 Pathlike indicating desired location of output. 

214 style_file: Path, optional 

215 Path to a style file. 

216 plot_resolution: int, optional 

217 Resolution of plots in dpi. 

218 skip_write: bool, optional 

219 Skip saving processed output alongside plots. 

220 

221 Raises 

222 ------ 

223 FileNotFoundError 

224 The recipe or input file cannot be found. 

225 FileExistsError 

226 The output directory as actually a file. 

227 ValueError 

228 The recipe is not well formed. 

229 TypeError 

230 The provided recipe is not a stream or Path. 

231 """ 

232 # Create output directory. 

233 try: 

234 output_directory.mkdir(parents=True, exist_ok=True) 

235 except (FileExistsError, NotADirectoryError): 

236 logger.error("Output directory is a file. %s", output_directory) 

237 raise 

238 steps = recipe["steps"] 

239 

240 # Execute the steps in a recipe. 

241 original_working_directory = Path.cwd() 

242 try: 

243 os.chdir(output_directory) 

244 diagnostic_log = logging.FileHandler( 

245 filename="CSET.log", mode="w", encoding="UTF-8" 

246 ) 

247 diagnostic_log.setFormatter( 

248 logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s") 

249 ) 

250 logger.addHandler(diagnostic_log) 

251 # Create metadata file used by some steps. 

252 if style_file: 

253 recipe["style_file_path"] = str(style_file) 

254 if plot_resolution: 

255 recipe["plot_resolution"] = plot_resolution 

256 if skip_write: 

257 recipe["skip_write"] = skip_write 

258 _write_metadata(recipe) 

259 

260 # Execute the recipe. 

261 step_input = None 

262 for step in steps: 

263 step_input = _step_parser(step, step_input) 

264 logger.info("Recipe output:\n%s", step_input) 

265 

266 logger.info("Creating diagnostic archive.") 

267 create_diagnostic_archive() 

268 finally: 

269 os.chdir(original_working_directory)