Plot a set of 2d figures

plot2d.py

import matplotlib.pyplot as plt
%matplotlib inline


def plot2d(fig_lst, labels=None, figsize=(10, 10), savepath=None):
    """Plots a set of 2d shapes represented as 2d point arrays
    
    Parameters
    ----------
    fig_lst: list
        list of 2d numpy arrays [n_points, 2], each containing points 
        to plot as a separate shape
    labels: list or None
        list of labels for every figure, should be of the same size as x_arr
    figsise: tuple
        size of the figure
    savepath: str or None
        path to save figure. If None, the resulting plot will be only shown
    
    """
    plt.figure(figsize=figsize)
    plt.rc('font', family='serif', size=18) #nice font for LaTeX
    
    
    if labels is not None:
        for ix, fig in enumerate(fig_lst):
            plt.scatter(fig[:,0], fig[:, 1], s=4, label=labels[ix])
            plt.legend()
    else:
        for ix, fig in enumerate(fig_lst):
            plt.scatter(fig[:,0], fig[:, 1], s=4)
    
    if savepath is not None:
        plt.savefig(savepath, dpi=200)
    
    return

# Demo

import numpy as np
x1 = np.random.normal(size=[100, 2])
x2 = np.random.uniform(size=[100, 2])

plot2d([x1, x2], ['normal samples', 'uniform'])
plot2d([x1, x2], ['normal samples', 'uniform'], savepath='../logs/test.png')