Switch to using barotropic streamfunction from mpas_tools

Author: xylar

mpas_analysis/ocean/climatology_map_bsf.py

@@ -9,14 +9,15 @@
 # distributed with this code, or at
 # https://raw.githubusercontent.com/MPAS-Dev/MPAS-Analysis/main/LICENSE
 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
 
 
@@ -280,8 +281,21 @@ 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)
+
+        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.info('bsf on vertices computed.')
 
         climatology['barotropicStreamfunction'] = bsf_vertex
@@ -305,234 +319,12 @@ 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)
                 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