How to remove contour colours outside continental U.S. (within Mexico and Canada) from basemap?

Question

I have a data to visualize on U.S. map. I don't want any color on lands outside the United States (within Mexico in the southwest and Canada in the northeast). How can I mask those regions from contourf? Note that it's not important whether those state boundaries are drawn or not.

The code is:

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, cm, maskoceans
from scipy.interpolate import griddata
np.random.seed(77)



lllat = 24.396308
lllon = -124.848974
urlat =  49.384358
urlon = -66.885444
m = Basemap(llcrnrlat=lllat,
urcrnrlat=urlat,
llcrnrlon=lllon,
urcrnrlon=urlon,
resolution='i', projection='cyl')
m.drawcountries(linewidth=1.0)
m.drawstates(linewidth=1.0, color='lightgray')
m.drawlsmask(land_color='gray',ocean_color="#b0c4de", lakes=True)

#create 100 random latitudes
lats = np.random.randint(low=lllat-1, high=urlat+1, size=1000) + np.random.ranf(size=1000)
#create 100 random longitudes
lons = np.random.randint(low=lllon-1, high=urlon+1, size=1000) + np.random.ranf(size=1000)
#create 100 random values/probabilities
probabilities = np.random.random(size=1000)


#now use meshgrid and contourf to visualize it
mlon, mlat = m(*(lons, lats))
# grid data
numcols, numrows = 1000, 1000
xi = np.linspace(mlon.min(), mlon.max(), numcols)
yi = np.linspace(mlat.min(), mlat.max(), numrows)

xi, yi = np.meshgrid(xi, yi)
# interpolate
x, y, z = mlon, mlat, probabilities
zi = griddata((mlon, mlat), probabilities, (xi, yi), method='nearest', rescale=False)
data = maskoceans(xi, yi, zi)
con = m.contourf(xi, yi, data, cmap=plt.get_cmap('YlOrRd'))
cbar = m.colorbar(con,location='right',pad="3%")
plt.show()

which produces the following image Note the colors outside U.S. boundaries, I want them to be removed.

I tried removing points from lats and lons which don't exist in U.S. but still contourf colors some parts outside the boundary.

I could mask waters with oceanmask but couldn't mask those points in Mexico and Canada.

Note: There is a solution to loop over all grid points and set zi to nan for all those points that are not in the U.S. but that is so computationally expensive given my actual data size that is not a solution for me.

Answer 1

I ended up drawing Mexico and Canada polygons again over the contourf. Downloaded Countries shape file from here. The code is like:

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, cm, maskoceans
from scipy.interpolate import griddata
from matplotlib.patches import Polygon as MplPolygon
import shapefile
import pdb
np.random.seed(77)



lllat = 24.396308
lllon = -124.848974
urlat =  49.384358
urlon = -66.885444
m = Basemap(llcrnrlat=lllat,
urcrnrlat=urlat,
llcrnrlon=lllon,
urcrnrlon=urlon,
resolution='i', projection='cyl')
m.drawcountries(linewidth=1.0)
m.drawstates(linewidth=1, color='lightgray')
m.drawcoastlines()
m.drawlsmask(land_color='gray',ocean_color="#b0c4de", lakes=True)

#create 100 random latitudes
lats = np.random.randint(low=lllat-1, high=urlat+1, size=1000) + np.random.ranf(size=1000)
#create 100 random longitudes
lons = np.random.randint(low=lllon-1, high=urlon+1, size=1000) + np.random.ranf(size=1000)
#create 100 random values/probabilities
probabilities = np.random.random(size=1000)


#now use meshgrid and contourf to visualize it
mlon, mlat = m(*(lons, lats))
# grid data
numcols, numrows = 1000, 1000
xi = np.linspace(mlon.min(), mlon.max(), numcols)
yi = np.linspace(mlat.min(), mlat.max(), numrows)

xi, yi = np.meshgrid(xi, yi)
# interpolate
x, y, z = mlon, mlat, probabilities
zi = griddata((mlon, mlat), probabilities, (xi, yi), method='nearest', rescale=False)
data = maskoceans(xi, yi, zi)
con = m.contourf(xi, yi, data, cmap=plt.get_cmap('YlOrRd'))
cbar = m.colorbar(con,location='right',pad="3%")

world_shp_info = m.readshapefile('./data/CNTR_2014_10M_SH/Data/CNTR_RG_10M_2014','world',drawbounds=False)

ax = plt.gca()
for shapedict,state in zip(m.world_info, m.world):
    if shapedict['CNTR_ID'] not in ['CA', 'MX']: continue
    poly = MplPolygon(state,facecolor='gray',edgecolor='gray')
    ax.add_patch(poly)
plt.show()

and this is the result .

I think there should be a smarter way of doing it.