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