diff --git a/conda_package/docs/api.rst b/conda_package/docs/api.rst index a14f213d5..d4affcba7 100644 --- a/conda_package/docs/api.rst +++ b/conda_package/docs/api.rst @@ -225,6 +225,8 @@ Ocean Tools depth.compute_depth depth.compute_zmid + compute_barotropic_streamfunction + .. currentmodule:: mpas_tools.ocean.inject_bathymetry .. autosummary:: diff --git a/conda_package/docs/index.rst b/conda_package/docs/index.rst index b2deb349b..4824ffe04 100644 --- a/conda_package/docs/index.rst +++ b/conda_package/docs/index.rst @@ -50,6 +50,7 @@ analyzing simulations, and in other MPAS-related workflows. ocean/coastline_alteration ocean/moc ocean/depth + ocean/streamfunction ocean/visualization .. toctree:: diff --git a/conda_package/docs/ocean/streamfunction.rst b/conda_package/docs/ocean/streamfunction.rst new file mode 100644 index 000000000..21e6e8481 --- /dev/null +++ b/conda_package/docs/ocean/streamfunction.rst @@ -0,0 +1,14 @@ +.. _ocean_streamfunction: + +Computing streamfunctions +========================= + +Computing the barotropic streamfunction +--------------------------------------- + +The function :py:func:`mpas_tools.ocean.compute_barotropic_streamfunction()` +computes the barotproic streamfunction at vertices on the MPAS-Ocean grid. +The function takes a dataset containing an MPAS-Ocean mesh and another with +``normalVelocity`` and ``layerThickness`` variables (possibly with a +``timeMonthly_avg_`` prefix). The streamfunction is computed only over the +range of (positive-down) depths provided and at the given time index. diff --git a/conda_package/mpas_tools/ocean/__init__.py b/conda_package/mpas_tools/ocean/__init__.py index 1fc1dcb22..24faa794d 100644 --- a/conda_package/mpas_tools/ocean/__init__.py +++ b/conda_package/mpas_tools/ocean/__init__.py @@ -1,7 +1,16 @@ -from mpas_tools.ocean.build_mesh import build_spherical_mesh, \ - build_planar_mesh +from mpas_tools.ocean.build_mesh import ( + build_spherical_mesh, + build_planar_mesh, +) +from mpas_tools.ocean.barotropic_streamfunction import ( + compute_barotropic_streamfunction, +) from mpas_tools.ocean.inject_bathymetry import inject_bathymetry -from mpas_tools.ocean.inject_meshDensity import inject_meshDensity_from_file, \ - inject_spherical_meshDensity, inject_planar_meshDensity -from mpas_tools.ocean.inject_preserve_floodplain import \ - inject_preserve_floodplain +from mpas_tools.ocean.inject_meshDensity import ( + inject_meshDensity_from_file, + inject_spherical_meshDensity, + inject_planar_meshDensity, +) +from mpas_tools.ocean.inject_preserve_floodplain import ( + inject_preserve_floodplain, +) diff --git a/conda_package/mpas_tools/ocean/barotropic_streamfunction.py b/conda_package/mpas_tools/ocean/barotropic_streamfunction.py new file mode 100644 index 000000000..9eae28481 --- /dev/null +++ b/conda_package/mpas_tools/ocean/barotropic_streamfunction.py @@ -0,0 +1,158 @@ +import xarray as xr +import numpy as np +import scipy.sparse +import scipy.sparse.linalg + +import logging +import sys +from mpas_tools.ocean.depth import compute_zmid + + +def compute_barotropic_streamfunction(ds_mesh, ds, logger=None, + min_depth=-5., max_depth=1.e4, + prefix='timeMonthly_avg_', + time_index=0): + """ + Compute barotropic streamfunction. Returns BSF in Sv on vertices. + + Parameters + ---------- + ds_mesh : ``xarray.Dataset`` + A dataset containing MPAS mesh variables + + ds : ``xarray.Dataset`` + A dataset containing MPAS output variables ``normalVelocity`` and + ``layerThickness`` (possibly with a ``prefix``) + + logger : ``logging.Logger``, optional + A logger for the output if not stdout + + min_depth : float, optional + The minimum depth (positive down) to compute transport over + + max_depth : float, optional + The maximum depth (positive down) to compute transport over + + prefix : str, optional + The prefix on the ``normalVelocity`` and ``layerThickness`` variables + + time_index : int, optional + The time at which to index ``ds`` (if it has ``Time`` as a dimension) + """ + + useStdout = logger is None + if useStdout: + logger = logging.getLogger() + logger.addHandler(logging.StreamHandler(sys.stdout)) + logger.setLevel(logging.INFO) + + inner_edges, transport = _compute_transport( + ds_mesh, ds, min_depth=min_depth, max_depth=max_depth, prefix=prefix, + time_index=time_index) + logger.info('transport computed.') + + nvertices = ds_mesh.sizes['nVertices'] + + cells_on_vertex = ds_mesh.cellsOnVertex - 1 + vertices_on_edge = ds_mesh.verticesOnEdge - 1 + is_boundary_cov = cells_on_vertex == -1 + boundary_vertices = np.logical_or(is_boundary_cov.isel(vertexDegree=0), + is_boundary_cov.isel(vertexDegree=1)) + boundary_vertices = np.logical_or(boundary_vertices, + is_boundary_cov.isel(vertexDegree=2)) + + # convert from boolean mask to indices + boundary_vertices = np.flatnonzero(boundary_vertices.values) + + n_boundary_vertices = len(boundary_vertices) + n_inner_edges = len(inner_edges) + + indices = np.zeros((2, 2 * n_inner_edges + n_boundary_vertices), dtype=int) + data = np.zeros(2 * n_inner_edges + n_boundary_vertices, dtype=float) + + # The difference between the streamfunction at vertices on an inner + # edge should be equal to the transport + v0 = vertices_on_edge.isel(nEdges=inner_edges, TWO=0).values + v1 = vertices_on_edge.isel(nEdges=inner_edges, TWO=1).values + + ind = np.arange(n_inner_edges) + indices[0, 2 * ind] = ind + indices[1, 2 * ind] = v1 + data[2 * ind] = 1. + + indices[0, 2 * ind + 1] = ind + indices[1, 2 * ind + 1] = v0 + data[2 * ind + 1] = -1. + + # the streamfunction should be zero at all boundary vertices + ind = np.arange(n_boundary_vertices) + indices[0, 2 * n_inner_edges + ind] = n_inner_edges + ind + indices[1, 2 * n_inner_edges + ind] = boundary_vertices + data[2 * n_inner_edges + ind] = 1. + + rhs = np.zeros(n_inner_edges + n_boundary_vertices, dtype=float) + + # convert to Sv + ind = np.arange(n_inner_edges) + rhs[ind] = 1e-6 * transport + + ind = np.arange(n_boundary_vertices) + rhs[n_inner_edges + ind] = 0. + + matrix = scipy.sparse.csr_matrix( + (data, indices), + shape=(n_inner_edges + n_boundary_vertices, nvertices)) + + solution = scipy.sparse.linalg.lsqr(matrix, rhs) + bsf_vertex = xr.DataArray(-solution[0], + dims=('nVertices',)) + + return bsf_vertex + +def _compute_transport(ds_mesh, ds, min_depth, max_depth, prefix, + time_index): + + cells_on_edge = ds_mesh.cellsOnEdge - 1 + inner_edges = np.logical_and(cells_on_edge.isel(TWO=0) >= 0, + cells_on_edge.isel(TWO=1) >= 0) + + if 'Time' in ds.dims: + ds = ds.isel(Time=time_index) + + # convert from boolean mask to indices + inner_edges = np.flatnonzero(inner_edges.values) + + cell0 = cells_on_edge.isel(nEdges=inner_edges, TWO=0) + cell1 = cells_on_edge.isel(nEdges=inner_edges, TWO=1) + + normal_velocity = \ + ds[f'{prefix}normalVelocity'].isel(nEdges=inner_edges) + layer_thickness = ds[f'{prefix}layerThickness'] + layer_thickness_edge = 0.5 * (layer_thickness.isel(nCells=cell0) + + layer_thickness.isel(nCells=cell1)) + + n_vert_levels = ds.sizes['nVertLevels'] + + vert_index = xr.DataArray.from_dict( + {'dims': ('nVertLevels',), 'data': np.arange(n_vert_levels)}) + mask_bottom = (vert_index < ds_mesh.maxLevelCell).T + mask_bottom_edge = 0.5 * (mask_bottom.isel(nCells=cell0) + + mask_bottom.isel(nCells=cell1)) + + if 'zMid' not in ds.keys(): + z_mid = compute_zmid(ds_mesh.bottomDepth, ds_mesh.maxLevelCell, + ds_mesh.layerThickness) + else: + z_mid = ds.zMid + z_mid_edge = 0.5 * (z_mid.isel(nCells=cell0) + + z_mid.isel(nCells=cell1)) + + mask = np.logical_and(np.logical_and(z_mid_edge >= -max_depth, + z_mid_edge <= -min_depth), + mask_bottom_edge) + normal_velocity = normal_velocity.where(mask) + layer_thickness_edge = layer_thickness_edge.where(mask) + transport = ds_mesh.dvEdge[inner_edges] * \ + (layer_thickness_edge * normal_velocity).sum(dim='nVertLevels') + + return inner_edges, transport diff --git a/conda_package/mpas_tools/ocean/depth.py b/conda_package/mpas_tools/ocean/depth.py index 848b8fd76..7d00b7b68 100644 --- a/conda_package/mpas_tools/ocean/depth.py +++ b/conda_package/mpas_tools/ocean/depth.py @@ -37,7 +37,7 @@ def compute_depth(refBottomDepth): depth_bnds[0, 0] = 0. depth_bnds[1:, 0] = refBottomDepth[0:-1] depth_bnds[:, 1] = refBottomDepth - depth = 0.5*(depth_bnds[:, 0] + depth_bnds[:, 1]) + depth = 0.5 * (depth_bnds[:, 0] + depth_bnds[:, 1]) return depth, depth_bnds @@ -82,11 +82,11 @@ def compute_zmid(bottomDepth, maxLevelCell, layerThickness, thicknessSum = layerThickness.sum(dim=depth_dim) thicknessCumSum = layerThickness.cumsum(dim=depth_dim) - zSurface = -bottomDepth+thicknessSum + zSurface = -bottomDepth + thicknessSum zLayerBot = zSurface - thicknessCumSum - zMid = zLayerBot + 0.5*layerThickness + zMid = zLayerBot + 0.5 * layerThickness zMid = zMid.where(vertIndex < maxLevelCell) if 'Time' in zMid.dims: @@ -150,8 +150,7 @@ def add_depth(inFileName, outFileName, coordFileName=None): history = '{}: {}'.format(time, ' '.join(sys.argv)) if 'history' in ds.attrs: - ds.attrs['history'] = '{}\n{}'.format(history, - ds.attrs['history']) + ds.attrs['history'] = f'{history}\n{ds.attrs["history"]}' else: ds.attrs['history'] = history @@ -229,8 +228,7 @@ def add_zmid(inFileName, outFileName, coordFileName=None): history = '{}: {}'.format(time, ' '.join(sys.argv)) if 'history' in ds.attrs: - ds.attrs['history'] = '{}\n{}'.format(history, - ds.attrs['history']) + ds.attrs['history'] = f'{history}\n{ds.attrs["history"]}' else: ds.attrs['history'] = history @@ -288,8 +286,8 @@ def write_time_varying_zmid(inFileName, outFileName, coordFileName=None, dsIn = xarray.open_dataset(inFileName) dsIn = dsIn.rename({'nVertLevels': 'depth'}) - inVarName = '{}layerThickness'.format(prefix) - outVarName = '{}zMid'.format(prefix) + inVarName = f'{prefix}layerThickness' + outVarName = f'{prefix}zMid' layerThickness = dsIn[inVarName] zMid = compute_zmid(dsCoord.bottomDepth, dsCoord.maxLevelCell,