Inconsistent Return Types

[sgdecker]

Is this a bug?

import datetime as dt
import xarray as xr
from xarray.backends import NetCDF4DataStore
from metpy.calc import geospatial_laplacian, geospatial_gradient, gradient
from metpy.units import units
from siphon.catalog import TDSCatalog


def get_nam212(init_time, valid_time):
    """Obtain NAM data on the 212 grid."""
    ymd = init_time.strftime('%Y%m%d')
    hr = init_time.strftime('%H')
    filename = f'{ymd}_{hr}00.grib2'
    ds_name = 'NAM_CONUS_40km_conduit_' + filename
    cat_name = ('https://thredds.ucar.edu/thredds/catalog/grib/NCEP/NAM/'
                'CONUS_40km/conduit/' + ds_name + '/catalog.xml')

    cat = TDSCatalog(cat_name)
    ds = cat.datasets[ds_name]
    ncss = ds.subset()
    query = ncss.query()
    query.time(valid_time).variables('all')
    nc = ncss.get_data(query)
    data = xr.open_dataset(NetCDF4DataStore(nc)).metpy.parse_cf()
    data = data.metpy.assign_latitude_longitude()
    return data


init_time = dt.datetime(2024, 3, 5, 12)
plot_time = dt.datetime(2024, 3, 5, 18)

nam = get_nam212(init_time, plot_time)

p = 300*units('hPa')
T = (nam['Temperature_isobaric'].metpy.sel(vertical=p).squeeze()
     .metpy.quantify())
u = (nam['u-component_of_wind_isobaric'].metpy.sel(vertical=p).squeeze()
     .metpy.quantify())

lap = geospatial_laplacian(T)
dudx, dudy = geospatial_gradient(u)
a, b = gradient(u)

print(type(T), type(lap))
print(type(u), type(dudx), type(dudy))
print(type(a), type(b))

I get

<class 'xarray.core.dataarray.DataArray'> <class 'xarray.core.dataarray.DataArray'>
<class 'xarray.core.dataarray.DataArray'> <class 'pint.Quantity'> <class 'pint.Quantity'>
<class 'xarray.core.dataarray.DataArray'> <class 'xarray.core.dataarray.DataArray'>

Shouldn't geospatial_gradient return DataArrays just like geospatial_laplacian and gradient do?

[dopplershift]

That's...tricky. This is controlled by the wrap_like argument to our preprocess_and_wrap() decorator, and would need to somehow change depending on what the value of return_only passed to geospatial_gradient().

I think this hasn't been noticed because it's not really used much outside, and internal usage of this function is in contexts where we aren't using DataArrays.

Is this impacting you in ways beyond the work you're doing exploring #3422?

[sgdecker]

This is related to #3422, yes. I tried to compute 2*J(u,v) manually (part of a larger calculation) with:

dudx, dudy = mpcalc.geospatial_gradient(u)
dvdx, dvdy = mpcalc.geospatial_gradient(v)
term2 = 2 * (dudx*dvdy - dudy*dvdx)

but when I went to get a quick look at the result for comparison purposes:

term2.plot()

I got an error since term2 is not a DataArray and thus has no plot method.

Luckily, the other terms are all DataArrays, so when I add everything together, I get something I can plot for my actual use case.

[dopplershift]

Yeah makes sense. In an ideal world, geospatial_gradient would return DataArrays, but doing so would require a lot of effort that IMO isn't worth it at the moment, unless someone (@jthielen ?) has thoughts on an easy way to handle the fact that geospatial_gradient can return anything from a single array to a tuple of 4 arrays.

[jthielen]

If I recall correctly, there are a couple issues:

• using preprocess_and_wrap with "multiple" (i.e., tuple of) return values is supposed to be supported with a tuple wrap_like argument (but that doesn't look to be tested right now in tests/test_xarray.py when I looked quick)
• having an undetermined count of return values isn't supported right now

However, I think, for the case of identically laid out (i.e., same array shapes and coordinates) output that match a single input, which should be the case here, it should be fairly easy to add such functionality, by changing

https://github.com/Unidata/MetPy/blob/main/src/metpy/xarray.py#L1340-L1343

to something like

if isinstance(match, tuple):
    return tuple(wrapping(*args) for args in zip(result, match))
elif isinstance(result, tuple):
    return tuple(wrapping(single_result, match) for single_result in result)
else:
    return wrapping(result, match)

Of course, tests should be added, and the wrap_like argument may need to be added to/altered for geospatial_gradient and vector_derivative.