Source code for aeolus.region

"""Subsetting variables over geographical regions."""
from dataclasses import dataclass, field

import iris
from iris.analysis.cartography import wrap_lons

from .coord_utils import UM_LATLON
from .exceptions import BoundaryError
from .util import fmt_lonlat


__all__ = ("Region",)


@dataclass
class BoundsRect:
    """Bounding longitudes and latitudes of a given lon-lat rectangle."""

    west: float = field(metadata={"coord": UM_LATLON[1]})
    east: float = field(metadata={"coord": UM_LATLON[1]})
    south: float = field(metadata={"coord": UM_LATLON[0]})
    north: float = field(metadata={"coord": UM_LATLON[0]})

    def __post_init__(self):  # noqa
        # if self.west > self.east:
        #     raise BoundaryError("West boundary value should be less than east")
        if self.south > self.north:
            raise BoundaryError("South boundary value should be less than north")

    def __repr__(self):  # noqa
        return (
            f"BoundsRect(west={self.west}, east={self.east}, south={self.south}, "
            f"north={self.north})"
        )


[docs]class Region(object): """ Rectangular geographical region. Attributes ---------- name : str The region's name description : str A description of the region constraint: iris.Constraint A constraint object associated with the region """
[docs] def __init__(self, west_bound, east_bound, south_bound, north_bound, name="", description=""): """ Instantiate a `Region` object. Parameters ---------- name: str The region's name. description : str, optional A description of the region. west_bound, east_bound, south_bound, north_bound : scalar, optional The western, eastern, southern, and northern boundaries, respectively, of the region. """ self.name = name self.description = description self.bounds = BoundsRect(west_bound, east_bound, south_bound, north_bound) self._sides = [ (key, f.metadata["coord"]) for key, f in self.bounds.__dataclass_fields__.items() ] self.lon_size = abs(self.bounds.east - self.bounds.west) self.lat_size = self.bounds.north - self.bounds.south
def __repr__(self): # noqa txt = ( f"Geographical region '{self.name}' (west={self.bounds.west}, " f"east={self.bounds.east}, south={self.bounds.south}, north={self.bounds.north})" ) if self.description: txt += "\n\n" txt += self.description return txt def __getitem__(self, index): # noqa return { "value": getattr(self.bounds, self._sides[index][0]), "name": self._sides[index][0], "coord": self._sides[index][1], } def _perpendicular_side_limits(self, side): """Get minimum and maximum values of the region boundary perpendicular to the given one.""" if side in ["west", "east"]: coord_name = UM_LATLON[0] _min, _max = self.bounds.south, self.bounds.north elif side in ["south", "north"]: coord_name = UM_LATLON[1] _min, _max = self.bounds.west, self.bounds.east else: raise BoundaryError(f"Boundary name '{side}' is not valid") return coord_name, (_min, _max) def to_str(self, sep="_"): # noqa return sep.join([fmt_lonlat(i["value"], i["coord"]) for i in self])
[docs] @classmethod def from_cube(cls, cube, name=None, margin=None, margin_units="points", shift_lons=False): """ Create a Region from limits of longitude and latitude of the cube. Parameters ---------- cube: iris.cube.Cube Source cube. name: str, optional Name for the region. If not given, created automatically from `cube`'s name. margin: scalar, optional Use `margin` number of points or degrees to create a region smaller than the cube. margin_units: str, optional Units of margin. Can be "points" or "degrees". shift_lons: bool, optional Shift longitudes to -180...180. Returns ------- aeolus.region.Region """ if name is None: name = f"extent_of_{cube.name()}" lons = cube.coord(UM_LATLON[1]).points if shift_lons: lons = sorted(wrap_lons(lons, -180, 360)) lats = cube.coord(UM_LATLON[0]).points idx0, idx1 = 0, -1 if margin is not None: if margin_units == "points": idx0 += margin idx1 -= margin lon0 = lons[idx0] lon1 = lons[idx1] lat0 = lats[idx0] lat1 = lats[idx1] else: lon0 = lons[idx0] + margin lon1 = lons[idx1] - margin lat0 = lats[idx0] + margin lat1 = lats[idx1] - margin else: lon0 = lons[idx0] lon1 = lons[idx1] lat0 = lats[idx0] lat1 = lats[idx1] return cls(lon0, lon1, lat0, lat1, name=name)
@property def constraint(self): """Constraint to select data within the region.""" cnstr = iris.Constraint(latitude=lambda x: self.bounds.south <= x <= self.bounds.north) if self.bounds.west < self.bounds.east: # Western boundary is to the west cnstr &= iris.Constraint(longitude=lambda x: self.bounds.west <= x <= self.bounds.east) else: # Region wrapping around dateline (180deg) cnstr &= iris.Constraint( longitude=lambda x: (self.bounds.west <= x) or (x <= self.bounds.east) ) return cnstr
[docs] def add_to_ax(self, ax, **kwargs): """Add a Rectangle patch to matplotlib axes `ax` with given keyword arguments `kwargs`.""" from matplotlib.patches import Rectangle # noqa xy = (self.bounds.west, self.bounds.south) width = self.lon_size height = self.lat_size ax.add_patch(Rectangle(xy, width, height, **kwargs))