def kinetic_energy(sshuv, gravity=default_gravity, format_axes=None, dxy=None):
"""Compute kinetic energy in m2.s-2 either from SSH or velocity on C-grid
.. todo:: Rewrite it using :mod:`vacumm.data.misc.arakawa` and defining
a limited number of algorithms for different staggering configurations.
:Params:
- **sshuv**: SSH or (U,V).
- If SSH, geostrophic velocity is computed at U and V points
using :func:`barotropic_geostrophic_velocity`.
- If (U,V), velocities are supposed to be at V and U points.
- **dxy**, optional: Horizontal resolutions (see :func:`barotropic_geostrophic_velocity`).
:Return: KE at T points.
"""
# Init and get velocities
if cdms2.isVariable(sshuv): # from SSH
ke = sshuv*MV2.masked
u, v = barotropic_geostrophic_velocity(sshuv, dxy=dxy, gravity=gravity, format_axes=format_axes)
if format_axes is None: format_axes = False
else: # from (U,V)
u, v = sshuv
ke = u*MV2.masked
if format_axes is None: format_axes = True
gridt = shiftgrid(u.getGrid(), jshift=1)
set_grid(ke, gridt)
# Sum contributions
uf = u.filled(0.)
vf = v.filled(0.)
ke[..., 1:, :] = uf[..., 1:, :]**2
ke[..., 1:, :] += uf[..., :-1, :]**2
ke[..., 1:] += vf[..., :-1]**2
ke[..., 1:] += vf[..., :-1]**2
# Weight and mask
count = N.zeros(ke.shape, 'i')
gu = 1-N.ma.getmaskarray(u).astype('i')
gv = 1-N.ma.getmaskarray(v).astype('i')
count[1:] = gu[:-1]
count[1:] += gu[1:]
count[:, 1:] += gv[:, :-1]
count[:, 1:] += gv[:, 1:]
del gu, gv
mask = count==0
count[mask] = 1
ke[:] /= count
ke[:] = MV2.masked_where(mask, ke, copy=0)
del mask, count
# Format
if format_axes:
format_grid(gridt, 't')
return format_var(ke, "ke", format_axes=False)
def barotropic_geostrophic_velocity(ssh,
dxy=None,
gravity=default_gravity,
cyclic=False,
format_axes=True,
getu=True,
getv=True,
filter=None):
"""Get barotropic geostropic velocity from SSH on a C-grid
.. note:: ssh is supposed to be at T points,
ubt is computed at V points,
and vbt is computed at U points.
.. todo:: Rewrite it using :mod:`vacumm.data.misc.arakawa` and defining
a limited number of algorithms for different staggering configurations.
:Params:
- **ssh**: Sea surface height.
- **dxy**, optional: Horizontal resolutions (m).
Resolution along X and Y are respectively at U and V points.
Possible forms:
- ``res``: A scalar meaning a constant resolution along X and Y.
- ``(dx,dy)``: A tuple of resolutions along X and Y.
- ``None``: Resolution is estimated using
:func:`~vacumm.misc.grid.misc.resol`.
:Return: ``(ubt,vbt)``
"""
if not getu and not getv: return
# Init masked
if getu: ugbt = format_var(ssh * MV2.masked, 'ugbt', format_axes=False)
if getv: vgbt = format_var(ssh * MV2.masked, 'vgbt', format_axes=False)
# Grid
tgrid = ssh.getGrid()
if getu: ugrid = shiftgrid(tgrid, ishift=0.5)
if getv: vgrid = shiftgrid(tgrid, jshift=0.5)
if format_axes:
if getv: format_grid(ugrid, 'u')
if getu: format_grid(vgrid, 'v')
if getu: set_grid(ugbt, vgrid)
if getv: set_grid(vgbt, ugrid)
# Resolutions
if dxy is None:
dxt, dyt = resol(ssh, proj=True, mode='local')
dxu = 0.5 * (dxt[:, 1:] + dxt[:, :-1])
del dxt
dyv = 0.5 * (dyt[1:, :] + dyt[:-1, :])
del dyt
elif not isinstance(dxy, (list, tuple)):
dxu = dyv = dxy
else:
dxu, dyv = dxy
if getv and isinstance(dxu, N.ndarray):
if cdms2.isVariable(dxu): dxu = dxu.asma()
if dxu.ndim == 1: dxu.shape = 1, -1
if dxu.shape[1] == ssh.shape[-1]:
dxu = dxu[:, :-1]
if getu and isinstance(dyv, N.ndarray):
if cdms2.isVariable(dyv): dyv = dyv.asma()
if dyv.ndim == 1: dyv.shape = -1, 1
if dyv.shape[0] == ssh.shape[-2]:
dyv = dyv[:-1]
# Get geostrophic factor
f0 = coriolis_parameter(ssh, gravity=gravity, fromvar=True).asma()
bad = f0 == 0.
f0[bad] = 1.
f0[bad] = N.ma.masked
del bad
gf = gravity / f0
del f0
# Computes
sshm = ssh.asma()
sshm = sshm * gf
del gf
if getu:
ugbt[..., :-1, :] = -N.ma.diff(sshm, axis=-2) / dyv
del dyv
if getv:
vgbt[..., :-1] = N.ma.diff(sshm, axis=-1) / dxu
del dxu
del sshm
if getu and cyclic:
ugbt[..., -1] = ugbt[..., 0]
if not getu: return vgbt
elif not getv: return ugbt
return ugbt, vgbt
def kinetic_energy(sshuv, gravity=default_gravity, format_axes=None, dxy=None):
"""Compute kinetic energy in m2.s-2 either from SSH or velocity on C-grid
.. todo:: Rewrite it using :mod:`vacumm.data.misc.arakawa` and defining
a limited number of algorithms for different staggering configurations.
:Params:
- **sshuv**: SSH or (U,V).
- If SSH, geostrophic velocity is computed at U and V points
using :func:`barotropic_geostrophic_velocity`.
- If (U,V), velocities are supposed to be at V and U points.
- **dxy**, optional: Horizontal resolutions (see :func:`barotropic_geostrophic_velocity`).
:Return: KE at T points.
"""
# Init and get velocities
if cdms2.isVariable(sshuv): # from SSH
ke = sshuv * MV2.masked
u, v = barotropic_geostrophic_velocity(sshuv,
dxy=dxy,
gravity=gravity,
format_axes=format_axes)
if format_axes is None: format_axes = False
else: # from (U,V)
u, v = sshuv
ke = u * MV2.masked
if format_axes is None: format_axes = True
gridt = shiftgrid(u.getGrid(), jshift=1)
set_grid(ke, gridt)
# Sum contributions
uf = u.filled(0.)
vf = v.filled(0.)
ke[..., 1:, :] = uf[..., 1:, :]**2
ke[..., 1:, :] += uf[..., :-1, :]**2
ke[..., 1:] += vf[..., 1:]**2
ke[..., 1:] += vf[..., :-1]**2
# Weight and mask
count = N.zeros(ke.shape, 'i')
gu = 1 - N.ma.getmaskarray(u).astype('i')
gv = 1 - N.ma.getmaskarray(v).astype('i')
count[1:] = gu[:-1]
count[1:] += gu[1:]
count[:, 1:] += gv[:, :-1]
count[:, 1:] += gv[:, 1:]
del gu, gv
mask = count == 0
count[mask] = 1
ke[:] /= count
ke[:] = MV2.masked_where(mask, ke, copy=0)
del mask, count
# Format
if format_axes:
format_grid(gridt, 't')
return format_var(ke, "ke", format_axes=False)
def barotropic_geostrophic_velocity(ssh, dxy=None, gravity=default_gravity, cyclic=False,
format_axes=True, getu=True, getv=True, filter=None):
"""Get barotropic geostropic velocity from SSH on a C-grid
.. note:: ssh is supposed to be at T points,
ubt is computed at V points,
and vbt is computed at U points.
.. todo:: Rewrite it using :mod:`vacumm.data.misc.arakawa` and defining
a limited number of algorithms for different staggering configurations.
:Params:
- **ssh**: Sea surface height.
- **dxy**, optional: Horizontal resolutions (m).
Resolution along X and Y are respectively at U and V points.
Possible forms:
- ``res``: A scalar meaning a constant resolution along X and Y.
- ``(dx,dy)``: A tuple of resolutions along X and Y.
- ``None``: Resolution is estimated using
:func:`~vacumm.misc.grid.misc.resol`.
:Return: ``(ubt,vbt)``
"""
if not getu and not getv: return
# Init masked
if getu: ugbt = format_var(ssh*MV2.masked, 'ugbt', format_axes=False)
if getv: vgbt = format_var(ssh*MV2.masked, 'vgbt', format_axes=False)
# Grid
tgrid = ssh.getGrid()
if getu: ugrid = shiftgrid(tgrid, ishift=0.5)
if getv: vgrid = shiftgrid(tgrid, jshift=0.5)
if format_axes:
if getv: format_grid(ugrid, 'u')
if getu: format_grid(vgrid, 'v')
if getu: set_grid(ugbt, vgrid)
if getv: set_grid(vgbt, ugrid)
# Resolutions
if dxy is None:
dxt, dyt = resol(ssh, proj=True, mode='local')
dxu = 0.5*(dxt[:, 1:]+dxt[:, :-1]) ; del dxt
dyv = 0.5*(dyt[1:, :]+dyt[:-1, :]) ; del dyt
elif not isinstance(dxy, (list, tuple)):
dxu = dyv = dxy
else:
dxu, dyv = dxy
if getv and isinstance(dxu, N.ndarray):
if cdms2.isVariable(dxu): dxu = dxu.asma()
if dxu.ndim==1: dxu.shape = 1, -1
if dxu.shape[1]==ssh.shape[-1]:
dxu = dxu[:, :-1]
if getu and isinstance(dyv, N.ndarray):
if cdms2.isVariable(dyv): dyv = dyv.asma()
if dyv.ndim==1: dyv.shape = -1, 1
if dyv.shape[0]==ssh.shape[-2]:
dyv = dyv[:-1]
# Get geostrophic factor
f0 = coriolis_parameter(ssh, gravity=gravity, fromvar=True).asma()
bad = f0==0.
f0[bad] = 1.
f0[bad] = N.ma.masked ; del bad
gf = gravity/f0 ; del f0
# Computes
sshm = ssh.asma()
sshm = sshm*gf ; del gf
if getu: ugbt[..., :-1, :] = -N.ma.diff(sshm, axis=-2)/dyv ; del dyv
if getv: vgbt[..., :-1] = N.ma.diff(sshm, axis=-1)/dxu ; del dxu
del sshm
if getu and cyclic:
ugbt[..., -1] = ugbt[..., 0]
if not getu: return vgbt
elif not getv: return ugbt
return ugbt, vgbt
