ENH: Replace use of numpy.gradient wrapper

Use our own first_derivative and gradient functions. Some tests needed
to be updated to account for changes to how values at the edge of the
grid are treated.

metpy/calc/kinematics.py

@@ -10,7 +10,7 @@
 import numpy as np
 from pyproj import Geod
 
-from ..calc.tools import get_layer_heights
+from ..calc.tools import get_layer_heights, first_derivative, gradient
 from ..cbook import is_string_like, iterable
 from ..constants import Cp_d, g
 from ..deprecation import deprecated
@@ -19,29 +19,10 @@
 
 exporter = Exporter(globals())
 
-
-def _gradient(f, *args, **kwargs):
-    """Wrap :func:`numpy.gradient` to handle units."""
-    if len(args) < f.ndim:
-        args = list(args)
-        args.extend([units.Quantity(1., 'dimensionless')] * (f.ndim - len(args)))
-    grad = np.gradient(f, *(a.magnitude for a in args), **kwargs)
-    if f.ndim == 1:
-        return units.Quantity(grad, f.units / args[0].units)
-    return [units.Quantity(g, f.units / dx.units) for dx, g in zip(args, grad)]
-
-
 def _stack(arrs):
     return concatenate([a[np.newaxis] for a in arrs], axis=0)
 
 
-def _get_gradients(u, v, dx, dy):
-    """Return derivatives for components to simplify calculating divergence and vorticity."""
-    dudy, dudx = _gradient(u, dy, dx)
-    dvdy, dvdx = _gradient(v, dy, dx)
-    return dudx, dudy, dvdx, dvdy
-
-
 def _is_x_first_dim(dim_order):
     """Determine whether x is the first dimension based on the value of dim_order."""
     if dim_order is None:
@@ -112,8 +93,6 @@ def wrapper(*args, **kwargs):
 def vorticity(u, v, dx, dy):
     r"""Calculate the vertical vorticity of the horizontal wind.
 
-    The grid must have a constant spacing in each direction.
-
     Parameters
     ----------
     u : (M, N) ndarray
@@ -135,7 +114,8 @@ def vorticity(u, v, dx, dy):
     divergence, divergence_vorticity
 
     """
-    _, dudy, dvdx, _ = _get_gradients(u, v, dx, dy)
+    dudy = first_derivative(u, delta=dy, axis=0)
+    dvdx = first_derivative(v, delta=dx, axis=1)
     return dvdx - dudy
 
 
@@ -155,8 +135,6 @@ def v_vorticity(u, v, dx, dy, dim_order='xy'):
 def divergence(u, v, dx, dy):
     r"""Calculate the horizontal divergence of the horizontal wind.
 
-    The grid must have a constant spacing in each direction.
-
     Parameters
     ----------
     u : (M, N) ndarray
@@ -178,7 +156,8 @@ def divergence(u, v, dx, dy):
     vorticity, divergence_vorticity
 
     """
-    dudx, _, _, dvdy = _get_gradients(u, v, dx, dy)
+    dudx = first_derivative(u, delta=dx, axis=1)
+    dvdy = first_derivative(v, delta=dy, axis=0)
     return dudx + dvdy
 
 
@@ -200,8 +179,6 @@ def h_convergence(u, v, dx, dy, dim_order='xy'):
 def convergence_vorticity(u, v, dx, dy, dim_order='xy'):
     r"""Calculate the horizontal divergence and vertical vorticity of the horizontal wind.
 
-    The grid must have a constant spacing in each direction.
-
     Parameters
     ----------
     u : (M, N) ndarray
@@ -228,7 +205,10 @@ def convergence_vorticity(u, v, dx, dy, dim_order='xy'):
     the horizontal divergence and vertical vorticity separately.
 
     """
-    dudx, dudy, dvdx, dvdy = _get_gradients(u, v, dx, dy)
+    dudx = first_derivative(u, delta=dx, axis=1)
+    dudy = first_derivative(u, delta=dy, axis=0)
+    dvdx = first_derivative(v, delta=dx, axis=1)
+    dvdy = first_derivative(v, delta=dy, axis=0)
     return dudx + dvdy, dvdx - dudy
 
 
@@ -237,8 +217,6 @@ def convergence_vorticity(u, v, dx, dy, dim_order='xy'):
 def shearing_deformation(u, v, dx, dy):
     r"""Calculate the shearing deformation of the horizontal wind.
 
-    The grid must have a constant spacing in each direction.
-
     Parameters
     ----------
     u : (M, N) ndarray
@@ -260,7 +238,8 @@ def shearing_deformation(u, v, dx, dy):
     stretching_deformation, total_deformation
 
     """
