Coverage for src/CSET/operators/regrid.py: 99%

1# © Crown copyright, Met Office (2022-2024) and CSET contributors.

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

7# http://www.apache.org/licenses/LICENSE-2.0

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 warnings

20import iris

21import iris.coord_systems

22import iris.cube

23import numpy as np

24from iris.analysis.cartography import rotate_pole

26from CSET._common import iter_maybe

27from CSET.operators._utils import get_cube_yxcoordname

29logger = logging.getLogger(__name__)

32class BoundaryWarning(UserWarning):

33 """Selected gridpoint is close to the domain edge.

35 In many cases gridpoints near the domain boundary contain non-physical

36 values, so caution is advised when interpreting them.

37 """

40def regrid_onto_cube(

41 toregrid: iris.cube.Cube | iris.cube.CubeList,

42 target: iris.cube.Cube,

43 method: str,

44 **kwargs,

45) -> iris.cube.Cube | iris.cube.CubeList:

46 """Regrid a cube or CubeList, projecting onto a target cube.

48 All cubes must have at least 2 spatial (map) dimensions.

50 Arguments

51 ----------

52 toregrid: iris.cube | iris.cube.CubeList

53 An iris Cube of data to regrid, or multiple cubes to regrid in a

54 CubeList. A minimum requirement is that the cube(s) need to be 2D with a

55 latitude, longitude coordinates.

56 target: Cube

57 An iris cube of the data to regrid onto. It needs to be 2D with a

58 latitude, longitude coordinate.

59 method: str

60 Method used to regrid onto, etc. Linear will use iris.analysis.Linear()

62 Returns

63 -------

64 iris.cube | iris.cube.CubeList

65 An iris cube of the data that has been regridded, or a CubeList of the

66 cubes that have been regridded in the same order they were passed in

67 toregrid.

69 Raises

70 ------

71 ValueError

72 If a unique x/y coordinate cannot be found

73 NotImplementedError

74 If the cubes grid, or the method for regridding, is not yet supported.

76 Notes

77 -----

78 Currently rectlinear grids (uniform) are supported.

79 """

80 # To store regridded cubes.

81 regridded_cubes = iris.cube.CubeList()

83 # Iterate over all cubes and regrid.

84 for cube in iter_maybe(toregrid):

85 # Get y,x coord names

86 y_coord, x_coord = get_cube_yxcoordname(cube)

88 # List of supported grids - check if it is compatible

89 supported_grids = (iris.coord_systems.GeogCS, iris.coord_systems.RotatedGeogCS)

90 if not isinstance(cube.coord(x_coord).coord_system, supported_grids):

91 raise NotImplementedError(

92 f"Does not currently support {cube.coord(x_coord).coord_system} coordinate system"

93 )

94 if not isinstance(cube.coord(y_coord).coord_system, supported_grids):

95 raise NotImplementedError(

96 f"Does not currently support {cube.coord(y_coord).coord_system} coordinate system"

97 )

99 regrid_method = getattr(iris.analysis, method, None)

100 if callable(regrid_method):

101 regridded_cubes.append(cube.regrid(target, regrid_method()))

102 else:

103 raise NotImplementedError(

104 f"Does not currently support {method} regrid method"

105 )

106

107 # Preserve returning a cube if only a cube has been supplied to regrid.

108 if len(regridded_cubes) == 1:

109 return regridded_cubes[0]

110 else:

111 return regridded_cubes

112

113

114def regrid_onto_xyspacing(

115 toregrid: iris.cube.Cube | iris.cube.CubeList,

116 xspacing: float,

117 yspacing: float,

118 method: str,

119 **kwargs,

120) -> iris.cube.Cube | iris.cube.CubeList:

