#!/usr/bin/env python
# coding: utf-8

# > This is one of the 100 recipes of the [IPython Cookbook](http://ipython-books.github.io/), the definitive guide to high-performance scientific computing and data science in Python.
# 

# # 14.6. Manipulating geospatial data with Shapely and basemap

# In order to run this recipe, you will need the following packages:
# 
# * [GDAL/OGR](http://www.gdal.org/ogr/)
# * [fiona](http://toblerity.org/fiona/README.html)
# * [Shapely](http://toblerity.org/shapely/project.html)
# * [descartes](https://pypi.python.org/pypi/descartes)
# * [basemap](http://matplotlib.org/basemap/)
# 
# On Windows, you can find binary installers for all those packages except descartes on Chris Gohlke's webpage. (http://www.lfd.uci.edu/~gohlke/pythonlibs/)
# 
# Installing descartes is easy: `pip install descartes`.
# 
# On other systems, you can find installation instructions on the projects' websites. GDAL/OGR is a C++ library that is required by fiona. The other packages are regular Python packages.
# 
# Finally, you need to download the *Africa* dataset on the book's website. (http://ipython-books.github.io)

# 1. Let's import all the packages.

# In[ ]:


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as col
from mpl_toolkits.basemap import Basemap
import fiona
import shapely.geometry as geom
from descartes import PolygonPatch
get_ipython().run_line_magic('matplotlib', 'inline')


# 2. We load the Shapefile dataset with fiona. This dataset notably contains the borders of all countries in the world.

# In[ ]:


# natural earth data
countries = fiona.open("data/ne_10m_admin_0_countries.shp")


# 3. We select the countries in Africa.

# In[ ]:


africa = [c for c in countries 
          if c['properties']['CONTINENT'] == 'Africa']


# 4. Now, we create a basemap map showing the African continent.

# In[ ]:


m = Basemap(llcrnrlon=-23.03,
            llcrnrlat=-37.72,
            urcrnrlon=55.20,
            urcrnrlat=40.58)


# 5. We need to convert the geographical coordinates of the countries' borders in map coordinates, so that we can display then in basemap.

# In[ ]:


def _convert(poly, m):
    if isinstance(poly, list):
        return [_convert(_, m) for _ in poly]
    elif isinstance(poly, tuple):
        return m(*poly)


# In[ ]:


for _ in africa:
    _['geometry']['coordinates'] = _convert(
        _['geometry']['coordinates'], m)


# 6. The next step is to create matplotlib `PatchCollection` objects from the Shapefile dataset loaded with fiona. We use Shapely and descartes to do this.

# In[ ]:


def get_patch(shape, **kwargs):
    """Return a matplotlib PatchCollection from a geometry 
    object loaded with fiona."""
    # Simple polygon.
    if isinstance(shape, geom.Polygon):
        return col.PatchCollection([PolygonPatch(shape, **kwargs)],
                                   match_original=True)
    # Collection of polygons.
    elif isinstance(shape, geom.MultiPolygon):
        return col.PatchCollection([PolygonPatch(c, **kwargs)
                                    for c in shape],
                                   match_original=True)


# In[ ]:


def get_patches(shapes, fc=None, ec=None, **kwargs):
    """Return a list of matplotlib PatchCollection objects
    from a Shapefile dataset."""
    # fc and ec are the face and edge colors of the countries.
    # We ensure these are lists of colors, with one element
    # per country.
    if not isinstance(fc, list):
        fc = [fc] * len(shapes)
    if not isinstance(ec, list):
        ec = [ec] * len(shapes)
    # We convert each polygon to a matplotlib PatchCollection
    # object.
    return [get_patch(geom.shape(s['geometry']), 
                      fc=fc_, ec=ec_, **kwargs) 
            for s, fc_, ec_ in zip(shapes, fc, ec)]


# 7. We also define a function to get countries colors depending on a specific field in the Shapefile dataset. Indeed, our dataset not only contains countries borders, but also a few administrative, economical and geographical properties for each country. Here, we will choose the color according to the population and GDP of the countries.

# In[ ]:


def get_colors(field, cmap):
    """Return one color per country, depending on a specific
    field in the dataset."""
    values = [country['properties'][field] for country in africa]
    values_max = max(values)
    return [cmap(v / values_max) for v in values]


# 8. Finally, we display the maps. We display the coast lines with basemap, and the countries with our Shapefile dataset.

# In[ ]:


plt.figure(figsize=(8,6));
# Display the countries color-coded with their population.
ax = plt.subplot(121);
m.drawcoastlines();
patches = get_patches(africa, 
                      fc=get_colors('POP_EST', plt.cm.Reds), 
                      ec='k')
for p in patches:
    ax.add_collection(p)
plt.title("Population");
# Display the countries color-coded with their population.
ax = plt.subplot(122);
m.drawcoastlines();
patches = get_patches(africa, 
                      fc=get_colors('GDP_MD_EST', plt.cm.Blues), 
                      ec='k')
for p in patches:
    ax.add_collection(p)
plt.title("GDP");


# > You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).
# 
# > [IPython Cookbook](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages).