NCAR / wrf-python

A collection of diagnostic and interpolation routines for use with output from the Weather Research and Forecasting (WRF-ARW) Model.
https://wrf-python.readthedocs.io
Apache License 2.0
411 stars 155 forks source link

Improved documentation for wrf.get_cartopy, wrf.cartopy_xlim, and wrf.cartopy_ylim when passing a variable, not a file #172

Open jaredalee opened 2 years ago

jaredalee commented 2 years ago

On a related but separate note to the issue/request I posted in #171, I'd like to request that improved documentation be made when passing a variable and not a NetCDF4.Dataset to wrf.get_cartopy, wrf.cartopy_xlim, and wrf.cartopy_ylim. Right now the documentation for those three functions have no mention of what the variable requires in terms of attributes or anything else. By digging through the source code and going through lots of trial and error, I've come up with a working minimal example that I think could be added to the documentation to guide other users, particularly if they're trying to use non-WRF output.

First, to describe the file I'm using for context (for anyone curious, it's CAMS global composition forecast output, which is originally on a 40-km lat-lon rectangular grid, but which I regridded to the CMAQ 12-km curvilinear grid over CONUS, using geocat.comp.rgrid2rcm):

import xarray as xr
import wrf
import cartopy
import cartopy.crs as ccrs
import matplotlib as mpl
import matplotlib.pyplot as plt

ds_xr = xr.open_dataset(fname)
print(ds_xr)

Here's the output that describes the xarray.Dataset:

<xarray.Dataset>
Dimensions:    (time: 1, lat: 265, lon: 442)
Coordinates:
    latitude   (lat, lon) float32 ...
    longitude  (lat, lon) float32 ...
    times      (time) datetime64[ns] ...
Dimensions without coordinates: time, lat, lon
Data variables:
    cams_o3    (time, lat, lon) float64 ...
    cams_pm25  (time, lat, lon) float64 ...
Attributes:
    description:   CAMS forecast data regridded to CMAQ grid
    MAP_PROJ:      1
    CEN_LAT:       40.0
    CEN_LON:       -97.0
    TRUELAT1:      33.0
    TRUELAT2:      45.0
    STAND_LON:     -97.0
    MOAD_CEN_LAT:  40.0
    POLE_LAT:      90.0
    POLE_LON:      0.0
    DX:            12000.0
    DY:            12000.0

Now that the file/dataset is described, here's the code segment that would be needed to successfully pass a variable to wrf.cartopy_xlim and wrf.cartopy_ylim:

latitude = ds_xr.latitude
longitude = ds_xr.longitude
latitude = latitude.rename({'latitude':'XLAT', 'longitude':'XLONG'})
longitude = longitude.rename({'latitude':'XLAT', 'longitude':'XLONG'})

## Get the WRF-style projection attributes
map_proj = ds_xr.attrs['MAP_PROJ']
truelat1 = ds_xr.attrs['TRUELAT1']
truelat2 = ds_xr.attrs['TRUELAT2']
stand_lon = ds_xr.attrs['STAND_LON']
cen_lat = ds_xr.attrs['CEN_LAT']
cen_lon = ds_xr.attrs['CEN_LON']
pole_lat = ds_xr.attrs['POLE_LAT']
pole_lon = ds_xr.attrs['POLE_LON']
moad_cen_lat = ds_xr.attrs['MOAD_CEN_LAT']
dx = ds_xr.attrs['DX']
dy = ds_xr.attrs['DY']

## Create a dictionary of these projection attributes
dict_proj = {
   'MAP_PROJ':map_proj, 'CEN_LAT':cen_lat, 'CEN_LON':cen_lon,
   'TRUELAT1':truelat1, 'TRUELAT2':truelat2, 'MOAD_CEN_LAT':moad_cen_lat,
   'STAND_LON':stand_lon, 'POLE_LAT':pole_lat, 'POLE_LON':pole_lon, 'DX':dx, 'DY':dy,
   }

## Create an object of class wrf.WrfProj, then a cartopy mapping object
## This is essentially a manual reproduction of what wrf.get_cartopy() does
wrf_proj = wrf.util.getproj(**dict_proj)
proj_obj = wrf_proj.cartopy()
cart_proj = proj_obj

## Now get the cartopy projection x and y limits
latitude = latitude.assign_attrs(projection=wrf_proj)
longitude = longitude.assign_attrs(projection=wrf_proj)
cart_xlim = wrf.cartopy_xlim(var=latitude)
cart_ylim = wrf.cartopy_ylim(var=latitude)

Then to plot a data variable this code could be used (it wouldn't strictly be needed in the documentation, but would complete the example):

## Read in a data variable
o3 = ds_xr.cams_o3[0,:,:]
o3 = o3.rename({'latitude':'XLAT', 'longitude':'XLONG'})
o3 = o3.assign_attrs(projection=wrf_proj)

## Make a nice plot
borders = cartopy.feature.BORDERS
states  = cartopy.feature.NaturalEarthFeature(category='cultural',scale='10m',facecolor='none',name='admin_1_states_provinces_lakes')
lakes   = cartopy.feature.LAKES
data_crs = ccrs.PlateCarree()

mpl.rcParams['grid.color'] = 'gray'
mpl.rcParams['grid.linestyle'] = ':'
mpl.rcParams['figure.titlesize'] = 14
mpl.rcParams['font.size'] = 12
mpl.rcParams['savefig.bbox'] = 'tight'

fig = plt.figure()
ax = plt.subplot(projection=cart_proj)
ax.set_xlim(cart_xlim)
ax.set_ylim(cart_ylim)
ax.coastlines()
ax.gridlines()
ax.add_feature(borders, linestyle='-')
ax.add_feature(states, linewidth=0.25, edgecolor='black')
ax.add_feature(lakes, facecolor='none', linewidth=0.25, edgecolor='black')
extend = 'max'
cmap = mpl.rainbow
bounds = np.arange(0.0, 100.001, 5.0)
norm = mpl.colors.BoundaryNorm(bounds, cmap.N, extend=extend)
plt.contourf(wrf.to_np(longitude), wrf.to_np(latitude), wrf.to_np(o3),
   bounds, cmap=cmap, norm=norm, extend=extend, transform=data_crs, transform_first=(ax,True))

## Draw the colorbar to exactly match the height or width of the map
## Create colorbar axes (temporarily) anywhere
cax = fig.add_axes([0,0,0.1,0.1])
## Find the location of the main plot axes
posn = ax.get_position()
## Adjust the positioning and orientation of the colorbar and then draw it
## The colorbar will inherit the extend property from the plt.contourf call above and its norm/extend attributes
cbar_loc = 'bottom'
cbar_lab = 'Ozone Concentration [ppbv]'
if cbar_loc == 'bottom':
   cax.set_position([posn.x0, posn.y0-0.09, posn.width, 0.05])
   plt.colorbar(cax=cax, orientation='horizontal', label=cbar_lab)
elif cbar_loc == 'right':
   cax.set_position([posn.x0+posn.width+0.05, posn.y0, 0.04, posn.height])
   plt.colorbar(cax=cax, orientation='vertical', label=cbar_lab)

plt.suptitle('Regridded Forecast', y=0.92)
plt.show() # note that this method also works with plt.savefig