# -*- coding:utf-8 -*-
import math
from datetime import timedelta
import cartopy
import copy
import fiona
import numpy as np
import numpy.ma as ma
import scipy as sp
from PIL import Image, ImageDraw
from fiona import crs
from pyproj import Geod
from dypy.constants import rdv, L, R, cp, t_zero, g
from dypy.intergrid import Intergrid
__all__ = ['CrossSection', 'interpolate', 'rotate_points', 'interval',
'dewpoint', 'esat', 'moist_lapse', 'equivalent_pot_temp',
'mask_polygon', 'potential_temperature',
'great_circle_distance', 'read_shapefile']
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}.
If version is rotated: need to contain as least rlon, rlat
if version is regular: need to contain at least lon, lat
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
version: string (defautl rotated)
type of grid used as input (rotated or regular)
"""
def __init__(self, variables, coo, pressure, int2p=False, int2z=False,
flip=False, pollon=-170, pollat=43, nbre=1000, order=1,
version='rotated'):
variables = copy.deepcopy(variables)
# test the input
if version == 'rotated':
required = {'rlon', 'rlat'}
elif version == 'regular':
required = {'lon', 'lat'}
else:
msg = 'only rotated or regular are supported'
raise NotImplementedError(msg)
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)
if version == 'rotated':
rlon = variables.pop('rlon')
rlat = variables.pop('rlat')
elif version == 'regular':
lon = variables.pop('lon')
lat = variables.pop('lat')
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
if version == 'rotated':
(self.lon, self.lat, crlon, crlat,
self.distance) = find_cross_points(coo, nbre, pollon, pollat)
elif version == 'regular':
self.lon, self.lat, self.distance = find_cross_points(coo, nbre)
crlon, crlat = self.lon, self.lat
elif type(coo) == np.ndarray:
self.lon = coo[:, 0]
self.lat = coo[:, 1]
self.nbre = coo.shape[0]
if version == 'rotated':
crlon, crlat = rotate_points(pollon, pollat, self.lon, self.lat)
elif version == 'regular':
crlon, crlat = 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
if version == 'rotated':
self._lo = np.array([rlat.min(), rlon.min()])
self._hi = np.array([rlat.max(), rlon.max()])
elif version == 'regular':
self._lo = np.array([lat.min(), lon.min()])
self._hi = np.array([lat.max(), lon.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=None, pole_latitude=None):
""" 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])
distance = great_circle_distance(coo[0][0], coo[0][1],
coo[1][0], coo[1][1])
if (pole_latitude is None) and (pole_latitude is None):
return lon, lat, distance
else:
rlon, rlat = rotate_points(pole_longitude, pole_latitude, lon, lat)
return lon, lat, rlon, rlat, distance
[docs]def great_circle_distance(lon1, lat1, lon2, lat2):
"""Calculate the great circle distance between two points
based on : https://gist.github.com/gabesmed/1826175
Parameters
----------
lon1: float
longitude of the starting point
lat1: float
latitude of the starting point
lon2: float
longitude of the ending point
lat2: float
latitude of the ending point
Returns
-------
distance (km): float
Examples
--------
>>> 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):
"""interpolate `data` on `grid` for given `interplevels`
Interpolate the `data` array at every given levels (`interplevels`) using
the `grid` array as reference.
The `grid` array need to be increasing along the first axis.
Therefore ERA-Interim pressure must be flip: p[::-1, ...], but not
COSMO pressure since its level 0 is located at the top of the atmosphere.
Parameters
----------
data : array (nz, nlat, nlon)
data to interpolate
grid : array (nz, nlat, nlon)
grid use to perform the interpolation
interplevels: list, array
list of the new vertical levels, in the same unit as the grid
Returns
-------
interpolated array: array (len(interplevels), nlat, nlon)
Examples
--------
>>> print(qv.shape)
(60, 181, 361)
>>> print(p.shape)
(60, 181, 361)
>>> levels = np.arange(200, 1050, 50)
(17,)
>>> qv_int = interpolate(qv, p, levels)
>>> print(qv_int.shape)
(17, 181, 361)
"""
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 lon, lat from/to a rotated system
Parameters
----------
pole_longitude: float
longitudinal coordinate of the rotated pole
pole_latitude: float
latitudinal coordinate of the rotated pole
lon: array (1d)
longitudinal coordinates to rotate
lat: array (1d)
latitudinal coordinates to rotate
direction: string, optional
direction of the rotation;
n2r: from non-rotated to rotated (default)
r2n: from rotated to non-rotated
Returns
-------
rlon: array
rlat: array
"""
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
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
Parameters
----------
p: float, array
pressure in Pa
qv: float, array
specific humidity in kg/kg
Returns
-------
dewpoint (in °C): float, array
"""
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
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
Parameters
----------
t: float, numpy.array
temperature in K
Examples
--------
>>> esat(0)
610.94000000000005
"""
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
Parameters
----------
t: float, numpy.array
temperature in K
Returns
-------
esat_ice: float, numpy array
saturation vapor pressure over ice in Pa
Examples
--------
>>> esat_ice(273.15)
610.27696635637164
"""
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
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)
Parameters
----------
t : float, array
temperature in degree Celsius
p : float, array
temperature in Pa
Returns
-------
moist adiabatic lapse rate: float, array
"""
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
[docs]def read_shapefile(shapefile):
"""Return the geometry (list) of each shape and the projection
Parameters
----------
shapefile: string
path to a shapefile
Returns
-------
geometries : list of array (number of points, 2)
projection : projection string (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
[docs]def mask_polygon(polygon, array, lon, lat):
"""Mask value outside polygon, work only for regular grid
Parameters
----------
polygon: array (X x 2)
coordinates of the polygon
array: array (nlon x nlat)
array to mask
lon: array (nlon x nlat)
regular coordinates of the array
lat: array (nlon x nlat)
regular coordinates of the array
Returns
-------
masked array: MaskedArray (nlon x nlat)
Examples
--------
>>> import numpy as np
>>> polygon = np.array([[5, 5],
>>> [10, 5],
>>> [10, 10],
>>> [5, 10],
>>> [5, 5]])
>>> array = np.zeros((20, 20))
>>> lon = np.arange(0, 20, dtype=float)
>>> lat = np.arange(0, 20, dtype=float)
>>> lons, lats = np.meshgrid(lon, lat)
>>> clipped = mask_polygon(polygon, array, lons, lats)
"""
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
[docs]def potential_temperature(t, p, p0=1000):
"""Calculate the potential temperature
Parameters
----------
t: array
temperature in K
p: array
pressure in hPa
p0: float, optional
pressure at sea level in hPa
Returns
-------
potential temperature: array
"""
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)
[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).
Parameters
----------
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 (equivalent potential temperature in K) : (np.array, np.float)
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 path_along_domain_boundary(lon, lat, nbnd=0):
"""Construct a path parallel to the domain boundary.
Slice the lon and lat arrays. The distance to the boundary is given in
grid points as <nbnd>.
Returns
-------
x coordinates of path: array
y coordinates of path: array
"""
# SR_TODO: Find out why e.g. the first is not the southern boundary!
# SR_TODO: Somehow the values in lon,lat are stored as (y,x)...
ns, ne = nbnd, -(nbnd + 1)
pxs = (
lon[ns:ne, ns].tolist() + # west
lon[ne, ns:ne].tolist() + # north
lon[ne:ns:-1, ne].tolist() + # east
lon[ns, ne:ns:-1].tolist() # south
)
pys = (
lat[ns:ne, ns].tolist() + # west
lat[ne, ns:ne].tolist() + # north
lat[ne:ns:-1, ne].tolist() + # east
lat[ns, ne:ns:-1].tolist() # south
)
if (pxs[-1], pys[-1]) != (pxs[0], pys[0]):
pxs.append(pxs[0])
pys.append(pys[0])
return pxs, pys
def return_erainterim_pressure(ps):
"""Return the pressure in hPa
Parameters
----------
ps: np.ndarray (2D)
surface pressure in hPa
Returns
-------
pressure in hPa (60 x ps.shape): np.ndarray
"""
aklay = np.array([[[0.00000000e+00, 3.68387103e-02, 3.66284877e-01,
1.38141561e+00, 3.38863707e+00, 6.61347675e+00,
1.12063723e+01, 1.72484627e+01, 2.47573814e+01,
3.36930466e+01, 4.39634552e+01, 5.54304657e+01,
6.79156036e+01, 8.12057953e+01, 9.50592804e+01,
1.09211288e+02, 1.23379883e+02, 1.37271729e+02,
1.50587860e+02, 1.63029587e+02, 1.74304138e+02,
1.84130539e+02, 1.92245392e+02, 1.98408646e+02,
2.02409409e+02, 2.04071716e+02, 2.03260330e+02,
1.99886566e+02, 1.93914017e+02, 1.85364380e+02,
1.74323273e+02, 1.60996384e+02, 1.45775635e+02,
1.29263840e+02, 1.12267853e+02, 9.57058945e+01,
8.03584366e+01, 6.66232605e+01, 5.46236343e+01,
4.43349915e+01, 3.57835655e+01, 2.88815498e+01,
2.33108120e+01, 1.88145714e+01, 1.51855755e+01,
1.22565460e+01, 9.89247513e+00, 7.98439217e+00,
6.44434547e+00, 5.20134640e+00, 4.19295025e+00,
3.36233878e+00, 2.66637444e+00, 2.07681704e+00,
1.57533824e+00, 1.15060127e+00, 7.96423793e-01,
5.10365665e-01, 2.92126685e-01, 9.99999940e-02]]],
dtype=sp.float32).T
bklay = np.array([[[9.98815060e-01, 9.95824814e-01, 9.91144776e-01,
9.83966410e-01, 9.73653972e-01, 9.59733367e-01,
9.41880941e-01, 9.19912100e-01, 8.93770397e-01,
8.63515913e-01, 8.29314172e-01, 7.91424990e-01,
7.50191212e-01, 7.06027210e-01, 6.59408033e-01,
6.10857964e-01, 5.60939193e-01, 5.10240734e-01,
4.59367245e-01, 4.08927530e-01, 3.59523505e-01,
3.11738938e-01, 2.66128063e-01, 2.23204523e-01,
1.83430105e-01, 1.47203416e-01, 1.14848614e-01,
8.66042674e-02, 6.26121163e-02, 4.29057851e-02,
2.73995195e-02, 1.59103926e-02, 8.11201334e-03,
3.44813662e-03, 1.13827549e-03, 2.68609205e-04,
3.79117482e-05, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00]]],
dtype=sp.float32).T
ps = ps.reshape((1,) + ps.shape)
return aklay + ps * bklay