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