Allow `DataArray`s in `interpolate_to_isosurface`

[sgdecker]

Discussed in https://github.com/Unidata/MetPy/discussions/2237

^{Originally posted by **sgdecker** December 3, 2021} I don't have time at the moment to simplify my code to something reproducible, but given a variable `ds` that looks like this: ``` Dimensions: (pres: 37, x: 185, y: 129) Coordinates: * pres (pres) int64 100 125 150 175 200 225 ... 875 900 925 950 975 1000 time datetime64[ns] 2019-11-30T12:00:00 * x (x) float32 -4.226e+06 -4.185e+06 ... 3.21e+06 3.251e+06 * y (y) float32 2.035e+06 2.076e+06 2.117e+06 ... 7.196e+06 7.237e+06 metpy_crs object Projection: lambert_conformal_conic latitude (y, x) float64 12.27 12.37 12.47 12.56 ... 57.64 57.54 57.44 longitude (y, x) float64 -133.4 -133.1 -132.7 ... -50.48 -49.95 -49.42 Data variables: z (pres, y, x) float32 ----> 1 theta_dyn_trop = interpolate_to_isosurface(ds['pv'], ds['theta'], 2*pvu) ~/local/miniconda3/envs/met433a/lib/python3.9/site-packages/metpy/pandas.py in wrapper(*args, **kwargs) 19 kwargs = {name: (v.values if isinstance(v, pd.Series) else v) 20 for name, v in kwargs.items()} ---> 21 return func(*args, **kwargs) 22 return wrapper ~/local/miniconda3/envs/met433a/lib/python3.9/site-packages/metpy/interpolate/grid.py in interpolate_to_isosurface(level_var, interp_var, level, bottom_up_search) 349 350 # Linear interpolation of variable to interpolated surface value --> 351 interp_level = (((level - level_var[above]) / (level_var[below] - level_var[above])) 352 * (interp_var[below] - interp_var[above])) + interp_var[above] 353 ~/local/miniconda3/envs/met433a/lib/python3.9/site-packages/xarray/core/dataarray.py in __getitem__(self, key) 748 else: 749 # xarray-style array indexing --> 750 return self.isel(indexers=self._item_key_to_dict(key)) 751 752 def __setitem__(self, key: Any, value: Any) -> None: ~/local/miniconda3/envs/met433a/lib/python3.9/site-packages/xarray/core/dataarray.py in isel(self, indexers, drop, missing_dims, **indexers_kwargs) 1181 1182 if any(is_fancy_indexer(idx) for idx in indexers.values()): -> 1183 ds = self._to_temp_dataset()._isel_fancy( 1184 indexers, drop=drop, missing_dims=missing_dims 1185 ) ~/local/miniconda3/envs/met433a/lib/python3.9/site-packages/xarray/core/dataset.py in _isel_fancy(self, indexers, drop, missing_dims) 2370 # Note: we need to preserve the original indexers variable in order to merge the 2371 # coords below -> 2372 indexers_list = list(self._validate_indexers(indexers, missing_dims)) 2373 2374 variables: Dict[Hashable, Variable] = {} ~/local/miniconda3/envs/met433a/lib/python3.9/site-packages/xarray/core/dataset.py in _validate_indexers(self, indexers, missing_dims) 2203 2204 if v.ndim > 1: -> 2205 raise IndexError( 2206 "Unlabeled multi-dimensional array cannot be " 2207 "used for indexing: {}".format(k) IndexError: Unlabeled multi-dimensional array cannot be used for indexing: pres ``` I am not sure if this is a shortcoming of the MetPy function, or if I need to massage the xarray Dataset in some way. I don't see anything in the `interpolate_to_isosurface` docstring that indicates this shouldn't work. Ideas?

[sgdecker]

Oops, clumsy keyboarding there! Here is the working example. The feature request is for the commented out code to work:

from datetime import datetime, timedelta

import xarray as xr
from xarray.backends import NetCDF4DataStore
import metpy.calc as mpcalc
from metpy.interpolate import interpolate_to_isosurface
from metpy.plots import FilledContourPlot, MapPanel, PanelContainer
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()
    return data

def extract_and_normalize(xda, name):
    var = xda.squeeze()
    var = mpcalc.smooth_gaussian(var, 16)
    var = var.rename({var.metpy.vertical.name: 'pres'})
    var = var.rename(name).metpy.assign_latitude_longitude()
    return var

init_time = datetime(2021, 12, 6, 12)
plot_time = init_time + timedelta(hours=6)

nam = get_nam212(init_time, plot_time)

tmpk = extract_and_normalize(nam['Temperature_isobaric'], 'tmpk')
urel = extract_and_normalize(nam['u-component_of_wind_isobaric'], 'urel')
vrel = extract_and_normalize(nam['v-component_of_wind_isobaric'], 'vrel')
ds = xr.merge((tmpk, urel, vrel), combine_attrs='drop_conflicts')

ds['theta'] = mpcalc.potential_temperature(ds['pres'], ds['tmpk'])
ds['pv'] = mpcalc.potential_vorticity_baroclinic(ds['theta'], ds['pres'], ds['urel'], ds['vrel'])

# Feature request: Allow this
#pvu = 1e-6 * units('degK') / units('kg') * units('m2') / units('s')
#theta_dt = interpolate_to_isosurface(ds['pv'], ds['theta'], 2*pvu)

# Workaround
pv = ds['pv'].values
theta = ds['theta'].values
theta_dyn_trop = interpolate_to_isosurface(pv, theta, 2e-6)
theta_dt = xr.DataArray(theta_dyn_trop, dims=['y', 'x'], 
                        coords={'y': ds['theta'].y, 'x': ds['theta'].x, 'metpy_crs': ds['theta'].metpy_crs},
                        attrs=ds.attrs)

fcp = FilledContourPlot()
fcp.data = theta_dt
fcp.colorbar = 'horizontal'

mp = MapPanel()
mp.area = [-97, -60, 33, 49]
mp.layers = ['coastline', 'borders', 'states']
mp.plots = [fcp]

pc = PanelContainer()
pc.size = (12, 8)
pc.panels = [mp]
pc.show()

[dopplershift]

Completely reasonable request for that to work. However, the problem is xarray, not units. If you try:

theta_dt = interpolate_to_isosurface(ds['pv'].metpy.unit_array, ds['theta'].metpy.unit_array, 2*pvu)

you'll see that call succeeds (though it later fails in declarative because theta_dt isn't a DataArray).

Units isn't always the problem. 😉

I've updated the issue title to reflect this.