How to use set clipped path for Basemap polygon

Question

I want to use imshow (for example) to display some data inside the boundaries of a country (for the purposes of example I chose the USA) The simple example below illustrates what I want:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import RegularPolygon

data = np.arange(100).reshape(10, 10)
fig = plt.figure()
ax = fig.add_subplot(111)
im = ax.imshow(data)
poly = RegularPolygon([ 0.5,  0.5], 6, 0.4, fc='none', 
                      ec='k', transform=ax.transAxes)
im.set_clip_path(poly)
ax.add_patch(poly)
ax.axis('off')
plt.show()

The result is:

Now I want to do this but instead of a simple polygon, I want to use the complex shape of the USA. I have created some example data contained in the array of "Z" as can be seen in the code below. It is this data that I want to display, using a colourmap but only within the boundaries of mainland USA.

So far I have tried the following. I get a shape file from here contained in "nationp010g.shp.tar.gz" and I use the Basemap module in python to plot the USA. Note that this is the only method I have found which gives me the ability get a polygon of the area I need. If there are alternative methods I would also be interested in them. I then create a polygon called "mainpoly" which is almost the polygon I want coloured in blue:

Notice how only one body has been coloured, all other disjoint polygons remain white:

So the area coloured blue is almost what I want, note that there are unwanted borderlines near canada because the border actually goes through some lakes, but that is a minor problem. The real problem is, why doesn't my imshow data display inside the USA? Comparing my first and second example codes I can't see why I don't get a clipped imshow in my second example, the way I do in the first. Any help would be appreciated in understanding what I am missing.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap as Basemap
from matplotlib.patches import Polygon

# Lambert Conformal map of lower 48 states.
m = Basemap(llcrnrlon=-119,llcrnrlat=22,urcrnrlon=-64,urcrnrlat=49,
            projection='lcc',lat_1=33,lat_2=45,lon_0=-95)


shp_info = m.readshapefile('nationp010g/nationp010g', 'borders', drawbounds=True) # draw     country boundaries.

for nshape,seg in enumerate(m.borders):
    if nshape == 1873: #This nshape denotes the large continental body of the USA, which we want
        mainseg = seg
        mainpoly =  Polygon(mainseg,facecolor='blue',edgecolor='k')



nx, ny = 10, 10
lons, lats = m.makegrid(nx, ny) # get lat/lons of ny by nx evenly space grid.
x, y = m(lons, lats) # compute map proj coordinates.

Z = np.zeros((nx,ny))
Z[:] = np.NAN

for i in np.arange(len(x)):
    for j in np.arange(len(y)):
        Z[i,j] = x[0,i] 

ax = plt.gca()
im = ax.imshow(Z, cmap = plt.get_cmap('coolwarm') )
im.set_clip_path(mainpoly)
ax.add_patch(mainpoly)
plt.show()

Update

I realise that the line

ax.add_patch(mainpoly)

does not even add the polygon shape to a plot. Am I not using it correctly? As far as I know mainpoly was calculated correctly using the Polygon() method. I checked that the coordinate inputs are a sensible:

plt.plot(mainseg[:,0], mainseg[:,1] ,'.') 

which gives

Answer 1

I have also considered about this problem for so long.
And I found NCL language has the function to mask the data outside some border.
Here is the example:

http://i5.tietuku.com/bdb1a6c007b82645.png

The contourf plot only show within China border. Click here for the code.

I know python has a package called PyNCL which support all NCL code in Python framework.
But I really want to plot this kind of figure using basemap. If you have figured it out, please post on the internet. I'll learn at the first time.

Thanks!

Add 2016-01-16

In a way, I have figured it out.
This is my idea and code, and it's inspired from this question I have asked today.

My method:
1. Make the shapefile of the interesting area(like U.S) into shapely.polygon.
2. Test each value point within/out of the polygon.
3. If the value point is out of the study area, mask it as np.nan

Intro * the polygon xxx was a city in China in ESRI shapefile format. * fiona, shapely package were used here.

# generate the shapely.polygon
shape = fiona.open("xxx.shp")
pol = shape.next()
geom = shape(pol['geometry'])
poly_data = pol["geometry"]["coordinates"][0]
poly = Polygon(poly_data)

It shows like:

http://i4.tietuku.com/2012307faec02634.png

### test the value point 
### generate the grid network which represented by the grid midpoints.
lon_med  = np.linspace((xi[0:2].mean()),(xi[-2:].mean()),len(x_grid))
lat_med  = np.linspace((yi[0:2].mean()),(yi[-2:].mean()),len(y_grid))

value_test_mean = dsu.mean(axis = 0)
value_mask =  np.zeros(len(lon_med)*len(lat_med)).reshape(len(lat_med),len(lon_med))
for i in range(0,len(lat_med),1):
    for j in range(0,len(lon_med),1):
        points = np.array([lon_med[j],lat_med[i]])
        mask = np.array([poly.contains(Point(points[0], points[1]))])
        if mask == False:
            value_mask[i,j] = np.nan
        if mask == True:
            value_mask[i,j] = value_test_mean[i,j]


# Mask the np.nan value 
Z_mask = np.ma.masked_where(np.isnan(so2_mask),so2_mask)

# plot!
fig=plt.figure(figsize=(6,4))
ax=plt.subplot()

map = Basemap(llcrnrlon=x_map1,llcrnrlat=y_map1,urcrnrlon=x_map2,urcrnrlat=y_map2)
map.drawparallels(np.arange(y_map1+0.1035,y_map2,0.2),labels=  [1,0,0,1],size=14,linewidth=0,color= '#FFFFFF')
lon_grid  = np.linspace(x_map1,x_map2,len(x_grid))
lat_grid  = np.linspace(y_map1,y_map2,len(y_grid))
xx,yy = np.meshgrid(lon_grid,lat_grid)
pcol =plt.pcolor(xx,yy,Z_mask,cmap = plt.cm.Spectral_r ,alpha =0.75,zorder =2)

result

http://i4.tietuku.com/c6620c5b6730a5f0.png

http://i4.tietuku.com/a22ad484fee627b9.png

original result

http://i4.tietuku.com/011584fbc36222c9.png