Skip to content

Commit

Permalink
Solve barotropic streamfunction from vorticity
Browse files Browse the repository at this point in the history 
This merge is a port of:
MPAS-Dev/MPAS-Analysis#1069

Instead of computing the BSF using a least-squares solve based on the transport between vertices, the new approach inverts a Poisson equation for the BSF based on the (vertically integrated) vorticity.

Using a SORRM G-Case simulation, the vertically integrated vorticity computed from the streamfunciton is within machine precision of the original vertically integrated vorticity. The velocity produced by the streamfunciton has errors that are on the order of 1e-3 compared with the original vertically integrated velocity (likely caused by the field not being perfectly divergence free).
  • Loading branch information
@xylar
xylar committed Mar 9, 2025
1 parent 287d230 commit a0bb205
Showing 1 changed file with 242 additions and 86 deletions.
328 changes: 242 additions & 86 deletions conda_package/mpas_tools/ocean/barotropic_streamfunction.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand All @@ -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
Expand All @@ -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

0 comments on commit a0bb205

Please sign in to comment.