121 """Regrid cube or cubelist onto a set x,y spacing.

122

123 Regrid cube(s) using specified x,y spacing, which is performed linearly.

124

125 Parameters

126 ----------

127 toregrid: iris.cube | iris.cube.CubeList

128 An iris cube of the data to regrid, or multiple cubes to regrid in a

129 cubelist. A minimum requirement is that the cube(s) need to be 2D with a

130 latitude, longitude coordinates.

131 xspacing: float

132 Spacing of points in longitude direction (could be degrees, meters etc.)

133 yspacing: float

134 Spacing of points in latitude direction (could be degrees, meters etc.)

135 method: str

136 Method used to regrid onto, etc. Linear will use iris.analysis.Linear()

137

138 Returns

139 -------

140 iris.cube | iris.cube.CubeList

141 An iris cube of the data that has been regridded, or a cubelist of the

142 cubes that have been regridded in the same order they were passed in

143 toregrid.

144

145 Raises

146 ------

147 ValueError

148 If a unique x/y coordinate cannot be found

149 NotImplementedError

150 If the cubes grid, or the method for regridding, is not yet supported.

151

152 Notes

153 -----

154 Currently rectlinear grids (uniform) are supported.

155

156 """

157 # To store regridded cubes.

158 regridded_cubes = iris.cube.CubeList()

159

160 # Iterate over all cubes and regrid.

161 for cube in iter_maybe(toregrid):

162 # Get x,y coord names

163 y_coord, x_coord = get_cube_yxcoordname(cube)

164

165 # List of supported grids - check if it is compatible

166 supported_grids = (iris.coord_systems.GeogCS, iris.coord_systems.RotatedGeogCS)

167 if not isinstance(cube.coord(x_coord).coord_system, supported_grids):

168 raise NotImplementedError(

169 f"Does not currently support {cube.coord(x_coord).coord_system} regrid method"

170 )

171 if not isinstance(cube.coord(y_coord).coord_system, supported_grids):

172 raise NotImplementedError(

173 f"Does not currently support {cube.coord(y_coord).coord_system} regrid method"

174 )

175

176 # Get axis

177 lat, lon = cube.coord(y_coord), cube.coord(x_coord)

178

179 # Get bounds

180 lat_min, lon_min = lat.points.min(), lon.points.min()

181 lat_max, lon_max = lat.points.max(), lon.points.max()

182

183 # Generate new mesh

184 latout = np.arange(lat_min, lat_max, yspacing)

185 lonout = np.arange(lon_min, lon_max, xspacing)

186

187 regrid_method = getattr(iris.analysis, method, None)

188 if callable(regrid_method):

189 regridded_cubes.append(

190 cube.interpolate(

191 [(y_coord, latout), (x_coord, lonout)], regrid_method()

192 )

193 )

194 else:

195 raise NotImplementedError(

196 f"Does not currently support {method} regrid method"

197 )

198

199 # Preserve returning a cube if only a cube has been supplied to regrid.

200 if len(regridded_cubes) == 1:

201 return regridded_cubes[0]

202 else:

203 return regridded_cubes

204

205

206def regrid_to_single_point(

207 cubes: iris.cube.Cube | iris.cube.CubeList,

208 lat_pt: float,

209 lon_pt: float,

210 latlon_in_type: str = "rotated",

211 method: str = "Nearest",

212 boundary_margin: int = 8,

213 **kwargs,

214) -> iris.cube.Cube:

