From e9d0bc6e387c40d699cad145af083b1e4e31f830 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sun, 9 Mar 2025 06:56:08 -0500 Subject: [PATCH 1/2] Switch to using barotropic streamfunction from `mpas_tools` --- mpas_analysis/ocean/climatology_map_bsf.py | 267 +++------------------ 1 file changed, 36 insertions(+), 231 deletions(-) diff --git a/mpas_analysis/ocean/climatology_map_bsf.py b/mpas_analysis/ocean/climatology_map_bsf.py index ddae4ade9..5a55cf516 100644 --- a/mpas_analysis/ocean/climatology_map_bsf.py +++ b/mpas_analysis/ocean/climatology_map_bsf.py @@ -11,14 +11,15 @@ import os import xarray as xr -import numpy as np -import scipy.sparse -import scipy.sparse.linalg + +from mpas_tools.ocean.barotropic_streamfunction import ( + compute_barotropic_streamfunction, + shift_barotropic_streamfunction +) from mpas_analysis.shared import AnalysisTask from mpas_analysis.shared.climatology import RemapMpasClimatologySubtask from mpas_analysis.shared.plot import PlotClimatologyMapSubtask -from mpas_analysis.ocean.utility import compute_zmid from mpas_analysis.shared.projection import comparison_grid_option_suffixes @@ -315,8 +316,30 @@ def customize_masked_climatology(self, climatology, season): 'edgesOnVertex', 'dcEdge', 'dvEdge', 'bottomDepth', 'maxLevelCell', 'latVertex', 'areaTriangle',]] ds_mesh.load() - bsf_vertex = self._compute_barotropic_streamfunction_vertex( - ds_mesh, climatology) + + cells_on_vertex = ds_mesh.cellsOnVertex - 1 + lat_vertex = ds_mesh.latVertex + bsf_vertex = compute_barotropic_streamfunction( + ds_mesh=ds_mesh, + ds=climatology, + min_depth=self.min_depth, + max_depth=self.max_depth, + include_bolus=self.include_bolus, + include_submesoscale=self.include_submesoscale, + logger=logger, + ) + + lat_range = config.getexpression( + self.taskName, 'latitudeRangeForZeroBSF') + + bsf_vertex = shift_barotropic_streamfunction( + bsf_vertex=bsf_vertex, + lat_range=lat_range, + cells_on_vertex=cells_on_vertex, + lat_vertex=lat_vertex, + logger=logger, + ) + logger.info('bsf on vertices computed.') climatology['barotropicStreamfunction'] = bsf_vertex @@ -340,234 +363,16 @@ def customize_masked_climatology(self, climatology, season): lat_range = config.getexpression( config_section_name, 'latitudeRangeForZeroBSF') - climatology[mpas_field_name] = _shift_bsf( - bsf_vertex, lat_range, ds_mesh.cellsOnVertex - 1, - ds_mesh.latVertex) + climatology[mpas_field_name] = shift_barotropic_streamfunction( + bsf_vertex=bsf_vertex, + lat_range=lat_range, + cells_on_vertex=cells_on_vertex, + lat_vertex=lat_vertex, + logger=logger, + ) climatology[mpas_field_name].attrs['units'] = 'Sv' climatology[mpas_field_name].attrs['description'] = \ f'barotropic streamfunction at vertices, offset for ' \ f'{grid_suffix} plots' return climatology - - def _compute_vert_integ_velocity(self, ds_mesh, ds): - - 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) - - # 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) - 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)}) - 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.timeMonthly_avg_normalVelocity - if self.include_bolus: - normal_velocity += ds.timeMonthly_avg_normalGMBolusVelocity - if self.include_submesoscale: - normal_velocity += ds.timeMonthly_avg_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, - z_mid_edge <= self.min_depth, - z_mid_edge >= self.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(self, 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(self, 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(self, ds_mesh, ds): - edge_sign_on_vertex = self._compute_edge_sign_on_vertex(ds_mesh) - vert_integ_velocity = self._compute_vert_integ_velocity(ds_mesh, ds) - vert_integ_vorticity = self._compute_vert_integ_vorticity( - ds_mesh, vert_integ_velocity, edge_sign_on_vertex) - self.logger.info('vertically integrated vorticity computed.') - - config = self.config - lat_range = config.getexpression( - 'climatologyMapBSF', 'latitudeRangeForZeroBSF') - - nvertices = ds_mesh.sizes['nVertices'] - vertex_degree = ds_mesh.sizes['vertexDegree'] - - cells_on_vertex = ds_mesh.cellsOnVertex - 1 - 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',)) - - bsf_vertex = _shift_bsf(bsf_vertex, lat_range, cells_on_vertex, - ds_mesh.latVertex) - - return bsf_vertex - - -def _shift_bsf(bsf_vertex, lat_range, cells_on_vertex, lat_vertex): - """ - Shift the barotropic streamfunction to be zero at the boundary over - the given latitude range - """ - is_boundary_cov = cells_on_vertex == -1 - boundary_vertices = is_boundary_cov.sum(dim='vertexDegree') > 0 - - 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]) - ) - - # convert from boolean mask to indices - boundary_vertices = np.flatnonzero(boundary_vertices.values) - - mean_boundary_bsf = bsf_vertex.isel(nVertices=boundary_vertices).mean() - - bsf_shifted = bsf_vertex - mean_boundary_bsf - - return bsf_shifted From 56f1e89eb53d5564b407b038616690e54d9641d8 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sun, 13 Apr 2025 12:11:14 -0500 Subject: [PATCH 2/2] Update to mpas_tools >=1.1.0 --- ci/recipe/meta.yaml | 2 +- dev-spec.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ci/recipe/meta.yaml b/ci/recipe/meta.yaml index bfd22bb0d..f1b412d0f 100644 --- a/ci/recipe/meta.yaml +++ b/ci/recipe/meta.yaml @@ -32,7 +32,7 @@ requirements: - lxml - mache >=1.11.0 - matplotlib-base >=3.9.0 - - mpas_tools >=1.0.0,<2.0.0 + - mpas_tools >=1.1.0,<2.0.0 - nco >=4.8.1,!=5.2.6 - netcdf4 - numpy >=2.0,<3.0 diff --git a/dev-spec.txt b/dev-spec.txt index 9f826e7bb..1ce19783b 100644 --- a/dev-spec.txt +++ b/dev-spec.txt @@ -15,7 +15,7 @@ gsw lxml mache >=1.11.0 matplotlib-base>=3.9.0 -mpas_tools >=1.0.0,<2.0.0 +mpas_tools >=1.1.0,<2.0.0 nco>=4.8.1,!=5.2.6 netcdf4 numpy>=2.0,<3.0