2-stage-regrid - different coord system regrid

2stage_regrid.py

import warnings

import iris
import biggus
from iris.analysis._area_weighted import AreaWeightedRegridder \
    as _AreaWeightedRegridder
from iris.experimental.regrid import regrid_weighted_curvilinear_to_rectilinear
import iris.experimental.regrid as eregrid
from iris.analysis._regrid import RectilinearRegridder
import numpy as np
from shapely.geometry import Polygon


def transform_bbox(points, src_crs, tgt_crs):
    """
    >>> from iris.coord_systems import GeogCS, OSGB
    >>> bbox = [(0, 45), (10, 45), (10, 55), (0, 55)]
    >>> print transform_bbox(bbox, GeogCS(6371229), OSGB())
    (527915.8213500988, 0.0, 700000.0, 577349.2303676724)
    """
    cartopy_src_crs = src_crs.as_cartopy_projection()
    cartopy_tgt_crs = tgt_crs.as_cartopy_projection()

    src_geom = Polygon(points)
    tgt_geom = cartopy_tgt_crs.project_geometry(src_geom, cartopy_src_crs)

    # (minx, miny, maxx, maxy)
    result = tgt_geom.bounds
    if not result:
        msg = 'Unable to succesfully project geometry from {} to {}'
        raise RuntimeError(msg.format(src_crs, tgt_crs))
    return result


def gen_regular_target(cube, target_grid, shape=None):
    # Generate a grid cube which covers the same area as the source cube only
    # that it has the coordinate system of the target grid and specified shape.
    src_x_coord, src_y_coord = (cube.coord(axis='x', dim_coords=True),
                                cube.coord(axis='y', dim_coords=True))
    dim_x = cube.coord_dims(src_x_coord)[0]
    dim_y = cube.coord_dims(src_y_coord)[0]
    dim_y, dim_x = np.argsort([dim_y, dim_x])

    if shape is None:
        shape = cube.shape

    dtype = target_grid.lazy_data().dtype
    data = biggus.ConstantArray(shape, dtype=dtype)
    grid_cube = iris.cube.Cube(data)

    src_x_points = src_x_coord.points
    if src_x_coord.has_bounds():
        src_x_points = src_x_coord.bounds
    src_y_points = src_y_coord.points
    if src_y_coord.has_bounds():
        src_y_points = src_y_coord.bounds

    # Determine bbox around source data.
    src_bounds = ((src_x_points.min(), src_x_points.max()),
                  (src_y_points.min(), src_y_points.max()))
    src_bbox = ((src_bounds[0][0], src_bounds[1][0]),
                (src_bounds[0][0], src_bounds[1][1]),
                (src_bounds[0][1], src_bounds[1][1]),
                (src_bounds[0][1], src_bounds[1][0]))
    src_crs = src_x_coord.coord_system
    tgt_crs = target_grid.coord(axis='x').coord_system
    xmin, ymin, xmax, ymax = transform_bbox(src_bbox, src_crs, tgt_crs)

    # Generate points by utilising bounds if available.
    if src_x_coord.has_bounds():
        x_bound = np.linspace(xmin, xmax, endpoint=True,
                              num=shape[dim_x]+1)
        x_bounds = np.array([x_bound[:-1], x_bound[1:]]).T
        x_points = x_bounds.mean(axis=1)
    else:
        x_points = np.linspace(xmin, xmax, num=shape[dim_x])

    if src_y_coord.has_bounds():
        y_bound = np.linspace(ymin, ymax, endpoint=True,
                              num=shape[dim_y]+1)
        y_bounds = np.array([y_bound[:-1], y_bound[1:]]).T
        y_points = y_bounds.mean(axis=1)
    else:
        y_points = np.linspace(ymin, ymax, num=shape[dim_y])

    # Add these dimensions to our grid cube
    grid_x_coord = target_grid.coord(axis='x').copy(x_points)
    grid_x_coord.bounds = x_bounds
    grid_y_coord = target_grid.coord(axis='y').copy(y_points)
    grid_y_coord.bounds = y_bounds
    grid_cube.add_dim_coord(grid_y_coord, dim_y)
    grid_cube.add_dim_coord(grid_x_coord, dim_x)

    return grid_cube


class _TwoStageRegrid(_AreaWeightedRegridder):
    def __call__(self, cube):
        try:
            result = super(_TwoStageRegrid, self).__call__(cube)
        except ValueError as err:
            message = ('The horizontal grid coordinates of both the source '
                       'and grid cubes must have the same coordinate '
                       'system.')
            if message in err.message:
                warnings.warn('Single-stage area weighted regrid not yet '
                              'possible: ({}).  Two-stage regrid being '
                              'performed, taking a bilinear interpolation '
                              'to a similar resolution intermediate grid, '
                              'then performing an area weighted '
                              'regrid.'.format(err.message))
                intermediate_target = gen_regular_target(
                    cube, self._target_grid_cube)

                int_tar = cube.regrid(intermediate_target,
                                      iris.analysis.Linear())

                result = eregrid.regrid_area_weighted_rectilinear_src_and_grid(
                    int_tar, self._target_grid_cube, mdtol=self._mdtol)
            else:
                raise err
        return result


class AreaWeighted(iris.analysis.AreaWeighted):
    def regridder(self, src_grid_cube, target_grid_cube):
        return _TwoStageRegrid(src_grid_cube, target_grid_cube,
                               mdtol=self.mdtol)


if __name__ == '__main__':
    import iris
    import iris.tests.stock as stock

    cube = stock.lat_lon_cube()
    cube.coord(axis='x').guess_bounds()
    cube.coord(axis='y').guess_bounds()

    cube2 = cube.copy()
    cube2.coord(axis='x').coord_system = stock.RotatedGeogCS(30, 30)
    cube2.coord(axis='y').coord_system = stock.RotatedGeogCS(30, 30)

    cube.regrid(cube2, AreaWeighted())