215 """Select data at a single point by longitude and latitude.

216

217 Selection of model grid point is performed by a regrid function, either

218 selecting the nearest gridpoint to the selected longitude and latitude

219 values or using linear interpolation across the surrounding points.

220

221 Parameters

222 ----------

223 cubes: Cube | CubeList

224 An iris cube or CubeList of the data to regrid. As a minimum, it needs

225 to be 2D with latitude, longitude coordinates.

226 lon_pt: float

227 Selected value of longitude: this should be in the range -180 degrees to

228 180 degrees.

229 lat_pt: float

230 Selected value of latitude: this should be in the range -90 degrees to

231 90 degrees.

232 latlon_in_type: str, optional

233 Specify whether the input longitude and latitude point is in standard

234 geographic realworld coordinates ("realworld") or on the rotated grid

235 of the cube ("rotated"). Default is "rotated".

236 method: str

237 Method used to determine the values at the selected longitude and

238 latitude. The recommended approach is to use iris.analysis.Nearest(),

239 which selects the nearest gridpoint. An alternative is

240 iris.analysis.Linear(), which obtains the values at the selected

241 longitude and latitude by linear interpolation.

242 boundary_margin: int, optional

243 Number of grid points from the domain boundary considered "unreliable".

244 Defaults to 8.

245

246 Returns

247 -------

248 regridded_cubes: Cube | CubeList

249 An iris cube or CubeList of the data at the specified point (this may

250 have time and/or height dimensions).

251

252 Raises

253 ------

254 ValueError

255 If a unique x/y coordinate cannot be found; also if, for selecting a

256 single gridpoint, the chosen longitude and latitude point is outside the

257 domain; also (currently) if the difference between the actual and target

258 points exceed 0.1 degrees.

259 NotImplementedError

260 If the cubes grid, or the method for regridding, is not yet supported.

261

262 Notes

263 -----

264 The acceptable coordinate names for X and Y coordinates are currently

265 described in X_COORD_NAMES and Y_COORD_NAMES. These cover commonly used

266 coordinate types, though a user can append new ones. Currently rectilinear

267 grids (uniform) are supported. Warnings are raised if the selected gridpoint

268 is within boundary_margin grid lengths of the domain boundary as data here

269 is potentially unreliable.

270 """

271 # To store regridded cubes.

272 regridded_cubes = iris.cube.CubeList()

273

274 # Iterate over all cubes and regrid.

275 for cube in iter_maybe(cubes):

276 # Get x and y coordinate names.

277 y_coord, x_coord = get_cube_yxcoordname(cube)

278

279 # List of supported grids - check if it is compatible

280 # NOTE: The "RotatedGeogCS" option below seems to be required for rotated grids --

281 # this may need to be added in other places in these Operators.

282 supported_grids = (iris.coord_systems.GeogCS, iris.coord_systems.RotatedGeogCS)

283 if not isinstance(cube.coord(x_coord).coord_system, supported_grids):

284 raise NotImplementedError(

285 f"Does not currently support {cube.coord(x_coord).coord_system} regrid method"

286 )

287 if not isinstance(cube.coord(y_coord).coord_system, supported_grids):

288 raise NotImplementedError(

289 f"Does not currently support {cube.coord(y_coord).coord_system} regrid method"

290 )

291

292 # Transform input coordinates onto rotated grid if requested

293 if latlon_in_type == "realworld":

294 lon_tr, lat_tr = transform_lat_long_points(lon_pt, lat_pt, cube)

295 elif latlon_in_type == "rotated": 295 ↛ 299line 295 didn't jump to line 299 because the condition on line 295 was always true

296 lon_tr, lat_tr = lon_pt, lat_pt

297

298 # Get axis

299 lat, lon = cube.coord(y_coord), cube.coord(x_coord)

300

301 # Get bounds

302 lat_min, lon_min = lat.points.min(), lon.points.min()

303 lat_max, lon_max = lat.points.max(), lon.points.max()

304

305 # Use different logic for single point obs data.

306 if len(cube.coord(x_coord).points) > 1:

307 # Get boundaries of frame to avoid selecting gridpoint close to domain edge

308 lat_min_bound, lon_min_bound = (

309 lat.points[boundary_margin - 1],

310 lon.points[boundary_margin - 1],

311 )

312 lat_max_bound, lon_max_bound = (

313 lat.points[-boundary_margin],

314 lon.points[-boundary_margin],

315 )

316

317 # Check to see if selected point is outside the domain

318 if (lat_tr < lat_min) or (lat_tr > lat_max):

319 raise ValueError("Selected point is outside the domain.")

320 else:

321 if (lon_tr < lon_min) or (lon_tr > lon_max):

322 if (lon_tr + 360.0 >= lon_min) and (lon_tr + 360.0 <= lon_max):

323 lon_tr += 360.0

324 elif (lon_tr - 360.0 >= lon_min) and (lon_tr - 360.0 <= lon_max):

325 lon_tr -= 360.0

326 else:

327 raise ValueError("Selected point is outside the domain.")

328

329 # Check to see if selected point is near the domain boundaries

