<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Create 3D boolean masks
In this tutorial we will show how to create 3D boolean masks for arbitrary latitude and longitude grids. It uses the same algorithm to determine if a gridpoint is in a region as for the 2D mask. However, it returns a xarray.Dataset with shape region x lat x lon, gridpoints that do not fall in a region are False, the gridpoints that fall in a region are True.
3D masks are convenient as they can be used to directly calculate weighted regional means (over all regions) using xarray v0.15.1 or later. Further, the mask includes the region names and abbreviations as non-dimension coordinates.
Import regionmask and check the version
Step1:  Load xarray and numpy
Step2:  Creating a mask
Define a lon/ lat grid with a 1° grid spacing, where the points define the center of the grid
Step3:  We will create a mask with the SREX regions (Seneviratne et al., 2012).
Step4:  The function mask_3D determines which gripoints lie within the polygon making up each region
Step5:  As mentioned, mask is a boolean xarray.Dataset with shape region x lat x lon. It contains region (=numbers) as dimension coordinate as well as abbrevs and names as non-dimension coordinates (see the xarray docs for the details on the terminology).
Plotting
Plotting individual layers
The four first layers look as follows
Step6:  Plotting flattened masks
A 3D mask cannot be directly plotted - it needs to be flattened first. To do this regionmask offers a convenience function
Step7:  Working with a 3D mask
masks can be used to select data in a certain region and to calculate regional averages - let's illustrate this with a 'real' dataset
Step8:  The example data is a temperature field over North America. Let's plot the first time step
Step9:  An xarray object can be passed to the mask_3D function
Step10:  Per default this creates a mask containing one layer (slice) for each region containing (at least) one gridpoint. As the example data only has values over Northern America we only get only 6 layers even though there are 26 SREX regions. To obtain all layers specify drop=False
Step11:  Note mask_full now has 26 layers.
Select a region
As mask_3D contains region, abbrevs, and names as (non-dimension) coordinates we can use each of those to select an individual region
Step12:  This also applies to the regionally-averaged data below. 
It is currently not possible to use sel with a non-dimension coordinate - to directly select abbrev or name you need to create a MultiIndex
Step13:  Mask out a region
Using where a specific region can be 'masked out' (i.e. all data points outside of the region become NaN)
Step14:  Which looks as follows
Step15:  We could now use airtemps_cna to calculate the regional average for 'Central North America'. However, there is a more elegant way.
Calculate weighted regional averages
Using the 3-dimensional mask it is possible to calculate weighted averages of all regions in one go, using the weighted method (requires xarray 0.15.1 or later). As proxy of the grid cell area we use cos(lat).
Step16:  Let's break down what happens here. By multiplying mask_3D * weights we get a DataArray where gridpoints not in the region get a weight of 0. Gridpoints within a region get a weight proportional to the gridcell area. airtemps.weighted(mask_3D * weights) creates an xarray object which can be used for weighted operations. From this we calculate the weighted mean over the lat and lon dimensions. The resulting dataarray has the dimensions region x time
Step17:  The regionally-averaged time series can be plotted
Step18:  Restrict the mask to land points
Combining the mask of the regions with a land-sea mask we can create a land-only mask using the land_110 region from NaturalEarth.
With this caveat in mind we can create the land-sea mask
Step19:  and plot it
Step20:  To create the combined mask we multiply the two
Step21:  Note the .squeeze(drop=True). This is required to remove the region dimension from land_mask.
Finally, we compare the original mask with the one restricted to land points
<ASSISTANT_TASK:>
Python Code:
import regionmask
regionmask.__version__
Explanation: Create 3D boolean masks
In this tutorial we will show how to create 3D boolean masks for arbitrary latitude and longitude grids. It uses the same algorithm to determine if a gridpoint is in a region as for the 2D mask. However, it returns a xarray.Dataset with shape region x lat x lon, gridpoints that do not fall in a region are False, the gridpoints that fall in a region are True.
3D masks are convenient as they can be used to directly calculate weighted regional means (over all regions) using xarray v0.15.1 or later. Further, the mask includes the region names and abbreviations as non-dimension coordinates.
Import regionmask and check the version:
End of explanation
import xarray as xr
import numpy as np
# don't expand data
xr.set_options(display_style="text", display_expand_data=False)
Explanation: Load xarray and numpy:
End of explanation
lon = np.arange(-179.5, 180)
lat = np.arange(-89.5, 90)
Explanation: Creating a mask
Define a lon/ lat grid with a 1° grid spacing, where the points define the center of the grid:
End of explanation
regionmask.defined_regions.srex
Explanation: We will create a mask with the SREX regions (Seneviratne et al., 2012).
End of explanation
mask = regionmask.defined_regions.srex.mask_3D(lon, lat)
mask
Explanation: The function mask_3D determines which gripoints lie within the polygon making up each region:
End of explanation
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
from matplotlib import colors as mplc
cmap1 = mplc.ListedColormap(["none", "#9ecae1"])
fg = mask.isel(region=slice(4)).plot(
    subplot_kws=dict(projection=ccrs.PlateCarree()),
    col="region",
    col_wrap=2,
    transform=ccrs.PlateCarree(),
    add_colorbar=False,
    aspect=1.5,
    cmap=cmap1,
)
for ax in fg.axes.flatten():
    ax.coastlines()
