Skip to content

Commit

Permalink
Merge pull request #1225 from kthyng/wetdry_roms
Browse files Browse the repository at this point in the history 
changes to ROMS reader for wetdry masks
  • Loading branch information
@knutfrode
knutfrode authored Feb 12, 2024
2 parents 954dc14 + 424c5d1 commit f07a1ed
Showing 1 changed file with 110 additions and 47 deletions.
157 changes: 110 additions & 47 deletions opendrift/readers/reader_ROMS_native.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -286,15 +286,78 @@ def drop_non_essential_vars_pop(ds):
# Run constructor of parent Reader class
super(Reader, self).__init__()

@property
def mask_rho(self):
"""Mask for the rho-points.
Uses wetdry_mask_rho (which should be 3D) if available, otherwise mask_rho (2D).
If this mask is 2D, read it in this one time and use going forward in simulation. If 3D,
will read in parts of the mask each loop.
"""
if not hasattr(self, '_mask_rho'):
if 'wetdry_mask_rho' in self.Dataset.data_vars:
self._mask_rho = self.Dataset.variables['wetdry_mask_rho']
logger.info("Using wetdry_mask_rho for mask_rho")
else:
# Read landmask for whole domain, for later re-use
self._mask_rho = self.Dataset.variables['mask_rho'][:]
# load in once if static mask
if hasattr(self._mask_rho, "chunks"): # to see if dask array
self._mask_rho = self._mask_rho.compute()
logger.info("Using mask_rho for mask_rho")
return self._mask_rho

@property
def mask_u(self):
"""Mask for the u-points.
Uses wetdry_mask_u (which should be 3D) if available, otherwise mask_u (2D).
If this mask is 2D, read it in this one time and use going forward in simulation. If 3D,
will read in parts of the mask each loop.
"""
if not hasattr(self, '_mask_u'):
if 'wetdry_mask_u' in self.Dataset.data_vars:
self._mask_u = self.Dataset.variables['wetdry_mask_u']
logger.info("Using wetdry_mask_u for mask_u")
else:
# Read landmask for whole domain, for later re-use
self._mask_u = self.Dataset.variables['mask_u'][:]
# load in once if static mask
if hasattr(self._mask_u, "chunks"): # to see if dask array
self._mask_u = self._mask_u.compute()
logger.info("Using mask_u for mask_u")
return self._mask_u

@property
def mask_v(self):
"""Mask for the v-points.
Uses wetdry_mask_v (which should be 3D) if available, otherwise mask_v (2D).
If this mask is 2D, read it in this one time and use going forward in simulation. If 3D,
will read in parts of the mask each loop.
"""
if not hasattr(self, '_mask_v'):
if 'wetdry_mask_v' in self.Dataset.data_vars:
self._mask_v = self.Dataset.variables['wetdry_mask_v']
logger.info("Using wetdry_mask_v for mask_v")
else:
# Read landmask for whole domain, for later re-use
self._mask_v = self.Dataset.variables['mask_v'][:]
# load in once if static mask
if hasattr(self._mask_v, "chunks"): # to see if dask array
self._mask_v = self._mask_v.compute()
logger.info("Using mask_v for mask_v")
return self._mask_v

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)

# land_binary_mask should be based on the rho grid
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'][:]
self.land_binary_mask = 1 - self.mask_rho

if 'sea_floor_depth_below_sea_level' in requested_variables and not hasattr(
self, 'sea_floor_depth_below_sea_level'):
Expand Down Expand Up @@ -380,16 +443,29 @@ def get_variables(self, requested_variables, time=None,
bisect_right(-np.array(self.zlevels),
-z.min()) + self.verticalbuffer)
variables['z'] = np.array(self.zlevels[zi1:zi2])

# use these same indices for all mask subsetting since if one is
# 3D they all should be
if self.mask_rho.ndim == 2:
imask = (indy,indx)
elif self.mask_rho.ndim == 3:
imask = (indxTime,indy,indx)

def get_mask(mask_name, imask):
if mask_name in masks_for_loop:
mask = masks_for_loop[mask_name]
else:
mask = getattr(self, mask_name)[imask]
return mask, mask_name

#read_masks = {} # To store maskes for various grids
mask_values = {}
masks_for_loop = {} # 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]]

if par == 'land_binary_mask':
variables[par] = self.land_binary_mask[indy, indx]
variables[par] = self.land_binary_mask[imask]
elif par == 'sea_floor_depth_below_sea_level':
variables[par] = self.sea_floor_depth_below_sea_level[indy, indx]
elif var.ndim == 2:
Expand All @@ -402,49 +478,30 @@ def get_variables(self, requested_variables, time=None,
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]
# If Xarray, load in so can be used in future loop iterations too
variables[par] = np.asarray(variables[par])

if par != 'land_binary_mask':

# make sure that var has matching horizontal dimensions with the mask
# make sure coord names also match
if self.mask_rho.shape[-2:] == var.shape[-2:] and set(self.mask_rho.dims).issubset(var.dims):
mask, mask_name = get_mask("mask_rho", imask)
elif self.mask_u.shape[-2:] == var.shape[-2:] and set(self.mask_u.dims).issubset(var.dims):
mask, mask_name = get_mask("mask_u", imask)
elif self.mask_v.shape[-2:] == var.shape[-2:] and set(self.mask_v.dims).issubset(var.dims):
mask, mask_name = get_mask("mask_v", imask)
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]
raise Exception('No mask found for ' + par)

masks_for_loop[mask_name] = np.asarray(mask)
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 mask.min() == 0:
# Not using the fill value directly allows the mask to have more control
# which is necessary when using a wetdry mask that changes in time
# since the fill value will not cover all masked locations then.
variables[par][...,mask==0] = np.nan

if var.ndim == 4:
# Regrid from sigma to z levels
Expand Down Expand Up @@ -572,7 +629,10 @@ def get_variables(self, requested_variables, time=None,

if 'x_sea_water_velocity' in variables.keys() or \
'sea_ice_x_velocity' in variables.keys() or \
'x_wind' in variables.keys():
'x_wind' in variables.keys() and \
'x_sea_water_velocity' not in self.do_not_rotate and \
'sea_ice_x_velocity' not in self.do_not_rotate and \
'x_wind' not in self.do_not_rotate:
# We must rotate current vectors
if not hasattr(self, 'angle_xi_east'):
if 'angle' in self.Dataset.variables:
Expand All @@ -589,18 +649,21 @@ def get_variables(self, requested_variables, time=None,
variables['y_sea_water_velocity'] = rotate_vectors_angle(
variables['x_sea_water_velocity'],
variables['y_sea_water_velocity'], rad)
logger.debug('Rotated x_sea_water_velocity and y_sea_water_velocity')
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)
logger.debug('Rotated sea_ice_x_velocity and sea_ice_y_velocity')
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)
logger.debug('Rotated x_wind and y_wind')

# Masking NaN
for var in requested_variables:
Expand Down

0 comments on commit f07a1ed

Please sign in to comment.