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