diff --git a/conda_package/mpas_tools/ocean/barotropic_streamfunction.py b/conda_package/mpas_tools/ocean/barotropic_streamfunction.py index 4ee94e78..99286d4d 100644 --- a/conda_package/mpas_tools/ocean/barotropic_streamfunction.py +++ b/conda_package/mpas_tools/ocean/barotropic_streamfunction.py @@ -10,11 +10,12 @@ def compute_barotropic_streamfunction(ds_mesh, ds, logger=None, - min_depth=-5., max_depth=1.e4, + min_depth=None, max_depth=None, prefix='timeMonthly_avg_', - time_index=0): + time_index=None, include_bolus=False, + include_submesoscale=False): """ - Compute barotropic streamfunction. Returns BSF in Sv on vertices. + Compute barotropic streamfunction Parameters ---------- @@ -29,16 +30,27 @@ def compute_barotropic_streamfunction(ds_mesh, ds, logger=None, A logger for the output if not stdout min_depth : float, optional - The minimum depth (positive down) to compute transport over + The minimum depth (positive up) to compute BSF over max_depth : float, optional - The maximum depth (positive down) to compute transport over + The maximum depth (positive up) to compute BSF 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) + + include_bolus : bool, optional + Whether to include the GM bolus velocity in the computation + + include_submesoscale : bool, optional + Whether to include the submesoscale velocity in the computation + + Returns + ------- + bsf_vertex : ``xarray.DataArray`` + The barotropic streamfunction in Sv on vertices """ useStdout = logger is None @@ -47,114 +59,258 @@ def compute_barotropic_streamfunction(ds_mesh, ds, logger=None, 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'] + if time_index is not None: + ds = ds.isel(Time=time_index) - 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)) + bsf_vertex = _compute_barotropic_streamfunction_vertex( + ds_mesh, ds, prefix, include_bolus, include_submesoscale, min_depth, + max_depth) - # convert from boolean mask to indices - boundary_vertices = np.flatnonzero(boundary_vertices.values) + return bsf_vertex - 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) +def shift_barotropic_streamfunction(bsf_vertex, lat_range, cells_on_vertex, + lat_vertex): + """ + Shift the barotropic streamfunction to be zero on average at the boundary + over the given latitude range - # 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 + Parameters + ---------- + bsf_vertex : xarray.DataArray + The barotropic streamfunction in Sv on vertices - ind = np.arange(n_inner_edges) - indices[0, 2 * ind] = ind - indices[1, 2 * ind] = v1 - data[2 * ind] = 1. + lat_range : list of float + The latitude range over which to set the mean boundary value of the BSF + to zero - indices[0, 2 * ind + 1] = ind - indices[1, 2 * ind + 1] = v0 - data[2 * ind + 1] = -1. + cells_on_vertex : xarray.DataArray + The zero-based cell indices on each vertex - # 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. + lat_vertex : xarray.DataArray + The latitude of each vertex - rhs = np.zeros(n_inner_edges + n_boundary_vertices, dtype=float) + Returns + ------- + bsf_shifted : ``xarray.DataArray`` + The shifted barotropic streamfunction in Sv on vertices + """ + is_boundary_cov = cells_on_vertex == -1 + boundary_vertices = is_boundary_cov.sum(dim='vertexDegree') > 0 - # convert to Sv - ind = np.arange(n_inner_edges) - rhs[ind] = 1e-6 * transport + boundary_vertices = np.logical_and( + boundary_vertices, + lat_vertex >= np.deg2rad(lat_range[0]) + ) + boundary_vertices = np.logical_and( + boundary_vertices, + lat_vertex <= np.deg2rad(lat_range[1]) + ) - ind = np.arange(n_boundary_vertices) - rhs[n_inner_edges + ind] = 0. + # convert from boolean mask to indices + boundary_vertices = np.flatnonzero(boundary_vertices.values) - matrix = scipy.sparse.csr_matrix( - (data, indices), - shape=(n_inner_edges + n_boundary_vertices, nvertices)) + mean_boundary_bsf = bsf_vertex.isel(nVertices=boundary_vertices).mean() - solution = scipy.sparse.linalg.lsqr(matrix, rhs) - bsf_vertex = xr.DataArray(-solution[0], - dims=('nVertices',)) + bsf_shifted = bsf_vertex - mean_boundary_bsf - return bsf_vertex + return bsf_shifted -def _compute_transport(ds_mesh, ds, min_depth, max_depth, prefix, - time_index): +def _compute_vert_integ_velocity(ds_mesh, ds, prefix, include_bolus, + include_submesoscale, min_depth, max_depth): 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'] + layer_thickness = ds.timeMonthly_avg_layerThickness + max_level_cell = ds_mesh.maxLevelCell - 1 + 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 + z_mid = compute_zmid(ds_mesh.bottomDepth, max_level_cell, + layer_thickness) + z_mid_edge = 0.5*(z_mid.isel(nCells=cell0) + + z_mid.isel(nCells=cell1)) + + normal_velocity = ds[f'{prefix}normalVelocity'] + if include_bolus: + normal_velocity += ds[f'{prefix}normalGMBolusVelocity'] + if include_submesoscale: + normal_velocity += ds[f'{prefix}normalMLEvelocity'] + normal_velocity = normal_velocity.isel(nEdges=inner_edges) + + layer_thickness_edge = 0.5*(layer_thickness.isel(nCells=cell0) + + layer_thickness.isel(nCells=cell1)) + mask_bottom = (vert_index <= max_level_cell).T + mask_bottom_edge = np.logical_and(mask_bottom.isel(nCells=cell0), + mask_bottom.isel(nCells=cell1)) + masks = [mask_bottom_edge] + if min_depth is not None: + masks.append(z_mid_edge <= min_depth) + if max_depth is not None: + masks.append(z_mid_edge >= max_depth) + for mask in masks: + normal_velocity = normal_velocity.where(mask) + layer_thickness_edge = layer_thickness_edge.where(mask) + + vert_integ_velocity = np.zeros(ds_mesh.dims['nEdges'], dtype=float) + inner_vert_integ_vel = ( + (layer_thickness_edge * normal_velocity).sum(dim='nVertLevels')) + vert_integ_velocity[inner_edges] = inner_vert_integ_vel.values + + vert_integ_velocity = xr.DataArray(vert_integ_velocity, + dims=('nEdges',)) + + return vert_integ_velocity + + +def _compute_edge_sign_on_vertex(ds_mesh): + edges_on_vertex = ds_mesh.edgesOnVertex - 1 + vertices_on_edge = ds_mesh.verticesOnEdge - 1 + + nvertices = ds_mesh.sizes['nVertices'] + vertex_degree = ds_mesh.sizes['vertexDegree'] + + edge_sign_on_vertex = np.zeros((nvertices, vertex_degree), dtype=int) + vertices = np.arange(nvertices) + for iedge in range(vertex_degree): + eov = edges_on_vertex.isel(vertexDegree=iedge) + valid_edge = eov >= 0 + + v0_on_edge = vertices_on_edge.isel(nEdges=eov, TWO=0) + v1_on_edge = vertices_on_edge.isel(nEdges=eov, TWO=1) + valid_edge = np.logical_and(valid_edge, v0_on_edge >= 0) + valid_edge = np.logical_and(valid_edge, v1_on_edge >= 0) + + mask = np.logical_and(valid_edge, v0_on_edge == vertices) + edge_sign_on_vertex[mask, iedge] = -1 + + mask = np.logical_and(valid_edge, v1_on_edge == vertices) + edge_sign_on_vertex[mask, iedge] = 1 + + return edge_sign_on_vertex + + +def _compute_vert_integ_vorticity(ds_mesh, vert_integ_velocity, + edge_sign_on_vertex): + + area_vertex = ds_mesh.areaTriangle + dc_edge = ds_mesh.dcEdge + edges_on_vertex = ds_mesh.edgesOnVertex - 1 + + vertex_degree = ds_mesh.sizes['vertexDegree'] + + vert_integ_vorticity = xr.zeros_like(ds_mesh.latVertex) + for iedge in range(vertex_degree): + eov = edges_on_vertex.isel(vertexDegree=iedge) + edge_sign = edge_sign_on_vertex[:, iedge] + dc = dc_edge.isel(nEdges=eov) + vert_integ_vel = vert_integ_velocity.isel(nEdges=eov) + vert_integ_vorticity += ( + dc / area_vertex * edge_sign * vert_integ_vel) + + return vert_integ_vorticity + + +def _compute_barotropic_streamfunction_vertex(ds_mesh, ds, prefix, + include_bolus, + include_submesoscale, min_depth, + max_depth): + edge_sign_on_vertex = _compute_edge_sign_on_vertex(ds_mesh) + vert_integ_velocity = _compute_vert_integ_velocity(ds_mesh, ds, prefix, + include_bolus, + include_submesoscale, + min_depth, max_depth) + vert_integ_vorticity = _compute_vert_integ_vorticity( + ds_mesh, vert_integ_velocity, edge_sign_on_vertex) + + nvertices = ds_mesh.sizes['nVertices'] + vertex_degree = ds_mesh.sizes['vertexDegree'] + + edges_on_vertex = ds_mesh.edgesOnVertex - 1 + vertices_on_edge = ds_mesh.verticesOnEdge - 1 + area_vertex = ds_mesh.areaTriangle + dc_edge = ds_mesh.dcEdge + dv_edge = ds_mesh.dvEdge + + # one equation involving vertex degree + 1 vertices for each vertex + # plus 2 entries for the boundary condition and Lagrange multiplier + ndata = (vertex_degree + 1) * nvertices + 2 + indices = np.zeros((2, ndata), dtype=int) + data = np.zeros(ndata, dtype=float) + + # the laplacian on the dual mesh of the streamfunction is the + # vertically integrated vorticity + vertices = np.arange(nvertices, dtype=int) + idata = (vertex_degree + 1) * vertices + 1 + indices[0, idata] = vertices + indices[1, idata] = vertices + for iedge in range(vertex_degree): + eov = edges_on_vertex.isel(vertexDegree=iedge) + dc = dc_edge.isel(nEdges=eov) + dv = dv_edge.isel(nEdges=eov) + + v0 = vertices_on_edge.isel(nEdges=eov, TWO=0) + v1 = vertices_on_edge.isel(nEdges=eov, TWO=1) + + edge_sign = edge_sign_on_vertex[:, iedge] + + mask = v0 == vertices + # the difference is v1 - v0, so we want to subtract this vertex + # when it is v0 and add it when it is v1 + this_vert_sign = np.where(mask, -1., 1.) + # the other vertex is obviously whichever one this is not + other_vert_index = np.where(mask, v1, v0) + # if there are invalid vertices, we need to make sure we don't + # index out of bounds. The edge_sign will mask these out + other_vert_index = np.where(other_vert_index >= 0, + other_vert_index, 0) + + idata_other = idata + iedge + 1 + + indices[0, idata] = vertices + indices[1, idata] = vertices + indices[0, idata_other] = vertices + indices[1, idata_other] = other_vert_index + + this_data = this_vert_sign * edge_sign * dc / (dv * area_vertex) + data[idata] += this_data + data[idata_other] = -this_data + + # Now, the boundary condition: To begin with, we set the BSF at the + # frist vertext to zero + indices[0, -2] = nvertices + indices[1, -2] = 0 + data[-2] = 1. + + # The same in the final column + indices[0, -1] = 0 + indices[1, -1] = nvertices + data[-1] = 1. + + # one extra spot for the Lagrange multiplier + rhs = np.zeros(nvertices + 1, dtype=float) + + rhs[0:-1] = vert_integ_vorticity.values + + matrix = scipy.sparse.csr_matrix( + (data, indices), + shape=(nvertices + 1, nvertices + 1)) + + solution = scipy.sparse.linalg.spsolve(matrix, rhs) + + # drop the Lagrange multiplier and convert to Sv with the desired sign + # convention + bsf_vertex = xr.DataArray(-1e-6 * solution[0:-1], + dims=('nVertices',)) + + return bsf_vertex