Skip to content

Commit

Permalink
Fix up lat/lon handling to allow proper xarray broadcasting
Browse files Browse the repository at this point in the history 
This fixes auto-broadcasting of latitude for e.g. absolute_vorticity.
  • Loading branch information
@dopplershift
dopplershift committed Nov 30, 2022
1 parent f905901 commit 24f5655
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 41 deletions.
6 changes: 2 additions & 4 deletions src/metpy/calc/tools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1123,8 +1123,6 @@ def wrapper(*args, **kwargs):
proj = grid_prototype.metpy.pyproj_proj
latitude, longitude = grid_prototype.metpy.coordinates('latitude',
'longitude')
latitude = latitude.metpy.unit_array
longitude = longitude.metpy.unit_array
calculate_scales = True
except AttributeError:
# Fall back to basic cartesian calculation if we don't have a CRS or we are
Expand All @@ -1147,8 +1145,8 @@ def wrapper(*args, **kwargs):

# Do we have everything we need to sensibly calculate the scale arrays?
if calculate_scales:
scale_lat = latitude.squeeze().m_as('degrees')
scale_lon = longitude.squeeze().m_as('degrees')
scale_lat = latitude.metpy.unit_array.squeeze().m_as('degrees')
scale_lon = longitude.metpy.unit_array.squeeze().m_as('degrees')
if scale_lat.ndim == 1 and scale_lon.ndim == 1:
scale_lon, scale_lat = np.meshgrid(scale_lon, scale_lat)
elif scale_lat.ndim != 2 or scale_lon.ndim != 2:
Expand Down
84 changes: 47 additions & 37 deletions tests/calc/test_kinematics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1276,46 +1276,56 @@ def data_4d():
})


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)
truth = 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')
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)


Expand Down

0 comments on commit 24f5655

Please sign in to comment.