#!/usr/bin/python
########################################################################
"""
Spatial computation package
"""
# 2013, G. SERAZIN, PhD student at MEOM/LGGE Grenoble
########################################################################
#=======================================================================
# Main module importation
#=======================================================================
try:
from netCDF4 import Dataset
except:
print 'Package numpy is not available...'
exit()
try:
import numpy as np
import numpy.ma as ma
import numpy.linalg as li
import numpy.fft as fft
except:
print 'Package numpy is not available...'
exit()
try:
import scipy.stats as st
import scipy.interpolate
except:
print 'Package scipy is not available...'
exit()
try:
from mpl_toolkits import basemap
except:
print 'Package scipy is not available...'
#=======================================================================
# Useful constants
#=======================================================================
R_earth = 6371 * 1000
G_earth = 9.81
omega_earth = 2 * np.pi / 86164.
#=======================================================================
# Functions definition
#=======================================================================
[docs]def zomean(var, lat, lat_ref, mask, dx=1, dy=1):
'''
Compute the spatial mean of a variable using:
- e1,e2: non-uniform gridto weight by surface
- lat_ref: a vector of referenced latitude to decompose the zonal
mean into bands
'''
l = np.size(lat_ref)
var_mean = np.zeros(l-1)
lat_mean = np.zeros(l-1)
if np.size(np.shape(lat)) == 1 :
lat = np.reshape(np.resize(lat, np.size(var)), np.shape(var),
order='F')
for i in range(l-1):
I = np.where((lat_ref[i] < lat) & (lat <= lat_ref[i+1])
& (mask == 1))
w = dx[I] * dy[I] / (ma.sum(dx[I] * dy[I]))
var_mean[i] = ma.sum(w * var[I])
lat_mean[i] = ma.sum(w * lat[I])
return var_mean, lat_mean
[docs]def zonal_corr(var1, var2, lat, lat_ref, e1=1, e2=1):
'''
Compute the spatial correlation of two variables over zonal bands:
- e1,e2: non-uniform gridto weight by surface
- lat_ref: a vector of referenced latitude to decompose the zonal
mean into bands
'''
l = np.size(lat_ref)
corr = np.zeros(l - 1)
lat_mean = np.zeros(l - 1)
for i in range(l - 1):
I = np.where((lat_ref[i] < lat) & (lat <= lat_ref[i + 1]))
w = e1[I] * e2[I] / (ma.sum(e1[I] * e2[I]))
X = w * var1[I]
Y = w * var2[I]
corr[i], pval = st.pearsonr(X.compressed(),Y.compressed())
lat_mean[i] = ma.sum(w*lat[I])
return corr, lat_mean
[docs]def global_av(var, e1=1, e2=1):
'''
Compute the spatial mean of a variable using a non-uniform grid to
weight by surface
'''
w=e1*e2/ma.mean(e1*e2)
var_mean=ma.mean(w*var)
return var_mean
[docs]def vect_mod(var_x, var_y, mode='same'):
'''
Compute the modulus of a 2D vector from its x-y components
calculated on a C grid
'''
n,m=np.shape(var_x)
var_x_E=np.array(var_x)
var_x_E[:,1:]=var_x[:,:-1]
if mode=='same':
var_x_E[:,0]=var_x[:,1]
elif mode=='cylindric':
var_x_E[:,0]=var_x[:,-1]
var_y_S=np.array(var_y)
var_y_S[1:,:]=var_y[:-1,:]
var_y_S[0,:]=var_y[1,:]
var_mod=np.sqrt(1./2*(var_x_E**2+var_x**2+var_y_S**2+var_y**2))
return var_mod
[docs]def detrend(x, y, var):
"""
Fit a 2D linear surface from the input field in least-square sense
and return the detrended field
"""
mask = ma.getmask(var)
if not mask:
B = np.ravel(var)
X = np.ravel(x)
Y = np.ravel(y)
else:
ind = ma.where(mask == 0)
B = var[ind]
X = x[ind]
Y = y[ind]
n = np.size(B)
A = np.vstack([X, Y, np.ones(n)]).T
#A = np.vstack([X ** 2, Y ** 2, X * Y, X, Y, np.ones(n)]).T
#A = np.vstack([X, Y, np.ones(n)])
#Ainv = np.dot(li.pinv(np.outer(A, A)), A)
#a, b, c = np.dot(Ainv, B)
# Solve the least square system
fit = li.lstsq(A, B)
a, b, c = fit[0]
#a, b, c, d, e, f = fit[0]
#var_dtr = var - (a * x ** 2 + b * y **2 + c * x * y + d * x +
#e * y + c)
var_dtr = var - (a * x + b * y + c)
return var_dtr
[docs]def latlon2yx(lat, lon):
"""
Convert latitude and longitude arrays to y and x arrays in km
"""
if np.size(np.shape(lat)) == 1:
lon2d, lat2d = np.meshgrid(lon, lat)
else:
lon2d, lat2d = lon, lat
y = (1 / 180.) * np.pi * R_earth * lat2d
x = np.cos(np.pi / 180. * lat2d) * np.pi / 180. * R_earth * lon2d
return y, x
[docs]def latlon2lk(lat, lon, n_l, n_k):
"""
Define wavenumber vectors according to the latitude and longitude
arrays
"""
dy = (1. / n_lat * (1 / 180.) * np.pi * R_earth *
(np.max(lat)-np.min(lat)))
dx = (1. / n_lon * np.cos(np.pi / 180. * np.mean(lat)) *
np.pi / 180. * R_earth * (np.max(lon)-np.min(lon)))
l = fft.fftshift(fft.fftfreq(n_l, d=dy))
k = fft.fftshift(fft.fftfreq(n_k, d=dx))
return l, k
[docs]def latlon2dydx(lat, lon, n_lat, n_lon):
"""
Define wavenumber vectors according to the latitude and longitude
arrays
"""
delta_lon = np.max(lon)-np.min(lon)
if delta_lon > 350:
indx = np.where(lon < 0)
if not indx:
indx = np.where(lon > 180)
lon[indx] -= 360
else:
lon[indx] += 360
delta_lon = np.max(lon)-np.min(lon)
dy = 1. / n_lat * (1 / 180.) * np.pi * R_earth * (np.max(lat)-np.min(lat))
dx = (1. / n_lon * np.cos(np.pi / 180. * np.mean(lat)) *
np.pi / 180. * R_earth * (delta_lon))
return dy, dx
[docs]def lat2f(lat):
"""
Return the coriolis parameter computed from the latitude
"""
f = 2 * omega_earth * np.sin(np.pi / 180. * lat)
return f
[docs]def lat2dy(lat):
"""
Compute dy on a regular grid
"""
dy = (1 / 180.) * np.pi * R_earth * np.diff(lat, n=1, axis=0)
return dy
[docs]def lon2dx(lat, lon):
"""
Compute dx on a regular grid
"""
delta_lon = np.max(lon)-np.min(lon)
if delta_lon > 350:
indx = np.where(lon < 0)
if not indx:
indx = np.where(lon > 180)
lon[indx] -= 360
else:
lon[indx] += 360
delta_lon = np.max(lon)-np.min(lon)
dx = (1. / n_lon * np.cos(np.pi / 180. * np.mean(lat)) *
np.pi / 180. * R_earth * (delta_lon))
return dx
[docs]def ddy(val, dy=1, axis=-2):
"""
Compute the first derivative along the y axis
"""
dval = np.diff(val, n=1, axis=axis)
dvaldy = dval / dy
return dvaldy
[docs]def ddx(val, dx=1, axis=-1):
"""
Compute the first derivative along the x axis
"""
dval = np.diff(val, n=1, axis=axis)
dvaldx = dval / dx
return dvaldx
[docs]def zogeo(ssh, lat, dy=None, axis=-1):
"""
Compute the zonal surface geostrophy current from a sea level field
"""
if dy == None:
dy = lat2dy(lat)
f = lat2f(lat)[:-1]
u_s = - G_earth / f * ddy(ssh, dy, axis=axis)
return u_s
def megeo(ssh, lat, lon, dx=None, axis=-1):
if dx == None:
dx = lon2dx(lon)
f = lat2f(lat)[:-1]
v_s = G_earth / f * ddx(ssh, dx, axis=axis)
return v_s
[docs]def hgrad(var, e1, e2):
"""
Compute the horizontal gradient
"""
n,m=np.shape(var)
var_E=np.array(var)
e1_E=np.array(var)
var_E[:,1:]=var[:,:-1]
if mode=='same':
var_E[:,0]=var[:,1]
elif mode=='cylindric':
var_E[:,0]=var[:,-1]
dvar_dx=(var_E-var)
var_S=np.array(var)
var_y_S[1:,:]=var_y[:-1,:]
var_y_S[0,:]=var_y[1,:]
return dvar_dx,dvar_dy
# /!\ Test this function !!!!!!
[docs]def interpolate(x, y, var, interp=None, kind='linear'):
"""
Interpolate a 2D field on a regular grid
Parameters
----------
var : 2D array_like
The values of the variable to interpolate at the data points.
y, x : 1D or 2D array_like
The y and x coordinates on which the variable is evaluated.
interp : {None, 'basemap', 'scipy'}, optional
The algorithm to use for the interpolation.
kind : {'linear', 'cubic', 'quintic'}
The kind of spline interpolation to use. Default is 'linear'.
Returns
-------
var_reg : 2D array_like
The values of the variable evaluated on the regular grid
y_reg, x_reg : 2D array_like
The y and x coordinates of the new grid
"""
orderdict = {'linear':1, 'cubic':3}
if np.size(np.shape(y)) == 1:
x2d, y2d = np.meshgrid(x, y)
else:
x2d, y2d = x, y
x_reg = np.linspace(np.min(x2d), np.max(x2d), len(x1d_in))
y_reg = np.linspace(np.min(y2d), np.max(y2d), len(y1d_in))
x2d_reg, y2d_reg = np.meshgrid(x_reg ,y_reg)
if interp is None:
return x2d_in, y2d_in, var.copy()
elif interp == 'basemap':
var_reg = basemap.interp(var, x1d_in, y1d_in, x2d_reg, y2d_reg,
checkbounds=False,
order=orderdict[kind])
return x2d_reg, y2d_reg, var_reg
elif interp == 'scipy':
interp = scipy.interpolate.interp2d(x1d_in, y1d_in, var,
kind=kind)
a1d = interp(x2d_reg[0,:], y2d_reg[:,0])
var_reg = np.reshape(a1d, y2d_reg.shape)
return x2d_reg, y2d_reg, var_reg