-    _, dudy, dvdx, _ = _get_gradients(u, v, dx, dy)
+    dudy = first_derivative(u, delta=dy, axis=0)
+    dvdx = first_derivative(v, delta=dx, axis=1)
     return dvdx + dudy
 
 
@@ -269,8 +248,6 @@ def shearing_deformation(u, v, dx, dy):
 def stretching_deformation(u, v, dx, dy):
     r"""Calculate the stretching deformation of the horizontal wind.
 
-    The grid must have a constant spacing in each direction.
-
     Parameters
     ----------
     u : (M, N) ndarray
@@ -304,8 +281,6 @@ def stretching_deformation(u, v, dx, dy):
 def shearing_stretching_deformation(u, v, dx, dy):
     r"""Calculate the horizontal shearing and stretching deformation of the horizontal wind.
 
-    The grid must have a constant spacing in each direction.
-
     Parameters
     ----------
     u : (M, N) ndarray
@@ -327,7 +302,10 @@ def shearing_stretching_deformation(u, v, dx, dy):
     shearing_deformation, stretching_deformation
 
     """
-    dudx, dudy, dvdx, dvdy = _get_gradients(u, v, dx, dy)
+    dudx = first_derivative(u, delta=dx, axis=1)
+    dudy = first_derivative(u, delta=dy, axis=0)
+    dvdx = first_derivative(v, delta=dx, axis=1)
+    dvdy = first_derivative(v, delta=dy, axis=0)
     return dvdx + dudy, dudx - dvdy
 
 
@@ -336,8 +314,6 @@ def shearing_stretching_deformation(u, v, dx, dy):
 def total_deformation(u, v, dx, dy):
     r"""Calculate the horizontal total deformation of the horizontal wind.
 
-    The grid must have a constant spacing in each direction.
-
     Parameters
     ----------
     u : (M, N) ndarray
@@ -359,7 +335,10 @@ def total_deformation(u, v, dx, dy):
     shearing_deformation, stretching_deformation
 
     """
-    dudx, dudy, dvdx, dvdy = _get_gradients(u, v, dx, dy)
+    dudx = first_derivative(u, delta=dx, axis=1)
+    dudy = first_derivative(u, delta=dy, axis=0)
+    dvdx = first_derivative(v, delta=dx, axis=1)
+    dvdy = first_derivative(v, delta=dy, axis=0)
     return np.sqrt((dvdx + dudy)**2 + (dudx - dvdy)**2)
 
 
@@ -403,7 +382,7 @@ def advection(scalar, wind, deltas):
     # Gradient returns a list of derivatives along each dimension. We convert
     # this to an array with dimension as the first index. Reverse the deltas to line up
     # with the order of the dimensions.
-    grad = _stack(_gradient(scalar, *deltas[::-1]))
+    grad = _stack(gradient(scalar, deltas=deltas[::-1]))
 
     # Make them be at least 2D (handling the 1D case) so that we can do the
     # multiply and sum below
@@ -455,8 +434,8 @@ def frontogenesis(thta, u, v, dx, dy, dim_order='yx'):
 
     """
     # Get gradients of potential temperature in both x and y
-    grad = _gradient(thta, dy, dx)
-    ddy_thta, ddx_thta = grad[-2:]  # Throw away unused gradient components
+    ddy_thta = first_derivative(thta, delta=dy, axis=-2)
+    ddx_thta = first_derivative(thta, delta=dx, axis=-1)
 
     # Compute the magnitude of the potential temperature gradient
     mag_thta = np.sqrt(ddx_thta**2 + ddy_thta**2)
@@ -505,16 +484,9 @@ def geostrophic_wind(heights, f, dx, dy):
     else:
         norm_factor = g / f
 
-    # If heights has more than 2 dimensions, we need to pass in some dummy
-    # grid deltas so that we can still use np.gradient. It may be better to
-    # to loop in this case, but that remains to be done.
-    deltas = [dy, dx]
-    if heights.ndim > 2:
-        deltas = [units.Quantity(1., units.m)] * (heights.ndim - 2) + deltas
-
-    grad = _gradient(heights, *deltas)
-    dy, dx = grad[-2:]  # Throw away unused gradient components
-    return -norm_factor * dy, norm_factor * dx
+    dhdy = first_derivative(heights, delta=dy, axis=-2)
+    dhdx = first_derivative(heights, delta=dx, axis=-1)
+    return -norm_factor * dhdy, norm_factor * dhdx
 
 
 @exporter.export

