Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
137 lines (112 sloc) 4.65 KB
'''
countourf_to_shp.py
an example for converting the results of a matplotlib contourf plot to a
shapefile, preserving contour level data
to run:
python contourf_to_shp
outputs data in data/shaped_contour, plots a figure
Copyright (C) 2019 Chris Havlin, <https://chrishavlin.wordpress.com>
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
'''
from shapely import geometry
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
import fiona
import os,json
from descartes.patch import PolygonPatch
# create some test data with multiple peaks
lon = np.linspace(0,45,100)
lat = np.linspace(-20,32,90)
long,latg=np.meshgrid(lon,lat)
C1=np.sqrt((long-5.)**2+(latg-25)**2)/30.
C2=np.sqrt((long-30.)**2+(latg-1)**2)/10.
m = 30*np.exp(-C1**2)+20.*np.exp(-C2**2)
# make the contourf plot, storing the resulting ContourSet in cs
plt.figure(figsize=[10,5])
plt.subplot(1,2,1)
Nlevels=10
cs = plt.contourf(lon,lat,m,Nlevels,cmap='gist_heat')
plt.title('contourf figure with Nlevels='+str(Nlevels))
# create lookup table for levels
lvl_lookup = dict(zip(cs.collections, cs.levels))
# loop over collections (and polygons in each collection), store in list for fiona
PolyList=[]
for col in cs.collections:
z=lvl_lookup[col] # the value of this level
for contour_path in col.get_paths():
# create the polygon for this level
for ncp,cp in enumerate(contour_path.to_polygons()):
lons = cp[:,0]
lats = cp[:,1]
new_shape = geometry.Polygon([(i[0], i[1]) for i in zip(lons,lats)])
if ncp == 0:
poly = new_shape # first shape
else:
poly = poly.difference(new_shape) # Remove the holes
PolyList.append({'poly':poly,'props':{'z': z}})
## write the fiona collection
# clean up directories
outname=os.path.join('..','data','shaped_contour')
if os.path.isdir(outname):
for file in os.listdir(outname):
os.remove(os.path.join(outname,file))
os.rmdir(outname)
os.mkdir(outname)
# define ESRI schema, write each polygon to the file
outfi=os.path.join(outname,'shaped_contour.shp')
schema = {'geometry': 'Polygon','properties': {'z': 'float'}}
with fiona.collection(outfi, "w", "ESRI Shapefile", schema) as output:
for p in PolyList:
output.write({'properties': p['props'],
'geometry': geometry.mapping(p['poly'])})
# save the levels and global min/max as a separate json for convenience
Lvls={'levels':cs.levels.tolist(),'min':m.min(),'max':m.max()}
with open(os.path.join(outname,'levels.json'), 'w') as fp:
json.dump(Lvls, fp)
## Plotting the results: reads data back in, plots the polygons with data only
## from shapefile and levels.txt
ax=plt.subplot(1,2,2)
# read in levels, define colormap
with open(os.path.join(outname,'levels.json')) as jfile:
Lvls=json.load(jfile)
levels=np.array(Lvls['levels'])
cmap=plt.cm.gist_heat
lv_range=[Lvls['min'],Lvls['max']]
# loop over each shape, pull out level value ('z'), plot a polygon with a color
# matching the colormap.
ishp=0
with fiona.open(outfi) as shape:
for shp in shape:
# pull this shape's level and set the color from the map
lv=shp['properties']['z'] # this shape's level
clr=cmap((lv - lv_range[0])/(lv_range[1]-lv_range[0]))
# build the polygon and add the patch
coords=shp['geometry']['coordinates'][0] # coords of this polygon
poly=geometry.Polygon(coords)
patch = PolygonPatch(poly, facecolor=clr, edgecolor=clr)
ax.add_patch(patch)
# track max/min coordinate values
bnds=poly.bounds
rng_C={'lon':{'min':bnds[0],'max':bnds[2]},
'lat':{'min':bnds[1],'max':bnds[3]}}
if ishp==0:
rngs=rng_C
else:
for ll in ['lon','lat']:
rngs[ll]['max']=max([rngs[ll]['max'],rng_C[ll]['max']])
rngs[ll]['min']=min([rngs[ll]['min'],rng_C[ll]['min']])
ishp=ishp+1
ax.set_xlim([rngs['lon']['min'],rngs['lon']['max']])
ax.set_ylim([rngs['lat']['min'],rngs['lat']['max']])
plt.title('polygon patches from shapefile')
plt.show()
You can’t perform that action at this time.