import numpy as np
import xarray as xr
import dask.array as dsar
from scipy.sparse import coo_matrix, csc_matrix, eye
from scipy.sparse.linalg import spsolve
from scipy.interpolate import PchipInterpolator as pchip
# import xrft
import warnings
__all__ = ['w_rigid']
[docs]def w_rigid(N2, f0, beta, Frhs, kx, ky, dZ, dZ0=None, dZ1=None, zdim='Zl',
dim=None, coord=None):
"""
Inverts the Omega equation given by Giordani and Planton (2000)
to get the balanced vertical velocity (:math:`w_b`)
for rigid-lid boundary conditions given the right-hand side fo the equation.
Parameters
----------
N2 : float or xarray.DataArray
The buoyancy frequency.
f0 : float
Coriolis parameter.
beta : float
Meridional gradient of the Coriolis parameter
for a beta-plane approximation.
Frhs : xarray.DataArray
The Fourier transform of the right-hand side of
the Omega equation. The last two dimensions should be
the meridional and zonal wavenumber.
kx : xarray.DataArray
Zonal wavenumber.
ky : xarray.DataArray
Meridional wavenumber.
dZ : float or xarray.DataArray
Vertical distance between grid. This should take constant
value(s) for best numerical percision.
dZ0 : float or xarray.DataArray, optional
Top vertical distance between grids.
dZ1 : float or xarray.DataArray, optional
Bottom vertical distance between grids.
zdim : str, optional
Dimension name of the vertical axis of `Frhs`.
dim : list, optional
List of the xarray.DataArray output.
coord : dict, optional
Dictionary of the xarray.DataArray output.
Returns
-------
wa : xarray.DataArray
The inverted ageostrophic vertical velocity.
"""
Zl = Frhs[zdim]
kdims = Frhs.dims[-2:]
N = Frhs.shape
nz = N[0]
if isinstance(N2, float):
enu = eye(nz-2) * N2
else:
if nz-2 != len(N2):
raise ValueError("N2 should have two elements less than Frhs.")
row = range(nz-2)
col = range(nz-2)
# if len(N2) != nz-1:
# raise ValueError("N2 should have one element less than psi.")
if N2.dims != Zl.dims:
raise ValueError("N2 and psi should be on the same vertical grid.")
enu = coo_matrix((N2,(row,col)),
shape=(nz,nz), dtype=np.float64
)
### Delta matrix ###
row = np.repeat(range(1,nz-1),3)
row = np.append(np.array([0]), row)
# row = np.append(0, row)
row = np.append(row, np.array([nz-1]))
# row = np.append(row, nz)
col = np.arange(3)
for i in range(nz-3):
col = np.append(col, np.arange(i+1,i+4))
col = np.append(0, col)
# col = np.append(0, col)
col = np.append(col, np.array([nz-1]))
# col = np.append(col, nz-1)
if dZ0 == None:
dZ0 = dZ
if dZ1 == None:
dZ1 = dZ
if isinstance(dZ, float):
dZ = dZ*np.ones(nz)
DZ = dZ
dZ = np.append(dZ0, dZ)
dZ = np.append(dZ, dZ1)
else:
warnings.warn("The data is not on uniform vertical gridding. "
"The numerical errors for vertical derivatives "
"may be significant.")
tmp = np.zeros((nz-2,3))
tmp[:,0] = DZ[:-2]**-1
tmp[:,-1] = DZ[1:-1]**-1
tmp[:,1] = -(DZ[:-2]**-1 + DZ[1:-1]**-1)
data = np.zeros(len(tmp.ravel())+2)
data[1:-1] = (tmp * dZ[1:nz-1,np.newaxis]**-1).ravel()
# data[0] = dZ[0]**-1 * DZ[0]
data[0] = 1
data[-1] = 1
# data[-1] = dZ[-1]**-1 * DZ[-1]**-1
data *= f0**2
delta = coo_matrix((data,(row,col)),
shape=(nz,nz),dtype=np.float64
)
# ky = Frhs[kdims[0]]
# kx = Frhs[kdims[1]]
# if wvnm == False:
# warnings.warn("The coordinates are in inverse wavelenths so "
# "converting them to wavenumbers.")
# ky *= 2*np.pi
# kx *= 2*np.pi
nk, nl = (len(kx),len(ky))
wahat = np.zeros_like(Frhs, dtype=np.complex128)
for i in range(nk):
for j in range(nl):
kappa2 = kx[i].data**2+ky[j].data**2
### Normal inversion ###
# A = csc_matrix(delta, dtype=np.float64)
# A[1:-1] -= csc_matrix(kappa2*enu, dtype=np.float64)
A = csc_matrix(-kappa2*enu+delta, dtype=np.float64)
### Rigid lid solution ###
wahat[:,j,i] = spsolve(A, Frhs[:,j,i])
# wahat[1:-1,j,i], res, rnk, s = lstsq(A.todense(), Frhs[:,j,i])
wahat = xr.DataArray(wahat, dims=[dim[0],kdims[-2],kdims[-1]],
coords={dim[0]:Zl.data,kdims[-2]:ky,kdims[-1]:kx}
)
wa = dsar.fft.ifft2(wahat.chunk(chunks={dim[0]:1,
kdims[-1]:N[-1],
kdims[-2]:N[-2]}
).data, axes=[-2,-1]
).real
return xr.DataArray(wa, dims=dim, coords=coord)