Source code for opendrift.readers.reader_ROMS_native

# This file is part of OpenDrift.
#
# OpenDrift is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 2
#
# OpenDrift is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with OpenDrift.  If not, see <https://www.gnu.org/licenses/>.
#
# Copyright 2015, Knut-Frode Dagestad, MET Norway

from bisect import bisect_left, bisect_right
from datetime import datetime
import logging
logger = logging.getLogger(__name__)

import numpy as np
from netCDF4 import num2date
import xarray as xr

from opendrift.readers.basereader import BaseReader, vector_pairs_xy, StructuredReader
from opendrift.readers.roppy import depth


[docs] class Reader(BaseReader, StructuredReader): def __init__(self, filename=None, name=None, gridfile=None, standard_name_mapping={}): if filename is None: raise ValueError('Need filename as argument to constructor') # Map ROMS variable names to CF standard_name self.ROMS_variable_mapping = { # Removing (temoprarily) land_binary_mask from ROMS-variables, # as this leads to trouble with linearNDFast interpolation 'mask_rho': 'land_binary_mask', 'mask_psi': 'land_binary_mask', 'h': 'sea_floor_depth_below_sea_level', 'zeta': 'sea_surface_height', 'u': 'x_sea_water_velocity', 'v': 'y_sea_water_velocity', #'u_eastward': 'x_sea_water_velocity', # these are wrognly rotated below #'v_northward': 'y_sea_water_velocity', 'w': 'upward_sea_water_velocity', 'temp': 'sea_water_temperature', 'salt': 'sea_water_salinity', 'uice': 'sea_ice_x_velocity', 'vice': 'sea_ice_y_velocity', 'aice': 'sea_ice_area_fraction', 'hice': 'sea_ice_thickness', 'gls': 'turbulent_generic_length_scale', 'tke': 'turbulent_kinetic_energy', 'AKs': 'ocean_vertical_diffusivity', 'sustr': 'surface_downward_x_stress', 'svstr': 'surface_downward_y_stress', 'tair': 'air_temperature', 'wspd': 'wind_speed', 'uwnd': 'x_wind', 'vwnd': 'y_wind', 'uwind': 'x_wind', 'vwind': 'y_wind', 'Uwind': 'x_wind', 'Vwind': 'y_wind'} # Add user provided variable mappings self.ROMS_variable_mapping.update(standard_name_mapping) # z-levels to which sigma-layers may be interpolated self.zlevels = np.array([ 0, -.5, -1, -3, -5, -10, -25, -50, -75, -100, -150, -200, -250, -300, -400, -500, -600, -700, -800, -900, -1000, -1500, -2000, -2500, -3000, -3500, -4000, -4500, -5000, -5500, -6000, -6500, -7000, -7500, -8000]) gls_param = ['gls_cmu0', 'gls_p', 'gls_m', 'gls_n'] filestr = str(filename) if name is None: self.name = filestr else: self.name = name try: # Open file, check that everything is ok logger.info('Opening dataset: ' + filestr) if ('*' in filestr) or ('?' in filestr) or ('[' in filestr): logger.info('Opening files with MFDataset') def drop_non_essential_vars_pop(ds): dropvars = [v for v in ds.variables if v not in list(self.ROMS_variable_mapping.keys()) + gls_param + ['ocean_time', 'time', 'bulk_time', 's_rho', 'Cs_r', 'hc', 'angle', 'Vtransform'] and v[0:3] not in ['lon', 'lat', 'mas']] logger.debug('Dropping variables: %s' % dropvars) ds = ds.drop_vars(dropvars) return ds self.Dataset = xr.open_mfdataset(filename, chunks={'ocean_time': 1}, compat='override', decode_times=False, preprocess=drop_non_essential_vars_pop, data_vars='minimal', coords='minimal') else: logger.info('Opening file with Dataset') self.Dataset = xr.open_dataset(filename, decode_times=False) except Exception as e: raise ValueError(e) if gridfile is not None: # Merging gridfile dataset with main dataset gf = xr.open_dataset(gridfile) self.Dataset = xr.merge([self.Dataset, gf]) if 'Vtransform' in self.Dataset.variables: self.Vtransform = self.Dataset.variables['Vtransform'].data # scalar else: logger.warning('Vtransform not found, using 1') self.Vtransform = 1 self.Vtransform = np.asarray(self.Vtransform) if 's_rho' not in self.Dataset.variables: dimensions = 2 else: dimensions = 3 if dimensions == 3: # Read sigma-coordinate values try: self.sigma = self.Dataset.variables['s_rho'][:] self.sigma = self.sigma.data # Extract data except: num_sigma = len(self.Dataset.dimensions['s_rho']) logger.warning( 's_rho not available in dataset, constructing from' ' number of layers (%s).' % num_sigma) self.sigma = (np.arange(num_sigma)+.5-num_sigma)/num_sigma # Read sigma-coordinate transform parameters try: self.Dataset.variables['Cs_r'].set_auto_mask(False) except: pass self.Cs_r = self.Dataset.variables['Cs_r'][:] try: self.hc = self.Dataset.variables['hc'][:] except: self.hc = self.Dataset.variables['hc'].data # scalar self.num_layers = len(self.sigma) else: self.num_layers = 1 self.ROMS_variable_mapping['ubar'] = 'x_sea_water_velocity' self.ROMS_variable_mapping['vbar'] = 'y_sea_water_velocity' del self.ROMS_variable_mapping['u'] del self.ROMS_variable_mapping['v'] if 'lat_rho' in self.Dataset.variables: # Horizontal oordinates and directions self.lat = self.Dataset.variables['lat_rho'][:] self.lon = self.Dataset.variables['lon_rho'][:] self.lon = self.lon.data # Extract, could be avoided downstream self.lat = self.lat.data if self.lat.ndim == 1: self.lon, self.lat = np.meshgrid(self.lon, self.lat) self.angle_xi_east = 0 else: raise ValueError(filename + ' does not contain lon/lat ' 'arrays, please supply a grid-file: "gridfile=<grid_file>"') for var in list(self.ROMS_variable_mapping): # Remove unused variables if var not in self.Dataset.variables: del self.ROMS_variable_mapping[var] try: # Check for GLS parameters (diffusivity) self.gls_parameters = {} for gls_par in gls_param: self.gls_parameters[gls_par] = \ self.Dataset.variables[gls_par][()] logger.info('Read GLS parameters from file.') except Exception as e: logger.info(e) logger.info('Did not find complete set of GLS parameters') # Get time coverage ocean_time = None for tv in ['ocean_time', 'time', 'bulk_time']: if tv in self.Dataset.variables: ocean_time = self.Dataset.variables[tv] if ocean_time is None: raise ValueError('Time variable not found in ROMS file') time_units = ocean_time.attrs['units'] if time_units == 'second': logger.info('Ocean time given as seconds relative to start ' 'Setting artifical start time of 1 Jan 2000.') time_units = 'seconds since 2000-01-01 00:00:00' if time_units == 'days': logger.info('Ocean time given as days relative to start ' 'Setting artifical start time of 1 Jan 2000.') time_units = 'days since 2000-01-01 00:00:00' self.times = num2date(ocean_time[:], time_units) self.start_time = self.times[0] self.end_time = self.times[-1] if len(self.times) > 1: self.time_step = self.times[1] - self.times[0] else: self.time_step = None self.name = 'roms native' self.precalculate_s2z_coefficients = True # Find all variables having standard_name self.variables = [] for var_name in list(self.Dataset.variables): var = self.Dataset.variables[var_name] if 'standard_name' in var.attrs and var_name not in self.ROMS_variable_mapping.keys(): self.ROMS_variable_mapping[var_name] = var.attrs['standard_name'] if var_name in self.ROMS_variable_mapping.keys(): self.variables.append(self.ROMS_variable_mapping[var_name]) # A bit hackish solution: # If variable names or their standard_name contain "east" or "north", # these should not be rotated from xi-direction to east-direction self.do_not_rotate = [] for var, stdname in self.ROMS_variable_mapping.items(): if 'east' in var.lower() or 'east' in stdname.lower() or \ 'north' in var.lower() or 'north' in stdname.lower(): self.do_not_rotate.append(stdname) if len(self.do_not_rotate)>0: logger.debug('The following ROMS vectors are considered east-north, and will not be rotated %s' % self.do_not_rotate) # Run constructor of parent Reader class super(Reader, self).__init__()
[docs] def get_variables(self, requested_variables, time=None, x=None, y=None, z=None): start_time = datetime.now() requested_variables, time, x, y, z, outside = self.check_arguments( requested_variables, time, x, y, z) if 'land_binary_mask' in requested_variables and 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 'sea_floor_depth_below_sea_level' in requested_variables and not hasattr( self, 'sea_floor_depth_below_sea_level'): self.sea_floor_depth_below_sea_level = self.Dataset.variables['h'][:] # 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.copy() indy_el = indy.copy() # 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'): logger.debug('Reading sea floor depth...') self.sea_floor_depth_below_sea_level = \ self.Dataset.variables['h'][:] if not hasattr(self, 'z_rho_tot'): Htot = self.sea_floor_depth_below_sea_level self.z_rho_tot = depth.sdepth(Htot, self.hc, self.Cs_r, Vtransform=self.Vtransform) H = self.sea_floor_depth_below_sea_level[indy, indx] z_rho = depth.sdepth(H, self.hc, self.Cs_r, Vtransform=self.Vtransform) # Element indices must be relative to extracted subset indx_el = np.clip(indx_el - indx.min(), 0, z_rho.shape[2]-1) indy_el = np.clip(indy_el - indy.min(), 0, z_rho.shape[1]-1) # 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': variables[par] = self.land_binary_mask[indy, indx] elif par == 'sea_floor_depth_below_sea_level': variables[par] = self.sea_floor_depth_below_sea_level[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: indxgrid = indx indygrid = indy if par in ['x_sea_water_velocity', 'sea_water_x_velocity', 'eastward_sea_water_velocity']: if not hasattr(self, 'mask_u'): if 'mask_u' in self.Dataset.variables: self.mask_u = self.Dataset.variables['mask_u'][:] elif 'mask_rho' in self.Dataset.variables: self.mask_u = self.Dataset.variables['mask_rho'][:] else: continue mask = self.mask_u[indygrid, indxgrid] elif par in ['y_sea_water_velocity', 'sea_water_y_velocity', 'northward_sea_water_velocity']: if not hasattr(self, 'mask_v'): if 'mask_v' in self.Dataset.variables: self.mask_v = self.Dataset.variables['mask_v'][:] elif 'mask_rho' in self.Dataset.variables: self.mask_v = self.Dataset.variables['mask_rho'][:] else: continue mask = self.mask_v[indygrid, indxgrid] else: if not hasattr(self, 'land_binary_mask'): # For ROMS-Agrif this must perhaps be mask_psi? if 'mask_rho' in self.Dataset.variables: self.land_binary_mask = 1 - self.Dataset.variables['mask_rho'][:] elif 'mask_psi' in self.Dataset.variables: self.land_binary_mask = 1 - self.Dataset.variables['mask_psi'][:] mask = 1 - self.land_binary_mask[indygrid, indxgrid] 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: logger.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'): logger.debug('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 logger.info('Time: ' + str(datetime.now() - starttime)) if 'A' not in locals(): logger.debug('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(): logger.debug('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: logger.debug('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 # Mask values outside domain variables[par] = np.ma.array(variables[par], ndmin=2, mask=False) # Skipping de-staggering, as it leads to invalid values at later interpolation #if block is True: # # Unstagger grid for vectors # logger.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::] # 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 variables['x'] = variables['x'].astype(np.float32) variables['y'] = variables['y'].astype(np.float32) variables['time'] = nearestTime if 'x_sea_water_velocity' in variables.keys() or \ 'sea_ice_x_velocity' in variables.keys() or \ 'x_wind' in variables.keys(): # We must rotate current vectors if not hasattr(self, 'angle_xi_east'): if 'angle' in self.Dataset.variables: logger.debug('Reading angle between xi and east...') self.angle_xi_east = self.Dataset.variables['angle'][:] if isinstance(self.angle_xi_east, int): rad = self.angle_xi_east else: rad = self.angle_xi_east[indy, indx] rad = np.ma.asarray(rad) if 'x_sea_water_velocity' in variables.keys() and \ 'x_sea_water_velocity' not in self.do_not_rotate: 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() and \ 'sea_ice_x_velocity' not in self.do_not_rotate: 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() and \ 'x_wind' not in self.do_not_rotate: 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]) logger.debug('Time for ROMS native reader: ' + str(datetime.now()-start_time)) return variables
[docs] def rotate_vectors_angle(u, v, radians): u2 = u*np.cos(radians) - v*np.sin(radians) v2 = u*np.sin(radians) + v*np.cos(radians) return u2, v2