高校生っぽいグラフ

highschoollike_plot.py

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns


def decorate_arrows(plotting):
    def inner(x, y,figsize=(6,4), delete_labels=True, fontdict='default',xlim='default',ylim='default'):
        fig, ax = plotting(x, y, figsize)
        incs = fig.get_size_inches()
        incs = np.array(incs)
        incs /= np.linalg.norm(incs)
        for key, sp in ax.spines.items():
            sp.set_color('none')

        for seter,lims in zip([ax.set_xlim,ax.set_ylim],[xlim,ylim]):
            if lims == 'default':
                continue
            seter(*lims)

        xl = ax.set_xlim()
        yl = ax.set_ylim()

        x_length = (xl[1] - xl[0])
        y_length = (yl[1] - yl[0])

        print(x_length, y_length)

        xarrow_params = {'width': 0.004 * y_length * incs[0],
                         'color': '0',
                         'clip_on': False,
                         'head_width': 0.05 * y_length * incs[0],
                         'head_length': 0.05 * x_length * incs[1]}

        yarrow_params = {'width': 0.004 * x_length * incs[1],
                         'color': '0',
                         'clip_on': False,
                         'head_width': 0.05 * x_length * incs[1],
                         'head_length': 0.05 * y_length * incs[0]}

        ax.arrow(xl[0], 0, x_length, 0, **xarrow_params)
        ax.arrow(0, yl[0], 0, y_length, **yarrow_params)

        if fontdict == 'default' or fontdict is None:
            fd = {'size': 20}
        else:
            fd = fontdict

        ax.text(x_length / 20. * incs[0], yl[1],
                '$y$', ha='center', fontdict=fd)
        ax.text(xl[1], y_length / 20. * incs[1],
                '$x$', ha='right', fontdict=fd)
        ax.text(- x_length / 20. * incs[0], - y_length / 20. * incs[1],
                '${\\rm O}$', va='top', ha='right', fontdict=fd)

        if delete_labels:
            ax.xaxis.set_ticklabels('')
            ax.yaxis.set_ticklabels('')

        fig.tight_layout()
        return fig, ax
    return inner


@decorate_arrows
def highschoollike_plot(x, y, figsize=(6,5)):
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111)
    ax.plot(x, y)
    return fig, ax

if __name__ == '__main__':
    sns.set_style('white')
    x = np.linspace(-5, 10, 1000)
    fig, ax = highschoollike_plot(
        x, (2 * x - 1) / (x**2 + 2) * .2, fontdict={'size': 15},ylim=[-.21,.11])
    fig.savefig('test1.png')

    t = np.linspace(0, 2 * np.pi, 1000)
    x = (1 + np.cos(t)) * np.cos(t)
    y = (1 + np.cos(t)) * np.sin(t)
    fig, ax = highschoollike_plot(x, y, figsize=(4,4), fontdict='default',xlim=[-.5,2.5],ylim=[-1.5,1.5])
    fig.savefig('test2.png')