analysis/barotropicStreamfunction.py at master · milenaveneziani/analysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
from __future__ import absolute_import, division, print_function, \
    unicode_literals
import xarray as xr
import numpy as np
import scipy.sparse
import scipy.sparse.linalg

from mpas_analysis.ocean.utility import compute_zmid


def compute_barotropic_streamfunction_vertex(dsMesh, ds, min_lat, min_depth, max_depth):

        inner_edges, transport = _compute_transport(dsMesh, ds, min_depth, max_depth)

        nvertices = dsMesh.sizes['nVertices']
        cells_on_vertex = dsMesh.cellsOnVertex - 1
        vertices_on_edge = dsMesh.verticesOnEdge - 1
        is_boundary_cov = cells_on_vertex == -1
        boundary_vertices = is_boundary_cov.sum(dim='vertexDegree') > 0
        boundary_vertices = np.logical_and(boundary_vertices,
                                           dsMesh.latVertex > np.deg2rad(min_lat))

        # 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.

        # Make the average of the stream function at boundary vertices north
        # of min_lat equal to zero
        ind = np.arange(n_boundary_vertices)
        indices[0, 2*n_inner_edges + ind] = n_inner_edges
        indices[1, 2*n_inner_edges + ind] = ind
        data[2*n_inner_edges + ind] = 1.

        rhs = np.zeros(n_inner_edges + 1, dtype=float)

        # convert to Sv
        ind = np.arange(n_inner_edges)
        rhs[ind] = 1e-6*np.squeeze(transport)

        rhs[n_inner_edges] = 0.

        matrix = scipy.sparse.csr_matrix(
            (data, indices),
            shape=(n_inner_edges + 1, nvertices))

        solution = scipy.sparse.linalg.lsqr(matrix, rhs)
        bsf_vertex = xr.DataArray(-solution[0],
                                  dims=('nVertices',))

        return bsf_vertex


def _compute_transport(dsMesh, ds, min_depth, max_depth):

    cells_on_edge = dsMesh.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']

    vert_index = xr.DataArray.from_dict(
        {'dims': ('nVertLevels',), 'data': np.arange(n_vert_levels)})
    z_mid = compute_zmid(dsMesh.bottomDepth, dsMesh.maxLevelCell-1,
                         dsMesh.layerThickness)
    z_mid_edge = 0.5*(z_mid.isel(nCells=cell0) + z_mid.isel(nCells=cell1))

    normal_velocity = ds.timeMonthly_avg_normalVelocity.isel(nEdges=inner_edges)
    if 'timeMonthly_avg_normalGMBolusVelocity' in ds.keys():
        normal_velocity = normal_velocity + ds.timeMonthly_avg_normalGMBolusVelocity.isel(nEdges=inner_edges)
    if 'timeMonthly_avg_normalMLEvelocity' in ds.keys():
        normal_velocity = normal_velocity + ds.timeMonthly_avg_normalMLEvelocity.isel(nEdges=inner_edges)

    layer_thickness = ds.timeMonthly_avg_layerThickness
    layer_thickness_edge = 0.5*(layer_thickness.isel(nCells=cell0) +
                                layer_thickness.isel(nCells=cell1))
    mask_bottom = (vert_index < dsMesh.maxLevelCell).T
    mask_bottom_edge = np.logical_and(mask_bottom.isel(nCells=cell0),
                                      mask_bottom.isel(nCells=cell1))
    masks = [mask_bottom_edge,
             z_mid_edge <= max_depth,
             z_mid_edge >= min_depth]
    for mask in masks:
        normal_velocity = normal_velocity.where(mask)
        layer_thickness_edge = layer_thickness_edge.where(mask)

    transport = dsMesh.dvEdge[inner_edges] * (layer_thickness_edge * normal_velocity).sum(dim='nVertLevels')

    return inner_edges, transport