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