Subversion Repositories lagranto.ecmwf

# -*- coding:utf-8 -*-

import cartopy

import copy

import math

import numpy as np

from pyproj import Geod

from datetime import timedelta

from dypy.intergrid import Intergrid

from dypy.constants import rdv, L, R, cp, t_zero

__all__ = ['CrossSection', 'interpolate', 'rotate_points', 'interval',

           'dewpoint', 'esat', 'moist_lapse']

class CrossSection(object):

    """

    Create a cross-section

    return the distance, the pressure levels,

    and the variables.

    """

    def __init__(self, variables, coo, pressure,

                 int2p=False, pollon=-170, pollat=43, nbre=1000, order=1):

        """

            variables: dictionary of variables {'name': np.array},

                need to contain rlat, rlon

            coo: list of coordinates [(startlon, startlat), (endlon, endlat)]

            pressure: pressure coordinate 1d np.array

            int2p: True if variables need to be interpolated,

                require P in the variables

            pollon: pole_longitude of the rotated grid

            pollat: pole_latitude of the rotated grid

            nre: nbre of points along the cross-section

            order: order of interpolation see for details

        """

        variables = copy.deepcopy(variables)

        # test the input

        required = ['rlon', 'rlat']

        if int2p:

            required.append('p')

        for var in required:

            if var not in variables.keys():

                raise Exception('{} not in variables dictionary'.format(var))

        rlon = variables.pop('rlon')

        rlat = variables.pop('rlat')

        self.coo = coo

        self.nbre = nbre

        self.pressure = pressure

        self.nz = pressure.size

        if int2p:

            p = variables['p']

        # find the coordinate of the cross-section

        self.lon, self.lat, crlon, crlat, self.distance = find_cross_points(

            coo, nbre, pollon, pollat)

        self.distances = np.linspace(0, self.distance, nbre)

        # get the information for intergrid

        self._lo = np.array([rlat.min(), rlon.min()])

        self._hi = np.array([rlat.max(), rlon.max()])

        self._query_points = [[lat, lon] for lat, lon in zip(crlat, crlon)]

        if int2p:

            for var, data in variables.items():

                if data.ndim > 2:

                    dataint = interpolate(data, p, pressure)

                    variables[var] = dataint

        variables = self._int2cross(variables)

        self.__dict__.update(variables)

    def _int2cross(self, variables):

        for var, data in variables.items():

            if data.squeeze().ndim > 2:

                datacross = np.zeros((self.nz, self.nbre))

                for i, datah in enumerate(data.squeeze()):

                    intfunc = Intergrid(datah, lo=self._lo, hi=self._hi,

                                        verbose=False, order=1)

                    datacross[i, :] = intfunc.at(self._query_points)

            else:

                datacross = np.zeros((self.nbre, ))

                intfunc = Intergrid(data.squeeze(), lo=self._lo, hi=self._hi,

                                    verbose=False, order=1)

                datacross[:] = intfunc.at(self._query_points)

            variables[var] = datacross

        return variables

def find_cross_points(coo, nbre, pole_longitude, pole_latitude):

    """ give nbre points along a great circle between coo[0] and coo[1]

        iand rotated the points

    """

    g = Geod(ellps='WGS84')

    cross_points = g.npts(coo[0][0], coo[0][1], coo[1][0], coo[1][1], nbre)

    lat = np.array([point[1] for point in cross_points])

    lon = np.array([point[0] for point in cross_points])

    rlon, rlat = rotate_points(

        pole_longitude, pole_latitude, lon, lat)

    distance = great_circle_distance(coo[0][0], coo[0][1],

                                     coo[1][0], coo[1][1])

    return lon, lat, rlon, rlat, distance

def great_circle_distance(lon1, lat1, lon2, lat2):

    """

    return the distance (km) between points following a great circle

    based on : https://gist.github.com/gabesmed/1826175

    >>> great_circle_distance(0, 55, 8, 45.5)

    1199.13065955064

    """

    EARTH_CIRCUMFERENCE = 6378137     # earth circumference in meters

    dLat = math.radians(lat2 - lat1)

    dLon = math.radians(lon2 - lon1)

    a = (math.sin(dLat / 2) * math.sin(dLat / 2) +

         math.cos(math.radians(lat1)) * math.cos(math.radians(lat2)) *

         math.sin(dLon / 2) * math.sin(dLon / 2))

    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))

    distance = EARTH_CIRCUMFERENCE * c

    return distance/1000

