lib.plotting_functions module

. Generic computation and plotting helper functions

Functions

load_dataset(): generalized load dataset method used for plotting/analysis functions
use_this_norm(): switches matplotlib color normalization method
get_difference_colors(values): Provide a color norm and colormap assuming values is a difference field.
mask_land_or_ocean(arr, msk, use_nan=False): Apply a land or ocean mask to provided variable.
get_central_longitude(*args): Determine central longitude for maps.
global_average(fld, wgt, verbose=False): pure numpy global average.
spatial_average(indata, weights=None, spatial_dims=None): Compute spatial average
wgt_rmse(fld1, fld2, wgt):: Calculate the area-weighted RMSE.
annual_mean(data, whole_years=False, time_name=’time’):: Calculate annual averages from time series data.
seasonal_mean(data, season=None, is_climo=None):: Calculates the time-weighted seasonal average (or average over all time).
domain_stats(data, domain):: Provides statistics in specified region.
make_polar_plot(wks, case_nickname, base_nickname,: case_climo_yrs, baseline_climo_yrs, d1:xr.DataArray, d2:xr.DataArray, difference:Optional[xr.DataArray]=None, domain:Optional[list]=None, hemisphere:Optional[str]=None, **kwargs):

Make a stereographic polar plot for the given data and hemisphere.
plot_map_vect_and_save(wks, case_nickname, base_nickname,: case_climo_yrs, baseline_climo_yrs, plev, umdlfld_nowrap, vmdlfld_nowrap, uobsfld_nowrap, vobsfld_nowrap, udiffld_nowrap, vdiffld_nowrap, **kwargs):

Plots a vector field on a map.
plot_map_and_save(wks, case_nickname, base_nickname,: case_climo_yrs, baseline_climo_yrs, mdlfld, obsfld, diffld, **kwargs):

Map plots of mdlfld, obsfld, and their difference, diffld.
pres_from_hybrid(psfc, hya, hyb, p0=100000.):: Converts a hybrid level to a pressure
vert_remap(x_mdl, p_mdl, plev): Interpolates to specified pressure levels.
lev_to_plev(data, ps, hyam, hybm, P0=100000., new_levels=None, convert_to_mb=False): Interpolate model hybrid levels to specified pressure levels.
pmid_to_plev(data, pmid, new_levels=None, convert_to_mb=False): Interpolate data from hybrid-sigma levels to isobaric levels using provided mid-level pressures.
zonal_mean_xr(fld): Average over all dimensions except lev and lat.
validate_dims(fld, list_of_dims): Checks if specified dimensions are in a DataArray
lat_lon_validate_dims(fld): Check if input field has lat and lon.
zm_validate_dims(fld): Check for dimensions for zonal average.
zonal_plot(lat, data, ax=None, color=None, **kwargs): Make a line plot or pressure-latitude plot of data.
meridional_plot(lon, data, ax=None, color=None, **kwargs): Make a line plot or pressure-longitude plot of data.
prep_contour_plot: Preparation for making contour plots.
plot_zonal_mean_and_save: zonal mean plot
plot_meridional_mean_and_save: meridioanl mean plot
square_contour_difference: Produce filled contours of fld1, fld2, and their difference with square axes.

Notes

This module includes several “private” methods intended for internal use only.

_plot_line(axobject, xdata, ydata, color, **kwargs): Create a generic line plot

_meridional_plot_line

_zonal_plot_line

_zonal_plot_preslat

_meridional_plot_preslon

lib.plotting_functions.annual_mean(data, whole_years=False, time_name='time')[source]

Calculate annual averages from monthly time series data.

Parameters:

data (xr.DataArray or xr.Dataset) – monthly data values with temporal dimension
whole_years (bool, optional) – whether to restrict endpoints of the average to start at first January and end at last December
time_name (str, optional) – name of the time dimension, defaults to time

Returns:

result – data reduced to annual averages

Return type:

xr.DataArray or xr.Dataset

Notes

This function assumes monthly data, and weights the average by the number of days in each month.

