Search code examples
pythonmetpy

Calculating vorticity for multiple vertical levels in MetPy


I'm trying to calculate vorticity in MetPy for multiple (consecutive) vertical levels. When I try to calculate it for a single level, everything works fine.

Here's the code; I've used the example for cross sections from https://unidata.github.io/MetPy/latest/examples/cross_section.html#sphx-glr-examples-cross-section-py.

import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

import metpy.calc as mpcalc
from metpy.cbook import get_test_data
from metpy.interpolate import cross_section
from metpy.units import units

data = xr.open_dataset(get_test_data('narr_example.nc', False))
data = data.metpy.parse_cf().squeeze()

data_crs = data['Temperature'].metpy.cartopy_crs
lat = data['lat']
lon = data['lon']
f = mpcalc.coriolis_parameter(lat)
dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat, initstring=data_crs.proj4_init)

Then the calculation of vorticity is performed.

vort = mpcalc.vorticity(data['u_wind'], data['v_wind'], dx, dy)

The traceback:

Traceback (most recent call last):
  File "E:\Временные файлы\cross_section (1).py", line 63, in <module>
    vort = mpcalc.vorticity(data['u_wind'], data['v_wind'], dx, dy)
  File "C:\ProgramData\Miniconda3\lib\site-packages\metpy\xarray.py", line 436, in wrapper
    return func(*args, **kwargs)
  File "C:\ProgramData\Miniconda3\lib\site-packages\metpy\calc\kinematics.py", line 60, in wrapper
    ret = func(*args, **kwargs)
  File "C:\ProgramData\Miniconda3\lib\site-packages\metpy\calc\kinematics.py", line 121, in vorticity
    dudy = first_derivative(u, delta=dy, axis=-2)
  File "C:\ProgramData\Miniconda3\lib\site-packages\metpy\calc\tools.py", line 920, in wrapper
    return preprocess_xarray(func)(f, **kwargs)
  File "C:\ProgramData\Miniconda3\lib\site-packages\metpy\xarray.py", line 436, in wrapper
    return func(*args, **kwargs)
  File "C:\ProgramData\Miniconda3\lib\site-packages\metpy\calc\tools.py", line 1014, in first_derivative
    combined_delta = delta[tuple(delta_slice0)] + delta[tuple(delta_slice1)]
  File "C:\ProgramData\Miniconda3\lib\site-packages\pint\quantity.py", line 1400, in __getitem__
    value = self._magnitude[key]
IndexError: too many indices for array

I'm absolutely stuck. Searching "metpy multiple levels calculations" (no actual quotes) gives no relevant results. The doc says:

metpy.calc.vorticity(u, v, dx, dy)[source]
Calculate the vertical vorticity of the horizontal wind.

Parameters: 
u ((M, N) ndarray) – x component of the wind
v ((M, N) ndarray) – y component of the wind
dx (float or ndarray) – The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of u along the applicable axis.
dy (float or ndarray) – The grid spacing(s) in the y-direction. If an array, there should be one item less than the size of u along the applicable axis.
dim_order (str or None, optional) – The ordering of dimensions in passed in arrays. Can be one of None, 'xy', or 'yx'. 'xy' indicates that the dimension corresponding to x is the leading dimension, followed by y. 'yx' indicates that x is the last dimension, preceded by y. None indicates that the default ordering should be assumed, which is ‘yx’. Can only be passed as a keyword argument, i.e. func(…, dim_order=’xy’).
Returns:    
(M, N) ndarray – vertical vorticity

I conclude that the input can have more than 2 dimensions, but 3-dimensional input (as it is in my case) gives errors. What can be done to fix them?

I'm absolutely new to Python, so I could've made a stupid mistake.


Solution

  • Unfortunately, the error message that comes up isn't that helpful in this case if you don't know what to look for!

    The problem with the vorticity function call in your example is that the dimensionality of your input variables do not match. data['u_wind'] and data['v_wind'] are 3D arrays with shape (29, 118, 292), but dx and dy, since they were computed from lat_lon_grid_deltas, are 2D arrays with shapes (118, 291) and (117, 292) respectively. And so, we need to obtain arrays that broadcast appropriately...there are many different ways that you could do this, but here are two options that I would recommend:

    Option 1: Manual Broadcasting

    Since the "extra" dimension that dx and dy are missing is the first dimension (in the vertical), we can just make dx and dy into properly aligned 3D arrays by inserting a size-one leading dimension:

    dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat, initstring=data_crs.proj4_init)
    dx = dx[None, :]
    dy = dy[None, :]
    
    vort = mpcalc.vorticity(data['u_wind'], data['v_wind'], dx, dy)
    

    Option 2: Use the grid_deltas_from_dataarray() helper function

    MetPy also has a helper function to make pulling the grid deltas from an xarray DataArray easy. It also ensures that the broadcasting occurs properly, so you don't have to do it yourself. Using it in your example, it would be:

    dx, dy = mpcalc.grid_deltas_from_dataarray(data['u_wind'])
    
    vort = mpcalc.vorticity(data['u_wind'], data['v_wind'], dx, dy)