From 424c5d1cb000377805413590f8fbf0d0c8a2f27a Mon Sep 17 00:00:00 2001 From: Kristen Thyng Date: Mon, 12 Feb 2024 10:36:28 -0800 Subject: [PATCH] changes to ROMS reader for wetdry masks Main changes are to allow for the land_binary_mask as well as mask_rho, mask_u, and mask_v to have 3D or 2D dimensions since masks in ROMS can vary in time if they represented cells that can wet and dry in time. With the changes, if a variable `wetdry_mask_rho` is available, it will be used for the rho mask in place of `mask_rho`, otherwise will use `mask_rho`, and the same for the u and v masks. The masks for different grids use the same indices, same as before these changes, but now they are managed a little more simply in one location in the function. Additionally, more effort is made to match variables to the appropriate mask since sometimes they may not be a typical ROMS- output variable. More changes: * masks are saved in a dict in case they are reused across different variables * the mask is applied to the variable array as a bool index directly to convert to nans. This is better than the previous approach of searching for the missing value and replacing that with nan because with wet/dry masks the dry cells are not always represented by a fill value and instead can only be found using the mask. * added some more variables to check for whether to rotate or not and added log statements too. --- opendrift/readers/reader_ROMS_native.py | 157 +++++++++++++++++------- 1 file changed, 110 insertions(+), 47 deletions(-) diff --git a/opendrift/readers/reader_ROMS_native.py b/opendrift/readers/reader_ROMS_native.py index 6c2b3c2cc..7766245fd 100644 --- a/opendrift/readers/reader_ROMS_native.py +++ b/opendrift/readers/reader_ROMS_native.py @@ -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'): @@ -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: @@ -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 @@ -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: @@ -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: