# -*- 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__)