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