330 if (

331 (lat_tr < lat_min_bound)

332 or (lat_tr > lat_max_bound)

333 or (lon_tr < lon_min_bound)

334 or (lon_tr > lon_max_bound)

335 ):

336 warnings.warn(

337 f"Selected point is within {boundary_margin} gridlengths of the domain edge, data may be unreliable.",

338 category=BoundaryWarning,

339 stacklevel=2,

340 )

341

342 regrid_method = getattr(iris.analysis, method, None)

343 if not callable(regrid_method):

344 raise NotImplementedError(

345 f"Does not currently support {method} regrid method"

346 )

347

348 cube_rgd = cube.interpolate(((lat, lat_tr), (lon, lon_tr)), regrid_method())

349 regridded_cubes.append(cube_rgd)

350 else:

351 if (

352 np.abs(lat_tr - lat.points[0]) > 0.1

353 or np.abs(lon_tr - lon.points[0]) > 0.1

354 ):

355 raise ValueError(

356 "Selected point is too far from the specified coordinates. It should be within 0.1 degrees."

357 )

358 else:

359 print(

360 "*** lat/long diffs",

361 np.abs(lat_tr - lat_pt),

362 np.abs(lon_tr - lon_pt),

363 )

364 regridded_cubes.append(cube)

365

366 # Preserve returning a cube if only a cube has been supplied to regrid.

367 if len(regridded_cubes) == 1:

368 return regridded_cubes[0]

369 else:

370 return regridded_cubes

371

372

373def transform_lat_long_points(lon, lat, cube):

374 """Transform a selected point in longitude and latitude.

375

376 Transform the coordinates of a point from the real world

377 grid to the corresponding point on the rotated grid of a cube.

378

379 Parameters

380 ----------

381 cube: Cube

382 An iris cube of data defining the rotated grid to be used in

383 the longitude-latitude transformation.

384 lon: float

385 Selected value of longitude: this should be in the range -180 degrees to

386 180 degrees.

387 lat: float

388 Selected value of latitude: this should be in the range -90 degrees to

389 90 degrees.

390

391 Returns

392 -------

393 lon_rot, lat_rot: float

394 Coordinates of the selected point on the rotated grid specified within

395 the selected cube.

396

397 """

398 import cartopy.crs as ccrs

399

400 rot_pole = cube.coord_system().as_cartopy_crs()

401 true_grid = ccrs.Geodetic()

402 rot_coords = rot_pole.transform_point(lon, lat, true_grid)

403 lon_rot = rot_coords[0]

404 lat_rot = rot_coords[1]

405

406 return lon_rot, lat_rot

407

408

409def interpolate_to_point_cube(

410 fld: iris.cube.Cube | iris.cube.CubeList,

411 point_cube: iris.cube.Cube | iris.cube.CubeList,

412 **kwargs,

413) -> iris.cube.Cube | iris.cube.CubeList:

414 """Interpolate a 2D field in cube or CubeList to a set of points.

415

416 Regrid cube(s) to set of sample points specified by point_cube.

417 Ensures only matching times between input and point_cube are included.

418

419 Parameters

420 ----------

421 fld: Cube | CubeList

422 An iris cube or CubeList containing a two-dimensional field(s).

423 point_cube: Cube

424 An iris cube specifying the coord point(s) to which the data

425 will be interpolated.

426

427 Returns

428 -------

429 fld_point_cube: Cube | CubeList

430 An iris cube or CubeList containing interpolated values at the

431 points specified by the point cube for matching times common to

432 both fld and point_cube.

433 """

434 # Empty CubeList To store regridded cubes.

435 regridded_cubes = iris.cube.CubeList()

436

437 # Generate array of point cube lat and lon points.

438 point_lat_name, point_lon_name = get_cube_yxcoordname(point_cube)

439 point_lats = point_cube.coord(point_lat_name).points

440 point_lons = point_cube.coord(point_lon_name).points

441

442 # Iterate over all cubes and regrid.

443 for cube in iter_maybe(fld):

444 # Get forecast field cube spatial names.

445 y_coord, x_coord = get_cube_yxcoordname(cube)

446

447 # Rotate point_cube coords if required to match model coord rotation.

