tests/calc/test_kinematics.py
# Copyright (c) 2008,2015,2017,2018,2019 MetPy Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
"""Test the `kinematics` module."""
import numpy as np
import pytest
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,
potential_vorticity_baroclinic, potential_vorticity_barotropic,
q_vector, shearing_deformation, static_stability,
storm_relative_helicity, stretching_deformation, total_deformation,
vorticity, wind_components)
from metpy.constants import g, Re
from metpy.testing import (assert_almost_equal, assert_array_almost_equal, assert_array_equal,
get_test_data)
from metpy.units import concatenate, units
@pytest.fixture()
def basic_dataset():
"""Fixture to create a dataset for use in basic tests using xarray integration."""
lon = xr.DataArray([-100, -96, -90],
attrs={'standard_name': 'longitude', 'units': 'degrees_east'})
lat = xr.DataArray([30, 31, 33],
attrs={'standard_name': 'latitude', 'units': 'degrees_north'})
u = xr.DataArray(np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s'),
coords=(lat, lon), dims=('lat', 'lon'))
v = xr.DataArray(np.array([[6, 4, 8], [2, 6, 0], [2, 2, 6]]) * units('m/s'),
coords=(lat, lon), dims=('lat', 'lon'))
z = xr.DataArray(np.array([[1, 2, 4], [4, 8, 4], [8, 6, 4]]) * 20. * units.meter,
coords=(lat, lon), dims=('lat', 'lon'))
t = xr.DataArray(np.arange(9).reshape(3, 3) * units.kelvin, coords=(lat, lon),
dims=('lat', 'lon'))
return xr.Dataset({'u': u, 'v': v, 'height': z, 'temperature': t}).metpy.parse_cf()
def test_default_order():
"""Test using the default array ordering."""
u = np.ones((3, 3)) * units('m/s')
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)
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, dx=1 * units.meter, dy=1 * units.meter)
true_v = np.zeros_like(u) / units.sec
assert_array_equal(v, true_v)
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, dx=1 * units.meter, dy=1 * units.meter)
true_v = np.ones_like(u) / units.sec
assert_array_equal(v, true_v)
@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'),
indirect=True)
def test_vorticity_geographic(geog_data):
"""Test vorticity for simple case on geographic coordinates."""
crs, lons, lats, u, v, mx, my, dx, dy = geog_data
vort = vorticity(u, v, longitude=lons, latitude=lats, crs=crs)
# Calculate the true field using known map-correct approach
truth = (mx * first_derivative(v, delta=dx, axis=1)
- my * first_derivative(u, delta=dy, axis=0)
- (v * mx / my) * first_derivative(my, delta=dx, axis=1)
+ (u * my / mx) * first_derivative(mx, delta=dy, axis=0))
assert_array_almost_equal(vort, truth, 12)
@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'),
indirect=True)
def test_abs_vorticity_geographic(geog_data):
"""Test absolute_vorticity for simple case on geographic coordinates."""
# Generate a field of u and v on a lat/lon grid
crs, lons, lats, u, v, mx, my, dx, dy = geog_data
vort = absolute_vorticity(u, v, longitude=lons, latitude=lats[:, None], crs=crs)
# Calculate the true field using known map-correct approach
truth = ((mx * first_derivative(v, delta=dx, axis=1)
- my * first_derivative(u, delta=dy, axis=0)
- (v * mx / my) * first_derivative(my, delta=dx, axis=1)
+ (u * my / mx) * first_derivative(mx, delta=dy, axis=0)
)
+ coriolis_parameter(lats[:, None]))
assert_array_almost_equal(vort, truth, 12)
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, 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."""
# pylint: disable=too-many-function-args
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_vorticity_xarray(basic_dataset):
"""Test vorticity calculation using xarray support."""
d = vorticity(basic_dataset.u, basic_dataset.v)
truth = np.array([[2.03538383e-5, 3.02085059e-5, 9.64086237e-5],
[2.48885712e-5, 8.32454044e-6, 4.33065628e-6],
[-4.46930721e-5, -3.87823454e-5, -6.92512555e-5]]) / units.sec
truth = xr.DataArray(truth, coords=basic_dataset.coords)
assert_array_almost_equal(d, truth)
def test_vorticity_grid_pole():
"""Test vorticity consistency at a pole (#2582)."""
xy = [-25067.525, 0., 25067.525]
us = np.ones((len(xy), len(xy)))
vs = us * np.linspace(-1, 0, len(xy))[None, :]
grid = {'grid_mapping_name': 'lambert_azimuthal_equal_area',
'longitude_of_projection_origin': 0, 'latitude_of_projection_origin': 90,
'false_easting': 0, 'false_northing': 0}
x = xr.DataArray(
xy, name='x', attrs={'standard_name': 'projection_x_coordinate', 'units': 'm'})
y = xr.DataArray(
xy, name='y', attrs={'standard_name': 'projection_y_coordinate', 'units': 'm'})
u = xr.DataArray(us, name='u', coords=(y, x), dims=('y', 'x'), attrs={'units': 'm/s'})
v = xr.DataArray(vs, name='v', coords=(y, x), dims=('y', 'x'), attrs={'units': 'm/s'})
ds = xr.merge((u, v)).metpy.assign_crs(grid).metpy.assign_latitude_longitude()
vort = vorticity(ds.u, ds.v)
assert_array_almost_equal(vort.isel(y=0), vort.isel(y=-1), decimal=9)
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, dx=1 * units.meter, dy=1 * units.meter)
true_c = 2. * np.ones_like(u) / units.sec
assert_array_equal(c, true_c)
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, dx=1 * units.meter, dy=1 * units.meter)
true_c = np.ones_like(u) / units.sec
assert_array_equal(c, true_c)
@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'),
indirect=True)
def test_divergence_geographic(geog_data):
"""Test divergence for simple case on geographic coordinates."""
# Generate a field of u and v on a lat/lon grid
crs, lons, lats, u, v, mx, my, dx, dy = geog_data
div = divergence(u, v, longitude=lons, latitude=lats, crs=crs)
# Calculate the true field using known map-correct approach
truth = (mx * first_derivative(u, delta=dx, axis=1)
+ my * first_derivative(v, delta=dy, axis=0)
- (u * mx / my) * first_derivative(my, delta=dx, axis=1)
- (v * my / mx) * first_derivative(mx, delta=dy, axis=0))
assert_array_almost_equal(div, truth, 12)
def test_horizontal_divergence():
"""Test taking the horizontal divergence of a 3D field."""
u = np.array([[[1., 1., 1.],
[1., 0., 1.],
[1., 1., 1.]],
[[0., 0., 0.],
[0., 1., 0.],
[0., 0., 0.]]]) * units('m/s')
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.]],
[[0., 2., 0.],
[2., 0., -2.],
[0., -2., 0.]]]) * units('s^-1')
assert_array_equal(c, true_c)
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, 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."""
# pylint: disable=too-many-function-args
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_divergence_xarray(basic_dataset):
"""Test divergence calculation using xarray support."""
d = divergence(basic_dataset.u, basic_dataset.v)
truth = np.array([[-4.361528313e-5, 3.593794959e-5, -9.728275045e-5],
[-1.672604193e-5, 9.158127775e-6, -4.223658565e-5],
[3.011996772e-5, -3.745237046e-5, 9.570189256e-5]]) / units.sec
truth = xr.DataArray(truth, coords=basic_dataset.coords)
assert_array_almost_equal(d, truth, 4)
@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'),
indirect=True)
def test_shearing_deformation_geographic(geog_data):
"""Test shearing deformation for simple case on geographic coordinates."""
# Generate a field of u and v on a lat/lon grid
crs, lons, lats, u, v, mx, my, dx, dy = geog_data
shear = shearing_deformation(u, v, longitude=lons, latitude=lats, crs=crs)
# Calculate the true field using known map-correct approach
truth = (mx * first_derivative(v, delta=dx, axis=1)
+ my * first_derivative(u, delta=dy, axis=0)
+ (v * mx / my) * first_derivative(my, delta=dx, axis=1)
+ (u * my / mx) * first_derivative(mx, delta=dy, axis=0))
assert_array_almost_equal(shear, truth, 12)
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')
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)
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)
@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'),
indirect=True)
def test_stretching_deformation_geographic(geog_data):
"""Test divergence for simple case on geographic coordinates."""
# Generate a field of u and v on a lat/lon grid
crs, lons, lats, u, v, mx, my, dx, dy = geog_data
stretch = stretching_deformation(u, v, longitude=lons, latitude=lats, crs=crs)
# Calculate the true field using known map-correct approach
truth = (mx * first_derivative(u, delta=dx, axis=1)
- my * first_derivative(v, delta=dy, axis=0)
+ (u * mx / my) * first_derivative(my, delta=dx, axis=1)
- (v * my / mx) * first_derivative(mx, delta=dy, axis=0))
assert_array_almost_equal(stretch, truth, 12)
def test_stretching_deformation_asym():
"""Test stretching deformation 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')
st = stretching_deformation(u, v, 1 * units.meters, 2 * units.meters)
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)
def test_total_deformation_asym():
"""Test total deformation 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')
tdef = total_deformation(u, v, 1 * units.meters, 2 * units.meters)
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)
def test_frontogenesis_asym():
"""Test frontogenesis 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')
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)
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)
def test_frontogenesis_nan():
"""Test that frontogenesis calculation does not result in nan."""
x = np.array([-4142340.8, -4061069.8, -3979798.8, -3898527.8, -3817256.8],
dtype=np.float32)
y = np.array([-832207.56, -750936.56, -669665.56, -588394.5, -507123.56],
dtype=np.float32)
lat = np.array([[12.38805122, 12.58268367, 12.77387305, 12.96159447, 13.14582354],
[13.07545967, 13.27159197, 13.46425116, 13.65341235, 13.83905111],
[13.76423766, 13.96186003, 14.15597929, 14.34657051, 14.53360928],
[14.45429168, 14.65339375, 14.84896275, 15.04097373, 15.22940228],
[15.14552468, 15.3460955, 15.54310332, 15.73652321, 15.92633074]])
lon = np.array([[-132.75696788, -132.05286812, -131.34671228, -130.63852983,
-129.92835084],
[-132.9590417, -132.25156385, -131.54199837, -130.83037505,
-130.11672431],
[-133.16323241, -132.45234731, -131.73934239, -131.02424779,
-130.30709426],
[-133.36957268, -132.65525085, -131.93877637, -131.22017972,
-130.49949199],
[-133.57809517, -132.86030681, -132.14033233, -131.41820252,
-130.69394884]])
uvals = np.array([[0.165024, 0.055023, -0.454977, -1.534977, -2.744976],
[0.155024, -0.434977, -2.114977, -3.474977, -4.034977],
[-1.554977, -2.714977, -2.084976, -5.274977, -3.334976],
[-3.424976, -7.644977, -7.654977, -5.384976, -3.224977],
[-9.564977, -9.934977, -7.454977, -6.004977, -4.144977]]) * units('m/s')
vvals = (np.array([[2.6594982, 1.9194984, 2.979498, 2.149498, 2.6394978],
[3.4994984, 4.0794983, 4.8494987, 5.2594986, 3.1694984],
[5.159498, 6.4594975, 6.559498, 5.9694977, 3.189499],
[6.5994987, 9.799498, 7.4594975, 4.2894993, 3.729498],
[11.3394985, 6.779499, 4.0994987, 4.819498, 4.9994984]])
* units('m/s'))
tvals = np.array([[290.2, 290.1, 290.2, 290.30002, 290.30002],
[290.5, 290.30002, 290.30002, 290.30002, 290.2],
[290.80002, 290.40002, 290.2, 290.40002, 289.90002],
[290.7, 290.90002, 290.7, 290.1, 289.7],
[290.90002, 290.40002, 289.7, 289.7, 289.30002]]) * units('degK')
x = xr.DataArray(data=x, coords={'x': x}, dims='x', attrs={'units': 'meters'})
y = xr.DataArray(data=y, coords={'y': y}, dims='y', attrs={'units': 'meters'})
dims = ['y', 'x']
coords = {'x': x, 'y': y, 'latitude': (dims, lat), 'longitude': (dims, lon)}
u = xr.DataArray(data=uvals, coords=coords, dims=dims)
v = xr.DataArray(data=vvals, coords=coords, dims=dims)
t = xr.DataArray(data=tvals, coords=coords, dims=dims)
th = potential_temperature(850 * units('hPa'), t)
f = frontogenesis(th, u, v)
assert not np.any(np.isnan(f))
def test_advection_uniform():
"""Test advection calculation for a uniform 1D field."""
u = np.ones((3,)) * units('m/s')
s = np.ones_like(u) * units.kelvin
a = advection(s.T, u.T, dx=1 * units.meter)
truth = np.zeros_like(u) * units('K/sec')
assert_array_equal(a, truth)
def test_advection_1d_uniform_wind():
"""Test advection for simple 1D case with uniform wind."""
u = np.ones((3,)) * units('m/s')
s = np.array([1, 2, 3]) * units('kg')
a = advection(s.T, u.T, dx=1 * units.meter)
truth = -np.ones_like(u) * units('kg/sec')
assert_array_equal(a, truth)
def test_advection_1d():
"""Test advection calculation with varying wind and field."""
u = np.array([1, 2, 3]) * units('m/s')
s = np.array([1, 2, 3]) * units('Pa')
a = advection(s.T, u.T, dx=1 * units.meter)
truth = np.array([-1, -2, -3]) * units('Pa/sec')
assert_array_equal(a, truth)
def test_advection_2d_uniform():
"""Test advection for uniform 2D field."""
u = np.ones((3, 3)) * units('m/s')
s = np.ones_like(u) * units.kelvin
a = advection(s.T, u.T, u.T, dx=1 * units.meter, dy=1 * units.meter)
truth = np.zeros_like(u) * units('K/sec')
assert_array_equal(a, truth)
def test_advection_2d():
"""Test advection in varying 2D field."""
u = np.ones((3, 3)) * units('m/s')
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.T, v.T, u.T, dx=1 * units.meter, dy=1 * units.meter)
truth = np.array([[-6, -4, 2], [-8, 0, 8], [-2, 4, 6]]) * units('K/sec')
assert_array_equal(a, truth)
def test_advection_z_y():
"""Test advection in varying 2D z-y field."""
v = 2 * np.ones((3, 3)) * units('m/s')
w = np.ones((3, 3)) * units('m/s')
s = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) * units.kelvin
a = advection(s.T, v=v.T, w=w.T,
dy=1 * units.meter, dz=1 * units.meter,
y_dim=-1, vertical_dim=-2)
truth = np.array([[-6, -4, 2], [-8, 0, 8], [-2, 4, 6]]) * units('K/sec')
assert_array_equal(a, truth)
def test_advection_4d_vertical(data_4d):
"""Test 4-d vertical advection with parsed dims."""
data_4d['w'] = -abs(data_4d['u'])
data_4d['w'].attrs['units'] = 'Pa/s'
a = advection(data_4d.temperature, w=data_4d.w)
assert (a < 0).sum() == 0
assert a.data.units == units.Unit('K/sec')
@pytest.mark.filterwarnings('ignore:Horizontal dimension numbers not found.')
def test_advection_1d_vertical():
"""Test 1-d vertical advection with parsed dims."""
pressure = xr.DataArray(
np.array([1000., 950., 900.]), dims='pressure', attrs={'units': 'hPa'})
omega = xr.DataArray(
np.array([20., 30., 40.]),
coords=[pressure], dims=['pressure'], attrs={'units': 'hPa/sec'})
s = xr.DataArray(
np.array([25., 20., 15.]),
coords=[pressure], dims=['pressure'], attrs={'units': 'degC'})
a = advection(s, w=omega)
truth = xr.DataArray(
-np.array([2, 3, 4]) * units('K/sec'), coords=[pressure], dims=['pressure'])
assert_array_almost_equal(a, truth)
def test_advection_2d_asym():
"""Test advection in asymmetric varying 2D field."""
u = np.arange(9).reshape(3, 3) * units('m/s')
v = 2 * u
s = np.array([[1, 2, 4], [4, 8, 4], [8, 6, 4]]) * units.kelvin
a = advection(s, u, v, dx=2 * units.meter, dy=1 * units.meter)
truth = np.array([[0, -20.75, -2.5], [-33., -16., 20.], [-48, 91., 8]]) * units('K/sec')
assert_array_equal(a, truth)
@pytest.mark.filterwarnings('ignore:Vertical dimension number not found.')
def test_advection_xarray(basic_dataset):
"""Test advection calculation using xarray support."""
a = advection(basic_dataset.temperature, basic_dataset.u, basic_dataset.v)
truth = np.array([[-0.0001953019, -0.000135269, -0.000262247],
[-4.51008510e-5, -0.000139840, -2.44370835e-6],
[-2.11362160e-5, -2.29201236e-5, -3.70146905e-5]]) * units('K/sec')
truth = xr.DataArray(truth, coords=basic_dataset.coords)
assert_array_almost_equal(a, truth, 4)
def test_geostrophic_wind():
"""Test geostrophic wind calculation with basic conditions."""
z = np.array([[48, 49, 48], [49, 50, 49], [48, 49, 48]]) * 10. * units.meter
latitude = 30 * units.degrees
ug, vg = geostrophic_wind(
z.T,
100. * units.kilometer,
100. * units.kilometer,
latitude
)
true_u = np.array([[-26.897, 0, 26.897]] * 3) * units('m/s')
true_v = -true_u.T
assert_array_almost_equal(ug, true_u.T, 2)
assert_array_almost_equal(vg, true_v.T, 2)
def test_geostrophic_wind_asym():
"""Test geostrophic wind calculation with a complicated field."""
z = np.array([[1, 2, 4], [4, 8, 4], [8, 6, 4]]) * 20. * units.meter
latitude = 30 * units.degrees
ug, vg = geostrophic_wind(
z,
2000. * units.kilometer,
1000. * units.kilometer,
latitude
)
true_u = np.array(
[[-6.724, -26.897, 0], [-9.414, -5.379, 0], [-12.103, 16.138, 0]]
) * units('m/s')
true_v = np.array(
[[0.672, 2.017, 3.362], [10.759, 0, -10.759], [-2.690, -2.690, -2.690]]
) * units('m/s')
assert_array_almost_equal(ug, true_u, 2)
assert_array_almost_equal(vg, true_v, 2)
def test_geostrophic_geopotential():
"""Test geostrophic wind calculation with geopotential."""
z = np.array([[48, 49, 48], [49, 50, 49], [48, 49, 48]]) * 100. * units('m^2/s^2')
latitude = 30 * units.degrees
ug, vg = geostrophic_wind(
z.T,
100. * units.kilometer,
100. * units.kilometer,
latitude
)
true_u = np.array([[-27.427, 0, 27.427]] * 3) * units('m/s')
true_v = -true_u.T
assert_array_almost_equal(ug, true_u.T, 2)
assert_array_almost_equal(vg, true_v.T, 2)
def test_geostrophic_3d():
"""Test geostrophic wind calculation with 3D array."""
z = np.array([[48, 49, 48], [49, 50, 49], [48, 49, 48]]) * 10.
latitude = 30 * units.degrees
z3d = np.dstack((z, z)) * units.meter
ug, vg = geostrophic_wind(
z3d.T,
100. * units.kilometer,
100. * units.kilometer,
latitude
)
true_u = np.array([[-26.897, 0, 26.897]] * 3) * units('m/s')
true_v = -true_u.T
true_u = concatenate((true_u[..., None], true_u[..., None]), axis=2)
true_v = concatenate((true_v[..., None], true_v[..., None]), axis=2)
assert_array_almost_equal(ug, true_u.T, 2)
assert_array_almost_equal(vg, true_v.T, 2)
def test_geostrophic_gempak():
"""Test of geostrophic wind calculation against gempak values."""
z = np.array([[5586387.00, 5584467.50, 5583147.50],
[5594407.00, 5592487.50, 5591307.50],
[5604707.50, 5603247.50, 5602527.50]]).T * (9.80616 * units('m/s^2')) * 1e-3
dx = np.deg2rad(0.25) * Re * np.cos(np.deg2rad(44))
dy = -np.deg2rad(0.25) * Re
lat = 44 * units.degrees
ug, vg = geostrophic_wind(z.T * units.m, dx.T, dy.T, lat)
true_u = np.array([[21.97512, 21.97512, 22.08005],
[31.89402, 32.69477, 33.73863],
[38.43922, 40.18805, 42.14609]])
true_v = np.array([[-10.93621, -7.83859, -4.54839],
[-10.74533, -7.50152, -3.24262],
[-8.66612, -5.27816, -1.45282]])
assert_almost_equal(ug[1, 1], true_u[1, 1] * units('m/s'), 2)
assert_almost_equal(vg[1, 1], true_v[1, 1] * units('m/s'), 2)
def test_no_ageostrophic_geopotential():
"""Test the updated ageostrophic wind function."""
z = np.array([[48, 49, 48], [49, 50, 49], [48, 49, 48]]) * 100. * units('m^2/s^2')
u = np.array([[-27.427, 0, 27.427]] * 3) * units('m/s')
v = -u.T
latitude = 30 * units.degrees
uag, vag = ageostrophic_wind(z.T, u.T, v.T, 100. * units.kilometer,
100. * units.kilometer, latitude)
true = np.array([[0, 0, 0]] * 3) * units('m/s')
assert_array_almost_equal(uag, true.T, 2)
assert_array_almost_equal(vag, true.T, 2)
def test_ageostrophic_geopotential():
"""Test ageostrophic wind calculation with future input variable order."""
z = np.array([[48, 49, 48], [49, 50, 49], [48, 49, 48]]) * 100. * units('m^2/s^2')
u = v = np.array([[0, 0, 0]] * 3) * units('m/s')
latitude = 30 * units.degrees
uag, vag = ageostrophic_wind(z.T, u.T, v.T, 100. * units.kilometer,
100. * units.kilometer, latitude)
u_true = np.array([[27.427, 0, -27.427]] * 3) * units('m/s')
v_true = -u_true.T
assert_array_almost_equal(uag, u_true.T, 2)
assert_array_almost_equal(vag, v_true.T, 2)
def test_streamfunc():
"""Test of Montgomery Streamfunction calculation."""
t = 287. * units.kelvin
hgt = 5000. * units.meter
msf = montgomery_streamfunction(hgt, t)
assert_almost_equal(msf, 337372.45469 * units('m^2 s^-2'), 4)
def test_storm_relative_helicity_no_storm_motion():
"""Test storm relative helicity with no storm motion and differing input units."""
u = np.array([0, 20, 10, 0]) * units('m/s')
v = np.array([20, 0, 0, 10]) * units('m/s')
u = u.to('knots')
heights = np.array([0, 250, 500, 750]) * units.m
positive_srh, negative_srh, total_srh = storm_relative_helicity(heights, u, v,
depth=750 * units.meters)
assert_almost_equal(positive_srh, 400. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(negative_srh, -100. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(total_srh, 300. * units('meter ** 2 / second ** 2 '), 6)
def test_storm_relative_helicity_storm_motion():
"""Test storm relative helicity with storm motion and differing input units."""
u = np.array([5, 25, 15, 5]) * units('m/s')
v = np.array([30, 10, 10, 20]) * units('m/s')
u = u.to('knots')
heights = np.array([0, 250, 500, 750]) * units.m
pos_srh, neg_srh, total_srh = storm_relative_helicity(heights, u, v,
depth=750 * units.meters,
storm_u=5 * units('m/s'),
storm_v=10 * units('m/s'))
assert_almost_equal(pos_srh, 400. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(neg_srh, -100. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(total_srh, 300. * units('meter ** 2 / second ** 2 '), 6)
def test_storm_relative_helicity_with_interpolation():
"""Test storm relative helicity with interpolation."""
u = np.array([-5, 15, 25, 15, -5]) * units('m/s')
v = np.array([40, 20, 10, 10, 30]) * units('m/s')
u = u.to('knots')
heights = np.array([0, 100, 200, 300, 400]) * units.m
pos_srh, neg_srh, total_srh = storm_relative_helicity(heights, u, v,
bottom=50 * units.meters,
depth=300 * units.meters,
storm_u=5 * units('m/s'),
storm_v=10 * units('m/s'))
assert_almost_equal(pos_srh, 400. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(neg_srh, -100. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(total_srh, 300. * units('meter ** 2 / second ** 2 '), 6)
def test_storm_relative_helicity():
"""Test function for SRH calculations on an eigth-circle hodograph."""
# Create larger arrays for everything except pressure to make a smoother graph
hgt_int = np.arange(0, 2050, 50)
hgt_int = hgt_int * units('meter')
dir_int = np.arange(180, 272.25, 2.25)
spd_int = np.zeros(hgt_int.shape[0])
spd_int[:] = 2.
u_int, v_int = wind_components(spd_int * units('m/s'), dir_int * units.degree)
# Put in the correct value of SRH for a eighth-circle, 2 m/s hodograph
# (SRH = 2 * area under hodo, in this case...)
srh_true_p = (.25 * np.pi * (2 ** 2)) * units('m^2/s^2')
# Since there's only positive SRH in this case, total SRH will be equal to positive SRH and
# negative SRH will be zero.
srh_true_t = srh_true_p
srh_true_n = 0 * units('m^2/s^2')
p_srh, n_srh, t_srh = storm_relative_helicity(hgt_int, u_int, v_int,
1000 * units('meter'),
bottom=0 * units('meter'),
storm_u=0 * units.knot,
storm_v=0 * units.knot)
assert_almost_equal(p_srh, srh_true_p, 2)
assert_almost_equal(n_srh, srh_true_n, 2)
assert_almost_equal(t_srh, srh_true_t, 2)
def test_storm_relative_helicity_agl():
"""Test storm relative helicity with heights above ground."""
u = np.array([-5, 15, 25, 15, -5]) * units('m/s')
v = np.array([40, 20, 10, 10, 30]) * units('m/s')
u = u.to('knots')
heights = np.array([100, 200, 300, 400, 500]) * units.m
pos_srh, neg_srh, total_srh = storm_relative_helicity(heights, u, v,
bottom=50 * units.meters,
depth=300 * units.meters,
storm_u=5 * units('m/s'),
storm_v=10 * units('m/s'))
# Check that heights isn't modified--checks for regression of #789
assert_almost_equal(heights[0], 100 * units.m, 6)
assert_almost_equal(pos_srh, 400. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(neg_srh, -100. * units('meter ** 2 / second ** 2 '), 6)
assert_almost_equal(total_srh, 300. * units('meter ** 2 / second ** 2 '), 6)
def test_storm_relative_helicity_masked():
"""Test that srh does not return masked values."""
h = units.Quantity(np.ma.array([20.72, 234.85, 456.69, 683.21]), units.meter)
u = units.Quantity(np.ma.array(np.zeros((4,))), units.knot)
v = units.Quantity(np.zeros_like(u), units.knot)
pos, neg, com = storm_relative_helicity(h, u, v, depth=500 * units.meter,
storm_u=15.77463015050421 * units('m/s'),
storm_v=21.179437759755647 * units('m/s'))
assert not np.ma.is_masked(pos)
assert not np.ma.is_masked(neg)
assert not np.ma.is_masked(com)
def test_absolute_vorticity_asym():
"""Test absolute 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')
lats = np.array([[30, 30, 30], [20, 20, 20], [10, 10, 10]]) * units.degrees
vort = absolute_vorticity(u, v, 1 * units.meters, 2 * units.meters, lats)
true_vort = np.array([[-2.499927, 3.500073, 13.00007],
[8.500050, -1.499950, -10.99995],
[-5.499975, -1.499975, 2.532525e-5]]) / units.sec
assert_almost_equal(vort, true_vort, 5)
@pytest.fixture
def pv_data():
"""Test data for all PV testing."""
u = np.array([[[100, 90, 80, 70],
[90, 80, 70, 60],
[80, 70, 60, 50],
[70, 60, 50, 40]],
[[100, 90, 80, 70],
[90, 80, 70, 60],
[80, 70, 60, 50],
[70, 60, 50, 40]],
[[100, 90, 80, 70],
[90, 80, 70, 60],
[80, 70, 60, 50],
[70, 60, 50, 40]]]) * units('m/s')
v = np.zeros_like(u) * units('m/s')
lats = np.array([[40, 40, 40, 40],
[40.1, 40.1, 40.1, 40.1],
[40.2, 40.2, 40.2, 40.2],
[40.3, 40.3, 40.3, 40.3]]) * units.degrees
lons = np.array([[40, 39.9, 39.8, 39.7],
[40, 39.9, 39.8, 39.7],
[40, 39.9, 39.8, 39.7],
[40, 39.9, 39.8, 39.7]]) * units.degrees
dx, dy = lat_lon_grid_deltas(lons, lats)
return u, v, lats, lons, dx, dy
def test_potential_vorticity_baroclinic_unity_axis0(pv_data):
"""Test potential vorticity calculation with unity stability and height on axis 0."""
u, v, lats, _, dx, dy = pv_data
potential_temperature = np.ones((3, 4, 4)) * units.kelvin
potential_temperature[0] = 200 * units.kelvin
potential_temperature[1] = 300 * units.kelvin
potential_temperature[2] = 400 * units.kelvin
pressure = np.ones((3, 4, 4)) * units.hPa
pressure[2] = 1000 * units.hPa
pressure[1] = 900 * units.hPa
pressure[0] = 800 * units.hPa
pvor = potential_vorticity_baroclinic(potential_temperature, pressure,
u, v, dx[None, :, :], dy[None, :, :],
lats[None, :, :])
abs_vorticity = absolute_vorticity(u, v, dx[None, :, :], dy[None, :, :],
lats[None, :, :])
vort_difference = pvor - (abs_vorticity * g * (-1 * (units.kelvin / units.hPa)))
true_vort = np.zeros_like(u) * (units.kelvin * units.meter**2
/ (units.second * units.kilogram))
assert_almost_equal(vort_difference, true_vort, 10)
def test_potential_vorticity_baroclinic_non_unity_derivative(pv_data):
"""Test potential vorticity calculation with unity stability and height on axis 0."""
u, v, lats, _, dx, dy = pv_data
potential_temperature = np.ones((3, 4, 4)) * units.kelvin
potential_temperature[0] = 200 * units.kelvin
potential_temperature[1] = 300 * units.kelvin
potential_temperature[2] = 400 * units.kelvin
pressure = np.ones((3, 4, 4)) * units.hPa
pressure[2] = 1000 * units.hPa
pressure[1] = 999 * units.hPa
pressure[0] = 998 * units.hPa
pvor = potential_vorticity_baroclinic(potential_temperature, pressure,
u, v, dx[None, :, :], dy[None, :, :],
lats[None, :, :])
abs_vorticity = absolute_vorticity(u, v, dx[None, :, :], dy[None, :, :],
lats[None, :, :])
vort_difference = pvor - (abs_vorticity * g * (-100 * (units.kelvin / units.hPa)))
true_vort = np.zeros_like(u) * (units.kelvin * units.meter ** 2
/ (units.second * units.kilogram))
assert_almost_equal(vort_difference, true_vort, 10)
def test_potential_vorticity_baroclinic_wrong_number_of_levels_axis_0(pv_data):
"""Test that potential vorticity calculation errors without 3 levels on axis 0."""
u, v, lats, _, dx, dy = pv_data
potential_temperature = np.ones((3, 4, 4)) * units.kelvin
potential_temperature[0] = 200 * units.kelvin
potential_temperature[1] = 300 * units.kelvin
potential_temperature[2] = 400 * units.kelvin
pressure = np.ones((3, 4, 4)) * units.hPa
pressure[2] = 1000 * units.hPa
pressure[1] = 900 * units.hPa
pressure[0] = 800 * units.hPa
with pytest.raises(ValueError):
potential_vorticity_baroclinic(potential_temperature[:1, :, :], pressure, u, v,
dx[None, :, :], dy[None, :, :],
lats[None, :, :])
with pytest.raises(ValueError):
potential_vorticity_baroclinic(u, v, dx[None, :, :], dy[None, :, :],
lats[None, :, :], potential_temperature,
pressure[:1, :, :])
def test_potential_vorticity_baroclinic_isentropic_real_data():
"""Test potential vorticity calculation with real isentropic data."""
isentlevs = [328, 330, 332] * units.K
isentprs = np.array([[[245.88052, 245.79416, 245.68776, 245.52525, 245.31844],
[245.97734, 245.85878, 245.74838, 245.61089, 245.4683],
[246.4308, 246.24358, 246.08649, 245.93279, 245.80148],
[247.14348, 246.87215, 246.64842, 246.457, 246.32005],
[248.05727, 247.72388, 247.44029, 247.19205, 247.0112]],
[[239.66074, 239.60431, 239.53738, 239.42496, 239.27725],
[239.5676, 239.48225, 239.4114, 239.32259, 239.23781],
[239.79681, 239.6465, 239.53227, 239.43031, 239.35794],
[240.2442, 240.01723, 239.84442, 239.71255, 239.64021],
[240.85277, 240.57112, 240.34885, 240.17174, 240.0666]],
[[233.63297, 233.60493, 233.57542, 233.51053, 233.41898],
[233.35995, 233.3061, 233.27275, 233.23009, 233.2001],
[233.37685, 233.26152, 233.18793, 233.13496, 233.11841],
[233.57312, 233.38823, 233.26366, 233.18817, 233.17694],
[233.89297, 233.66039, 233.49615, 233.38635, 233.35281]]]) * units.hPa
isentu = np.array([[[28.94226812, 28.53362902, 27.98145564, 27.28696092, 26.46488305],
[28.15024259, 28.12645242, 27.95788749, 27.62007338, 27.10351611],
[26.27821641, 26.55765132, 26.7329775, 26.77170719, 26.64779014],
[24.07215607, 24.48837805, 24.86738637, 25.17622757, 25.38030319],
[22.25524153, 22.65568001, 23.07333679, 23.48542321, 23.86341343]],
[[28.50078095, 28.12605738, 27.6145395, 26.96565679, 26.1919881],
[27.73718892, 27.73189078, 27.58886228, 27.28329365, 26.80468118],
[25.943111, 26.23034592, 26.41833632, 26.47466534, 26.37320009],
[23.82858821, 24.24937503, 24.63505859, 24.95235053, 25.16669265],
[22.09498322, 22.5008718, 22.9247538, 23.34295878, 23.72623895]],
[[28.05929378, 27.71848573, 27.24762337, 26.64435265, 25.91909314],
[27.32413525, 27.33732915, 27.21983708, 26.94651392, 26.50584625],
[25.60800559, 25.90304052, 26.10369515, 26.17762349, 26.09861004],
[23.58502035, 24.01037201, 24.4027308, 24.72847348, 24.95308212],
[21.9347249, 22.34606359, 22.77617081, 23.20049435,
23.58906447]]]) * units('m/s')
isentv = np.array([[[-2.22336191, -2.82451946, -3.27190475, -3.53076527, -3.59311591],
[-2.12438321, -2.98895919, -3.73633746, -4.32254411, -4.70849598],
[-1.24050415, -2.31904635, -3.32284815, -4.20895826, -4.93036136],
[0.32254009, -0.89843808, -2.09621275, -3.2215678, -4.2290825],
[2.14238865, 0.88207403, -0.40652485, -1.67244834, -2.86837275]],
[[-1.99024801, -2.59146057, -3.04973279, -3.3296825, -3.42137476],
[-1.8711102, -2.71865804, -3.45952099, -4.05064148, -4.45309013],
[-0.99367383, -2.04299168, -3.02642031, -3.90252563, -4.62540783],
[0.547778, -0.63635567, -1.80391109, -2.90776869, -3.90375721],
[2.33967328, 1.12072805, -0.13066324, -1.3662872, -2.5404749]],
[[-1.75713411, -2.35840168, -2.82756083, -3.12859972, -3.24963361],
[-1.6178372, -2.44835688, -3.18270452, -3.77873886, -4.19768429],
[-0.7468435, -1.76693701, -2.72999246, -3.596093, -4.32045429],
[0.7730159, -0.37427326, -1.51160943, -2.59396958, -3.57843192],
[2.53695791, 1.35938207, 0.14519838, -1.06012605,
-2.21257705]]]) * units('m/s')
lats = np.array([57.5, 57., 56.5, 56., 55.5]) * units.degrees
lons = np.array([227.5, 228., 228.5, 229., 229.5]) * units.degrees
dx, dy = lat_lon_grid_deltas(lons, lats)
pvor = potential_vorticity_baroclinic(isentlevs[:, None, None], isentprs,
isentu, isentv, dx[None, :, :], dy[None, :, :],
lats[None, :, None])
true_pv = np.array([[[2.97116898e-06, 3.38486331e-06, 3.81432403e-06, 4.24722471e-06,
4.64995688e-06],
[2.04235589e-06, 2.35739554e-06, 2.71138003e-06, 3.11803005e-06,
3.54655984e-06],
[1.41179481e-06, 1.60663306e-06, 1.85439220e-06, 2.17827401e-06,
2.55309150e-06],
[1.25933892e-06, 1.31915377e-06, 1.43444064e-06, 1.63067920e-06,
1.88631658e-06],
[1.37533104e-06, 1.31658998e-06, 1.30424716e-06, 1.36777872e-06,
1.49289942e-06]],
[[3.07674708e-06, 3.48172482e-06, 3.90371030e-06, 4.33207155e-06,
4.73253199e-06],
[2.16369614e-06, 2.47112604e-06, 2.81747901e-06, 3.21722053e-06,
3.63944011e-06],
[1.53925419e-06, 1.72853221e-06, 1.97026966e-06, 2.28774012e-06,
2.65577906e-06],
[1.38675388e-06, 1.44221972e-06, 1.55296146e-06, 1.74439951e-06,
1.99486345e-06],
[1.50312413e-06, 1.44039769e-06, 1.42422805e-06, 1.48403040e-06,
1.60544869e-06]],
[[3.17979446e-06, 3.57430736e-06, 3.98713951e-06, 4.40950119e-06,
4.80650246e-06],
[2.28618901e-06, 2.58455503e-06, 2.92172357e-06, 3.31292186e-06,
3.72721632e-06],
[1.67022518e-06, 1.85294576e-06, 2.08747504e-06, 2.39710083e-06,
2.75677598e-06],
[1.51817109e-06, 1.56879550e-06, 1.67430213e-06, 1.85997008e-06,
2.10409000e-06],
[1.63449148e-06, 1.56773336e-06, 1.54753266e-06, 1.60313832e-06,
1.72018062e-06]]]) * (units.kelvin * units.meter ** 2
/ (units.second * units.kilogram))
assert_almost_equal(pvor, true_pv, 10)
def test_potential_vorticity_baroclinic_isobaric_real_data():
"""Test potential vorticity calculation with real isentropic data."""
pressure = [20000., 25000., 30000.] * units.Pa
theta = np.array([[[344.45776, 344.5063, 344.574, 344.6499, 344.735],
[343.98444, 344.02536, 344.08682, 344.16284, 344.2629],
[343.58792, 343.60876, 343.65628, 343.72818, 343.82834],
[343.21542, 343.2204, 343.25833, 343.32935, 343.43414],
[342.85272, 342.84982, 342.88556, 342.95645, 343.0634]],
[[326.70923, 326.67603, 326.63416, 326.57153, 326.49155],
[326.77695, 326.73468, 326.6931, 326.6408, 326.58405],
[326.95062, 326.88986, 326.83627, 326.78134, 326.7308],
[327.1913, 327.10928, 327.03894, 326.97546, 326.92587],
[327.47235, 327.3778, 327.29468, 327.2188, 327.15973]],
[[318.47897, 318.30374, 318.1081, 317.8837, 317.63837],
[319.155, 318.983, 318.79745, 318.58905, 318.36212],
[319.8042, 319.64206, 319.4669, 319.2713, 319.0611],
[320.4621, 320.3055, 320.13373, 319.9425, 319.7401],
[321.1375, 320.98648, 320.81473, 320.62186, 320.4186]]]) * units.K
uwnd = np.array([[[25.309322, 25.169882, 24.94082, 24.61212, 24.181437],
[24.849028, 24.964956, 24.989666, 24.898415, 24.673553],
[23.666418, 24.003235, 24.269922, 24.435743, 24.474638],
[22.219162, 22.669518, 23.09492, 23.460283, 23.731855],
[21.065105, 21.506243, 21.967466, 22.420042, 22.830257]],
[[29.227198, 28.803436, 28.23203, 27.516447, 26.670708],
[28.402836, 28.376076, 28.199024, 27.848948, 27.315084],
[26.454042, 26.739328, 26.916056, 26.952703, 26.822044],
[24.17064, 24.59482, 24.979027, 25.290913, 25.495026],
[22.297522, 22.70384, 23.125736, 23.541069, 23.921045]],
[[27.429195, 26.97554, 26.360558, 25.594944, 24.7073],
[26.959536, 26.842077, 26.56688, 26.118752, 25.50171],
[25.460867, 25.599699, 25.62171, 25.50819, 25.249628],
[23.6418, 23.920736, 24.130007, 24.255558, 24.28613],
[21.915337, 22.283215, 22.607704, 22.879448,
23.093569]]]) * units('m/s')
vwnd = np.array([[[-0.3050951, -0.90105104, -1.4307652, -1.856761, -2.156073],
[-0.10017005, -0.82312256, -1.5097888, -2.1251845, -2.631675],
[0.6832816, -0.16461015, -1.0023694, -1.7991445, -2.5169075],
[2.0360851, 1.0960612, 0.13380499, -0.81640035, -1.718524],
[3.6074955, 2.654059, 1.6466523, 0.61709386, -0.39874703]],
[[-2.3738103, -2.9788015, -3.423631, -3.6743853, -3.7226477],
[-2.2792664, -3.159968, -3.917221, -4.507328, -4.8893175],
[-1.3700132, -2.4722757, -3.4953287, -4.3956766, -5.123884],
[0.2314668, -1.0151587, -2.2366724, -3.382317, -4.403803],
[2.0903401, 0.8078297, -0.5038105, -1.7920332, -3.0061343]],
[[-1.4415079, -1.7622383, -1.9080431, -1.8903408, -1.7376306],
[-1.5708634, -2.288579, -2.823628, -3.1583376, -3.285275],
[-0.9814599, -1.999404, -2.8674111, -3.550859, -4.0168552],
[0.07641177, -1.1033016, -2.1928647, -3.1449537, -3.9159832],
[1.2759045, 0.05043932, -1.1469103, -2.264961,
-3.2550638]]]) * units('m/s')
lats = np.array([57.5, 57., 56.5, 56., 55.5]) * units.degrees
lons = np.array([227.5, 228., 228.5, 229., 229.5]) * units.degrees
dx, dy = lat_lon_grid_deltas(lons, lats)
pvor = potential_vorticity_baroclinic(theta, pressure[:, None, None],
uwnd, vwnd, dx[None, :, :], dy[None, :, :],
lats[None, :, None])
true_pv = np.array([[[4.29013406e-06, 4.61736108e-06, 4.97453387e-06, 5.36730237e-06,
5.75500645e-06],
[3.48415057e-06, 3.72492697e-06, 4.00658450e-06, 4.35128065e-06,
4.72701041e-06],
[2.87775662e-06, 3.01866087e-06, 3.21074864e-06, 3.47971854e-06,
3.79924194e-06],
[2.70274738e-06, 2.71627883e-06, 2.78699880e-06, 2.94197238e-06,
3.15685712e-06],
[2.81293318e-06, 2.70649941e-06, 2.65188277e-06, 2.68109532e-06,
2.77737801e-06]],
[[2.43090597e-06, 2.79248225e-06, 3.16783697e-06, 3.54497301e-06,
3.89481001e-06],
[1.61968826e-06, 1.88924405e-06, 2.19296648e-06, 2.54191855e-06,
2.91119712e-06],
[1.09089606e-06, 1.25384007e-06, 1.46192044e-06, 1.73476959e-06,
2.05268876e-06],
[9.72047256e-07, 1.02016741e-06, 1.11466014e-06, 1.27721014e-06,
1.49122340e-06],
[1.07501523e-06, 1.02474621e-06, 1.01290749e-06, 1.06385170e-06,
1.16674712e-06]],
[[6.10254835e-07, 7.31519400e-07, 8.55731472e-07, 9.74301226e-07,
1.08453329e-06],
[3.17052987e-07, 3.98799900e-07, 4.91789955e-07, 5.96021549e-07,
7.10773939e-07],
[1.81983099e-07, 2.26503437e-07, 2.83058115e-07, 3.56549337e-07,
4.47098851e-07],
[1.54729567e-07, 1.73825926e-07, 2.01823376e-07, 2.44513805e-07,
3.02525735e-07],
[1.55220676e-07, 1.63334569e-07, 1.76335524e-07, 1.98346439e-07,
2.30155553e-07]]]) * (units.kelvin * units.meter ** 2
/ (units.second * units.kilogram))
assert_almost_equal(pvor, true_pv, 10)
def test_potential_vorticity_baroclinic_4d(data_4d):
"""Test potential vorticity calculation with latlon+xarray spatial handling."""
theta = potential_temperature(data_4d.pressure, data_4d.temperature)
pvor = potential_vorticity_baroclinic(theta, data_4d.pressure, data_4d.u, data_4d.v)
truth = np.array([
[[[2.02341517e-07, 1.08253899e-06, 5.07866020e-07, 7.59602062e-07],
[5.10389680e-07, 6.85689387e-07, 8.21670367e-07, 7.07634816e-07],
[1.32493368e-06, 7.42556664e-07, 6.56995963e-07, 1.42860463e-06],
[3.98119942e-07, 1.44178504e-06, 1.00098404e-06, 1.32741769e-07]],
[[3.78824281e-07, 8.69275146e-07, 8.85194259e-07, 6.71317237e-07],
[6.98417346e-07, 9.07612472e-07, 9.43897715e-07, 7.86981464e-07],
[1.14118467e-06, 5.46283726e-07, 8.51417036e-07, 1.47484547e-06],
[6.09694315e-07, 8.92755943e-07, 8.21736234e-07, 2.19146777e-07]],
[[5.45372476e-07, 8.65038943e-07, 1.02542271e-06, 7.01655222e-07],
[9.09010760e-07, 1.14690318e-06, 9.52200248e-07, 8.39364616e-07],
[1.30601001e-06, 5.13731599e-07, 9.45482183e-07, 1.12678378e-06],
[1.41700436e-06, 5.34416471e-07, 5.77202761e-07, 8.00215780e-07]]],
[[[4.89875284e-07, 7.41732002e-07, 4.00156659e-07, 4.51659753e-07],
[4.92109734e-07, 5.00766168e-07, 4.65459579e-07, 6.57429624e-07],
[5.25432209e-07, 4.65439077e-07, 5.95175649e-07, 6.15264682e-07],
[5.31988096e-07, 6.02477834e-07, 5.69272740e-07, 4.23351696e-07]],
[[5.14269220e-07, 7.78503321e-07, 6.11304383e-07, 5.15249894e-07],
[4.46066171e-07, 5.87690456e-07, 5.40874995e-07, 5.20729202e-07],
[5.54138102e-07, 4.80436803e-07, 5.44944125e-07, 7.67293518e-07],
[5.50869543e-07, 5.67508510e-07, 6.15430155e-07, 7.11393271e-07]],
[[4.62763045e-07, 7.58095696e-07, 5.71561539e-07, 5.09461534e-07],
[4.00198925e-07, 5.65386246e-07, 6.59228506e-07, 5.21051149e-07],
[4.86756849e-07, 4.51122732e-07, 5.54841504e-07, 6.37263135e-07],
[4.97103017e-07, 3.76458794e-07, 3.84346823e-07, 6.33177143e-07]]],
[[[3.67414624e-07, 3.11634409e-07, 4.63243895e-07, 3.57094992e-07],
[3.09361430e-07, 3.77719588e-07, 2.44198465e-07, 4.83354174e-07],
[5.69920205e-08, 4.16754253e-07, 6.39950078e-07, 1.01328837e-07],
[2.56285156e-07, 2.35613341e-07, 4.95745172e-07, 5.31565087e-07]],
[[4.91680068e-07, 4.55365178e-07, 4.76828376e-07, 4.27773462e-07],
[3.43227964e-07, 3.21022454e-07, 2.81916434e-07, 4.21074000e-07],
[2.65819971e-07, 5.26528676e-07, 4.79102139e-07, 2.74517652e-07],
[2.22251840e-07, 3.44727929e-07, 7.41995750e-07, 4.76425941e-07]],
[[3.16830323e-07, 4.45198415e-07, 4.82149658e-07, 4.92118755e-07],
[2.47719020e-07, 1.13643951e-07, 4.11871361e-07, 4.19639595e-07],
[1.14884698e-07, 4.59177263e-07, 3.22239409e-07, 3.14475957e-07],
[6.39184081e-08, 3.11908917e-07, 6.38295102e-07, 4.58138799e-07]]]]
) * units('K * m ** 2 / s / kg')
assert_array_almost_equal(pvor, truth, 10)
def test_potential_vorticity_barotropic(pv_data):
"""Test the barotopic (Rossby) potential vorticity."""
u, v, lats, _, dx, dy = pv_data
u = u[0]
v = v[0]
heights = np.ones_like(u) * 3 * units.km
pv = potential_vorticity_barotropic(heights, u, v, dx, dy, lats)
avor = absolute_vorticity(u, v, dx, dy, lats)
truth = avor / heights
assert_almost_equal(pv, truth, 10)
def test_inertial_advective_wind_diffluent():
"""Test inertial advective wind with a diffluent flow."""
lats = np.array([[50., 50., 50., 50., 50., 50., 50., 50., 50., 50., 50.],
[48., 48., 48., 48., 48., 48., 48., 48., 48., 48., 48.],
[46., 46., 46., 46., 46., 46., 46., 46., 46., 46., 46.],
[44., 44., 44., 44., 44., 44., 44., 44., 44., 44., 44.],
[42., 42., 42., 42., 42., 42., 42., 42., 42., 42., 42.],
[40., 40., 40., 40., 40., 40., 40., 40., 40., 40., 40.],
[38., 38., 38., 38., 38., 38., 38., 38., 38., 38., 38.],
[36., 36., 36., 36., 36., 36., 36., 36., 36., 36., 36.],
[34., 34., 34., 34., 34., 34., 34., 34., 34., 34., 34.],
[32., 32., 32., 32., 32., 32., 32., 32., 32., 32., 32.],
[30., 30., 30., 30., 30., 30., 30., 30., 30., 30., 30.]]) * units.degrees
lons = np.array([[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286., 290.],
[250., 254., 258., 262., 266., 270., 274., 278., 282., 286.,
290.]]) * units.degrees
ug = np.array([[23.68206888, 23.28736773, 22.49796543, 21.70856314, 19.7350574,
17.36685051, 14.20924133, 10.26222985, 7.49932181, 4.34171263,
2.36820689],
[24.4118194, 24.00495574, 23.19122843, 22.37750111, 20.34318283,
17.90200089, 14.64709164, 10.57845507, 7.73040948, 4.47550022,
2.44118194],
[25.21967679, 24.79934884, 23.95869295, 23.11803706, 21.01639732,
18.49442965, 15.13180607, 10.92852661, 7.98623098, 4.62360741,
2.52196768],
[26.11573982, 25.68047749, 24.80995283, 23.93942817, 21.76311652,
19.15154253, 15.66944389, 11.31682059, 8.26998428, 4.78788563,
2.61157398],
[27.11207213, 26.66020426, 25.75646853, 24.85273279, 22.59339344,
19.88218623, 16.26724328, 11.74856459, 8.58548951, 4.97054656,
2.71120721],
[28.22319067, 27.75280415, 26.81203113, 25.87125811, 23.51932555,
20.69700649, 16.9339144, 12.23004929, 8.93734371, 5.17425162,
2.82231907],
[29.46670856, 28.97559675, 27.99337313, 27.01114951, 24.55559047,
21.60891961, 17.68002514, 12.76890704, 9.33112438, 5.4022299,
2.94667086],
[30.86419265, 30.34978944, 29.32098302, 28.2921766, 25.72016054,
22.63374128, 18.51851559, 13.37448348, 9.77366101, 5.65843532,
3.08641927],
[32.44232384, 31.90161845, 30.82020765, 29.73879686, 27.03526987,
23.79103749, 19.46539431, 14.05834033, 10.27340255, 5.94775937,
3.24423238],
[34.23449286, 33.66391798, 32.52276821, 31.38161845, 28.52874405,
25.10529476, 20.54069571, 14.8349469, 10.84092274, 6.27632369, 3.42344929],
[36.28303453, 35.67831729, 34.46888281, 33.25944832, 30.23586211,
26.60755866, 21.76982072, 15.7226483, 11.4896276, 6.65188966,
3.62830345]]) * units('m/s')
vg = np.array([[7.67648972e-01, 2.30294692e+00, 3.07059589e+00,
5.37354281e+00, 8.44413870e+00, 1.07470856e+01,
1.38176815e+01, 1.30500325e+01, 1.15147346e+01,
9.97943664e+00, 5.37354281e+00],
[6.08116408e-01, 1.82434923e+00, 2.43246563e+00,
4.25681486e+00, 6.68928049e+00, 8.51362972e+00,
1.09460954e+01, 1.03379789e+01, 9.12174613e+00,
7.90551331e+00, 4.25681486e+00],
[4.53862086e-01, 1.36158626e+00, 1.81544834e+00,
3.17703460e+00, 4.99248295e+00, 6.35406920e+00,
8.16951755e+00, 7.71565546e+00, 6.80793129e+00,
5.90020712e+00, 3.17703460e+00],
[3.02572579e-01, 9.07717738e-01, 1.21029032e+00,
2.11800806e+00, 3.32829837e+00, 4.23601611e+00,
5.44630643e+00, 5.14373385e+00, 4.53858869e+00,
3.93344353e+00, 2.11800806e+00],
[1.52025875e-01, 4.56077624e-01, 6.08103499e-01,
1.06418112e+00, 1.67228462e+00, 2.12836225e+00,
2.73646575e+00, 2.58443987e+00, 2.28038812e+00,
1.97633637e+00, 1.06418112e+00],
[-5.44403782e-13, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 3.62935855e-13],
[-1.55819455e-01, -4.67458366e-01, -6.23277822e-01,
-1.09073619e+00, -1.71401401e+00, -2.18147238e+00,
-2.80475020e+00, -2.64893074e+00, -2.33729183e+00,
-2.02565292e+00, -1.09073619e+00],
[-3.17940982e-01, -9.53822947e-01, -1.27176393e+00,
-2.22558688e+00, -3.49735080e+00, -4.45117375e+00,
-5.72293768e+00, -5.40499670e+00, -4.76911473e+00,
-4.13323277e+00, -2.22558688e+00],
[-4.89187491e-01, -1.46756247e+00, -1.95674996e+00,
-3.42431243e+00, -5.38106240e+00, -6.84862487e+00,
-8.80537483e+00, -8.31618734e+00, -7.33781236e+00,
-6.35943738e+00, -3.42431243e+00],
[-6.72847961e-01, -2.01854388e+00, -2.69139184e+00,
-4.70993572e+00, -7.40132757e+00, -9.41987145e+00,
-1.21112633e+01, -1.14384153e+01, -1.00927194e+01,
-8.74702349e+00, -4.70993572e+00],
[-8.72878488e-01, -2.61863546e+00, -3.49151395e+00,
-6.11014941e+00, -9.60166336e+00, -1.22202988e+01,
-1.57118128e+01, -1.48389343e+01, -1.30931773e+01,
-1.13474203e+01, -6.11014941e+00]]) * units('m/s')
uiaw_truth = np.array([[-1.42807415e+00, -8.84702475e-01, -1.16169714e+00,
-2.07178191e+00, -2.26651744e+00, -2.44307980e+00,
-2.13572115e+00, -1.07805246e+00, -7.66864343e-01,
-4.29350989e-01, 2.09863394e-01],
[-1.15466056e+00, -7.14539881e-01, -9.37868053e-01,
-1.67069607e+00, -1.82145232e+00, -1.95723406e+00,
-1.69677456e+00, -8.44795197e-01, -5.99128909e-01,
-3.31430392e-01, 1.74263065e-01],
[-8.85879800e-01, -5.47662808e-01, -7.18560490e-01,
-1.27868851e+00, -1.38965979e+00, -1.48894561e+00,
-1.28074427e+00, -6.29324678e-01, -4.45008769e-01,
-2.43265738e-01, 1.36907150e-01],
[-6.11536708e-01, -3.77851670e-01, -4.95655649e-01,
-8.81515449e-01, -9.56332337e-01, -1.02300759e+00,
-8.76092304e-01, -4.27259697e-01, -3.01610757e-01,
-1.63732062e-01, 9.57321746e-02],
[-3.20542252e-01, -1.98032517e-01, -2.59762867e-01,
-4.61930812e-01, -5.00960635e-01, -5.35715150e-01,
-4.58376210e-01, -2.23205436e-01, -1.57510625e-01,
-8.53847182e-02, 5.03061160e-02],
[-7.17595005e-13, -2.36529156e-13, -0.00000000e+00,
-0.00000000e+00, -0.00000000e+00, -0.00000000e+00,
-0.00000000e+00, -0.00000000e+00, -0.00000000e+00,
-2.93991041e-14, -6.68646808e-14],
[3.66014877e-01, 2.26381410e-01, 2.97076576e-01,
5.28911772e-01, 5.75670994e-01, 6.17640324e-01,
5.33241822e-01, 2.63664097e-01, 1.86703719e-01,
1.02644563e-01, -5.59431414e-02],
[7.98146331e-01, 4.94284987e-01, 6.48956331e-01,
1.15693361e+00, 1.26429103e+00, 1.36142948e+00,
1.18700769e+00, 5.96585261e-01, 4.23976442e-01,
2.36489429e-01, -1.18303960e-01],
[1.32422955e+00, 8.21548216e-01, 1.07935767e+00,
1.92781420e+00, 2.11849275e+00, 2.29274320e+00,
2.02576077e+00, 1.04018545e+00, 7.42658541e-01,
4.21848012e-01, -1.87693140e-01],
[1.98305622e+00, 1.23316526e+00, 1.62158055e+00,
2.90329131e+00, 3.21355659e+00, 3.50025687e+00,
3.14455187e+00, 1.65685097e+00, 1.18935849e+00,
6.89756719e-01, -2.64167230e-01],
[2.83017758e+00, 1.76405766e+00, 2.32173329e+00,
4.16683333e+00, 4.64487467e+00, 5.09076175e+00,
4.64597480e+00, 2.50595982e+00, 1.80748965e+00,
1.06712544e+00, -3.52915793e-01]]) * units('m/s')
viaw_truth = np.array([[-0.16767916, -0.49465351, -0.63718079, -1.07594125, -1.53705893,
-1.721506, -1.81093489, -1.23523645, -0.79647599, -0.39963532,
-0.11737541],
[-0.17337355, -0.51145198, -0.6588195, -1.1124803, -1.58925758,
-1.77996849, -1.87243438, -1.27718518, -0.82352438, -0.41320697,
-0.12136149],
[-0.18010801, -0.53131862, -0.68441043, -1.15569305, -1.65099008,
-1.84910889, -1.94516649, -1.32679566, -0.85551304, -0.42925742,
-0.12607561],
[-0.18806768, -0.55479966, -0.71465719, -1.20676763, -1.72395376,
-1.93082821, -2.03113097, -1.38543193, -0.89332149, -0.44822798,
-0.13164738],
[-0.19730148, -0.58203938, -0.74974564, -1.26601785, -1.80859693,
-2.02562856, -2.13085602, -1.45345426, -0.93718205, -0.4702352,
-0.13811104],
[-0.2078345, -0.61311178, -0.78977111, -1.33360472, -1.90514961,
-2.13376756, -2.24461263, -1.5310475, -0.98721389, -0.4953389,
-0.14548415],
[-0.21963486, -0.64792283, -0.83461247, -1.40932368, -2.01331954,
-2.25491789, -2.37205648, -1.6179768, -1.04326558, -0.52346308,
-0.1537444],
[-0.2325551, -0.68603755, -0.88370939, -1.49222857, -2.1317551,
-2.38756571, -2.5115951, -1.71315592, -1.10463673, -0.55425633,
-0.16278857],
[-0.24622751, -0.72637116, -0.93566454, -1.57995986, -2.25708551,
-2.52793577, -2.65925711, -1.81387599, -1.16958067, -0.58684223,
-0.17235926],
[-0.25987451, -0.76662981, -0.98752314, -1.66752812, -2.38218302,
-2.66804499, -2.80664473, -1.9144089, -1.23440393, -0.61936759,
-0.18191216],
[-0.27342538, -0.80660487, -1.03901645, -1.75447953, -2.50639932,
-2.80716724, -2.95299411, -2.01423364, -1.29877056, -0.65166382,
-0.19139777]]) * units('m/s')
dx, dy = lat_lon_grid_deltas(lons, lats)
uiaw, viaw = inertial_advective_wind(ug, vg, ug, vg, dx, dy, lats)
assert_almost_equal(uiaw, uiaw_truth, 5)
assert_almost_equal(viaw, viaw_truth, 5)
@pytest.fixture
def q_vector_data():
"""Define data for use in Q-vector tests."""
speed = np.ones((4, 4)) * 50. * units('knots')
wdir = np.array([[210., 190., 170., 150.],
[190., 180., 180., 170.],
[170., 180., 180., 190.],
[150., 170., 190., 210.]]) * units('degrees')
u, v = wind_components(speed, wdir)
temp = np.array([[[18., 18., 18., 18.],
[17., 17., 17., 17.],
[17., 17., 17., 17.],
[16., 16., 16., 16.]],
[[12., 11., 10., 9.],
[11., 10.5, 10.5, 10.],
[10., 10.5, 10.5, 11.],
[9., 10., 11., 12.]],
[[-10., -10., -10., -10.],
[-10., -10., -10., -10.],
[-11., -11., -11., -11.],
[-11., -11., -11., -11.]]]) * units('degC')
p = np.array([850., 700., 500.]) * units('hPa')
lats = np.linspace(35., 40., 4) * units('degrees')
lons = np.linspace(-100., -90., 4) * units('degrees')
dx, dy = lat_lon_grid_deltas(lons, lats)
return u, v, temp, p, dx, dy
def test_q_vector_without_static_stability(q_vector_data):
"""Test the Q-vector function without using static stability."""
u, v, temp, p, dx, dy = q_vector_data
# Treating as 700 hPa data
q1, q2 = q_vector(u, v, temp[1], p[1], dx, dy)
q1_truth = (np.array([[-2.7454089e-14, -3.0194267e-13, -3.0194267e-13, -2.7454089e-14],
[-1.8952185e-13, -2.2269905e-14, -2.2269905e-14, -1.8952185e-13],
[-1.9918390e-13, -2.3370829e-14, -2.3370829e-14, -1.9918390e-13],
[-5.6160772e-14, -3.5145951e-13, -3.5145951e-13, -5.6160772e-14]])
* units('m^2 kg^-1 s^-1'))
q2_truth = (np.array([[-4.4976059e-14, -4.3582378e-13, 4.3582378e-13, 4.4976059e-14],
[-3.0124244e-13, -3.5724617e-14, 3.5724617e-14, 3.0124244e-13],
[3.1216232e-13, 3.6662900e-14, -3.6662900e-14, -3.1216232e-13],
[8.6038280e-14, 4.6968342e-13, -4.6968342e-13, -8.6038280e-14]])
* units('m^2 kg^-1 s^-1'))
assert_almost_equal(q1, q1_truth, 16)
assert_almost_equal(q2, q2_truth, 16)
def test_q_vector_with_static_stability(q_vector_data):
"""Test the Q-vector function using static stability."""
u, v, temp, p, dx, dy = q_vector_data
sigma = static_stability(p[:, np.newaxis, np.newaxis], temp)
# Treating as 700 hPa data
q1, q2 = q_vector(u, v, temp[1], p[1], dx, dy, sigma[1])
q1_truth = (np.array([[-1.4158140e-08, -1.6197987e-07, -1.6875014e-07, -1.6010616e-08],
[-9.3971386e-08, -1.1252476e-08, -1.1252476e-08, -9.7617234e-08],
[-1.0785670e-07, -1.2403513e-08, -1.2403513e-08, -1.0364793e-07],
[-2.9186946e-08, -1.7577703e-07, -1.6937879e-07, -2.6112047e-08]])
* units('kg m^-2 s^-3'))
q2_truth = (np.array([[-2.31770213e-08, -2.33621439e-07, 2.43378967e-07, 2.62072251e-08],
[-1.4936626e-07, -1.8050836e-08, 1.8050836e-08, 1.5516129e-07],
[1.6903373e-07, 1.9457964e-08, -1.9457964e-08, -1.6243771e-07],
[4.46812456e-08, 2.34736724e-07, -2.26197708e-07, -3.99768328e-08]])
* units('kg m^-2 s^-3'))
assert_almost_equal(q1, q1_truth, 10)
assert_almost_equal(q2, q2_truth, 10)
@pytest.fixture
def data_4d():
"""Define 4D data (extracted from Irma GFS example) for testing kinematics functions."""
data = xr.open_dataset(get_test_data('irma_gfs_example.nc', False))
data = data.metpy.parse_cf()
data['Geopotential_height_isobaric'].attrs['units'] = 'm'
subset = data.drop_vars((
'LatLon_361X720-0p25S-180p00E', 'Vertical_velocity_pressure_isobaric', 'isobaric1',
'Relative_humidity_isobaric', 'reftime'
)).sel(
latitude=[46., 44., 42., 40.],
longitude=[262., 267., 272., 277.],
isobaric3=[50000., 70000., 85000.]
).isel(time1=[0, 1, 2])
return subset.rename({
'Geopotential_height_isobaric': 'height',
'Temperature_isobaric': 'temperature',
'isobaric3': 'pressure',
'u-component_of_wind_isobaric': 'u',
'v-component_of_wind_isobaric': 'v'
})
true_vort4d = np.array([[[[-5.72939079e-05, 3.36008149e-05, 4.80394116e-05, 2.24754927e-05],
[2.28437884e-05, 2.16350819e-05, 4.40912008e-05, 7.21109010e-05],
[6.56150935e-05, 7.12554707e-05, 8.63437939e-05, 8.77146299e-05],
[-4.12479588e-05, 1.60707608e-04, 1.47465661e-04, -5.63406909e-05]],
[[1.22453259e-07, 4.40258958e-05, -8.22480293e-06, 1.54493600e-05],
[1.29420183e-05, 1.25760315e-05, 2.98881935e-05, 6.40671857e-05],
[1.96998118e-05, 8.78628308e-06, 3.96330962e-05, 4.88475149e-05],
[3.37810678e-05, 2.94602756e-05, 3.98077989e-05, 5.71554040e-05]],
[[-2.76821428e-05, 5.08462417e-06, 5.55937962e-06, 5.23436098e-05],
[-2.11754797e-05, 6.40524521e-06, 2.11226065e-05, 3.52627761e-05],
[-1.92494063e-05, -1.43439529e-06, 2.99489927e-05, 3.13418677e-05],
[-2.32787494e-05, -1.76993463e-05, 3.10941039e-05, 3.53835159e-05]]],
[[[-3.57414525e-05, 2.61424456e-05, 8.46799855e-05, 2.62297854e-05],
[3.41307192e-05, 2.48272187e-05, 4.93974252e-05, 7.85589219e-05],
[7.95962242e-05, 5.17417889e-05, 6.89810168e-05, 1.03949044e-04],
[-6.39992725e-07, 9.11570311e-05, 1.15816379e-04, -4.01350495e-05]],
[[-1.85416639e-06, 4.06009696e-05, 3.90917706e-05, 3.92211904e-05],
[-3.72155456e-06, 2.21444097e-05, 3.05974559e-05, 3.17910074e-05],
[1.64244406e-05, 9.33099989e-06, 2.59450976e-05, 7.20713763e-05],
[2.19198952e-05, 2.29714884e-05, 3.55228162e-05, 9.42695439e-05]],
[[-6.29026250e-06, -1.66926104e-06, 2.06531086e-05, 6.30024082e-05],
[-1.71967796e-05, 8.10200354e-06, 1.52458021e-05, 1.94769674e-05],
[-2.22495255e-06, 3.57057325e-06, 2.35516080e-05, 3.85710155e-05],
[-1.44681821e-05, -5.45860797e-06, 3.80976184e-05, 1.24881360e-05]]],
[[[-2.07301156e-05, 3.23990819e-05, 9.57142159e-05, 6.38114024e-05],
[2.92811973e-05, 2.88056901e-05, 4.70659778e-05, 8.20235562e-05],
[7.50632852e-05, 3.26235585e-05, 3.92811088e-05, 8.12137436e-05],
[7.16082561e-05, 2.43401051e-05, 7.43764563e-05, 7.33103146e-05]],
[[1.28299480e-08, 5.67151478e-05, 3.02790507e-05, 3.75851668e-05],
[-5.47604749e-06, 2.78629076e-05, 3.41596648e-05, 3.01239273e-05],
[9.66906328e-06, 7.80152347e-06, 2.20928721e-05, 5.18810534e-05],
[1.64696390e-05, 2.44849598e-06, -5.61052143e-06, 6.28005847e-05]],
[[3.76422464e-06, 3.03913454e-05, 3.42662513e-05, 4.60870862e-05],
[-2.50531945e-06, 9.38416716e-06, 1.46413567e-05, 1.94701388e-05],
[-5.24048728e-06, 3.21705642e-07, 7.17758181e-06, 1.95403688e-05],
[-2.47265560e-06, 4.73080463e-06, 6.29036551e-06,
2.84689950e-05]]]]) * units('s^-1')
def test_vorticity_4d(data_4d):
"""Test vorticity on a 4D (time, pressure, y, x) grid."""
vort = vorticity(data_4d.u, data_4d.v)
assert_array_almost_equal(vort.data, true_vort4d, 12)
def test_absolute_vorticity_4d(data_4d):
"""Test absolute_vorticity on a 4D (time, pressure, y, x) grid."""
vort = absolute_vorticity(data_4d.u, data_4d.v)
f = coriolis_parameter(data_4d.latitude).broadcast_like(vort)
truth = true_vort4d + f.data
assert_array_almost_equal(vort.data, truth, 12)
def test_divergence_4d(data_4d):
"""Test divergence on a 4D (time, pressure, y, x) grid."""
div = divergence(data_4d.u, data_4d.v)
truth = np.array([[[[-5.69109693e-06, -1.97918528e-06, 1.47453542e-05, 2.69697704e-05],
[3.20267932e-05, -6.19720681e-06, 1.25570333e-05, -7.10519011e-06],
[-4.28862128e-05, -1.91207500e-05, 1.39780734e-05, 4.55906339e-05],
[-3.10230392e-05, 1.65756168e-05, -4.24591337e-05, 3.82235500e-05]],
[[1.98223791e-05, 4.33913279e-06, 4.19627202e-05, -5.09003830e-05],
[2.56348274e-05, -1.55289420e-06, 6.96268077e-06, 1.70390048e-05],
[-3.52183670e-06, 3.78206345e-07, -1.34093219e-05, 2.90710519e-06],
[-2.84461615e-05, -6.53843845e-06, -2.54072285e-05, -2.81482021e-05]],
[[-1.21348077e-05, -1.93861224e-05, -1.93459201e-05, -3.87356806e-05],
[4.81616405e-06, -7.66038273e-06, 1.88430179e-07, 4.20022198e-06],
[1.66798208e-05, 5.65659378e-06, 6.33736697e-06, 5.67003948e-06],
[1.30753908e-05, 1.80197572e-05, 1.26966380e-05, 4.18043296e-06]]],
[[[-7.75829235e-07, -5.08426457e-06, 7.57910544e-06, 4.72124287e-05],
[2.57585409e-05, 1.71301607e-06, 5.83802681e-06, -3.33138015e-05],
[-2.91819759e-05, -7.49775551e-06, 2.63853084e-06, 2.33586676e-05],
[-6.76888907e-05, 1.76394873e-05, -5.08169287e-05, 2.85916802e-05]],
[[1.93044895e-05, 1.51461678e-05, -2.09465009e-05, -1.91221470e-05],
[2.02601342e-05, 7.55251174e-07, 4.86519855e-06, -5.99451216e-06],
[6.46768008e-06, -3.39133854e-06, 4.95963402e-06, 3.75958887e-06],
[-1.45155227e-06, 5.60979108e-06, -2.09967347e-05, 2.89704581e-05]],
[[-4.39050924e-06, 2.12833521e-06, -4.50196821e-05, 6.49783523e-06],
[8.22035480e-07, -4.71231966e-06, 2.45757249e-06, 2.41048520e-06],
[7.57532808e-06, -7.32507793e-07, 1.78057678e-05, 1.29309987e-05],
[-2.29661166e-06, 1.96837178e-05, 1.45078799e-06, 1.41496820e-05]]],
[[[3.16631969e-07, -1.24957659e-05, 1.23451304e-05, 9.09226076e-06],
[2.53440942e-05, 3.33772853e-07, 4.20355495e-06, -1.38016966e-05],
[1.66685173e-06, -1.25348400e-05, -7.29217984e-07, 1.40404816e-05],
[-7.16330286e-05, -2.04996415e-05, -6.39953567e-06, -1.13599582e-05]],
[[-6.14675217e-07, 2.05951752e-05, -1.43773812e-05, 5.83203981e-06],
[2.44795938e-05, 4.42280257e-06, -7.63592160e-06, -4.90036880e-06],
[9.02514162e-06, 6.51518845e-08, 5.88086792e-06, -8.59999454e-06],
[-3.99115438e-06, 2.05745950e-07, 1.42084579e-05, 2.83814897e-05]],
[[5.23848091e-06, -1.63679904e-06, -1.97566839e-05, 9.19774945e-06],
[1.32383435e-05, 1.42742942e-06, -8.96735083e-06, -7.41887021e-06],
[3.32715273e-06, 9.54519710e-07, 7.33022680e-06, -9.09165376e-06],
[2.24746232e-06, -1.69640129e-06, 1.80208289e-05,
-5.73083897e-06]]]]) * units('s^-1')
assert_array_almost_equal(div.data, truth, 12)
def test_shearing_deformation_4d(data_4d):
"""Test shearing_deformation on a 4D (time, pressure, y, x) grid."""
shdef = shearing_deformation(data_4d.u, data_4d.v)
truth = np.array([[[[-2.22216294e-05, 5.27319738e-06, 2.91543418e-05, 1.30329364e-05],
[-6.25886934e-05, 1.21925428e-05, 1.98103919e-05, -2.09655345e-05],
[4.46342492e-06, -1.68748849e-05, -1.12290966e-05, 4.23005194e-05],
[6.66667593e-05, -1.03683458e-04, -9.12956532e-05, 7.72037279e-05]],
[[-2.32590651e-05, -8.58252633e-06, 5.74233121e-05, 1.06072378e-06],
[-1.44661146e-06, 1.12270967e-05, 2.17945891e-05, 3.52899261e-05],
[-5.92993188e-06, 1.86784643e-05, -3.53279109e-06, -1.09552232e-05],
[-2.33237922e-05, 1.05752016e-05, 2.39065363e-07, -5.03096678e-05]],
[[-2.49842754e-05, 1.13796431e-05, 4.69266814e-05, 1.53376235e-06],
[-1.48804543e-05, 7.30453364e-06, 2.47197602e-05, 2.67195275e-05],
[-2.16290910e-06, 1.25045903e-05, 9.26533963e-06, 3.17915141e-05],
[1.17935334e-05, 2.77147641e-05, -3.81014510e-07, 1.15523534e-05]]],
[[[6.21038574e-06, 1.20236024e-05, -1.57256625e-05, 3.85950088e-05],
[-3.56990901e-05, 9.80909130e-06, 5.73692616e-06, -2.15769374e-05],
[3.02174095e-06, 1.60641317e-06, 3.91744031e-06, 3.55580983e-05],
[2.10778238e-05, -2.83135572e-05, -4.87985007e-05, 6.74649144e-05]],
[[-2.47860362e-05, -1.78528251e-05, 8.96558477e-06, 9.09500677e-06],
[-1.49626706e-05, 1.04536473e-05, 2.11549168e-05, 7.95984267e-06],
[3.83438985e-06, 1.15792231e-05, 1.65025584e-05, 1.31679266e-05],
[-5.05877901e-06, 6.33465037e-06, 5.39663041e-06, -4.10734948e-05]],
[[-1.88803090e-05, 9.12220863e-06, 2.06531065e-05, 1.26422093e-05],
[-1.71967796e-05, 8.10200354e-06, 1.70443811e-05, -5.70312992e-06],
[1.12643867e-05, 4.46986382e-06, 8.26368981e-06, 4.30674619e-05],
[1.34097890e-05, 8.03073340e-06, -1.31618753e-05, 5.11575682e-05]]],
[[[-7.69041420e-06, 4.07147725e-06, -2.25423126e-05, 1.92965675e-05],
[-1.79314900e-05, 4.97452325e-06, 1.60404993e-05, 1.50265065e-05],
[2.67050064e-06, 2.04831498e-05, 1.90470999e-05, 1.33174097e-05],
[9.10766558e-06, 3.10847747e-05, -1.15056632e-05, 2.60976273e-05]],
[[-9.42970708e-06, -4.53541805e-05, 9.59539764e-06, 3.57865814e-05],
[-1.58178729e-05, -3.16257088e-06, 8.08027693e-06, 3.14524883e-06],
[-3.37063173e-06, 1.63447699e-05, 1.98446489e-05, 7.36623139e-06],
[-1.06650676e-06, 1.90853425e-05, 4.51993196e-05, 8.39356857e-06]],
[[-2.00669443e-05, -2.94113884e-05, -2.60460413e-06, 2.81012941e-05],
[-3.85425208e-06, 2.63949754e-06, 8.34633349e-06, 6.88009010e-06],
[6.45027402e-06, 1.15628217e-05, 1.66201167e-05, 1.68425036e-05],
[1.28152573e-05, -1.11457227e-06, 1.66321845e-05,
4.01597531e-05]]]]) * units('s^-1')
assert_array_almost_equal(shdef.data, truth, 12)
def test_stretching_deformation_4d(data_4d):
"""Test stretching_deformation on a 4D (time, pressure, y, x) grid."""
stdef = stretching_deformation(data_4d.u, data_4d.v)
truth = np.array([[[[3.74747989e-05, 2.58987773e-05, -5.48865461e-06, -2.92358076e-05],
[-7.54193179e-06, 2.70764949e-05, 5.81236371e-06, 2.93160254e-05],
[-5.59259142e-05, 6.05935158e-06, 4.00574629e-05, 5.32345919e-05],
[9.17299272e-05, 2.01727791e-05, 3.57790347e-05, -1.04313800e-04]],
[[5.88340213e-06, 1.64795501e-05, -4.25704731e-05, 1.78952481e-05],
[2.15880187e-05, 7.88964498e-06, -7.42594688e-06, 7.59646711e-06],
[2.77318655e-06, 1.47668340e-05, 2.84076306e-05, -3.38792021e-06],
[-1.13596428e-05, 2.04402443e-05, 5.86763185e-05, 5.00899658e-05]],
[[3.14807173e-05, 3.63698156e-05, 1.75249482e-05, 5.77913742e-06],
[1.87551496e-05, 1.75197146e-05, 6.48345772e-06, 1.18441808e-05],
[1.03847975e-05, 9.70339384e-06, -3.55481427e-06, -1.59129031e-05],
[-4.01111070e-06, 1.03758035e-05, 1.00587992e-06, -3.89854532e-05]]],
[[[4.00968607e-05, 2.17145390e-05, -1.86342944e-06, -3.61966573e-05],
[-3.60325721e-06, 2.94111210e-05, 7.90639233e-06, 5.31067383e-06],
[-3.65111798e-05, 6.84591093e-06, 4.82774623e-05, 6.04093871e-05],
[3.92815820e-05, 4.37497554e-06, 7.04522343e-05, -5.95386733e-05]],
[[-2.72810481e-06, 1.65683717e-06, 2.71654705e-05, 6.95724033e-06],
[1.53140500e-05, 1.60431672e-05, -1.87947106e-06, 3.89766962e-06],
[1.09641265e-05, 1.68426660e-05, 7.65750143e-06, 8.70568056e-06],
[4.84344526e-06, 6.95872586e-06, 5.54428478e-05, 4.02115752e-05]],
[[1.94406533e-05, 8.42336276e-06, 4.31106607e-05, -2.90240924e-05],
[2.17096597e-06, 1.23741775e-05, 6.05472618e-06, 7.35657635e-06],
[7.12568173e-06, 8.26038502e-06, -6.47504105e-06, 1.02331313e-05],
[1.61388203e-05, 1.69793215e-06, 5.94724298e-06, -4.43041213e-05]]],
[[[3.35903423e-05, 3.02204835e-05, -3.39244060e-06, 1.89794480e-06],
[1.06327675e-06, 3.13592514e-05, 1.27468014e-05, -4.80880378e-06],
[-1.04735570e-05, 1.80409922e-05, 3.16451980e-05, 5.00120508e-05],
[-1.93809208e-06, 2.08676689e-05, 4.93564018e-05, 6.23817547e-05]],
[[-1.63522334e-05, -9.08137285e-06, 3.10367228e-05, -8.10694416e-06],
[1.86342126e-05, 1.07178258e-05, 8.10163869e-06, 7.24003618e-06],
[1.12733648e-05, 1.85005839e-05, 6.33051213e-06, 9.83543530e-06],
[4.55210065e-06, 6.95042414e-06, 1.37588137e-05, 3.33275803e-05]],
[[-2.40546855e-06, 1.23021865e-05, 2.47581446e-05, -2.49752420e-05],
[1.00908318e-05, 1.08699686e-05, 7.66950217e-06, -9.21744893e-06],
[8.27324121e-06, 1.08467010e-05, -7.63377597e-07, 4.84732988e-06],
[1.88843132e-05, 1.35915105e-05, -1.16557148e-05,
7.30885665e-06]]]]) * units('s^-1')
assert_array_almost_equal(stdef.data, truth, 10)
def test_total_deformation_4d(data_4d):
"""Test total_deformation on a 4D (time, pressure, y, x) grid."""
totdef = total_deformation(data_4d.u, data_4d.v)
truth = np.array([[[[4.35678937e-05, 2.64301585e-05, 2.96664959e-05, 3.20092155e-05],
[6.30414568e-05, 2.96950278e-05, 2.06454644e-05, 3.60414065e-05],
[5.61037436e-05, 1.79297931e-05, 4.16015979e-05, 6.79945271e-05],
[1.13396809e-04, 1.05627651e-04, 9.80562880e-05, 1.29775901e-04]],
[[2.39916345e-05, 1.85805094e-05, 7.14820393e-05, 1.79266572e-05],
[2.16364331e-05, 1.37220333e-05, 2.30249604e-05, 3.60982714e-05],
[6.54634675e-06, 2.38105946e-05, 2.86264578e-05, 1.14671234e-05],
[2.59430292e-05, 2.30138757e-05, 5.86768055e-05, 7.09934317e-05]],
[[4.01901677e-05, 3.81085262e-05, 5.00922872e-05, 5.97920197e-06],
[2.39412522e-05, 1.89814807e-05, 2.55558559e-05, 2.92270041e-05],
[1.06076480e-05, 1.58278435e-05, 9.92387137e-06, 3.55516645e-05],
[1.24569836e-05, 2.95933345e-05, 1.07562376e-06, 4.06610677e-05]]],
[[[4.05749569e-05, 2.48211245e-05, 1.58356822e-05, 5.29128784e-05],
[3.58804752e-05, 3.10037467e-05, 9.76848819e-06, 2.22208794e-05],
[3.66360092e-05, 7.03186033e-06, 4.84361405e-05, 7.00975920e-05],
[4.45793376e-05, 2.86495712e-05, 8.57018727e-05, 8.99798216e-05]],
[[2.49357203e-05, 1.79295419e-05, 2.86067212e-05, 1.14508664e-05],
[2.14103162e-05, 1.91484192e-05, 2.12382418e-05, 8.86289591e-06],
[1.16152751e-05, 2.04390265e-05, 1.81926293e-05, 1.57855366e-05],
[7.00358530e-06, 9.41018921e-06, 5.57048740e-05, 5.74804554e-05]],
[[2.70999090e-05, 1.24164299e-05, 4.78025091e-05, 3.16579120e-05],
[1.73332721e-05, 1.47906299e-05, 1.80878588e-05, 9.30832458e-06],
[1.33289815e-05, 9.39221184e-06, 1.04983202e-05, 4.42665026e-05],
[2.09829446e-05, 8.20826733e-06, 1.44431528e-05, 6.76753422e-05]]],
[[[3.44594481e-05, 3.04935165e-05, 2.27961512e-05, 1.93896806e-05],
[1.79629867e-05, 3.17513547e-05, 2.04884983e-05, 1.57772143e-05],
[1.08086526e-05, 2.72953627e-05, 3.69352213e-05, 5.17547932e-05],
[9.31159348e-06, 3.74395890e-05, 5.06797266e-05, 6.76207769e-05]],
[[1.88763056e-05, 4.62544379e-05, 3.24861481e-05, 3.66933503e-05],
[2.44425650e-05, 1.11746877e-05, 1.14423522e-05, 7.89371358e-06],
[1.17664741e-05, 2.46864965e-05, 2.08299178e-05, 1.22880899e-05],
[4.67536705e-06, 2.03115410e-05, 4.72470469e-05, 3.43682935e-05]],
[[2.02106045e-05, 3.18806142e-05, 2.48947723e-05, 3.75958169e-05],
[1.08018585e-05, 1.11858466e-05, 1.13350142e-05, 1.15020435e-05],
[1.04905936e-05, 1.58540142e-05, 1.66376388e-05, 1.75261671e-05],
[2.28220968e-05, 1.36371342e-05, 2.03097329e-05,
4.08194213e-05]]]]) * units('s^-1')
assert_array_almost_equal(totdef.data, truth, 12)
def test_frontogenesis_4d(data_4d):
"""Test frontogenesis on a 4D (time, pressure, y, x) grid."""
theta = potential_temperature(data_4d.pressure, data_4d.temperature)
frnt = frontogenesis(theta, data_4d.u, data_4d.v).transpose(
'time1',
'pressure',
'latitude',
'longitude'
)
truth = np.array([[[[4.03241166e-10, -1.66671969e-11, -2.21695296e-10, -3.92052431e-10],
[-5.25244095e-10, -1.12160993e-11, -4.89245961e-11, 3.12029903e-10],
[6.64291257e-10, 3.35294072e-10, 4.40696926e-11, -3.80990599e-10],
[-3.81175558e-10, 2.01421545e-10, 1.69276538e-09, -1.46727967e-09]],
[[-1.14725815e-10, 1.00977715e-10, -1.36697064e-10, 3.42060878e-10],
[-6.07354303e-11, 1.47758507e-11, -4.61570931e-11, 1.07716080e-10],
[4.52780481e-11, 6.99776255e-11, 3.54918971e-10, -7.12011926e-11],
[5.82577150e-10, 1.41596752e-10, 8.60051223e-10, 1.39190722e-09]],
[[-2.76524144e-13, 3.17817935e-10, 5.61139182e-10, 1.41251234e-10],
[9.13538909e-11, 1.07222839e-10, 5.84889489e-11, -7.04354673e-12],
[-7.63814267e-11, -2.61136261e-11, -9.38996785e-12, 2.25155943e-10],
[-8.07125189e-11, -4.09501260e-11, -1.68556325e-10, -1.68395224e-10]]],
[[[1.70912241e-10, -3.59604596e-11, -5.63440060e-11, -3.72141552e-10],
[-1.88604573e-10, -2.84591125e-11, 8.04708643e-12, 3.35465874e-10],
[4.05009495e-10, 4.99273109e-11, 2.70840073e-10, 3.53292290e-10],
[7.61811501e-10, -1.15049239e-10, 1.39114133e-09, -2.13934119e-11]],
[[-7.11061577e-11, -5.56233487e-11, 1.38759396e-10, 2.10158880e-10],
[-9.85704771e-11, 2.43793585e-11, 2.41161028e-11, 8.41366288e-11],
[2.07079174e-11, 9.67316909e-11, 3.79484230e-11, 1.00231778e-10],
[9.09016673e-11, 4.70716770e-12, 7.27411299e-10, -3.68904210e-10]],
[[1.48452574e-11, 1.64659568e-12, 7.71858317e-10, 1.74891129e-10],
[7.51825243e-11, 6.34791773e-11, 6.26549997e-11, -2.97116232e-11],
[-9.19046148e-11, 3.17048878e-11, -6.59923945e-11, 2.25154449e-10],
[-3.68975988e-12, -1.20891474e-10, -3.53749041e-11, 2.42234202e-10]]],
[[[-1.34106978e-10, 1.19278109e-10, -1.70196541e-10, 2.48281391e-11],
[-4.99795205e-11, -2.30130765e-11, 4.96545465e-11, 3.90460132e-11],
[-6.23025651e-12, -2.90005871e-11, 5.57986734e-11, 3.82595360e-10],
[1.33830354e-09, -8.27063507e-11, -2.04614424e-10, 4.66009647e-10]],
[[1.13855928e-10, -3.71418369e-10, 2.37111014e-10, 1.60355663e-10],
[-3.01604394e-10, 3.21033959e-11, 9.52301632e-11, 4.26592524e-11],
[6.25482337e-12, 7.81804086e-11, 2.58199246e-11, 1.74886075e-10],
[4.73684042e-11, 1.42713420e-11, 2.25862198e-10, 3.35966198e-11]],
[[3.63828967e-11, 1.41447035e-10, 9.83470917e-11, 1.37432553e-10],
[-7.52505235e-11, 7.47348135e-12, 8.59892617e-11, 1.09800029e-10],
[-7.58453531e-11, -4.69966422e-12, 1.14342322e-11, 1.81473021e-10],
[-2.97566390e-11, 9.55288188e-11, 1.90872070e-12,
5.32192321e-10]]]]) * units('K/m/s')
assert_array_almost_equal(frnt.data, truth, 13)
def test_geostrophic_wind_4d(data_4d):
"""Test geostrophic_wind on a 4D (time, pressure, y, x) grid."""
u_g, v_g = geostrophic_wind(data_4d.height)
u_g_truth = np.array([[[[4.40486857, 12.51692362, 20.63729052, 3.17769103],
[14.10194385, 17.12263527, 22.04954906, 28.25627455],
[24.44520744, 22.83658981, 31.70185785, 41.43475568],
[35.55079058, 29.81196157, 50.61168553, 41.3453152]],
[[7.35973026, 11.15080483, 15.35393153, 8.90224492],
[8.36112125, 12.51333666, 13.38382965, 14.31962023],
[10.36996866, 13.03590323, 16.55132073, 20.5818555],
[13.5135907, 12.61987724, 25.47981975, 27.81300618]],
[[5.7532349, 8.87025457, 12.11513303, 6.95699048],
[5.63036393, 9.22723096, 9.46050119, 9.63463697],
[5.15111753, 8.92136337, 10.13229436, 10.02026917],
[4.27093407, 7.87208545, 14.52880097, 7.84194092]]],
[[[2.5637431, 12.12175173, 18.88903199, 9.31429705],
[11.13363928, 16.0692665, 22.88529458, 23.2247996],
[21.17380737, 18.19154369, 27.45449837, 37.89231093],
[32.89749798, 18.27860794, 32.68137607, 53.46238172]],
[[5.88868723, 10.23886179, 13.99207128, 7.62863391],
[7.72562524, 12.48283965, 13.87130359, 12.97472345],
[9.38948632, 12.47561185, 15.29521563, 18.71570682],
[10.86569541, 9.9484405, 18.45258492, 24.92010765]],
[[5.37666204, 9.31750379, 9.01145336, 3.68871571],
[5.42142755, 8.93123996, 9.3456061, 9.00788096],
[4.94868897, 8.34298027, 9.29367749, 11.09021722],
[3.89473037, 7.52596886, 8.80903478, 9.55782485]]],
[[[4.07701238, 9.91100559, 14.63521328, 11.44931302],
[9.21849096, 15.3989699, 20.84826449, 20.35213024],
[17.27879494, 16.28474382, 23.2252306, 32.43391015],
[28.63615274, 12.02290076, 21.31740598, 48.11881923]],
[[4.67797945, 7.67496476, 7.67070623, 7.43540912],
[6.36765831, 10.59388475, 12.09551703, 11.52096191],
[7.77187799, 11.17427747, 14.91109777, 16.17178096],
[8.86174464, 9.13936139, 15.93606235, 21.47254981]],
[[4.06859791, 6.49637561, 4.98326026, 5.11096512],
[4.19923606, 6.75503407, 8.50298015, 8.50994027],
[3.85339598, 6.92959314, 9.8142002, 10.51547453],
[2.97279588, 7.0103826, 8.65854182, 10.96893324]]]]) * units('m/s')
v_g_truth = np.array([[[[-2.34958057e+01, -1.94104519e+01, -7.44959497e+00,
1.23868322e+01],
[-2.05867367e+01, -1.59688225e+01, -7.24619436e+00,
5.58114910e+00],
[-2.13003979e+01, -1.50644426e+01, -1.26465809e+00,
2.00990219e+01],
[-2.83335381e+01, -1.22608318e+01, 2.75571752e+00,
1.67161713e+01]],
[[-2.12135105e+01, -1.57486000e+01, -7.18331385e+00,
4.48243952e+00],
[-1.85706921e+01, -1.38995152e+01, -7.25590754e+00,
1.36025941e+00],
[-1.48431730e+01, -1.30190716e+01, -6.20916080e+00,
5.58656025e+00],
[-1.64091930e+01, -1.07454290e+01, -3.26166773e+00,
6.04215336e+00]],
[[-1.84210243e+01, -1.51837034e+01, -8.32569885e+00,
2.15305471e+00],
[-1.60743446e+01, -1.37354202e+01, -8.54446602e+00,
-5.01543939e-01],
[-1.26119165e+01, -1.31178055e+01, -8.13879681e+00,
2.32514095e+00],
[-1.08224831e+01, -1.12312374e+01, -8.07368088e+00,
-1.34987926e+00]]],
[[[-2.47784901e+01, -2.06641865e+01, -7.55605650e+00,
1.45456514e+01],
[-2.05139866e+01, -1.66804104e+01, -6.96553278e+00,
8.63076687e+00],
[-2.04345818e+01, -1.41986904e+01, -3.59461641e+00,
1.13773890e+01],
[-3.07159233e+01, -1.35134182e+01, 3.63993049e+00,
2.07441883e+01]],
[[-2.20703144e+01, -1.61019173e+01, -6.81787109e+00,
5.78179121e+00],
[-1.89408665e+01, -1.40810776e+01, -7.12525749e+00,
1.92659533e+00],
[-1.49793730e+01, -1.27466383e+01, -6.57639217e+00,
3.53139591e+00],
[-1.57215986e+01, -1.10794334e+01, -3.83887053e+00,
6.00018406e+00]],
[[-1.89922485e+01, -1.49378052e+01, -8.35085773e+00,
7.68607914e-01],
[-1.58400993e+01, -1.38690310e+01, -9.15049839e+00,
-1.68443954e+00],
[-1.34329786e+01, -1.28181543e+01, -8.34892273e+00,
-2.52818279e-02],
[-1.10563183e+01, -1.17126417e+01, -7.79271078e+00,
7.03427792e-01]]],
[[[-2.87962914e+01, -2.08093758e+01, -7.41080666e+00,
1.13992844e+01],
[-2.51369133e+01, -1.76726551e+01, -6.50083351e+00,
8.37874126e+00],
[-2.16215363e+01, -1.44126577e+01, -4.67937374e+00,
7.57850361e+00],
[-3.09025593e+01, -1.47021618e+01, 2.18099499e+00,
1.97469769e+01]],
[[-2.14603043e+01, -1.55501490e+01, -7.21480083e+00,
3.54577303e+00],
[-1.86117344e+01, -1.43230457e+01, -7.12040138e+00,
2.99635530e+00],
[-1.53198442e+01, -1.24255934e+01, -6.73208141e+00,
1.76072347e+00],
[-1.53703299e+01, -1.06545277e+01, -4.50935888e+00,
3.06527138e+00]],
[[-1.62525253e+01, -1.41536722e+01, -9.22987461e+00,
-1.48113370e+00],
[-1.41632900e+01, -1.34236937e+01, -9.18536949e+00,
-1.44826770e+00],
[-1.30243769e+01, -1.18180895e+01, -8.29443932e+00,
-2.45343924e+00],
[-1.09246559e+01, -1.03824110e+01, -7.37222433e+00,
-1.89407897e+00]]]]) * units('m/s')
assert_array_almost_equal(u_g.data, u_g_truth, 4)
assert_array_almost_equal(v_g.data, v_g_truth, 4)
def test_inertial_advective_wind_4d(data_4d):
"""Test inertial_advective_wind on a 4D (time, pressure, y, x) grid."""
u_g, v_g = geostrophic_wind(data_4d.height)
u_i, v_i = inertial_advective_wind(u_g, v_g, u_g, v_g)
u_i_truth = np.array([[[[-4.76966186, -6.39706038, -7.24003746, -11.13794333],
[-1.89586566, -4.35883424, -6.85805714, -9.4212875],
[2.31372726, -6.96059926, -14.11458588, -20.68380008],
[-0.92883306, -13.81354883, -17.96354053, -23.79779997]],
[[-2.62082254, -3.50555985, -3.63842481, -4.20932821],
[-3.38566969, -2.58907076, -2.67708014, -3.36021786],
[-0.56713627, -2.3416945, -4.39000187, -6.69093397],
[1.70678751, -3.60860629, -5.96627063, -7.52914321]],
[[-1.61492016, -2.31780373, -2.40235407, -2.60787441],
[-2.19903344, -1.48707548, -1.58037953, -2.25343451],
[-1.11096954, -1.25163409, -2.02857574, -3.32734329],
[-0.26020197, -1.62905796, -1.75707467, -1.2223621]]],
[[[-6.72701434, -6.76960203, -7.94802076, -12.50171137],
[-2.22284799, -5.07983672, -7.76025363, -11.23189296],
[2.67509705, -4.83471753, -9.58547825, -12.94725576],
[8.58545145, -7.72587914, -12.41979585, -10.25605548]],
[[-3.19317899, -3.55857747, -3.56352137, -4.31615186],
[-3.70727146, -2.8684896, -2.7782166, -3.33031965],
[-1.17242459, -2.18140469, -3.58528354, -5.27404394],
[1.42344232, -2.45475499, -4.65221513, -6.1169067]],
[[-3.23907889, -1.91350728, -1.17379843, -1.09402307],
[-2.0340837, -1.38963467, -1.40556307, -1.93552382],
[-1.31936373, -1.1627646, -1.73546489, -2.82082041],
[-0.96507328, -0.94398947, -1.53168307, -2.57261637]]],
[[[-5.13667819, -5.35808776, -5.96105057, -8.09779516],
[-5.27868329, -6.04992134, -7.09615152, -9.11538451],
[0.32367483, -4.40754181, -7.26937211, -8.89052436],
[11.86601164, -3.52532263, -8.2149503, -3.91397366]],
[[-2.95853902, -1.94361543, -1.79128105, -2.22848035],
[-2.98114417, -2.49536376, -2.66131831, -3.4095258],
[-1.43210061, -2.24010995, -3.02803196, -3.96476269],
[0.38124008, -2.11580893, -3.41706461, -4.07935491]],
[[-1.85523484, -0.74020207, -0.62945585, -1.19060464],
[-0.90996905, -1.11068858, -1.44720476, -1.96113271],
[-0.97632032, -1.23447402, -1.48613628, -1.80024482],
[-1.30046767, -0.98449831, -1.25199805,
-1.96583328]]]]) * units('m/s')
v_i_truth = np.array([[[[1.03212922e+01, 5.87785876e+00, -3.24290351e+00, -1.88453875e+01],
[9.87498125e+00, 5.33624247e+00, 4.80855268e+00, 3.62780511e-02],
[6.37519841e+00, 6.45883096e+00, 8.14332496e+00, 4.38659798e+00],
[-1.31389541e+00, 1.00955857e+01, 4.19848197e+00,
-1.97713955e+01]],
[[1.10365470e+00, 2.30316727e+00, -1.82344497e+00, -3.54754121e+00],
[2.43595083e+00, 1.35702893e+00, 4.91118248e-01, -1.03105842e-02],
[2.33831643e+00, 1.03116363e+00, 3.27903073e+00, 4.52178657e-01],
[2.90828402e-01, 1.43477414e+00, 6.69517000e+00, -4.27716340e+00]],
[[4.77177073e-01, 1.14435024e+00, -1.82680726e+00, -1.95986760e+00],
[5.18719070e-01, 4.51688547e-01, -3.28412094e-01, 6.84697225e-02],
[2.50141134e-01, 1.41518671e-01, 1.08838497e+00, -9.61933095e-02],
[-3.39178295e-01, 2.45727962e-01, 2.41825249e+00,
-2.84771923e+00]]],
[[[9.01360331e+00, 6.74640647e+00, 5.47040255e-01, -1.25154925e+01],
[9.56977790e+00, 4.57707018e+00, 3.34473925e+00, -7.13502610e+00],
[5.46464641e+00, 2.13949666e+00, 7.51823914e+00, 2.43533142e+00],
[-5.48839487e+00, -6.52611598e-01, 1.34292069e+01,
1.61544754e+01]],
[[2.49507477e+00, 3.34927241e+00, -7.11661027e-01, -3.42627695e+00],
[2.69966530e+00, 1.64559616e+00, 2.90248174e-01, -1.12696139e+00],
[1.83330337e+00, 1.69378198e-01, 1.87762364e+00, 7.55276554e-01],
[-4.89132896e-01, -1.06737759e+00, 4.20052028e+00,
1.54873202e+00]],
[[1.05176368e+00, 2.35279690e-01, -4.37230320e-01, -9.41455734e-01],
[5.26256702e-01, 1.32552797e-01, 6.61475967e-02, 1.17988702e-01],
[9.40681182e-02, 3.45287932e-02, 2.13397644e-01, 6.10768896e-01],
[-2.44304796e-01, -6.00961285e-02, -3.78761065e-02,
2.27978276e-01]]],
[[[5.18728227e+00, 8.23825046e+00, 2.86046723e+00, -5.59088886e+00],
[8.85355614e+00, 4.70956220e+00, 2.51349179e+00, -5.64414232e+00],
[7.54622775e+00, 7.98092891e-02, 4.70152506e+00, 3.47162602e+00],
[-1.92789744e+00, -5.92225638e+00, 1.00594741e+01,
2.62864566e+01]],
[[2.20468164e+00, 3.00812312e+00, 1.59439971e+00, -6.42312367e-01],
[2.15609133e+00, 1.86103734e+00, 1.28243894e+00, -1.03944156e+00],
[1.50383253e+00, 5.72866867e-01, 1.51969207e+00, -3.94601885e-01],
[2.57841077e-02, -8.98532915e-01, 2.48926548e+00, 1.81145651e+00]],
[[6.98587642e-01, 2.55740716e-01, 1.74401316e+00, 3.79592864e-01],
[2.39095399e-01, 4.87795233e-01, 1.16885491e+00, -7.66586054e-03],
[-6.48645993e-02, 5.81727905e-01, 4.66123480e-01, 3.71778788e-02],
[-2.11967488e-01, 5.16025460e-01, -4.15578572e-01,
6.96366806e-01]]]]) * units('m/s')
assert_array_almost_equal(u_i.data, u_i_truth, 4)
assert_array_almost_equal(v_i.data, v_i_truth, 4)
def test_q_vector_4d(data_4d):
"""Test q_vector on a 4D (time, pressure, y, x) grid."""
u_g, v_g = geostrophic_wind(data_4d.height)
q1, q2 = q_vector(u_g, v_g, data_4d.temperature, data_4d.pressure)
q1_truth = np.array([[[[-1.11399407e-12, 2.50794237e-13, 3.16093168e-12, 2.32093331e-12],
[5.65649869e-13, 1.45620779e-13, 1.71343752e-12, 4.35800011e-12],
[-4.96503931e-13, 2.62116549e-12, 3.96540726e-12, 4.08996874e-12],
[-1.31411324e-12, 8.91776830e-12, 9.28518964e-12, -2.68490726e-12]],
[[7.62925715e-13, 1.16785318e-12, -2.01309755e-13, 1.26529742e-12],
[3.47435346e-13, 6.73455725e-13, 6.90294419e-13, 8.12267467e-13],
[3.45704077e-14, 3.82817753e-13, 1.54656386e-12, 3.07185369e-12],
[6.11434780e-13, 6.23632300e-13, 1.40617773e-12, 5.34947219e-12]],
[[3.06414278e-13, 6.53804262e-13, 1.75404505e-12, 5.51976164e-13],
[3.28229719e-13, 2.75782033e-13, 4.00407507e-13, 7.84750100e-13],
[1.32588098e-13, 2.51525423e-13, 5.49106514e-13, 1.78892467e-12],
[7.88840796e-14, 1.60673966e-13, 1.19208617e-12, 2.05418653e-12]]],
[[[-3.34132897e-13, 1.53374763e-12, 4.49316053e-12, 1.64643286e-12],
[1.07061926e-13, 1.48351071e-13, 2.40731954e-12, 5.09815184e-12],
[-6.72234608e-13, 1.09871184e-12, 1.78399997e-12, 2.83734147e-12],
[-2.04431842e-12, 6.47809851e-12, 4.82039700e-12, -2.09744034e-12]],
[[5.49758402e-13, 1.84138510e-13, 6.32622851e-13, 2.64607266e-12],
[6.72993111e-13, 6.48589900e-13, 8.38201872e-13, 1.02446030e-12],
[6.71507328e-14, 4.68905430e-13, 1.21606351e-12, 2.11104302e-12],
[3.15734101e-13, 1.95983121e-13, 1.20260143e-12, 4.19732652e-12]],
[[6.33871103e-13, 5.90709910e-13, 1.21586844e-12, 1.06166654e-12],
[3.43322382e-13, 2.76046202e-13, 4.90662239e-13, 5.62988991e-13],
[2.09877678e-13, 2.37809232e-13, 3.65590502e-13, 8.07598362e-13],
[2.35522003e-14, 2.21315437e-13, 5.14061506e-13, 1.17222164e-12]]],
[[[8.22178446e-13, 2.09477989e-12, 5.54298564e-12, 2.21511518e-12],
[1.54450727e-12, 5.45033765e-13, 2.37288896e-12, 3.30727156e-12],
[2.94161134e-13, 3.03848451e-13, 1.47235183e-13, 2.95945450e-12],
[-6.20823394e-12, 9.77981323e-13, -2.06881609e-12,
-3.58251099e-12]],
[[7.61196804e-13, -6.76343613e-13, 9.48323229e-13, 6.33365711e-13],
[1.14599786e-12, 6.99199729e-13, 7.41681860e-13, 9.28590425e-13],
[1.03166054e-13, 6.13187200e-13, 8.39627802e-13, 1.48207669e-12],
[3.22870872e-13, -2.18606955e-13, 9.98812765e-13, 1.91778451e-12]],
[[5.55381844e-13, 2.79538040e-13, 3.30236669e-13, -4.91571259e-14],
[2.50227841e-13, 2.70855069e-13, 4.03362348e-13, 5.22065702e-13],
[3.37119836e-13, 3.17667714e-13, 2.25387106e-13, 6.46265259e-13],
[2.05548507e-13, 3.55426850e-13, -1.74728156e-14,
5.04028133e-13]]]]) * units('m^2 kg^-1 s^-1')
q2_truth = np.array([[[[3.34318820e-12, -1.32561232e-13, 1.01510711e-12, 6.03331800e-12],
[2.51737448e-13, -1.71044158e-13, -8.25290924e-13, 1.68843717e-13],
[-3.50533924e-12, -1.68864979e-12, 7.74026063e-13, 1.53811977e-12],
[-1.75456351e-12, -3.86555813e-12, -1.89153040e-12,
-5.15241976e-12]],
[[-2.09823428e-13, -6.26774796e-13, -1.40145242e-13, 1.09119884e-12],
[-2.58383579e-13, -2.67580088e-13, -6.44081099e-14, 5.90687237e-13],
[-2.73791441e-13, -2.28454021e-13, -4.76780567e-13,
-8.48612071e-13],
[1.21028867e-12, -5.10570435e-13, 6.32933788e-14, 2.44873356e-12]],
[[-6.72615385e-14, -3.57492837e-13, -4.18039453e-14, 3.81431652e-13],
[-3.56221252e-13, -1.23227388e-13, -3.21550682e-14,
-4.69364017e-14],
[-2.82392676e-13, -1.20969658e-13, 1.13815813e-13, -6.93334063e-14],
[5.19714375e-14, -4.61213207e-13, 5.33263529e-13, 1.28188808e-12]]],
[[[1.72090175e-12, -1.35788214e-12, 1.48184196e-13, 3.22826220e-12],
[-2.13531998e-13, -1.17527920e-13, -6.94495723e-13, 1.76851853e-12],
[-2.67906083e-12, -3.78055500e-13, -9.90177606e-13, 2.87945080e-12],
[1.48308333e-12, 2.15094364e-13, -4.84492586e-12, 2.77186124e-12]],
[[-3.09679995e-13, -2.52107306e-13, 4.57596343e-14, 2.03410764e-12],
[-3.95694028e-13, -3.00120864e-13, 1.05194228e-14, 1.06118650e-12],
[-2.46886775e-13, -2.43771482e-13, -3.81434550e-13,
-1.70381858e-13],
[8.12779739e-13, -1.38604573e-13, -8.06018823e-13,
-7.80874251e-13]],
[[-2.19867650e-13, -1.53226258e-13, 4.07751663e-13, 1.52430853e-12],
[-2.56504387e-13, -1.21082762e-13, 6.29695620e-15, 3.49769107e-13],
[-2.44113387e-13, -1.21986494e-13, -9.12718030e-14,
-1.60274595e-13],
[-2.47713487e-13, -1.77307889e-13, -1.13295694e-13,
-6.07631206e-13]]],
[[[-6.49198232e-13, -1.97145105e-12, -5.54605298e-13, 1.94723486e-12],
[-2.00875018e-12, -3.72744112e-13, -4.59665809e-13, 1.12359119e-13],
[-3.83508199e-12, 1.18439125e-13, -4.24891463e-13, -5.88482481e-13],
[-1.84276119e-12, 1.55112390e-12, -7.38034738e-13, 1.03676199e-13]],
[[-4.58813210e-13, -1.88617051e-13, 2.58369931e-13, 8.15071067e-13],
[-6.09657914e-13, -3.51811097e-13, 2.39365587e-13, 5.80541301e-13],
[-1.69115858e-13, -3.49864908e-13, -2.26620147e-13, 7.79560474e-13],
[2.23317058e-13, 1.20352864e-13, -1.01565643e-12, -2.16675768e-13]],
[[-1.68140233e-13, -5.07963999e-14, 2.77741196e-13, 8.37842279e-13],
[-1.39578146e-13, -1.36744814e-13, 3.12352497e-14, 4.55339789e-13],
[-1.06614836e-13, -2.19878930e-13, -8.37992151e-14, 1.87868902e-13],
[-2.27057581e-13, -2.74474045e-13, -1.10759455e-13,
-3.90242255e-13]]]]) * units('m^2 kg^-1 s^-1')
assert_array_almost_equal(q1.data, q1_truth, 15)
assert_array_almost_equal(q2.data, q2_truth, 15)
@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'),
indirect=True)
def test_geospatial_laplacian_geographic(geog_data):
"""Test geospatial_laplacian across projections."""
crs, lons, lats, _, arr, mx, my, dx, dy = geog_data
laplac = geospatial_laplacian(arr, longitude=lons, latitude=lats, crs=crs)
# Calculate the true fields using known map-correct approach
u = mx * first_derivative(arr, delta=dx, axis=1)
v = my * first_derivative(arr, delta=dy, axis=0)
truth = (mx * first_derivative(u, delta=dx, axis=1)
+ my * first_derivative(v, delta=dy, axis=0)
- (u * mx / my) * first_derivative(my, delta=dx, axis=1)
- (v * my / mx) * first_derivative(mx, delta=dy, axis=0))
assert_array_almost_equal(laplac, truth)