result includes an attribute that reports the date range used for the average.

lib.plotting_functions.domain_stats(data, domain)[source]

Provides statistics in specified region.

Parameters:

data (xarray.DataArray) – data values
domain (list or tuple or numpy.ndarray) – the domain specification as: [west_longitude, east_longitude, south_latitude, north_latitude]

Returns:

x_region_mean (float) – the regional area-weighted average
x_region_max (float) – the maximum value in the region
x_region_min (float) – the minimum value in the region

Notes

Currently assumes ‘lat’ is a dimension and uses cos(lat) as weight. Should use spatial_average

See also

validate_dims

lib.plotting_functions.lev_to_plev(data, ps, hyam, hybm, P0=100000.0, new_levels=None, convert_to_mb=False)[source]

Interpolate model hybrid levels to specified pressure levels.

Parameters:

data
ps – surface pressure
hyam – hybrid-sigma A and B coefficients
hybm – hybrid-sigma A and B coefficients
P0 (float, optional) – reference pressure, defaults to 100000 Pa
new_levels (numpy.ndarray, optional) – 1-D array containing pressure levels in Pascals (Pa). If not specified, then the levels will be set to the GeoCAT defaults, which are (in hPa): 1000, 925, 850, 700, 500, 400, 300, 250, 200, 150, 100, 70, 50, 30, 20, 10, 7, 5, 3, 2, 1
convert_to_mb (bool, optional) – If True, then vertical (lev) dimension will have values of mb/hPa, otherwise the units are Pa.

Returns:

data interpolated to new pressure levels

Return type:

data_interp_rename

Notes

The function interp_hybrid_to_pressure used here is dask-enabled, and so can potentially be sped-up via the use of a DASK cluster.

lib.plotting_functions.load_dataset(fils)[source]

This method exists to get an xarray Dataset from input file information that can be passed into the plotting methods.

Parameters:: fils (list) – strings or paths to input file(s)
Return type:: xr.Dataset

Notes

When just one entry is provided, use open_dataset, otherwise open_mfdatset

lib.plotting_functions.make_polar_plot(wks, case_nickname, base_nickname, case_climo_yrs, baseline_climo_yrs, d1, d2, difference=None, pctchange=None, domain=None, hemisphere=None, obs=False, **kwargs)[source]

Make a stereographic polar plot for the given data and hemisphere.

Parameters:

wks (str or Path) – output file path
case_nickname (str) – short case name for d1
base_nickname (str) – short case name for d2
case_climo_yrs (list) – years for case d1, used for annotation
baseline_climo_yrs (list) – years for case d2, used for annotation
d1 (xr.DataArray) – input data, must contain dimensions lat and lon
d2 (xr.DataArray) – input data, must contain dimensions lat and lon
difference (xr.DataArray, optional) – data to use as the difference, otherwise d2 - d1
pctchange (xr.DataArray, optional data to use as the percent change)
domain (list, optional) – the domain to plot, specified as [west_longitude, east_longitude, south_latitude, north_latitude] Defaults to pole to 45-degrees, all longitudes
hemisphere ({'NH', 'SH'}, optional) – Hemsiphere to plot
kwargs (dict, optional) – variable-dependent options for plots, See Notes.

Notes

Uses contourf. No contour lines (yet).
kwargs is checked for:
- colormap
- contour_levels
- contour_levels_range
- diff_contour_levels
- diff_contour_range
- diff_colormap
- units

lib.plotting_functions.mask_land_or_ocean(arr, msk, use_nan=False)[source]

Apply a land or ocean mask to provided variable.

Parameters:

arr (xarray.DataArray) – the xarray variable to apply the mask to.
msk (xarray.DataArray) – the xarray variable that contains the land or ocean mask, assumed to be the same shape as “arr”.
use_nan (bool, optional) – argument for whether to set the missing values to np.nan values instead of the defaul “-999.” values.

Returns:

arr – Same as input arr but masked as specified.

Return type:

xarray.DataArray

