Geometry functions at the surface of the Earth

geometry.py

import xarray as xr
import numpy as np
import pandas as pd

EARTH_RADIUS = 6371 * 1e3


def latlon2yx(lat, lon):
    """
    Convert latitude and longitude arrays to y and x arrays in m

    Parameters
    ----------
    lat : array-like
        Latitudinal spherical coordinates
    lon : array-like
        Longitudinal spherical coordinates

    Returns
    -------
    y : array-like
        Zonal cartesian coordinates
    x : array-like
        Meridional cartesian coordinates
    """
    y = np.pi / 180. * EARTH_RADIUS * lat
    x = np.cos(np.pi / 180. * lat) * np.pi / 180. * EARTH_RADIUS * lon
    return y, x


def latlon2dydx(lat, lon, dim, label='upper'):
    """
    Convert latitude and longitude arrays to elementary displacements in dy
    and dx

    Parameters
    ----------
    lat : array-like
        Latitudinal spherical coordinates
    lon : array-like
        Longitudinal spherical coordinates
    dim : str
        Dimension along which the differentiation is performed, generally
        associated with the time dimension.
    label : {'upper', 'lower'}, optional
        The new coordinate in dimension dim will have the values of
        either the minuend’s or subtrahend’s coordinate for values ‘upper’
        and ‘lower’, respectively.

    Returns
    -------
    dy : array-like
        Zonal elementary displacement in cartesian coordinates
    dx : array-like
        Meridional elementary displacement in cartesian coordinates
    """
    dlat = lat.diff(dim, label=label)
    dlon = lon.diff(dim, label=label)
    dy = np.pi / 180. * EARTH_RADIUS * dlat
    # Need to slice the latitude data when there are duplicate values
    if label is 'upper':
        dx = (np.cos(np.pi / 180. * lat.isel(**{dim: slice(1, None)})) *
              np.pi / 180. * EARTH_RADIUS * dlon)
    elif label is 'lower':
        dx = (np.cos(np.pi / 180. * lat.isel(**{dim: slice(None, -1)})) *
              np.pi / 180. * EARTH_RADIUS * dlon)
    return dy, dx

def latlon2heading(lat, lon, dim, label='upper'):
    """
    Calculates the bearing between two points.
    The formulae used is the following:
        θ = atan2(sin(Δlong).cos(lat2),
                  cos(lat1).sin(lat2) − sin(lat1).cos(lat2).cos(Δlong))
    Parameters
    ----------

    Returns
    -------
      The bearing in degrees

    """
    dy, dx = latlon2dydx(lat, lon, dim, label=label)
    initial_heading = np.arctan2(dx, dy) * 180. / np.pi
    # Normalize the initial heading
    compass_heading = (initial_heading + 360) % 360
    return compass_heading


def latlon2vu(lat, lon, dim, label='upper'):
	"""
	Estimate the  meriodional and zonal velocity based on the
	latitude and longitude coordinates.

	Paramaters
	----------
    lat : xarray.DataArray
        Latitudinal spherical coordinates
    lon : xarray.DataArray
        Longitudinal spherical coordinates
    dim : str
        Name of the time dimension.

	Returns
	------
	v : xarray.DataArray
		The meridional velocity
	u : xarray.DataArray
		The zonal velocity
	"""
    # Compute spatial difference
	dy, dx = latlon2dydx(lat, lon, dim=dim, label=label)
    # Compute time difference in seconds
	dt = pd.to_numeric(lat[dim].diff(dim=dim, label=label)) * 1e-9
	v, u = dx / dt, dy / dt
	return v, u


def inpolygon(data, poly):
    """
    Mask the data outside a polygon using shapely. Data must have the
    longitudinal and latitudinal coordinates 'lon' and 'lat', respectively.

    Paramaters
    ----------
    data : xarray.DataArray or xarray.Dataset
        The data to mask
    poly : BaseGeometry
        A polygon defined using shapely

    Returns
    -------
    res : xarray.DataArray or xarray.Dataset
        The data masked outside the poylgon
    """
    from shapely.geometry import Point
    lon = data['lon']
    lat =  data['lat']
    def inpolygon(polygon, xp, yp):
        return np.array([Point(x, y).intersects(polygon)
                         for x, y in zip(xp, yp)], dtype=np.bool)
    mask = inpolygon(poly, lon.data.ravel(), lat.data.ravel())
    da_mask = xr.DataArray(mask, dims=lon.dims, coords=lon.coords)
    return da_mask


def inpolygon2d(data, poly):
    from shapely.geometry import Point 
    lon = data['lon']
    lat =  data['lat']
    lon2d, lat2d = xr.broadcast(lon, lat)
    def inpolygon(polygon, xp, yp):
        return np.array([Point(x, y).intersects(polygon) for x, y in zip(xp, yp)], dtype=np.bool)
    mask = inpolygon(poly, lon2d.data.ravel(), lat2d.data.ravel()).reshape(lon2d.shape)
    da_mask = xr.DataArray(mask, dims=lon2d.dims, coords=lon2d.coords)
    return da_mask