data-carpentry2020/plot_pr_data.py at master · robwarrenwx/data-carpentry2020 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import argparse
import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
import cartopy.crs as ccrs
import cmocean

def convert_pr_units(clim):
    """Convert units of precipitation climatology

    Args:
      clim (xarray dat array): Precipitation climatology (units of kg/m2/s)

    """

    clim.data *= 86400
    clim.attrs['units'] = 'mm/day'
    return clim

def create_plot(clim, model, season, gridlines=False, levels=None):
    """Create plot of precipitation climatology

    Args:
      clim (data array): Precipitation climatology
      model (str): Model ID
      season (str): Season (DJF, MAM, JJA, SON)
      gridlines (bool): Select whether to plot gridlines
      levels (list): Tick marks on the colorbar
    """

    if not levels:
        levels = np.arange(0, 13.5, 1.5)

    fig = plt.figure(figsize=[12, 5])
    ax = fig.add_subplot(111, projection=ccrs.PlateCarree(central_longitude=0))
    clim.sel(season=season).plot.contourf(
        ax = ax,
        levels = levels,
        extend = 'max',
        transform = ccrs.PlateCarree(),
        cbar_kwargs = {'label':clim.units},
        cmap = cmocean.cm.haline_r
    )
    ax.coastlines()
    if gridlines:
        plt.gca().gridlines()
    title = '%s precipitation climatology (%s)' %(model, season)
    plt.title(title)

def main(inargs):
    dset = xr.open_dataset(inargs.infile)

    clim = dset['pr'].groupby('time.season').mean('time', keep_attrs=True)
    clim = convert_pr_units(clim)

    create_plot(clim, dset.attrs['model_id'], inargs.season,
                gridlines=inargs.season, levels=inargs.levels)
    plt.savefig(inargs.outfile)

if __name__ == '__main__':
    description = "Plot the precipitation climatology for a given season"

    parser = argparse.ArgumentParser(description=description)

    parser.add_argument('infile', type=str, help='Input precipitation data file')
    parser.add_argument('season', type=str,
                        choices=['DJF','MAM','JJA','SON'],
                        help='Season to plot')

    parser.add_argument('outfile', type=str, help='Output plot file')

    parser.add_argument("--gridlines", action="store_true", default=False,
                        help="Include gridlines on the plot")

    parser.add_argument("--levels", type=float, nargs='*', default=None,
                        help='list of color levels')

    args = parser.parse_args()

    main(args)