modified the wrong exercise

exercises/03_naming/precipitation_climatology.py

@@ -13,21 +13,18 @@
 from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
 
 
-def convert_precipitation_units(precipitation_in_kg_per_m_squared_s):
+def convert_pr_units(darray):
     """Convert kg m-2 s-1 to mm day-1."""
 
-    precipitation_in_mm_per_day = xr.DataArray(
-        precipitation_in_kg_per_m_squared_s * 86400
-    )
-
-    precipitation_in_mm_per_day.attrs["units"] = "mm/day"
+    darray.data = darray.data * 86400
+    darray.attrs["units"] = "mm/day"
 
-    if precipitation_in_mm_per_day.data.min() < 0.0:
-        raise ValueError("There is at least one negative precipitation value")
-    if precipitation_in_mm_per_day.data.max() > 2000:
-        raise ValueError("There is a precipitation value/s > 2000 mm/day")
+    assert (
+        darray.data.min() >= 0.0
+    ), "There is at least one negative precipitation value"
+    assert darray.data.max() < 2000, "There is a precipitation value/s > 2000 mm/day"
 
-    return precipitation_in_mm_per_day
+    return darray
 
 
 def apply_mask(darray, sftlf_file, realm):
@@ -41,52 +38,44 @@ def apply_mask(darray, sftlf_file, realm):
     pass
 
 
-def plot_zonally_averaged_precipitation(precipitation_data):
+def plot_zonal(data):
     # print(data)
-    zonal_precipitation = precipitation_data["pr"].mean("lon", keep_attrs=True)
+    zonal_pr = data["pr"].mean("lon", keep_attrs=True)
 
-    figure, axes = plt.subplots(nrows=4, ncols=1, figsize=(12, 8))
+    fig, ax = plt.subplots(nrows=4, ncols=1, figsize=(12, 8))
 
-    zonal_precipitation.sel(lat=[0]).plot.line(ax=axes[0], hue="lat")
-    zonal_precipitation.sel(lat=[-20, 20]).plot.line(ax=axes[1], hue="lat")
-    zonal_precipitation.sel(lat=[-45, 45]).plot.line(ax=axes[2], hue="lat")
-    zonal_precipitation.sel(lat=[-70, 70]).plot.line(ax=axes[3], hue="lat")
+    zonal_pr.sel(lat=[0]).plot.line(ax=ax[0], hue="lat")
+    zonal_pr.sel(lat=[-20, 20]).plot.line(ax=ax[1], hue="lat")
+    zonal_pr.sel(lat=[-45, 45]).plot.line(ax=ax[2], hue="lat")
+    zonal_pr.sel(lat=[-70, 70]).plot.line(ax=ax[3], hue="lat")
 
     plt.tight_layout()
-    for axis in axes:
+    for axis in ax:
         axis.set_ylim(0.0, 1.0e-4)
         axis.grid()
     plt.savefig("zonal.png", dpi=200)  # Save figure to file
 
-    figure, axes = plt.subplots(nrows=1, ncols=1, figsize=(12, 5))
+    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 5))
 
-    zonal_precipitation.T.plot()
+    zonal_pr.T.plot()
 
     plt.savefig("zonal_map.png", dpi=200)  # Save figure to file
 
 
-def get_country_annual_average(data, countries):
-    annual_average_precipitation = (
-        data["pr"].groupby("time.year").mean("time", keep_attrs=True)
-    )
-    annual_average_precipitation = convert_precipitation_units(
-        annual_average_precipitation
-    )
+def get_country_ann_avg(data, countries):
+    data_avg = data["pr"].groupby("time.year").mean("time", keep_attrs=True)
+    data_avg = convert_pr_units(data_avg)
 
-    country_mask = regionmask.defined_regions.natural_earth_v5_0_0.countries_110.mask(
-        annual_average_precipitation
-    )
+    land = regionmask.defined_regions.natural_earth_v5_0_0.countries_110.mask(data_avg)
 
     # List possible locations to plot
     # [print(k, v) for k, v in regionmask.defined_regions.natural_earth_v5_0_0.countries_110.regions.items()]
 
     with open("data.txt", "w") as datafile:
-        for country_name, country_code in countries.items():
+        for k, v in countries.items():
             # land.plot(ax=geo_axes, add_label=False, fc="white", lw=2, alpha=0.5)
             # clim = clim.where(ocean == "South Pacific Ocean")
-            country_annual_average_precipitation = annual_average_precipitation.where(
-                country_mask.cf == country_code
-            )
+            data_avg_mask = data_avg.where(land.cf == v)
 
             # Debugging - plot countries to make sure mask works correctly
             # fig, geo_axes = plt.subplots(nrows=1, ncols=1, figsize=(12,5),
@@ -110,21 +99,17 @@ def get_country_annual_average(data, countries):
             # print("show %s" %k)
             # plt.show()
 
-            for year in country_annual_average_precipitation.year.values:
-                precipitation = (
-                    country_annual_average_precipitation.sel(year=year).mean().values
-                )
+            for yr in data_avg_mask.year.values:
+                precip = data_avg_mask.sel(year=yr).mean().values
                 datafile.write(
-                    "{} {} : {:2.3f} mm/day\n".format(
-                        country_name.ljust(25), year, precipitation
-                    )
+                    "{} {} : {:2.3f} mm/day\n".format(k.ljust(25), yr, precip)
                 )
             datafile.write("\n")
 
 
