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