metpy/calc/tests/test_kinematics.py

@@ -66,8 +66,8 @@ def test_vorticity_divergence_asym():
     with pytest.warns(MetpyDeprecationWarning):
         c, vort = convergence_vorticity(u, v, 1 * units.meters, 2 * units.meters,
                                         dim_order='yx')
-    true_c = np.array([[0., 4., 0.], [1., 0.5, -0.5], [2., 0., 5.]]) / units.sec
-    true_vort = np.array([[-1., 2., 7.], [3.5, -1.5, -6.], [-2., 0., 1.]]) / units.sec
+    true_c = np.array([[-2, 5.5, -2.5], [2., 0.5, -1.5], [3., -1.5, 8.5]]) / units.sec
+    true_vort = np.array([[-2.5, 3.5, 13.], [8.5, -1.5, -11.], [-5.5, -1.5, 0.]]) / units.sec
     assert_array_equal(c, true_c)
     assert_array_equal(vort, true_vort)
 
@@ -102,7 +102,7 @@ def test_vorticity_asym():
     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, dim_order='yx')
-    true_vort = np.array([[-1., 2., 7.], [3.5, -1.5, -6.], [-2., 0., 1.]]) / units.sec
+    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)
 
     # Now try for xy ordered
@@ -133,7 +133,7 @@ def test_divergence_asym():
     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, dim_order='yx')
-    true_c = np.array([[0., 4., 0.], [1., 0.5, -0.5], [2., 0., 5.]]) / units.sec
+    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)
 
     # Now try for xy ordered
@@ -185,8 +185,8 @@ def test_shst_deformation_asym():
     with pytest.warns(MetpyDeprecationWarning):
         sh, st = shearing_stretching_deformation(u, v, 1 * units.meters, 2 * units.meters,
                                                  dim_order='yx')
-    true_sh = np.array([[-3., 0., 1.], [4.5, -0.5, -6.], [2., 4., 7.]]) / units.sec
-    true_st = np.array([[4., 2., 8.], [3., 1.5, 0.5], [2., 4., -1.]]) / units.sec
+    true_sh = np.array([[-7.5, -1.5, 1.], [9.5, -0.5, -11.], [1.5, 5.5, 12.]]) / units.sec
+    true_st = np.array([[4., 0.5, 12.5], [4., 1.5, -0.5], [1., 5.5, -4.5]]) / units.sec
     assert_array_equal(sh, true_sh)
     assert_array_equal(st, true_st)
 
@@ -203,7 +203,7 @@ def test_shearing_deformation_asym():
     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')
     sh = shearing_deformation(u, v, 1 * units.meters, 2 * units.meters, dim_order='yx')
-    true_sh = np.array([[-3., 0., 1.], [4.5, -0.5, -6.], [2., 4., 7.]]) / units.sec
+    true_sh = np.array([[-7.5, -1.5, 1.], [9.5, -0.5, -11.], [1.5, 5.5, 12.]]) / units.sec
     assert_array_equal(sh, true_sh)
 
     # Now try for yx ordered
@@ -217,7 +217,7 @@ def test_stretching_deformation_asym():
     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')
     st = stretching_deformation(u, v, 1 * units.meters, 2 * units.meters, dim_order='yx')
-    true_st = np.array([[4., 2., 8.], [3., 1.5, 0.5], [2., 4., -1.]]) / units.sec
+    true_st = np.array([[4., 0.5, 12.5], [4., 1.5, -0.5], [1., 5.5, -4.5]]) / units.sec
     assert_array_equal(st, true_st)
 
     # Now try for yx ordered
@@ -232,8 +232,8 @@ def test_total_deformation_asym():
     v = np.array([[6, 4, 8], [2, 6, 0], [2, 2, 6]]) * units('m/s')
     tdef = total_deformation(u, v, 1 * units.meters, 2 * units.meters,
                              dim_order='yx')
-    true_tdef = np.array([[5., 2., 8.06225775], [5.40832691, 1.58113883, 6.02079729],
-                          [2.82842712, 5.65685425, 7.07106781]]) / units.sec
+    true_tdef = np.array([[8.5, 1.58113883, 12.5399362], [10.30776406, 1.58113883, 11.0113578],
+                          [1.80277562, 7.7781746, 12.8160056]]) / units.sec
     assert_almost_equal(tdef, true_tdef)
 
     # Now try for xy ordered