-def plot_enso_hovmoller_diagram(precipitation_data):
+def plot_enso(data):
     enso = (
-        precipitation_data["pr"]
+        data["pr"]
         .sel(lat=slice(-1, 1))
         .sel(lon=slice(120, 280))
         .mean(dim="lat", keep_attrs=True)
@@ -145,22 +130,14 @@ def plot_enso_hovmoller_diagram(precipitation_data):
     plt.savefig("enso.png", dpi=200)  # Save figure to file
 
 
-def create_precipitation_climatology_plot(
-    seasonal_average_precipitation,
-    model_name,
-    season,
-    mask=None,
-    plot_gridlines=False,
-    levels=None,
-):
+def create_plot(clim, model, season, mask=None, gridlines=False, levels=None):
     """Plot the precipitation climatology.
 
-    seasonal_average_precipitation : xarray.DataArray
-        Precipitation climatology data. Seasonally averaged precipitation data.
-    model_name : str
+    clim (xarray.DataArray): Precipitation climatology data
+    model (str): Name of the climate model
     season (str): Season
 
-    plot_gridlines (bool): Select whether to plot gridlines
+    gridlines (bool): Select whether to plot gridlines
     levels (list): Tick marks on the colorbar
 
     """
@@ -179,12 +156,12 @@ def create_precipitation_climatology_plot(
         subplot_kw={"projection": ccrs.PlateCarree(central_longitude=180)},
     )
 
-    seasonal_average_precipitation.sel(season=season).plot.contourf(
+    clim.sel(season=season).plot.contourf(
         ax=geo_axes,
         levels=levels,
         extend="max",
         transform=ccrs.PlateCarree(),
-        cbar_kwargs={"label": seasonal_average_precipitation.units},
+        cbar_kwargs={"label": clim.units},
         cmap=cmocean.cm.rain,
     )
 
@@ -216,76 +193,72 @@ def create_precipitation_climatology_plot(
             alpha=0.75,
         )
 
-    if plot_gridlines:
+    if gridlines:
         # If we want gridlines run the code to do this:
-        gridlines = geo_axes.gridlines(
+        gl = geo_axes.gridlines(
             crs=ccrs.PlateCarree(),
             draw_labels=True,
             linewidth=2,
             color="gray",
             alpha=0.5,
             linestyle="--",
         )
-        gridlines.top_labels = False
-        gridlines.left_labels = True
+        gl.top_labels = False
+        gl.left_labels = True
         # gl.xlines = False
-        gridlines.xlocator = mticker.FixedLocator([-180, -90, 0, 90, 180])
-        gridlines.ylocator = mticker.FixedLocator(
+        gl.xlocator = mticker.FixedLocator([-180, -90, 0, 90, 180])
+        gl.ylocator = mticker.FixedLocator(
             [-66, -23, 0, 23, 66]
         )  # Tropics & Polar Circles
-        gridlines.xformatter = LONGITUDE_FORMATTER
-        gridlines.yformatter = LATITUDE_FORMATTER
-        gridlines.xlabel_style = {"size": 15, "color": "gray"}
-        gridlines.ylabel_style = {"size": 15, "color": "gray"}
+        gl.xformatter = LONGITUDE_FORMATTER
+        gl.yformatter = LATITUDE_FORMATTER
+        gl.xlabel_style = {"size": 15, "color": "gray"}
+        gl.ylabel_style = {"size": 15, "color": "gray"}
 
-    title = "{} precipitation climatology ({})".format(model_name, season)
+    title = "{} precipitation climatology ({})".format(model, season)
     plt.title(title)
     # print("\n\n{}\n\n".format(clim.mean()))
 
 
 def main(
-    precipitation_netcdf_file,
+    pr_file,
     season="DJF",
     output_file="output.png",
-    plot_gridlines=False,
+    gridlines=False,
     mask=None,
     cbar_levels=None,
     countries={"United Kingdom": "GB"},
 ):
     """Run the program."""
 
-    precipitation_data = xr.open_dataset(precipitation_netcdf_file)
+    dset = xr.open_dataset(pr_file)
 
-    plot_zonally_averaged_precipitation(precipitation_data)
-    plot_enso_hovmoller_diagram(precipitation_data)
-    get_country_annual_average(precipitation_data, countries)
+    plot_zonal(dset)
+    plot_enso(dset)
+    get_country_ann_avg(dset, countries)
 
-    seasonal_average_precipitation = (
-        precipitation_data["pr"].groupby("time.season").mean("time", keep_attrs=True)
-    )
+    clim = dset["pr"].groupby("time.season").mean("time", keep_attrs=True)
 
     try:
-        input_units = seasonal_average_precipitation.attrs["units"]
+        input_units = clim.attrs["units"]
     except KeyError:
         raise KeyError(
             "Precipitation variable in {pr_file} must have a units attribute"
         )
 
     if input_units == "kg m-2 s-1":
-        seasonal_average_precipitation = convert_precipitation_units(
-            seasonal_average_precipitation
-        )
+        clim = convert_pr_units(clim)
     elif input_units == "mm/day":
         pass
     else:
         raise ValueError("""Input units are not 'kg m-2 s-1' or 'mm/day'""")
 
-    create_precipitation_climatology_plot(
-        seasonal_average_precipitation,
-        precipitation_data.attrs["source_id"],
+    create_plot(
+        clim,
+        dset.attrs["source_id"],
         season,
         mask=mask,
-        plot_gridlines=plot_gridlines,
+        gridlines=gridlines,
         levels=cbar_levels,
     )
 
@@ -315,6 +288,6 @@ def main(
         input_file,
         season=season_to_plot,
         mask=mask_id,
-        plot_gridlines=gridlines_on,
+        gridlines=gridlines_on,
         countries=countries,
     )