world-ecoregion/draw.py

import fire
import matplotlib.pyplot as plt
from utils import logger
from constants import BIOMES

import pandas as pd
import cartopy.crs as ccrs

def draw(df, path=None):
    logger.debug('draw(df, %s)', path)
    biomes = {}
    biome_numbers = df['biome_num'].unique()

    for i, row in df.iterrows():
        if row.biome_num in biomes:
            biomes[row.biome_num]['x'].append(row.longitude)
            biomes[row.biome_num]['y'].append(row.latitude)
        else:
            biomes[row.biome_num] = { 'x': [row.longitude], 'y': [row.latitude] }

    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.stock_img()


    for n in biome_numbers:
        xs = biomes[n]['x']
        ys = biomes[n]['y']
        scatter = ax.scatter(xs, ys, s=4, c=BIOMES[n]['color'], transform=ccrs.PlateCarree())
        scatter.set_label(BIOMES[n]['name'])

    ax.legend()
    figure = plt.gcf()
    figure.set_size_inches(20, 18)
    if path:
        plt.savefig(path)
    else:
        plt.show()

def draw_cmd(path=None):
    draw(pd.read_pickle('data.p'), path=path)

if __name__ == "__main__":
    fire.Fire(draw_cmd)