Merge pull request #1490 from jthielen/xarray-output-compat-3

Improved xarray compatibility on function input and output

conftest.py

@@ -57,3 +57,90 @@ def cfeature():
     to their parameter list.
     """
     return pytest.importorskip('cartopy.feature')
+
+
+@pytest.fixture()
+def test_da_lonlat():
+    """Return a DataArray with a lon/lat grid and no time coordinate for use in tests."""
+    pytest.importorskip('cartopy')
+
+    data = numpy.linspace(300, 250, 3 * 4 * 4).reshape((3, 4, 4))
+    ds = xarray.Dataset(
+        {'temperature': (['isobaric', 'lat', 'lon'], data)},
+        coords={
+            'isobaric': xarray.DataArray(
+                numpy.array([850., 700., 500.]),
+                name='isobaric',
+                dims=['isobaric'],
+                attrs={'units': 'hPa'}
+            ),
+            'lat': xarray.DataArray(
+                numpy.linspace(30, 40, 4),
+                name='lat',
+                dims=['lat'],
+                attrs={'units': 'degrees_north'}
+            ),
+            'lon': xarray.DataArray(
+                numpy.linspace(260, 270, 4),
+                name='lon',
+                dims=['lon'],
+                attrs={'units': 'degrees_east'}
+            )
+        }
+    )
+    ds['temperature'].attrs['units'] = 'kelvin'
+
+    return ds.metpy.parse_cf('temperature')
+
+
+@pytest.fixture()
+def test_da_xy():
+    """Return a DataArray with a x/y grid and a time coordinate for use in tests."""
+    pytest.importorskip('cartopy')
+
+    data = numpy.linspace(300, 250, 3 * 3 * 4 * 4).reshape((3, 3, 4, 4))
+    ds = xarray.Dataset(
+        {'temperature': (['time', 'isobaric', 'y', 'x'], data),
+         'lambert_conformal': ([], '')},
+        coords={
+            'time': xarray.DataArray(
+                numpy.array([numpy.datetime64('2018-07-01T00:00'),
+                             numpy.datetime64('2018-07-01T06:00'),
+                             numpy.datetime64('2018-07-01T12:00')]),
+                name='time',
+                dims=['time']
+            ),
+            'isobaric': xarray.DataArray(
+                numpy.array([850., 700., 500.]),
+                name='isobaric',
+                dims=['isobaric'],
+                attrs={'units': 'hPa'}
+            ),
+            'y': xarray.DataArray(
+                numpy.linspace(-1000, 500, 4),
+                name='y',
+                dims=['y'],
+                attrs={'units': 'km'}
+            ),
+            'x': xarray.DataArray(
+                numpy.linspace(0, 1500, 4),
+                name='x',
+                dims=['x'],
+                attrs={'units': 'km'}
+            )
+        }
+    )
+    ds['temperature'].attrs = {
+        'units': 'kelvin',
+        'grid_mapping': 'lambert_conformal'
+    }
+    ds['lambert_conformal'].attrs = {
+        'grid_mapping_name': 'lambert_conformal_conic',
+        'standard_parallel': 50.0,
+        'longitude_of_central_meridian': -107.0,
+        'latitude_of_projection_origin': 50.0,
+        'earth_shape': 'spherical',
+        'earth_radius': 6367470.21484375
+    }
+
+    return ds.metpy.parse_cf('temperature')

docs/_templates/overrides/metpy.calc.rst

@@ -81,6 +81,7 @@ Soundings
       most_unstable_parcel
       parcel_profile
       parcel_profile_with_lcl
+      parcel_profile_with_lcl_as_dataset
       significant_tornado
       storm_relative_helicity
       supercell_composite
@@ -140,7 +141,6 @@ Mathematical Functions
       cross_section_components
       first_derivative
       gradient
-      grid_deltas_from_dataarray
       laplacian
       lat_lon_grid_deltas
       normal_component
@@ -195,6 +195,7 @@ Other
       get_layer_heights
       get_perturbation
       isentropic_interpolation
+      isentropic_interpolation_as_dataset
       nearest_intersection_idx
       parse_angle
       reduce_point_density

docs/_templates/overrides/metpy.xarray.rst

@@ -11,3 +11,11 @@ Accessors
 
 .. autoclass:: MetPyDatasetAccessor()
     :members:
+
+Functions
+---------
+
+   .. autosummary::
+      :toctree: ./
+
+      grid_deltas_from_dataarray

examples/cross_section.py

@@ -50,27 +50,21 @@
 # For this example, we will be plotting potential temperature, relative humidity, and
 # tangential/normal winds. And so, we need to calculate those, and add them to the dataset:
 
-temperature, pressure, specific_humidity = xr.broadcast(cross['Temperature'],
-                                                        cross['isobaric'],
-                                                        cross['Specific_humidity'])
-
-theta = mpcalc.potential_temperature(pressure, temperature)
-rh = mpcalc.relative_humidity_from_specific_humidity(pressure, temperature, specific_humidity)
-
-# These calculations return unit arrays, so put those back into DataArrays in our Dataset
-cross['Potential_temperature'] = xr.DataArray(theta,
-                                              coords=temperature.coords,
-                                              dims=temperature.dims,
-                                              attrs={'units': theta.units})
-cross['Relative_humidity'] = xr.DataArray(rh,
-                                          coords=specific_humidity.coords,
-                                          dims=specific_humidity.dims,
-                                          attrs={'units': rh.units})
-
+cross['Potential_temperature'] = mpcalc.potential_temperature(
+    cross['isobaric'],
+    cross['Temperature']
+)
+cross['Relative_humidity'] = mpcalc.relative_humidity_from_specific_humidity(
+    cross['isobaric'],
+    cross['Temperature'],
+    cross['Specific_humidity']
+)
 cross['u_wind'] = cross['u_wind'].metpy.convert_units('knots')
 cross['v_wind'] = cross['v_wind'].metpy.convert_units('knots')
-cross['t_wind'], cross['n_wind'] = mpcalc.cross_section_components(cross['u_wind'],
-                                                                   cross['v_wind'])
+cross['t_wind'], cross['n_wind'] = mpcalc.cross_section_components(
+    cross['u_wind'],
+    cross['v_wind']
+)
 
 print(cross)
 

examples/isentropic_example.py

@@ -37,22 +37,9 @@
 
 #############################
 # We will reduce the dimensionality of the data as it is pulled in to remove an empty time
-# dimension.
+# dimension, as well as add longitude and latitude as coordinates (instead of data variables).
 
-# Assign data to variable names
-lat = data['lat']
-lon = data['lon']
-lev = data['isobaric']
-times = data['time']
-
-tmp = data['Temperature'][0]
-uwnd = data['u_wind'][0]
-vwnd = data['v_wind'][0]
-spech = data['Specific_humidity'][0]
-
-# pint doesn't understand gpm
-data['Geopotential_height'].attrs['units'] = 'meter'
-hgt = data['Geopotential_height'][0]
+data = data.squeeze().set_coords(['lon', 'lat'])
 
 #############################
 # To properly interpolate to isentropic coordinates, the function must know the desired output
@@ -65,55 +52,55 @@
 #
 # Once three dimensional data in isobaric coordinates has been pulled and the desired
 # isentropic levels created, the conversion to isentropic coordinates can begin. Data will be
-# passed to the function as below. The function requires that isentropic levels, isobaric
-# levels, and temperature be input. Any additional inputs (in this case relative humidity, u,
-# and v wind components) will be linearly interpolated to isentropic space.
-
-isent_ana = mpcalc.isentropic_interpolation(isentlevs,
-                                            lev,
-                                            tmp,
-                                            spech,
-                                            uwnd,
-                                            vwnd,
-                                            hgt,
-                                            temperature_out=True)
+# passed to the function as below. The function requires that isentropic levels, as well as a
+# DataArray of temperature on isobaric coordinates be input. Any additional inputs (in this
+# case specific humidity, geopotential height, and u and v wind components) will be
+# logarithmicaly interpolated to isentropic space.
+
+isent_data = mpcalc.isentropic_interpolation_as_dataset(
+    isentlevs,
+    data['Temperature'],
+    data['u_wind'],
+    data['v_wind'],
+    data['Specific_humidity'],
+    data['Geopotential_height']
+)
 
 #####################################
-# The output is a list, so now we will separate the variables to different names before
-# plotting.
+# The output is an xarray Dataset:
 
-isentprs, isenttmp, isentspech, isentu, isentv, isenthgt = isent_ana
-isentu.ito('kt')
-isentv.ito('kt')
+isent_data
 
 ########################################
-# A quick look at the shape of these variables will show that the data is now in isentropic
-# coordinates, with the number of vertical levels as specified above.
+# Note that the units on our wind variables are not ideal for plotting. Instead, let us
+# convert them to more appropriate values.
 
-print(isentprs.shape)
-print(isentspech.shape)
-print(isentu.shape)
-print(isentv.shape)
-print(isenttmp.shape)
-print(isenthgt.shape)
+isent_data['u_wind'] = isent_data['u_wind'].metpy.convert_units('kt')
+isent_data['v_wind'] = isent_data['v_wind'].metpy.convert_units('kt')
 
 #################################
 # **Converting to Relative Humidity**
 #
 # The NARR only gives specific humidity on isobaric vertical levels, so relative humidity will
 # have to be calculated after the interpolation to isentropic space.
 
-isentrh = 100 * mpcalc.relative_humidity_from_specific_humidity(isentprs, isenttmp, isentspech)
+isent_data['Relative_humidity'] = mpcalc.relative_humidity_from_specific_humidity(
+    isent_data['pressure'],
+    isent_data['temperature'],
+    isent_data['Specific_humidity']
+).metpy.convert_units('percent')
 
 #######################################
 # **Plotting the Isentropic Analysis**
 
-# Set up our projection
+# Set up our projection and coordinates
 crs = ccrs.LambertConformal(central_longitude=-100.0, central_latitude=45.0)
+lon = isent_data['pressure'].metpy.longitude
+lat = isent_data['pressure'].metpy.latitude
 
 # Coordinates to limit map area
 bounds = [(-122., -75., 25., 50.)]
-# Choose a level to plot, in this case 296 K
+# Choose a level to plot, in this case 296 K (our sole level in this example)
 level = 0
 
 fig = plt.figure(figsize=(17., 12.))
@@ -125,26 +112,28 @@
 
 # Plot the surface
 clevisent = np.arange(0, 1000, 25)
-cs = ax.contour(lon, lat, isentprs[level, :, :], clevisent,
-                colors='k', linewidths=1.0, linestyles='solid', transform=ccrs.PlateCarree())
+cs = ax.contour(lon, lat, isent_data['pressure'].isel(isentropic_level=level),
+                clevisent, colors='k', linewidths=1.0, linestyles='solid',
+                transform=ccrs.PlateCarree())
 cs.clabel(fontsize=10, inline=1, inline_spacing=7, fmt='%i', rightside_up=True,
           use_clabeltext=True)
 
 # Plot RH
-cf = ax.contourf(lon, lat, isentrh[level, :, :], range(10, 106, 5),
-                 cmap=plt.cm.gist_earth_r, transform=ccrs.PlateCarree())
+cf = ax.contourf(lon, lat, isent_data['Relative_humidity'].isel(isentropic_level=level),
+                 range(10, 106, 5), cmap=plt.cm.gist_earth_r, transform=ccrs.PlateCarree())
 cb = fig.colorbar(cf, orientation='horizontal', aspect=65, shrink=0.5, pad=0.05,
                   extendrect='True')
 cb.set_label('Relative Humidity', size='x-large')
 
 # Plot wind barbs
-ax.barbs(lon.values, lat.values, isentu[level, :, :].m, isentv[level, :, :].m, length=6,
+ax.barbs(lon.values, lat.values, isent_data['u_wind'].isel(isentropic_level=level).values,
+         isent_data['v_wind'].isel(isentropic_level=level).values, length=6,
          regrid_shape=20, transform=ccrs.PlateCarree())
 
 # Make some titles
 ax.set_title(f'{isentlevs[level]:~.0f} Isentropic Pressure (hPa), Wind (kt), '
              'Relative Humidity (percent)', loc='left')
-add_timestamp(ax, times[0].values.astype('datetime64[ms]').astype('O'),
+add_timestamp(ax, isent_data['time'].values.astype('datetime64[ms]').astype('O'),
               y=0.02, high_contrast=True)
 fig.tight_layout()
 
@@ -157,7 +146,10 @@
 
 
 # Calculate Montgomery Streamfunction and scale by 10^-2 for plotting
-msf = mpcalc.montgomery_streamfunction(isenthgt, isenttmp) / 100.
+msf = mpcalc.montgomery_streamfunction(
+    isent_data['Geopotential_height'],
+    isent_data['temperature']
+).values / 100.
 
 # Choose a level to plot, in this case 296 K
 level = 0
@@ -177,20 +169,21 @@
           use_clabeltext=True)
 
 # Plot RH
-cf = ax.contourf(lon, lat, isentrh[level, :, :], range(10, 106, 5),
-                 cmap=plt.cm.gist_earth_r, transform=ccrs.PlateCarree())
+cf = ax.contourf(lon, lat, isent_data['Relative_humidity'].isel(isentropic_level=level),
+                 range(10, 106, 5), cmap=plt.cm.gist_earth_r, transform=ccrs.PlateCarree())
 cb = fig.colorbar(cf, orientation='horizontal', aspect=65, shrink=0.5, pad=0.05,
                   extendrect='True')
 cb.set_label('Relative Humidity', size='x-large')
 
-# Plot wind barbs.
-ax.barbs(lon.values, lat.values, isentu[level, :, :].m, isentv[level, :, :].m, length=6,
+# Plot wind barbs
+ax.barbs(lon.values, lat.values, isent_data['u_wind'].isel(isentropic_level=level).values,
+         isent_data['v_wind'].isel(isentropic_level=level).values, length=6,
          regrid_shape=20, transform=ccrs.PlateCarree())
 
 # Make some titles
 ax.set_title(f'{isentlevs[level]:~.0f} Montgomery Streamfunction '
              r'($10^{-2} m^2 s^{-2}$), Wind (kt), Relative Humidity (percent)', loc='left')
-add_timestamp(ax, times[0].values.astype('datetime64[ms]').astype('O'),
+add_timestamp(ax, isent_data['time'].values.astype('datetime64[ms]').astype('O'),
               y=0.02, pretext='Valid: ', high_contrast=True)
 
 fig.tight_layout()