448 if (

449 isinstance(

450 cube.coord(x_coord).coord_system, iris.coord_systems.RotatedGeogCS

451 )

452 and point_lat_name == "latitude"

453 and point_lon_name == "longitude"

454 ):

455 point_lons_rp, point_lats_rp = rotate_pole(

456 point_lons,

457 point_lats,

458 pole_lon=cube.coord(x_coord).coord_system.grid_north_pole_longitude,

459 pole_lat=cube.coord(x_coord).coord_system.grid_north_pole_latitude,

460 )

461 sample_points = [

462 ("grid_latitude", point_lats_rp),

463 ("grid_longitude", point_lons_rp),

464 ]

465 # Default, sample points based on point cube dimensions

466 else:

467 sample_points = [

468 (point_lat_name, np.array(point_lats)),

469 (point_lon_name, np.array(point_lons)),

470 ]

471

472 # Interpolate fld cube to required sample points

473 fld_point_cube = cube.interpolate(

474 sample_points, iris.analysis.Linear(extrapolation_mode="mask")

475 )

476

477 diag_data = np.diagonal(

478 fld_point_cube.data,

479 axis1=-2,

480 axis2=-1,

481 )

482

483 fv_cube = iris.cube.Cube(

484 diag_data,

485 standard_name=cube.standard_name,

486 long_name=cube.long_name,

487 units=cube.units,

488 )

489

490 #

491 # Add all non-horizontal dimension coordinates

492 # from the forecast cube.

493 #

494 out_dim = 0

495

496 for coord in cube.coords(dim_coords=True):

497 if coord.name() in (

498 "latitude",

499 "longitude",

500 "grid_latitude",

501 "grid_longitude",

502 ):

503 continue

504

505 fv_cube.add_dim_coord(

506 coord.copy(),

507 out_dim,

508 )

509

510 out_dim += 1

511

512 #

513 # Station dimension is always last.

514 #

515 station_dim = out_dim

516

517 fv_cube.add_dim_coord(

518 point_cube.coord("station").copy(),

519 station_dim,

520 )

521

522 #

523 # Copy forecast auxiliary coordinates.

524 #

525 for coord in cube.coords():

526 if coord.name() in (

527 "latitude",

528 "longitude",

529 "grid_latitude",

530 "grid_longitude",

531 ):

532 continue

533

534 if coord.name() in [c.name() for c in fv_cube.coords()]:

535 continue

536

537 dims = cube.coord_dims(coord)

538

539 if dims:

540 fv_cube.add_aux_coord(

541 coord.copy(),

542 dims,

543 )

544 else:

545 fv_cube.add_aux_coord(

546 coord.copy(),

547 )

548

549 #

550 # Copy observation auxiliary coordinates.

551 #

552 for coord in point_cube.coords():

553 if coord.name() in (

554 "station",

555 "forecast_period",

556 "forecast_reference_time",

557 ):

558 continue

559

560 if coord.name() in [c.name() for c in fv_cube.coords()]:

561 continue

562

563 dims = point_cube.coord_dims(coord)

564

565 if dims: 565 ↛ 571line 565 didn't jump to line 571 because the condition on line 565 was always true

566 fv_cube.add_aux_coord(

567 coord.copy(),

568 dims,

569 )

570 else:

571 fv_cube.add_aux_coord(

572 coord.copy(),

573 )

574

575 fv_cube.attributes = cube.attributes.copy()

576 fv_cube.cell_methods = cube.cell_methods

577 fv_cube.units = cube.units

578

579 regridded_cubes.append(fv_cube)

580

581 # Preserve returning a cube if only a cube supplied.

582 if len(regridded_cubes) == 1:

583 return regridded_cubes[0]

584

585 return regridded_cubes

586

587 # # Ensure matching times in fld cube and point_cube

588 # base_time_coord = point_cube.coord("time")

589 # other_time_coord = cube.coord("time")

590 # base_times = base_time_coord.units.num2date(base_time_coord.points)

591 # other_times = other_time_coord.units.num2date(other_time_coord.points)

592 # shared_times = set.intersection(set(base_times), set(other_times))

