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