lib.plotting_functions.meridional_plot(lon, data, ax=None, color=None, **kwargs)[source]

Make meridional plot

Determine which kind of meridional plot is needed based on the input variable’s dimensions.

Parameters:

lon – longitude
data – input data
ax (Axes, optional) – axes object to use
color (str or mpl color specification) – color for the curve
kwargs (dict, optional) – plotting options

Notes

Checks if there is a lev dimension to determine if it is a lon-pres plot or a line plot.

lib.plotting_functions.my_formatwarning(msg, *args, **kwargs)[source]: Issue msg as warning.

lib.plotting_functions.plot_map_and_save(wks, case_nickname, base_nickname, case_climo_yrs, baseline_climo_yrs, mdlfld, obsfld, diffld, pctld, obs=False, **kwargs)[source]

Plots mdlfld, obsfld, diffld in a 3-row panel plot of maps.

Parameters:

wks (str or Path) – Output file path.
case_nickname (str) – Short name for case.
base_nickname (str) – Short name for base case.
case_climo_yrs (list) – List of years in case climatology, used for annotation.
baseline_climo_yrs (list) – List of years in base case climatology, used for annotation.
mdlfld (xarray.DataArray) – Input data for case.
obsfld (xarray.DataArray) – Input data for base case.
diffld (xarray.DataArray) – Input difference data.
pctld (xarray.DataArray) – Input percent difference data.
obs (bool, optional) – If True, use observational data. Default is False.
**kwargs (dict, optional) – Variable-specific options. See Notes.

Notes

The kwargs parameter is expected to be a variable-specific section, possibly provided by an ADF Variable Defaults YAML file. Currently it is inspected for:

colormap (str): Name of matplotlib colormap.

contour_levels (list or tuple): Explicit values or a tuple (min, max, step).

diff_colormap

diff_contour_levels

tiString (str): Title string.

tiFontSize (int): Title font size.

mpl (dict): Any matplotlib kwargs that should be passed along.

Example YAML for mpl options:

mpl:
    subplots:
        figsize: (3, 9)
    contourf:
        levels: 15
        cmap: Blues
    colorbar:
        shrink: 0.4

This is experimental, and if you find yourself doing much with this, you probably should write a new plotting script that does not rely on this module.

When these are not provided, colormap is set to ‘coolwarm’ and limits/levels are set by data range.

lib.plotting_functions.plot_map_vect_and_save(wks, case_nickname, base_nickname, case_climo_yrs, baseline_climo_yrs, plev, umdlfld_nowrap, vmdlfld_nowrap, uobsfld_nowrap, vobsfld_nowrap, udiffld_nowrap, vdiffld_nowrap, obs=False, **kwargs)[source]

This plots a vector plot.

Vector fields constructed from x- and y-components (u, v).

Parameters:

wks (str or Path) – output file path
case_nickname (str) – short name for case
base_nickname (str) – short name for base case
case_climo_yrs (list) – list of years in case climatology, used for annotation
baseline_climo_yrs (list) – list of years in base case climatology, used for annotation
plev (str or float or None) – if not None, label denoting the pressure level
umdlfld_nowrap (xarray.DataArray) – input data for case, the x- and y- components of the vectors
vmdlfld_nowrap (xarray.DataArray) – input data for case, the x- and y- components of the vectors
uobsfld_nowrap (xarray.DataArray) – input data for base case, the x- and y- components of the vectors
vobsfld_nowrap (xarray.DataArray) – input data for base case, the x- and y- components of the vectors
udiffld_nowrap (xarray.DataArray) – input difference data, the x- and y- components of the vectors
vdiffld_nowrap (xarray.DataArray) – input difference data, the x- and y- components of the vectors
kwargs (dict, optional) – variable-specific options, See Notes

Notes

kwargs expected to be a variable-specific section, possibly provided by an ADF Variable Defaults YAML file. Currently it is inspected for: - central_longitude - var_name - case_name - baseline - tiString - tiFontSize - units

_Note_ The title string constructed by kwargs appears to not be used.