593 # logger.debug("Shared times: %s", shared_times)

594 # time_constraint = iris.Constraint(

595 # coord_values={

596 # "time": lambda cell, shared_times=shared_times: (

597 # cell.point in shared_times

598 # )

599 # }

600 # )

601

602 # # Extract points matching the shared times.

603 # cube = cube.extract(time_constraint)

604 # point_cube = point_cube.extract(time_constraint)

605 # if cube is None or point_cube is None:

606 # raise ValueError("No common time points found!")

607

608 # # Retain only diagonal elements of 2D interpolated cube to vector points

609 # od_index = point_cube.coord_dims("station")[0]

610 # fv_cube = iris.cube.Cube(

611 # fld_point_cube.data.diagonal(axis1=od_index, axis2=od_index + 1),

612 # standard_name=cube.standard_name,

613 # long_name=cube.long_name,

614 # units=cube.units,

615 # )

616 # # Copy all non-lat/lon coordinates and cube attributes

617 # if "time" in [coord.name() for coord in fld_point_cube.coords(dim_coords=True)]:

618 # fv_cube.add_dim_coord(

619 # point_cube.coord("time"), point_cube.coord_dims("time")[0]

620 # )

621 # for coord in cube.coords():

622 # if coord.name() not in [

623 # "time",

624 # "latitude",

625 # "longitude",

626 # "grid_latitude",

627 # "grid_longitude",

628 # ] and coord.name() not in [coord.name() for coord in fv_cube.coords()]:

629 # fv_cube.add_aux_coord(coord.copy(), cube.coord_dims(coord))

630 # for coord in point_cube.coords():

631 # if coord.name() not in [

632 # "time",

633 # "forecast_period",

634 # "forecast_reference_time",

635 # "realization",

636 # "station",

637 # ] and coord.name() not in [coord.name() for coord in fv_cube.coords()]:

638 # fv_cube.add_aux_coord(coord.copy(), point_cube.coord_dims(coord))

639 # fv_cube.add_dim_coord(point_cube.coord("station"), od_index)

640 # fv_cube.attributes = cube.attributes.copy()

641 # fv_cube.cell_methods = cube.cell_methods

642 # fv_cube.units = cube.units

643 # regridded_cubes.append(fv_cube)

644

645 # # Preserve returning a cube if only a cube has been supplied to regrid.

646 # if len(regridded_cubes) == 1:

647 # return regridded_cubes[0]

648 # else:

649 # return regridded_cubes

650

651

652def vertical_interpolation(

653 cubes: iris.cube.Cube | iris.cube.CubeList,

654 coordinate: str,

655 target: iris.cube.Cube | iris.cube.CubeList,

656) -> iris.cube.Cube | iris.cube.CubeList:

657 """Vertical interpolation of a cube to match that off a different cube.

658

659 Acts as a wrapper around the `cube.interpolate` functionality and uses

660 linear interpolation as the method.

661

662 Parameters

663 ----------

664 cubes: iris.cube.Cube | iris.cube.CubeList

665 An iris cube or cubelist of data defining field that should be

666 vertically interpolated.

667 coordinate: str

668 The coordinate the interpolation occurs over.

669 target: iris.cube.Cube | iris.cube.CubeList

670 The target cube or cubelist that provides the vertical coordinate

671 information. It will use `cube.coord(coordinate).points` to provide

672 the vertical target. The number of target cubes should match the number

673 of cubes used as input.

674

675 Returns

676 -------

677 interpolated_cubes: iris.cube.Cube | iris.cube.CubeList

678 Coordinates of the selected point on the rotated grid specified within

679 the selected cube.

680 """

681 interpolated_cubes = iris.cube.CubeList([])

682 for cube, cube_t in zip(iter_maybe(cubes), iter_maybe(target), strict=True):

683 target_levels = cube_t.coord(coordinate).points

684 new_cube = cube.interpolate(

685 [(coordinate, target_levels)], iris.analysis.Linear()

686 )

687 interpolated_cubes.append(new_cube)

688 if len(interpolated_cubes) == 1:

689 return interpolated_cubes[0]

690 else:

691 return interpolated_cubes