气象编程 | PyAOS基础教程三：自定义函数
1、概要
文章解答以下疑问：
第一：如何自定义函数。
文章的目标
2、内容回顾
我们先来回顾上节完整绘图代码：
1 ACCESS-CM2 historical precipitation climatology using the following commands:
2
3import xarray as xr
4import cartopy.crs as ccrs
5import matplotlib.pyplot as plt
6import numpy as np
7import cmocean
8
9accesscm2_pr_file = 'data/pr_Amon_ACCESS-CM2_historical_r1i1p1f1_gn_201001-201412.nc'
10
11dset = xr.open_dataset(accesscm2_pr_file)
12
13clim = dset['pr'].groupby('time.season').mean('time', keep_attrs=True)
14
15clim.data = clim.data * 86400
16clim.attrs['units'] = 'mm/day'
17
18fig = plt.figure(figsize=[12,5])
19ax = fig.add_subplot(111, projection=ccrs.PlateCarree(central_longitude=180))
20clim.sel(season='JJA').plot.contourf(ax=ax,
21 levels=np.arange(0, 13.5, 1.5),
22 extend='max',
23 transform=ccrs.PlateCarree(),
24 cbar_kwargs={'label': clim.units},
25 cmap=cmocean.cm.haline_r)
26ax.coastlines()
27
28model = dset.attrs['source_id']
29title = f'{model} precipitation climatology (JJA)'
30plt.title(title)
31
32plt.show()
3、自定义函数
如果需要为其他不同的模式绘制一个相似的图片，那么我们需要复制、粘贴这段代码，然后做相应的修改就可以了，但是这种方式却可能提升出错的概率，比如数据替换成其他模式了，同时我们也要调整季节JJA为DJF，而忘记了调整title，这样就出现标题与数据内容不匹配的问题。
另外，复制、粘贴代码的方法可能更加消耗时间，未来，如果我们找到了更好的方法来绘图，比如用函数：plt.gca().gridlines()打开网格线，我们就需要对所有的代码段进行调整，这种方法太累了，因此我们推荐使用自定义函数来实现这个功能。
1def plot_pr_climatology(pr_file, season, gridlines=False): 2 '''Plot the precipitation climatology. 3 4 Args: 5 pr_file (str): Precipitation data file 6 season (str): Season (3 letter abbreviation, e.g. JJA) 7 gridlines (bool): Select whether to plot gridlines 8 9 '''1011 dset = xr.open_dataset(pr_file)1213 clim = dset['pr'].groupby('time.season').mean('time', keep_attrs=True)1415 clim.data = clim.data * 8640016 clim.attrs['units'] = 'mm/day'1718 fig = plt.figure(figsize=[12,5])19 ax = fig.add_subplot(111, projection=ccrs.PlateCarree(central_longitude=180))20 clim.sel(season=season).plot.contourf(ax=ax,21 levels=np.arange(0, 13.5, 1.5),22 extend='max',23 transform=ccrs.PlateCarree(),24 cbar_kwargs={'label': clim.units},25 cmap=cmocean.cm.haline_r)26 ax.coastlines()27 if gridlines:28 plt.gca().gridlines()2930 model = dset.attrs['source_id']31 title = f'{model} precipitation climatology ({season})'32 plt.title(title)然后，我们使用帮助命令help，我们就可以看到对于当前函数实现的功能，以及几个参数的解释。
1help(plot_pr_climatology)
2
3Help on function plot_pr_climatology in module __main__:
4
5plot_pr_climatology(pr_file, season, gridlines=False)
6 Plot the precipitation climatology.
7
8 Args:
9 pr_file (str): Precipitation data file
10 season (str): Season (3 letter abbreviation, e.g. JJA)
11 gridlines (bool): Select whether to plot gridlines然后，我们就可以用这个自定义函数来绘图了：
1plot_pr_climatology('data/pr_Amon_ACCESS-CM2_historical_r1i1p1f1_gn_201001-201412.nc', 'JJA')2plt.show()如果要绘制DJF的图片，只需要修改参数就可以了：
1plot_pr_climatology('data/pr_Amon_ACCESS-ESM1-5_historical_r1i1p1f1_gn_201001-201412.nc', 'DJF')
2plt.show()
gridlines这个参数模式是关闭的，我们也可以打开：
1plot_pr_climatology('data/pr_Amon_ACCESS-ESM1-5_historical_r1i1p1f1_gn_201001-201412.nc',2 'DJF', gridlines=True)3plt.show()4、代码改进
上面的代码超过16行，稍微有点长了。一般来说，人对于7-12行的代码理解起来比较容易，而且理解速度也比较快点，而且越精简的代码，其可靠性也越高。人们对于代码理解的速度多数依赖于代码量，代码量越大，理解速度越慢，代码量越小，理解速度越快，而且帮助文档写得好，也能帮助快速理解。因此我们把上面那段代码进一步细分，提取出功能更加独立、代码量更少的自定义函数。
1def plot_pr_climatology(pr_file, season, gridlines=False):
2 '''Plot the precipitation climatology.
3
4 Args:
5 pr_file (str): Precipitation data file
6 season (str): Season (3 letter abbreviation, e.g. JJA)
7 gridlines (bool): Select whether to plot gridlines
8
9 '''
10
11 dset = xr.open_dataset(pr_file)
12 clim = dset['pr'].groupby('time.season').mean('time', keep_attrs=True)
13 clim = convert_pr_units(clim)
14 create_plot(clim, dset.attrs['source_id'], season, gridlines=gridlines)
15 plt.show()这也就意味着，我们需要定义新的函数：convert_pr_units 和create_plot，
1def convert_pr_units(darray): 2 '''Convert kg m-2 s-1 to mm day-1. 3 4 Args: 5 darray (xarray.DataArray): Precipitation data 6 7 ''' 8 9 darray.data = darray.data * 8640010 darray.attrs['units'] = 'mm/day'1112 return darray131415def create_plot(clim, model, season, gridlines=False):16 '''Plot the precipitation climatology.1718 Args:19 clim (xarray.DataArray): Precipitation climatology data20 model (str) : Name of the climate model21 season (str): Season 2223 Kwargs: 24 gridlines (bool): Select whether to plot gridlines 2526 '''2728 fig = plt.figure(figsize=[12,5])29 ax = fig.add_subplot(111, projection=ccrs.PlateCarree(central_longitude=180))30 clim.sel(season=season).plot.contourf(ax=ax,31 levels=np.arange(0, 13.5, 1.5),32 extend='max',33 transform=ccrs.PlateCarree(),34 cbar_kwargs={'label': clim.units},35 cmap=cmocean.cm.haline_r)36 ax.coastlines()37 if gridlines:38 plt.gca().gridlines()3940 title = f'{model} precipitation climatology ({season})'41 plt.title(title)424344def plot_pr_climatology(pr_file, season, gridlines=False):45 '''Plot the precipitation climatology.4647 Args:48 pr_file (str): Precipitation data file49 season (str): Season (3 letter abbreviation, e.g. JJA)50 gridlines (bool): Select whether to plot gridlines5152 '''5354 dset = xr.open_dataset(pr_file)55 clim = dset['pr'].groupby('time.season').mean('time', keep_attrs=True)56 clim = convert_pr_units(clim)57 create_plot(clim, dset.attrs['source_id'], season, gridlines=gridlines)58 plt.show()5、自定义模块
对于我们经常用的自定义函数，比如上面的函数：convert_pr_units，如何避免在不同的项目或者python编辑器中反复复制粘贴呢？答案就是把这块代码放到独立的py文件中去，比如命名一个叫unit_conversion.py的python文件，把convert_pr_units函数拷贝其中，然后下次用的时候，我们只要导入unit_conversion模块就可以了。
1import unit_conversion
2clim.data = unit_conversion.convert_pr_units(clim.data)6、总结
本文的主要知识点：
第一：使用def自定义函数；
第二：函数体需要整体缩进；
第三：函数的调用；
第四：help命令的使用；
第五：将文档字符串放到函数中，提供帮助命令；
第六：缺省参数的默认值和非默认值；
第七：短小且功能单一的函数可靠性较高；
第八：写自定义模块，与如何使用自定义模块
赞 (0)