def get_time_idx(self, utc):
"""
Get closest index to desired time
Args:
utcs (list,tuple, int) : times
Returns:
tidx (int): closest index to desired time
"""
dn = utils.ensure_datenum(utc)
dns = utils.get_sequence(dn)
tidx = []
for t in dns:
tidx.append(utils.closest(N.array(self.utcs), t))
return N.array(tidx)
def get(self,
vrbl,
utc=None,
level=None,
lats=None,
lons=None,
smooth=1,
other=False):
"""
Get data.
Will interpolate onto pressure, height coordinates if needed.
Will smooth if needed.
Level:
* indices: integer or N.ndarray of integers
* pressure: string ending in 'hPa'
* height: string ending in 'm' or 'km'
* isentropic: string ending in 'K'
Lats:
* indices: integer or N.ndarray of integers
* lats: float or N.ndarray of floats
Lons:
* indices: integer or N.ndarray of integers
* lons: float or N.ndarray of floats
Args:
vrbl : WRF or computed variable required
utc : indices (<500): integer/N.ndarray of integers.
time tuple: 6-item tuple or list/tuple of these.
datenum: integer >500 or list of them.
"""
# import pdb; pdb.set_trace()
# Time
if utc is None:
tidx = None
elif isinstance(utc, int) and (utc < 500):
# elif isinstance(utc,(int,N.int64)) and utc<500:
if utc < 0:
# Logic to allow negative indices
tidx = self.t_dim + utc
else:
tidx = utc
elif isinstance(
utc,
(list, tuple, int, datetime.datetime)): # and len(utc[0])==6:
tidx = self.get_time_idx(utc)
elif isinstance(utc, N.ndarray): #and isinstance(utc[0],int):
tidx = utc
else:
print("Invalid time selection.")
raise Exception
# Level
# if not level:
coords = utils.check_vertical_coordinate(level)
# import pdb; pdb.set_trace()
if (level is None) or (coords is 'eta'):
lvidx = None
elif coords == 'index':
lvidx = level
elif isinstance(coords, str):
if coords == 'surface':
lvidx = 0
else:
lvidx = coords
else:
print("Invalid level selection.")
raise Exception
# Lat/lon
if not type(lats) == type(lons):
# What about case where all lats with one lon?
raise Exception
if lats is None:
lonidx = None
latidx = None
elif isinstance(lons, (list, tuple, N.ndarray)):
if isinstance(lons[0], int):
lonidx = lons
latidx = lats
elif isinstance(lons[0], float):
# Interpolate to lat/lon
lonidx = None
latidx = None
elif isinstance(lons, (int, N.int64)):
lonidx = lons
latidx = lats
elif isinstance(lons, float):
# Interpolate to lat/lon
lonidx = utils.closest(self.lons1D, lons)
latidx = utils.closest(self.lats1D, lats)
else:
print("Invalid lat/lon selection.")
raise Exception
# Check if computing required
# When data is loaded from nc, it is destaggered
if debug_get:
print(("Computing {0} for level {1} of index {2}".format(
vrbl, level, lvidx)))
if self.check_compute(vrbl):
if debug_get:
print(("Variable {0} exists in dataset.".format(vrbl)))
if lvidx is 'isobaric':
data = self.get_p(vrbl, tidx, level, lonidx, latidx)
elif isinstance(lvidx, (tuple, list, N.ndarray, int, type(None))):
data = self.load(vrbl, tidx, lvidx, lonidx, latidx)
else:
raise Exception
else:
if debug_get:
print(("Variable {0} needs to be computed.".format(vrbl)))
if lvidx is 'isobaric':
data = self.compute(vrbl, tidx, level, lonidx, latidx, other)
else:
data = self.compute(vrbl, tidx, lvidx, lonidx, latidx, other)
data_4d = self.make_4D(data, vrbl=vrbl)
return data_4d
def DE_z(nc0, nc1, t, energy, lower, upper):
"""
Computation for difference kinetic energy (DKE).
Sums DKE over all levels between lower and upper,
for each grid point, and returns a 2D array.
Destaggering is not enabled as it introduces
computational cost that is of miniscule value considering
the magnitudes of output values.
Method finds levels nearest lower/upper hPa and sums between
them inclusively.
Inputs:
nc0 : netCDF file
nc1 : netCDF file
t : times index to difference
energy : kinetic or total
lower : lowest level, hPa
upper : highest level, hPa
Outputs:
data : 2D array.
"""
# Speed up script by only referencing data, not
# loading it to a variable yet
# WIND
U0 = nc0.variables['U'][t, ...]
U1 = nc1.variables['U'][t, ...]
Ud = U0 - U1
#del U0, U1
V0 = nc0.variables['V'][t, ...]
V1 = nc1.variables['V'][t, ...]
Vd = V0 - V1
#del V0, V1
# PERT and BASE PRESSURE
if lower or upper:
P0 = nc0.variables['P'][t, ...]
PB0 = nc0.variables['PB'][t, ...]
Pr = P0 + PB0
#del P0, PB1
# Here we assume pressure columns are
# roughly the same between the two...
if energy == 'DTE':
T0 = nc0.variables['T'][t, ...]
T1 = nc1.variables['T'][t, ...]
Td = T0 - T1
#del T0, T1
R = 287.0 # Universal gas constant (J / deg K * kg)
Cp = 1004.0 # Specific heat of dry air at constant pressure (J / deg K * kg)
kappa = R / Cp
xlen = Ud.shape[1] # 1 less than in V
ylen = Vd.shape[2] # 1 less than in U
zlen = Ud.shape[0] # identical in U & V
# Generator for lat/lon points
def latlon(nlats, nlons):
for i in range(nlats): # y-axis
for j in range(nlons): # x-axis
yield i, j
DKE = []
DKE2D = N.zeros((xlen, ylen))
print_time = ''.join((nc0.variables['Times'][t]))
print(("Calculating 3D grid for time {0}...".format(print_time)))
gridpts = latlon(xlen, ylen)
for gridpt in gridpts:
i, j = gridpt
# Find closest level to 'lower', 'upper'
if lower or upper:
P_col = Pr[:, j, i]
if lower:
low_idx = utils.closest(P_col, lower * 100.0)
else:
low_idx = None
if upper:
upp_idx = utils.closest(P_col, upper * 100.0) + 1
else:
upp_idx = None
zidx = slice(low_idx, upp_idx)
if energy == 'DKE':
DKE2D[j,
i] = N.sum(0.5 * ((Ud[zidx, j, i])**2 + (Vd[zidx, j, i])**2))
elif energy == 'DTE':
DKE2D[j, i] = N.sum(
0.5 * ((Ud[zidx, j, i])**2 + (Vd[zidx, j, i])**2 + kappa *
(Td[zidx, j, i])**2))
DKE.append(DKE2D)
return DKE
def interpolate(self,data,lats=None,lons=None,grid=None,method=2,
Nlim=None,Elim=None,Slim=None,Wlim=None):
""" Interpolate data and lat/lon grid to current grid.
Note:
User specifies grid or lat/lons.
Args:
Nlim, Elim, Slim, Wlim: Cut down data to this bounding box to
speed up intepolation.
"""
if lats is None:
lats = grid.lats
lons = grid.lons
data = utils.enforce_2d(data)
if Nlim or Elim or Wlim or Slim:
# og_data = data
# og_lats = lats
# og_lons = lons
assert Nlim and Elim and Slim and Wlim
Nidx = utils.closest(lats[:,0],Nlim)
Eidx = utils.closest(lons[0,:],Elim)
Sidx = utils.closest(lats[:,0],Slim)
Widx = utils.closest(lons[0,:],Wlim)
data = data[Sidx:Nidx+1,Widx:Eidx+1]
lats = lats[Sidx:Nidx+1,Widx:Eidx+1]
lons = lons[Sidx:Nidx+1,Widx:Eidx+1]
if method == 1:
# First cut before interpolating
# cut_data, cut_lats, cut_lons = self.cut(data=data,lats=lats,lons=lons)
# xx,yy = self.convert_latlon_xy(cut_lats,cut_lons)
xx,yy = self.convert_latlon_xy(lats,lons)
# data_reproj = reproject_tools.reproject(data_orig=cut_data,xx_orig=xx,
data_reproj = reproject_tools.reproject(data_orig=data,xx_orig=xx,
yy_orig=yy,xx_new=self.xx,yy_new=self.yy)
elif method == 2:
print("Geog limits done.")
from scipy.spatial import cKDTree
xog,yog,zog = self.lon_lat_to_cartesian(lons,lats)
xnew,ynew,znew = self.lon_lat_to_cartesian(self.lons,self.lats)
print("lonlat converted to cartesian.")
# xnew = self.xx
# ynew = self.yy
# znew = N.ones_like(xnew)
# pdb.set_trace()
# tree = cKDTree(zip(xog,yog,zog))
# tree = cKDTree((xog,yog,zog))
tree = cKDTree(N.c_[xog.ravel(),yog.ravel()])
print("Tree created.")
# d, inds = tree.query(zip(xnew,ynew,znew), k=10)
# d, inds = tree.query((xnew,ynew,znew), k=10)
d, inds = tree.query(list(zip(xnew,ynew)),k=10)
print("Query done.")
w = 1/(d**2)
# Inverse Distance Weighting
data_reproj = N.sum(w*data.flatten()[inds],axis=1)/N.sum(w,axis=1)
data_reproj.shape = self.shape
# pdb.set_trace()
else:
raise Exception
return data_reproj