def get_variables(self,
requested_variables,
time=None,
x=None,
y=None,
z=None,
block=False):
start_time = datetime.now()
requested_variables, time, x, y, z, outside = self.check_arguments(
requested_variables, time, x, y, z)
# If one vector component is requested, but not the other
# we must add the other for correct rotation
for vector_pair in vector_pairs_xy:
if (vector_pair[0] in requested_variables
and vector_pair[1] not in requested_variables):
requested_variables.extend([vector_pair[1]])
if (vector_pair[1] in requested_variables
and vector_pair[0] not in requested_variables):
requested_variables.extend([vector_pair[0]])
nearestTime, dummy1, dummy2, indxTime, dummy3, dummy4 = \
self.nearest_time(time)
variables = {}
if z is None:
z = np.atleast_1d(0)
# Find horizontal indices corresponding to requested x and y
if hasattr(self, 'clipped'):
clipped = self.clipped
else:
clipped = 0
indx = np.floor((x - self.xmin) / self.delta_x).astype(int) + clipped
indy = np.floor((y - self.ymin) / self.delta_y).astype(int) + clipped
indx[outside] = 0 # To be masked later
indy[outside] = 0
indx_el = indx
indy_el = indy
if block is True:
# Adding buffer, to cover also future positions of elements
buffer = self.buffer
# Avoiding the last pixel in each dimension, since there are
# several grids which are shifted (rho, u, v, psi)
indx = np.arange(
np.max([0, indx.min() - buffer]),
np.min([indx.max() + buffer, self.lon.shape[1] - 1]))
indy = np.arange(
np.max([0, indy.min() - buffer]),
np.min([indy.max() + buffer, self.lon.shape[0] - 1]))
# Find depth levels covering all elements
if z.min() == 0 or not hasattr(self, 'hc'):
indz = self.num_layers - 1 # surface layer
variables['z'] = 0
else:
# Find the range of indices covering given z-values
if not hasattr(self, 'sea_floor_depth_below_sea_level'):
logging.debug('Reading sea floor depth...')
self.sea_floor_depth_below_sea_level = \
self.Dataset.variables['h'][:]
Htot = self.sea_floor_depth_below_sea_level
self.z_rho_tot = depth.sdepth(Htot, self.hc, self.Cs_r)
if has_xarray is False:
indxgrid, indygrid = np.meshgrid(indx, indy)
H = self.sea_floor_depth_below_sea_level[indygrid, indxgrid]
else:
H = self.sea_floor_depth_below_sea_level[indy, indx]
z_rho = depth.sdepth(H, self.hc, self.Cs_r)
# Element indices must be relative to extracted subset
indx_el = indx_el - indx.min()
indy_el = indy_el - indy.min()
# Loop to find the layers covering the requested z-values
indz_min = 0
indz_max = self.num_layers
for i in range(self.num_layers):
if np.min(z - z_rho[i, indy_el, indx_el]) > 0:
indz_min = i
if np.max(z - z_rho[i, indy_el, indx_el]) > 0:
indz_max = i
indz = range(
np.maximum(0, indz_min - self.verticalbuffer),
np.minimum(self.num_layers,
indz_max + 1 + self.verticalbuffer))
z_rho = z_rho[indz, :, :]
# Determine the z-levels to which to interpolate
zi1 = np.maximum(
0,
bisect_left(-np.array(self.zlevels), -z.max()) -
self.verticalbuffer)
zi2 = np.minimum(
len(self.zlevels),
bisect_right(-np.array(self.zlevels), -z.min()) +
self.verticalbuffer)
variables['z'] = np.array(self.zlevels[zi1:zi2])
#read_masks = {} # To store maskes for various grids
mask_values = {}
for par in requested_variables:
varname = [
name for name, cf in self.ROMS_variable_mapping.items()
if cf == par
]
var = self.Dataset.variables[varname[0]]
if par == 'land_binary_mask':
if not hasattr(self, 'land_binary_mask'):
# Read landmask for whole domain, for later re-use
self.land_binary_mask = \
1 - self.Dataset.variables['mask_rho'][:]
if has_xarray is False:
indxgrid, indygrid = np.meshgrid(indx, indy)
variables[par] = self.land_binary_mask[indygrid, indxgrid]
else:
variables[par] = self.land_binary_mask[indy, indx]
elif var.ndim == 2:
variables[par] = var[indy, indx]
elif var.ndim == 3:
variables[par] = var[indxTime, indy, indx]
elif var.ndim == 4:
variables[par] = var[indxTime, indz, indy, indx]
else:
raise Exception('Wrong dimension of variable: ' +
self.variable_mapping[par])
variables[par] = np.asarray(variables[par]) # If Xarray
start = datetime.now()
if par not in mask_values:
if has_xarray is False:
indxgrid, indygrid = np.meshgrid(indx, indy)
else:
indxgrid = indx
indygrid = indy
if par == 'x_sea_water_velocity':
if not hasattr(self, 'mask_u'):
self.mask_u = self.Dataset.variables['mask_u'][:]
mask = self.mask_u[indygrid, indxgrid]
elif par == 'y_sea_water_velocity':
if not hasattr(self, 'mask_v'):
self.mask_v = self.Dataset.variables['mask_v'][:]
mask = self.mask_v[indygrid, indxgrid]
else:
if not hasattr(self, 'mask_rho'):
# For ROMS-Agrif this must perhaps be mask_psi?
self.mask_rho = self.Dataset.variables['mask_rho'][:]
mask = self.mask_rho[indygrid, indxgrid]
if has_xarray is True:
mask = np.asarray(mask)
if mask.min() == 0 and par != 'land_binary_mask':
first_mask_point = np.where(mask.ravel() == 0)[0][0]
if variables[par].ndim == 3:
upper = variables[par][0, :, :]
else:
upper = variables[par]
mask_values[par] = upper.ravel()[first_mask_point]
variables[par][variables[par] == mask_values[par]] = np.nan
if var.ndim == 4:
# Regrid from sigma to z levels
if len(np.atleast_1d(indz)) > 1:
logging.debug('sigma to z for ' + varname[0])
if self.precalculate_s2z_coefficients is True:
M = self.sea_floor_depth_below_sea_level.shape[0]
N = self.sea_floor_depth_below_sea_level.shape[1]
O = len(self.z_rho_tot)
if not hasattr(self, 's2z_A'):
logging.info(
'Calculating sigma2z-coefficients for whole domain'
)
starttime = datetime.now()
dummyvar = np.ones((O, M, N))
dummy, self.s2z_total = depth.multi_zslice(
dummyvar, self.z_rho_tot, self.zlevels)
# Store arrays/coefficients
self.s2z_A = self.s2z_total[0].reshape(
len(self.zlevels), M, N)
self.s2z_C = self.s2z_total[1].reshape(
len(self.zlevels), M, N)
#self.s2z_I = self.s2z_total[2].reshape(M, N)
self.s2z_kmax = self.s2z_total[3]
del self.s2z_total # Free memory
logging.info('Time: ' +
str(datetime.now() - starttime))
if 'A' not in locals():
logging.info('Re-using sigma2z-coefficients')
# Select relevant subset of full arrays
zle = np.arange(zi1,
zi2) # The relevant depth levels
A = self.s2z_A.copy(
) # Awkward subsetting to prevent losing one dimension
A = A[:, :, indx]
A = A[:, indy, :]
A = A[zle, :, :]
C = self.s2z_C.copy()
C = C[:, :, indx]
C = C[:, indy, :]
C = C[zle, :, :]
C = C - C.max() + variables[par].shape[0] - 1
C[C < 1] = 1
A = A.reshape(len(zle), len(indx) * len(indy))
C = C.reshape(len(zle), len(indx) * len(indy))
I = np.arange(len(indx) * len(indy))
## Check
#dummyvar2, s2z = depth.multi_zslice(
# variables[par].copy(), z_rho.copy(), variables['z'])
#print len(zle), variables[par].shape, 'zle, varshape'
#Ac,Cc,Ic,kmaxc = s2z
#print C, 'C'
#print Cc, 'Cc'
#print C.shape, Cc.shape
#if C.max() != Cc.max():
# print 'WARNING!!'
# import sys; sys.exit('stop')
kmax = len(
zle
) # Must be checked. Or number of sigma-layers?
if 'A' not in locals():
logging.info('Calculating new sigma2z-coefficients')
variables[par], s2z = depth.multi_zslice(
variables[par], z_rho, variables['z'])
A, C, I, kmax = s2z
# Reshaping to compare with subset of full array
#zle = np.arange(zi1, zi2)
#A = A.reshape(len(zle), len(indx), len(indy))
#C = C.reshape(len(zle), len(indx), len(indy))
#I = I.reshape(len(indx), len(indy))
else:
logging.info('Applying sigma2z-coefficients')
# Re-using sigma2z koefficients:
F = np.asarray(variables[par])
Fshape = F.shape
N = F.shape[0]
M = F.size // N
F = F.reshape((N, M))
R = (1 - A) * F[(C - 1, I)] + A * F[(C, I)]
variables[par] = R.reshape((kmax, ) + Fshape[1:])
# Nan in input to multi_zslice gives extreme values in output
variables[par][variables[par] > 1e+9] = np.nan
# If 2D array is returned due to the fancy slicing methods
# of netcdf-python, we need to take the diagonal
if variables[par].ndim > 1 and block is False:
variables[par] = variables[par].diagonal()
# Mask values outside domain
variables[par] = np.ma.array(variables[par], ndmin=2, mask=False)
if block is False:
variables[par].mask[outside] = True
# Skipping de-staggering, as it leads to invalid values at later interpolation
#if block is True:
# # Unstagger grid for vectors
# logging.debug('Unstaggering ' + par)
# if 'eta_v' in var.dimensions:
# variables[par] = np.ma.array(variables[par],
# mask=variables[par].mask)
# variables[par][variables[par].mask] = 0
# if variables[par].ndim == 2:
# variables[par] = \
# (variables[par][0:-1,0:-1] +
# variables[par][0:-1,1::])/2
# elif variables[par].ndim == 3:
# variables[par] = \
# (variables[par][:,0:-1,0:-1] +
# variables[par][:,0:-1,1::])/2
# variables[par] = np.ma.masked_where(variables[par]==0,
# variables[par])
# elif 'eta_u' in var.dimensions:
# variables[par] = np.ma.array(variables[par],
# mask=variables[par].mask)
# variables[par][variables[par].mask] = 0
# if variables[par].ndim == 2:
# variables[par] = \
# (variables[par][0:-1,0:-1] +
# variables[par][1::,0:-1])/2
# elif variables[par].ndim == 3:
# variables[par] = \
# (variables[par][:,0:-1,0:-1] +
# variables[par][:,1::,0:-1])/2
# variables[par] = np.ma.masked_where(variables[par]==0,
# variables[par])
# else:
# if variables[par].ndim == 2:
# variables[par] = variables[par][1::, 1::]
# elif variables[par].ndim == 3:
# variables[par] = variables[par][:,1::, 1::]
if block is True:
# TODO: should be midpoints, but angle array below needs integer
#indx = indx[0:-1] # Only if de-staggering has been performed
#indy = indy[1::]
variables['x'] = indx
variables['y'] = indy
else:
variables['x'] = self.xmin + (indx - 1) * self.delta_x
variables['y'] = self.ymin + (indy - 1) * self.delta_y
variables['x'] = variables['x'].astype(np.float)
variables['y'] = variables['y'].astype(np.float)
variables['time'] = nearestTime
if 'x_sea_water_velocity' or 'sea_ice_x_velocity' \
or 'x_wind' in variables.keys():
# We must rotate current vectors
if not hasattr(self, 'angle_xi_east'):
logging.debug('Reading angle between xi and east...')
self.angle_xi_east = self.Dataset.variables['angle'][:]
if has_xarray is False:
rad = self.angle_xi_east[tuple(np.meshgrid(indy, indx))].T
else:
rad = self.angle_xi_east[indy, indx]
if 'x_sea_water_velocity' in variables.keys():
variables['x_sea_water_velocity'], \
variables['y_sea_water_velocity'] = rotate_vectors_angle(
variables['x_sea_water_velocity'],
variables['y_sea_water_velocity'], rad)
if 'sea_ice_x_velocity' in variables.keys():
variables['sea_ice_x_velocity'], \
variables['sea_ice_y_velocity'] = rotate_vectors_angle(
variables['sea_ice_x_velocity'],
variables['sea_ice_y_velocity'], rad)
if 'x_wind' in variables.keys():
variables['x_wind'], \
variables['y_wind'] = rotate_vectors_angle(
variables['x_wind'],
variables['y_wind'], rad)
# Masking NaN
for var in requested_variables:
variables[var] = np.ma.masked_invalid(variables[var])
logging.debug('Time for ROMS native reader: ' +
str(datetime.now() - start_time))
return variables
def get_variables(self,
requested_variables,
time=None,
x=None,
y=None,
z=None,
block=False):
requested_variables, time, x, y, z, outside = self.check_arguments(
requested_variables, time, x, y, z)
# If one vector component is requested, but not the other
# we must add the other for correct rotation
for vector_pair in vector_pairs_xy:
if (vector_pair[0] in requested_variables
and vector_pair[1] not in requested_variables):
requested_variables.extend([vector_pair[1]])
if (vector_pair[1] in requested_variables
and vector_pair[0] not in requested_variables):
requested_variables.extend([vector_pair[0]])
nearestTime, dummy1, dummy2, indxTime, dummy3, dummy4 = \
self.nearest_time(time)
variables = {}
if z is None:
z = np.atleast_1d(0)
# Find horizontal indices corresponding to requested x and y
if hasattr(self, 'clipped'):
clipped = self.clipped
else:
clipped = 0
indx = np.floor((x - self.xmin) / self.delta_x).astype(int) + clipped
indy = np.floor((y - self.ymin) / self.delta_y).astype(int) + clipped
indx[outside] = 0 # To be masked later
indy[outside] = 0
indx_el = indx
indy_el = indy
if block is True:
# Adding buffer, to cover also future positions of elements
buffer = self.buffer
# Avoiding the last pixel in each dimension, since there are
# several grids which are shifted (rho, u, v, psi)
indx = np.arange(
np.max([0, indx.min() - buffer]),
np.min([indx.max() + buffer, self.lon.shape[1] - 1]))
indy = np.arange(
np.max([0, indy.min() - buffer]),
np.min([indy.max() + buffer, self.lon.shape[0] - 1]))
# Find depth levels covering all elements
if z.min() == 0 or not hasattr(self, 'hc'):
indz = self.num_layers - 1 # surface layer
variables['z'] = 0
else:
# Find the range of indices covering given z-values
if not hasattr(self, 'sea_floor_depth_below_sea_level'):
logging.debug('Reading sea floor depth...')
self.sea_floor_depth_below_sea_level = \
self.Dataset.variables['h'][:]
indxgrid, indygrid = np.meshgrid(indx, indy)
H = self.sea_floor_depth_below_sea_level[indygrid, indxgrid]
z_rho = depth.sdepth(H, self.hc, self.Cs_r)
# Element indices must be relative to extracted subset
indx_el = indx_el - indx.min()
indy_el = indy_el - indy.min()
# Loop to find the layers covering the requested z-values
indz_min = 0
indz_max = self.num_layers
for i in range(self.num_layers):
if np.min(z - z_rho[i, indy_el, indx_el]) > 0:
indz_min = i
if np.max(z - z_rho[i, indy_el, indx_el]) > 0:
indz_max = i
indz = range(
np.maximum(0, indz_min - self.verticalbuffer),
np.minimum(self.num_layers,
indz_max + 1 + self.verticalbuffer))
z_rho = z_rho[indz, :, :]
# Determine the z-levels to which to interpolate
zi1 = np.maximum(
0,
bisect_left(-np.array(self.zlevels), -z.max()) -
self.verticalbuffer)
zi2 = np.minimum(
len(self.zlevels),
bisect_right(-np.array(self.zlevels), -z.min()) +
self.verticalbuffer)
variables['z'] = np.array(self.zlevels[zi1:zi2])
#read_masks = {} # To store maskes for various grids
for par in requested_variables:
varname = [
name for name, cf in self.ROMS_variable_mapping.items()
if cf == par
]
var = self.Dataset.variables[varname[0]]
# Automatic masking may lead to trouble for ROMS files
# with valid_min/max, _Fill_value or missing_value
# https://github.com/Unidata/netcdf4-python/issues/703
var.set_auto_maskandscale(False)
try:
FillValue = getattr(var, '_FillValue')
except:
FillValue = None
try:
scale = getattr(var, 'scale_factor')
except:
scale = 1
try:
offset = getattr(var, 'add_offset')
except:
offset = 0
if var.ndim == 2:
variables[par] = var[indy, indx]
elif var.ndim == 3:
variables[par] = var[indxTime, indy, indx]
elif var.ndim == 4:
variables[par] = var[indxTime, indz, indy, indx]
else:
raise Exception('Wrong dimension of variable: ' +
self.variable_mapping[par])
# Manual scaling, offsetting and masking due to issue with ROMS files
logging.debug(
'Manually masking %s, FillValue %s, scale %s, offset %s' %
(par, FillValue, scale, offset))
if FillValue is not None:
if var.dtype != FillValue.dtype:
mask = variables[par] == 0
if not 'already_warned' in locals():
logging.warning(
'Data type of variable (%s) and _FillValue (%s) is not the same. Masking 0-values instead'
% (var.dtype, FillValue.dtype))
already_warned = True
else:
logging.warning('Masking ' + str(FillValue))
mask = variables[par] == FillValue
variables[par] = variables[par] * scale + offset
if FillValue is not None:
variables[par][mask] = np.nan
if var.ndim == 4:
# Regrid from sigma to z levels
if len(np.atleast_1d(indz)) > 1:
logging.debug('sigma to z for ' + varname[0])
variables[par] = depth.multi_zslice(
variables[par], z_rho, variables['z'])
# Nan in input to multi_zslice gives extreme values in output
variables[par][variables[par] > 1e+9] = np.nan
# If 2D array is returned due to the fancy slicing methods
# of netcdf-python, we need to take the diagonal
if variables[par].ndim > 1 and block is False:
variables[par] = variables[par].diagonal()
# Mask values outside domain
variables[par] = np.ma.array(variables[par], ndmin=2, mask=False)
if block is False:
variables[par].mask[outside] = True
# Skipping de-staggering, as it leads to invalid values at later interpolation
#if block is True:
# # Unstagger grid for vectors
# logging.debug('Unstaggering ' + par)
# if 'eta_v' in var.dimensions:
# variables[par] = np.ma.array(variables[par],
# mask=variables[par].mask)
# variables[par][variables[par].mask] = 0
# if variables[par].ndim == 2:
# variables[par] = \
# (variables[par][0:-1,0:-1] +
# variables[par][0:-1,1::])/2
# elif variables[par].ndim == 3:
# variables[par] = \
# (variables[par][:,0:-1,0:-1] +
# variables[par][:,0:-1,1::])/2
# variables[par] = np.ma.masked_where(variables[par]==0,
# variables[par])
# elif 'eta_u' in var.dimensions:
# variables[par] = np.ma.array(variables[par],
# mask=variables[par].mask)
# variables[par][variables[par].mask] = 0
# if variables[par].ndim == 2:
# variables[par] = \
# (variables[par][0:-1,0:-1] +
# variables[par][1::,0:-1])/2
# elif variables[par].ndim == 3:
# variables[par] = \
# (variables[par][:,0:-1,0:-1] +
# variables[par][:,1::,0:-1])/2
# variables[par] = np.ma.masked_where(variables[par]==0,
# variables[par])
# else:
# if variables[par].ndim == 2:
# variables[par] = variables[par][1::, 1::]
# elif variables[par].ndim == 3:
# variables[par] = variables[par][:,1::, 1::]
if block is True:
# TODO: should be midpoints, but angle array below needs integer
#indx = indx[0:-1] # Only if de-staggering has been performed
#indy = indy[1::]
variables['x'] = indx
variables['y'] = indy
else:
variables['x'] = self.xmin + (indx - 1) * self.delta_x
variables['y'] = self.ymin + (indy - 1) * self.delta_y
variables['x'] = variables['x'].astype(np.float)
variables['y'] = variables['y'].astype(np.float)
variables['time'] = nearestTime
if 'x_sea_water_velocity' or 'sea_ice_x_velocity' \
or 'x_wind' in variables.keys():
# We must rotate current vectors
if not hasattr(self, 'angle_xi_east'):
logging.debug('Reading angle between xi and east...')
self.angle_xi_east = self.Dataset.variables['angle'][:]
rad = self.angle_xi_east[np.meshgrid(indy, indx)].T
if 'x_sea_water_velocity' in variables.keys():
variables['x_sea_water_velocity'], \
variables['y_sea_water_velocity'] = rotate_vectors_angle(
variables['x_sea_water_velocity'],
variables['y_sea_water_velocity'], rad)
if 'sea_ice_x_velocity' in variables.keys():
variables['sea_ice_x_velocity'], \
variables['sea_ice_y_velocity'] = rotate_vectors_angle(
variables['sea_ice_x_velocity'],
variables['sea_ice_y_velocity'], rad)
if 'x_wind' in variables.keys():
variables['x_wind'], \
variables['y_wind'] = rotate_vectors_angle(
variables['x_wind'],
variables['y_wind'], rad)
if 'land_binary_mask' in requested_variables:
variables['land_binary_mask'] = \
1 - variables['land_binary_mask']
# Masking NaN
for var in requested_variables:
variables[var] = np.ma.masked_invalid(variables[var])
return variables