world-ecoregion/biomes/draw.py

import fire
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Circle, Patch
from utils import logger, to_range
from constants import BIOMES

import pandas as pd
import cartopy.crs as ccrs

def draw(df, earth=True, width=23.22, height=13, longitude_max=None, latitude_max=None, only_draw=False, path=None, alpha=1):
    logger.debug('draw(df, %s)', path)
    biomes = {}
    biome_numbers = df['biome_num'].unique()

    for i, row in df.iterrows():
        if earth:
            p = (row.longitude, row.latitude)
        else:
            p = (to_range(-180, 180, 0, longitude_max)(row.longitude), to_range(-90, 90, 0, latitude_max)(row.latitude))

        if row.biome_num in biomes:
            biomes[row.biome_num].append(p)
        else:
            biomes[row.biome_num] = [p]

    if earth:
        ax = plt.axes(projection=ccrs.PlateCarree())
        ax.stock_img()
    else:
        ax = plt.gca()

    legend_handles = []
    for n in biome_numbers:
        color = BIOMES[n]['color']

        patches = [Circle(p, radius=0.4) for p in biomes[n]]
        collection = PatchCollection(patches, color=color)
        collection.set_alpha(alpha)

        legend_handles.append(Patch(color=color, label=BIOMES[n]['name']))
        ax.add_collection(collection)

    ax.legend(handles=legend_handles, loc='center left', bbox_to_anchor=(1, 0.5), markerscale=4)

    ax.autoscale_view()
    figure = plt.gcf()
    figure.set_size_inches(width, height)
    figure.subplots_adjust(left=0.02, right=0.79)

    if only_draw: return

    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)