Coverage for src/CSET/operators/read.py: 94%
453 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 12:33 +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 """
434 # List of cube attributes to preserve
435 attrs_to_preserve = [
436 "model_name",
437 "cset_comparison_base",
438 "STASH",
439 "um_stash_source",
440 ]
442 # Iterate through each cube and check each attribute.
443 for cube in iter_maybe(cubes):
444 for attr in list(cube.attributes):
445 if attr not in attrs_to_preserve:
446 cube.attributes.pop(attr, None)
448 return cubes
451def _loading_callback(cube: iris.cube.Cube, field, filename: str) -> iris.cube.Cube:
452 """Compose together the needed callbacks into a single function."""
453 # Most callbacks operate in-place, but save the cube when returned!
454 _remove_cset_comparison_base_attribute_callback(cube)
455 _realization_callback(cube)
456 _um_normalise_callback(cube)
457 _lfric_stash_source_callback(cube)
458 cube = _lfric_time_coord_fix_callback(cube)
459 _normalise_var0_varname(cube)
460 cube = _fix_no_spatial_coords_callback(cube)
461 _fix_spatial_coords_callback(cube)
462 _fix_pressure_coord_callback(cube)
463 _fix_um_radtime(cube)
464 _fix_cell_methods(cube)
465 cube = _convert_cube_units_callback(cube)
466 cube = _grid_longitude_fix_callback(cube)
467 _fix_lfric_cloud_base_altitude(cube)
468 _proleptic_gregorian_fix(cube)
469 _lfric_time_callback(cube)
470 _lfric_forecast_period_callback(cube)
471 cube = _fix_no_time_coords_callback(cube)
472 _normalise_longname(cube)
473 return cube
476def _remove_cset_comparison_base_attribute_callback(cube):
477 """Remove ``cset_comparison_base`` attribute if present.
479 This allows for reprocessing output previously saved by CSET.
480 """
481 cube.attributes.pop("cset_comparison_base", None)
484def _realization_callback(cube):
485 """Add a realization coordinate initialised to 0 if missing.
487 This means deterministic and ensemble cubes can assume realization coordinate through the rest
488 of the code.
489 """
490 # Only add if realization coordinate does not exist.
491 if not cube.coords("realization"):
492 cube.add_aux_coord(
493 iris.coords.DimCoord(0, standard_name="realization", units="1")
494 )
497@functools.lru_cache(None)
498def _log_once(msg, level=logging.WARNING):
499 """Print a warning message, skipping recent duplicates."""
500 logger.log(level, msg)
503def _um_normalise_callback(cube: iris.cube.Cube):
504 """Normalise UM STASH variable long names to LFRic variable names.
506 Note standard names will remain associated with cubes where different.
507 Long name will be used consistently in output filename and titles.
508 """
509 # Convert STASH to LFRic variable name
510 if "STASH" in cube.attributes:
511 stash = cube.attributes["STASH"]
512 try:
513 (name, grid) = STASH_TO_LFRIC[str(stash)]
514 cube.long_name = name
515 except KeyError:
516 # Don't change cubes with unknown stash codes.
517 _log_once(
518 f"Unknown STASH code: {stash}. Please check file stash_to_lfric.py to update.",
519 level=logging.WARNING,
520 )
523def _lfric_stash_source_callback(cube: iris.cube.Cube):
524 """Add um stash source to LFRic data.
526 The um_stash_source is a list of STASH codes for when an LFRic field maps to
527 multiple UM fields, however it can be encoded in any order. This attribute
528 is sorted to prevent this. This attribute is only present in LFRic data that
529 has been converted to look like UM data.
530 """
531 # Sort STASH code list.
532 stash_list = cube.attributes.get("um_stash_source")
533 if stash_list:
534 # Parse the string as a list, sort, then re-encode as a string.
535 cube.attributes["um_stash_source"] = str(sorted(ast.literal_eval(stash_list)))
538def _lfric_time_coord_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube:
539 """Ensure the time coordinate is a DimCoord rather than an AuxCoord.
541 The coordinate is converted and replaced if not. SLAMed LFRic data has this
542 issue, though the coordinate satisfies all the properties for a DimCoord.
543 Scalar time values are left as AuxCoords.
544 """
545 # This issue seems to come from iris's handling of NetCDF files where time
546 # always ends up as an AuxCoord.
547 if cube.coords("time"):
548 time_coord = cube.coord("time")
549 if (
550 not isinstance(time_coord, iris.coords.DimCoord)
551 and len(cube.coord_dims(time_coord)) == 1
552 ):
553 # Fudge the bounds to foil checking for strict monotonicity.
554 if ( 554 ↛ 558line 554 didn't jump to line 558 because the condition on line 554 was never true
555 time_coord.has_bounds()
556 and (time_coord.bounds[-1][0] - time_coord.bounds[0][0]) < 1.0e-8
557 ):
558 time_coord.bounds = [
559 [
560 time_coord.bounds[i][0] + 1.0e-8 * float(i),
561 time_coord.bounds[i][1],
562 ]
563 for i in range(len(time_coord.bounds))
564 ]
565 iris.util.promote_aux_coord_to_dim_coord(cube, time_coord)
566 return cube
569def _grid_longitude_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube:
570 """Check grid_longitude coordinates are in the range -180 deg to 180 deg.
572 This is necessary if comparing two models with different conventions --
573 for example, models where the prime meridian is defined as 0 deg or
574 360 deg. If not in the range -180 deg to 180 deg, we wrap the grid_longitude
575 so that it falls in this range. Checks are for near-180 bounds given
576 model data bounds may not extend exactly to 0. or 360.
577 Input cubes on non-rotated grid coordinates are not impacted.
578 """
579 try:
580 y, x = get_cube_yxcoordname(cube)
581 except ValueError:
582 # Don't modify non-spatial cubes.
583 return cube
585 long_coord = cube.coord(x)
586 # Wrap longitudes if rotated pole coordinates
587 coord_system = long_coord.coord_system
588 if x == "grid_longitude" and isinstance(
589 coord_system, iris.coord_systems.RotatedGeogCS
590 ):
591 long_points = long_coord.points.copy()
592 long_centre = np.median(long_points)
593 while long_centre < -175.0:
594 long_centre += 360.0
595 long_points += 360.0
596 while long_centre >= 175.0:
597 long_centre -= 360.0
598 long_points -= 360.0
599 long_coord.points = long_points
601 # Update coord bounds to be consistent with wrapping.
602 if long_coord.has_bounds():
603 long_coord.bounds = None
604 long_coord.guess_bounds()
606 return cube
609def _fix_no_spatial_coords_callback(cube: iris.cube.Cube):
610 import CSET.operators._utils as utils
612 # Don't modify spatial cubes that already have spatial dimensions
613 if utils.is_spatialdim(cube):
614 return cube
616 else:
617 # attempt to get lat/long from cube attributes
618 try:
619 lat_min = cube.attributes.get("geospatial_lat_min")
620 lat_max = cube.attributes.get("geospatial_lat_max")
621 lon_min = cube.attributes.get("geospatial_lon_min")
622 lon_max = cube.attributes.get("geospatial_lon_max")
624 lon_val = (lon_min + lon_max) / 2.0
625 lat_val = (lat_min + lat_max) / 2.0
627 lat_coord = iris.coords.DimCoord(
628 lat_val,
629 standard_name="latitude",
630 units="degrees_north",
631 var_name="latitude",
632 coord_system=iris.coord_systems.GeogCS(6371229.0),
633 circular=True,
634 )
636 lon_coord = iris.coords.DimCoord(
637 lon_val,
638 standard_name="longitude",
639 units="degrees_east",
640 var_name="longitude",
641 coord_system=iris.coord_systems.GeogCS(6371229.0),
642 circular=True,
643 )
645 cube.add_aux_coord(lat_coord)
646 cube.add_aux_coord(lon_coord)
647 return cube
649 # if lat/long are not in attributes, then return cube unchanged:
650 except TypeError:
651 return cube
654def _fix_spatial_coords_callback(cube: iris.cube.Cube):
655 """Check latitude and longitude coordinates name.
657 This is necessary as some models define their grid as on rotated
658 'grid_latitude' and 'grid_longitude' coordinates while others define
659 the grid on non-rotated 'latitude' and 'longitude'.
660 Cube dimensions need to be made consistent to avoid recipe failures,
661 particularly where comparing multiple input models with differing spatial
662 coordinates.
663 """
664 # Check if cube is spatial.
665 if not is_spatialdim(cube):
666 # Don't modify non-spatial cubes.
667 return
669 # Get spatial coords and dimension index.
670 y_name, x_name = get_cube_yxcoordname(cube)
671 ny = get_cube_coordindex(cube, y_name)
672 nx = get_cube_coordindex(cube, x_name)
674 # Remove spatial coords bounds if erroneous values detected.
675 # Aims to catch some errors in input coord bounds by setting
676 # invalid threshold of 10000.0
677 if cube.coord(x_name).has_bounds() and cube.coord(y_name).has_bounds():
678 bx_max = np.max(np.abs(cube.coord(x_name).bounds))
679 by_max = np.max(np.abs(cube.coord(y_name).bounds))
680 if bx_max > 10000.0 or by_max > 10000.0:
681 cube.coord(x_name).bounds = None
682 cube.coord(y_name).bounds = None
684 # Translate [grid_latitude, grid_longitude] to an unrotated 1-d DimCoord
685 # [latitude, longitude] for instances where rotated_pole=90.0
686 if "grid_latitude" in [coord.name() for coord in cube.coords(dim_coords=True)]:
687 coord_system = cube.coord("grid_latitude").coord_system
688 pole_lat = getattr(coord_system, "grid_north_pole_latitude", None)
689 if pole_lat == 90.0: 689 ↛ 690line 689 didn't jump to line 690 because the condition on line 689 was never true
690 lats = cube.coord("grid_latitude").points
691 lons = cube.coord("grid_longitude").points
693 cube.remove_coord("grid_latitude")
694 cube.add_dim_coord(
695 iris.coords.DimCoord(
696 lats,
697 standard_name="latitude",
698 var_name="latitude",
699 units="degrees",
700 coord_system=iris.coord_systems.GeogCS(6371229.0),
701 circular=True,
702 ),
703 ny,
704 )
705 y_name = "latitude"
706 cube.remove_coord("grid_longitude")
707 cube.add_dim_coord(
708 iris.coords.DimCoord(
709 lons,
710 standard_name="longitude",
711 var_name="longitude",
712 units="degrees",
713 coord_system=iris.coord_systems.GeogCS(6371229.0),
714 circular=True,
715 ),
716 nx,
717 )
718 x_name = "longitude"
720 # Create additional AuxCoord [grid_latitude, grid_longitude] with
721 # rotated pole attributes for cases with [lat, lon] inputs
722 if y_name in ["latitude"] and cube.coord(y_name).units in [
723 "degrees",
724 "degrees_north",
725 "degrees_south",
726 ]:
727 # Add grid_latitude AuxCoord
728 if "grid_latitude" not in [
729 coord.name() for coord in cube.coords(dim_coords=False)
730 ]:
731 cube.add_aux_coord(
732 iris.coords.AuxCoord(
733 cube.coord(y_name).points,
734 var_name="grid_latitude",
735 units="degrees",
736 ),
737 ny,
738 )
739 # Ensure input latitude DimCoord has CoordSystem
740 # This attribute is sometimes lost on iris.save
741 if not cube.coord(y_name).coord_system:
742 cube.coord(y_name).coord_system = iris.coord_systems.GeogCS(6371229.0)
744 if x_name in ["longitude"] and cube.coord(x_name).units in [
745 "degrees",
746 "degrees_west",
747 "degrees_east",
748 ]:
749 # Add grid_longitude AuxCoord
750 if "grid_longitude" not in [
751 coord.name() for coord in cube.coords(dim_coords=False)
752 ]:
753 cube.add_aux_coord(
754 iris.coords.AuxCoord(
755 cube.coord(x_name).points,
756 var_name="grid_longitude",
757 units="degrees",
758 ),
759 nx,
760 )
762 # Ensure input longitude DimCoord has CoordSystem
763 # This attribute is sometimes lost on iris.save
764 if not cube.coord(x_name).coord_system:
765 cube.coord(x_name).coord_system = iris.coord_systems.GeogCS(6371229.0)
768def _fix_pressure_coord_callback(cube: iris.cube.Cube):
769 """Rename pressure coordinate to "pressure" if it exists and ensure hPa units.
771 This problem was raised because the AIFS model data from ECMWF
772 defines the pressure coordinate with the name "pressure_level" rather
773 than compliant CF coordinate names.
775 Additionally, set the units of pressure to be hPa to be consistent with the UM,
776 and approach the coordinates in a unified way.
777 """
778 for coord in cube.dim_coords:
779 if coord.name() in ["pressure_level", "pressure_levels"]:
780 coord.rename("pressure")
782 if coord.name() == "pressure" and str(cube.coord("pressure").units) != "hPa":
783 cube.coord("pressure").convert_units("hPa")
786def _fix_um_radtime(cube: iris.cube.Cube):
787 """Move radiation diagnostics from timestamps which are output N minutes or seconds past every hour.
789 This callback does not have any effect for output diagnostics with
790 timestamps exactly 00 or 30 minutes past the hour. Only radiation
791 diagnostics are checked.
792 Note this callback does not interpolate the data in time, only adjust
793 timestamps to sit on the hour to enable time-to-time difference plotting
794 with models which may output radiation data on the hour.
795 """
796 try:
797 if cube.attributes["STASH"] in [
798 "m01s01i207",
799 "m01s01i208",
800 "m01s02i205",
801 "m01s02i201",
802 "m01s01i207",
803 "m01s02i207",
804 "m01s01i235",
805 ]:
806 time_coord = cube.coord("time")
808 # Convert time points to datetime objects
809 time_unit = time_coord.units
810 time_points = time_unit.num2date(time_coord.points)
811 # Skip if times don't need fixing.
812 if time_points[0].minute == 0 and time_points[0].second == 0:
813 return
814 if time_points[0].minute == 30 and time_points[0].second == 0: 814 ↛ 815line 814 didn't jump to line 815 because the condition on line 814 was never true
815 return
817 # Subtract time difference from the hour from each time point
818 n_minute = time_points[0].minute
819 n_second = time_points[0].second
820 # If times closer to next hour, compute difference to add on to following hour
821 if n_minute > 30:
822 n_minute = n_minute - 60
823 # Compute new diagnostic time stamp
824 new_time_points = (
825 time_points
826 - datetime.timedelta(minutes=n_minute)
827 - datetime.timedelta(seconds=n_second)
828 )
830 # Convert back to numeric values using the original time unit.
831 new_time_values = time_unit.date2num(new_time_points)
833 # Replace the time coordinate with updated values.
834 time_coord.points = new_time_values
836 # Recompute forecast_period with corrected values.
837 if cube.coord("forecast_period"): 837 ↛ exitline 837 didn't return from function '_fix_um_radtime' because the condition on line 837 was always true
838 fcst_prd_points = cube.coord("forecast_period").points
839 new_fcst_points = (
840 time_unit.num2date(fcst_prd_points)
841 - datetime.timedelta(minutes=n_minute)
842 - datetime.timedelta(seconds=n_second)
843 )
844 cube.coord("forecast_period").points = time_unit.date2num(
845 new_fcst_points
846 )
847 except KeyError:
848 pass
851def _fix_cell_methods(cube: iris.cube.Cube):
852 """To fix the assumed cell_methods in accumulation STASH from UM.
854 Lightning (m01s21i104), rainfall amount (m01s04i201, m01s05i201) and snowfall amount
855 (m01s04i202, m01s05i202) in UM is being output as a time accumulation,
856 over each hour (TAcc1hr), but input cubes show cell_methods as "mean".
857 For UM and LFRic inputs to be compatible, we assume accumulated cell_methods are
858 "sum". This callback changes "mean" cube attribute cell_method to "sum",
859 enabling the cell_method constraint on reading to select correct input.
860 """
861 # Shift "mean" cell_method to "sum" for selected UM inputs.
862 if cube.attributes.get("STASH") in [
863 "m01s21i104",
864 "m01s04i201",
865 "m01s04i202",
866 "m01s05i201",
867 "m01s05i202",
868 ] and {cm.method for cm in cube.cell_methods} == {"mean"}:
869 # Retrieve interval and any comment information.
870 for cell_method in cube.cell_methods:
871 interval_str = cell_method.intervals
872 comment_str = cell_method.comments
874 # Remove input aggregation method.
875 cube.cell_methods = ()
877 # Replace "mean" with "sum" cell_method to indicate aggregation.
878 cube.add_cell_method(
879 iris.coords.CellMethod(
880 method="sum",
881 coords="time",
882 intervals=interval_str,
883 comments=comment_str,
884 )
885 )
888def _convert_cube_units_callback(cube: iris.cube.Cube):
889 """Adjust diagnostic units for specific variables.
891 Some precipitation diagnostics are output with unit kg m-2 s-1 and are
892 converted here to mm hr-1.
894 Visibility diagnostics are converted here from m to km to improve output
895 formatting.
896 """
897 # Convert precipitation diagnostic units if required.
898 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
899 if any("surface_microphysical" in name for name in varnames):
900 if cube.units == "kg m-2 s-1":
901 _log_once(
902 "Converting precipitation rate units from kg m-2 s-1 to mm hr-1",
903 level=logging.DEBUG,
904 )
905 # Convert from kg m-2 s-1 to mm s-1 assuming 1kg water = 1l water = 1dm^3 water.
906 # This is a 1:1 conversion, so we just change the units.
907 cube.units = "mm s-1"
908 # Convert the units to per hour.
909 cube.convert_units("mm hr-1")
910 elif cube.units == "kg m-2": 910 ↛ 920line 910 didn't jump to line 920 because the condition on line 910 was always true
911 _log_once(
912 "Converting precipitation amount units from kg m-2 to mm",
913 level=logging.DEBUG,
914 )
915 # Convert from kg m-2 to mm assuming 1kg water = 1l water = 1dm^3 water.
916 # This is a 1:1 conversion, so we just change the units.
917 cube.units = "mm"
919 # Convert visibility diagnostic units if required.
920 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
921 if any("visibility" in name for name in varnames) and cube.units == "m":
922 _log_once("Converting visibility units m to km.", level=logging.DEBUG)
923 # Convert the units to km.
924 cube.convert_units("km")
926 return cube
929def _fix_lfric_cloud_base_altitude(cube: iris.cube.Cube):
930 """Mask cloud_base_altitude diagnostic in regions with no cloud."""
931 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
932 if any("cloud_base_altitude" in name for name in varnames):
933 # Mask cube where set > 144kft to catch default 144.35695538058164
934 cube.data = dask.array.ma.masked_greater(cube.core_data(), 144.0)
937def _compute_winds(
938 cubes: iris.cube.CubeList,
939 constraint: iris.Constraint | ConstraintCombination | None = None,
940):
941 """To compute wind_speed from vector components if not available as diagnostic.
943 Diagnostics of wind are also not always consistent between the UM
944 and LFRic. Here, winds from the UM are adjusted to make them
945 consistent with LFRic.
946 """
947 # Check whether we have components of the wind identified by varname
948 # but not the wind speed and calculate it if it is missing. Note that
949 # this will be biased low in general because the components will mostly
950 # be time averages. For simplicity, we do this only if there is just one
951 # cube of a component. A more complicated approach would be to consider
952 # the cell methods, but it may not be warranted.
953 #
954 # A check on UM STASH attributes is also conducted to adjust directions.
956 if constraint is None:
957 return cubes
958 filter_windspeed = None
959 for constr in _flatten_combined_constraint(constraint):
960 filter_windspeed = getattr(constr, "varname", None)
961 if filter_windspeed:
962 constraint = constr
963 break
965 u_constr = iris.Constraint("eastward_wind_at_10m")
966 v_constr = iris.Constraint("northward_wind_at_10m")
967 sp_constr = iris.Constraint("wind_speed_at_10m")
969 try:
970 if (
971 cubes.extract(u_constr)
972 and cubes.extract(v_constr)
973 and not cubes.extract(sp_constr)
974 ):
975 if "wind_speed_at_10m" in constraint.varname:
976 _add_wind_speed_um(cubes)
977 # Convert winds in the UM to be relative to true east and true north.
978 _convert_wind_true_dirn_um(cubes)
979 except (KeyError, AttributeError):
980 pass
982 if filter_windspeed:
983 filter_windspeed_constraint = iris.Constraint(
984 cube_func=lambda cube: (
985 cube.long_name in filter_windspeed
986 or cube.standard_name in filter_windspeed
987 or cube.var_name in filter_windspeed
988 )
989 )
990 cubes = cubes.extract(filter_windspeed_constraint)
991 return cubes
994def _flatten_combined_constraint(
995 con: iris.Constraint | ConstraintCombination,
996) -> Iterator[iris.Constraint]:
997 # yields constraints of a possibly nested constraint combination
998 if isinstance(con, ConstraintCombination):
999 yield from _flatten_combined_constraint(con.lhs)
1000 yield from _flatten_combined_constraint(con.rhs)
1001 else:
1002 yield con
1005def _add_wind_speed_um(cubes: iris.cube.CubeList):
1006 """Add windspeeds to cubes from components."""
1007 u_wind = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m"))
1008 v_wind = cubes.extract_cube(iris.Constraint("northward_wind_at_10m"))
1009 wspd10 = (u_wind**2 + v_wind**2) ** 0.5
1010 wspd10.attributes["STASH"] = "m01s03i227"
1011 wspd10.standard_name = "wind_speed"
1012 wspd10.long_name = "wind_speed_at_10m"
1013 wspd10.units = "m s-1"
1014 cubes.append(wspd10)
1017def _convert_wind_true_dirn_um(cubes: iris.cube.CubeList):
1018 """To convert winds to true directions.
1020 Convert from the components relative to the grid to true directions.
1021 This functionality only handles the simplest case.
1022 Constrains using STASH code only to ensure applied to UM outputs only.
1023 """
1024 u_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i225"))
1025 v_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i226"))
1026 for u, v in zip(u_grids, v_grids, strict=True): 1026 ↛ 1027line 1026 didn't jump to line 1027 because the loop on line 1026 never started
1027 true_u, true_v = rotate_winds(u, v, iris.coord_systems.GeogCS(6371229.0))
1028 u.data = true_u.core_data()
1029 v.data = true_v.core_data()
1032def _normalise_var0_varname(cube: iris.cube.Cube):
1033 """Fix varnames for consistency to allow merging.
1035 Some model data netCDF sometimes have a coordinate name end in
1036 "_0" etc, where duplicate coordinates of same name are defined but
1037 with different attributes. This can be inconsistently managed in
1038 different model inputs and can cause cubes to fail to merge.
1039 """
1040 for coord in cube.coords():
1041 if coord.var_name and coord.var_name.endswith("_0"):
1042 coord.var_name = coord.var_name.removesuffix("_0")
1043 if coord.var_name and coord.var_name.endswith("_1"):
1044 coord.var_name = coord.var_name.removesuffix("_1")
1045 if coord.var_name and coord.var_name.endswith("_2"): 1045 ↛ 1046line 1045 didn't jump to line 1046 because the condition on line 1045 was never true
1046 coord.var_name = coord.var_name.removesuffix("_2")
1047 if coord.var_name and coord.var_name.endswith("_3"): 1047 ↛ 1048line 1047 didn't jump to line 1048 because the condition on line 1047 was never true
1048 coord.var_name = coord.var_name.removesuffix("_3")
1050 if cube.var_name and cube.var_name.endswith("_0"):
1051 cube.var_name = cube.var_name.removesuffix("_0")
1054def _proleptic_gregorian_fix(cube: iris.cube.Cube):
1055 """Convert the calendars of time units to use a standard calendar."""
1056 try:
1057 time_coord = cube.coord("time")
1058 if time_coord.units.calendar == "proleptic_gregorian":
1059 logger.debug(
1060 "Changing proleptic Gregorian calendar to standard calendar for %s",
1061 repr(time_coord.units),
1062 )
1063 time_coord.units = time_coord.units.change_calendar("standard")
1064 except iris.exceptions.CoordinateNotFoundError:
1065 pass
1068def _lfric_time_callback(cube: iris.cube.Cube):
1069 """Fix time coordinate metadata if missing dimensions.
1071 Some model data does not contain forecast_reference_time or forecast_period as
1072 expected coordinates, and so we cannot aggregate over case studies without this
1073 metadata. This callback fixes these issues.
1075 This callback also ensures all time coordinates are referenced as hours since
1076 1970-01-01 00:00:00 for consistency across different model inputs.
1078 Notes
1079 -----
1080 Some parts of the code have been adapted from Paul Earnshaw's scripts.
1081 """
1082 # Construct forecast_reference time if it doesn't exist.
1083 try:
1084 tcoord = cube.coord("time")
1085 # Set time coordinate to common basis "hours since 1970"
1086 try:
1087 tcoord.convert_units("hours since 1970-01-01 00:00:00")
1088 except ValueError:
1089 logger.warning("Unrecognised base time unit: %s", tcoord.units)
1091 if not cube.coords("forecast_reference_time"):
1092 try:
1093 init_time = datetime.datetime.fromisoformat(
1094 tcoord.attributes["time_origin"]
1095 )
1096 frt_point = tcoord.units.date2num(init_time)
1097 frt_coord = iris.coords.AuxCoord(
1098 frt_point,
1099 units=tcoord.units,
1100 standard_name="forecast_reference_time",
1101 long_name="forecast_reference_time",
1102 )
1103 cube.add_aux_coord(frt_coord)
1104 except KeyError:
1105 logger.warning(
1106 "Cannot find forecast_reference_time, but no `time_origin` attribute to construct it from."
1107 )
1109 # Construct forecast_period axis (forecast lead time) if it doesn't exist.
1110 if not cube.coords("forecast_period"):
1111 try:
1112 # Create array of forecast lead times.
1113 init_coord = cube.coord("forecast_reference_time")
1114 init_time_points_in_tcoord_units = tcoord.units.date2num(
1115 init_coord.units.num2date(init_coord.points)
1116 )
1117 lead_times = tcoord.points - init_time_points_in_tcoord_units
1119 # Get unit for lead time from time coordinate's unit.
1120 # Convert all lead time to hours for consistency between models.
1121 if "seconds" in str(tcoord.units): 1121 ↛ 1122line 1121 didn't jump to line 1122 because the condition on line 1121 was never true
1122 lead_times = lead_times / 3600.0
1123 units = "hours"
1124 elif "hours" in str(tcoord.units): 1124 ↛ 1127line 1124 didn't jump to line 1127 because the condition on line 1124 was always true
1125 units = "hours"
1126 else:
1127 raise ValueError(f"Unrecognised base time unit: {tcoord.units}")
1129 # Create lead time coordinate.
1130 lead_time_coord = iris.coords.AuxCoord(
1131 lead_times,
1132 standard_name="forecast_period",
1133 long_name="forecast_period",
1134 units=units,
1135 )
1137 # Associate lead time coordinate with time dimension.
1138 cube.add_aux_coord(lead_time_coord, cube.coord_dims("time"))
1139 except iris.exceptions.CoordinateNotFoundError:
1140 logger.warning(
1141 "Cube does not have both time and forecast_reference_time coordinate, so cannot construct forecast_period"
1142 )
1143 except iris.exceptions.CoordinateNotFoundError:
1144 logger.warning("No time coordinate on cube.")
1147def _lfric_forecast_period_callback(cube: iris.cube.Cube):
1148 """Check forecast_period name and units."""
1149 try:
1150 coord = cube.coord("forecast_period")
1151 if coord.units != "hours":
1152 cube.coord("forecast_period").convert_units("hours")
1153 if not coord.standard_name:
1154 coord.standard_name = "forecast_period"
1155 except iris.exceptions.CoordinateNotFoundError:
1156 pass
1159def _fix_no_time_coords_callback(cube: iris.cube.Cube):
1160 """Add dummy time coord to process cubes that don't have sequence coord."""
1161 # Only add if time coordinate does not exist.
1162 if not cube.coords("time"):
1163 cube.add_aux_coord(
1164 iris.coords.DimCoord(
1165 0, standard_name="time", units="hours since 0001-01-01 00:00:00"
1166 )
1167 )
1169 return cube
1172def _normalise_longname(cube: iris.cube.Cube):
1173 """Fix plev variable names to standard names."""
1174 if cube.coords("pressure"):
1175 if cube.name() == "x_wind":
1176 cube.long_name = "zonal_wind_at_pressure_levels"
1177 if cube.name() == "y_wind":
1178 cube.long_name = "meridional_wind_at_pressure_levels"
1179 if cube.name() == "air_temperature":
1180 cube.long_name = "temperature_at_pressure_levels"
1181 if cube.name() == "specific_humidity": 1181 ↛ 1182line 1181 didn't jump to line 1182 because the condition on line 1181 was never true
1182 cube.long_name = (
1183 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging"
1184 )
1185 else:
1186 if cube.name() == "x_wind" and cube.var_name == "u_wind_at_10m": 1186 ↛ 1187line 1186 didn't jump to line 1187 because the condition on line 1186 was never true
1187 cube.long_name = "eastward_wind_at_10m"
1188 if cube.name() == "y_wind" and cube.var_name == "v_wind_at_10m": 1188 ↛ 1189line 1188 didn't jump to line 1189 because the condition on line 1188 was never true
1189 cube.long_name = "northward_wind_at_10m"
1192def _check_combine_point_observations(cubes: iris.cube.CubeList):
1193 """Enable cubes containing different point observation sources to be concatenated."""
1194 nstation = 0
1195 for cube in cubes:
1196 if "station" in [coord.name() for coord in cube.coords(dim_coords=True)]:
1197 if "obs_source" in [coord.name() for coord in cube.coords()]:
1198 cube.remove_coord("obs_source")
1199 cube.coord("station").points = cube.coord("station").points + nstation
1200 nstation = nstation + len(cube.coord("station").points)
1202 return cubes