lib.plotting_functions.plot_meridional_mean_and_save(wks, case_nickname, base_nickname, case_climo_yrs, baseline_climo_yrs, adata, bdata, has_lev, latbounds=None, obs=False, **kwargs)[source]

Default meridional mean plot

Parameters:

wks – the figure object to plot in
case_nickname (str) – short name of adata case, use for annotation
base_nickname (str) – short name of bdata case, use for annotation
case_climo_yrs (list) – years in the adata case, use for annotation
baseline_climo_yrs (list:) – years in the bdata case, use for annotation
adata (xarray.DataArray) – data to plot ([lev], [lat], lon). The vertical coordinate (lev) must be pressure levels.
bdata (xarray.DataArray) – baseline or observations to plot adata against. It must have the same dimensions and vertical levels as adata.
has_lev (bool) – whether lev dimension is present
latbounds (numbers.Number or slice, optional) – indicates latitude bounds to average over if it is a number, assume symmetric about equator, otherwise expects slice(south, north) defaults to slice(-5,5)
kwargs (dict, optional) – optional dictionary of plotting options, See Notes

Notes

For 2-d variables (reduced to (lon,)):
- 2 panels: (top) meridional mean, (bottom) difference
For 3-D variables (reduced to (lev,lon)):
- 3 panels: (top) meridonal mean adata, (middle) meridional mean bdata, (bottom) difference
- pcolormesh/contour plot
kwargs -> optional dictionary of plotting options
** Expecting this to be variable-specific section, possibly provided by ADF Variable Defaults YAML file.** - colormap -> str, name of matplotlib colormap - contour_levels -> list of explicit values or a tuple: (min, max, step) - diff_colormap -> str, name of matplotlib colormap used for different plot - diff_contour_levels -> list of explicit values or a tuple: (min, max, step) - tiString -> str, Title String - tiFontSize -> int, Title Font Size - mpl -> dict, This should be any matplotlib kwargs that should be passed along. Keep reading:
Organize these by the mpl function. In this function (plot_meridional_mean_and_save)

we will check for an entry called subplots, contourf, and colorbar. So the YAML might looks something like: ``` mpl:

subplots: figsize: (3, 9) contourf: levels: 15 cmap: Blues colorbar: shrink: 0.4

```

lib.plotting_functions.plot_zonal_mean_and_save(wks, case_nickname, base_nickname, case_climo_yrs, baseline_climo_yrs, adata, bdata, has_lev, log_p, obs=False, **kwargs)[source]

This is the default zonal mean plot

Parameters:

adata (data to plot ([lev], lat, [lon]).) – The vertical coordinate (lev) must be pressure levels.
bdata (baseline or observations to plot adata against.) –
- For 2-d variables (reduced to (lat,)): + 2 panels: (top) zonal mean, (bottom) difference
- For 3-D variables (reduced to (lev,lat)): + 3 panels: (top) zonal mean adata, (middle) zonal mean bdata, (bottom) difference + pcolormesh/contour plot
options (kwargs -> optional dictionary of plotting) – ** Expecting this to be variable-specific section, possibly provided by ADF Variable Defaults YAML file.**
str (- tiString ->)
colormap (name of matplotlib)
tuple (- contour_levels -> list of explict values or a)
diff_colormap (-)
diff_contour_levels (-)
str
String (Title)
int (- tiFontSize ->)
Size (Title Font)
dict (- mpl ->) –
- Organize these by the mpl function. In this function (plot_map_and_save) we will check for an entry called subplots, contourf, and colorbar. So the YAML might looks something like: ```
  
  mpl:
  
  subplots:
  figsize: (3, 9)
  
  contourf:
  levels: 15 cmap: Blues
  
  colorbar:
  shrink: 0.4
  
```
reading (This should be any matplotlib kwargs that should be passed along. Keep) –
- Organize these by the mpl function. In this function (plot_map_and_save) we will check for an entry called subplots, contourf, and colorbar. So the YAML might looks something like: ```
  
  mpl:
  
  subplots:
  figsize: (3, 9)
  
  contourf:
  levels: 15 cmap: Blues
  
  colorbar:
  shrink: 0.4
  
```

