Search code examples
pythonanimationmatplotlibmatplotlib-basemapcartopy

Animation Using Cartopy


I'm trying to write a python program that displays an animation of a map of the world where countries change color based on how much renewable energy use they have. I'm trying to have it display the colors for all countries in year 1960, then the colors for all countries in the year 1961, then 1962...

I'm using cartopy to add countries to the figure and basing their color off of values that I pull into a pandas dataframe from a SQL database. I was able to get the map to show what I want for one year like this: enter image description here

However, I can't figure out how to animate it. I tried using FuncAnimate, but I'm really struggling to understand how it works. All the examples seem to have functions that return lines, but I'm not graphing lines or contours. Here is what I tried:

import sqlite3
import pandas as pd
import os
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.animation as animation
import cartopy.crs as ccrs
import cartopy.io.shapereader as shpreader
from math import log
from math import exp
from matplotlib import colors
path = 'H:/USER/DVanLunen/indicator_data/world-development-indicators/'
os.chdir(path)
con = sqlite3.connect('database.sqlite')

# Grab :
# % of electricity from renewable sources EG.ELC.RNWX.ZS
# 1960 - 2013

Indicator_df = pd.read_sql('SELECT * '
                           'FROM Indicators '
                           'WHERE IndicatorCode in('
                                                   '"EG.ELC.RNWX.ZS"'
                                                   ')'
                           , con)
# setup colorbar stuff and shape files
norm = mpl.colors.Normalize(vmin=0, vmax=30)
colors_in_map = []
for i in range(30):
    val = log(i + 1, logbase) / log(31, logbase)
    colors_in_map.append((1 - val, val, 0))
cmap = colors.ListedColormap(colors_in_map)

shpfilename = shpreader.natural_earth(resolution='110m',
                                      category='cultural',
                                      name='admin_0_countries')
reader = shpreader.Reader(shpfilename)
countries_map = reader.records()
logbase = exp(1)
fig, ax = plt.subplots(figsize=(12, 6),
                       subplot_kw={'projection': ccrs.PlateCarree()})


def run(data):
    """Update the Dist"""
    year = 1960 + data % 54
    logbase = exp(1)
    for n, country in enumerate(countries_map):
        facecolor = 'gray'
        edgecolor = 'black'
        indval = Indicator_df.loc[(Indicator_df['CountryName'] ==
                                   country.attributes['name_long']) &
                                  (Indicator_df['Year'] == year), 'Value']
        if indval.any():
                greenamount = (log(float(indval) + 1, logbase) /
                               log(31, logbase))
                facecolor = 1 - greenamount, greenamount, 0
        ax.add_geometries(country.geometry, ccrs.PlateCarree(),
                          facecolor=facecolor, edgecolor=edgecolor)
        ax.set_title('Percent of Electricity from Renewable Sources ' +
                     str(year))
    ax.figure.canvas.draw()

cax = fig.add_axes([0.92, 0.2, 0.02, 0.6])
cb = mpl.colorbar.ColorbarBase(cax, cmap=cmap, norm=norm,
                               spacing='proportional')
cb.set_label('%')

ani = animation.FuncAnimation(fig, run, interval=200, blit=False)
plt.show()

Any help would be greatly appreciated. Thanks!

Some example data for Indicator_df (not real):

CountryName     Year     Value
United States     1960     5
United States     1961     10
United States     1962     20
United States     1963     30

