Source code for dypy.small_tools

# -*- coding:utf-8 -*-
from datetime import timedelta
import copy
import math

import fiona
from fiona import crs
import cartopy
import numpy as np
import numpy.ma as ma
import scipy as sp
from pyproj import Geod
from PIL import Image, ImageDraw

from dypy.intergrid import Intergrid
from dypy.constants import rdv, L, R, cp, t_zero, g

__all__ = ['CrossSection', 'interpolate', 'rotate_points', 'interval',
           'dewpoint', 'esat', 'moist_lapse', 'equivalent_pot_temp']


class NotAllRequiredVarsException(Exception):
    pass


[docs]class CrossSection(object): """Create a cross-section return the distance, the pressure levels, and the variables. Parameters ---------- variables : dictionary A dictionary of the variables with the following structure {'name': np.array}, need to contain as least rlat, rlon coo : list or numpy.ndarray If `coo` is a list of coordinates, it is used of the starting and end points: [(startlon, startlat), (endlon, endlat)] If `coo` is a numpy.ndarray it is used as the cross-section points. coo need then to be similar to : np.array([[10, 45], [11, 46], [12, 47]]) pressure : np.array pressure coordinate as 1D array int2p : bool, optional (default: False) True if variables need to be interpolated on pressure levels, requires p in the variables, the levels are given by `pressure` int2z : bool, optional (default: False) True if variables need to be interpolated on heights levels, requires z in the variables, the levels are given by `pressure` may require `flip`=True flip : bool, optional (default: False) True if variables need to be flip vertically, the first dimension of 3D array will be flipped pollon : float, optional (default: -170) pole_longitude of the rotated grid pollat : float, optional (default: 43) pole_latitude of the rotated grid nbre : int, optional (default: 10000) nbre of points along the cross-section order: {1, 2, 3, 4}, optional (default: 1) order of interpolation see for details """ def __init__(self, variables, coo, pressure, int2p=False, int2z=False, flip=False, pollon=-170, pollat=43, nbre=1000, order=1): variables = copy.deepcopy(variables) # test the input required = {'rlon', 'rlat'} if int2p: required.add('p') if int2z: required.add('z') diff = required.difference(set(variables.keys())) if diff: msg = '{} not in variables dictionary'.format(diff) raise NotAllRequiredVarsException(msg) rlon = variables.pop('rlon') rlat = variables.pop('rlat') self.pressure = pressure self.nz = pressure.size self.order = order if type(coo) in [list, tuple]: self.coo = coo self.nbre = nbre # find the coordinate of the cross-section self.lon, self.lat, crlon, crlat, self.distance = find_cross_points( coo, nbre, pollon, pollat) elif type(coo) == np.ndarray: self.lon = coo[:, 0] self.lat = coo[:, 1] self.nbre = coo.shape[0] crlon, crlat = rotate_points(pollon, pollat, self.lon, self.lat) self.distance = great_circle_distance(coo[0, 0], coo[0, 1], coo[-1, 0], coo[-1, 1]) else: msg = '<coo> should be of type list, tuple or nump.ndarray' raise TypeError(msg) self.distances = np.linspace(0, self.distance, self.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 or int2z: if int2p: grid = variables['p'] if int2z: grid = variables['z'] if flip and grid.ndim == 3: grid = grid[::-1, ...] for var, data in variables.items(): if data.ndim > 2: if flip: data = data[::-1, ...] dataint = interpolate(data, grid, 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=self.order) 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=self.order) 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.3240879770135 """ 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
[docs]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
[docs]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
[docs]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
[docs]def esat(t): """ Calculate the saturation vapor pressure for t in °C Examples -------- >>> esat(0) 610.94000000000005 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 esat_ice(t): """calculate the saturation vapor pressure over ice for t in K Follows COSMO, line 2468 in src_gscp.f90 constant values are from data_constants.f90 Examples -------- >>> esat_ice(273.15) 610.27696635637164 Parameters ---------- t: float, numpy.array temperature in K Returns ------- esat_ice: float, numpy array saturation vapor pressure over ice in Pa """ b1 = 610.78 b2i = 21.8745584 b3 = 273.16 b4i = 7.66 return b1*np.exp(b2i * (t - b3) / (t - b4i))
[docs]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
def read_shapefile(shapefile): """ Return the geometry (list) of each shape as numpy array (nre pts, 2) and the projections of the shapefile in a PROJ4.strings """ with fiona.open(shapefile) as shapef: geometries = [np.array(shape['geometry']['coordinates']).squeeze() for shape in shapef] projection = crs.to_string(shapef.crs) return geometries, projection def mask_polygon(polygon, array, lon, lat): "Mask value outside polygon, work only for regular grid" xsize = array.shape[-1] ysize = array.shape[-2] lon0 = lon[0, 0] lat0 = lat[0, 0] lon1 = lon[-1, -1] lat1 = lat[-1, -1] nx, ny = lon.T.shape dx = (lon1 - lon0) / nx dy = (lat1 - lat0) / ny def return_indices(x, y): i = int((x - lon0) / dx) j = int((y - lat0) / dy) return i, j # transform the polygon coordinates into pixels coordinates pixels = [return_indices(x, y) for x, y in polygon] # rasterize of the polygone raster_poly = Image.new("L", (xsize, ysize), 1) rasterize = ImageDraw.Draw(raster_poly) rasterize.polygon(pixels, 0) mask = sp.fromstring(raster_poly.tobytes(), 'b') mask.shape = raster_poly.im.size[1], raster_poly.im.size[0] # clip the raster if array.ndim >= 3: mask.shape = (1,) + mask.shape mask = np.broadcast_arrays(mask, array)[0] clipped = ma.masked_where(mask != 0, array) return clipped def potential_temperature(t, p, p0=1000): "for p in hPa and t in K" return t * (p0 / p)**(2/7.) def virtual_temperature(t, qv): """ Based on http://glossary.ametsoc.org/wiki/Virtual_potential_temperature """ return t * (1 + 0.61 * qv) def brunt_vaisala(theta, z): """ Based on http://glossary.ametsoc.org/wiki/Brunt-väisälä_frequency """ dthdz = np.gradient(theta)[0] / np.gradient(z)[0] return np.sqrt((g / theta) * dthdz) def merge_dict(dict_1, dict_2): """Merge two dictionaries. Values that evaluate to true take priority over falsy values. `dict_1` takes priority over `dict_2`. """ return dict((str(key), dict_1.get(key) or dict_2.get(key)) for key in set(dict_2) | set(dict_1))
[docs]def equivalent_pot_temp(t, p, qv): """Return the equivalent potential temperature in K computation of equivalent potential temperature according to Bolton (1980) except constant 0.1998 (4.805 according to Bolton). Args: t (np.array, np.float): Temperature in K p (np.array, np.float): Pressure in hPa qv (np.array, np.float): Specific humidity in kg/kg Returns: the (np.array, np.float): equivalent potential temperature in K Examples: >>> t = 16 + 273 >>> qv = 0.01156 >>> p = 850 >>> equivalent_pot_temp(t, p, qv) 337.59612858187029 """ r = 1 / ((1 / qv) - 1) e = 100 * p * qv / (0.622 + 0.378 * qv) tl = 2840 / (3.5 * np.log(t) - np.log(e) - 0.1998) + 55 the = t * (1000 / p) ** (0.2854 * (1 - 0.28 * r)) the *= np.exp(((3.376 / tl) - 0.00254) * 1e3 * r * (1 + 0.81 * r)) return the
def ipython(__globals__, __locals__, __msg__=None, __err__=66): """Drop into an iPython shell with all global and local variables. To pass the global and local variables, call with globals() and locals() as arguments, respectively. Notice how both functions must called in the call to ipython. To exit the program after leaving the iPython shell, pass an integer, which will be returned as the error code. To display a message when dropping into the iPython shell, pass a string after the error code. Examples -------- >>> ipython(globals(), locals(), "I'm here!") """ import warnings with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=DeprecationWarning) warnings.filterwarnings("ignore", category=PendingDeprecationWarning) import IPython print('\n----------\n') globals().update(__globals__) locals().update(__locals__) if __msg__ is not None: print("\n{l}\n{m}\n{l}\n".format(l="*"*60, m=__msg__)) IPython.embed(display_banner=None) if __err__ is not None: exit(__err__)