def density(temp,
sal,
depth=None,
lat=None,
potential=False,
getdepth=False,
getlat=False,
format_axes=False):
"""Compute density from temperature, salinity and depth (and latitude)
:Params:
- **temp**: Insitu or potential temperature.
- **sal**: Salinity.
- **depth**, optional: Depth at temperature and salinty points.
Assumed to be 0 if not found.
- **lat**, optional: Latitude. Error when not found.
- **potential**, optional: True to get the potential density (at atmospheric
pressure).
:Algo:
>>> pressure = seawater.csiro.pres(depth, lat)
>>> density = seawater.csiro.dens(sal, temp, depth)
"""
# Compute
if not potential and depth is not False: # In-situ
# Get depth and latitude
lat = grow_lat(temp, lat, mode='raise', getvar=False)
if lat is None: raise VACUMMError('No latitude found for density')
depth = grow_depth(temp, depth, mode='raise', getvar=False)
if N.abs(depth.max()) < N.abs(depth.min()): # positive
depth = -depth
if (depth.asma() < 0).any():
depth = depth - depth.min() # top=0
# Get density
pres = sw_pres(depth, lat)
dens = sw_dens(sal, temp, pres)
del pres
else: # Potential
dens = sw_dens0(sal, temp)
getdepth = getlat = False
# Format
dens.setAxisList(temp.getAxisList())
set_grid(dens, get_grid(temp))
format_var(dens, 'dens', format_axes=format_axes)
# Out
if not getdepth and not getlat: return dens
dens = dens,
if getdepth: dens += depth,
if getlat: dens += lat,
return dens
def density(temp, sal, depth=None, lat=None, potential=False,
getdepth=False, getlat=False, format_axes=False):
"""Compute density from temperature, salinity and depth (and latitude)
:Params:
- **temp**: Insitu or potential temperature.
- **sal**: Salinity.
- **depth**, optional: Depth at temperature and salinty points.
Assumed to be 0 if not found.
- **lat**, optional: Latitude. Error when not found.
- **potential**, optional: True to get the potential density (at atmospheric
pressure).
:Algo:
>>> pressure = seawater.csiro.pres(depth, lat)
>>> density = seawater.csiro.dens(sal, temp, depth)
"""
# Compute
if not potential and depth is not False: # In-situ
# Get depth and latitude
lat = grow_lat(temp, lat, mode='raise', getvar=False)
if lat is None: raise VACUMMError('No latitude found for density')
depth = grow_depth(temp, depth, mode='raise', getvar=False)
if N.abs(depth.max())<N.abs(depth.min()): # positive
depth = -depth
if (depth.asma()<0).any():
depth = depth-depth.min() # top=0
# Get density
pres = sw_pres(depth, lat)
dens = sw_dens(sal, temp, pres) ; del pres
else: # Potential
dens = sw_dens0(sal, temp)
getdepth = getlat = False
# Format
dens.setAxisList(temp.getAxisList())
set_grid(dens, get_grid(temp))
format_var(dens, 'dens', format_axes=format_axes)
# Out
if not getdepth and not getlat: return dens
dens = dens,
if getdepth: dens += depth,
if getlat: dens += lat,
return dens
def ws2w(us, vs, rhoa=1.25, cd=0.016, format_axes=False, alongxy=None):
"""Convert from wind stress to 10m wind components
This function is the reverse one of :func:`wind_stress`.
Output variables are formatted using :func:`~vacumm.data.cf.format_var`.
:Formula:
.. math::
U_{10m} = (\\rho_a C_d)^{-\\frac{1}{2}}
(\\tau_x^2+\\tau_y^2)^{-\\frac{1}{4}} \\tau_x
V_{10m} = (\\rho_a C_d)^{-\\frac{1}{2}}
(\\tau_x^2+\\tau_y^2)^{-\\frac{1}{4}} \\tau_y
:Params:
- **us/vs**: Wind stress components.
- **rhoa**, optional: Air density (in kg.m-3).
- **cd**, optional: Drag coefficient.
- **format_axes**, optional: Also format axes using
:func:`~vacumm.data.cf.format_axis`.
- **alongxy**, optional: Format variables considering components are along X/Y
direction and not along zonal/meridional direction.
:Return: ``u, v``
"""
# Init and format variables
u = MV2.asarray(us).clone()
v = MV2.asarray(vs).clone()
if alongxy is None:
alongxy = hasattr(u, 'long_name') and 'x' in us.long_name.lower()
if alongxy:
uname, vname = 'ux10m', 'vy10m'
else:
uname, vname = 'u10m', 'v10m'
format_var(u, uname, format_axes=format_axes)
format_var(v, vname, format_axes=format_axes)
# Compute
zero = us.filled(1)==0.
zero &= vs.filled(1)==0.
uvsmod = (us**2+vs**2)**-0.25
uvsmod /= N.sqrt(rhoa*cd)
u.assignValue(uvsmod*us)
v.assignValue(uvsmod*vs)
u.assignValue(MV2.where(zero, 0., u))
v.assignValue(MV2.where(zero, 0., v))
del uvsmod
return u, v
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 wind_stress(u, v, rhoa=1.25, cd=0.016, format_axes=False, alongxy=None):
"""Compute the sea surface zonal and meridional wind stress from 10m wind components
Output variables are formatted using :func:`~vacumm.data.cf.format_var`.
:Formula:
.. math::
\\tau_x = \\rho_a C_d (U_{10m}^2+V_{10m}^2)^{\\frac{1}{2}}U_{10m}
\\tau_y = \\rho_a C_d (U_{10m}^2+V_{10m}^2)^{\\frac{1}{2}}V_{10m}
:Params:
- **u/v**: Wind at 10m above sea surface.
- **rhoa**, optional: Air density (in kg.m-3).
- **cd**, optional: Drag coefficient.
- **format_axes**, optional: Also format axes using
:func:`~vacumm.data.cf.format_axis`.
- **alongxy**, optional: Format variables considering components are along X/Y
direction and not along zonal/meridional direction.
:Return: ``us, vs``
"""
# Init and format variables
us = MV2.asarray(u).clone()
vs = MV2.asarray(v).clone()
if alongxy is None:
alongxy = hasattr(u, 'long_name') and 'x' in u.long_name.lower()
if alongxy:
usname, vsname = 'taux', 'tauy'
else:
usname, vsname = 'tauu', 'tauv'
format_var(us, usname, format_axes=format_axes)
format_var(vs, vsname, format_axes=format_axes)
# Compute
uvmod = N.ma.sqrt(u**2+v**2)
us.assignValue(cd*rhoa*uvmod*u)
vs.assignValue(cd*rhoa*uvmod*v)
del uvmod
return us, vs
def coriolis_parameter(lat,
gravity=default_gravity,
fromvar=False,
format_axes=False):
"""Get the coriolis parameters computed at each latitude
:Params:
- **lat**: Latitude or a variable with latitude coordinates.
- **gravity**, optional: Gravity.
- **fromvar**, optional: If True, lat is supposed to be a MV2
array with latitude coordinates.
"""
# Latitude
if fromvar:
if not cdms2.isVariable(lat):
raise VACUMMError('lat must a MV2 array because fromvar is True')
latv = lat * 0
lat = lat.getLatitude()
if lat is None:
raise VACUMMError(
'lat must a MV2 array with a latitude axis because fromvar is True'
)
if cdms2.isVariable(lat): lat = lat.asma() # 2D axes
if lat.shape != latv.shape:
if len(lat.shape) == 2:
latv[:] = N.ma.resize(lat, latv.shape)
else:
yaxis = latv.getOrder().index('y')
new_shape = len(latv.shape) * [1]
new_shape[yaxis] = latv.shape[yaxis]
tile_shape = list(latv.shape)
tile_shape[yaxis] = 1
latv[:] = N.tile(lat[:].reshape(new_shape), tile_shape)
else:
latv = lat if not N.ndim(lat) else lat[:]
# Compute
f0 = 2 * N.ma.sin(N.pi * latv / 180.)
# f0 *= 2*N.pi/(24.*3600.)
f0 *= 2 * N.pi / (86164.) # 86164 = sidereal day....
# Format
if N.isscalar(f0): return f0
f0 = MV2.asarray(f0)
if not fromvar and isaxis(lat) and f0.ndim == 1:
f0.setAxis(0, lat)
return format_var(f0, 'corio', format_axes=format_axes)
def coriolis_parameter(lat, gravity=default_gravity, fromvar=False, format_axes=False):
"""Get the coriolis parameters computed at each latitude
:Params:
- **lat**: Latitude or a variable with latitude coordinates.
- **gravity**, optional: Gravity.
- **fromvar**, optional: If True, lat is supposed to be a MV2
array with latitude coordinates.
"""
# Latitude
if fromvar:
if not cdms2.isVariable(lat):
raise VACUMMError('lat must a MV2 array because fromvar is True')
latv = lat*0
lat = lat.getLatitude()
if lat is None:
raise VACUMMError('lat must a MV2 array with a latitude axis because fromvar is True')
if cdms2.isVariable(lat): lat=lat.asma() # 2D axes
if lat.shape!=latv.shape:
if len(lat.shape)==2:
latv[:] = N.ma.resize(lat, latv.shape)
else:
yaxis = latv.getOrder().index('y')
new_shape = len(latv.shape)*[1]
new_shape[yaxis] = latv.shape[yaxis]
tile_shape = list(latv.shape)
tile_shape[yaxis] = 1
latv[:] = N.tile(lat[:].reshape(new_shape), tile_shape)
else:
latv = lat if N.isscalar(lat) else lat[:]
# Compute
f0 = 2*N.ma.sin(N.pi*latv/180.)
# f0 *= 2*N.pi/(24.*3600.)
f0 *= 2*N.pi/(86164.) # 86164 = sidereal day....
# Format
if N.isscalar(f0): return f0
f0 = MV2.asarray(f0)
if not fromvar and isaxis(lat) and f0.ndim==1:
f0.setAxis(0, lat)
return format_var(f0, 'corio', format_axes=format_axes)
def mixed_layer_depth(data, depth=None, lat=None, zaxis=None,
mode=None, deltatemp=.2, deltadens=.01, kzmax=0.0005,
potential=True, format_axes=False):
"""Get mixed layer depth from temperature and salinity
:Params:
- **temp**: Insitu or potential temperature.
- **sal**: Salinity.
- **depth**, optional: Depth at temperature and salinty points.
- **lat**, optional: Latitude.
- **mode**, optional: ``"deltatemp"``, ``"deltadens"``, ``"kz"``
or ``"twolayers"``
:Raise: :class:`~vacumm.VACUMMError` if can't get depth (and latitude for density).
"""
# TODO: positive up
# Inspection
if isinstance(data, tuple): # data = temp,sal
temp, sal=data
# Get density
if mode!='deltatemp':
res = density(temp, sal, depth=depth, lat=lat,
format_axes=False, potential=potential, getdepth=True)
if isinstance(res, tuple):
dens, depth = res
else:
dens = res
dens = dens.asma()
if mode is None:
mode = 'deltadens'
else:
temp = data[0]
# Check mode
if mode == 'kz':
warn("Switching MLD computation mode to 'deltadens'")
mode = "deltadens"
elif match_var(data, 'temp', mode='nslu'):
if mode is not None and mode!='deltatemp':
warn("Switching MLD computation mode to 'deltatemp'")
mode = 'deltatemp'
temp = data
elif match_var(data, 'dens', mode='nslu'):
if mode in ['kz', 'deltatemp']:
warn("Switching MLD computation mode to 'deltadens'")
mode = None
if mode is None:
mode = "deltadens"
dens = data
elif match_var(data, 'kz', mode='nslu'):
if mode is None:
mode = "kz"
if mode != "kz":
warn("Switching MLD computation mode to 'kz'")
kz = data
else:
if mode in ['deltadens', 'twolayers']:
dens = data
elif mode == "deltatemp":
temp = data
elif mode == "kz":
kz = data
elif mode is not None:
raise VACUMMError("Invalid MLD computation mode : '%s'"%mode)
else:
raise VACUMMError("Can't guess MLD computation mode")
temp = delta
# Find Z dim
data0 = data[0] if isinstance(data, tuple) else data
depth = grow_depth(data0, depth, mode='raise', getvar=False)
zaxis = get_zdim(data0, axis=zaxis)
if zaxis is None:
raise VACUMMError("Can't guess zaxis")
slices = get_axis_slices(data0, zaxis)
# Init MLD
axes = data0.getAxisList()
del axes[zaxis]
mld = MV2.array(data0.asma()[slices['first']], copy=1, axes=axes, copyaxes=False)
set_grid(mld, get_grid(data0))
format_var(mld, 'mld', format_axes=format_axes)
mld[:] = MV2.masked
# Two-layers
if mode=='twolayers':
densbot = dens[slices['first']]
denstop = dens[slices['last']]
del dens
H = 1.5*depth[slices['first']] - 0.5*depth[slices['firstp1']]
H = -1.5*depth[slices['last']] + 0.5*depth[slices['lastm1']]
mld[:] = -H*(densbot-denstop)/(densbot-denstop)
del H
elif mode=='deltadens':
denscrit = dens[slices['last']]+deltadens
mld[:] = -_val2z_(dens, depth, denscrit, zaxis, -1)
del dens
elif mode=='deltatemp':
tempcrit = temp[slices['last']]-deltatemp
mld[:] = -_val2z_(temp, depth, tempcrit, zaxis, 1)
elif mode=='kz':
mld[:] = -_valmin2z_(kz, depth, kzmax, zaxis, 1)
else:
raise VACUMMError("Invalid mode for computing MLD (%s)."%mode +
"Please choose one of: deltadens, twolayers")
# Mask zeros
mld[:] = MV2.masked_values(mld, 0., copy=0)
return mld
#!/usr/bin/env python # -*- coding: utf8 -*- """Utilitaires et conventions de formatage """ from vcmq import MV2, cdms2, N from vacumm.data.cf import VAR_SPECS, format_var, match_var, format_axis, AXIS_SPECS, GENERIC_VAR_NAMES # Création d'une variable 1D sst = MV2.arange(5.) # -> VERIFIER LES INFOS DE CET VARIABLE sst.getAxis(0).designateLongitude() # -> VERIFIER LES INFOS DE CET AXE # Formatage format_var(sst, 'sst') # -> VERIFIER LES NOUVELLES INFOS DE VAR+AXE # -> FORMATER L'AXE EN LATITUDE AU POINT U # Verification print match_var(sst,'sst')
def mixed_layer_depth(data,
depth=None,
lat=None,
zaxis=None,
mode=None,
deltatemp=.2,
deltadens=.03,
kzmax=0.0005,
potential=True,
format_axes=False):
"""Get mixed layer depth from temperature and salinity
:Params:
- **temp**: Insitu or potential temperature.
- **sal**: Salinity.
- **depth**, optional: Depth at temperature and salinty points.
- **lat**, optional: Latitude.
- **mode**, optional: ``"deltatemp"``, ``"deltadens"``, ``"kz"``
or ``"twolayers"``
:Raise: :class:`~vacumm.VACUMMError` if can't get depth (and latitude for density).
"""
# TODO: positive up
# Inspection
if isinstance(data, tuple): # data = temp,sal
temp, sal = data
# Get density
if mode != 'deltatemp':
res = density(temp,
sal,
depth=depth,
lat=lat,
format_axes=False,
potential=potential,
getdepth=True)
if isinstance(res, tuple):
dens, depth = res
else:
dens = res
dens = dens.asma()
if mode is None:
mode = 'deltadens'
else:
temp = data[0]
# Check mode
if mode == 'kz':
warn("Switching MLD computation mode to 'deltadens'")
mode = "deltadens"
elif match_var(data, 'temp', mode='nslu'):
if mode is not None and mode != 'deltatemp':
warn("Switching MLD computation mode to 'deltatemp'")
mode = 'deltatemp'
temp = data
elif match_var(data, 'dens', mode='nslu'):
if mode in ['kz', 'deltatemp']:
warn("Switching MLD computation mode to 'deltadens'")
mode = None
if mode is None:
mode = "deltadens"
dens = data
elif match_var(data, 'kz', mode='nslu'):
if mode is None:
mode = "kz"
if mode != "kz":
warn("Switching MLD computation mode to 'kz'")
kz = data
else:
if mode in ['deltadens', 'twolayers']:
dens = data
elif mode == "deltatemp":
temp = data
elif mode == "kz":
kz = data
elif mode is not None:
raise VACUMMError("Invalid MLD computation mode : '%s'" % mode)
else:
raise VACUMMError("Can't guess MLD computation mode")
temp = delta
# Find Z dim
data0 = data[0] if isinstance(data, tuple) else data
depth = grow_depth(data0, depth, mode='raise', getvar=False)
zaxis = get_zdim(data0, axis=zaxis)
if zaxis is None:
raise VACUMMError("Can't guess zaxis")
slices = get_axis_slices(data0, zaxis)
# Init MLD
axes = data0.getAxisList()
del axes[zaxis]
mld = MV2.array(data0.asma()[slices['first']],
copy=1,
axes=axes,
copyaxes=False)
set_grid(mld, get_grid(data0))
format_var(mld, 'mld', format_axes=format_axes)
mld[:] = MV2.masked
# Two-layers
if mode == 'twolayers':
densbot = dens[slices['first']]
denstop = dens[slices['last']]
del dens
H = 1.5 * depth[slices['first']] - 0.5 * depth[slices['firstp1']]
H = -1.5 * depth[slices['last']] + 0.5 * depth[slices['lastm1']]
mld[:] = -H * (densbot - denstop) / (densbot - denstop)
del H
elif mode == 'deltadens':
denscrit = dens[slices['last']] + deltadens
mld[:] = -_val2z_(dens, depth, denscrit, zaxis, -1)
del dens
elif mode == 'deltatemp':
tempcrit = temp[slices['last']] - deltatemp
mld[:] = -_val2z_(temp, depth, tempcrit, zaxis, 1)
elif mode == 'kz':
mld[:] = -_valmin2z_(kz, depth, kzmax, zaxis, 1)
else:
raise VACUMMError("Invalid mode for computing MLD (%s)." % mode +
"Please choose one of: deltadens, twolayers")
# Mask zeros
mld[:] = MV2.masked_values(mld, 0., copy=0)
return mld
#!/usr/bin/env python # -*- coding: utf8 -*- """Utilitaires et conventions de formatage """ from vcmq import MV2, cdms2, N from vacumm.data.cf import VAR_SPECS, format_var, match_var, format_axis, AXIS_SPECS, GENERIC_VAR_NAMES # Création d'une variable 1D sst = MV2.arange(5.) # -> VERIFIER LES INFOS DE CET VARIABLE sst.getAxis(0).designateLongitude() # -> VERIFIER LES INFOS DE CET AXE # Formatage format_var(sst, 'sst') # -> VERIFIER LES NOUVELLES INFOS DE VAR+AXE # -> FORMATER L'AXE EN LATITUDE AU POINT U # Verification print match_var(sst, 'sst')
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
