Coverage for src/CSET/operators/regrid.py: 56%
206 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-05-27 12:45 +0000
1# © Crown copyright, Met Office (2022-2024) 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 to regrid cubes."""
17import logging
18import os
19import re
20import time
21import warnings
22from multiprocessing import Pool
24import iris
25import iris.coord_systems
26import iris.coords as icoords
27import iris.cube
28import numpy as np
29from scipy.interpolate import LinearNDInterpolator
31from CSET._common import iter_maybe
32from CSET.operators._utils import get_cube_yxcoordname
35class BoundaryWarning(UserWarning):
36 """Selected gridpoint is close to the domain edge.
38 In many cases gridpoints near the domain boundary contain non-physical
39 values, so caution is advised when interpreting them.
40 """
43def regrid_onto_cube(
44 toregrid: iris.cube.Cube | iris.cube.CubeList,
45 target: iris.cube.Cube,
46 method: str,
47 **kwargs,
48) -> iris.cube.Cube | iris.cube.CubeList:
49 """Regrid a cube or CubeList, projecting onto a target cube.
51 All cubes must have at least 2 spatial (map) dimensions.
53 Arguments
54 ----------
55 toregrid: iris.cube | iris.cube.CubeList
56 An iris Cube of data to regrid, or multiple cubes to regrid in a
57 CubeList. A minimum requirement is that the cube(s) need to be 2D with a
58 latitude, longitude coordinates.
59 target: Cube
60 An iris cube of the data to regrid onto. It needs to be 2D with a
61 latitude, longitude coordinate.
62 method: str
63 Method used to regrid onto, etc. Linear will use iris.analysis.Linear()
65 Returns
66 -------
67 iris.cube | iris.cube.CubeList
68 An iris cube of the data that has been regridded, or a CubeList of the
69 cubes that have been regridded in the same order they were passed in
70 toregrid.
72 Raises
73 ------
74 ValueError
75 If a unique x/y coordinate cannot be found
76 NotImplementedError
77 If the cubes grid, or the method for regridding, is not yet supported.
79 Notes
80 -----
81 Currently rectlinear grids (uniform) are supported.
82 """
83 # To store regridded cubes.
84 regridded_cubes = iris.cube.CubeList()
86 # Iterate over all cubes and regrid.
87 for cube in iter_maybe(toregrid):
88 # Get y,x coord names
89 y_coord, x_coord = get_cube_yxcoordname(cube)
91 # List of supported grids - check if it is compatible
92 supported_grids = (iris.coord_systems.GeogCS, iris.coord_systems.RotatedGeogCS)
93 if not isinstance(cube.coord(x_coord).coord_system, supported_grids):
94 raise NotImplementedError(
95 f"Does not currently support {cube.coord(x_coord).coord_system} coordinate system"
96 )
97 if not isinstance(cube.coord(y_coord).coord_system, supported_grids):
98 raise NotImplementedError(
99 f"Does not currently support {cube.coord(y_coord).coord_system} coordinate system"
100 )
102 regrid_method = getattr(iris.analysis, method, None)
103 if callable(regrid_method):
104 regridded_cubes.append(cube.regrid(target, regrid_method()))
105 else:
106 raise NotImplementedError(
107 f"Does not currently support {method} regrid method"
108 )
110 # Preserve returning a cube if only a cube has been supplied to regrid.
111 if len(regridded_cubes) == 1:
112 return regridded_cubes[0]
113 else:
114 return regridded_cubes
117def regrid_onto_xyspacing(
118 toregrid: iris.cube.Cube | iris.cube.CubeList,
119 xspacing: float,
120 yspacing: float,
121 method: str,
122 **kwargs,
123) -> iris.cube.Cube | iris.cube.CubeList:
124 """Regrid cube or cubelist onto a set x,y spacing.
126 Regrid cube(s) using specified x,y spacing, which is performed linearly.
128 Parameters
129 ----------
130 toregrid: iris.cube | iris.cube.CubeList
131 An iris cube of the data to regrid, or multiple cubes to regrid in a
132 cubelist. A minimum requirement is that the cube(s) need to be 2D with a
133 latitude, longitude coordinates.
134 xspacing: float
135 Spacing of points in longitude direction (could be degrees, meters etc.)
136 yspacing: float
137 Spacing of points in latitude direction (could be degrees, meters etc.)
138 method: str
139 Method used to regrid onto, etc. Linear will use iris.analysis.Linear()
141 Returns
142 -------
143 iris.cube | iris.cube.CubeList
144 An iris cube of the data that has been regridded, or a cubelist of the
145 cubes that have been regridded in the same order they were passed in
146 toregrid.
148 Raises
149 ------
150 ValueError
151 If a unique x/y coordinate cannot be found
152 NotImplementedError
153 If the cubes grid, or the method for regridding, is not yet supported.
155 Notes
156 -----
157 Currently rectlinear grids (uniform) are supported.
159 """
160 # To store regridded cubes.
161 regridded_cubes = iris.cube.CubeList()
163 # Iterate over all cubes and regrid.
164 for cube in iter_maybe(toregrid):
165 # Get x,y coord names
166 y_coord, x_coord = get_cube_yxcoordname(cube)
168 # List of supported grids - check if it is compatible
169 supported_grids = (iris.coord_systems.GeogCS, iris.coord_systems.RotatedGeogCS)
170 if not isinstance(cube.coord(x_coord).coord_system, supported_grids):
171 raise NotImplementedError(
172 f"Does not currently support {cube.coord(x_coord).coord_system} regrid method"
173 )
174 if not isinstance(cube.coord(y_coord).coord_system, supported_grids):
175 raise NotImplementedError(
176 f"Does not currently support {cube.coord(y_coord).coord_system} regrid method"
177 )
179 # Get axis
180 lat, lon = cube.coord(y_coord), cube.coord(x_coord)
182 # Get bounds
183 lat_min, lon_min = lat.points.min(), lon.points.min()
184 lat_max, lon_max = lat.points.max(), lon.points.max()
186 # Generate new mesh
187 latout = np.arange(lat_min, lat_max, yspacing)
188 lonout = np.arange(lon_min, lon_max, xspacing)
190 regrid_method = getattr(iris.analysis, method, None)
191 if callable(regrid_method):
192 regridded_cubes.append(
193 cube.interpolate(
194 [(y_coord, latout), (x_coord, lonout)], regrid_method()
195 )
196 )
197 else:
198 raise NotImplementedError(
199 f"Does not currently support {method} regrid method"
200 )
202 # Preserve returning a cube if only a cube has been supplied to regrid.
203 if len(regridded_cubes) == 1:
204 return regridded_cubes[0]
205 else:
206 return regridded_cubes
209def regrid_to_single_point(
210 cubes: iris.cube.Cube | iris.cube.CubeList,
211 lat_pt: float,
212 lon_pt: float,
213 latlon_in_type: str = "rotated",
214 method: str = "Nearest",
215 boundary_margin: int = 8,
216 **kwargs,
217) -> iris.cube.Cube:
218 """Select data at a single point by longitude and latitude.
220 Selection of model grid point is performed by a regrid function, either
221 selecting the nearest gridpoint to the selected longitude and latitude
222 values or using linear interpolation across the surrounding points.
224 Parameters
225 ----------
226 cubes: Cube | CubeList
227 An iris cube or CubeList of the data to regrid. As a minimum, it needs
228 to be 2D with latitude, longitude coordinates.
229 lon_pt: float
230 Selected value of longitude: this should be in the range -180 degrees to
231 180 degrees.
232 lat_pt: float
233 Selected value of latitude: this should be in the range -90 degrees to
234 90 degrees.
235 latlon_in_type: str, optional
236 Specify whether the input longitude and latitude point is in standard
237 geographic realworld coordinates ("realworld") or on the rotated grid
238 of the cube ("rotated"). Default is "rotated".
239 method: str
240 Method used to determine the values at the selected longitude and
241 latitude. The recommended approach is to use iris.analysis.Nearest(),
242 which selects the nearest gridpoint. An alternative is
243 iris.analysis.Linear(), which obtains the values at the selected
244 longitude and latitude by linear interpolation.
245 boundary_margin: int, optional
246 Number of grid points from the domain boundary considered "unreliable".
247 Defaults to 8.
249 Returns
250 -------
251 regridded_cubes: Cube | CubeList
252 An iris cube or CubeList of the data at the specified point (this may
253 have time and/or height dimensions).
255 Raises
256 ------
257 ValueError
258 If a unique x/y coordinate cannot be found; also if, for selecting a
259 single gridpoint, the chosen longitude and latitude point is outside the
260 domain; also (currently) if the difference between the actual and target
261 points exceed 0.1 degrees.
262 NotImplementedError
263 If the cubes grid, or the method for regridding, is not yet supported.
265 Notes
266 -----
267 The acceptable coordinate names for X and Y coordinates are currently
268 described in X_COORD_NAMES and Y_COORD_NAMES. These cover commonly used
269 coordinate types, though a user can append new ones. Currently rectilinear
270 grids (uniform) are supported. Warnings are raised if the selected gridpoint
271 is within boundary_margin grid lengths of the domain boundary as data here
272 is potentially unreliable.
273 """
274 # To store regridded cubes.
275 regridded_cubes = iris.cube.CubeList()
277 # Iterate over all cubes and regrid.
278 for cube in iter_maybe(cubes):
279 # Get x and y coordinate names.
280 y_coord, x_coord = get_cube_yxcoordname(cube)
282 # List of supported grids - check if it is compatible
283 # NOTE: The "RotatedGeogCS" option below seems to be required for rotated grids --
284 # this may need to be added in other places in these Operators.
285 supported_grids = (iris.coord_systems.GeogCS, iris.coord_systems.RotatedGeogCS)
286 if not isinstance(cube.coord(x_coord).coord_system, supported_grids):
287 raise NotImplementedError(
288 f"Does not currently support {cube.coord(x_coord).coord_system} regrid method"
289 )
290 if not isinstance(cube.coord(y_coord).coord_system, supported_grids):
291 raise NotImplementedError(
292 f"Does not currently support {cube.coord(y_coord).coord_system} regrid method"
293 )
295 # Transform input coordinates onto rotated grid if requested
296 if latlon_in_type == "realworld":
297 lon_tr, lat_tr = transform_lat_long_points(lon_pt, lat_pt, cube)
298 elif latlon_in_type == "rotated": 298 ↛ 302line 298 didn't jump to line 302 because the condition on line 298 was always true
299 lon_tr, lat_tr = lon_pt, lat_pt
301 # Get axis
302 lat, lon = cube.coord(y_coord), cube.coord(x_coord)
304 # Get bounds
305 lat_min, lon_min = lat.points.min(), lon.points.min()
306 lat_max, lon_max = lat.points.max(), lon.points.max()
308 # Use different logic for single point obs data.
309 if len(cube.coord(x_coord).points) > 1:
310 # Get boundaries of frame to avoid selecting gridpoint close to domain edge
311 lat_min_bound, lon_min_bound = (
312 lat.points[boundary_margin - 1],
313 lon.points[boundary_margin - 1],
314 )
315 lat_max_bound, lon_max_bound = (
316 lat.points[-boundary_margin],
317 lon.points[-boundary_margin],
318 )
320 # Check to see if selected point is outside the domain
321 if (lat_tr < lat_min) or (lat_tr > lat_max):
322 raise ValueError("Selected point is outside the domain.")
323 else:
324 if (lon_tr < lon_min) or (lon_tr > lon_max):
325 if (lon_tr + 360.0 >= lon_min) and (lon_tr + 360.0 <= lon_max):
326 lon_tr += 360.0
327 elif (lon_tr - 360.0 >= lon_min) and (lon_tr - 360.0 <= lon_max):
328 lon_tr -= 360.0
329 else:
330 raise ValueError("Selected point is outside the domain.")
332 # Check to see if selected point is near the domain boundaries
333 if (
334 (lat_tr < lat_min_bound)
335 or (lat_tr > lat_max_bound)
336 or (lon_tr < lon_min_bound)
337 or (lon_tr > lon_max_bound)
338 ):
339 warnings.warn(
340 f"Selected point is within {boundary_margin} gridlengths of the domain edge, data may be unreliable.",
341 category=BoundaryWarning,
342 stacklevel=2,
343 )
345 regrid_method = getattr(iris.analysis, method, None)
346 if not callable(regrid_method):
347 raise NotImplementedError(
348 f"Does not currently support {method} regrid method"
349 )
351 cube_rgd = cube.interpolate(((lat, lat_tr), (lon, lon_tr)), regrid_method())
352 regridded_cubes.append(cube_rgd)
353 else:
354 if (
355 np.abs((lat_tr - lat.points[0])) > 0.1
356 or np.abs((lon_tr - lon.points[0])) > 0.1
357 ):
358 raise ValueError(
359 "Selected point is too far from the specified coordinates. It should be within 0.1 degrees."
360 )
361 else:
362 print(
363 "*** lat/long diffs",
364 np.abs(lat_tr - lat_pt),
365 np.abs(lon_tr - lon_pt),
366 )
367 regridded_cubes.append(cube)
369 # Preserve returning a cube if only a cube has been supplied to regrid.
370 if len(regridded_cubes) == 1:
371 return regridded_cubes[0]
372 else:
373 return regridded_cubes
376def transform_lat_long_points(lon, lat, cube):
377 """Transform a selected point in longitude and latitude.
379 Transform the coordinates of a point from the real world
380 grid to the corresponding point on the rotated grid of a cube.
382 Parameters
383 ----------
384 cube: Cube
385 An iris cube of data defining the rotated grid to be used in
386 the longitude-latitude transformation.
387 lon: float
388 Selected value of longitude: this should be in the range -180 degrees to
389 180 degrees.
390 lat: float
391 Selected value of latitude: this should be in the range -90 degrees to
392 90 degrees.
394 Returns
395 -------
396 lon_rot, lat_rot: float
397 Coordinates of the selected point on the rotated grid specified within
398 the selected cube.
400 """
401 import cartopy.crs as ccrs
403 rot_pole = cube.coord_system().as_cartopy_crs()
404 true_grid = ccrs.Geodetic()
405 rot_coords = rot_pole.transform_point(lon, lat, true_grid)
406 lon_rot = rot_coords[0]
407 lat_rot = rot_coords[1]
409 return lon_rot, lat_rot
412def interpolate_to_point_cube(
413 fld: iris.cube.Cube | iris.cube.CubeList, point_cube: iris.cube.Cube, **kwargs
414) -> iris.cube.Cube | iris.cube.CubeList:
415 """Interpolate from a 2D field to a set of points.
417 Interpolate the 2D field in fld to the set of points
418 specified in point_cube.
420 Parameters
421 ----------
422 fld: Cube
423 An iris cube containing a two-dimensional field.
424 point_cube: Cube
425 An iris cube specifying the point to which the data
426 will be interpolated.
428 Returns
429 -------
430 fld_point_cube: Cube
431 An iris cube containing interpolated values at the points
432 specified by the point cube.
434 """
435 #
436 # As a basis, create a copy of the point cube.
437 fld_point_cube = point_cube.copy()
438 # Get indices of positional coordinates. We assume that the
439 # point cube is unrotated.
440 klon = None
441 klat = None
442 for kc in range(len(fld_point_cube.aux_coords)):
443 if fld_point_cube.aux_coords[kc].standard_name == "latitude":
444 klat = kc
445 elif fld_point_cube.aux_coords[kc].standard_name == "longitude":
446 klon = kc
447 #
448 # The input may have a rotated coordinate system.
449 if len(fld.coords("grid_latitude")) > 0:
450 # Interpolate in rotated coordinates.
451 rot_csyst = fld.coords("grid_latitude")[0].coord_system
452 rotpt = iris.analysis.cartography.rotate_pole(
453 fld_point_cube.aux_coords[klon].points,
454 fld_point_cube.aux_coords[klat].points,
455 rot_csyst.grid_north_pole_longitude,
456 rot_csyst.grid_north_pole_latitude,
457 )
458 # Add other interpolation options later.
459 fld_interpolator = iris.analysis.Linear(extrapolation_mode="mask").interpolator(
460 fld, ["time", "grid_latitude", "grid_longitude"]
461 )
462 for jt in range(len(fld_point_cube.coords("time")[0].points)):
463 fld_point_cube.data[jt, :] = np.ma.masked_invalid(
464 [
465 fld_interpolator(
466 [
467 fld_point_cube.coord("time").points[jt],
468 rotpt[1][k],
469 rotpt[0][k],
470 ]
471 ).data
472 if ~point_cube.data.mask[jt][k]
473 else np.nan
474 for k in range(len(rotpt[0]))
475 ]
476 )
477 else:
478 # Add other interpolation options later.
479 fld_interpolator = iris.analysis.Linear(extrapolation_mode="mask").interpolator(
480 fld, ["time", "latitude", "longitude"]
481 )
482 for jt in range(len(fld_point_cube.coords("time")[0].points)):
483 fld_point_cube.data[jt, :] = np.ma.masked_invalid(
484 [
485 fld_interpolator(
486 [
487 fld_point_cube.coords("time")[0].points[jt],
488 fld_point_cube.coord("latitude").points[k],
489 fld_point_cube.coord("longitude").points[k],
490 ]
491 ).data
492 if ~point_cube.data.mask[jt][k]
493 else np.nan
494 for k in range(fld_point_cube.coord("latitude").points)
495 ]
496 )
497 return fld_point_cube
500def _rebuild_ugrid_meta(data, origcube, lat, lon):
501 """
502 Build a structured iris cube (time, lat, lon) from regridded data.
504 The cube will have metadata transferred and an additional pressure auxiliary
505 coordinate inferred from the cube name if present.
507 Parameters
508 ----------
509 data : np.ndarray
510 Regridded data array with shape (time, lat, lon)
511 origcube : iris.cube.Cube
512 Original unstructured source cube, used for plucking metadata.
513 lat : np.ndarray
514 1D latitude coordinate values of regridded data
515 lon : np.ndarray
516 1D longitude coordinate values of regridded data
518 Returns
519 -------
520 iris.cube.Cube
521 A structured iris cube with appropriate metadata.
522 """
523 # Lookup dictionary for standard anemoi ML variable names and their translation.
524 UGRID_VAR_LOOKUP = {
525 "t": {
526 "standard_name": "air_temperature",
527 "long_name": "temperature_at_pressure_levels",
528 "units": "K",
529 },
530 "u": {
531 "standard_name": "eastward_wind",
532 "long_name": "zonal_wind_at_pressure_levels",
533 "units": "m s-1",
534 },
535 "v": {
536 "standard_name": "northward_wind",
537 "long_name": "meridional_wind_at_pressure_levels",
538 "units": "m s-1",
539 },
540 "w": {
541 "standard_name": "upward_air_velocity",
542 "long_name": "vertical_wind_at_pressure_levels",
543 "units": "m s-1",
544 },
545 "q": {
546 "standard_name": "specific_humidity",
547 "long_name": "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",
548 "units": "kg kg-1",
549 },
550 "z": {
551 "standard_name": "geopotential_height",
552 "long_name": "geopotential_height_at_pressure_levels",
553 "units": "m",
554 },
555 "sp": {
556 "standard_name": "surface_air_pressure",
557 "long_name": "surface_air_pressure",
558 "units": "Pa",
559 },
560 "10u": {
561 "long_name": "eastward_wind_at_10m",
562 "units": "m s-1",
563 },
564 "10v": {
565 "long_name": "northward_wind_at_10m",
566 "units": "m s-1",
567 },
568 "lsm": {
569 "long_name": "land_binary_mask",
570 },
571 "2t": {
572 "long_name": "temperature_at_screen_level",
573 "units": "K",
574 },
575 "2d": {
576 "long_name": "dew_point_temperature_at_screen_level",
577 "units": "K",
578 },
579 "skt": {
580 "long_name": "grid_surface_temperature",
581 "units": "K",
582 },
583 "tp": {
584 "long_name": "surface_microphysical_rainfall_rate",
585 "units": "mm 6hr-1",
586 },
587 }
589 # Get original cube time coordinate dimension.
590 time_coord = origcube.coord("time")
592 # Create new latitude coordinate.
593 lat_coord = icoords.DimCoord(
594 lat,
595 standard_name="latitude",
596 units="degrees",
597 )
599 # Create new longitude coordinate.
600 lon_coord = icoords.DimCoord(
601 lon,
602 standard_name="longitude",
603 units="degrees",
604 )
606 # Parse cube name, to determine if it contains a likely pressure variable/level.
607 m = re.match(r"^([a-zA-Z][a-zA-Z0-9]*|\d+[a-zA-Z]+)(?:_(\d+))?$", origcube.name())
609 # If pattern is not None
610 if m:
611 # Extract variable and pressure from cube name components.
612 var_key, pressure_hpa = m.group(1), m.group(2)
614 # If there is a number in cube name that can be split.
615 if pressure_hpa is not None:
616 # Create new pressure coordinate dimension.
617 pressure_coord = icoords.DimCoord(
618 [int(pressure_hpa)],
619 long_name="pressure",
620 units="hPa",
621 )
623 # Create new cube with these dimensions.
624 out_cube = iris.cube.Cube(
625 data[:, np.newaxis, :, :],
626 dim_coords_and_dims=[
627 (time_coord, 0),
628 (pressure_coord, 1),
629 (lat_coord, 2),
630 (lon_coord, 3),
631 ],
632 )
634 else:
635 # Not a pressure level variable, so only 3 dimensions.
636 out_cube = iris.cube.Cube(
637 data,
638 dim_coords_and_dims=[
639 (time_coord, 0),
640 (lat_coord, 1),
641 (lon_coord, 2),
642 ],
643 )
645 # Rename cube using lookup dictionary, if a lookup exists.
646 meta = UGRID_VAR_LOOKUP.get(var_key)
648 if meta is not None:
649 if "standard_name" in meta:
650 out_cube.standard_name = meta["standard_name"]
652 if "long_name" in meta:
653 out_cube.long_name = meta["long_name"]
655 if "units" in meta:
656 out_cube.units = meta["units"]
658 # Also rename cube itself.
659 out_cube.rename(meta["long_name"])
660 else:
661 # Fallback: keep original name.
662 out_cube.rename(var_key)
664 # Add forecast reference time as 'time_origin' to mimic lfric where it will
665 # reconstruct forecast_period in a later callback.
666 # Extract the origin string from the units
667 time_origin = time_coord.units.origin
669 # Strip the "seconds since " part.
670 time_origin = time_origin.split("since ")[1]
672 # Add to cube attributes as str.
673 out_cube.coord("time").attributes["time_origin"] = time_origin
675 # Change units, geopot in m2 s-2.
676 if out_cube.long_name == "geopotential_height":
677 out_cube.data = out_cube.data / 9.81
679 # Raw data in units of 6h accum in meters.
680 if out_cube.long_name == "surface_microphysical_rainfall_rate":
681 out_cube.data = (out_cube.data * 1000.0) / 6
683 return out_cube
686def _restructure_ugrid_regrid(cube, tri, lat_grid, lon_grid, xy):
687 """
688 Restructure a flattened/unstructured cube.
690 Parameters
691 ----------
692 cube : iris.cube
693 An iris cube to restructure.
694 tri : scipy.spatial._qhull.Delaunay
695 A scipy object containing the triangulation mapping of cell points.
696 lat_grid : np.ndarray
697 1D latitude coordinate values of target grid.
698 lon_grid : np.ndarray
699 1D longitude coordinate values of target grid.
700 xy : np.ndarray
701 Meshed and flattened target grid points.
703 Returns
704 -------
705 iris.cube.Cube
706 A structured iris cube with appropriate metadata.
708 Notes
709 -----
710 This function uses a pre-calculated triangulation, to save rebuilding for
711 every cube. This therefore assumes all cubes being restructured have the
712 same flattened structure.
713 """
714 # Only process cubes that have 2 dimensions (i.e. time and space), not
715 # 1 dimension (to avoid trying to restructure latitude/longitude).
716 if cube.ndim > 1:
717 # Create empty numpy array to store regridded data.
718 out = np.empty((cube.shape[0], lat_grid.size, lon_grid.size))
720 logging.info("Interpolating", cube.name())
722 # Extract and transpose source data values.
723 src_vals = cube.data.T
725 # Build linear interpolator object mapping target triangulation to source values.
726 interp = LinearNDInterpolator(tri, src_vals)
728 # Interpolate values onto target grid using linear interpolation.
729 out_flat = interp(xy)
731 # Transpose, and reshape to target 2D regular lat/lon grid.
732 out = out_flat.T.reshape(cube.shape[0], lat_grid.size, lon_grid.size)
734 # Rebuild metadata using lookup table (mostly for anemoi ML models).
735 out_cube = _rebuild_ugrid_meta(out, cube, lat_grid, lon_grid)
737 # Reduce precision to reduce filesize.
738 if out_cube is not None:
739 out_cube.data = np.asarray(out_cube.data, dtype=np.float32)
741 # Return restructured cube with appropriate metadata
742 return out_cube
745def restructure_ugrid(cubes):
746 """
747 Restructured unstructured or flattened data onto a regular grid.
749 Parameters
750 ----------
751 cubes : iris.cube.CubeList
752 A cubelist containing unstructured cubes, along with cubes containing
753 latitude and longitude information.
755 Returns
756 -------
757 iris.cube.CubeList
758 A list of iris cubes, that have been restructured onto a regular grid,
759 with appropriate corrections to metadata.
760 """
761 # Setup folder containing data location to write lock/restructured data.
762 dataloc = os.environ["ROSE_DATAC"] + "/data/1/"
764 # If this directory doesn't exist, we are probably not running in a rose suite.
765 # Currently, command line use of restructuring is not supported, as it is not
766 # clear where a sensible location to store regridded data is.
767 if not os.path.isdir(dataloc):
768 raise NotImplementedError(
769 "Restructuring data outside a cylc workflow is not currently supported."
770 )
772 logging.info("Restructuring UGRID...")
774 # Define location of lock and done hidden file.
775 lock_path = os.path.join(dataloc, ".lock")
776 done_path = os.path.join(dataloc, ".done")
778 while True:
779 try:
780 # Try and create file lock in data directory. If not possible, it will
781 # raise FileExistsError and not do any data restructuring
782 fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
783 os.close(fd)
785 logging.info("Running preprocess restructure...")
787 # First, extract latitude and longitude coordinates
788 lat = cubes.extract("latitude")[0].data
789 lon = cubes.extract("longitude")[0].data
790 points = np.column_stack((lon, lat))
792 # Create output mesh, using standard grid ~2km resolution
793 # TODO: discussions with ML developers to include metadata so
794 # we don't have to a) guess lat/lon resolution and b) know if
795 # data is truly unstructured or just flattened (where np.reshape
796 # would be substantially faster and preserve original data).
797 lon_grid = np.arange(lon.data.min(), lon.data.max(), 0.02)
798 lat_grid = np.arange(lat.data.min(), lat.data.max(), 0.02)
799 Lon2d, Lat2d = np.meshgrid(lon_grid, lat_grid)
801 # Flatten target points
802 xy = np.column_stack((Lon2d.ravel(), Lat2d.ravel()))
804 # Build triangulation via a dummy interpolator
805 tri_interp = LinearNDInterpolator(points, np.zeros(points.shape[0]))
806 tri = tri_interp.tri
808 # Utilise multiprocessing to speed up code.
809 with Pool(processes=int(os.cpu_count() / 2)) as pool:
810 results = pool.starmap(
811 _restructure_ugrid_regrid,
812 [(cube, tri, lat_grid, lon_grid, xy) for cube in cubes],
813 )
815 # Filter results to only collect cubes, not None which is sometimes returned
816 # from _restructure_ugrid_regrid.
817 fixed_cubes = iris.cube.CubeList(c for c in results if c is not None)
819 # Save concatenated cubes to data location, for other processes to use.
820 iris.save(fixed_cubes.concatenate(), dataloc + "/restructured_cubes.nc")
822 # Write done file.
823 fd = os.open(done_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
824 os.close(fd)
826 # Return cubes as no more work to do for this process in read.
827 return fixed_cubes.concatenate()
829 # File lock exists, so wait for .done to appear.
830 except FileExistsError:
831 logging.info("Lock exists...")
833 # If the .done file exists, can proceed, and load restructured data.
834 if os.path.isfile(done_path):
835 logging.info("Done file found, proceeding")
837 return iris.load(dataloc + "/restructured_cubes.nc")
839 else:
840 # Otherwise, wait 60 seconds before trying again.
841 logging.info("Waiting 60 seconds...")
842 time.sleep(60)
845def vertical_interpolation(
846 cubes: iris.cube.Cube | iris.cube.CubeList,
847 coordinate: str,
848 target: iris.cube.Cube | iris.cube.CubeList,
849) -> iris.cube.Cube | iris.cube.CubeList:
850 """Vertical interpolation of a cube to match that off a different cube.
852 Acts as a wrapper around the `cube.interpolate` functionality and uses
853 linear interpolation as the method.
855 Parameters
856 ----------
857 cubes: iris.cube.Cube | iris.cube.CubeList
858 An iris cube or cubelist of data defining field that should be
859 vertically interpolated.
860 coordinate: str
861 The coordinate the interpolation occurs over.
862 target: iris.cube.Cube | iris.cube.CubeList
863 The target cube or cubelist that provides the vertical coordinate
864 information. It will use `cube.coord(coordinate).points` to provide
865 the vertical target. The number of target cubes should match the number
866 of cubes used as input.
868 Returns
869 -------
870 interpolated_cubes: iris.cube.Cube | iris.cube.CubeList
871 Coordinates of the selected point on the rotated grid specified within
872 the selected cube.
873 """
874 interpolated_cubes = iris.cube.CubeList([])
875 for cube, cube_t in zip(iter_maybe(cubes), iter_maybe(target), strict=True):
876 target_levels = cube_t.coord(coordinate).points
877 new_cube = cube.interpolate(
878 [(coordinate, target_levels)], iris.analysis.Linear()
879 )
880 interpolated_cubes.append(new_cube)
881 if len(interpolated_cubes) == 1:
882 return interpolated_cubes[0]
883 else:
884 return interpolated_cubes