lib.plotting_functions.pmid_to_plev(data, pmid, new_levels=None, convert_to_mb=False)[source]

Interpolate data from hybrid-sigma levels to isobaric levels.

Parameters:

data (xarray.DataArray) – field with a ‘lev’ coordinate
pmid (xarray.DataArray) – the pressure field (Pa), same shape as data
new_levels (optional) – the output pressure levels (Pa), defaults to standard levels
convert_to_mb (bool, optional) – flag to convert output to mb (i.e., hPa), defaults to False

Returns:

output – data interpolated onto new_levels

Return type:

xarray.DataArray

lib.plotting_functions.prep_contour_plot(adata, bdata, diffdata, pctdata, **kwargs)[source]

Preparation for making contour plots.

Prepares for making contour plots of adata, bdata, diffdata, and pctdata, which is presumably the difference between adata and bdata. - set colormap from kwargs or defaults to coolwarm - set contour levels from kwargs or 12 evenly spaced levels to span the data - normalize colors based on specified contour levels or data range - set option for linear or log pressure when applicable - similar settings for difference, defaults to symmetric about zero - separates Matplotlib kwargs into their own dicts

Parameters:

adata – the data to be plotted
bdata – the data to be plotted
diffdata – the data to be plotted
pctdata – the data to be plotted
kwargs (dict, optional) – plotting options

Returns:

a dict with the following: - ‘subplots_opt’: mpl kwargs for subplots - ‘contourf_opt’: mpl kwargs for contourf - ‘colorbar_opt’: mpl kwargs for colorbar - ‘diff_colorbar_opt’ : mpl kwargs for difference colorbar - ‘normdiff’: color normalization for difference panel - ‘cmapdiff’: colormap for difference panel - ‘levelsdiff’: contour levels for difference panel - ‘cmap1’: color map for a and b panels - ‘norm1’: color normalization for a and b panels - ‘levels1’ : contour levels for a and b panels - ‘plot_log_p’ : true/false whether to plot log(pressure) axis

Return type:

dict

lib.plotting_functions.pres_from_hybrid(psfc, hya, hyb, p0=100000.0)[source]

Calculates pressure field

pressure derived with the formula: `p = a(k)*p0 + b(k)*ps`

Parameters:

psfc – surface pressure
hya – hybrid-sigma A and B coefficients
hyb – hybrid-sigma A and B coefficients
p0 (optional) – reference pressure, defaults to 100000 Pa

Return type:

pressure, size is same as psfc with len(hya) levels

lib.plotting_functions.seasonal_mean(data, season=None, is_climo=None)[source]

Calculates the time-weighted seasonal average (or average over all time).

Parameters:

data (xarray.DataArray or xarray.Dataset) – data to be averaged
season (str, optional) – the season to extract from data If season is ANN or None, average all available time.
is_climo (bool, optional) – If True, expects data to have time or month dimenion of size 12. If False, then ‘time’ must be a coordinate, and the time.dt.days_in_month attribute must be available.

Returns:

the average of data in season season

Return type:

xarray.DataArray or xarray.Dataset

Notes

If the data is a climatology, the code will make an attempt to understand the time or month dimension, but will assume that it is ordered from January to December. If the data is a climatology and is just a numpy array with one dimension that is size 12, it will assume that dimension is time running from January to December.

lib.plotting_functions.spatial_average(indata, weights=None, spatial_dims=None)[source]

Compute spatial average.

Parameters:

indata (xr.DataArray) – input data
weights (np.ndarray or xr.DataArray, optional) – the weights to apply, see Notes for default behavior
spatial_dims (list, optional) – list of dimensions to average, see Notes for default behavior

Returns:

weighted average of indata

Return type:

xr.DataArray

Notes

When weights is not provided, tries to find sensible values. If there is a ‘lat’ dimension, use cos(lat). If there is a ‘ncol’ dimension, looks for area in indata. Otherwise, set to equal weights.

