Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: Add shear and curvature vorticity calculations #3361

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/_templates/overrides/metpy.calc.rst
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ Dynamic/Kinematic
advection
ageostrophic_wind
coriolis_parameter
curvature_vorticity
divergence
exner_function
frontogenesis
Expand All @@ -123,6 +124,7 @@ Dynamic/Kinematic
potential_vorticity_baroclinic
potential_vorticity_barotropic
q_vector
shear_vorticity
shearing_deformation
stretching_deformation
total_deformation
Expand Down
156 changes: 156 additions & 0 deletions src/metpy/calc/kinematics.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,162 @@ def vorticity(
return dvdx - dudy


@exporter.export
@parse_grid_arguments
@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale'))
@check_units('[speed]', '[speed]', dx='[length]', dy='[length]')
def shear_vorticity(
u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2,
parallel_scale=None, meridional_scale=None
):
r"""Calculate the vertical shear vorticity of the horizontal wind.

Parameters
----------
u : (..., M, N) `xarray.DataArray` or `pint.Quantity`
x component of the wind
v : (..., M, N) `xarray.DataArray` or `pint.Quantity`
y component of the wind

Returns
-------
(..., M, N) `xarray.DataArray` or `pint.Quantity`
vertical vorticity

Other Parameters
----------------
dx : `pint.Quantity`, optional
The grid spacing(s) in the x-direction. If an array, there should be one item less than
the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
latitude/longitude coordinates used as input. Also optional if one-dimensional
longitude and latitude arguments are given for your data on a non-projected grid.
Keyword-only argument.
dy : `pint.Quantity`, optional
The grid spacing(s) in the y-direction. If an array, there should be one item less than
the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
latitude/longitude coordinates used as input. Also optional if one-dimensional
longitude and latitude arguments are given for your data on a non-projected grid.
Keyword-only argument.
x_dim : int, optional
Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically
parsed from input if using `xarray.DataArray`. Keyword-only argument.
y_dim : int, optional
Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically
parsed from input if using `xarray.DataArray`. Keyword-only argument.
parallel_scale : `pint.Quantity`, optional
Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray`
with latitude/longitude coordinates and MetPy CRS used as input. Also optional if
longitude, latitude, and crs are given. If otherwise omitted, calculation will be
carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument.
meridional_scale : `pint.Quantity`, optional
Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray`
with latitude/longitude coordinates and MetPy CRS used as input. Also optional if
longitude, latitude, and crs are given. If otherwise omitted, calculation will be
carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument.

See Also
--------
divergence, absolute_vorticity, vorticity, curvature_vorticity

Notes
-----
This implements a numerical version of the vertical shear vorticity:

.. math:: \zeta_s = \frac{1}{u^2 + v^2}(v u \frac{\partial u}{\partial x} +
v v \frac{\partial v}{\partial x} -
u u \frac{\partial u}{\partial y} -
u v \frac{\partial v}{\partial y})

If sufficient grid projection information is provided, these partial derivatives are
taken from the projection-correct derivative matrix of the vector wind. Otherwise, they
are evaluated as scalar derivatives on a Cartesian grid.
"""
dudx, dudy, dvdx, dvdy = vector_derivative(
u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, parallel_scale=parallel_scale,
meridional_scale=meridional_scale
)

return (v * u * dudx + v * v * dvdx - u * u * dudy - u * v * dvdy) / (u ** 2 + v ** 2)


@exporter.export
@parse_grid_arguments
@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale'))
@check_units('[speed]', '[speed]', dx='[length]', dy='[length]')
def curvature_vorticity(
u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2,
parallel_scale=None, meridional_scale=None
):
r"""Calculate the vertical curvature vorticity of the horizontal wind.

Parameters
----------
u : (..., M, N) `xarray.DataArray` or `pint.Quantity`
x component of the wind
v : (..., M, N) `xarray.DataArray` or `pint.Quantity`
y component of the wind

Returns
-------
(..., M, N) `xarray.DataArray` or `pint.Quantity`
vertical vorticity

Other Parameters
----------------
dx : `pint.Quantity`, optional
The grid spacing(s) in the x-direction. If an array, there should be one item less than
the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
latitude/longitude coordinates used as input. Also optional if one-dimensional
longitude and latitude arguments are given for your data on a non-projected grid.
Keyword-only argument.
dy : `pint.Quantity`, optional
The grid spacing(s) in the y-direction. If an array, there should be one item less than
the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
latitude/longitude coordinates used as input. Also optional if one-dimensional
longitude and latitude arguments are given for your data on a non-projected grid.
Keyword-only argument.
x_dim : int, optional
Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically
parsed from input if using `xarray.DataArray`. Keyword-only argument.
y_dim : int, optional
Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically
parsed from input if using `xarray.DataArray`. Keyword-only argument.
parallel_scale : `pint.Quantity`, optional
Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray`
with latitude/longitude coordinates and MetPy CRS used as input. Also optional if
longitude, latitude, and crs are given. If otherwise omitted, calculation will be
carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument.
meridional_scale : `pint.Quantity`, optional
Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray`
with latitude/longitude coordinates and MetPy CRS used as input. Also optional if
longitude, latitude, and crs are given. If otherwise omitted, calculation will be
carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument.

See Also
--------
divergence, absolute_vorticity

Notes
-----
This implements a numerical version of the vertical curvature vorticity:

.. math:: \zeta_c = \frac{1}{u^2 + v^2}(u u \frac{\partial v}{\partial x} -
v v \frac{\partial u}{\partial y} -
v u \frac{\partial u}{\partial x} +
u v \frac{\partial v}{\partial y})

If sufficient grid projection information is provided, these partial derivatives are
taken from the projection-correct derivative matrix of the vector wind. Otherwise, they
are evaluated as scalar derivatives on a Cartesian grid.
"""
dudx, dudy, dvdx, dvdy = vector_derivative(
u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, parallel_scale=parallel_scale,
meridional_scale=meridional_scale
)

return (u * u * dvdx - v * v * dudy - v * u * dudx + u * v * dvdy) / (u ** 2 + v ** 2)


@exporter.export
@parse_grid_arguments
@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale'))
Expand Down
16 changes: 12 additions & 4 deletions tests/calc/test_kinematics.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,11 @@
import xarray as xr

from metpy.calc import (absolute_vorticity, advection, ageostrophic_wind, coriolis_parameter,
divergence, first_derivative, frontogenesis, geospatial_laplacian,
geostrophic_wind, inertial_advective_wind, lat_lon_grid_deltas,
montgomery_streamfunction, potential_temperature,
curvature_vorticity, divergence, first_derivative, frontogenesis,
geospatial_laplacian, geostrophic_wind, inertial_advective_wind,
lat_lon_grid_deltas, montgomery_streamfunction, potential_temperature,
potential_vorticity_baroclinic, potential_vorticity_barotropic,
q_vector, shearing_deformation, static_stability,
q_vector, shear_vorticity, shearing_deformation, static_stability,
storm_relative_helicity, stretching_deformation, total_deformation,
vorticity, wind_components)
from metpy.constants import g, Re
Expand Down Expand Up @@ -151,6 +151,14 @@ def test_vorticity_grid_pole():
assert_array_almost_equal(vort.isel(y=0), vort.isel(y=-1), decimal=9)


def test_sum_shear_curvature(basic_dataset):
"""Test the sum of shear and curvature vorticity equals the total vorticity."""
d = vorticity(basic_dataset.u, basic_dataset.v)
s = shear_vorticity(basic_dataset.u, basic_dataset.v)
c = curvature_vorticity(basic_dataset.u, basic_dataset.v)
assert_array_almost_equal(s + c, d)


def test_zero_divergence():
"""Test divergence calculation when zeros should be returned."""
a = np.arange(3)
Expand Down