Coverage for src/CSET/operators/__init__.py: 100%
89 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-06 10:09 +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.
15"""Subpackage contains all of CSET's operators."""
17import inspect
18import json
19import logging
20import os
21import zipfile
22from pathlib import Path
24from iris import FUTURE
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 curvature,
36 ensembles,
37 feature,
38 filters,
39 fluxes,
40 humidity,
41 imageprocessing,
42 mesoscale,
43 misc,
44 plot,
45 power_spectrum,
46 precipitation,
47 pressure,
48 read,
49 regrid,
50 scoreswrappers,
51 temperature,
52 transect,
53 wind,
54 write,
55)
57# Exported operators & functions to use elsewhere.
58__all__ = [
59 "ageofair",
60 "aggregate",
61 "aviation",
62 "collapse",
63 "constraints",
64 "convection",
65 "curvature",
66 "ensembles",
67 "execute_recipe",
68 "feature",
69 "filters",
70 "fluxes",
71 "get_operator",
72 "humidity",
73 "imageprocessing",
74 "mesoscale",
75 "misc",
76 "plot",
77 "power_spectrum",
78 "precipitation",
79 "pressure",
80 "read",
81 "regrid",
82 "scoreswrappers",
83 "temperature",
84 "transect",
85 "wind",
86 "write",
87]
89logger = logging.getLogger(__name__)
91# Stop iris giving a warning whenever it loads something.
92FUTURE.datum_support = True
93# Stop iris giving a warning whenever it saves something.
94FUTURE.save_split_attrs = True
95# Accept microsecond precision in iris times.
96FUTURE.date_microseconds = True
99def get_operator(name: str):
100 """Get an operator by its name.
102 Parameters
103 ----------
104 name: str
105 The name of the desired operator.
107 Returns
108 -------
109 function
110 The named operator.
112 Raises
113 ------
114 ValueError
115 If name is not an operator.
117 Examples
118 --------
119 >>> CSET.operators.get_operator("read.read_cubes")
120 <function read_cubes at 0x7fcf9353c8b0>
121 """
122 logger.debug("get_operator(%s)", name)
123 try:
124 name_sections = name.split(".")
125 operator = CSET.operators
126 for section in name_sections:
127 operator = getattr(operator, section)
128 if callable(operator):
129 return operator
130 else:
131 raise TypeError
132 except (AttributeError, TypeError) as err:
133 raise ValueError(f"Unknown operator: {name}") from err
136def _write_metadata(recipe: dict):
137 """Write a meta.json file in the CWD."""
138 metadata = recipe.copy()
139 # Remove steps, as not needed, and might contain non-serialisable types.
140 metadata.pop("steps", None)
141 # To remove long variable names with suffix
142 if "title" in metadata:
143 metadata["title"] = metadata["title"].replace("_for_climate_averaging", "")
144 metadata["title"] = metadata["title"].replace("_radiative_timestep", "")
145 metadata["title"] = metadata["title"].replace("_maximum_random_overlap", "")
146 with open("meta.json", "wt", encoding="UTF-8") as fp:
147 json.dump(metadata, fp, indent=2)
150def _step_parser(step: dict, step_input: any) -> str:
151 """Execute a recipe step, recursively executing any sub-steps."""
152 logger.debug("Executing step: %s", step)
153 kwargs = {}
154 for key, value in step.items():
155 if key == "operator":
156 operator = get_operator(value)
157 logger.info("operator: %s", value)
158 elif isinstance(value, dict) and "operator" in value:
159 logger.debug("Recursing into argument: %s", key)
160 kwargs[key] = _step_parser(value, step_input)
161 else:
162 kwargs[key] = value
163 logger.debug("args: %s", kwargs)
164 logger.debug("step_input: %s", step_input)
165 # If first argument of operator is explicitly defined, use that rather
166 # than step_input. This is known through introspection of the operator.
167 first_arg = next(iter(inspect.signature(operator).parameters.keys()))
168 logger.debug("first_arg: %s", first_arg)
169 if first_arg not in kwargs:
170 logger.debug("first_arg not in kwargs, using step_input.")
171 return operator(step_input, **kwargs)
172 else:
173 logger.debug("first_arg in kwargs.")
174 return operator(**kwargs)
177def create_diagnostic_archive():
178 """Create archive for easy download of plots and data."""
179 output_directory: Path = Path.cwd()
180 archive_path = output_directory / "diagnostic.zip"
181 with zipfile.ZipFile(
182 archive_path, "w", compression=zipfile.ZIP_DEFLATED
183 ) as archive:
184 for file in output_directory.rglob("*"):
185 # Check the archive doesn't add itself.
186 if not file.samefile(archive_path):
187 archive.write(file, arcname=file.relative_to(output_directory))
190def execute_recipe(
191 recipe: dict,
192 output_directory: Path,
193 style_file: Path | None = None,
194 plot_resolution: int | None = None,
195 skip_write: bool | None = None,
196) -> None:
197 """Parse and executes the steps from a recipe file.
199 Parameters
200 ----------
201 recipe: dict
202 Parsed recipe.
203 output_directory: Path
204 Pathlike indicating desired location of output.
205 style_file: Path, optional
206 Path to a style file.
207 plot_resolution: int, optional
208 Resolution of plots in dpi.
209 skip_write: bool, optional
210 Skip saving processed output alongside plots.
212 Raises
213 ------
214 FileNotFoundError
215 The recipe or input file cannot be found.
216 FileExistsError
217 The output directory as actually a file.
218 ValueError
219 The recipe is not well formed.
220 TypeError
221 The provided recipe is not a stream or Path.
222 """
223 # Create output directory.
224 try:
225 output_directory.mkdir(parents=True, exist_ok=True)
226 except (FileExistsError, NotADirectoryError):
227 logger.error("Output directory is a file. %s", output_directory)
228 raise
229 steps = recipe["steps"]
231 # Execute the steps in a recipe.
232 original_working_directory = Path.cwd()
233 try:
234 os.chdir(output_directory)
235 diagnostic_log = logging.FileHandler(
236 filename="CSET.log", mode="w", encoding="UTF-8"
237 )
238 diagnostic_log.setFormatter(
239 logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s")
240 )
241 logger.addHandler(diagnostic_log)
242 # Create metadata file used by some steps.
243 if style_file:
244 recipe["style_file_path"] = str(style_file)
245 if plot_resolution:
246 recipe["plot_resolution"] = plot_resolution
247 if skip_write:
248 recipe["skip_write"] = skip_write
249 _write_metadata(recipe)
251 # Execute the recipe.
252 step_input = None
253 for step in steps:
254 step_input = _step_parser(step, step_input)
255 logger.info("Recipe output:\n%s", step_input)
257 logger.info("Creating diagnostic archive.")
258 create_diagnostic_archive()
259 finally:
260 os.chdir(original_working_directory)