Coverage for src/CSET/operators/read.py: 94%
430 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 08:32 +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)
170 # Check we have appropriate number of model names.
171 if model_names != (None,) and len(model_names) != len(paths):
172 raise ValueError(
173 f"The number of model names ({len(model_names)}) should equal "
174 f"the number of paths given ({len(paths)})."
175 )
177 # Load the data for each model into a CubeList per model.
178 model_cubes = (
179 _load_model(path, name, constraint)
180 for path, name in itertools.zip_longest(paths, model_names, fillvalue=None)
181 )
183 # Split out first model's cubes and mark it as the base for comparisons.
184 cubes = next(model_cubes)
185 for cube in cubes:
186 # Use 1 to indicate True, as booleans can't be saved in NetCDF attributes.
187 cube.attributes["cset_comparison_base"] = 1
189 # Load the rest of the models.
190 cubes.extend(itertools.chain.from_iterable(model_cubes))
192 # Enable different point-based observation sources to be concatenated.
193 cubes = _check_combine_point_observations(cubes)
195 # Unify time units so different case studies can merge.
196 iris.util.unify_time_units(cubes)
198 # Select sub region.
199 cubes = _cutout_cubes(cubes, subarea_type, subarea_extent)
201 # Merge and concatenate cubes now metadata has been fixed.
202 cubes = _merge_cubes_check_ensemble(cubes)
203 cubes = cubes.concatenate()
205 # Squeeze single valued coordinates into scalar coordinates.
206 cubes = iris.cube.CubeList(iris.util.squeeze(cube) for cube in cubes)
208 # Ensure dimension coordinates are bounded.
209 for cube in cubes:
210 for dim_coord in cube.coords(dim_coords=True):
211 if (dim_coord.standard_name == "time") and (
212 dim_coord.name()
213 not in itertools.chain.from_iterable(
214 m.coord_names for m in cube.cell_methods if m.method != "point"
215 )
216 ):
217 # Instantaneous time coordinate
218 continue
219 # Iris can't guess the bounds of a scalar coordinate.
220 if not dim_coord.has_bounds() and dim_coord.shape[0] > 1:
221 dim_coord.guess_bounds()
223 logger.info("Loaded cubes: %s", cubes)
224 if len(cubes) == 0:
225 raise NoDataError("No cubes loaded, check your constraints!")
226 return cubes
229def _load_model(
230 paths: str | list[str],
231 model_name: str | None,
232 constraint: iris.Constraint | None,
233) -> iris.cube.CubeList:
234 """Load a single model's data into a CubeList."""
235 input_files = _check_input_files(paths)
236 # If unset, a constraint of None lets everything be loaded.
237 logger.debug("Constraint: %s", constraint)
238 cubes = iris.load(input_files, constraint, callback=_loading_callback)
239 # If required, compute wind_speed from components.
240 cubes = _compute_winds(cubes)
242 # Add model_name attribute to each cube to make it available at any further
243 # step without needing to pass it as function parameter.
244 if model_name is not None:
245 for cube in cubes:
246 cube.attributes["model_name"] = model_name
247 return cubes
250def _check_input_files(input_paths: str | list[str]) -> list[Path]:
251 """Get an iterable of files to load, and check that they all exist.
253 Arguments
254 ---------
255 input_paths: list[str]
256 List of paths to input files or directories. The path may itself contain
257 glob patterns, but unlike in shells it will match directly first.
259 Returns
260 -------
261 list[Path]
262 A list of files to load.
264 Raises
265 ------
266 FileNotFoundError:
267 If the provided arguments don't resolve to at least one existing file.
268 """
269 files = []
270 for raw_filename in iter_maybe(input_paths):
271 # Match glob-like files first, if they exist.
272 raw_path = Path(raw_filename)
273 if raw_path.is_file():
274 files.append(raw_path)
275 else:
276 for input_path in glob.glob(raw_filename):
277 # Convert string paths into Path objects.
278 input_path = Path(input_path)
279 # Get the list of files in the directory, or use it directly.
280 if input_path.is_dir():
281 logger.debug("Checking directory '%s' for files", input_path)
282 files.extend(p for p in input_path.iterdir() if p.is_file())
283 else:
284 files.append(input_path)
286 files.sort()
287 logger.info("Loading files:\n%s", "\n".join(str(path) for path in files))
288 if len(files) == 0:
289 raise FileNotFoundError(f"No files found for {input_paths}")
290 return files
293def _merge_cubes_check_ensemble(cubes: iris.cube.CubeList):
294 """Merge CubeList, renumbering realizations of 0 if required.
296 An unsuccessful merge indicates common input cube attributes, most
297 commonly from ensemble members missing an explicit realization
298 coordinate. Therefore the members are renumbered before being merged
299 again.
300 """
301 try:
302 cubes = cubes.merge()
303 except iris.exceptions.MergeError:
304 _log_once(
305 "Attempt to merge input CubeList failed. Attempting to iterate realization coords to enable merge.",
306 level=logging.WARNING,
307 )
308 for ir, cube in enumerate(cubes):
309 if cube.coord("realization").points == 0: 309 ↛ 308line 309 didn't jump to line 308 because the condition on line 309 was always true
310 cube.coord("realization").points = ir + 1
311 cubes = cubes.merge()
312 return cubes
315def _cutout_cubes(
316 cubes: iris.cube.CubeList,
317 subarea_type: Literal["gridcells", "realworld", "modelrelative"] | None,
318 subarea_extent: list[float],
319):
320 """Cut out a subarea from a CubeList."""
321 if subarea_type is None:
322 logger.debug("Subarea selection is disabled.")
323 return cubes
325 # If selected, cutout according to number of grid cells to trim from each edge.
326 cutout_cubes = iris.cube.CubeList()
327 # Find spatial coordinates
328 for cube in cubes:
329 # Find dimension coordinates.
330 lat_name, lon_name = get_cube_yxcoordname(cube)
332 # Compute cutout based on number of cells to trim from edges.
333 if subarea_type == "gridcells":
334 logger.debug(
335 "User requested LowerTrim: %s LeftTrim: %s UpperTrim: %s RightTrim: %s",
336 subarea_extent[0],
337 subarea_extent[1],
338 subarea_extent[2],
339 subarea_extent[3],
340 )
341 lat_points = np.sort(cube.coord(lat_name).points)
342 lon_points = np.sort(cube.coord(lon_name).points)
343 # Define cutout region using user provided cell points.
344 lats = [lat_points[subarea_extent[0]], lat_points[-subarea_extent[2] - 1]]
345 lons = [lon_points[subarea_extent[1]], lon_points[-subarea_extent[3] - 1]]
347 # Compute cutout based on specified coordinate values.
348 elif subarea_type == "realworld" or subarea_type == "modelrelative":
349 # If not gridcells, cutout by requested geographic area,
350 logger.debug(
351 "User requested LLat: %s ULat: %s LLon: %s ULon: %s",
352 subarea_extent[0],
353 subarea_extent[1],
354 subarea_extent[2],
355 subarea_extent[3],
356 )
357 # Define cutout region using user provided coordinates.
358 lats = np.array(subarea_extent[0:2])
359 lons = np.array(subarea_extent[2:4])
360 # Ensure cutout longitudes are within +/- 180.0 bounds.
361 while lons[0] < -180.0:
362 lons += 360.0
363 while lons[1] > 180.0:
364 lons -= 360.0
365 # If the coordinate system is rotated we convert coordinates into
366 # model-relative coordinates to extract the appropriate cutout.
367 coord_system = cube.coord(lat_name).coord_system
368 if subarea_type == "realworld" and isinstance(
369 coord_system, iris.coord_systems.RotatedGeogCS
370 ):
371 lons, lats = rotate_pole(
372 lons,
373 lats,
374 pole_lon=coord_system.grid_north_pole_longitude,
375 pole_lat=coord_system.grid_north_pole_latitude,
376 )
377 else:
378 raise ValueError("Unknown subarea_type:", subarea_type)
380 # Do cutout and add to cutout_cubes.
381 intersection_args = {lat_name: lats, lon_name: lons}
382 logger.debug("Cutting out coords: %s", intersection_args)
383 try:
384 cutout_cubes.append(cube.intersection(**intersection_args))
385 except IndexError as err:
386 raise ValueError(
387 "Region cutout error. Check and update SUBAREA_EXTENT."
388 "Cutout region requested should be contained within data area. "
389 "Also check if cutout region requested is smaller than input grid spacing."
390 ) from err
392 return cutout_cubes
395def _loading_callback(cube: iris.cube.Cube, field, filename: str) -> iris.cube.Cube:
396 """Compose together the needed callbacks into a single function."""
397 # Most callbacks operate in-place, but save the cube when returned!
398 _realization_callback(cube)
399 _um_normalise_callback(cube)
400 _lfric_normalise_callback(cube)
401 _nimrod_normalise_callback(cube)
402 cube = _lfric_time_coord_fix_callback(cube)
403 _normalise_var0_varname(cube)
404 cube = _fix_no_spatial_coords_callback(cube)
405 _fix_spatial_coords_callback(cube)
406 _fix_pressure_coord_callback(cube)
407 _fix_um_radtime(cube)
408 _fix_cell_methods(cube)
409 cube = _convert_cube_units_callback(cube)
410 cube = _grid_longitude_fix_callback(cube)
411 _fix_lfric_cloud_base_altitude(cube)
412 _proleptic_gregorian_fix(cube)
413 _lfric_time_callback(cube)
414 _lfric_forecast_period_callback(cube)
415 cube = _fix_no_time_coords_callback(cube)
416 _normalise_ML_varname(cube)
417 return cube
420def _realization_callback(cube):
421 """Add a realization coordinate initialised to 0 if missing.
423 This means deterministic and ensemble cubes can assume realization coordinate through the rest
424 of the code.
425 """
426 # Only add if realization coordinate does not exist.
427 if not cube.coords("realization"):
428 cube.add_aux_coord(
429 iris.coords.DimCoord(0, standard_name="realization", units="1")
430 )
433@functools.lru_cache(None)
434def _log_once(msg, level=logging.WARNING):
435 """Print a warning message, skipping recent duplicates."""
436 logger.log(level, msg)
439def _um_normalise_callback(cube: iris.cube.Cube):
440 """Normalise UM STASH variable long names to LFRic variable names.
442 Note standard names will remain associated with cubes where different.
443 Long name will be used consistently in output filename and titles.
444 """
445 # Convert STASH to LFRic variable name
446 if "STASH" in cube.attributes:
447 stash = cube.attributes["STASH"]
448 try:
449 (name, grid) = STASH_TO_LFRIC[str(stash)]
450 cube.long_name = name
451 except KeyError:
452 # Don't change cubes with unknown stash codes.
453 _log_once(
454 f"Unknown STASH code: {stash}. Please check file stash_to_lfric.py to update.",
455 level=logging.WARNING,
456 )
459def _lfric_normalise_callback(cube: iris.cube.Cube):
460 """Normalise attributes that prevents LFRic cube from merging.
462 The uuid and timeStamp relate to the output file, as saved by XIOS, and has
463 no relation to the data contained. These attributes are removed.
465 The um_stash_source is a list of STASH codes for when an LFRic field maps to
466 multiple UM fields, however it can be encoded in any order. This attribute
467 is sorted to prevent this. This attribute is only present in LFRic data that
468 has been converted to look like UM data.
469 """
470 # Remove unwanted attributes.
471 cube.attributes.pop("timeStamp", None)
472 cube.attributes.pop("uuid", None)
473 cube.attributes.pop("name", None)
474 cube.attributes.pop("source", None)
475 cube.attributes.pop("analysis_source", None)
476 cube.attributes.pop("history", None)
478 # Sort STASH code list.
479 stash_list = cube.attributes.get("um_stash_source")
480 if stash_list:
481 # Parse the string as a list, sort, then re-encode as a string.
482 cube.attributes["um_stash_source"] = str(sorted(ast.literal_eval(stash_list)))
485def _nimrod_normalise_callback(cube: iris.cube.Cube):
486 """Normalise attributes that prevents NIMROD radar cubes from merging."""
487 # Remove unwanted attributes.
488 cube.attributes.pop("radar_sites", None)
489 cube.attributes.pop("additional_radar_sites", None)
490 cube.attributes.pop("recursive_filter_iterations", None)
491 cube.attributes.pop("Probability methods", None)
494def _lfric_time_coord_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube:
495 """Ensure the time coordinate is a DimCoord rather than an AuxCoord.
497 The coordinate is converted and replaced if not. SLAMed LFRic data has this
498 issue, though the coordinate satisfies all the properties for a DimCoord.
499 Scalar time values are left as AuxCoords.
500 """
501 # This issue seems to come from iris's handling of NetCDF files where time
502 # always ends up as an AuxCoord.
503 if cube.coords("time"):
504 time_coord = cube.coord("time")
505 if (
506 not isinstance(time_coord, iris.coords.DimCoord)
507 and len(cube.coord_dims(time_coord)) == 1
508 ):
509 # Fudge the bounds to foil checking for strict monotonicity.
510 if ( 510 ↛ 514line 510 didn't jump to line 514 because the condition on line 510 was never true
511 time_coord.has_bounds()
512 and (time_coord.bounds[-1][0] - time_coord.bounds[0][0]) < 1.0e-8
513 ):
514 time_coord.bounds = [
515 [
516 time_coord.bounds[i][0] + 1.0e-8 * float(i),
517 time_coord.bounds[i][1],
518 ]
519 for i in range(len(time_coord.bounds))
520 ]
521 iris.util.promote_aux_coord_to_dim_coord(cube, time_coord)
522 return cube
525def _grid_longitude_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube:
526 """Check grid_longitude coordinates are in the range -180 deg to 180 deg.
528 This is necessary if comparing two models with different conventions --
529 for example, models where the prime meridian is defined as 0 deg or
530 360 deg. If not in the range -180 deg to 180 deg, we wrap the grid_longitude
531 so that it falls in this range. Checks are for near-180 bounds given
532 model data bounds may not extend exactly to 0. or 360.
533 Input cubes on non-rotated grid coordinates are not impacted.
534 """
535 try:
536 y, x = get_cube_yxcoordname(cube)
537 except ValueError:
538 # Don't modify non-spatial cubes.
539 return cube
541 long_coord = cube.coord(x)
542 # Wrap longitudes if rotated pole coordinates
543 coord_system = long_coord.coord_system
544 if x == "grid_longitude" and isinstance(
545 coord_system, iris.coord_systems.RotatedGeogCS
546 ):
547 long_points = long_coord.points.copy()
548 long_centre = np.median(long_points)
549 while long_centre < -175.0:
550 long_centre += 360.0
551 long_points += 360.0
552 while long_centre >= 175.0:
553 long_centre -= 360.0
554 long_points -= 360.0
555 long_coord.points = long_points
557 # Update coord bounds to be consistent with wrapping.
558 if long_coord.has_bounds():
559 long_coord.bounds = None
560 long_coord.guess_bounds()
562 return cube
565def _fix_no_spatial_coords_callback(cube: iris.cube.Cube):
566 import CSET.operators._utils as utils
568 # Don't modify spatial cubes that already have spatial dimensions
569 if utils.is_spatialdim(cube):
570 return cube
572 else:
573 # attempt to get lat/long from cube attributes
574 try:
575 lat_min = cube.attributes.get("geospatial_lat_min")
576 lat_max = cube.attributes.get("geospatial_lat_max")
577 lon_min = cube.attributes.get("geospatial_lon_min")
578 lon_max = cube.attributes.get("geospatial_lon_max")
580 lon_val = (lon_min + lon_max) / 2.0
581 lat_val = (lat_min + lat_max) / 2.0
583 lat_coord = iris.coords.DimCoord(
584 lat_val,
585 standard_name="latitude",
586 units="degrees_north",
587 var_name="latitude",
588 coord_system=iris.coord_systems.GeogCS(6371229.0),
589 circular=True,
590 )
592 lon_coord = iris.coords.DimCoord(
593 lon_val,
594 standard_name="longitude",
595 units="degrees_east",
596 var_name="longitude",
597 coord_system=iris.coord_systems.GeogCS(6371229.0),
598 circular=True,
599 )
601 cube.add_aux_coord(lat_coord)
602 cube.add_aux_coord(lon_coord)
603 return cube
605 # if lat/long are not in attributes, then return cube unchanged:
606 except TypeError:
607 return cube
610def _fix_spatial_coords_callback(cube: iris.cube.Cube):
611 """Check latitude and longitude coordinates name.
613 This is necessary as some models define their grid as on rotated
614 'grid_latitude' and 'grid_longitude' coordinates while others define
615 the grid on non-rotated 'latitude' and 'longitude'.
616 Cube dimensions need to be made consistent to avoid recipe failures,
617 particularly where comparing multiple input models with differing spatial
618 coordinates.
619 """
620 # Check if cube is spatial.
621 if not is_spatialdim(cube):
622 # Don't modify non-spatial cubes.
623 return
625 # Get spatial coords and dimension index.
626 y_name, x_name = get_cube_yxcoordname(cube)
627 ny = get_cube_coordindex(cube, y_name)
628 nx = get_cube_coordindex(cube, x_name)
630 # Remove spatial coords bounds if erroneous values detected.
631 # Aims to catch some errors in input coord bounds by setting
632 # invalid threshold of 10000.0
633 if cube.coord(x_name).has_bounds() and cube.coord(y_name).has_bounds():
634 bx_max = np.max(np.abs(cube.coord(x_name).bounds))
635 by_max = np.max(np.abs(cube.coord(y_name).bounds))
636 if bx_max > 10000.0 or by_max > 10000.0:
637 cube.coord(x_name).bounds = None
638 cube.coord(y_name).bounds = None
640 # Translate [grid_latitude, grid_longitude] to an unrotated 1-d DimCoord
641 # [latitude, longitude] for instances where rotated_pole=90.0
642 if "grid_latitude" in [coord.name() for coord in cube.coords(dim_coords=True)]:
643 coord_system = cube.coord("grid_latitude").coord_system
644 pole_lat = getattr(coord_system, "grid_north_pole_latitude", None)
645 if pole_lat == 90.0: 645 ↛ 646line 645 didn't jump to line 646 because the condition on line 645 was never true
646 lats = cube.coord("grid_latitude").points
647 lons = cube.coord("grid_longitude").points
649 cube.remove_coord("grid_latitude")
650 cube.add_dim_coord(
651 iris.coords.DimCoord(
652 lats,
653 standard_name="latitude",
654 var_name="latitude",
655 units="degrees",
656 coord_system=iris.coord_systems.GeogCS(6371229.0),
657 circular=True,
658 ),
659 ny,
660 )
661 y_name = "latitude"
662 cube.remove_coord("grid_longitude")
663 cube.add_dim_coord(
664 iris.coords.DimCoord(
665 lons,
666 standard_name="longitude",
667 var_name="longitude",
668 units="degrees",
669 coord_system=iris.coord_systems.GeogCS(6371229.0),
670 circular=True,
671 ),
672 nx,
673 )
674 x_name = "longitude"
676 # Create additional AuxCoord [grid_latitude, grid_longitude] with
677 # rotated pole attributes for cases with [lat, lon] inputs
678 if y_name in ["latitude"] and cube.coord(y_name).units in [
679 "degrees",
680 "degrees_north",
681 "degrees_south",
682 ]:
683 # Add grid_latitude AuxCoord
684 if "grid_latitude" not in [
685 coord.name() for coord in cube.coords(dim_coords=False)
686 ]:
687 cube.add_aux_coord(
688 iris.coords.AuxCoord(
689 cube.coord(y_name).points,
690 var_name="grid_latitude",
691 units="degrees",
692 ),
693 ny,
694 )
695 # Ensure input latitude DimCoord has CoordSystem
696 # This attribute is sometimes lost on iris.save
697 if not cube.coord(y_name).coord_system:
698 cube.coord(y_name).coord_system = iris.coord_systems.GeogCS(6371229.0)
700 if x_name in ["longitude"] and cube.coord(x_name).units in [
701 "degrees",
702 "degrees_west",
703 "degrees_east",
704 ]:
705 # Add grid_longitude AuxCoord
706 if "grid_longitude" not in [
707 coord.name() for coord in cube.coords(dim_coords=False)
708 ]:
709 cube.add_aux_coord(
710 iris.coords.AuxCoord(
711 cube.coord(x_name).points,
712 var_name="grid_longitude",
713 units="degrees",
714 ),
715 nx,
716 )
718 # Ensure input longitude DimCoord has CoordSystem
719 # This attribute is sometimes lost on iris.save
720 if not cube.coord(x_name).coord_system:
721 cube.coord(x_name).coord_system = iris.coord_systems.GeogCS(6371229.0)
724def _fix_pressure_coord_callback(cube: iris.cube.Cube):
725 """Rename pressure coordinate to "pressure" if it exists and ensure hPa units.
727 This problem was raised because the AIFS model data from ECMWF
728 defines the pressure coordinate with the name "pressure_level" rather
729 than compliant CF coordinate names.
731 Additionally, set the units of pressure to be hPa to be consistent with the UM,
732 and approach the coordinates in a unified way.
733 """
734 for coord in cube.dim_coords:
735 if coord.name() in ["pressure_level", "pressure_levels"]:
736 coord.rename("pressure")
738 if coord.name() == "pressure" and str(cube.coord("pressure").units) != "hPa":
739 cube.coord("pressure").convert_units("hPa")
742def _fix_um_radtime(cube: iris.cube.Cube):
743 """Move radiation diagnostics from timestamps which are output N minutes or seconds past every hour.
745 This callback does not have any effect for output diagnostics with
746 timestamps exactly 00 or 30 minutes past the hour. Only radiation
747 diagnostics are checked.
748 Note this callback does not interpolate the data in time, only adjust
749 timestamps to sit on the hour to enable time-to-time difference plotting
750 with models which may output radiation data on the hour.
751 """
752 try:
753 if cube.attributes["STASH"] in [
754 "m01s01i207",
755 "m01s01i208",
756 "m01s02i205",
757 "m01s02i201",
758 "m01s01i207",
759 "m01s02i207",
760 "m01s01i235",
761 ]:
762 time_coord = cube.coord("time")
764 # Convert time points to datetime objects
765 time_unit = time_coord.units
766 time_points = time_unit.num2date(time_coord.points)
767 # Skip if times don't need fixing.
768 if time_points[0].minute == 0 and time_points[0].second == 0:
769 return
770 if time_points[0].minute == 30 and time_points[0].second == 0: 770 ↛ 771line 770 didn't jump to line 771 because the condition on line 770 was never true
771 return
773 # Subtract time difference from the hour from each time point
774 n_minute = time_points[0].minute
775 n_second = time_points[0].second
776 # If times closer to next hour, compute difference to add on to following hour
777 if n_minute > 30:
778 n_minute = n_minute - 60
779 # Compute new diagnostic time stamp
780 new_time_points = (
781 time_points
782 - datetime.timedelta(minutes=n_minute)
783 - datetime.timedelta(seconds=n_second)
784 )
786 # Convert back to numeric values using the original time unit.
787 new_time_values = time_unit.date2num(new_time_points)
789 # Replace the time coordinate with updated values.
790 time_coord.points = new_time_values
792 # Recompute forecast_period with corrected values.
793 if cube.coord("forecast_period"): 793 ↛ exitline 793 didn't return from function '_fix_um_radtime' because the condition on line 793 was always true
794 fcst_prd_points = cube.coord("forecast_period").points
795 new_fcst_points = (
796 time_unit.num2date(fcst_prd_points)
797 - datetime.timedelta(minutes=n_minute)
798 - datetime.timedelta(seconds=n_second)
799 )
800 cube.coord("forecast_period").points = time_unit.date2num(
801 new_fcst_points
802 )
803 except KeyError:
804 pass
807def _fix_cell_methods(cube: iris.cube.Cube):
808 """To fix the assumed cell_methods in accumulation STASH from UM.
810 Lightning (m01s21i104), rainfall amount (m01s04i201, m01s05i201) and snowfall amount
811 (m01s04i202, m01s05i202) in UM is being output as a time accumulation,
812 over each hour (TAcc1hr), but input cubes show cell_methods as "mean".
813 For UM and LFRic inputs to be compatible, we assume accumulated cell_methods are
814 "sum". This callback changes "mean" cube attribute cell_method to "sum",
815 enabling the cell_method constraint on reading to select correct input.
816 """
817 # Shift "mean" cell_method to "sum" for selected UM inputs.
818 if cube.attributes.get("STASH") in [
819 "m01s21i104",
820 "m01s04i201",
821 "m01s04i202",
822 "m01s05i201",
823 "m01s05i202",
824 ] and {cm.method for cm in cube.cell_methods} == {"mean"}:
825 # Retrieve interval and any comment information.
826 for cell_method in cube.cell_methods:
827 interval_str = cell_method.intervals
828 comment_str = cell_method.comments
830 # Remove input aggregation method.
831 cube.cell_methods = ()
833 # Replace "mean" with "sum" cell_method to indicate aggregation.
834 cube.add_cell_method(
835 iris.coords.CellMethod(
836 method="sum",
837 coords="time",
838 intervals=interval_str,
839 comments=comment_str,
840 )
841 )
844def _convert_cube_units_callback(cube: iris.cube.Cube):
845 """Adjust diagnostic units for specific variables.
847 Some precipitation diagnostics are output with unit kg m-2 s-1 and are
848 converted here to mm hr-1.
850 Visibility diagnostics are converted here from m to km to improve output
851 formatting.
852 """
853 # Convert precipitation diagnostic units if required.
854 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
855 if any("surface_microphysical" in name for name in varnames):
856 if cube.units == "kg m-2 s-1":
857 _log_once(
858 "Converting precipitation rate units from kg m-2 s-1 to mm hr-1",
859 level=logging.DEBUG,
860 )
861 # Convert from kg m-2 s-1 to mm s-1 assuming 1kg water = 1l water = 1dm^3 water.
862 # This is a 1:1 conversion, so we just change the units.
863 cube.units = "mm s-1"
864 # Convert the units to per hour.
865 cube.convert_units("mm hr-1")
866 elif cube.units == "kg m-2": 866 ↛ 876line 866 didn't jump to line 876 because the condition on line 866 was always true
867 _log_once(
868 "Converting precipitation amount units from kg m-2 to mm",
869 level=logging.DEBUG,
870 )
871 # Convert from kg m-2 to mm assuming 1kg water = 1l water = 1dm^3 water.
872 # This is a 1:1 conversion, so we just change the units.
873 cube.units = "mm"
875 # Convert visibility diagnostic units if required.
876 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
877 if any("visibility" in name for name in varnames) and cube.units == "m":
878 _log_once("Converting visibility units m to km.", level=logging.DEBUG)
879 # Convert the units to km.
880 cube.convert_units("km")
882 return cube
885def _fix_lfric_cloud_base_altitude(cube: iris.cube.Cube):
886 """Mask cloud_base_altitude diagnostic in regions with no cloud."""
887 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
888 if any("cloud_base_altitude" in name for name in varnames):
889 # Mask cube where set > 144kft to catch default 144.35695538058164
890 cube.data = dask.array.ma.masked_greater(cube.core_data(), 144.0)
893def _compute_winds(cubes: iris.cube.CubeList):
894 """To compute wind_speed from vector components if not available as diagnostic.
896 Diagnostics of wind are also not always consistent between the UM
897 and LFRic. Here, winds from the UM are adjusted to make them
898 consistent with LFRic.
899 """
900 # Check whether we have components of the wind identified by varname
901 # but not the wind speed and calculate it if it is missing. Note that
902 # this will be biased low in general because the components will mostly
903 # be time averages. For simplicity, we do this only if there is just one
904 # cube of a component. A more complicated approach would be to consider
905 # the cell methods, but it may not be warranted.
906 #
907 # A check on UM STASH attributes is also conducted to adjust directions.
908 u_constr = iris.Constraint("eastward_wind_at_10m")
909 v_constr = iris.Constraint("northward_wind_at_10m")
910 speed_constr = iris.Constraint("wind_speed_at_10m")
911 try:
912 if cubes.extract(u_constr) and cubes.extract(v_constr):
913 if len(cubes) == 2:
914 wind_only = True
915 else:
916 wind_only = False
917 if len(cubes.extract(u_constr)) == 1 and not cubes.extract(speed_constr): 917 ↛ 920line 917 didn't jump to line 920 because the condition on line 917 was always true
918 _add_wind_speed_um(cubes)
919 # Convert winds in the UM to be relative to true east and true north.
920 if cubes.extract(u_constr) and cubes.extract(v_constr): 920 ↛ 923line 920 didn't jump to line 923 because the condition on line 920 was always true
921 _convert_wind_true_dirn_um(cubes)
922 # Return only wind_speed cube
923 if wind_only:
924 cubes = cubes.extract(speed_constr)
925 except (KeyError, AttributeError):
926 pass
928 return cubes
931def _add_wind_speed_um(cubes: iris.cube.CubeList):
932 """Add windspeeds to cubes from components."""
933 u_wind = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
934 v_wind = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
935 wspd10 = (u_wind**2 + v_wind**2) ** 0.5
936 wspd10.attributes["STASH"] = "m01s03i227"
937 wspd10.standard_name = "wind_speed"
938 wspd10.long_name = "wind_speed_at_10m"
939 wspd10.units = "ms-1"
940 cubes.append(wspd10)
943def _convert_wind_true_dirn_um(cubes: iris.cube.CubeList):
944 """To convert winds to true directions.
946 Convert from the components relative to the grid to true directions.
947 This functionality only handles the simplest case.
948 Constrains using STASH code only to ensure applied to UM outputs only.
949 """
950 u_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i225"))
951 v_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i226"))
952 for u, v in zip(u_grids, v_grids, strict=True): 952 ↛ 953line 952 didn't jump to line 953 because the loop on line 952 never started
953 true_u, true_v = rotate_winds(u, v, iris.coord_systems.GeogCS(6371229.0))
954 u.data = true_u.core_data()
955 v.data = true_v.core_data()
958def _normalise_var0_varname(cube: iris.cube.Cube):
959 """Fix varnames for consistency to allow merging.
961 Some model data netCDF sometimes have a coordinate name end in
962 "_0" etc, where duplicate coordinates of same name are defined but
963 with different attributes. This can be inconsistently managed in
964 different model inputs and can cause cubes to fail to merge.
965 """
966 for coord in cube.coords():
967 if coord.var_name and coord.var_name.endswith("_0"):
968 coord.var_name = coord.var_name.removesuffix("_0")
969 if coord.var_name and coord.var_name.endswith("_1"):
970 coord.var_name = coord.var_name.removesuffix("_1")
971 if coord.var_name and coord.var_name.endswith("_2"): 971 ↛ 972line 971 didn't jump to line 972 because the condition on line 971 was never true
972 coord.var_name = coord.var_name.removesuffix("_2")
973 if coord.var_name and coord.var_name.endswith("_3"): 973 ↛ 974line 973 didn't jump to line 974 because the condition on line 973 was never true
974 coord.var_name = coord.var_name.removesuffix("_3")
976 if cube.var_name and cube.var_name.endswith("_0"):
977 cube.var_name = cube.var_name.removesuffix("_0")
980def _proleptic_gregorian_fix(cube: iris.cube.Cube):
981 """Convert the calendars of time units to use a standard calendar."""
982 try:
983 time_coord = cube.coord("time")
984 if time_coord.units.calendar == "proleptic_gregorian":
985 logger.debug(
986 "Changing proleptic Gregorian calendar to standard calendar for %s",
987 repr(time_coord.units),
988 )
989 time_coord.units = time_coord.units.change_calendar("standard")
990 except iris.exceptions.CoordinateNotFoundError:
991 pass
994def _lfric_time_callback(cube: iris.cube.Cube):
995 """Fix time coordinate metadata if missing dimensions.
997 Some model data does not contain forecast_reference_time or forecast_period as
998 expected coordinates, and so we cannot aggregate over case studies without this
999 metadata. This callback fixes these issues.
1001 This callback also ensures all time coordinates are referenced as hours since
1002 1970-01-01 00:00:00 for consistency across different model inputs.
1004 Notes
1005 -----
1006 Some parts of the code have been adapted from Paul Earnshaw's scripts.
1007 """
1008 # Construct forecast_reference time if it doesn't exist.
1009 try:
1010 tcoord = cube.coord("time")
1011 # Set time coordinate to common basis "hours since 1970"
1012 try:
1013 tcoord.convert_units("hours since 1970-01-01 00:00:00")
1014 except ValueError:
1015 logger.warning("Unrecognised base time unit: %s", tcoord.units)
1017 if not cube.coords("forecast_reference_time"):
1018 try:
1019 init_time = datetime.datetime.fromisoformat(
1020 tcoord.attributes["time_origin"]
1021 )
1022 frt_point = tcoord.units.date2num(init_time)
1023 frt_coord = iris.coords.AuxCoord(
1024 frt_point,
1025 units=tcoord.units,
1026 standard_name="forecast_reference_time",
1027 long_name="forecast_reference_time",
1028 )
1029 cube.add_aux_coord(frt_coord)
1030 except KeyError:
1031 logger.warning(
1032 "Cannot find forecast_reference_time, but no `time_origin` attribute to construct it from."
1033 )
1035 # Remove time_origin to allow multiple case studies to merge.
1036 tcoord.attributes.pop("time_origin", None)
1038 # Construct forecast_period axis (forecast lead time) if it doesn't exist.
1039 if not cube.coords("forecast_period"):
1040 try:
1041 # Create array of forecast lead times.
1042 init_coord = cube.coord("forecast_reference_time")
1043 init_time_points_in_tcoord_units = tcoord.units.date2num(
1044 init_coord.units.num2date(init_coord.points)
1045 )
1046 lead_times = tcoord.points - init_time_points_in_tcoord_units
1048 # Get unit for lead time from time coordinate's unit.
1049 # Convert all lead time to hours for consistency between models.
1050 if "seconds" in str(tcoord.units): 1050 ↛ 1051line 1050 didn't jump to line 1051 because the condition on line 1050 was never true
1051 lead_times = lead_times / 3600.0
1052 units = "hours"
1053 elif "hours" in str(tcoord.units): 1053 ↛ 1056line 1053 didn't jump to line 1056 because the condition on line 1053 was always true
1054 units = "hours"
1055 else:
1056 raise ValueError(f"Unrecognised base time unit: {tcoord.units}")
1058 # Create lead time coordinate.
1059 lead_time_coord = iris.coords.AuxCoord(
1060 lead_times,
1061 standard_name="forecast_period",
1062 long_name="forecast_period",
1063 units=units,
1064 )
1066 # Associate lead time coordinate with time dimension.
1067 cube.add_aux_coord(lead_time_coord, cube.coord_dims("time"))
1068 except iris.exceptions.CoordinateNotFoundError:
1069 logger.warning(
1070 "Cube does not have both time and forecast_reference_time coordinate, so cannot construct forecast_period"
1071 )
1072 except iris.exceptions.CoordinateNotFoundError:
1073 logger.warning("No time coordinate on cube.")
1076def _lfric_forecast_period_callback(cube: iris.cube.Cube):
1077 """Check forecast_period name and units."""
1078 try:
1079 coord = cube.coord("forecast_period")
1080 if coord.units != "hours":
1081 cube.coord("forecast_period").convert_units("hours")
1082 if not coord.standard_name:
1083 coord.standard_name = "forecast_period"
1084 except iris.exceptions.CoordinateNotFoundError:
1085 pass
1088def _fix_no_time_coords_callback(cube: iris.cube.Cube):
1089 """Add dummy time coord to process cubes that don't have sequence coord."""
1090 # Only add if time coordinate does not exist.
1091 if not cube.coords("time"):
1092 cube.add_aux_coord(
1093 iris.coords.DimCoord(
1094 0, standard_name="time", units="hours since 0001-01-01 00:00:00"
1095 )
1096 )
1098 return cube
1101def _normalise_ML_varname(cube: iris.cube.Cube):
1102 """Fix plev variable names to standard names."""
1103 if cube.coords("pressure"):
1104 if cube.name() == "x_wind":
1105 cube.long_name = "zonal_wind_at_pressure_levels"
1106 if cube.name() == "y_wind":
1107 cube.long_name = "meridional_wind_at_pressure_levels"
1108 if cube.name() == "air_temperature":
1109 cube.long_name = "temperature_at_pressure_levels"
1110 if cube.name() == "specific_humidity": 1110 ↛ 1111line 1110 didn't jump to line 1111 because the condition on line 1110 was never true
1111 cube.long_name = (
1112 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging"
1113 )
1114 else:
1115 if cube.name() == "x_wind" and cube.var_name == "u_wind_at_10m": 1115 ↛ 1116line 1115 didn't jump to line 1116 because the condition on line 1115 was never true
1116 cube.long_name = "eastward_wind_at_10m"
1117 if cube.name() == "y_wind" and cube.var_name == "v_wind_at_10m": 1117 ↛ 1118line 1117 didn't jump to line 1118 because the condition on line 1117 was never true
1118 cube.long_name = "northward_wind_at_10m"
1121def _check_combine_point_observations(cubes: iris.cube.CubeList):
1122 """Enable cubes containing different point observation sources to be concatenated."""
1123 nstation = 0
1124 for cube in cubes:
1125 if "station" in [coord.name() for coord in cube.coords(dim_coords=True)]:
1126 if "obs_source" in [coord.name() for coord in cube.coords()]:
1127 cube.remove_coord("obs_source")
1128 cube.coord("station").points = cube.coord("station").points + nstation
1129 nstation = nstation + len(cube.coord("station").points)
1131 return cubes