def interpolate(data, grid, interplevels):

    """

    intrerpolate data to grid at interplevels

    data and grid are numpy array in the form (z,lat,lon)

    """

    data = data.squeeze()

    grid = grid.squeeze()

    shape = list(data.shape)

    if (data.ndim > 3) | (grid.ndim > 3):

        message = "data and grid need to be 3d array"

        raise IndexError(message)

    try:

        nintlev = len(interplevels)

    except:

        interplevels = [interplevels]

        nintlev = len(interplevels)

    shape[-3] = nintlev

    outdata = np.ones(shape)*np.nan

    if nintlev > 20:

            for idx, _ in np.ndenumerate(data[0]):

                column = grid[:, idx[0], idx[1]]

                column_GRID = data[:, idx[0], idx[1]]

                value = np.interp(

                    interplevels,

                    column,

                    column_GRID,

                    left=np.nan,

                    right=np.nan)

                outdata[:, idx[0], idx[1]] = value[:]

    else:

        for j, intlevel in enumerate(interplevels):

            for lev in range(grid.shape[0]):

                cond1 = grid[lev, :, :] > intlevel

                cond2 = grid[lev-1, :, :] < intlevel

                right = np.where(cond1 & cond2)

                if right[0].size > 0:

                    sabove = grid[lev, right[0], right[1]]

                    sbelow = grid[lev-1, right[0], right[1]]

                    dabove = data[lev, right[0], right[1]]

                    dbelow = data[lev-1, right[0], right[1]]

                    result = (intlevel-sbelow)/(sabove-sbelow) * \

                        (dabove-dbelow)+dbelow

                    outdata[j, right[0], right[1]] = result

    return outdata

def rotate_points(pole_longitude, pole_latitude, lon, lat, direction='n2r'):

    """

        rotate points from non-rotated system to rotated system

        n2r : non-rotated to rotated (default)

        r2n : rotated to non-rotated

        return rlon,rlat,

    """

    lon = np.array(lon)

    lat = np.array(lat)

    rotatedgrid = cartopy.crs.RotatedPole(

        pole_longitude=pole_longitude,

        pole_latitude=pole_latitude

        )

    standard_grid = cartopy.crs.Geodetic()

    if direction == 'n2r':

        rotated_points = rotatedgrid.transform_points(standard_grid, lon, lat)

    elif direction == 'r2n':

        rotated_points = standard_grid.transform_points(rotatedgrid, lon, lat)

    rlon, rlat, _ = rotated_points.T

    return rlon, rlat

def interval(starting_date, ending_date, delta=timedelta(hours=1)):

    curr = starting_date

    while curr < ending_date:

        yield curr

        curr += delta

def dewpoint(p, qv):

    """

    Calculate the dew point temperature based on p and qv following (eq.8):

    Lawrence, M.G., 2005. The Relationship between Relative Humidity

    and the Dewpoint Temperature in Moist Air:

    A Simple Conversion and Applications.

    Bulletin of the American Meteorological Society 86,

    225–233. doi:10.1175/BAMS-86-2-225

    """

    B1 = 243.04  # °C

    C1 = 610.94  # Pa

    A1 = 17.625

    e = p*qv/(rdv+qv)

    td = (B1*np.log(e/C1))/(A1-np.log(e/C1))

    return td

def esat(t):

    """

    Calculate the saturation vapor pressure for t in °C

    Following eq. 6 of

    Lawrence, M.G., 2005. The Relationship between Relative Humidity

    and the Dewpoint Temperature in Moist Air:

    A Simple Conversion and Applications.

    Bulletin of the American Meteorological Society 86,

    225–233. doi:10.1175/BAMS-86-2-225

    """

    C1 = 610.94  # Pa

    A1 = 17.625    #

    B1 = 243.03  # °C

    return C1*np.exp((A1 * t) / (B1 + t))

def moist_lapse(t, p):

    """ Calculates moist adiabatic lapse rate for T (Celsius) and p (Pa)

    Note: We calculate dT/dp, not dT/dz

    See formula 3.16 in Rogers&Yau for dT/dz, but this must be combined with

    the dry adiabatic lapse rate (gamma = g/cp) and the

    inverse of the hydrostatic equation (dz/dp = -RT/pg)"""

    a = 2. / 7.

    b = rdv * L * L / (R * cp)

    c = a * L / R

    e = esat(t)

    wsat = rdv * e / (p - e)  # Rogers&Yau 2.18

    numer = a * (t + t_zero) + c*wsat

    denom = p * (1 + b * wsat / ((t + t_zero)**2))

    return numer / denom