extract_overlap

extract_overlap.py

import iris
import numpy as np
from shapely.geometry import Polygon


def horizontal_grid(cube):
    return cube.coord(axis='x'), cube.coord(axis='y')


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)

    if len(tgt_geom) > 1:
        msg = ('Currently bounding box transformation does not supported '
               'across the dateline')
        raise RuntimeError(msg)

    # (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 extract_overlap(source, target):
    """
    Return the source which overlaps the provided target.

    :param `~iris.cube.Cube` source: Source cube.
    :param `~iris.cube.Cube` target: Target cube used to define the area
        overlap with the source.
    :returns `~iris.cube.Cube`: An Iris cube representing the source which
        overlaps the area of the target.

    .. note::
        Currently limited to coordinates with associated modulus (geodetic
        coordinates).

    """
    # Handles cubes of different coordinate systems by determining a suitable
    # bounding box from the target points.
    target_x, target_y = horizontal_grid(target)
    xpnts = target_x.bounds if target_x.has_bounds() else target_x.points
    ypnts = target_y.bounds if target_y.has_bounds() else target_y.points
    source_x, source_y = horizontal_grid(source)

    # Make sure both x and y are described by the same coordinate system.
    if source_x.coord_system != source_y.coord_system:
        msg = 'Source coordinates system for the horizontal axes do not agree'
        raise ValueError(msg)
    if target_x.coord_system != target_y.coord_system:
        msg = 'Target coordinates system for the horizontal axes do not agree'
        raise ValueError(msg)

    # Use intersection if possible - its currently very limited
    # example it is has no support for coordinates with no modulus.
    tgt_bounds = ((xpnts.min(), xpnts.max()), (ypnts.min(), ypnts.max()))
    if source_x.coord_system != target_x.coord_system:
        tgt_bbox = ((tgt_bounds[0][0], tgt_bounds[1][0]),
                    (tgt_bounds[0][0], tgt_bounds[1][1]),
                    (tgt_bounds[0][1], tgt_bounds[1][1]),
                    (tgt_bounds[0][1], tgt_bounds[1][0]))
        (xmin, ymin, xmax, ymax) = transform_bbox(
            tgt_bbox, target_x.coord_system, source_x.coord_system)
    else:
        ((xmin, xmax), (ymin, ymax)) = tgt_bounds

    try:
        x_ext = iris.coords.CoordExtent(source_x, xmin, xmax)
        y_ext = iris.coords.CoordExtent(source_y, ymin, ymax)
        overlap = source.intersection(x_ext, y_ext)
    except ValueError as err:
        msg = 'coordinate units with no modulus are not yet supported'
        if msg in err.message:
            msg = ('{}.  Attempting to slice cube which is not wraparound '
                   'safe'.format(err.message))
            warnings.warn(msg)

            # super fast way by slicing - making assumptions like the
            # source covers at least the region of the target.
            xyes = (source_x.bounds >= xmin).sum(axis=1) > 0
            xyes *= (source_x.bounds <= xmax).sum(axis=1) > 0
            yyes = (source_y.bounds >= ymin).sum(axis=1) > 0
            yyes *= (source_y.bounds <= ymax).sum(axis=1) > 0
            xyes = np.where(xyes == True)[0]
            yyes = np.where(yyes == True)[0]
            overlap = source[yyes.min():yyes.max()+1, xyes.min():xyes.max()+1]

            # # using coord value constraint extraction
            # x_con = iris.Constraint(
            #     coord_values={source_x.name(): lambda cell:
            #                   float(xmin) <= cell <= float(xmax)})
            # y_con = iris.Constraint(
            #     coord_values={source_y.name(): lambda cell:
            #                   float(ymin) <= cell <= float(ymax)})
            # overlap = source.extract(y_con & x_con)
        else:
            raise err

    return overlap


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()

    cube3 = extract_overlap(cube, cube2)