Solution

  • There are actually several problems with how you've set up your run(), but the major problem appeared to actually be the enumate(countries_map). The records() function returns a generator, which once you've run through it once doesn't appear to like being run through again - I tried it separate from the animation to make sure.

    That said, the problem can be avoided entirely by moving a lot of code out of the run(). Currently, even if it worked you're re-drawing every single country every frame, not just the ones with colors. It's both intensive and unnecessary - you don't need to draw any gray ones more than once.

    I've restructured your code a bit and with the fake data I put in for the US and Argentina it works fine for me.

    import pandas as pd
    import matplotlib.pyplot as plt
    import matplotlib as mpl
    import matplotlib.animation as animation
    import cartopy.crs as ccrs
    import cartopy.io.shapereader as shpreader
    from math import log
    from math import exp
    from matplotlib import colors
    from shapely.geometry.multipolygon import MultiPolygon
    
    # Grab :
    # % of electricity from renewable sources EG.ELC.RNWX.ZS
    # 1960 - 2013
    
    # Make fake data
    Indicator_df = pd.DataFrame({
        'CountryName': ['United States'] * 4 + ['Argentina'] * 4,
        'Year': [1960, 1961, 1962, 1963] * 2,
        'Value': [5, 10, 20, 30] * 2
    })
    
    # setup colorbar stuff and shape files
    norm = mpl.colors.Normalize(vmin=0, vmax=30)
    colors_in_map = []
    logbase = exp(1)
    for i in range(30):
        val = log(i + 1, logbase) / log(31, logbase)
        colors_in_map.append((1 - val, val, 0))
    cmap = colors.ListedColormap(colors_in_map)
    
    shpfilename = shpreader.natural_earth(resolution='110m',
                                          category='cultural',
                                          name='admin_0_countries')
    reader = shpreader.Reader(shpfilename)
    countries_map = reader.records()
    
    # These don't need to constantly be redefined, especially edgecolor
    facecolor = 'gray'
    edgecolor = 'black'
    
    fig, ax = plt.subplots(figsize=(12, 6),
                           subplot_kw={'projection': ccrs.PlateCarree()})
    
    # Draw all the gray countries just once in an init function
    # I also make a dictionary for easy lookup of the geometries by country name later
    geom_dict = {}
    
    
    def init_run():
        for n, country in enumerate(countries_map):
            if country.geometry.type == "Polygon":
                geom = MultiPolygon([country.geometry])
            else:
                geom = country.geometry
            ax.add_geometries(geom,
                              ccrs.PlateCarree(),
                              facecolor=facecolor,
                              edgecolor=edgecolor)
            geom_dict[country.attributes['NAME_LONG']] = geom
    
    
    def run(data):
        """Update the Dist"""
        # "data" in this setup is a frame number starting from 0, so it corresponds nicely
        # with your years
        # data = 0
        year = 1960 + data
    
        # get a subset of the df for the current year
        year_df = Indicator_df[Indicator_df['Year'] == year]
        for i, row in year_df.iterrows():
            # This loops over countries, gets the value and geometry and adds
            # the new-colored shape
            geom = geom_dict[row['CountryName']]
            value = row['Value']
            greenamount = (log(float(value) + 1, logbase) / log(31, logbase))
            facecolor = 1 - greenamount, greenamount, 0
            ax.add_geometries(geom,
                              ccrs.PlateCarree(),
                              facecolor=facecolor,
                              edgecolor=edgecolor)
    
        # I decreased the indent of this, you only need to do it once per call to run()
        ax.set_title('Percent of Electricity from Renewable Sources ' + str(year))
    
    
    cax = fig.add_axes([0.92, 0.2, 0.02, 0.6])
    cb = mpl.colorbar.ColorbarBase(cax,
                                   cmap=cmap,
                                   norm=norm,
                                   spacing='proportional')
    cb.set_label('%')
    
    ani = animation.FuncAnimation(fig,
                                  run,
                                  init_func=init_run,
                                  frames=4,
                                  interval=500,
                                  blit=False)
    ani.save(filename="test.gif")
    

    The primary difference is that I'm not accessing the shpreader at all inside the run function. When making an animation, the only thing that should be in the run function are things that change, you don't need to re-draw everything every frame.

    That said, this could be even better if you just keep the artist from the very first draw and just change the color of it in the run function, instead of doing a whole new ax.add_geometries. You'll have to look into how to change the color of a cartopy FeatureArtist for that.