diff --git a/lib/plotting_functions.py b/lib/plotting_functions.py index e232a31c2..fa8bcc2a7 100644 --- a/lib/plotting_functions.py +++ b/lib/plotting_functions.py @@ -1,9 +1,86 @@ -""" -This module contains generic -plotting helper functions that -can be used across multiple -different user-provided -plotting scripts. +""" . +Generic computation and plotting helper functions + +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 + """ #import statements: @@ -62,10 +139,22 @@ def use_this_norm(): def get_difference_colors(values): """Provide a color norm and colormap assuming this is a difference field. - Values can be either the data field or a set of specified contour levels. - - Uses 'OrRd' colormap for positive definite, 'BuPu_r' for negative definite, - and 'RdBu_r' centered on zero if there are values of both signs. + Parameters + ---------- + values : array-like + can be either the data field or a set of specified contour levels. + + Returns + ------- + dnorm + Matplotlib color nomalization + cmap + Matplotlib colormap + + Notes + ----- + Uses 'OrRd' colormap for positive definite, 'BuPu_r' for negative definite, + and 'RdBu_r' centered on zero if there are values of both signs. """ normfunc, mplv = use_this_norm() dmin = np.min(values) @@ -85,17 +174,24 @@ def get_difference_colors(values): return dnorm, cmap -def mask_land_or_ocean(arr,msk, use_nan=False): - """ - Apply a land or ocean mask to provided variable. - - Inputs: - arr -> the xarray variable to apply the mask to. - msk -> the xarray variable that contains the land or ocean mask, - assumed to be the same shape as "arr". - - use_nan -> Optional argument for whether to set the missing values - to np.nan values instead of the defaul "-999." values. +def mask_land_or_ocean(arr, msk, use_nan=False): + """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 : xarray.DataArray + Same as input `arr` but masked as specified. """ if use_nan: @@ -111,15 +207,26 @@ def mask_land_or_ocean(arr,msk, use_nan=False): def get_central_longitude(*args): """Determine central longitude for maps. - Can provide multiple arguments. - If any of the arguments is an instance of AdfDiag, then check whether it has a central_longitude in `diag_basic_info`. - --> This takes precedence. - --> Else: if any of the arguments are scalars in [-180, 360], assumes the FIRST ONE is the central longitude. - There are no other possible conditions, so if none of those are met, - RETURN the default value of 180. - - This allows a script to, for example, allow a config file to specify, but also have a preference: - get_central_longitude(AdfObj, 30.0) + + Allows an arbitrary number of arguments. + If any of the arguments is an instance of `AdfDiag`, then check + whether it has a `central_longitude` in `diag_basic_info`. + _This case takes precedence._ + _Else_, if any of the arguments are scalars in [-180, 360], + assumes the FIRST ONE is the central longitude. + There are no other possible conditions, so if none of those are met, + returns the default value of 180. + + Parameters + ---------- + *args : tuple + Any number of objects to check for `central_longitude`. + After that, looks for the first number between -180 and 360 in the args. + + Notes + ----- + This allows a script to, for example, allow a config file to specify, but also have a preference: + `get_central_longitude(AdfObj, 30.0)` """ chk_for_adf = [isinstance(arg, AdfDiag) for arg in args] # preference is to get value from AdfDiag: @@ -162,11 +269,20 @@ def get_central_longitude(*args): ####### def global_average(fld, wgt, verbose=False): - """ - A simple, pure numpy global average. - fld: an input ndarray - wgt: a 1-dimensional array of weights - wgt should be same size as one dimension of fld + """A simple, pure numpy global average. + + Parameters + ---------- + fld : np.ndarray + an input ndarray + wgt : np.ndarray + a 1-dimensional array of weights, should be same size as one dimension of `fld` + verbose : bool, optional + prints information when `True` + + Returns + ------- + weighted average of `fld` """ s = fld.shape @@ -184,6 +300,33 @@ def global_average(fld, wgt, verbose=False): def spatial_average(indata, weights=None, spatial_dims=None): + """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 + ------- + xr.DataArray + weighted average of `indata` + + 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. + """ import warnings if weights is None: @@ -229,14 +372,24 @@ def spatial_average(indata, weights=None, spatial_dims=None): def wgt_rmse(fld1, fld2, wgt): - """Calculated the area-weighted RMSE. - - Inputs are 2-d spatial fields, fld1 and fld2 with the same shape. - They can be xarray DataArray or numpy arrays. - - Input wgt is the weight vector, expected to be 1-d, matching length of one dimension of the data. - - Returns a single float value. + """Calculate the area-weighted RMSE. + + Parameters + ---------- + fld1, 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 ) )``` """ assert len(fld1.shape) == 2, "Input fields must have exactly two dimensions." assert fld1.shape == fld2.shape, "Input fields must have the same array shape." @@ -263,7 +416,30 @@ def wgt_rmse(fld1, fld2, wgt): # Time-weighted averaging def annual_mean(data, whole_years=False, time_name='time'): - """Calculate annual averages from time series data.""" + """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 : xr.DataArray or xr.Dataset + `data` reduced to annual averages + + 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. + """ assert time_name in data.coords, f"Did not find the expected time coordinate '{time_name}' in the data" if whole_years: first_january = np.argwhere((data.time.dt.month == 1).values)[0].item() @@ -286,17 +462,29 @@ def annual_mean(data, whole_years=False, time_name='time'): def seasonal_mean(data, season=None, is_climo=None): """Calculates the time-weighted seasonal average (or average over all time). - If season is `ANN` or None, it will default to averaging over all available time. - - If is_climo is True, expects data to have time or month dimenion of size 12. - If is_climo is False, then time must be a coordinate, - and the time.dt.days_in_month attribute must be available. - + 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 + ------- + xarray.DataArray or xarray.Dataset + the average of `data` in season `season` + + Notes + ----- If the data is a climatology, the code will make an attempt to understand the time or month - dimension, but basically will assume that it is ordered from January to December. + 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. - """ if season is not None: assert season in ["ANN", "DJF", "JJA", "MAM", "SON"], f"Unrecognized season string provided: '{season}'" @@ -341,6 +529,35 @@ def seasonal_mean(data, season=None, is_climo=None): #Polar Plot functions def domain_stats(data, domain): + """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 + -------- + spatial_average + + """ x_region = data.sel(lat=slice(domain[2],domain[3]), lon=slice(domain[0],domain[1])) x_region_mean = x_region.weighted(np.cos(np.deg2rad(x_region['lat']))).mean().item() x_region_min = x_region.min().item() @@ -351,15 +568,46 @@ def 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, obs=False, **kwargs): - ''' - Make a stereographic polar plot for the given data and hemisphere. + + """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, d2 : xr.DataArray + input data, must contain dimensions `lat` and `lon` + difference : xr.DataArray, optional + data to use as the difference, otherwise `d2 - d1` + 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). - d1, d2 -> the data to be plotted. Any tranformations/operations should be done, and dimensions should be [lat, lon] - difference -> optional, the difference between the data (d2 - d1). If not supplied, it will be derived as d2 - d1. - domain -> optional, a list of [west_lon, east_lon, south_lat, north_lat] that defines the domain to be plotted. If not provided, defaults to all longitudes, 45deg to pole of the given hemisphere - hemisphere -> must be provided as NH or SH to determine which hemisphere to plot - kwargs -> expected to be variable-dependent options for plots. - ''' + - kwargs is checked for: + + `colormap` + + `contour_levels` + + `contour_levels_range` + + `diff_contour_levels` + + `diff_contour_range` + + `diff_colormap` + + `units` + """ if difference is None: dif = d2 - d1 else: @@ -569,11 +817,48 @@ def plot_map_vect_and_save(wks, case_nickname, base_nickname, plev, umdlfld_nowrap, vmdlfld_nowrap, uobsfld_nowrap, vobsfld_nowrap, udiffld_nowrap, vdiffld_nowrap, obs=False, **kwargs): - """This plots a vector plot. bringing in two variables which respresent a vector pair - kwargs -> optional dictionary of plotting options - ** Expecting this to be a variable-specific section, - possibly provided by an ADF Variable Defaults YAML file.** + """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, vmdlfld_nowrap : xarray.DataArray + input data for case, the x- and y- components of the vectors + uobsfld_nowrap, vobsfld_nowrap : xarray.DataArray + input data for base case, the x- and y- components of the vectors + udiffld_nowrap, 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. """ # specify the central longitude for the plot: @@ -771,9 +1056,32 @@ def plot_map_and_save(wks, case_nickname, base_nickname, mdlfld, obsfld, diffld, obs=False, **kwargs): """This plots mdlfld, obsfld, diffld in a 3-row panel plot of maps. - kwargs -> optional dictionary of plotting options - ** Expecting this to be variable-specific section, possibly provided by ADF Variable Defaults YAML file.** - + 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 + 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: - colormap -> str, name of matplotlib colormap - contour_levels -> list of explict values or a tuple: (min, max, step) - diff_colormap @@ -966,24 +1274,48 @@ def plot_map_and_save(wks, case_nickname, base_nickname, # def pres_from_hybrid(psfc, hya, hyb, p0=100000.): - """ - Converts a hybrid level to a pressure - level using the forumla: - - p = a(k)*p0 + b(k)*ps - + """Calculates pressure field + + pressure derived with the formula: + ```p = a(k)*p0 + b(k)*ps``` + + Parameters + ---------- + psfc + surface pressure + hya, hyb + hybrid-sigma A and B coefficients + p0 : optional + reference pressure, defaults to 100000 Pa + + Returns + ------- + pressure, size is same as `psfc` with `len(hya)` levels """ return hya*p0 + hyb*psfc ##### def vert_remap(x_mdl, p_mdl, plev): - """ - Apply simple 1-d interpolation to a field, x - given the pressure p and the new pressures plev. - x_mdl, p_mdl are numpy arrays of shape (nlevel, spacetime). - - Andrew G.: changed to do interpolation in log pressure + """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 + ------- + output + `x_mdl` interpolated to `plev` + + Notes + ----- + Interpolation done in log pressure """ #Determine array shape of output array: @@ -1004,22 +1336,35 @@ def vert_remap(x_mdl, p_mdl, plev): def lev_to_plev(data, ps, hyam, hybm, P0=100000., new_levels=None, convert_to_mb=False): - """ - Interpolate model hybrid levels to specified pressure levels. - - new_levels-> 1-D numpy array (ndarray) containing list of pressure levels - in Pascals (Pa). - - If "new_levels" is 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 - - If "convert_to_mb" is True, then vertical (lev) dimension will have - values of mb/hPa, otherwise the units are Pa. - - The function "interp_hybrid_to_pressure" used here is dask-enabled, + """Interpolate model hybrid levels to specified pressure levels. + + Parameters + ---------- + data : + ps : + surface pressure + hyam, 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_interp_rename + data interpolated to new pressure levels + + 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. """ @@ -1060,12 +1405,23 @@ def lev_to_plev(data, ps, hyam, hybm, P0=100000., new_levels=None, ##### def pmid_to_plev(data, pmid, new_levels=None, convert_to_mb=False): - """ - Interpolate data from hybrid-sigma levels to isobaric levels. - - data : DataArray with a 'lev' coordinate - pmid : like data array but the pressure of each point (Pa) - new_levels : the output pressure levels (Pa) + """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 : xarray.DataArray + `data` interpolated onto `new_levels` """ # determine pressure levels to interpolate to: @@ -1112,14 +1468,20 @@ def zonal_mean_xr(fld): def validate_dims(fld, list_of_dims): - """Generalized function to check if specified dimensions are in a DataArray. + """Check if specified dimensions are in a DataArray. - input - fld -> DataArray with named dimensions (fld.dims) - list_of_dims -> a list of strings that specifiy the dimensions to check for + Parameters + ---------- + fld : xarray.DataArray + field to check for named dimensions + list_of_dims : list + list of strings that specifiy the dimensions to check for - return - dict with keys that are "has_{x}" where x is the name from `list_of_dims` and values that are boolean + Returns + ------- + dict + dict with keys that are "has_{x}" where x is the name from + `list_of_dims` and values that are boolean """ if not isinstance(list_of_dims, list): @@ -1128,9 +1490,21 @@ def validate_dims(fld, list_of_dims): def lat_lon_validate_dims(fld): - """ - Check if input field has the correct - dimensions needed to plot on lat/lon map. + """Check if input field has lat and lon. + + Parameters + ---------- + fld : xarray.DataArray + data with named dimensions + + Returns + ------- + bool + True if lat and lon are both dimensions, False otherwise. + + See Also + -------- + validate_dims """ # note: we can only handle variables that reduce to (lat,lon) if len(fld.dims) > 3: @@ -1143,9 +1517,21 @@ def lat_lon_validate_dims(fld): def zm_validate_dims(fld): - """ - Check if input field has the correct - dimensions needed to zonally average. + """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. """ # note: we can only handle variables that reduce to (lev, lat) or (lat,) if len(fld.dims) > 4: @@ -1239,9 +1625,28 @@ def _meridional_plot_preslon(ax, lon, lev, data, **kwargs): return img, ax def zonal_plot(lat, data, ax=None, color=None, **kwargs): - """ + """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. """ if ax is None: ax = plt.gca() @@ -1253,9 +1658,29 @@ def zonal_plot(lat, data, ax=None, color=None, **kwargs): return ax def meridional_plot(lon, data, ax=None, color=None, **kwargs): - """ + """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. """ if ax is None: ax = plt.gca() @@ -1267,7 +1692,8 @@ def meridional_plot(lon, data, ax=None, color=None, **kwargs): return ax def prep_contour_plot(adata, bdata, diffdata, **kwargs): - """ + """Preparation for making contour plots. + Prepares for making contour plots of adata, bdata, and diffdata, which is presumably the difference between adata and bdata. - set colormap from kwargs or defaults to coolwarm @@ -1277,18 +1703,27 @@ def prep_contour_plot(adata, bdata, diffdata, **kwargs): - similar settings for difference, defaults to symmetric about zero - separates Matplotlib kwargs into their own dicts - return + Parameters + ---------- + adata, bdata, diffdata + the data to be plotted + kwargs : dict, optional + plotting options + + Returns + ------- + dict a dict with the following: - 'subplots_opt': mpl kwargs for subplots - 'contourf_opt': mpl kwargs for contourf - 'colorbar_opt': mpl kwargs for 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 + - 'subplots_opt': mpl kwargs for subplots + - 'contourf_opt': mpl kwargs for contourf + - 'colorbar_opt': mpl kwargs for 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 """ # determine levels & color normalization: minval = np.min([np.min(adata), np.min(bdata)]) @@ -1392,11 +1827,14 @@ def prep_contour_plot(adata, bdata, diffdata, **kwargs): def 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): - """This is the default zonal mean plot: - adata: data to plot ([lev], lat, [lon]). - The vertical coordinate (lev) must be pressure levels. - bdata: baseline or observations to plot adata against. - It must have the same dimensions and verttical levels as adata. + + """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 @@ -1461,7 +1899,7 @@ def plot_zonal_mean_and_save(wks, case_nickname, base_nickname, levs = np.unique(np.array(cp_info['levels1'])) levs_diff = np.unique(np.array(cp_info['levelsdiff'])) - + if len(levs) < 2: img0, ax[0] = zonal_plot(adata['lat'], azm, ax=ax[0]) @@ -1485,7 +1923,7 @@ def plot_zonal_mean_and_save(wks, case_nickname, base_nickname, ax[0].set_title(case_title, loc='left', fontsize=tiFontSize) ax[1].set_title(base_title, loc='left', fontsize=tiFontSize) ax[2].set_title("$\mathbf{Test} - \mathbf{Baseline}$", loc='left', fontsize=tiFontSize) - + # style the plot: #Set Main title for subplots: @@ -1543,29 +1981,49 @@ def plot_zonal_mean_and_save(wks, case_nickname, base_nickname, def 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): - """This is the default meridional mean plot: - adata: data to plot ([lev], [lat], lon). - The vertical coordinate (lev) must be pressure levels. - bdata: baseline or observations to plot adata against. - It must have the same dimensions and vertical levels as adata. - - - 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 - - has_lev: boolean whether 'lev' is a dimension - - latbounds: the latitude bounds to average, defaults to 5S to 5N; - - if it is a number, assume symmetric about equator - - otherwise, can be a slice object. E.g., slice(-10, 20) - - if not a number and not a slice, print warning and skip plotting. - - kwargs -> optional dictionary of plotting options - ** Expecting this to be variable-specific section, - possibly provided by ADF Variable Defaults YAML file.** + """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 @@ -1651,7 +2109,7 @@ def plot_meridional_mean_and_save(wks, case_nickname, base_nickname, levs = np.unique(np.array(cp_info['levels1'])) levs_diff = np.unique(np.array(cp_info['levelsdiff'])) - + if len(levs) < 2: img0, ax[0] = pltfunc(adata[xdim], adata, ax=ax[0]) ax[0].text(0.4, 0.4, empty_message, transform=ax[0].transAxes, bbox=props) @@ -1692,7 +2150,7 @@ def plot_meridional_mean_and_save(wks, case_nickname, base_nickname, line2 = Line2D([0], [0], label=base_title, color="#ff7f0e") # #ff7f0e -> matplotlib standard orange - + fig, ax = plt.subplots(nrows=2) ax = [ax[0],ax[1]] @@ -1730,40 +2188,43 @@ def plot_meridional_mean_and_save(wks, case_nickname, base_nickname, # def square_contour_difference(fld1, fld2, **kwargs): - """Produce a figure with square axes that show fld1, fld2, - and their difference as filled contours. - - 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. - - Input: - fld1 and fld2 are 2-dimensional DataArrays with same shape - kwargs are optional keyword arguments - this function checks kwargs for `case1name`, `case2name` - - Returns: - figure object - - Assumptions: - fld1.shape == fld2.shape - len(fld1.shape) == 2 - - Annnotation: - Will try to label the cases by looking for - case1name and case2name in kwargs, - and then fld1['case'] & fld2['case'] (i.e., attributes) - If no case name, will proceed 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). - + """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, fld2 : xarray.DataArray + 2-dimensional DataArrays with same shape + **kwargs : dict, optional + optional keyword arguments + this function _only checks_ `kwargs` for `case1name`, `case2name` + + Returns + ------- + fig + figure object + + 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)`. """ + assert len(fld1.shape) == 2, "Input fields must have exactly two dimensions." # 2-Dimension => plot contourf assert fld1.shape == fld2.shape, "Input fields must have the same array shape." # Same shape => allows difference diff --git a/scripts/plotting/global_latlon_map.py b/scripts/plotting/global_latlon_map.py index 8010531cb..f077f2829 100644 --- a/scripts/plotting/global_latlon_map.py +++ b/scripts/plotting/global_latlon_map.py @@ -1,3 +1,17 @@ +""" +Generate global maps of 2-D fields + +Functions +--------- +global_latlon_map(adfobj) + use ADF object to make maps +my_formatwarning(msg, *args, **kwargs) + format warning messages + (private method) +_load_dataset(fils) + load files into dataset + (private methd) +""" #Import standard modules: from pathlib import Path import numpy as np @@ -6,7 +20,7 @@ #Format warning messages: def my_formatwarning(msg, *args, **kwargs): - # ignore everything except the message + """Issue `msg` as warning.""" return str(msg) + '\n' warnings.formatwarning = my_formatwarning @@ -14,25 +28,61 @@ def global_latlon_map(adfobj): """ This script/function is designed to generate global 2-D lat/lon maps of model fields with continental overlays. - Description of needed inputs: - case_name -> Name of CAM case provided by "cam_case_name". - model_rgrid_loc -> Location of re-gridded CAM climo files provided by "cam_regrid_loc". - data_name -> Name of data set CAM case is being compared against, - which is always either "obs" or the baseline CAM case name, - depending on whether "compare_obs" is true or false. - data_loc -> Location of comparison data, which is either the location listed - in each variable's ""obs_file", or the same as "model_rgrid_loc", - depending on whether "compare_obs" is true or false. - var_list -> List of CAM output variables provided by "diag_var_list" - data_list -> List of data sets CAM will be compared against, which - is simply the baseline case name in situations when - "compare_obs" is false. - plot_location -> Location where plot files will be written to, which is - specified by "cam_diag_plot_loc". - climo_yrs -> Dictionary containing the start and end years of the test - and baseline model data (if applicable). - variable_defaults -> optional, - Dict that has keys that are variable names and values that are plotting preferences/defaults. + + Parameters + ---------- + adfobj : AdfDiag + The diagnostics object that contains all the configuration information + + Returns + ------- + Does not return a value; produces plots and saves files. + + Notes + ----- + This function imports `pandas` and `plotting_functions` + + It uses the AdfDiag object's methods to get necessary information. + Specificially: + adfobj.diag_var_list + List of variables + adfobj.get_basic_info + Regrid data path, checks `compare_obs`, checks `redo_plot`, checks `plot_press_levels` + adfobj.plot_location + output plot path + adfobj.get_cam_info + Get `cam_case_name` and `case_nickname` + adfobj.climo_yrs + start and end climo years of the case(s), `syears` & `eyears` + start and end climo years of the reference, `syear_baseline` & `eyear_baseline` + adfobj.var_obs_dict + reference data (conditional) + adfobj.get_baseline_info + get reference case, `cam_case_name` + adfobj.variable_defaults + dict of variable-specific plot preferences + adfobj.read_config_var + dict of basic info, `diag_basic_info` + Then use to check `plot_type` + adfobj.compare_obs + Used to set data path + adfobj.debug_log + Issues debug message + adfobj.add_website_data + Communicates information to the website generator + + + The `plotting_functions` module is needed for: + pf.get_central_longitude() + determine central longitude for global plots + pf.lat_lon_validate_dims() + makes sure latitude and longitude are valid + pf.seasonal_mean() + calculate seasonal mean + pf.plot_map_and_save() + send information to make the plot and save the file + pf.zm_validate_dims() + Checks on pressure level dimension """ #Import necessary modules: @@ -433,6 +483,22 @@ def global_latlon_map(adfobj): ######### def _load_dataset(fils): + """ + loads datasets from input file information + + Parameters + ---------- + fils : list + strings or paths to input file(s) + + Returns + ------- + xr.Dataset + + Notes + ----- + When just one entry is provided, use `open_dataset`, otherwise `open_mfdatset` + """ if len(fils) == 0: warnings.warn(f"Input file list is empty.") return None