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.

Sign up for GitHub

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

Jump to bottom

Make grid args for divergence and vorticity keyword-only #1521

Merged
merged 4 commits into from
Oct 2, 2020
Merged

Make grid args for divergence and vorticity keyword-only #1521

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2dd4912
Make grid args for divergence and vorticity keyword-only
jthielen Oct 1, 2020
17d5821
Remove _stack
jthielen Oct 1, 2020
db22ca5
DOC: Fix missing periods.
dopplershift Oct 2, 2020
724d9f5
MNT: Silence Codacy's complaints on these tests
dopplershift Oct 2, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
32 changes: 14 additions & 18 deletions src/metpy/calc/kinematics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,21 +8,17 @@
from .tools import first_derivative, get_layer_heights, gradient
from .. import constants as mpconsts
from ..package_tools import Exporter
from ..units import check_units, concatenate, units
from ..units import check_units, units
from ..xarray import add_grid_arguments_from_xarray, preprocess_and_wrap

exporter = Exporter(globals())


def _stack(arrs):
return concatenate([a[np.newaxis] if np.iterable(a) else a for a in arrs], axis=0)


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

Parameters
Expand All @@ -34,17 +30,17 @@ def vorticity(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2):
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.
latitude/longitude coordinates used as input. 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.
latitude/longitude coordinates used as input. 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`.
parsed from input if using `xarray.DataArray`. Keyword-only argument
dopplershift marked this conversation as resolved.
Show resolved Hide resolved
y_dim : int, optional
Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically
parsed from input if using `xarray.DataArray`.
parsed from input if using `xarray.DataArray`. Keyword-only argument
dopplershift marked this conversation as resolved.
Show resolved Hide resolved

Returns
-------
Expand All @@ -65,7 +61,7 @@ def vorticity(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2):
@add_grid_arguments_from_xarray
@preprocess_and_wrap(wrap_like='u')
@check_units(dx='[length]', dy='[length]')
def divergence(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2):
def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2):
r"""Calculate the horizontal divergence of a vector.

Parameters
Expand All @@ -77,17 +73,17 @@ def divergence(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2):
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.
latitude/longitude coordinates used as input. 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.
latitude/longitude coordinates used as input. 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`.
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`.
parsed from input if using `xarray.DataArray`. Keyword-only argument.

Returns
-------
Expand Down Expand Up @@ -369,7 +365,7 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim
tdef = total_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim)

# Get the divergence of the wind field
div = divergence(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim)
div = divergence(u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim)

# Compute the angle (beta) between the wind field and the gradient of potential temperature
psi = 0.5 * np.arctan2(shrd, strd)
Expand Down Expand Up @@ -641,7 +637,7 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2

"""
f = coriolis_parameter(latitude)
relative_vorticity = vorticity(u, v, dx, dy)
relative_vorticity = vorticity(u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim)
return relative_vorticity + f


Expand Down
32 changes: 24 additions & 8 deletions tests/calc/test_kinematics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
def test_default_order():
"""Test using the default array ordering."""
u = np.ones((3, 3)) * units('m/s')
v = vorticity(u, u, 1 * units.meter, 1 * units.meter)
v = vorticity(u, u, dx=1 * units.meter, dy=1 * units.meter)
true_v = np.zeros_like(u) / units.sec
assert_array_equal(v, true_v)

Expand All @@ -32,7 +32,7 @@ def test_zero_vorticity():
"""Test vorticity calculation when zeros should be returned."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
v = vorticity(u.T, u, 1 * units.meter, 1 * units.meter)
v = vorticity(u.T, u, dx=1 * units.meter, dy=1 * units.meter)
true_v = np.zeros_like(u) / units.sec
assert_array_equal(v, true_v)

Expand All @@ -41,7 +41,7 @@ def test_vorticity():
"""Test vorticity for simple case."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
v = vorticity(u.T, u.T, 1 * units.meter, 1 * units.meter)
v = vorticity(u.T, u.T, dx=1 * units.meter, dy=1 * units.meter)
true_v = np.ones_like(u) / units.sec
assert_array_equal(v, true_v)

Expand All @@ -50,16 +50,24 @@ def test_vorticity_asym():
"""Test vorticity calculation with a complicated field."""
u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s')
v = np.array([[6, 4, 8], [2, 6, 0], [2, 2, 6]]) * units('m/s')
vort = vorticity(u, v, 1 * units.meters, 2 * units.meters)
vort = vorticity(u, v, dx=1 * units.meters, dy=2 * units.meters)
true_vort = np.array([[-2.5, 3.5, 13.], [8.5, -1.5, -11.], [-5.5, -1.5, 0.]]) / units.sec
assert_array_equal(vort, true_vort)


def test_vorticity_positional_grid_args_failure():
"""Test that old API of positional grid arguments to vorticity fails."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
with pytest.raises(TypeError, match='too many positional arguments'):
vorticity(u.T, u, 1 * units.meter, 1 * units.meter)


def test_zero_divergence():
"""Test divergence calculation when zeros should be returned."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
c = divergence(u.T, u, 1 * units.meter, 1 * units.meter)
c = divergence(u.T, u, dx=1 * units.meter, dy=1 * units.meter)
true_c = 2. * np.ones_like(u) / units.sec
assert_array_equal(c, true_c)

Expand All @@ -68,7 +76,7 @@ def test_divergence():
"""Test divergence for simple case."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
c = divergence(u.T, u.T, 1 * units.meter, 1 * units.meter)
c = divergence(u.T, u.T, dx=1 * units.meter, dy=1 * units.meter)
true_c = np.ones_like(u) / units.sec
assert_array_equal(c, true_c)

Expand All @@ -81,7 +89,7 @@ def test_horizontal_divergence():
[[0., 0., 0.],
[0., 1., 0.],
[0., 0., 0.]]]) * units('m/s')
c = divergence(u, u, 1 * units.meter, 1 * units.meter)
c = divergence(u, u, dx=1 * units.meter, dy=1 * units.meter)
true_c = np.array([[[0., -2., 0.],
[-2., 0., 2.],
[0., 2., 0.]],
Expand All @@ -95,11 +103,19 @@ def test_divergence_asym():
"""Test divergence calculation with a complicated field."""
u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s')
v = np.array([[6, 4, 8], [2, 6, 0], [2, 2, 6]]) * units('m/s')
c = divergence(u, v, 1 * units.meters, 2 * units.meters)
c = divergence(u, v, dx=1 * units.meters, dy=2 * units.meters)
true_c = np.array([[-2, 5.5, -2.5], [2., 0.5, -1.5], [3., -1.5, 8.5]]) / units.sec
assert_array_equal(c, true_c)


def test_divergence_positional_grid_args_failure():
"""Test that old API of positional grid arguments to divergence fails."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
with pytest.raises(TypeError, match='too many positional arguments'):
divergence(u, u, 1 * units.meter, 1 * units.meter)


def test_shearing_deformation_asym():
"""Test shearing deformation calculation with a complicated field."""
u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s')
Expand Down