Makes an attempt to identify the spatial variables when spatial_dims is None. Will average over ncol if present, and then will check for lat and lon. When none of those three are found, raise an AdfError.

lib.plotting_functions.square_contour_difference(fld1, fld2, **kwargs)[source]

Produce filled contours of fld1, fld2, and their difference with square axes.

Intended use is latitude-by-month to show the annual cycle. Example use case: use climo files to get data, take zonal averages, rename “time” to “month” if desired, and then provide resulting DataArrays to this function.

Parameters:

fld1 (xarray.DataArray) – 2-dimensional DataArrays with same shape
fld2 (xarray.DataArray) – 2-dimensional DataArrays with same shape
**kwargs (dict, optional) – optional keyword arguments this function _only checks_ kwargs for case1name, case2name

Returns:

figure object

Return type:

fig

Notes

Assumes fld1.shape == fld2.shape and len(fld1.shape) == 2

Will try to label the cases by looking for case1name and case2name in kwargs, and then fld1[‘case’] and fld2[‘case’] (i.e., attributes) If no case name is found proceeds with empty strings. IF THERE IS A BETTER CONVENTION WE SHOULD USE IT.

Each panel also puts the Min/Max values into the title string.

Axis labels are upper-cased names of the coordinates of fld1. Ticks are automatic with the exception that if the first coordinate is “month” and is length 12, use np.arange(1,13).

lib.plotting_functions.use_this_norm()[source]: Just use the right normalization; avoids a deprecation warning.

lib.plotting_functions.validate_dims(fld, list_of_dims)[source]

Check if specified dimensions are in a DataArray.

Parameters:

fld (xarray.DataArray) – field to check for named dimensions
list_of_dims (list) – list of strings that specifiy the dimensions to check for

Returns:

dict with keys that are “has_{x}” where x is the name from list_of_dims and values that are boolean

Return type:

dict

lib.plotting_functions.vert_remap(x_mdl, p_mdl, plev)[source]

Apply simple 1-d interpolation to a field

Parameters:

x_mdl (xarray.DataArray or numpy.ndarray) – input data
p_mdl (xarray.DataArray or numpy.ndarray) – pressure field, same shape as x_mdl
plev (xarray.DataArray or numpy.ndarray) – the new pressures

Returns:

x_mdl interpolated to plev

Return type:

output

Notes

Interpolation done in log pressure

lib.plotting_functions.wgt_rmse(fld1, fld2, wgt)[source]

Calculate the area-weighted RMSE.

Parameters:

fld1 (array-like) – 2-dimensional spatial fields with the same shape. They can be xarray DataArray or numpy arrays.
fld2 (array-like) – 2-dimensional spatial fields with the same shape. They can be xarray DataArray or numpy arrays.
wgt (array-like) – the weight vector, expected to be 1-dimensional, matching length of one dimension of the data.

Returns:

float – root mean squared error
Notes
`rmse = sqrt( mean( (fld1 - fld2)**2 ) )`

lib.plotting_functions.zm_validate_dims(fld)[source]

Check for dimensions for zonal average.

Looks for dimensions called ‘lev’ and ‘lat’.

Parameters:

fld (xarray.DataArray) – field to check for lat and/or lev dimensions

Returns:

tuple – (has_lat, has_lev) each are bool
None – If ‘lat’ is not in dimensions, returns None.

lib.plotting_functions.zonal_mean_xr(fld)[source]: Average over all dimensions except lev and lat.

lib.plotting_functions.zonal_plot(lat, data, ax=None, color=None, **kwargs)[source]

Make zonal plot

Determine which kind of zonal plot is needed based on the input variable’s dimensions.

Parameters:

lat – latitude
data – input data
ax (Axes, optional) – axes object to use
color (str or mpl color specification) – color for the curve
kwargs (dict, optional) – plotting options

Notes

Checks if there is a lev dimension to determine if it is a lat-pres plot or a line plot.