import numpy as np
[docs]def build_grid(dims, e1u=1.0, e2v=1.0, e2u=1.0, e1v=1.0):
"""
Build pairs of adjacent grid points and associated geometry for a tripolar grid
Parameters
----------
dims : tuple of int
Dimensions of the grid. Letting `dims == (nj, ni)`, then
`nj` is the number of grid cells in the y (latitude-like) direction, &
`ni` is the number of grid cells in the x (longitude-like) direction.
e1u : float or 2D array, Default 1.0
Distance in the x dimension that lives on the U cell.
e2v : float or 2D array, Default 1.0
Distance in the y dimension that lives on the V cell.
e2u : float or 2D array, Default 1.0
Distance in the y dimension that lives on the U cell.
e1v : float or 2D array, Default 1.0
Distance in the x dimension that lives on the V cell.
Notes
-----
`e1u`, `e2v`, `e2u`, `e1v` can have shape `(nj,ni)` or `(nj+1, ni+2)`.
The latter form has 1 cell of padding in the y dimension to handle the
tripolar grid, and 1 cell of padding in the x dimension to handle the
longitudinal periodicity. If given the latter form, this extra padding is
removed internally. When of shape `(nj, ni)`,
- `e1u[j,i]` is the x distance on the U grid at [j,i+1/2],
- `e2v[j,i]` is the y distance on the V grid at [j+1/2,i],
- `e2u[j,i]` is the y distance on the U grid at [j,i+1/2],
- `e1v[j,i]` is the x distance on the V grid at [j+1/2,i].
Returns
-------
grid : dict
Containing the following:
edges : tuple of length 2
Each element is an array of int of length `E`, where `E` is the number of
edges in the grid's graph, i.e. the number of pairs of adjacent water
columns (including land) in the grid.
If `edges = (a, b)`, the nodes (water columns) whose linear indices are
`a[i]` and `b[i]` are adjacent.
dist : 1d array
Horizontal distance between adjacent water columns (nodes).
`dist[i]` is the distance between nodes whose linear indices are
`edges[0][i]` and `edges[1][i]`.
distperp : 1d array
Horizontal distance of the face between adjacent water columns (nodes).
`distperp[i]` is the distance of the interface between nodes whose
linear indices are `edges[0][i]` and `edges[1][i]`.
"""
grid = dict()
grid["edges"] = _build_edges(dims)
grid["dist"] = _build_edgedata(dims, (e2v, e1u))
grid["distperp"] = _build_edgedata(dims, (e1v, e2u))
return grid
def _build_edges(dims):
"""
Build list of pairs of adjacent grid points, numbered 0, ..., N-1, on a
tripolar grid (for ORCA)
"""
assert len(dims) == 2, "Need `len(dims) == 2`, for a 2D horizontal grid."
nj, ni = dims
Ei = ni * nj # number of edges in i dimension
Ej = ni * nj # number of edges in j dimension
E = Ei + Ej # number of edges in total
a = np.empty(E, dtype=int) # prealloc space
b = np.empty(E, dtype=int)
# Build linear indices to the entire grid with one extra row of padding
# for the northern latitudes. Do not add longitudinal padding.
idx = np.arange((nj + 1) * ni).reshape((nj + 1, ni))
# Handle tripolar grid's northern latitudes: using the extra row of padding,
# the cell indexed (here using 0-indexing) by [nj - 1, i] is actually in the
# domain (not the padded row) and is adjacent to [nj - 1, ni - i - 1]. So,
# make the linear index at cell [nj, i] have the same value as the linear
# index at cell [nj - 1, ni - i - 1].
idx[-1, :] = np.flip(idx[-2, :])
# Handle j dimension (non-periodic).
# Cells [j, i] and [j+1, i] are adjacent, for 0 <= j <= nj-1. This is true
# even for j == nj-1, thanks to the extra padded row and np.flip above.
a[:Ej] = idx[:nj, :].reshape(-1) # [j, i]
b[:Ej] = idx[1:, :].reshape(-1) # [j+1, i]
# Handle i dimension (zonally periodic).
# Cells [j, i] and [j, mod(i+1, ni)] are adjacent.
a[Ej:] = idx[0:-1, :].reshape(-1) # [j, i]
b[Ej:] = np.roll(idx[0:-1, :], -1, axis=1).reshape(-1) # [j, i+1]
return a, b
def _build_edgedata(dims, data):
"""
Build a 1D array of `data` in the same order as constructed by `_build_edges`
Parameters
----------
dims :
See `build_grid`
data : tuple of ndarray
Data that lives on edges, i.e. between nodes (water columns).
With `dims == (nj, ni)`, each element of `data` should be a 2D array of
shape `(nj, ni)` --- but if they are of shape `(nj+1, ni+2)` they will
be trimmed to remove the last row and the first and last columns.
- `data[0][j, i]` lives between nodes [i, j] and [i, j+1], for `0 <= j <= nj-1`.
- `data[1][j, i]` lives between nodes [i, j] and [i+1, j], for `0 <= i <= ni-1`.
Example: `data = (e2v, e1u)` where `e1u` is the distance in the 1st (x)
dimension on the U grid and `e2v` is the distance in the 2nd (y)
dimension on the V grid. These can be of shape `(nj, ni)` or be
of shape `(nj+1, ni+2)` as in the mesh_mask data file.
Returns
-------
edgedata : array
1D array of `data`.
If `a, b = _build_edges(dims)`, then `edgedata[i]` is the
value from `data` that corresponds to the interface between the grid
cells indexed by `a[i]` and `b[i]`.
"""
nj, ni = dims
v, u = data
if v.shape == (nj + 1, ni + 2):
v = v[:-1, 1:-1]
if u.shape == (nj + 1, ni + 2):
u = u[:-1, 1:-1]
return np.concatenate((v.reshape(-1), u.reshape(-1)))
[docs]def edgedata_to_maps(edgedata, dims):
"""
Convert 1D array of data living on edges into two 2D arrays, one for each
spatial dimension, for a tripolar grid
Parameters
----------
edgedata : array
1D array of data that lives on edges of the grid's graph, given in the
same order as the edges cosntructed from `_build_edges`.
If `a, b = _build_edges(dims, periodic)`, then `edgedata[i]` is the
data that lives at the interface between the grid cells indexed by
`a[i]` and `b[i]`.
dims :
See `build_grid`
Returns
-------
v, u : 2D array
2D arrays of shape `dims`. `v` and `u` contains the values from `edgedata`
that correspond to data living on the V and U grids, respectively.
"""
nj, ni = dims
Ej = ni * nj # number of edges in j dimension
v = np.reshape(edgedata[:Ej], dims)
u = np.reshape(edgedata[Ej:], dims)
return v, u