fg.fig.subplots_adjust(hspace=0, wspace=0.1);
Explanation: As mentioned, mask is a boolean xarray.Dataset with shape region x lat x lon. It contains region (=numbers) as dimension coordinate as well as abbrevs and names as non-dimension coordinates (see the xarray docs for the details on the terminology).
Plotting
Plotting individual layers
The four first layers look as follows:
End of explanation
regionmask.plot_3D_mask(mask, add_colorbar=False, cmap="plasma");
Explanation: Plotting flattened masks
A 3D mask cannot be directly plotted - it needs to be flattened first. To do this regionmask offers a convenience function: regionmask.plot_3D_mask. The function takes a 3D mask as argument, all other keyword arguments are passed through to xr.plot.pcolormesh.
End of explanation
airtemps = xr.tutorial.load_dataset("air_temperature")
Explanation: Working with a 3D mask
masks can be used to select data in a certain region and to calculate regional averages - let's illustrate this with a 'real' dataset:
End of explanation
# choose a good projection for regional maps
proj = ccrs.LambertConformal(central_longitude=-100)
ax = plt.subplot(111, projection=proj)
airtemps.isel(time=1).air.plot.pcolormesh(ax=ax, transform=ccrs.PlateCarree())
ax.coastlines();
Explanation: The example data is a temperature field over North America. Let's plot the first time step:
End of explanation
mask_3D = regionmask.defined_regions.srex.mask_3D(airtemps)
mask_3D
Explanation: An xarray object can be passed to the mask_3D function:
End of explanation
mask_full = regionmask.defined_regions.srex.mask_3D(airtemps, drop=False)
mask_full
Explanation: Per default this creates a mask containing one layer (slice) for each region containing (at least) one gridpoint. As the example data only has values over Northern America we only get only 6 layers even though there are 26 SREX regions. To obtain all layers specify drop=False:
End of explanation
# 1) by the index of the region:
r1 = mask_3D.sel(region=3)
# 2) with the abbreviation
r2 = mask_3D.isel(region=(mask_3D.abbrevs == "WNA"))
# 3) with the long name:
r3 = mask_3D.isel(region=(mask_3D.names == "E. North America"))
Explanation: Note mask_full now has 26 layers.
Select a region
As mask_3D contains region, abbrevs, and names as (non-dimension) coordinates we can use each of those to select an individual region:
End of explanation
mask_3D.set_index(regions=["region", "abbrevs", "names"]);
Explanation: This also applies to the regionally-averaged data below. 
It is currently not possible to use sel with a non-dimension coordinate - to directly select abbrev or name you need to create a MultiIndex:
End of explanation
airtemps_cna = airtemps.where(r1)
Explanation: Mask out a region
Using where a specific region can be 'masked out' (i.e. all data points outside of the region become NaN):
End of explanation
proj = ccrs.LambertConformal(central_longitude=-100)
ax = plt.subplot(111, projection=proj)
airtemps_cna.isel(time=1).air.plot(ax=ax, transform=ccrs.PlateCarree())
ax.coastlines();
Explanation: Which looks as follows:
End of explanation
weights = np.cos(np.deg2rad(airtemps.lat))
ts_airtemps_regional = airtemps.weighted(mask_3D * weights).mean(dim=("lat", "lon"))
Explanation: We could now use airtemps_cna to calculate the regional average for 'Central North America'. However, there is a more elegant way.
Calculate weighted regional averages
Using the 3-dimensional mask it is possible to calculate weighted averages of all regions in one go, using the weighted method (requires xarray 0.15.1 or later). As proxy of the grid cell area we use cos(lat).
End of explanation
ts_airtemps_regional
Explanation: Let's break down what happens here. By multiplying mask_3D * weights we get a DataArray where gridpoints not in the region get a weight of 0. Gridpoints within a region get a weight proportional to the gridcell area. airtemps.weighted(mask_3D * weights) creates an xarray object which can be used for weighted operations. From this we calculate the weighted mean over the lat and lon dimensions. The resulting dataarray has the dimensions region x time:
End of explanation
ts_airtemps_regional.air.plot(col="region", col_wrap=3);
Explanation: The regionally-averaged time series can be plotted:
End of explanation
land_110 = regionmask.defined_regions.natural_earth_v5_0_0.land_110
land_mask = land_110.mask_3D(airtemps)
Explanation: Restrict the mask to land points
Combining the mask of the regions with a land-sea mask we can create a land-only mask using the land_110 region from NaturalEarth.
With this caveat in mind we can create the land-sea mask:
End of explanation
proj = ccrs.LambertConformal(central_longitude=-100)
ax = plt.subplot(111, projection=proj)
land_mask.squeeze().plot.pcolormesh(
    ax=ax, transform=ccrs.PlateCarree(), cmap=cmap1, add_colorbar=False
)
ax.coastlines();
Explanation: and plot it
End of explanation
mask_lsm = mask_3D * land_mask.squeeze(drop=True)
Explanation: To create the combined mask we multiply the two:
End of explanation
f, axes = plt.subplots(1, 2, subplot_kw=dict(projection=proj))
ax = axes[0]
mask_3D.sel(region=2).plot(
    ax=ax, transform=ccrs.PlateCarree(), add_colorbar=False, cmap=cmap1
)
ax.coastlines()
ax.set_title("Regional mask: all points")
ax = axes[1]
mask_lsm.sel(region=2).plot(
    ax=ax, transform=ccrs.PlateCarree(), add_colorbar=False, cmap=cmap1
)
ax.coastlines()
ax.set_title("Regional mask: land only");
Explanation: Note the .squeeze(drop=True). This is required to remove the region dimension from land_mask.
Finally, we compare the original mask with the one restricted to land points:
End of explanation
<END_TASK>