Source code for salishsea_tools.grid_tools

# Copyright 2013-2021 The Salish Sea MEOPAR contributors
# and The University of British Columbia

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" Functions for calculating time-dependent scale factors and depth.

Functions developed/tested in this notebook:
http://nbviewer.jupyter.org/urls/bitbucket.org/salishsea/analysis-nancy/raw/tip/notebooks/energy_flux/Time-dependent%20depth%20and%20scale%20factors%20-%20development.ipynb

Examples of plotting with output:
http://nbviewer.jupyter.org/urls/bitbucket.org/salishsea/analysis-nancy/raw/tip/notebooks/energy_flux/Plotting%20with%20time%20dependent%20depths.ipynb

NKS nsoontie@eos.ubc.ca 08-2016
"""

from salishsea_tools import geo_tools
from tqdm import tqdm
import numpy as np
import netCDF4 as nc
import scipy.interpolate as spi
import scipy.sparse as sp

__all__ = [
    'calculate_H',
    'calculate_adjustment_factor', 'calculate_time_dependent_grid',
    'time_dependent_grid_U', 'time_dependent_grid_V', 'build_GEM_mask',
    'build_matrix', 'use_matrix',
]


[docs] def calculate_H(e3t0, tmask): """Calculate the initial water column thickness (H). :arg e3t0: initial vertical scale factors on T-grid. Dimensions: (depth, y, x). :type e3t0: :py:class:`numpy.ndarray` :arg tmask: T-grid mask. Dimensions: (depth, y, x) :type tmask: :py:class:`numpy.ndarray` :returns: the initial water column thickness. Dimensions: (y, x) """ H = np.sum(e3t0*tmask, axis=0) return H
[docs] def calculate_adjustment_factor(H, ssh): """Calculate the time-dependent adjustment factor for variable volume in NEMO. adj = (1+ssh/H) and e3t_t = e3t_0*adj :arg H: Water column thickness. Dimension: (y, x) :type H: :py:class:`numpy.array` :arg ssh: the model sea surface height. Dimensions: (time, y, x) :type ssh: :py:class:`numpy.ndarray` :returns: the adjustment factor with dimensions (time, y, x) """ with np.errstate(divide='ignore', invalid='ignore'): one_over_H = 1 / H one_over_H = np.nan_to_num(one_over_H) adj = (1 + ssh * one_over_H) return adj
[docs] def calculate_time_dependent_grid( e3t0, tmask, ssh, input_vars, ): """ Calculate the time dependent vertical grids and scale factors for variable volume in NEMO. :arg e3t0: initial vertical scale factors on T-grid. Dimensions: (depth, y, x). :type e3t0: :py:class:`numpy.ndarray` :arg tmask: T-grid mask. Dimensions: (depth, y, x) :type tmask: :py:class:`numpy.ndarray` :arg ssh: the model sea surface height. Dimensions: (time, y, x) :type ssh: :py:class:`numpy.ndarray` :arg input_vars: A dictionary of the initial grids/scale factors to be translated into time_dependent. Example keys can be 'e3t_0', 'gdept_0', 'e3w_0', 'gdepw_0'. A dictionary with corresponding time-dependent grids is returned, where the keys are now 'e3t_t', 'gdept_t', 'e3w_0', 'gdepw_0'. :typ input_vars: dictionary :type return_vars: list of strings :returns: A dictionary containing the desired time dependent vertical scale factors on t and w grids and depths on t and w grids. Dimensions of each: (time, depth, y, x) """ # adjustment factors H = calculate_H(e3t0, tmask) adj = calculate_adjustment_factor(H, ssh) adj = np.expand_dims(adj, axis=1) # expand to give depth dimension # Time-dependent grids return_vars = {} for key in input_vars: return_key = '{}t'.format(key[0:-1]) return_vars[return_key] = input_vars[key] * adj return return_vars
[docs] def time_dependent_grid_U(e3u0, e1u, e2u, e1t, e2t, umask, ssh, input_vars, return_ssh=False): """Calculate time-dependent vertical grid spacing and depths on U-grid for variable volume in NEMO. :arg e3u0: initial vertical scale factors on U-grid. Dimensions: (depth, y, x). :type e3u0: :py:class:`numpy.ndarray` :arg e1u: x-direction scale factors on U-grid. Dimensions: (y, x). :type e1u: :py:class:`numpy.ndarray` :arg e2u: y-direction scale factors on U-grid. Dimensions: (y, x). :type e2u: :py:class:`numpy.ndarray` :arg e1t: x-direction scale factors on T-grid. Dimensions: (y, x). :type e1t: :py:class:`numpy.ndarray` :arg e2t: y-direction scale factors on T-grid. Dimensions: (y, x). :type e2t: :py:class:`numpy.ndarray` :arg umask: U-grid mask. Dimensions: (depth, y, x) :type umask: :py:class:`numpy.ndarray` :arg ssh: the model sea surface height on T-grid. Dimensions: (time, y, x) :type ssh: :py:class:`numpy.ndarray` :arg input_vars: A dictionary of the initial grids/scale factors to be translated into time-dependent. Example keys can be 'e3u_0', 'gdepu_0'. A dictionary with corresponding time-dependent grids is returned, where the keys are now 'e3u_t', 'gdepu_t'. :type input_vars: dictionary :arg return_ssh: Specifies whether the ssh interpolated to the U-grid should be returned (ssh_u) :type return_ssh: boolean :returns: A dictionary containing the desired time dependent vertical scale factors on depths on u grid. Dimensions of each: (time, depth, y, x). If returned, ssh_u has dimensions (time, y, x)""" ssh_u = np.zeros_like(ssh) e1e2u = e1u * e2u e1e2t = e1t * e2t # Interpolate ssh to u grid ssh_u[:, :, 0:-1] = 0.5 * umask[0, :, 0:-1] / e1e2u[:, 0:-1] * ( e1e2t[:, 0:-1] * ssh[:, :, 0:-1] + e1e2t[:, 1:] * ssh[:, :, 1:] ) H = calculate_H(e3u0, umask) adj = calculate_adjustment_factor(H, ssh_u) adj = np.expand_dims(adj, axis=1) # Time-dependent grids return_vars = {} for key in input_vars: return_key = '{}t'.format(key[0:-1]) return_vars[return_key] = input_vars[key] * adj if return_ssh: return_vars['ssh_u'] = ssh_u return return_vars
[docs] def time_dependent_grid_V(e3v0, e1v, e2v, e1t, e2t, vmask, ssh, input_vars, return_ssh=False): """Calculate time-dependent vertical grid spacing and depths on V-grid for variable volume in NEMO. :arg e3v0: initial vertical scale factors on V-grid. Dimensions: (depth, y, x). :type e3v0: :py:class:`numpy.ndarray` :arg e1v: x-direction scale factors on V-grid. Dimensions: (y, x). :type e1v: :py:class:`numpy.ndarray` :arg e2v: y-direction scale factors on V-grid. Dimensions: (y, x). :type e2v: :py:class:`numpy.ndarray` :arg e1t: x-direction scale factors on T-grid. Dimensions: (y, x). :type e1t: :py:class:`numpy.ndarray` :arg e2t: y-direction scale factors on T-grid. Dimensions: (y, x). :type e2t: :py:class:`numpy.ndarray` :arg vmask: V-grid mask. Dimensions: (depth, y, x) :type vmask: :py:class:`numpy.ndarray` :arg ssh: the model sea surface height on T-grid. Dimensions: (time, y, x) :type ssh: :py:class:`numpy.ndarray` :arg input_vars: A dictionary of the initial grids/scale factors to be translated into time_dependent. Example keys can be 'e3v_0', 'gdepv_0'. A dictionary with corresponding time-dependent grids is returned, where the keys are now 'e3v_t', 'gdepv_t'. :type input_vars: dictionary :arg return_ssh: Specifies whether the ssh interpolated to the V-grid should be returned (ssh_v) :type return_ssh: boolean :returns: A dictionary containing the desired time dependent vertical scale factors and depths on v grid. Dimensions of each: (time, depth, y, x). If returned, ssh_uvhas dimensions (time, y, x)""" ssh_v = np.zeros_like(ssh) e1e2v = e1v * e2v e1e2t = e1t * e2t # Interpolate ssh to V-grid ssh_v[:, 0:-1, :] = 0.5 * vmask[0, 0:-1, :] / e1e2v[0:-1, :] * ( e1e2t[0:-1, :] * ssh[:, 0:-1, :] + e1e2t[1:, :] * ssh[:, 1:, :] ) H = calculate_H(e3v0, vmask) adj = calculate_adjustment_factor(H, ssh_v) adj = np.expand_dims(adj, axis=1) # Time-dependent grids return_vars = {} for key in input_vars: return_key = '{}t'.format(key[0:-1]) return_vars[return_key] = input_vars[key] * adj if return_ssh: return_vars['ssh_v'] = ssh_v return return_vars
[docs] def build_GEM_mask(grid_GEM, grid_NEMO, mask_NEMO): """ """ # Evaluate each point on GEM grid mask_GEM = [] msg = 'Building HRDPS mask' for lon, lat in zip( tqdm(grid_GEM['longitude'].values.flatten() - 360, desc=msg), grid_GEM['latitude'].values.flatten(), ): # Find closest NEMO ji point try: j, i = geo_tools.find_closest_model_point( lon, lat, grid_NEMO['longitude'], grid_NEMO['latitude'], ) except ValueError: j, i = np.nan, np.nan # Append to list if j is np.nan or i is np.nan: mask_GEM.append(0) else: mask_GEM.append(mask_NEMO[j, i].values) # Reshape mask_GEM = np.array(mask_GEM, dtype='int').reshape(grid_GEM['longitude'].shape) return mask_GEM
[docs] def build_matrix(weightsfile, opsfile): """ Given a NEMO weights file and an operational surface forcing file, we assemble the weights into a sparse interpolation matrix that interpolates the surface forcing data from the operational grid to the NEMO grid. This function returns the matrix and the NEMO grid shape. :arg weightsfile: Path to NEMO weights file. :type weightsfile: str :arg opsfile: Path to an operational file. :type opsfile: str :returns: Sparse interpolation matrix and NEMO grid shape :rtype: (:class:`~scipy:scipy.sparse.csr_matrix`, :class:`tuple`) """ # Weights with nc.Dataset(weightsfile) as f: s1 = f.variables['src01'][:]-1 # -1 for fortran-to-python indexing s2 = f.variables['src02'][:]-1 s3 = f.variables['src03'][:]-1 s4 = f.variables['src04'][:]-1 w1 = f.variables['wgt01'][:] w2 = f.variables['wgt02'][:] w3 = f.variables['wgt03'][:] w4 = f.variables['wgt04'][:] with nc.Dataset(opsfile) as f: NO = f.dimensions['x'].size * f.dimensions['y'].size # number of operational grid points NN, nemoshape = s1.size, s1.shape # number of NEMO grid points and shape of NEMO matrix # Build matrix n = np.array([x for x in range(0, NN)]) M1 = sp.csr_matrix((w1.flatten(), (n, s1.flatten())), (NN, NO)) M2 = sp.csr_matrix((w2.flatten(), (n, s2.flatten())), (NN, NO)) M3 = sp.csr_matrix((w3.flatten(), (n, s3.flatten())), (NN, NO)) M4 = sp.csr_matrix((w4.flatten(), (n, s4.flatten())), (NN, NO)) M = M1+M2+M3+M4 return M,nemoshape
[docs] def use_matrix(opsfile, matrix, nemoshape, variable, time): """ Given an operational surface forcing file, an interpolation matrix and NEMO grid shape (produced by grid_tools.build_matrix), a variable name and a time index, we return the operational data interpolated onto the NEMO grid. :arg opsfile: Path to operational file to interpolate. :type opsfile: str :arg matrix: Interpolation matrix (from build_matrix) :type matrix: :class:`~scipy:scipy.sparse.csr_matrix` :arg nemoshape: NEMO grid shape (from build_matrix) :type nemoshape: tuple :arg variable: Specified variable in ops file. :type variable: str :arg time index: time index in ops file. :type time index: integer :returns: Operational data interpolated onto the NEMO grid :rtype: :class:`~numpy:numpy.ndarray` """ with nc.Dataset(opsfile) as f: odata = f.variables[variable][time, ...] # Load the 2D field # Interpolate by matrix multiply - quite fast ndata = matrix*odata.flatten() # Reshape to NEMO shaped array ndata = ndata.reshape(nemoshape) return ndata