Coverage for src/CSET/operators/read.py: 94%
436 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 09:36 +0000
1# © Crown copyright, Met Office (2022-2025) 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"""Operators for reading various types of files from disk."""
17import ast
18import datetime
19import functools
20import glob
21import itertools
22import logging
23from pathlib import Path
24from typing import Literal
26import dask
27import iris
28import iris.coord_systems
29import iris.coords
30import iris.cube
31import iris.exceptions
32import iris.util
33import numpy as np
34from iris.analysis.cartography import rotate_pole, rotate_winds
36from CSET._common import iter_maybe
37from CSET.operators._stash_to_lfric import STASH_TO_LFRIC
38from CSET.operators._utils import (
39 get_cube_coordindex,
40 get_cube_yxcoordname,
41 is_spatialdim,
42)
44logger = logging.getLogger(__name__)
47class NoDataError(FileNotFoundError):
48 """Error that no data has been loaded."""
51def read_cube(
52 file_paths: list[str] | str,
53 constraint: iris.Constraint = None,
54 model_names: list[str] | str | None = None,
55 subarea_type: str | None = None,
56 subarea_extent: list[float] | None = None,
57 **kwargs,
58) -> iris.cube.Cube:
59 """Read a single cube from files.
61 Read operator that takes a path string (can include shell-style glob
62 patterns), and loads the cube matching the constraint. If any paths point to
63 directory, all the files contained within are loaded.
65 Ensemble data can also be loaded. If it has a realization coordinate
66 already, it will be directly used. If not, it will have its member number
67 guessed from the filename, based on one of several common patterns. For
68 example the pattern *emXX*, where XX is the realization.
70 Deterministic data will be loaded with a realization of 0, allowing it to be
71 processed in the same way as ensemble data.
73 Arguments
74 ---------
75 file_paths: str | list[str]
76 Path or paths to where .pp/.nc files are located
77 constraint: iris.Constraint | iris.ConstraintCombination, optional
78 Constraints to filter data by. Defaults to unconstrained.
79 model_names: str | list[str], optional
80 Names of the models that correspond to respective paths in file_paths.
81 subarea_type: "gridcells" | "modelrelative" | "realworld", optional
82 Whether to constrain data by model relative coordinates or real world
83 coordinates.
84 subarea_extent: list, optional
85 List of coordinates to constraint data by, in order lower latitude,
86 upper latitude, lower longitude, upper longitude.
88 Returns
89 -------
90 cubes: iris.cube.Cube
91 Cube loaded
93 Raises
94 ------
95 FileNotFoundError
96 If the provided path does not exist
97 ValueError
98 If the constraint doesn't produce a single cube.
99 """
100 cubes = read_cubes(
101 file_paths=file_paths,
102 constraint=constraint,
103 model_names=model_names,
104 subarea_type=subarea_type,
105 subarea_extent=subarea_extent,
106 )
107 # Check filtered cubes is a CubeList containing one cube.
108 if len(cubes) == 1:
109 return cubes[0]
110 else:
111 raise ValueError(
112 f"Constraint doesn't produce single cube: {constraint}\n{cubes}"
113 )
116def read_cubes(
117 file_paths: list[str] | str,
118 constraint: iris.Constraint | None = None,
119 model_names: str | list[str] | None = None,
120 subarea_type: str | None = None,
121 subarea_extent: list | None = None,
122 **kwargs,
123) -> iris.cube.CubeList:
124 """Read cubes from files.
126 Read operator that takes a path string (can include shell-style glob
127 patterns), and loads the cubes matching the constraint. If any paths point
128 to directory, all the files contained within are loaded.
130 Ensemble data can also be loaded. If it has a realization coordinate
131 already, it will be directly used. If not, it will have its member number
132 guessed from the filename, based on one of several common patterns. For
133 example the pattern *emXX*, where XX is the realization.
135 Deterministic data will be loaded with a realization of 0, allowing it to be
136 processed in the same way as ensemble data.
138 Data output by XIOS (such as LFRic) has its per-file metadata removed so
139 that the cubes merge across files.
141 Arguments
142 ---------
143 file_paths: str | list[str]
144 Path or paths to where .pp/.nc files are located. Can include globs.
145 constraint: iris.Constraint | iris.ConstraintCombination, optional
146 Constraints to filter data by. Defaults to unconstrained.
147 model_names: str | list[str], optional
148 Names of the models that correspond to respective paths in file_paths.
149 subarea_type: str, optional
150 Whether to constrain data by model relative coordinates or real world
151 coordinates.
152 subarea_extent: list[float], optional
153 List of coordinates to constraint data by, in order lower latitude,
154 upper latitude, lower longitude, upper longitude.
156 Returns
157 -------
158 cubes: iris.cube.CubeList
159 Cubes loaded after being merged and concatenated.
161 Raises
162 ------
163 FileNotFoundError
164 If the provided path does not exist
165 """
166 # Get iterable of paths. Each path corresponds to 1 model.
167 paths = iter_maybe(file_paths)
168 model_names = iter_maybe(model_names)
169 # Check we have appropriate number of model names.
170 if model_names != (None,) and len(model_names) != len(paths):
171 raise ValueError(
172 f"The number of model names ({len(model_names)}) should equal "
173 f"the number of paths given ({len(paths)})."
174 )
176 # Load the data for each model into a CubeList per model.
177 model_cubes = (
178 _load_model(path, name, constraint)
179 for path, name in itertools.zip_longest(paths, model_names, fillvalue=None)
180 )
182 # Split out first model's cubes and mark it as the base for comparisons.
183 cubes = next(model_cubes)
184 for cube in cubes:
185 # Use 1 to indicate True, as booleans can't be saved in NetCDF attributes.
186 cube.attributes["cset_comparison_base"] = 1
187 # Load the rest of the models.
188 cubes.extend(itertools.chain.from_iterable(model_cubes))
190 # Enable different point-based observation sources to be concatenated.
191 cubes = _check_combine_point_observations(cubes)
193 # Unify time units so different case studies can merge.
194 iris.util.unify_time_units(cubes)
196 # Select sub region.
197 cubes = _cutout_cubes(cubes, subarea_type, subarea_extent)
199 # Merge and concatenate cubes now metadata has been fixed.
200 cubes = _merge_cubes_check_ensemble(cubes)
201 cubes = cubes.concatenate()
203 # Squeeze single valued coordinates into scalar coordinates.
204 cubes = iris.cube.CubeList(iris.util.squeeze(cube) for cube in cubes)
206 # Ensure dimension coordinates are bounded.
207 for cube in cubes:
208 for dim_coord in cube.coords(dim_coords=True):
209 if (dim_coord.standard_name == "time") and (
210 dim_coord.name()
211 not in itertools.chain.from_iterable(
212 m.coord_names for m in cube.cell_methods if m.method != "point"
213 )
214 ):
215 # Instantaneous time coordinate
216 continue
217 # Iris can't guess the bounds of a scalar coordinate.
218 if not dim_coord.has_bounds() and dim_coord.shape[0] > 1:
219 dim_coord.guess_bounds()
221 logger.info("Loaded cubes: %s", cubes)
222 if len(cubes) == 0:
223 raise NoDataError("No cubes loaded, check your constraints!")
224 return cubes
227def _load_model(
228 paths: str | list[str],
229 model_name: str | None,
230 constraint: iris.Constraint | None,
231) -> iris.cube.CubeList:
232 """Load a single model's data into a CubeList."""
233 input_files = _check_input_files(paths)
234 # If unset, a constraint of None lets everything be loaded.
235 logger.debug("Constraint: %s", constraint)
236 cubes = iris.load(input_files, constraint, callback=_loading_callback)
237 # If required, compute wind_speed from components.
238 cubes = _compute_winds(cubes, constraint)
240 # Add model_name attribute to each cube to make it available at any further
241 # step without needing to pass it as function parameter.
242 if model_name is not None:
243 for cube in cubes:
244 cube.attributes["model_name"] = model_name
245 return cubes
248def _check_input_files(input_paths: str | list[str]) -> list[Path]:
249 """Get an iterable of files to load, and check that they all exist.
251 Arguments
252 ---------
253 input_paths: list[str]
254 List of paths to input files or directories. The path may itself contain
255 glob patterns, but unlike in shells it will match directly first.
257 Returns
258 -------
259 list[Path]
260 A list of files to load.
262 Raises
263 ------
264 FileNotFoundError:
265 If the provided arguments don't resolve to at least one existing file.
266 """
267 files = []
268 for raw_filename in iter_maybe(input_paths):
269 # Match glob-like files first, if they exist.
270 raw_path = Path(raw_filename)
271 if raw_path.is_file():
272 files.append(raw_path)
273 else:
274 for input_path in glob.glob(raw_filename):
275 # Convert string paths into Path objects.
276 input_path = Path(input_path)
277 # Get the list of files in the directory, or use it directly.
278 if input_path.is_dir():
279 logger.debug("Checking directory '%s' for files", input_path)
280 files.extend(p for p in input_path.iterdir() if p.is_file())
281 else:
282 files.append(input_path)
284 files.sort()
285 logger.info("Loading files:\n%s", "\n".join(str(path) for path in files))
286 if len(files) == 0:
287 raise FileNotFoundError(f"No files found for {input_paths}")
288 return files
291def _merge_cubes_check_ensemble(cubes: iris.cube.CubeList):
292 """Merge CubeList, renumbering realizations of 0 if required.
294 An unsuccessful merge indicates common input cube attributes, most
295 commonly from ensemble members missing an explicit realization
296 coordinate. Therefore the members are renumbered before being merged
297 again.
298 """
299 try:
300 cubes = cubes.merge()
301 except iris.exceptions.MergeError:
302 _log_once(
303 "Attempt to merge input CubeList failed. Attempting to iterate realization coords to enable merge.",
304 level=logging.WARNING,
305 )
306 for ir, cube in enumerate(cubes):
307 if cube.coord("realization").points == 0: 307 ↛ 306line 307 didn't jump to line 306 because the condition on line 307 was always true
308 cube.coord("realization").points = ir + 1
309 cubes = cubes.merge()
310 return cubes
313def _cutout_cubes(
314 cubes: iris.cube.CubeList,
315 subarea_type: Literal["gridcells", "realworld", "modelrelative"] | None,
316 subarea_extent: list[float],
317):
318 """Cut out a subarea from a CubeList."""
319 if subarea_type is None:
320 logger.debug("Subarea selection is disabled.")
321 return cubes
323 # If selected, cutout according to number of grid cells to trim from each edge.
324 cutout_cubes = iris.cube.CubeList()
325 # Find spatial coordinates
326 for cube in cubes:
327 # Find dimension coordinates.
328 lat_name, lon_name = get_cube_yxcoordname(cube)
330 # Compute cutout based on number of cells to trim from edges.
331 if subarea_type == "gridcells":
332 logger.debug(
333 "User requested LowerTrim: %s LeftTrim: %s UpperTrim: %s RightTrim: %s",
334 subarea_extent[0],
335 subarea_extent[1],
336 subarea_extent[2],
337 subarea_extent[3],
338 )
339 lat_points = np.sort(cube.coord(lat_name).points)
340 lon_points = np.sort(cube.coord(lon_name).points)
341 # Define cutout region using user provided cell points.
342 lats = [lat_points[subarea_extent[0]], lat_points[-subarea_extent[2] - 1]]
343 lons = [lon_points[subarea_extent[1]], lon_points[-subarea_extent[3] - 1]]
345 # Compute cutout based on specified coordinate values.
346 elif subarea_type == "realworld" or subarea_type == "modelrelative":
347 # If not gridcells, cutout by requested geographic area,
348 logger.debug(
349 "User requested LLat: %s ULat: %s LLon: %s ULon: %s",
350 subarea_extent[0],
351 subarea_extent[1],
352 subarea_extent[2],
353 subarea_extent[3],
354 )
355 # Define cutout region using user provided coordinates.
356 lats = np.array(subarea_extent[0:2])
357 lons = np.array(subarea_extent[2:4])
358 # Ensure cutout longitudes are within +/- 180.0 bounds.
359 while lons[0] < -180.0:
360 lons += 360.0
361 while lons[1] > 180.0:
362 lons -= 360.0
363 # If the coordinate system is rotated we convert coordinates into
364 # model-relative coordinates to extract the appropriate cutout.
365 coord_system = cube.coord(lat_name).coord_system
366 if subarea_type == "realworld" and isinstance(
367 coord_system, iris.coord_systems.RotatedGeogCS
368 ):
369 lons, lats = rotate_pole(
370 lons,
371 lats,
372 pole_lon=coord_system.grid_north_pole_longitude,
373 pole_lat=coord_system.grid_north_pole_latitude,
374 )
375 else:
376 raise ValueError("Unknown subarea_type:", subarea_type)
378 # Do cutout and add to cutout_cubes.
379 intersection_args = {lat_name: lats, lon_name: lons}
380 logger.debug("Cutting out coords: %s", intersection_args)
381 try:
382 cutout_cubes.append(cube.intersection(**intersection_args))
383 except IndexError as err:
384 raise ValueError(
385 "Region cutout error. Check and update SUBAREA_EXTENT."
386 "Cutout region requested should be contained within data area. "
387 "Also check if cutout region requested is smaller than input grid spacing."
388 ) from err
390 return cutout_cubes
393def _loading_callback(cube: iris.cube.Cube, field, filename: str) -> iris.cube.Cube:
394 """Compose together the needed callbacks into a single function."""
395 # Most callbacks operate in-place, but save the cube when returned!
396 _realization_callback(cube)
397 _um_normalise_callback(cube)
398 _lfric_normalise_callback(cube)
399 _nimrod_normalise_callback(cube)
400 cube = _lfric_time_coord_fix_callback(cube)
401 _normalise_var0_varname(cube)
402 cube = _fix_no_spatial_coords_callback(cube)
403 _fix_spatial_coords_callback(cube)
404 _fix_pressure_coord_callback(cube)
405 _fix_um_radtime(cube)
406 _fix_cell_methods(cube)
407 cube = _convert_cube_units_callback(cube)
408 cube = _grid_longitude_fix_callback(cube)
409 _fix_lfric_cloud_base_altitude(cube)
410 _proleptic_gregorian_fix(cube)
411 _lfric_time_callback(cube)
412 _lfric_forecast_period_callback(cube)
413 cube = _fix_no_time_coords_callback(cube)
414 _normalise_longname(cube)
415 return cube
418def _realization_callback(cube):
419 """Add a realization coordinate initialised to 0 if missing.
421 This means deterministic and ensemble cubes can assume realization coordinate through the rest
422 of the code.
423 """
424 # Only add if realization coordinate does not exist.
425 if not cube.coords("realization"):
426 cube.add_aux_coord(
427 iris.coords.DimCoord(0, standard_name="realization", units="1")
428 )
431@functools.lru_cache(None)
432def _log_once(msg, level=logging.WARNING):
433 """Print a warning message, skipping recent duplicates."""
434 logger.log(level, msg)
437def _um_normalise_callback(cube: iris.cube.Cube):
438 """Normalise UM STASH variable long names to LFRic variable names.
440 Note standard names will remain associated with cubes where different.
441 Long name will be used consistently in output filename and titles.
442 """
443 # Convert STASH to LFRic variable name
444 if "STASH" in cube.attributes:
445 stash = cube.attributes["STASH"]
446 try:
447 (name, grid) = STASH_TO_LFRIC[str(stash)]
448 cube.long_name = name
449 except KeyError:
450 # Don't change cubes with unknown stash codes.
451 _log_once(
452 f"Unknown STASH code: {stash}. Please check file stash_to_lfric.py to update.",
453 level=logging.WARNING,
454 )
457def _lfric_normalise_callback(cube: iris.cube.Cube):
458 """Normalise attributes that prevents LFRic cube from merging.
460 The uuid and timeStamp relate to the output file, as saved by XIOS, and has
461 no relation to the data contained. These attributes are removed.
463 The um_stash_source is a list of STASH codes for when an LFRic field maps to
464 multiple UM fields, however it can be encoded in any order. This attribute
465 is sorted to prevent this. This attribute is only present in LFRic data that
466 has been converted to look like UM data.
467 """
468 # Remove unwanted attributes.
469 cube.attributes.pop("timeStamp", None)
470 cube.attributes.pop("uuid", None)
471 cube.attributes.pop("name", None)
472 cube.attributes.pop("source", None)
473 cube.attributes.pop("analysis_source", None)
474 cube.attributes.pop("history", None)
476 # Sort STASH code list.
477 stash_list = cube.attributes.get("um_stash_source")
478 if stash_list:
479 # Parse the string as a list, sort, then re-encode as a string.
480 cube.attributes["um_stash_source"] = str(sorted(ast.literal_eval(stash_list)))
483def _nimrod_normalise_callback(cube: iris.cube.Cube):
484 """Normalise attributes that prevents NIMROD radar cubes from merging."""
485 # Remove unwanted attributes.
486 cube.attributes.pop("radar_sites", None)
487 cube.attributes.pop("additional_radar_sites", None)
488 cube.attributes.pop("recursive_filter_iterations", None)
489 cube.attributes.pop("Probability methods", None)
492def _lfric_time_coord_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube:
493 """Ensure the time coordinate is a DimCoord rather than an AuxCoord.
495 The coordinate is converted and replaced if not. SLAMed LFRic data has this
496 issue, though the coordinate satisfies all the properties for a DimCoord.
497 Scalar time values are left as AuxCoords.
498 """
499 # This issue seems to come from iris's handling of NetCDF files where time
500 # always ends up as an AuxCoord.
501 if cube.coords("time"):
502 time_coord = cube.coord("time")
503 if (
504 not isinstance(time_coord, iris.coords.DimCoord)
505 and len(cube.coord_dims(time_coord)) == 1
506 ):
507 # Fudge the bounds to foil checking for strict monotonicity.
508 if ( 508 ↛ 512line 508 didn't jump to line 512 because the condition on line 508 was never true
509 time_coord.has_bounds()
510 and (time_coord.bounds[-1][0] - time_coord.bounds[0][0]) < 1.0e-8
511 ):
512 time_coord.bounds = [
513 [
514 time_coord.bounds[i][0] + 1.0e-8 * float(i),
515 time_coord.bounds[i][1],
516 ]
517 for i in range(len(time_coord.bounds))
518 ]
519 iris.util.promote_aux_coord_to_dim_coord(cube, time_coord)
520 return cube
523def _grid_longitude_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube:
524 """Check grid_longitude coordinates are in the range -180 deg to 180 deg.
526 This is necessary if comparing two models with different conventions --
527 for example, models where the prime meridian is defined as 0 deg or
528 360 deg. If not in the range -180 deg to 180 deg, we wrap the grid_longitude
529 so that it falls in this range. Checks are for near-180 bounds given
530 model data bounds may not extend exactly to 0. or 360.
531 Input cubes on non-rotated grid coordinates are not impacted.
532 """
533 try:
534 y, x = get_cube_yxcoordname(cube)
535 except ValueError:
536 # Don't modify non-spatial cubes.
537 return cube
539 long_coord = cube.coord(x)
540 # Wrap longitudes if rotated pole coordinates
541 coord_system = long_coord.coord_system
542 if x == "grid_longitude" and isinstance(
543 coord_system, iris.coord_systems.RotatedGeogCS
544 ):
545 long_points = long_coord.points.copy()
546 long_centre = np.median(long_points)
547 while long_centre < -175.0:
548 long_centre += 360.0
549 long_points += 360.0
550 while long_centre >= 175.0:
551 long_centre -= 360.0
552 long_points -= 360.0
553 long_coord.points = long_points
555 # Update coord bounds to be consistent with wrapping.
556 if long_coord.has_bounds():
557 long_coord.bounds = None
558 long_coord.guess_bounds()
560 return cube
563def _fix_no_spatial_coords_callback(cube: iris.cube.Cube):
564 import CSET.operators._utils as utils
566 # Don't modify spatial cubes that already have spatial dimensions
567 if utils.is_spatialdim(cube):
568 return cube
570 else:
571 # attempt to get lat/long from cube attributes
572 try:
573 lat_min = cube.attributes.get("geospatial_lat_min")
574 lat_max = cube.attributes.get("geospatial_lat_max")
575 lon_min = cube.attributes.get("geospatial_lon_min")
576 lon_max = cube.attributes.get("geospatial_lon_max")
578 lon_val = (lon_min + lon_max) / 2.0
579 lat_val = (lat_min + lat_max) / 2.0
581 lat_coord = iris.coords.DimCoord(
582 lat_val,
583 standard_name="latitude",
584 units="degrees_north",
585 var_name="latitude",
586 coord_system=iris.coord_systems.GeogCS(6371229.0),
587 circular=True,
588 )
590 lon_coord = iris.coords.DimCoord(
591 lon_val,
592 standard_name="longitude",
593 units="degrees_east",
594 var_name="longitude",
595 coord_system=iris.coord_systems.GeogCS(6371229.0),
596 circular=True,
597 )
599 cube.add_aux_coord(lat_coord)
600 cube.add_aux_coord(lon_coord)
601 return cube
603 # if lat/long are not in attributes, then return cube unchanged:
604 except TypeError:
605 return cube
608def _fix_spatial_coords_callback(cube: iris.cube.Cube):
609 """Check latitude and longitude coordinates name.
611 This is necessary as some models define their grid as on rotated
612 'grid_latitude' and 'grid_longitude' coordinates while others define
613 the grid on non-rotated 'latitude' and 'longitude'.
614 Cube dimensions need to be made consistent to avoid recipe failures,
615 particularly where comparing multiple input models with differing spatial
616 coordinates.
617 """
618 # Check if cube is spatial.
619 if not is_spatialdim(cube):
620 # Don't modify non-spatial cubes.
621 return
623 # Get spatial coords and dimension index.
624 y_name, x_name = get_cube_yxcoordname(cube)
625 ny = get_cube_coordindex(cube, y_name)
626 nx = get_cube_coordindex(cube, x_name)
628 # Remove spatial coords bounds if erroneous values detected.
629 # Aims to catch some errors in input coord bounds by setting
630 # invalid threshold of 10000.0
631 if cube.coord(x_name).has_bounds() and cube.coord(y_name).has_bounds():
632 bx_max = np.max(np.abs(cube.coord(x_name).bounds))
633 by_max = np.max(np.abs(cube.coord(y_name).bounds))
634 if bx_max > 10000.0 or by_max > 10000.0:
635 cube.coord(x_name).bounds = None
636 cube.coord(y_name).bounds = None
638 # Translate [grid_latitude, grid_longitude] to an unrotated 1-d DimCoord
639 # [latitude, longitude] for instances where rotated_pole=90.0
640 if "grid_latitude" in [coord.name() for coord in cube.coords(dim_coords=True)]:
641 coord_system = cube.coord("grid_latitude").coord_system
642 pole_lat = getattr(coord_system, "grid_north_pole_latitude", None)
643 if pole_lat == 90.0: 643 ↛ 644line 643 didn't jump to line 644 because the condition on line 643 was never true
644 lats = cube.coord("grid_latitude").points
645 lons = cube.coord("grid_longitude").points
647 cube.remove_coord("grid_latitude")
648 cube.add_dim_coord(
649 iris.coords.DimCoord(
650 lats,
651 standard_name="latitude",
652 var_name="latitude",
653 units="degrees",
654 coord_system=iris.coord_systems.GeogCS(6371229.0),
655 circular=True,
656 ),
657 ny,
658 )
659 y_name = "latitude"
660 cube.remove_coord("grid_longitude")
661 cube.add_dim_coord(
662 iris.coords.DimCoord(
663 lons,
664 standard_name="longitude",
665 var_name="longitude",
666 units="degrees",
667 coord_system=iris.coord_systems.GeogCS(6371229.0),
668 circular=True,
669 ),
670 nx,
671 )
672 x_name = "longitude"
674 # Create additional AuxCoord [grid_latitude, grid_longitude] with
675 # rotated pole attributes for cases with [lat, lon] inputs
676 if y_name in ["latitude"] and cube.coord(y_name).units in [
677 "degrees",
678 "degrees_north",
679 "degrees_south",
680 ]:
681 # Add grid_latitude AuxCoord
682 if "grid_latitude" not in [
683 coord.name() for coord in cube.coords(dim_coords=False)
684 ]:
685 cube.add_aux_coord(
686 iris.coords.AuxCoord(
687 cube.coord(y_name).points,
688 var_name="grid_latitude",
689 units="degrees",
690 ),
691 ny,
692 )
693 # Ensure input latitude DimCoord has CoordSystem
694 # This attribute is sometimes lost on iris.save
695 if not cube.coord(y_name).coord_system:
696 cube.coord(y_name).coord_system = iris.coord_systems.GeogCS(6371229.0)
698 if x_name in ["longitude"] and cube.coord(x_name).units in [
699 "degrees",
700 "degrees_west",
701 "degrees_east",
702 ]:
703 # Add grid_longitude AuxCoord
704 if "grid_longitude" not in [
705 coord.name() for coord in cube.coords(dim_coords=False)
706 ]:
707 cube.add_aux_coord(
708 iris.coords.AuxCoord(
709 cube.coord(x_name).points,
710 var_name="grid_longitude",
711 units="degrees",
712 ),
713 nx,
714 )
716 # Ensure input longitude DimCoord has CoordSystem
717 # This attribute is sometimes lost on iris.save
718 if not cube.coord(x_name).coord_system:
719 cube.coord(x_name).coord_system = iris.coord_systems.GeogCS(6371229.0)
722def _fix_pressure_coord_callback(cube: iris.cube.Cube):
723 """Rename pressure coordinate to "pressure" if it exists and ensure hPa units.
725 This problem was raised because the AIFS model data from ECMWF
726 defines the pressure coordinate with the name "pressure_level" rather
727 than compliant CF coordinate names.
729 Additionally, set the units of pressure to be hPa to be consistent with the UM,
730 and approach the coordinates in a unified way.
731 """
732 for coord in cube.dim_coords:
733 if coord.name() in ["pressure_level", "pressure_levels"]:
734 coord.rename("pressure")
736 if coord.name() == "pressure" and str(cube.coord("pressure").units) != "hPa":
737 cube.coord("pressure").convert_units("hPa")
740def _fix_um_radtime(cube: iris.cube.Cube):
741 """Move radiation diagnostics from timestamps which are output N minutes or seconds past every hour.
743 This callback does not have any effect for output diagnostics with
744 timestamps exactly 00 or 30 minutes past the hour. Only radiation
745 diagnostics are checked.
746 Note this callback does not interpolate the data in time, only adjust
747 timestamps to sit on the hour to enable time-to-time difference plotting
748 with models which may output radiation data on the hour.
749 """
750 try:
751 if cube.attributes["STASH"] in [
752 "m01s01i207",
753 "m01s01i208",
754 "m01s02i205",
755 "m01s02i201",
756 "m01s01i207",
757 "m01s02i207",
758 "m01s01i235",
759 ]:
760 time_coord = cube.coord("time")
762 # Convert time points to datetime objects
763 time_unit = time_coord.units
764 time_points = time_unit.num2date(time_coord.points)
765 # Skip if times don't need fixing.
766 if time_points[0].minute == 0 and time_points[0].second == 0:
767 return
768 if time_points[0].minute == 30 and time_points[0].second == 0: 768 ↛ 769line 768 didn't jump to line 769 because the condition on line 768 was never true
769 return
771 # Subtract time difference from the hour from each time point
772 n_minute = time_points[0].minute
773 n_second = time_points[0].second
774 # If times closer to next hour, compute difference to add on to following hour
775 if n_minute > 30:
776 n_minute = n_minute - 60
777 # Compute new diagnostic time stamp
778 new_time_points = (
779 time_points
780 - datetime.timedelta(minutes=n_minute)
781 - datetime.timedelta(seconds=n_second)
782 )
784 # Convert back to numeric values using the original time unit.
785 new_time_values = time_unit.date2num(new_time_points)
787 # Replace the time coordinate with updated values.
788 time_coord.points = new_time_values
790 # Recompute forecast_period with corrected values.
791 if cube.coord("forecast_period"): 791 ↛ exitline 791 didn't return from function '_fix_um_radtime' because the condition on line 791 was always true
792 fcst_prd_points = cube.coord("forecast_period").points
793 new_fcst_points = (
794 time_unit.num2date(fcst_prd_points)
795 - datetime.timedelta(minutes=n_minute)
796 - datetime.timedelta(seconds=n_second)
797 )
798 cube.coord("forecast_period").points = time_unit.date2num(
799 new_fcst_points
800 )
801 except KeyError:
802 pass
805def _fix_cell_methods(cube: iris.cube.Cube):
806 """To fix the assumed cell_methods in accumulation STASH from UM.
808 Lightning (m01s21i104), rainfall amount (m01s04i201, m01s05i201) and snowfall amount
809 (m01s04i202, m01s05i202) in UM is being output as a time accumulation,
810 over each hour (TAcc1hr), but input cubes show cell_methods as "mean".
811 For UM and LFRic inputs to be compatible, we assume accumulated cell_methods are
812 "sum". This callback changes "mean" cube attribute cell_method to "sum",
813 enabling the cell_method constraint on reading to select correct input.
814 """
815 # Shift "mean" cell_method to "sum" for selected UM inputs.
816 if cube.attributes.get("STASH") in [
817 "m01s21i104",
818 "m01s04i201",
819 "m01s04i202",
820 "m01s05i201",
821 "m01s05i202",
822 ] and {cm.method for cm in cube.cell_methods} == {"mean"}:
823 # Retrieve interval and any comment information.
824 for cell_method in cube.cell_methods:
825 interval_str = cell_method.intervals
826 comment_str = cell_method.comments
828 # Remove input aggregation method.
829 cube.cell_methods = ()
831 # Replace "mean" with "sum" cell_method to indicate aggregation.
832 cube.add_cell_method(
833 iris.coords.CellMethod(
834 method="sum",
835 coords="time",
836 intervals=interval_str,
837 comments=comment_str,
838 )
839 )
842def _convert_cube_units_callback(cube: iris.cube.Cube):
843 """Adjust diagnostic units for specific variables.
845 Some precipitation diagnostics are output with unit kg m-2 s-1 and are
846 converted here to mm hr-1.
848 Visibility diagnostics are converted here from m to km to improve output
849 formatting.
850 """
851 # Convert precipitation diagnostic units if required.
852 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
853 if any("surface_microphysical" in name for name in varnames):
854 if cube.units == "kg m-2 s-1":
855 _log_once(
856 "Converting precipitation rate units from kg m-2 s-1 to mm hr-1",
857 level=logging.DEBUG,
858 )
859 # Convert from kg m-2 s-1 to mm s-1 assuming 1kg water = 1l water = 1dm^3 water.
860 # This is a 1:1 conversion, so we just change the units.
861 cube.units = "mm s-1"
862 # Convert the units to per hour.
863 cube.convert_units("mm hr-1")
864 elif cube.units == "kg m-2": 864 ↛ 874line 864 didn't jump to line 874 because the condition on line 864 was always true
865 _log_once(
866 "Converting precipitation amount units from kg m-2 to mm",
867 level=logging.DEBUG,
868 )
869 # Convert from kg m-2 to mm assuming 1kg water = 1l water = 1dm^3 water.
870 # This is a 1:1 conversion, so we just change the units.
871 cube.units = "mm"
873 # Convert visibility diagnostic units if required.
874 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
875 if any("visibility" in name for name in varnames) and cube.units == "m":
876 _log_once("Converting visibility units m to km.", level=logging.DEBUG)
877 # Convert the units to km.
878 cube.convert_units("km")
880 return cube
883def _fix_lfric_cloud_base_altitude(cube: iris.cube.Cube):
884 """Mask cloud_base_altitude diagnostic in regions with no cloud."""
885 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
886 if any("cloud_base_altitude" in name for name in varnames):
887 # Mask cube where set > 144kft to catch default 144.35695538058164
888 cube.data = dask.array.ma.masked_greater(cube.core_data(), 144.0)
891def get_filter_windspeed(constraint: iris.Constraint):
892 """Get the windspeed filter by using the hijacked constraint."""
893 if hasattr(constraint, "varname"):
894 return constraint.varname
895 else:
896 return None
899def _compute_winds(
900 cubes: iris.cube.CubeList, constraint: iris.Constraint | None = None
901):
902 """To compute wind_speed from vector components if not available as diagnostic.
904 Diagnostics of wind are also not always consistent between the UM
905 and LFRic. Here, winds from the UM are adjusted to make them
906 consistent with LFRic.
907 """
908 # Check whether we have components of the wind identified by varname
909 # but not the wind speed and calculate it if it is missing. Note that
910 # this will be biased low in general because the components will mostly
911 # be time averages. For simplicity, we do this only if there is just one
912 # cube of a component. A more complicated approach would be to consider
913 # the cell methods, but it may not be warranted.
914 #
915 # A check on UM STASH attributes is also conducted to adjust directions.
917 if constraint is None:
918 return cubes
920 filter_windspeed = get_filter_windspeed(constraint)
922 u_constr = iris.Constraint("eastward_wind_at_10m")
923 v_constr = iris.Constraint("northward_wind_at_10m")
924 sp_constr = iris.Constraint("wind_speed_at_10m")
926 try:
927 if (
928 cubes.extract(u_constr)
929 and cubes.extract(v_constr)
930 and not cubes.extract(sp_constr)
931 ):
932 if "wind_speed_at_10m" in constraint.varname:
933 _add_wind_speed_um(cubes)
934 # Convert winds in the UM to be relative to true east and true north.
935 _convert_wind_true_dirn_um(cubes)
936 except (KeyError, AttributeError):
937 pass
939 if filter_windspeed:
940 filter_windspeed_constraint = iris.Constraint(
941 cube_func=lambda cube: (
942 cube.long_name in filter_windspeed
943 or cube.standard_name in filter_windspeed
944 or cube.var_name in filter_windspeed
945 )
946 )
947 cubes = cubes.extract(filter_windspeed_constraint)
948 return cubes
951def _add_wind_speed_um(cubes: iris.cube.CubeList):
952 """Add windspeeds to cubes from components."""
953 u_wind = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
954 v_wind = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
955 wspd10 = (u_wind**2 + v_wind**2) ** 0.5
956 wspd10.attributes["STASH"] = "m01s03i227"
957 wspd10.standard_name = "wind_speed"
958 wspd10.long_name = "wind_speed_at_10m"
959 wspd10.units = "ms-1"
960 cubes.append(wspd10)
963def _convert_wind_true_dirn_um(cubes: iris.cube.CubeList):
964 """To convert winds to true directions.
966 Convert from the components relative to the grid to true directions.
967 This functionality only handles the simplest case.
968 Constrains using STASH code only to ensure applied to UM outputs only.
969 """
970 u_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i225"))
971 v_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i226"))
972 for u, v in zip(u_grids, v_grids, strict=True): 972 ↛ 973line 972 didn't jump to line 973 because the loop on line 972 never started
973 true_u, true_v = rotate_winds(u, v, iris.coord_systems.GeogCS(6371229.0))
974 u.data = true_u.core_data()
975 v.data = true_v.core_data()
978def _normalise_var0_varname(cube: iris.cube.Cube):
979 """Fix varnames for consistency to allow merging.
981 Some model data netCDF sometimes have a coordinate name end in
982 "_0" etc, where duplicate coordinates of same name are defined but
983 with different attributes. This can be inconsistently managed in
984 different model inputs and can cause cubes to fail to merge.
985 """
986 for coord in cube.coords():
987 if coord.var_name and coord.var_name.endswith("_0"):
988 coord.var_name = coord.var_name.removesuffix("_0")
989 if coord.var_name and coord.var_name.endswith("_1"):
990 coord.var_name = coord.var_name.removesuffix("_1")
991 if coord.var_name and coord.var_name.endswith("_2"): 991 ↛ 992line 991 didn't jump to line 992 because the condition on line 991 was never true
992 coord.var_name = coord.var_name.removesuffix("_2")
993 if coord.var_name and coord.var_name.endswith("_3"): 993 ↛ 994line 993 didn't jump to line 994 because the condition on line 993 was never true
994 coord.var_name = coord.var_name.removesuffix("_3")
996 if cube.var_name and cube.var_name.endswith("_0"):
997 cube.var_name = cube.var_name.removesuffix("_0")
1000def _proleptic_gregorian_fix(cube: iris.cube.Cube):
1001 """Convert the calendars of time units to use a standard calendar."""
1002 try:
1003 time_coord = cube.coord("time")
1004 if time_coord.units.calendar == "proleptic_gregorian":
1005 logger.debug(
1006 "Changing proleptic Gregorian calendar to standard calendar for %s",
1007 repr(time_coord.units),
1008 )
1009 time_coord.units = time_coord.units.change_calendar("standard")
1010 except iris.exceptions.CoordinateNotFoundError:
1011 pass
1014def _lfric_time_callback(cube: iris.cube.Cube):
1015 """Fix time coordinate metadata if missing dimensions.
1017 Some model data does not contain forecast_reference_time or forecast_period as
1018 expected coordinates, and so we cannot aggregate over case studies without this
1019 metadata. This callback fixes these issues.
1021 This callback also ensures all time coordinates are referenced as hours since
1022 1970-01-01 00:00:00 for consistency across different model inputs.
1024 Notes
1025 -----
1026 Some parts of the code have been adapted from Paul Earnshaw's scripts.
1027 """
1028 # Construct forecast_reference time if it doesn't exist.
1029 try:
1030 tcoord = cube.coord("time")
1031 # Set time coordinate to common basis "hours since 1970"
1032 try:
1033 tcoord.convert_units("hours since 1970-01-01 00:00:00")
1034 except ValueError:
1035 logger.warning("Unrecognised base time unit: %s", tcoord.units)
1037 if not cube.coords("forecast_reference_time"):
1038 try:
1039 init_time = datetime.datetime.fromisoformat(
1040 tcoord.attributes["time_origin"]
1041 )
1042 frt_point = tcoord.units.date2num(init_time)
1043 frt_coord = iris.coords.AuxCoord(
1044 frt_point,
1045 units=tcoord.units,
1046 standard_name="forecast_reference_time",
1047 long_name="forecast_reference_time",
1048 )
1049 cube.add_aux_coord(frt_coord)
1050 except KeyError:
1051 logger.warning(
1052 "Cannot find forecast_reference_time, but no `time_origin` attribute to construct it from."
1053 )
1055 # Remove time_origin to allow multiple case studies to merge.
1056 tcoord.attributes.pop("time_origin", None)
1058 # Construct forecast_period axis (forecast lead time) if it doesn't exist.
1059 if not cube.coords("forecast_period"):
1060 try:
1061 # Create array of forecast lead times.
1062 init_coord = cube.coord("forecast_reference_time")
1063 init_time_points_in_tcoord_units = tcoord.units.date2num(
1064 init_coord.units.num2date(init_coord.points)
1065 )
1066 lead_times = tcoord.points - init_time_points_in_tcoord_units
1068 # Get unit for lead time from time coordinate's unit.
1069 # Convert all lead time to hours for consistency between models.
1070 if "seconds" in str(tcoord.units): 1070 ↛ 1071line 1070 didn't jump to line 1071 because the condition on line 1070 was never true
1071 lead_times = lead_times / 3600.0
1072 units = "hours"
1073 elif "hours" in str(tcoord.units): 1073 ↛ 1076line 1073 didn't jump to line 1076 because the condition on line 1073 was always true
1074 units = "hours"
1075 else:
1076 raise ValueError(f"Unrecognised base time unit: {tcoord.units}")
1078 # Create lead time coordinate.
1079 lead_time_coord = iris.coords.AuxCoord(
1080 lead_times,
1081 standard_name="forecast_period",
1082 long_name="forecast_period",
1083 units=units,
1084 )
1086 # Associate lead time coordinate with time dimension.
1087 cube.add_aux_coord(lead_time_coord, cube.coord_dims("time"))
1088 except iris.exceptions.CoordinateNotFoundError:
1089 logger.warning(
1090 "Cube does not have both time and forecast_reference_time coordinate, so cannot construct forecast_period"
1091 )
1092 except iris.exceptions.CoordinateNotFoundError:
1093 logger.warning("No time coordinate on cube.")
1096def _lfric_forecast_period_callback(cube: iris.cube.Cube):
1097 """Check forecast_period name and units."""
1098 try:
1099 coord = cube.coord("forecast_period")
1100 if coord.units != "hours":
1101 cube.coord("forecast_period").convert_units("hours")
1102 if not coord.standard_name:
1103 coord.standard_name = "forecast_period"
1104 except iris.exceptions.CoordinateNotFoundError:
1105 pass
1108def _fix_no_time_coords_callback(cube: iris.cube.Cube):
1109 """Add dummy time coord to process cubes that don't have sequence coord."""
1110 # Only add if time coordinate does not exist.
1111 if not cube.coords("time"):
1112 cube.add_aux_coord(
1113 iris.coords.DimCoord(
1114 0, standard_name="time", units="hours since 0001-01-01 00:00:00"
1115 )
1116 )
1118 return cube
1121def _normalise_longname(cube: iris.cube.Cube):
1122 """Normalise long_name to the LFRic standard list."""
1123 if cube.coords("pressure"):
1124 if cube.name() == "x_wind":
1125 cube.long_name = "zonal_wind_at_pressure_levels"
1126 if cube.name() == "y_wind":
1127 cube.long_name = "meridional_wind_at_pressure_levels"
1128 if cube.name() == "air_temperature":
1129 cube.long_name = "temperature_at_pressure_levels"
1130 if cube.name() == "specific_humidity": 1130 ↛ 1131line 1130 didn't jump to line 1131 because the condition on line 1130 was never true
1131 cube.long_name = (
1132 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging"
1133 )
1134 else:
1135 if cube.name() == "x_wind" and cube.var_name == "u_wind_at_10m": 1135 ↛ 1136line 1135 didn't jump to line 1136 because the condition on line 1135 was never true
1136 cube.long_name = "eastward_wind_at_10m"
1137 if cube.name() == "y_wind" and cube.var_name == "v_wind_at_10m": 1137 ↛ 1138line 1137 didn't jump to line 1138 because the condition on line 1137 was never true
1138 cube.long_name = "northward_wind_at_10m"
1139 if cube.name() == "air_pressure_at_sea_level":
1140 cube.long_name = "air_pressure_at_mean_sea_level"
1143def _check_combine_point_observations(cubes: iris.cube.CubeList):
1144 """Enable cubes containing different point observation sources to be concatenated."""
1145 nstation = 0
1146 for cube in cubes:
1147 if "station" in [coord.name() for coord in cube.coords(dim_coords=True)]:
1148 if "obs_source" in [coord.name() for coord in cube.coords()]:
1149 cube.remove_coord("obs_source")
1150 cube.coord("station").points = cube.coord("station").points + nstation
1151 nstation = nstation + len(cube.coord("station").points)
1153 return cubes