@@ -249,9 +249,9 @@ def test_frontogenesis_asym():
     theta = np.array([[303, 295, 305], [308, 310, 312], [299, 293, 289]]) * units('K')
     fronto = frontogenesis(theta, u, v, 1 * units.meters, 2 * units.meters,
                            dim_order='yx')
-    true_fronto = np.array([[-20.93890452, -7.83070042, -36.43293256],
-                            [0.89442719, -2.12218672, -8.94427191],
-                            [-16.8, -7.65600391, -61.65921479]]
+    true_fronto = np.array([[-52.4746386, -37.3658646, -50.3996939],
+                            [3.5777088, -2.1221867, -16.9941166],
+                            [-23.1417334, 26.0499143, -158.4839684]]
                            ) * units.K / units.meter / units.sec
     assert_almost_equal(fronto, true_fronto)
 
@@ -303,7 +303,7 @@ def test_advection_2d():
     v = 2 * np.ones((3, 3)) * units('m/s')
     s = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) * units.kelvin
     a = advection(s, [u, v], (1 * units.meter, 1 * units.meter), dim_order='xy')
-    truth = np.array([[-3, -2, 1], [-4, 0, 4], [-1, 2, 3]]) * units('K/sec')
+    truth = np.array([[-6, -4, 2], [-8, 0, 8], [-2, 4, 6]]) * units('K/sec')
     assert_array_equal(a, truth)
 
 
@@ -313,7 +313,7 @@ def test_advection_2d_asym():
     v = 2 * u
     s = np.array([[1, 2, 4], [4, 8, 4], [8, 6, 4]]) * units.kelvin
     a = advection(s, [u, v], (2 * units.meter, 1 * units.meter), dim_order='yx')
-    truth = np.array([[0, -12.75, -2], [-27., -16., 10.], [-42, 35, 8]]) * units('K/sec')
+    truth = np.array([[0, -20.75, -2.5], [-33., -16., 20.], [-48, 91., 8]]) * units('K/sec')
     assert_array_equal(a, truth)
 
     # Now try xy ordered
@@ -327,7 +327,7 @@ def test_geostrophic_wind():
     # Using g as the value for f allows it to cancel out
     ug, vg = geostrophic_wind(z, g.magnitude / units.sec,
                               100. * units.meter, 100. * units.meter, dim_order='xy')
-    true_u = np.array([[-1, 0, 1]] * 3) * units('m/s')
+    true_u = np.array([[-2, 0, 2]] * 3) * units('m/s')
     true_v = -true_u.T
     assert_array_equal(ug, true_u)
     assert_array_equal(vg, true_v)
@@ -339,8 +339,8 @@ def test_geostrophic_wind_asym():
     # Using g as the value for f allows it to cancel out
     ug, vg = geostrophic_wind(z, g.magnitude / units.sec,
                               200. * units.meter, 100. * units.meter, dim_order='yx')
-    true_u = -np.array([[6, 12, 0], [7, 4, 0], [8, -4, 0]]) * units('m/s')
-    true_v = np.array([[1, 1.5, 2], [4, 0, -4], [-2, -2, -2]]) * units('m/s')
+    true_u = -np.array([[5, 20, 0], [7, 4, 0], [9, -12, 0]]) * units('m/s')
+    true_v = np.array([[0.5, 1.5, 2.5], [8, 0, -8], [-2, -2, -2]]) * units('m/s')
     assert_array_equal(ug, true_u)
     assert_array_equal(vg, true_v)
 
@@ -356,7 +356,7 @@ def test_geostrophic_geopotential():
     z = np.array([[48, 49, 48], [49, 50, 49], [48, 49, 48]]) * 100. * units('m^2/s^2')
     ug, vg = geostrophic_wind(z, 1 / units.sec, 100. * units.meter, 100. * units.meter,
                               dim_order='xy')
-    true_u = np.array([[-1, 0, 1]] * 3) * units('m/s')
+    true_u = np.array([[-2, 0, 2]] * 3) * units('m/s')
     true_v = -true_u.T
     assert_array_equal(ug, true_u)
     assert_array_equal(vg, true_v)
@@ -369,7 +369,7 @@ def test_geostrophic_3d():
     z3d = np.dstack((z, z)) * units.meter
     ug, vg = geostrophic_wind(z3d, g.magnitude / units.sec,
                               100. * units.meter, 100. * units.meter, dim_order='xy')
-    true_u = np.array([[-1, 0, 1]] * 3) * units('m/s')
+    true_u = np.array([[-2, 0, 2]] * 3) * units('m/s')
     true_v = -true_u.T
 
     true_u = concatenate((true_u[..., None], true_u[..., None]), axis=2)